# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation.  All rights reserved.
# Licensed under the MIT License.  See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import onnx_ir as ir
import torch
from pathlib import Path

from .base import Model


class GPTOSSModel(Model):
    def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
        super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
        self.is_local = lambda layer_id: bool(layer_id % 2 == 0)
        self.attention_attrs["sinks"] = True

        self.moe_attrs["activation_alpha"] = 1.702
        self.moe_attrs["activation_beta"] = 1.0
        self.moe_attrs["activation_type"] = "swiglu"
        self.moe_attrs["normalize_routing_weights"] = True
        self.moe_attrs["swiglu_fusion"] = 1
        self._mxfp4_weight_map = None
        self._mxfp4_weight_map_loaded = False
        self._mxfp4_snapshot_dir = None
        self._mxfp4_tensor_cache = {}

    def _get_mxfp4_snapshot_dir(self):
        if self._mxfp4_snapshot_dir is not None:
            return self._mxfp4_snapshot_dir

        model_path = Path(self.model_name_or_path)
        if model_path.is_dir():
            self._mxfp4_snapshot_dir = model_path
            return self._mxfp4_snapshot_dir

        try:
            from huggingface_hub import snapshot_download
        except ImportError as e:
            raise RuntimeError("huggingface_hub is required to locate original GPT-OSS MXFP4 weights.") from e

        self._mxfp4_snapshot_dir = Path(
            snapshot_download(
                self.model_name_or_path,
                cache_dir=self.cache_dir,
                token=self.hf_token,
                local_files_only=True,
            )
        )
        return self._mxfp4_snapshot_dir

    def _get_mxfp4_weight_map(self):
        if self._mxfp4_weight_map_loaded:
            return self._mxfp4_weight_map

        snapshot_dir = self._get_mxfp4_snapshot_dir()
        index_path = snapshot_dir / "model.safetensors.index.json"
        if index_path.exists():
            import json

            with open(index_path) as f:
                self._mxfp4_weight_map = json.load(f)["weight_map"]
        else:
            safetensors_files = sorted(snapshot_dir.glob("*.safetensors"))
            if len(safetensors_files) != 1:
                raise RuntimeError(
                    f"Could not locate original GPT-OSS MXFP4 safetensors index in {snapshot_dir}."
                )
            self._mxfp4_weight_map = None
        self._mxfp4_weight_map_loaded = True
        return self._mxfp4_weight_map

    def _load_original_mxfp4_tensor(self, tensor_name):
        if tensor_name in self._mxfp4_tensor_cache:
            return self._mxfp4_tensor_cache[tensor_name]

        snapshot_dir = self._get_mxfp4_snapshot_dir()
        weight_map = self._get_mxfp4_weight_map()

        if weight_map is None:
            candidate_files = sorted(snapshot_dir.glob("*.safetensors"))
        else:
            if tensor_name not in weight_map:
                raise RuntimeError(f"Original GPT-OSS MXFP4 tensor '{tensor_name}' was not found in safetensors index.")
            candidate_files = [snapshot_dir / weight_map[tensor_name]]

        try:
            from safetensors import safe_open
        except ImportError as e:
            raise RuntimeError("safetensors is required to load original GPT-OSS MXFP4 weights.") from e

        for tensor_file in candidate_files:
            with safe_open(tensor_file, framework="pt", device="cpu") as f:
                if tensor_name in f.keys():
                    tensor = f.get_tensor(tensor_name)
                    self._mxfp4_tensor_cache[tensor_name] = tensor
                    return tensor

        raise RuntimeError(f"Original GPT-OSS MXFP4 tensor '{tensor_name}' was not found in {snapshot_dir}.")

    def pack_original_mxfp4_blocks_for_qmoe(self, blocks):
        """Pack GPT-OSS checkpoint MXFP4 blocks to QMoE's column-major layout.

        GPT-OSS stores blocks as [E, N, K/32, 16], with each byte holding two
        adjacent K-axis FP4 e2m1 codes for the same output row N. CUDA QMoE
        expects [E, K, N/2], with each byte holding two adjacent N-axis codes for
        the same K (even N in low nibble, odd N in high nibble).
        """
        if blocks.dtype != torch.uint8:
            blocks = blocks.to(torch.uint8)
        if blocks.ndim != 4 or blocks.shape[-1] != 16:
            raise ValueError(f"GPT-OSS MXFP4 blocks must have shape [E, N, K/32, 16], got {tuple(blocks.shape)}.")
        if blocks.shape[1] % 2 != 0:
            raise ValueError(f"GPT-OSS MXFP4 output dimension N must be even, got {blocks.shape[1]}.")

        even_n = blocks[:, 0::2, :, :]
        odd_n = blocks[:, 1::2, :, :]
        packed_even_k = ((odd_n & 0x0F) << 4) | (even_n & 0x0F)
        packed_odd_k = ((odd_n >> 4) << 4) | (even_n >> 4)
        packed = torch.stack((packed_even_k, packed_odd_k), dim=-1)
        return packed.permute(0, 2, 3, 4, 1).reshape(blocks.shape[0], blocks.shape[2] * 32, blocks.shape[1] // 2).contiguous()

    def make_original_mxfp4_qmoe_initializers(self, layer_id, proj, weight_name, scales_name, global_scales_name):
        blocks = self._load_original_mxfp4_tensor(f"model.layers.{layer_id}.mlp.experts.{proj}_blocks")
        scales = self._load_original_mxfp4_tensor(f"model.layers.{layer_id}.mlp.experts.{proj}_scales")
        if scales.dtype != torch.uint8:
            scales = scales.to(torch.uint8)
        expected_scale_shape = blocks.shape[:-1]
        if tuple(scales.shape) != tuple(expected_scale_shape):
            raise ValueError(
                f"GPT-OSS MXFP4 scales for layer {layer_id} {proj} must have shape {tuple(expected_scale_shape)}, "
                f"got {tuple(scales.shape)}."
            )

        self.make_initializer(self.pack_original_mxfp4_blocks_for_qmoe(blocks), weight_name)
        self.make_fp8e8m0_initializer(scales, scales_name)
        self.make_initializer(torch.ones(blocks.shape[0], dtype=torch.float32), global_scales_name)

    def make_layer(self, layer_id, layer):
        # Each LLM decoder layer is typically defined as:
        # input_layernorm --> attention --> output_layernorm --> MoE
        self.make_layernorm(
            layer_id,
            layer.input_layernorm,
            skip=not self.layernorm_attrs["first_layernorm"],
            simple=self.layernorm_attrs["simple"],
            location="input",
        )
        self.make_attention(layer_id, layer.self_attn, root_input=self.layernorm_attrs["output_0"])
        self.make_layernorm(
            layer_id,
            layer.post_attention_layernorm,
            skip=True,
            simple=self.layernorm_attrs["simple"],
            location="post_attention",
        )
        self.make_moe(layer_id, layer.mlp, root_input=self.layernorm_attrs["output_0"])

        self.layernorm_attrs["first_layernorm"] = False
        if layer_id == self.num_layers - 1:
            # Norm after last decoder layer of model (last layer --> norm)
            self.layernorm_attrs["last_layernorm"] = True

    def make_layernorm(self, layer_id, layernorm, skip, simple, location):
        if "final_norm" in location:
            # Set cast for final LayerNorm since it is a special case and not covered in `make_layer`
            self.layernorm_attrs["cast"]["root_input"] = False
        super().make_layernorm(layer_id, layernorm, skip, simple, location)

    def make_rotary_embedding_caches_from_scratch(self):
        inv_freq = 1.0 / (
            self.rope_attrs["theta"] ** (torch.arange(0, self.head_size, 2, dtype=torch.float) / self.head_size)
        )
        inv_freq = self.make_inv_freq_rescaled(inv_freq)

        t = torch.arange(self.rope_attrs["cache_length"], dtype=torch.float32)
        freqs = torch.einsum("i,j->ij", t, inv_freq)
        cos_cache, sin_cache = freqs.cos() * self.rope_attrs["mscale"], freqs.sin() * self.rope_attrs["mscale"]
        return cos_cache, sin_cache

    def make_attention(self, layer_id, attention, root_input, **kwargs):
        original_window_size = self.window_size
        self.window_size = (
            original_window_size if self.is_local(layer_id) else -1
        )  # default is -1 in GroupQueryAttention kernel
        super().make_attention(layer_id, attention, root_input, **kwargs)
        self.window_size = original_window_size

    def make_moe(self, layer_id, mlp, root_input):
        if self.ep in {"cpu", "cuda", "trt-rtx", "webgpu"}:
            self.make_moe_fused(layer_id, mlp, root_input)
        else:
            self.make_moe_decomposed(layer_id, mlp, root_input)

    def make_moe_decomposed(self, layer_id, mlp, root_input):
        # Make nodes for the MoE subgraph
        #
        #                                              root_input
        #                                                  |
        #          +---------------------------------------+
        #          |                                       |
        #      Unsqueeze                                 MatMul
        #       (dim=2)                                 (Router)
        #          |                                       |
        #       Expand                                    Add
        #          |                                       |
        #      Unsqueeze                                 Cast
        #      (dim=-1)                               (to = FP32)
        #          |                                       |
        #          |                                     TopK
        #          |                                       |
        #          |                                     Cast
        #          |                                 (to = io_dtype)
        #          |                                       |
        #          |           +-------------+-------------+-------------+-------------+
        #          |           |             |             |             |             |
        #          |         Gather        Gather        Gather        Gather       Softmax
        #          |      (MLP1 Weight)  (MLP1 Bias)  (MLP2 Weight)  (MLP2 Bias)   (Expert Weights)
        #          |           |             |             |             |             |
        #          +-----+-----+         Unsqueeze         |         Unsqueeze     Unsqueeze
        #                |               (dim=-1)          |         (dim=-1)      (dim=-1)
        #             MatMul                 |             |             |             |
        #          (gate_up_proj)            |             |             |         Unsqueeze
        #                |                   |             |             |         (dim=-1)
        #                +---------+---------+             |             |             |
        #                          |                       |             |           Cast
        #                         Add                      |             |        (to = FP32)
        #                          |                       |             |             |
        #                    +-----+-----+                 |             |             |
        #                    |           |                 |             |             |
        #                  Slice       Slice               |             |             |
        #                    |           |                 |             |             |
        #                  Clip        Clip                |             |             |
        #                    |  \        |                 |             |             |
        #                   Mul  \       |                 |             |             |
        #                    |    \      |                 |             |             |
        #                 Sigmoid  |     |                 |             |             |
        #                    |     |    Add                |             |             |
        #                    \    /      |                 |             |             |
        #                     \  /       |                 |             |             |
        #                      Mul       |                 |             |             |
        #                       |        |                 |             |             |
        #                       +----+---+                 |             |             |
        #                            |                     |             |             |
        #                           Mul                    |             |             |
        #                            |                     |             |             |
        #                            +----------+----------+             |             |
        #                                       |                        |             |
        #                                    MatMul                      |             |
        #                                  (down_proj)                   |             |
        #                                       |                        |             |
        #                                       +------------+-----------+             |
        #                                                    |                         |
        #                                                   Add                        |
        #                                                    |                         |
        #                                                   Cast                       |
        #                                               (to = FP32)                    |
        #                                                    |                         |
        #                                                    +------------+------------+
        #                                                                 |
        #                                                                Mul
        #                                                                 |
        #                                                             ReduceSum
        #                                                              (dim=2)
        #                                                                 |
        #                                                              Squeeze
        #                                                              (dim=-1)
        #                                                                 |
        #                                                                Cast
        #                                                           (to = io_dtype)
        basename = f"/model/layers.{layer_id}/moe"
        use_cast = self.io_dtype != ir.DataType.FLOAT

        # Make root_input expansion nodes (root_input --> Unsqueeze --> Expand --> Unsqueeze)
        expand_root_input_unsqueeze_1_name = f"{basename}/expand_root_input/Unsqueeze_1"
        expand_root_input_unsqueeze_1_inputs = [root_input, "/model/constants/INT64/[2]"]
        self.make_unsqueeze(
            expand_root_input_unsqueeze_1_name,
            expand_root_input_unsqueeze_1_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", 1, self.hidden_size],
        )
        expand_name = f"{basename}/expand_root_input/Expand"
        expand_inputs = [
            f"{expand_root_input_unsqueeze_1_name}/output_0",
            f"/model/constants/INT64/[1, 1, {self.moe_attrs['top_k']}, 1]",
        ]
        self.make_expand(
            expand_name,
            expand_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.hidden_size],
        )
        expand_root_input_unsqueeze_2_name = f"{basename}/expand_root_input/Unsqueeze_2"
        expand_root_input_unsqueeze_2_inputs = [f"{expand_name}/output_0", "/model/constants/INT64/[-1]"]
        self.make_unsqueeze(
            expand_root_input_unsqueeze_2_name,
            expand_root_input_unsqueeze_2_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.hidden_size, 1],
        )

        # Make router nodes
        #                                          +--> Gather --> (MLP1 weight)
        #                                          |
        #                                          +--> Gather --> Unsqueeze --> (MLP1 bias)
        #                                          |
        # root_input --> MatMul --> Add --> TopK --+--> Gather --> (MLP2 weight)
        #                                          |
        #                                          +--> Gather --> Unsqueeze --> (MLP2 bias)
        #                                          |
        #                                          +--> Softmax --> Unsqueeze --> Unsqueeze --> Cast
        #
        router_basename = f"{basename}/router/MatMul"
        router_matmul_name = self.make_matmul(mlp.router, router_basename, root_input)
        router_add_name = f"{basename}/router/Add"
        self.make_add_bias(mlp.router.bias, router_add_name, root_input=f"{router_matmul_name}/output_0")

        if use_cast:
            topk_fp32_name = f"{basename}/topk_fp32/Cast"
            self.make_cast(
                topk_fp32_name,
                f"{router_add_name}/output_0",
                ir.DataType.FLOAT,
                shape=["batch_size", "sequence_length", self.moe_attrs["num_experts"]],
            )
        topk_name = f"{basename}/TopK"
        topk_inputs = [
            f"{topk_fp32_name if use_cast else router_add_name}/output_0",
            f"/model/constants/INT64/[{self.moe_attrs['top_k']}]",
        ]
        topk_outputs = [f"{topk_name}/output_0", f"{topk_name}/output_1"]
        self.make_node(
            "TopK", inputs=topk_inputs, outputs=topk_outputs, name=topk_name, axis=-1, largest=True, sorted=True
        )
        self.make_value(
            topk_outputs[0], ir.DataType.FLOAT, shape=["batch_size", "sequence_length", self.moe_attrs["top_k"]]
        )
        self.make_value(
            topk_outputs[1], ir.DataType.INT64, shape=["batch_size", "sequence_length", self.moe_attrs["top_k"]]
        )
        if use_cast:
            topk_io_name = f"{basename}/topk_io/Cast"
            self.make_cast(
                topk_io_name,
                topk_outputs[0],
                self.io_dtype,
                shape=["batch_size", "sequence_length", self.moe_attrs["top_k"]],
            )

        # Save initializers to use with Gather nodes
        gate_up_proj_weight = f"model.layers.{layer_id}.moe.experts.gate_up_proj.weight"
        self.make_initializer(mlp.experts.gate_up_proj, gate_up_proj_weight, to=self.io_dtype)
        gate_up_proj_bias = f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias"
        self.make_initializer(mlp.experts.gate_up_proj_bias, gate_up_proj_bias, to=self.io_dtype)
        down_proj_weight = f"model.layers.{layer_id}.moe.experts.down_proj.weight"
        self.make_initializer(mlp.experts.down_proj, down_proj_weight, to=self.io_dtype)
        down_proj_bias = f"model.layers.{layer_id}.moe.experts.down_proj.bias"
        self.make_initializer(mlp.experts.down_proj_bias, down_proj_bias, to=self.io_dtype)

        # Make Gather nodes + Unsqueeze nodes for biases
        mlp1_weight_gather_name = f"{basename}/mlp1/weight/Gather"
        mlp1_weight_gather_inputs = [gate_up_proj_weight, f"{topk_name}/output_1"]
        self.make_gather(
            mlp1_weight_gather_name,
            mlp1_weight_gather_inputs,
            dtype=self.io_dtype,
            shape=[
                "batch_size",
                "sequence_length",
                self.moe_attrs["top_k"],
                2 * self.intermediate_size,
                self.hidden_size,
            ],
            axis=0,
        )
        mlp1_bias_gather_name = f"{basename}/mlp1/bias/Gather"
        mlp1_bias_gather_inputs = [gate_up_proj_bias, f"{topk_name}/output_1"]
        self.make_gather(
            mlp1_bias_gather_name,
            mlp1_bias_gather_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], 2 * self.intermediate_size],
            axis=0,
        )
        mlp1_bias_unsqueeze_name = f"{basename}/mlp1/bias/Unsqueeze"
        mlp1_bias_unsqueeze_inputs = [f"{mlp1_bias_gather_name}/output_0", "/model/constants/INT64/[-1]"]
        self.make_unsqueeze(
            mlp1_bias_unsqueeze_name,
            mlp1_bias_unsqueeze_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], 2 * self.intermediate_size, 1],
        )
        mlp2_weight_gather_name = f"{basename}/mlp2/weight/Gather"
        mlp2_weight_gather_inputs = [down_proj_weight, f"{topk_name}/output_1"]
        self.make_gather(
            mlp2_weight_gather_name,
            mlp2_weight_gather_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.hidden_size, self.intermediate_size],
            axis=0,
        )
        mlp2_bias_gather_name = f"{basename}/mlp2/bias/Gather"
        mlp2_bias_gather_inputs = [down_proj_bias, f"{topk_name}/output_1"]
        self.make_gather(
            mlp2_bias_gather_name,
            mlp2_bias_gather_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.hidden_size],
            axis=0,
        )
        mlp2_bias_unsqueeze_name = f"{basename}/mlp2/bias/Unsqueeze"
        mlp2_bias_unsqueeze_inputs = [f"{mlp2_bias_gather_name}/output_0", "/model/constants/INT64/[-1]"]
        self.make_unsqueeze(
            mlp2_bias_unsqueeze_name,
            mlp2_bias_unsqueeze_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.hidden_size, 1],
        )

        # Make expert_weights path (Softmax --> Unsqueeze --> Unsqueeze --> Cast)
        softmax_name = f"{basename}/expert_weights/Softmax"
        self.make_softmax(
            softmax_name,
            f"{topk_io_name if use_cast else topk_name}/output_0",
            self.io_dtype,
            ["batch_size", "sequence_length", "num_experts_per_token"],
        )
        expert_weights_unsqueeze_1_name = f"{basename}/expert_weights/Unsqueeze_1"
        expert_weights_unsqueeze_1_inputs = [f"{softmax_name}/output_0", "/model/constants/INT64/[-1]"]
        self.make_unsqueeze(
            expert_weights_unsqueeze_1_name,
            expert_weights_unsqueeze_1_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", "num_experts_per_token", 1],
        )
        expert_weights_unsqueeze_2_name = f"{basename}/expert_weights/Unsqueeze_2"
        expert_weights_unsqueeze_2_inputs = [
            f"{expert_weights_unsqueeze_1_name}/output_0",
            "/model/constants/INT64/[-1]",
        ]
        self.make_unsqueeze(
            expert_weights_unsqueeze_2_name,
            expert_weights_unsqueeze_2_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", "num_experts_per_token", 1, 1],
        )
        if use_cast:
            expert_weights_cast_name = f"{basename}/expert_weights/Cast"
            self.make_cast(
                expert_weights_cast_name,
                f"{expert_weights_unsqueeze_2_name}/output_0",
                ir.DataType.FLOAT,
                shape=["batch_size", "sequence_length", "num_experts_per_token", 1, 1],
            )

        # Make Gate/Up proj nodes (MatMul --> Add)
        gate_up_proj_weight_name = f"{basename}/gate_up_proj/MatMul"
        gate_up_proj_weight_output = f"{gate_up_proj_weight_name}/output_0"
        self.make_node(
            "MatMul",
            inputs=[f"{mlp1_weight_gather_name}/output_0", f"{expand_root_input_unsqueeze_2_name}/output_0"],
            outputs=[gate_up_proj_weight_output],
            name=gate_up_proj_weight_name,
        )
        self.make_value(
            gate_up_proj_weight_output,
            self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], 2 * self.intermediate_size, 1],
        )
        gate_up_proj_bias_name = f"{basename}/gate_up_proj/Add"
        self.make_add(
            gate_up_proj_bias_name,
            [gate_up_proj_weight_output, f"{mlp1_bias_unsqueeze_name}/output_0"],
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], 2 * self.intermediate_size, 1],
        )

        # Make activation nodes
        #
        #         +--> Slice --> Clamp --> Mul --> Sigmoid --+
        #         |      |                                   |
        #         |      |                                  Mul --+
        # Add ----+      |                                   |    |
        #         |      +-----------------------------------+    +--> Mul
        #         |                                               |
        #         +---> Slice --> Clamp --> Add ------------------+
        glu_slice_name = f"{basename}/act_fn/Slice_1"
        glu_slice_inputs = [
            f"{gate_up_proj_bias_name}/output_0",
            "/model/constants/INT64/[0]",
            f"/model/constants/INT64/[{torch.iinfo(torch.int64).max}]",
            "/model/constants/INT64/[3]",
            "/model/constants/INT64/[2]",
        ]
        self.make_slice(
            glu_slice_name,
            glu_slice_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        glu_clip_name = f"{basename}/act_fn/Clip_1"
        glu_clip_inputs = [
            f"{glu_slice_name}/output_0",
            "",
            f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{self.moe_attrs['swiglu_limit']}",
        ]
        self.make_clip(
            glu_clip_name,
            glu_clip_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        linear_slice_name = f"{basename}/act_fn/Slice_2"
        linear_slice_inputs = [
            f"{gate_up_proj_bias_name}/output_0",
            "/model/constants/INT64/[1]",
            f"/model/constants/INT64/[{torch.iinfo(torch.int64).max}]",
            "/model/constants/INT64/[3]",
            "/model/constants/INT64/[2]",
        ]
        self.make_slice(
            linear_slice_name,
            linear_slice_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        linear_clip_name = f"{basename}/act_fn/Clip_2"
        linear_clip_inputs = [
            f"{linear_slice_name}/output_0",
            f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{-self.moe_attrs['swiglu_limit']}",
            f"/model/constants/{self.to_str_dtype(self.io_dtype)}/{self.moe_attrs['swiglu_limit']}",
        ]
        self.make_clip(
            linear_clip_name,
            linear_clip_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )

        # Make Mul node after activation
        act_fn_mul_1_name = f"{basename}/act_fn/Mul_1"
        act_fn_mul_1_inputs = [
            f"{glu_clip_name}/output_0",
            f"/model/constants/{self.to_str_dtype(self.io_dtype)}/1.703125",
        ]
        self.make_mul(
            act_fn_mul_1_name,
            act_fn_mul_1_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        sigmoid_name = f"{basename}/act_fn/Sigmoid"
        self.make_sigmoid(
            sigmoid_name,
            f"{act_fn_mul_1_name}/output_0",
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        act_fn_mul_2_name = f"{basename}/act_fn/Mul_2"
        act_fn_mul_2_inputs = [f"{glu_clip_name}/output_0", f"{sigmoid_name}/output_0"]
        self.make_mul(
            act_fn_mul_2_name,
            act_fn_mul_2_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        act_fn_add_name = f"{basename}/act_fn/Add"
        self.make_add(
            act_fn_add_name,
            [f"{linear_clip_name}/output_0", f"/model/constants/{self.to_str_dtype(self.io_dtype)}/1"],
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        act_fn_mul_3_name = f"{basename}/act_fn/Mul_3"
        act_fn_mul_3_inputs = [f"{act_fn_mul_2_name}/output_0", f"{act_fn_add_name}/output_0"]
        self.make_mul(
            act_fn_mul_3_name,
            act_fn_mul_3_inputs,
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )

        # Make Down proj nodes (MatMul --> Add --> Cast)
        down_proj_weight_name = f"{basename}/down_proj/MatMul"
        down_proj_weight_output = f"{down_proj_weight_name}/output_0"
        self.make_node(
            "MatMul",
            inputs=[f"{mlp2_weight_gather_name}/output_0", f"{act_fn_mul_3_name}/output_0"],
            outputs=[down_proj_weight_output],
            name=down_proj_weight_name,
        )
        self.make_value(
            down_proj_weight_output,
            self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        down_proj_bias_name = f"{basename}/down_proj/Add"
        self.make_add(
            down_proj_bias_name,
            [down_proj_weight_output, f"{mlp2_bias_unsqueeze_name}/output_0"],
            dtype=self.io_dtype,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        if use_cast:
            down_proj_cast_name = f"{basename}/down_proj/Cast"
            self.make_cast(
                down_proj_cast_name,
                f"{down_proj_bias_name}/output_0",
                ir.DataType.FLOAT,
                shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
            )

        # Make weighted sum nodes
        #
        # Cast (from Down proj) ------->
        #                               \
        #                                Mul --> ReduceSum --> Squeeze --> Cast (created in LayerNorm)
        #                               /
        # Cast (from expert weights) -->
        weighted_sum_mul_name = f"{basename}/weighted_sum/Mul"
        weighted_sum_mul_inputs = [
            f"{down_proj_cast_name if use_cast else down_proj_bias_name}/output_0",
            f"{expert_weights_cast_name if use_cast else expert_weights_unsqueeze_2_name}/output_0",
        ]
        self.make_mul(
            weighted_sum_mul_name,
            weighted_sum_mul_inputs,
            dtype=ir.DataType.FLOAT,
            shape=["batch_size", "sequence_length", self.moe_attrs["top_k"], self.intermediate_size, 1],
        )
        reduce_sum_name = f"{basename}/weighted_sum/ReduceSum"
        reduce_sum_inputs = [f"{weighted_sum_mul_name}/output_0", "/model/constants/INT64/[2]"]
        self.make_reduce_sum(
            reduce_sum_name,
            reduce_sum_inputs,
            dtype=ir.DataType.FLOAT,
            shape=["batch_size", "sequence_length", self.intermediate_size, 1],
        )
        weighted_sum_squeeze_name = f"{basename}/weighted_sum/Squeeze"
        weighted_sum_squeeze_inputs = [f"{reduce_sum_name}/output_0", "/model/constants/INT64/[-1]"]
        self.make_squeeze(
            weighted_sum_squeeze_name,
            weighted_sum_squeeze_inputs,
            dtype=ir.DataType.FLOAT,
            shape=["batch_size", "sequence_length", self.intermediate_size],
        )

        # Assign output 0 of previous MoE as root input to next SkipLayerNorm
        self.layernorm_attrs["skip_input"] = f"{weighted_sum_squeeze_name}/output_0"

    def make_moe_fused(self, layer_id, mlp, root_input):
        # Make nodes for the fused MoE subgraph
        #
        #               root_input
        #               /        \
        #             MatMul      |
        #            (router)     |
        #               |         |
        #              Add        |
        #            (router)     |
        #               |         |
        #            Reshape      |
        #               |         |
        #               +----+----+
        #                    |
        #                 MoE/QMoE
        basename = f"/model/layers.{layer_id}/moe"
        op_type = self.moe_attrs["op_type"]
        moe_weight_type = f"{'q' if op_type == 'QMoE' else ''}weight"

        has_quark_experts = self.has_quark_experts(mlp.experts)
        moe_quant_type = self.moe_attrs.get("quant_type")
        if moe_quant_type == "nvfp4":
            # GPT-OSS checkpoints ship MXFP4 experts; there is no NVFP4 source to preserve and
            # re-quantizing MXFP4 to NVFP4 is not implemented. Fail loudly instead of silently
            # falling through to the integer QMoE path.
            raise ValueError("moe_quant_type=nvfp4 is not supported for GPT-OSS. Use moe_quant_type=mxfp4.")
        is_fp4_moe = moe_quant_type == "fp4"
        if is_fp4_moe and has_quark_experts:
            raise ValueError("moe_quant_type=mxfp4 is not supported with pre-quantized Quark GPT-OSS experts.")

        # Make router nodes
        router_basename = f"{basename}/router/MatMul"
        router_matmul_name = self.make_matmul(mlp.router, router_basename, root_input)
        router_add_name = f"{basename}/router/Add"
        self.make_add_bias(mlp.router.bias, router_add_name, root_input=f"{router_matmul_name}/output_0")
        router_reshape_name = f"{basename}/router/Reshape"
        router_reshape_inputs = [
            f"{router_add_name}/output_0",
            f"/model/constants/INT64/{[-1, self.moe_attrs['num_experts']]}",
        ]
        self.make_reshape(
            router_reshape_name,
            router_reshape_inputs,
            dtype=self.io_dtype,
            shape=["batch_size * sequence_length", self.moe_attrs["num_experts"]],
        )

        gate_up_proj_weight = f"model.layers.{layer_id}.moe.experts.gate_up_proj.{moe_weight_type}"
        gate_up_proj_scales = f"model.layers.{layer_id}.moe.experts.gate_up_proj.scales"
        gate_up_proj_bias = f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias"
        gate_up_proj_zero_points = f"model.layers.{layer_id}.moe.experts.gate_up_proj.zero_points"
        down_proj_weight = f"model.layers.{layer_id}.moe.experts.down_proj.{moe_weight_type}"
        down_proj_scales = f"model.layers.{layer_id}.moe.experts.down_proj.scales"
        down_proj_bias = f"model.layers.{layer_id}.moe.experts.down_proj.bias"
        down_proj_zero_points = f"model.layers.{layer_id}.moe.experts.down_proj.zero_points"

        # FP4 (MXFP4) per-expert float32 global scales (empty for the integer QMoE path).
        gate_up_proj_global_scales = ""
        down_proj_global_scales = ""

        # Per-expert quantized weight lists are only populated on the non-Quark QMoE
        # path below; initialize them so the shape helper never sees an undefined name.
        gate_up_proj_qweight_list = None
        down_proj_qweight_list = None

        # HF GptOssExperts stores the expert weights input-major:
        #   gate_up_proj = [E, hidden, 2*inter], down_proj = [E, inter, hidden].
        # Every downstream consumer (non-quant MoE, and the QMoE quantizers in
        # make_qmoe_weights) expects output-major [E, N, K] with the contraction
        # axis (K) last: gate_up = [E, 2*inter, hidden], down = [E, hidden, inter].
        # Transpose for all EPs/ops; the CUDA QMoE path is no exception (keeping
        # the original orientation swaps N and K, which silently corrupts the
        # quantized weights/scales and yields garbage output).
        if not has_quark_experts and not is_fp4_moe:
            gate_up_proj_layout = mlp.experts.gate_up_proj.transpose(-1, -2)
            down_proj_layout = mlp.experts.down_proj.transpose(-1, -2)

        if op_type == "MoE" and not has_quark_experts:
            # Save non-quantized MoE weights as initializers
            self.make_initializer(
                gate_up_proj_layout.view(self.moe_attrs["num_experts"], -1, self.hidden_size),
                gate_up_proj_weight,
                to=self.io_dtype,
            )
            self.make_initializer(
                down_proj_layout.view(self.moe_attrs["num_experts"], self.hidden_size, self.intermediate_size),
                down_proj_weight,
                to=self.io_dtype,
            )
        else:
            if is_fp4_moe and not has_quark_experts:
                # Preserve the original GPT-OSS checkpoint MXFP4 experts. The
                # Hugging Face loader may expose dequantized BF16 weights here;
                # re-quantizing those changes many FP4 codes/scales. Read the
                # original *_blocks/*_scales tensors and only repack their byte
                # layout for CUDA QMoE.
                self.moe_attrs["block_size"] = 32
                gate_up_proj_global_scales = f"model.layers.{layer_id}.moe.experts.gate_up_proj.global_scales"
                down_proj_global_scales = f"model.layers.{layer_id}.moe.experts.down_proj.global_scales"
                self.make_original_mxfp4_qmoe_initializers(
                    layer_id, "gate_up_proj", gate_up_proj_weight, gate_up_proj_scales, gate_up_proj_global_scales
                )
                self.make_original_mxfp4_qmoe_initializers(
                    layer_id, "down_proj", down_proj_weight, down_proj_scales, down_proj_global_scales
                )
            else:
                if has_quark_experts:
                    # Use pre-quantized Quark experts
                    gate_up_proj_qweight_tensor, gate_up_proj_scales_tensor, gate_up_proj_zero_points_tensor = (
                        mlp.experts.fc1_weights,
                        mlp.experts.fc1_scales,
                        mlp.experts.fc1_zero_points,
                    )
                    down_proj_qweight_tensor, down_proj_scales_tensor, down_proj_zero_points_tensor = (
                        mlp.experts.fc2_weights,
                        mlp.experts.fc2_scales,
                        mlp.experts.fc2_zero_points,
                    )

                    # Save zero point as initializers
                    self.make_initializer(gate_up_proj_zero_points_tensor, gate_up_proj_zero_points)
                    self.make_initializer(down_proj_zero_points_tensor, down_proj_zero_points)
                else:
                    # Create and save quantized MoE weights as initializers
                    gate_up_proj_qweight_list, gate_up_proj_scales_list = [], []
                    down_proj_qweight_list, down_proj_scales_list = [], []

                    for i in range(self.moe_attrs["num_experts"]):
                        qweight1, scales1 = self.make_qmoe_weights(gate_up_proj_layout[i])
                        gate_up_proj_qweight_list.append(qweight1)
                        gate_up_proj_scales_list.append(scales1)
                        qweight2, scales2 = self.make_qmoe_weights(down_proj_layout[i])
                        down_proj_qweight_list.append(qweight2)
                        down_proj_scales_list.append(scales2)

                    gate_up_proj_qweight_tensor = torch.stack(gate_up_proj_qweight_list, dim=0).to(torch.uint8)
                    gate_up_proj_scales_tensor = torch.stack(gate_up_proj_scales_list, dim=0)
                    down_proj_qweight_tensor = torch.stack(down_proj_qweight_list, dim=0).to(torch.uint8)
                    down_proj_scales_tensor = torch.stack(down_proj_scales_list, dim=0)

                gate_up_proj_qweight_shape, down_proj_qweight_shape = self.make_qmoe_weight_initializer_shapes(
                    gate_up_proj_qweight_list if not has_quark_experts else None,
                    down_proj_qweight_list if not has_quark_experts else None,
                    has_quark_experts,
                )

                # Save qweight tensors
                self.make_initializer(
                    gate_up_proj_qweight_tensor.view(gate_up_proj_qweight_shape),
                    gate_up_proj_weight,
                )
                self.make_initializer(
                    down_proj_qweight_tensor.view(down_proj_qweight_shape),
                    down_proj_weight,
                )

                # Save scales tensors
                self.make_initializer(gate_up_proj_scales_tensor, gate_up_proj_scales, to=self.io_dtype)
                self.make_initializer(down_proj_scales_tensor, down_proj_scales, to=self.io_dtype)

        # Save biases (shared for all paths)
        if has_quark_experts:
            gate_up_bias = self.combine_quark_gate_up_biases_from_experts(mlp.experts)
            down_bias = self.combine_quark_down_biases_from_experts(mlp.experts)
        else:
            gate_up_bias = mlp.experts.gate_up_proj_bias
            down_bias = mlp.experts.down_proj_bias

        self.make_initializer(gate_up_bias, gate_up_proj_bias, to=self.io_dtype)
        self.make_initializer(down_bias, down_proj_bias, to=self.io_dtype)

        # Single make_moe_op call with EP-based zero_points
        # TRT-RTX doesn't support zero_points inputs
        moe_name = f"{basename}/{op_type}"
        use_zero_points = has_quark_experts and self.ep != "trt-rtx"

        self.make_moe_op(
            moe_name,
            root_input=root_input,
            router_probs=f"{router_reshape_name}/output_0",
            weight1=gate_up_proj_weight,
            scales1=gate_up_proj_scales,
            bias1=gate_up_proj_bias,
            weight2=down_proj_weight,
            scales2=down_proj_scales,
            bias2=down_proj_bias,
            zero_points1=gate_up_proj_zero_points if use_zero_points else "",
            zero_points2=down_proj_zero_points if use_zero_points else "",
            global_scales1=gate_up_proj_global_scales,
            global_scales2=down_proj_global_scales,
        )

        # Assign output 0 of previous MoE as root input to next SkipLayerNorm
        self.layernorm_attrs["skip_input"] = f"{moe_name}/output_0"

    def make_qmoe_weight_initializer_shapes(self, gate_up_proj_qweight_list, down_proj_qweight_list, has_quark_experts):
        pack_size = 8 // self.moe_attrs["expert_weight_bits"]
        if has_quark_experts:
            hidden_size_packed = self.hidden_size // pack_size
            intermediate_size_packed = self.intermediate_size // pack_size
            return (
                (self.moe_attrs["num_experts"], -1, hidden_size_packed),
                (self.moe_attrs["num_experts"], self.hidden_size, intermediate_size_packed),
            )
        else:
            return (
                (self.moe_attrs["num_experts"], *gate_up_proj_qweight_list[0].shape),
                (self.moe_attrs["num_experts"], *down_proj_qweight_list[0].shape),
            )

    def has_quark_experts(self, experts):
        return hasattr(experts, "fc1_weights") and hasattr(experts, "fc2_weights")

    def combine_quark_gate_up_biases_from_experts(self, experts):
        """Combine Quark gate_proj and up_proj biases from individual experts"""
        assert(self.has_quark_experts(experts))
        combined_biases = []

        for expert_id in sorted(experts.keys()):
            expert = experts[expert_id]

            if expert.gate_up_proj.qweight is not None:
                # Fused gate_up projection
                gate_up_proj = expert.gate_up_proj.bias if hasattr(expert.gate_up_proj, 'bias') and expert.gate_up_proj.bias is not None else torch.zeros(expert.gate_up_proj.qweight.shape[0])
                combined_biases.append(gate_up_proj)
            else:
                # Get biases from individual projections
                gate_bias = expert.gate_proj.bias if hasattr(expert.gate_proj, 'bias') and expert.gate_proj.bias is not None else torch.zeros(expert.gate_proj.qweight.shape[0])
                up_bias = expert.up_proj.bias if hasattr(expert.up_proj, 'bias') and expert.up_proj.bias is not None else torch.zeros(expert.up_proj.qweight.shape[0])

                # Combine gate and up biases (interleaved pattern: even=gate, odd=up)
                gate_out_dim = gate_bias.shape[0]
                up_out_dim = up_bias.shape[0]

                combined_bias = torch.zeros(gate_out_dim + up_out_dim, dtype=gate_bias.dtype, device="cpu")
                combined_bias[::2] = gate_bias   # Even indices = gate
                combined_bias[1::2] = up_bias    # Odd indices = up

                combined_biases.append(combined_bias)

        return torch.stack(combined_biases, dim=0)

    def combine_quark_down_biases_from_experts(self, experts):
        """Combine Quark down_proj biases from individual experts"""
        assert(self.has_quark_experts(experts))
        combined_biases = []

        for expert_id in sorted(experts.keys()):
            expert = experts[expert_id]
            down_bias = expert.down_proj.bias if hasattr(expert.down_proj, 'bias') and expert.down_proj.bias is not None else torch.zeros(expert.down_proj.qweight.shape[0])
            combined_biases.append(down_bias)

        return torch.stack(combined_biases, dim=0)
