[None][fix] Apply the NVFP4_AWQ pre-quant scale in MegaMoE CuteDSL

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -1104,7 +1104,10 @@ def _convert_moe_mlps_to_fp8_weight_read(
 # ---------------------------------------------------------------------------
 
 
-def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None:
+def _load_packed_mxfp4_expert(
+    backend, base, expert_idx, local_slot_id, get_tensor, get_optional_tensor
+) -> None:
+    del get_optional_tensor  # this layout stores no optional per-expert tensor
     backend.quant_method.load_packed_mxfp4_expert(
         backend,
         global_expert_id=expert_idx,
@@ -1118,7 +1121,18 @@ def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tens
     )
 
 
-def _load_nvfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None:
+def _load_nvfp4_expert(
+    backend, base, expert_idx, local_slot_id, get_tensor, get_optional_tensor
+) -> None:
+    # Only an NVFP4_AWQ / NVFP4_ARC calibration writes a pre-quant scale; a
+    # plain NVFP4 checkpoint has none and the pair stays None, which leaves
+    # ``fc31_act_scale`` unset exactly as before. Dropping it for a checkpoint
+    # that HAS one would be silently wrong -- the activation would be quantized
+    # unsmoothed against an input scale calibrated on the smoothed one.
+    pre_quant_scales = {
+        f"{w}_pre_quant_scale": get_optional_tensor(f"{base}.{expert_idx}.{w}.pre_quant_scale")
+        for w in ("w1", "w3")
+    }
     backend.quant_method.load_streaming_nvfp4_expert(
         backend,
         global_expert_id=expert_idx,
@@ -1128,21 +1142,37 @@ def _load_nvfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) ->
             for w in ("w1", "w2", "w3")
             for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale")
         },
+        **pre_quant_scales,
     )
 
 
 class _K3ExpertCkptSpec(NamedTuple):
     """How one routed-expert quantization is spelled and loaded."""
 
-    # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores.
+    # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. Every
+    # combination is required: a missing one fails the key plan.
     kinds: Tuple[str, ...]
     loader: Callable[..., None]
     # Set of filled slots the loader maintains, checked after the load.
     loaded_slots_attr: str
     # NVFP4 defers cat/pad/interleave and the alpha computation to
     # ``process_weights_after_loading``; the MXFP4 loaders write through.
     needs_layer_finalize: bool
-
+    # Whole per-expert key suffixes the loader passes through WHEN PRESENT, for
+    # a calibration artifact only some checkpoints of this layout carry. Spelled
+    # per weight rather than per kind, because they are not stored for every
+    # ``w``, and kept out of ``kinds`` because demanding them would reject the
+    # checkpoints that have none.
+    optional_suffixes: Tuple[str, ...] = ()
+
+
+_K3_NVFP4_EXPERT_CKPT_SPEC = _K3ExpertCkptSpec(
+    kinds=("weight", "weight_scale", "weight_scale_2", "input_scale"),
+    loader=_load_nvfp4_expert,
+    loaded_slots_attr="_streamed_expert_slots",
+    needs_layer_finalize=True,
+    optional_suffixes=("w1.pre_quant_scale", "w3.pre_quant_scale"),
+)
 
 _K3_EXPERT_CKPT_SPECS = {
     QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec(
@@ -1151,12 +1181,12 @@ class _K3ExpertCkptSpec(NamedTuple):
         loaded_slots_attr="_packed_mxfp4_loaded_slots",
         needs_layer_finalize=False,
     ),
-    QuantAlgo.NVFP4: _K3ExpertCkptSpec(
-        kinds=("weight", "weight_scale", "weight_scale_2", "input_scale"),
-        loader=_load_nvfp4_expert,
-        loaded_slots_attr="_streamed_expert_slots",
-        needs_layer_finalize=True,
-    ),
+    QuantAlgo.NVFP4: _K3_NVFP4_EXPERT_CKPT_SPEC,
+    # The two NVFP4 calibration recipes store plain NVFP4 expert weights plus
+    # the per-channel pre-quant scale above, so they share one layout; the MoE
+    # quantization method folds them onto the NVFP4 pipeline as well.
+    QuantAlgo.NVFP4_AWQ: _K3_NVFP4_EXPERT_CKPT_SPEC,
+    QuantAlgo.NVFP4_ARC: _K3_NVFP4_EXPERT_CKPT_SPEC,
 }
 
 
@@ -2813,6 +2843,11 @@ def checkpoint_name_plan(
                 continue
             moe = layer.block_sparse_moe
             base = f"{prefix}model.layers.{layer_idx}.block_sparse_moe.experts"
+            # Required suffixes only: ``optional_suffixes`` are per-checkpoint
+            # (an AWQ pre-quant scale), and demanding them here would reject
+            # every checkpoint of the same layout that does not carry them.
+            # Unmatched expert keys only warn, and this validator already
+            # exempts expert keys from that, so the present ones stay quiet.
             kinds = moe.expert_ckpt_spec.kinds
             for expert_idx in moe.local_expert_ids:
                 for w in ("w1", "w2", "w3"):
@@ -3164,7 +3199,12 @@ def maybe_finalize_layer(moe: KimiK3MoERuntime):
             backend.process_weights_after_loading()
 
         def load_expert(
-            moe: KimiK3MoERuntime, base: str, local_slot_id: int, expert_idx: int, get_tensor
+            moe: KimiK3MoERuntime,
+            base: str,
+            local_slot_id: int,
+            expert_idx: int,
+            get_tensor,
+            get_optional_tensor,
         ):
             if device.type == "cuda":
                 torch.cuda.set_device(device)
@@ -3173,7 +3213,9 @@ def load_expert(
             # split-file and the shared-dict fallback alike), so this is the
             # one place preparation has to be hooked.
             ensure_prepared(moe)
-            moe.expert_ckpt_spec.loader(backend, base, expert_idx, local_slot_id, get_tensor)
+            moe.expert_ckpt_spec.loader(
+                backend, base, expert_idx, local_slot_id, get_tensor, get_optional_tensor
+            )
             maybe_finalize_layer(moe)
 
         def load_experts_from_weights(layer_idx: int, moe: KimiK3MoERuntime, base: str):
@@ -3185,6 +3227,7 @@ def load_experts_from_weights(layer_idx: int, moe: KimiK3MoERuntime, base: str):
                     local_slot_id,
                     expert_idx,
                     lambda key: _materialize(weights[key]),
+                    lambda key: _materialize(weights[key]) if key in weights else None,
                 )
 
         # ---- backend expert slots: file-grouped streaming ----
@@ -3216,11 +3259,23 @@ def checkpoint_key(key: str) -> str:
             split_file_jobs = []
             for layer_idx, moe, base in expert_jobs:
                 del layer_idx
+                spec = moe.expert_ckpt_spec
                 for local_slot_id, expert_idx in enumerate(moe.local_expert_ids):
                     keys = [
                         f"{base}.{expert_idx}.{w}.{kind}"
                         for w in ("w1", "w2", "w3")
-                        for kind in moe.expert_ckpt_spec.kinds
+                        for kind in spec.kinds
+                    ]
+                    # An optional tensor this checkpoint does store has to be
+                    # counted here: it decides which files the job opens, and
+                    # ``get_optional_tensor`` below can only read a file that
+                    # ``files`` named.
+                    keys += [
+                        key
+                        for key in (
+                            f"{base}.{expert_idx}.{suffix}" for suffix in spec.optional_suffixes
+                        )
+                        if checkpoint_key(key) in weight_map
                     ]
                     files = {weight_map[checkpoint_key(key)] for key in keys}
                     job = (moe, base, local_slot_id, expert_idx)
@@ -3245,13 +3300,22 @@ def load_expert_file(file_name: str, jobs: list):
                     torch.cuda.set_device(device)
                 path = os.path.join(ckpt_dir, file_name)
                 with safe_open(path, framework="pt", device="cpu") as fh:
+                    # A job lands here only when every key it needs -- optional
+                    # ones included, see the grouping above -- is in this file,
+                    # so presence in this file is the whole question.
+                    file_keys = set(fh.keys())
                     for moe, base, local_slot_id, expert_idx in jobs:
                         load_expert(
                             moe,
                             base,
                             local_slot_id,
                             expert_idx,
                             lambda key: fh.get_tensor(checkpoint_key(key)),
+                            lambda key: (
+                                fh.get_tensor(checkpoint_key(key))
+                                if checkpoint_key(key) in file_keys
+                                else None
+                            ),
                         )
                 # Handle closed -> pages unmapped -> the drop takes effect.
                 drop_file_pages(file_name)
@@ -3273,7 +3337,16 @@ def get_tensor(key):
                         source_key = checkpoint_key(key)
                         return handles[weight_map[source_key]].get_tensor(source_key)
 
-                    load_expert(*job, get_tensor)
+                    def get_optional_tensor(key):
+                        source_key = checkpoint_key(key)
+                        # Absent from the index means the checkpoint does not
+                        # store it; present means the grouping above put its
+                        # file in ``handles``.
+                        if source_key not in weight_map:
+                            return None
+                        return handles[weight_map[source_key]].get_tensor(source_key)
+
+                    load_expert(*job, get_tensor, get_optional_tensor)
                 for file_name in files:
                     drop_file_pages(file_name)
 

tensorrt_llm/_torch/moe/fused_moe/cutlass/w4a16_nvfp4.py

@@ -18,6 +18,7 @@
 
 import torch
 
+from ....utils import Fp4QuantizedTensor
 from ..impl_contract import (
     MoEDeployment,
     MoEEligibility,
@@ -56,6 +57,34 @@ def can_implement(cls, p: MoEProblem, d: MoEDeployment) -> MoEEligibility:
     def _get_quant_method(self) -> object:
         return W4A16NVFP4CutlassFusedMoEMethod()
 
+    def quantize_input(
+        self,
+        x: torch.Tensor | Fp4QuantizedTensor,
+        post_quant_comm: bool = True,
+        **kwargs: object,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Apply the AWQ pre-quant scale, keeping the activations high precision.
+
+        This path keeps activations high precision -- ``run_moe`` dequantizes
+        the NVFP4 weights and runs the unquantized kernel -- so unlike the
+        ``nvfp4`` leaf there is no FP4 activation quantization here. A pre-quant
+        scale is not part of that quantization, though: an AWQ-style checkpoint
+        stores ``W / s`` and ``load_quant_scales`` fills ``fc31_act_scale`` from
+        its ``*.pre_quant_scale`` keys regardless of ``quant_algo``, so the
+        dequantized weights expect the scaled activation ``x * s``. Matches the
+        dense ``W4A16NVFP4LinearMethod``, which applies ``pre_quant_scale`` the
+        same way.
+
+        See :meth:`CutlassFusedMoEBase.quantize_input` for the arguments.
+        """
+        del post_quant_comm, kwargs
+        if self.fc31_act_scale is not None:
+            assert not isinstance(x, Fp4QuantizedTensor), (
+                "Fp4QuantizedTensor is not expected for AWQ quantization."
+            )
+            x = x * self.fc31_act_scale
+        return x, None
+
     def run_moe(self, ctx: MoERunContext, *, workspace: dict | None = None) -> torch.Tensor:
         """Dequantize the active experts into a static workspace, then run unquantized.
 

tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py

@@ -1041,6 +1041,18 @@ def quantize_input(self,
         """
         x_sf = None
         if self.has_nvfp4:
+            # An NVFP4_AWQ / NVFP4_ARC checkpoint carries a per-channel
+            # pre-quant scale, which the quantization method loads into
+            # ``fc31_act_scale``, and its ``fc31_input_scale`` was calibrated
+            # on the scaled activation. This is the only place the activation
+            # is quantized, so apply the scale here, as
+            # ``CutlassFusedMoE.quantize_input`` does.
+            fc31_act_scale = getattr(self, "fc31_act_scale", None)
+            if fc31_act_scale is not None:
+                assert not isinstance(
+                    x, Fp4QuantizedTensor
+                ), "Fp4QuantizedTensor is not expected for AWQ quantization."
+                x = x * fc31_act_scale
             if isinstance(x, Fp4QuantizedTensor):
                 assert not x.is_sf_swizzled, "Fp4QuantizedTensor should not be swizzled before communication"
                 x_row = x.shape[0]

tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py

@@ -472,10 +472,22 @@ def quantize_input(
         post_quant_comm: bool = True,
         **kwargs: object,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Always a pass-through: b12x consumes bf16 / fp16 and produces its
-        own scale factors."""
+        """Hand the activation over in high precision, with an AWQ pre-quant
+        scale applied to it.
+
+        The wrapper runs in ``quant_mode="w4a16"``: it dequantizes the NVFP4
+        weights and leaves the bf16 / fp16 activation alone, so there is no
+        activation quantization here at all. The pre-quant scale is a separate
+        concern. An AWQ-style checkpoint stores ``W / s``, and
+        ``load_quant_scales`` fills ``fc31_act_scale`` from its
+        ``*.pre_quant_scale`` keys regardless of ``quant_algo``, so the weights
+        the kernel dequantizes expect the scaled activation ``x * s``. Same as
+        the NVFP4 sibling's decode path above and the CUTLASS W4A16 leaf.
+        """
         del post_quant_comm, kwargs
         self._reject_fp4_input(x)
+        if self.fc31_act_scale is not None:
+            x = x * self.fc31_act_scale
         return x, None
 
     @nvtx_range("[b12x] run_moe")

tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py

@@ -300,6 +300,17 @@ def quantize_input(
         """
         del post_quant_comm  # single path: always quantized ahead of comm
         assert self.has_nvfp4
+        # An NVFP4_AWQ / NVFP4_ARC checkpoint carries a per-channel pre-quant
+        # scale, which the quantization method loads into ``fc31_act_scale``,
+        # and its ``fc31_input_scale`` was calibrated on the scaled activation.
+        # This is the only place the activation is quantized, so apply the
+        # scale here, as ``CutlassFusedMoE.quantize_input`` does.
+        fc31_act_scale = getattr(self, "fc31_act_scale", None)
+        if fc31_act_scale is not None:
+            assert not isinstance(x, Fp4QuantizedTensor), (
+                "Fp4QuantizedTensor is not expected for AWQ quantization."
+            )
+            x = x * fc31_act_scale
         if isinstance(x, Fp4QuantizedTensor):
             assert not x.is_sf_swizzled, (
                 "Fp4QuantizedTensor should not be swizzled before communication"

tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py

@@ -388,6 +388,18 @@ def quantize_input(
         """
         x_sf = None
         if self.has_nvfp4:
+            # An NVFP4_AWQ / NVFP4_ARC checkpoint carries a per-channel
+            # pre-quant scale, which the quantization method loads into
+            # ``fc31_act_scale``, and its ``fc31_input_scale`` was calibrated
+            # on the scaled activation. This is the only place the activation
+            # is quantized, so apply the scale here, as
+            # ``CutlassFusedMoE.quantize_input`` does.
+            fc31_act_scale = getattr(self, "fc31_act_scale", None)
+            if fc31_act_scale is not None:
+                assert not isinstance(x, Fp4QuantizedTensor), (
+                    "Fp4QuantizedTensor is not expected for AWQ quantization."
+                )
+                x = x * fc31_act_scale
             if isinstance(x, Fp4QuantizedTensor):
                 assert not x.is_sf_swizzled, (
                     "Fp4QuantizedTensor should not be swizzled before communication"

tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py

@@ -1167,6 +1167,12 @@ def quantize_input(
         registered by the quantization method's ``create_weights``; the
         value defaults to 1.0 until the checkpoint loader sets it.
 
+        A checkpoint that carries a per-channel pre-quant scale (NVFP4_AWQ /
+        NVFP4_ARC) gets it applied before quantization: the quantization
+        method loads it into ``self.fc31_act_scale`` (``[1, hidden]``), and
+        ``fc31_input_scale`` was calibrated on the scaled activation. The
+        kernel itself only ever sees the quantized input.
+
         Empty input (``x.shape[0] == 0``) short-circuits to empty NVFP4
         + empty SF without launching a quantization kernel, so the
         ``FusedCommMoEScheduler`` can call ``quantize_input`` uniformly
@@ -1187,6 +1193,13 @@ def quantize_input(
             empty_x = torch.empty((0, hidden // 2), dtype=torch.uint8, device=x_bf16.device)
             empty_sf = torch.empty((0, sf_cols), dtype=torch.uint8, device=x_bf16.device)
             return empty_x, empty_sf
+        fc31_act_scale = getattr(self, "fc31_act_scale", None)
+        if fc31_act_scale is not None:
+            # Product stays bf16, which is what ``fp4_quantize`` below is given
+            # on every other route into it: ``_SUPPORTED_ACTIVATION_DTYPES``
+            # admits bf16 activations only, so ``fc31_act_scale``, created at
+            # ``module.dtype``, is bf16 as well.
+            x_bf16 = x_bf16 * fc31_act_scale
         x_fp4, x_sf = torch.ops.trtllm.fp4_quantize(
             x_bf16,
             self.fc31_input_scale,