diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index f244d1cf2763..ec7aee2a0f05 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/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,13 +1142,15 @@ 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. @@ -1142,7 +1158,21 @@ class _K3ExpertCkptSpec(NamedTuple): # 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,6 +3300,10 @@ 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, @@ -3252,6 +3311,11 @@ def load_expert_file(file_name: str, jobs: list): 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) diff --git a/tensorrt_llm/_torch/moe/fused_moe/cutlass/w4a16_nvfp4.py b/tensorrt_llm/_torch/moe/fused_moe/cutlass/w4a16_nvfp4.py index e976a3dbacf6..71aca7d52c4c 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/cutlass/w4a16_nvfp4.py +++ b/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. diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py index 91879d73d94d..96641ca6c7ae 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl.py +++ b/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] diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py index c9c88f31197a..19d4793c371d 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_b12x.py +++ b/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") diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py index 2472709f992a..8f877b4a48c5 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cute_dsl_fc12.py +++ b/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" diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py index 387cf57da534..c973beaffd69 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_densegemm.py +++ b/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" diff --git a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py index a75c39ca9bed..6a4db9569abe 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/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, diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index 61209a66f33e..e0fffead9c3b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -2543,6 +2543,11 @@ def prepare_streaming_expert_load(self, module: torch.nn.Module) -> None: """ module.tmp_raw_input_scales = {} module.tmp_weight_scale_2 = {} + # Created unconditionally even though only an NVFP4_AWQ / NVFP4_ARC + # checkpoint fills it: whether this checkpoint carries a pre-quant + # scale is not known until the first expert arrives, by which time + # several loader threads are already running. + module.tmp_pre_quant_scales = {} module._streamed_expert_slots = set() def _finalize_staged_w3_w1_expert( @@ -2599,6 +2604,8 @@ def load_streaming_nvfp4_expert( w3_weight_scale: torch.Tensor, w3_weight_scale_2: torch.Tensor, w3_input_scale: torch.Tensor, + w1_pre_quant_scale: Optional[torch.Tensor] = None, + w3_pre_quant_scale: Optional[torch.Tensor] = None, ) -> None: """Load one NVFP4 checkpoint expert into a local slot. @@ -2613,6 +2620,15 @@ def load_streaming_nvfp4_expert( ``module.process_weights_after_loading()`` must run once the layer's last slot is in — nothing else calls it on this path. + ``w1_pre_quant_scale`` / ``w3_pre_quant_scale`` are the per-channel + pre-quant scale an NVFP4_AWQ / NVFP4_ARC checkpoint carries and a plain + NVFP4 one does not; passing them stages them exactly as + ``load_quant_scales`` does, so ``_finalize_pre_quant_scales`` lands the + same ``fc31_act_scale`` on either path. Leaving them out of such a + checkpoint's load is silently wrong rather than merely unoptimized: the + activation would be quantized unsmoothed against an ``input_scale`` + calibrated on ``x * pre_quant_scale``. + One deliberate difference from ``load_quant_scales``: it reduces the activation ``input_scale`` over every expert in the checkpoint, whereas a streaming caller only ever holds its own rank's experts, so the @@ -2645,6 +2661,18 @@ def load_streaming_nvfp4_expert( raise TypeError(f"{name} must contain packed NVFP4 uint8 data, " f"got {value.dtype}.") + has_pre_quant_scale = (w1_pre_quant_scale is not None + or w3_pre_quant_scale is not None) + if has_pre_quant_scale: + if w1_pre_quant_scale is None or w3_pre_quant_scale is None: + raise ValueError( + "w1_pre_quant_scale and w3_pre_quant_scale must be loaded " + "together; fc31_act_scale is their elementwise maximum.") + if not module.is_gated_activation: + raise ValueError( + "pre_quant_scale (NVFP4_AWQ / NVFP4_ARC) is not supported " + "with non-gated activations.") + # Cutlass stages the two halves of w3_w1 per expert and needs the slot # id as the staging key; other backends write straight through. w3_w1_weight_kargs = {} @@ -2681,6 +2709,25 @@ def load_streaming_nvfp4_expert( 'w2': w2_input_scale[...].reshape([]), } + # Per-channel pre-quant scale (NVFP4_AWQ / NVFP4_ARC), keyed by global + # expert id like the input scales above and for the same reason: it is + # only cross-checked and reduced, never indexed positionally. Sharded + # with the same ROW mode ``load_quant_scales`` uses, which for this + # 1-D per-channel tensor leaves it whole. Staged on the host + # (``load_weight_shard``'s default), where a streaming checkpoint's + # tensors already are, so nothing holds device memory for the length of + # the layer load; ``_finalize_pre_quant_scales`` moves the reduction's + # result onto ``fc31_act_scale``. + if has_pre_quant_scale: + module.tmp_pre_quant_scales[global_expert_id] = { + 'w1': + load_weight_shard(w1_pre_quant_scale, module.tp_size, + module.tp_rank, TensorParallelMode.ROW), + 'w3': + load_weight_shard(w3_pre_quant_scale, module.tp_size, + module.tp_rank, TensorParallelMode.ROW), + } + # Backends that only stage their input here resolve this slot right # away, so the staged footprint stays proportional to the experts in # flight rather than growing with the load. @@ -2782,6 +2829,37 @@ def load_fp4_weight_block_scales( tmp_weight_scale_2.setdefault(expert_idx, {}) tmp_weight_scale_2[expert_idx]['w2'] = w2_weight_scale_2 + def _ensure_fc31_act_scale(self, module: torch.nn.Module) -> None: + """Register the per-channel act scale an AWQ checkpoint needs. + + ``create_weights`` registers the slot as ``None`` because only a + checkpoint carrying ``*.pre_quant_scale`` (NVFP4_AWQ / NVFP4_ARC) has a + value for it. The whole-checkpoint loader creates it as soon as it sees + those keys; the streaming loader has no such single point, so it + creates it from ``_finalize_pre_quant_scales`` instead, which is the + first place after a layer's experts are all in. + """ + if getattr(module, "fc31_act_scale", None) is not None: + return + # Shape: (1, hidden_size) - single vector for all experts (they share + # the same input). + # + # Named device, unlike the staging tensors around it, because this is + # a forward operand rather than load scratch: ``quantize_input`` + # multiplies the activation by it on every call, so it belongs on the + # compute device whichever loader allocates it. The current device is + # the right one for both callers -- the whole-checkpoint loader runs + # on it, and each streaming loader thread calls + # ``torch.cuda.set_device`` before it loads an expert. This is the one + # place the device is chosen: ``_finalize_pre_quant_scales`` moves the + # staged scales onto this parameter. + fc31_act_scale = nn.Parameter(torch.empty(1, + module.hidden_size, + dtype=module.dtype, + device='cuda'), + requires_grad=False) + module.register_parameter("fc31_act_scale", fc31_act_scale) + def load_quant_scales(self, module: torch.nn.Module, weights: Dict): # Check if pre_quant_scale exists in the checkpoint (for NVFP4_AWQ) has_pre_quant_scale = False @@ -2792,14 +2870,7 @@ def load_quant_scales(self, module: torch.nn.Module, weights: Dict): # If pre_quant_scale exists, we need a per-channel act scale for fc31 # All experts share the same input, so pre_quant_scale should be identical across experts if has_pre_quant_scale: - # Create fc31_act_scale parameter (for gate_up_proj / w3_w1) - # Shape: (1, hidden_size) - single vector for all experts (they share the same input) - fc31_act_scale = nn.Parameter(torch.empty(1, - module.hidden_size, - dtype=module.dtype, - device='cuda'), - requires_grad=False) - module.register_parameter("fc31_act_scale", fc31_act_scale) + self._ensure_fc31_act_scale(module) # Step1: Store raw per-expert input scales individually into tmp dict. # Consistency check + global max computation deferred to process_weights_after_loading(). @@ -2977,6 +3048,15 @@ def _finalize_pre_quant_scales(self, module: torch.nn.Module): """Verify pre_quant_scale consistency across experts and compute fc31_act_scale.""" if not hasattr(module, 'tmp_pre_quant_scales'): return + if not module.tmp_pre_quant_scales: + # Only the streaming loader gets here: it pre-creates the container + # before it knows whether the checkpoint carries a pre-quant scale, + # and a plain NVFP4 one leaves it empty. The whole-checkpoint loader + # creates it only after finding the keys, so an empty-but-present + # container never means "the scales went missing" there either. + delattr(module, 'tmp_pre_quant_scales') + return + self._ensure_fc31_act_scale(module) # Collect all loaded w3/w1 pre_quant_scales all_w3 = [] all_w1 = [] @@ -3011,10 +3091,14 @@ def _finalize_pre_quant_scales(self, module: torch.nn.Module): logger.error(msg) raise ValueError(msg) - # Take the max of w3 and w1 from the first expert - fc31_pre_quant_scale = torch.max(w3_reference, - w1_reference).to(dtype=module.dtype, - device='cuda') + # Take the max of w3 and w1 from the first expert, moved onto the + # destination parameter's device: the staged scales sit wherever the + # loader that staged them put them -- the host for the streaming + # loader, ``fc31_act_scale.device`` for ``load_quant_scales``, which + # keys its own staging off this same parameter -- and + # ``_ensure_fc31_act_scale`` is the one place that chooses the device. + fc31_pre_quant_scale = torch.max(w3_reference, w1_reference).to( + dtype=module.dtype, device=module.fc31_act_scale.device) module.fc31_act_scale.data.copy_(fc31_pre_quant_scale.unsqueeze(0)) delattr(module, 'tmp_pre_quant_scales') diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index d95c426ee9e3..bf3b1ab6eff6 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -209,6 +209,7 @@ l0_b200: - unittest/_torch/moe/test_moe_backend.py::test_enumerate_megamoe_candidate_tactics_curated_space - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_runner_cache_identity_includes_swiglu_constants - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_minimax_m3_swiglu_bias_numerics + - unittest/_torch/moe/test_moe_backend.py::test_megamoe_cutedsl_nvfp4_awq_matches_cutlass - unittest/_torch/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts - unittest/_torch/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fuse_shared_expert_layout diff --git a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py index 1afef22e0a1b..855b509a57ab 100644 --- a/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py +++ b/tests/unittest/_torch/moe/test_cute_dsl_b12x_moe_backend.py @@ -232,6 +232,42 @@ def test_nvfp4_decode_applies_the_awq_pre_quant_scale(): assert out_sf is None +@pytest.mark.parametrize("has_pre_quant_scale", [False, True], ids=["nvfp4", "nvfp4_awq"]) +def test_w4a16_applies_the_awq_pre_quant_scale(has_pre_quant_scale): + """This leaf never quantizes the activation, but an AWQ checkpoint still + needs its pre-quant scale: it stores ``W / s``, so the weights the b12x + W4A16 kernel dequantizes expect ``x * s``. ``load_quant_scales`` fills + ``fc31_act_scale`` from the ``*.pre_quant_scale`` keys without consulting + ``quant_algo``, so a plain NVFP4 checkpoint leaves it None and the input is + handed over untouched. + + ``__new__`` without ``__init__``: the real constructor wants a GPU, while + ``quantize_input`` reads one parameter. Registered the way + ``create_weights`` does -- the slot as ``None``, then a ``(1, hidden)`` + parameter of the module dtype -- so the ``nn.Module`` parameter lookup and + the broadcast shape are the production ones. + """ + hidden_size = 8 + moe = CuteDslB12xW4a16Nvfp4FusedMoE.__new__(CuteDslB12xW4a16Nvfp4FusedMoE) + torch.nn.Module.__init__(moe) + moe.register_parameter("fc31_act_scale", None) + pre_quant_scale = None + if has_pre_quant_scale: + pre_quant_scale = torch.linspace(0.25, 4.0, hidden_size, dtype=torch.bfloat16).unsqueeze(0) + moe.register_parameter( + "fc31_act_scale", torch.nn.Parameter(pre_quant_scale, requires_grad=False) + ) + x = torch.randn(3, hidden_size, dtype=torch.bfloat16) + + out, out_sf = moe.quantize_input(x) + + assert out_sf is None + if has_pre_quant_scale: + torch.testing.assert_close(out, x * pre_quant_scale, rtol=0, atol=0) + else: + assert out is x + + def test_w4a16_nvfp4_post_load_uses_modelopt_scale_contract(monkeypatch): class _RoutingMethod: experts_per_token = 4 diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index fdaedb1cefc0..55f591215c82 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -92,6 +92,7 @@ ) from tensorrt_llm._torch.moe.fused_moe.fused_moe_cutlass import CutlassFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_deepgemm import DeepGemmFusedMoE +from tensorrt_llm._torch.moe.fused_moe.fused_moe_densegemm import DenseGEMMFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_marlin import MarlinFusedMoE from tensorrt_llm._torch.moe.fused_moe.fused_moe_trtllm_gen import ( TRTLLMGenFusedMoE, @@ -135,6 +136,7 @@ FusedMoEMethodBase, NVFP4FusedMoEMethod, NVFP4MarlinFusedMoEMethod, + NVFP4MegaMoECuteDslMethod, NVFP4TRTLLMGenFusedMoEBaseMethod, NVFP4TRTLLMGenFusedMoEMethod, UnquantizedFusedMoEMethod, @@ -146,7 +148,12 @@ TrtllmTrtllmGenNvfp4Impl, TrtllmTrtllmGenW4a8Mxfp4Mxfp8Impl, ) -from tensorrt_llm._torch.utils import ActivationType, MxFp8QuantizedTensor, is_gated_activation +from tensorrt_llm._torch.utils import ( + ActivationType, + Fp4QuantizedTensor, + MxFp8QuantizedTensor, + is_gated_activation, +) from tensorrt_llm._utils import get_sm_version, is_sm_100f, mpi_rank from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -2592,6 +2599,437 @@ def test_megamoe_cutedsl_rejects_nonuniform_swiglu_bias( ) +@pytest.mark.cpu_only +@pytest.mark.parametrize("has_pre_quant_scale", [False, True], ids=["nvfp4", "nvfp4_awq"]) +def test_megamoe_cutedsl_quantize_input_applies_pre_quant_scale( + monkeypatch: pytest.MonkeyPatch, has_pre_quant_scale: bool +) -> None: + """An NVFP4_AWQ / NVFP4_ARC pre-quant scale reaches the activation quantizer. + + ``canonical_quant`` folds both recipes into ``nvfp4``, so ``can_implement`` + admits them, and ``load_quant_scales`` loads their per-channel pre-quant + scale into ``fc31_act_scale``. The fused kernel consumes only the quantized + activation, which leaves ``quantize_input`` as the one place to apply it. + + ``__new__`` without ``__init__``: the real constructor wants a GPU and a + symmetric-memory provider, and ``quantize_input`` reads two parameters. + """ + hidden_size = 64 + backend = TrtllmCutedslMegaMoeNvfp4Impl.__new__(TrtllmCutedslMegaMoeNvfp4Impl) + torch.nn.Module.__init__(backend) + backend.fc31_input_scale = torch.nn.Parameter(torch.tensor(0.5), requires_grad=False) + # Weight creation registers the slot as None, and only a checkpoint that + # carries ``*.pre_quant_scale`` fills it, as ``NVFP4FusedMoEMethod`` does. + backend.register_parameter("fc31_act_scale", None) + if has_pre_quant_scale: + backend.register_parameter( + "fc31_act_scale", + torch.nn.Parameter( + torch.linspace(0.25, 4.0, hidden_size, dtype=torch.bfloat16).unsqueeze(0), + requires_grad=False, + ), + ) + quantizer_calls = [] + + def fake_fp4_quantize(x, global_scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled): + quantizer_calls.append((x.clone(), global_scale)) + num_rows, num_cols = x.shape + return ( + torch.zeros(num_rows, num_cols // 2, dtype=torch.uint8), + torch.zeros(num_rows * num_cols // sf_vec_size, dtype=torch.uint8), + ) + + monkeypatch.setattr(torch.ops.trtllm, "fp4_quantize", fake_fp4_quantize, raising=False) + x = torch.randn(3, hidden_size, dtype=torch.bfloat16) + + backend.quantize_input(x) + + [(quantized, global_scale)] = quantizer_calls + expected = x * backend.fc31_act_scale if has_pre_quant_scale else x + torch.testing.assert_close(quantized, expected, rtol=0, atol=0) + assert global_scale is backend.fc31_input_scale + + +def _make_cutedsl_sibling_for_quantize_input_test( + backend_cls: type, hidden_size: int, pre_quant_scale: Optional[torch.Tensor] +): + """A CuteDSL sibling backend carrying only what ``quantize_input`` reads. + + ``__new__`` without ``__init__``: the real constructors want a GPU (and, + for the FC12 one, Rubin), while ``quantize_input`` reads the quant config, + the two scales and the scale geometry. The DenseGEMM backend reads the + same four, so it shares this fake. + """ + backend = backend_cls.__new__(backend_cls) + torch.nn.Module.__init__(backend) + # ``has_nvfp4`` asserts weights exist and reads the layer quant mode; + # NVFP4_AWQ folds onto the NVFP4 quant mode (quantization/mode.py), so + # either recipe clears that gate. Declare plain NVFP4 when no scale is + # given, so that case exercises a non-AWQ layer rather than an AWQ one + # whose checkpoint happened to carry no ``pre_quant_scale``. + backend._weights_created = True + backend.quant_config = QuantConfig( + quant_algo=QuantAlgo.NVFP4_AWQ if pre_quant_scale is not None else QuantAlgo.NVFP4 + ) + backend.hidden_size = hidden_size + backend.scaling_vector_size = 16 + backend.fc31_input_scale = torch.nn.Parameter(torch.tensor(0.5), requires_grad=False) + # ``create_weights`` registers the slot as None, and only a checkpoint + # that carries ``*.pre_quant_scale`` fills it, as + # ``NVFP4FusedMoEMethod.load_quant_scales`` does. + backend.register_parameter("fc31_act_scale", None) + if pre_quant_scale is not None: + backend.register_parameter( + "fc31_act_scale", torch.nn.Parameter(pre_quant_scale, requires_grad=False) + ) + return backend + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "backend_cls", + [CuteDslFusedMoE, CuteDslFc12FusedMoE, DenseGEMMFusedMoE], + ids=["grouped_gemm", "fused_fc12", "dense_gemm"], +) +@pytest.mark.parametrize("has_pre_quant_scale", [False, True], ids=["nvfp4", "nvfp4_awq"]) +def test_cutedsl_quantize_input_applies_pre_quant_scale( + monkeypatch: pytest.MonkeyPatch, backend_cls: type, has_pre_quant_scale: bool +) -> None: + """The CuteDSL siblings apply an NVFP4_AWQ / NVFP4_ARC pre-quant scale. + + All three select ``NVFP4CuteDslFusedMoEMethod``, which inherits the CUTLASS + NVFP4 pipeline and therefore loads a checkpoint's per-channel + ``pre_quant_scale`` into ``fc31_act_scale``, and all three calibrate + ``fc31_input_scale`` on the scaled activation. Their kernels consume only + the quantized activation, which leaves ``quantize_input`` as the one place + to apply it; before this it was dropped and the output silently wrong. + """ + hidden_size = 64 + pre_quant_scale = ( + torch.linspace(0.25, 4.0, hidden_size, dtype=torch.bfloat16).unsqueeze(0) + if has_pre_quant_scale + else None + ) + backend = _make_cutedsl_sibling_for_quantize_input_test( + backend_cls, hidden_size, pre_quant_scale + ) + quantizer_calls = [] + + def fake_fp4_quantize(x, global_scale, sf_vec_size, sf_use_ue8m0, is_sf_swizzled): + quantizer_calls.append((x.clone(), global_scale)) + num_rows, num_cols = x.shape + return ( + torch.zeros(num_rows, num_cols // 2, dtype=torch.uint8), + torch.zeros(num_rows * num_cols // sf_vec_size, dtype=torch.uint8), + ) + + monkeypatch.setattr(torch.ops.trtllm, "fp4_quantize", fake_fp4_quantize, raising=False) + x = torch.randn(3, hidden_size, dtype=torch.bfloat16) + + backend.quantize_input(x, post_quant_comm=False) + + [(quantized, global_scale)] = quantizer_calls + expected = x * pre_quant_scale if has_pre_quant_scale else x + torch.testing.assert_close(quantized, expected, rtol=0, atol=0) + assert global_scale is backend.fc31_input_scale + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "backend_cls", + [CuteDslFusedMoE, CuteDslFc12FusedMoE, DenseGEMMFusedMoE], + ids=["grouped_gemm", "fused_fc12", "dense_gemm"], +) +def test_cutedsl_quantize_input_rejects_prequantized_awq_input(backend_cls: type) -> None: + """An already-quantized input can no longer take the pre-quant scale. + + Each backend also accepts an ``Fp4QuantizedTensor`` produced upstream (a + fused norm+quant, or a post-quant dispatch) that never saw + ``fc31_act_scale``. Refuse that combination rather than serve the + checkpoint wrong, as ``CutlassFusedMoE.quantize_input`` does. + """ + hidden_size = 64 + backend = _make_cutedsl_sibling_for_quantize_input_test( + backend_cls, + hidden_size, + torch.linspace(0.25, 4.0, hidden_size, dtype=torch.bfloat16).unsqueeze(0), + ) + x = Fp4QuantizedTensor( + fp4_tensor=torch.zeros(3, hidden_size // 2, dtype=torch.uint8), + scaling_factor=torch.zeros(3 * hidden_size // 16, dtype=torch.uint8), + is_sf_swizzled=False, + ) + + with pytest.raises(AssertionError, match="not expected for AWQ"): + backend.quantize_input(x, post_quant_comm=False) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("has_pre_quant_scale", [False, True], ids=["nvfp4", "nvfp4_awq"]) +def test_w4a16_nvfp4_quantize_input_applies_pre_quant_scale(has_pre_quant_scale: bool) -> None: + """The CUTLASS W4A16_NVFP4 leaf applies the pre-quant scale too. + + This leaf keeps activations in high precision -- ``run_moe`` dequantizes the + NVFP4 weights and runs the unquantized kernel -- so unlike the ``nvfp4`` + leaf it does no FP4 activation quantization. A pre-quant scale is not part + of that quantization, though: ``NVFP4FusedMoEMethod.load_quant_scales`` + fills ``fc31_act_scale`` from the checkpoint's ``*.pre_quant_scale`` keys + without consulting ``quant_algo``, and the dequantized ``W / s`` fed into + the BF16 MoE expect the scaled activation. The dense counterpart, + ``W4A16NVFP4LinearMethod``, applies ``pre_quant_scale`` the same way. + """ + hidden_size = 64 + moe = TrtllmCutlassW4a16Nvfp4Impl.__new__(TrtllmCutlassW4a16Nvfp4Impl) + torch.nn.Module.__init__(moe) + # ``create_weights`` registers the slot as None, and only a checkpoint + # that carries ``*.pre_quant_scale`` fills it. + moe.register_parameter("fc31_act_scale", None) + pre_quant_scale = None + if has_pre_quant_scale: + pre_quant_scale = torch.linspace(0.25, 4.0, hidden_size, dtype=torch.bfloat16).unsqueeze(0) + moe.register_parameter( + "fc31_act_scale", torch.nn.Parameter(pre_quant_scale, requires_grad=False) + ) + x = torch.randn(3, hidden_size, dtype=torch.bfloat16) + + out, out_sf = moe.quantize_input(x) + + assert out_sf is None + expected = x * pre_quant_scale if has_pre_quant_scale else x + torch.testing.assert_close(out, expected, rtol=0, atol=0) + + +# --------------------------------------------------------------------------- +# NVFP4_AWQ on the per-expert STREAMING loader. +# +# Kimi K3 never calls the MoE backend's whole-checkpoint ``load_weights``, so +# ``load_quant_scales`` -- the only thing that read ``*.pre_quant_scale`` -- +# never runs for it. The scale therefore has to travel through +# ``load_streaming_nvfp4_expert``, or an AWQ checkpoint served over that path +# gets ``fc31_act_scale is None`` and the fix above does nothing for it. +# Nothing raises in that state: the loads are silent and only accuracy moves. +# --------------------------------------------------------------------------- + + +@pytest.mark.cpu_only +def test_kimi_k3_nvfp4_expert_spec_covers_the_awq_recipes() -> None: + """The streaming layout serves NVFP4 and both of its calibration recipes. + + ``_k3_expert_ckpt_spec`` used to know only plain ``NVFP4``, so an + NVFP4_AWQ / NVFP4_ARC checkpoint could not reach this loader at all. The + weights are byte-identical across the three; only the extra per-channel + scale differs, so one layout covers them. + + ``pre_quant_scale`` must stay OPTIONAL: ``nvidia/Kimi-K3-NVFP4`` stores + none, and requiring it would reject the one checkpoint on this path today. + """ + from tensorrt_llm._torch.models.modeling_kimi_linear import _k3_expert_ckpt_spec + + nvfp4 = _k3_expert_ckpt_spec(QuantAlgo.NVFP4) + for quant_algo in (QuantAlgo.NVFP4_AWQ, QuantAlgo.NVFP4_ARC): + assert _k3_expert_ckpt_spec(quant_algo) is nvfp4, quant_algo + assert "pre_quant_scale" not in nvfp4.kinds + assert nvfp4.optional_suffixes == ("w1.pre_quant_scale", "w3.pre_quant_scale") + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("has_pre_quant_scale", [False, True], ids=["nvfp4", "nvfp4_awq"]) +def test_kimi_k3_nvfp4_expert_loader_forwards_pre_quant_scale( + has_pre_quant_scale: bool, +) -> None: + """The per-expert loader passes a stored pre-quant scale on, and only then. + + The two getters are the seam every streaming route shares (file-grouped, + split-file and the shared-dict fallback): the required one raises on a + missing key, the optional one answers ``None``. A plain NVFP4 checkpoint + must still load, which is why ``None`` has to reach the quant method + rather than the key being demanded. + """ + from tensorrt_llm._torch.models.modeling_kimi_linear import _load_nvfp4_expert + + pre_quant_scale = torch.linspace(0.5, 2.0, 8) + stored = ( + { + "experts.5.w1.pre_quant_scale": pre_quant_scale, + "experts.5.w3.pre_quant_scale": pre_quant_scale, + } + if has_pre_quant_scale + else {} + ) + captured = {} + backend = SimpleNamespace( + quant_method=SimpleNamespace( + load_streaming_nvfp4_expert=lambda module, **kwargs: captured.update( + module=module, **kwargs + ) + ) + ) + + _load_nvfp4_expert(backend, "experts", 5, 2, lambda key: key, stored.get) + + assert captured["module"] is backend + assert captured["global_expert_id"] == 5 + assert captured["local_slot_id"] == 2 + # Required tensors still come from the required getter, unchanged. + assert captured["w1_weight"] == "experts.5.w1.weight" + assert captured["w2_input_scale"] == "experts.5.w2.input_scale" + for w in ("w1", "w3"): + got = captured[f"{w}_pre_quant_scale"] + if has_pre_quant_scale: + assert got is pre_quant_scale, w + else: + assert got is None, w + + +def _streaming_module_for_pre_quant_scale_test(hidden_size: int) -> torch.nn.Module: + """A module carrying only what the pre-quant-scale staging touches.""" + module = torch.nn.Module() + module.dtype = torch.bfloat16 + module.hidden_size = hidden_size + module.tp_size = 1 + module.tp_rank = 0 + module.is_gated_activation = True + # ``create_weights`` registers the slot as None; on the whole-checkpoint + # path ``load_quant_scales`` replaces it the moment it sees the keys. + module.register_parameter("fc31_act_scale", None) + return module + + +@pytest.mark.cpu_only +def test_nvfp4_streaming_stages_pre_quant_scales_for_finalization() -> None: + """A streamed AWQ expert's scale reaches ``fc31_act_scale``. + + Two halves, asserted together because either alone passes trivially: the + staging container has to exist before the (multi-threaded) loads, and + ``_finalize_pre_quant_scales`` has to consume it and allocate the + parameter, which on this path no ``load_quant_scales`` has created. + + The real allocation is ``device='cuda'``, so it is substituted here; what + is asserted is that the finalizer allocates at all and lands + ``max(w3, w1)``, which is the reduction the whole-checkpoint path uses. + """ + hidden_size = 8 + method = NVFP4MegaMoECuteDslMethod.__new__(NVFP4MegaMoECuteDslMethod) + module = _streaming_module_for_pre_quant_scale_test(hidden_size) + + # The base implementation, not MegaMoE's override: that one additionally + # rematerializes the raw source parameters, which needs real weights. + NVFP4FusedMoEMethod.prepare_streaming_expert_load(method, module) + assert module.tmp_pre_quant_scales == {}, "staging must predate the loader threads" + + w1 = torch.linspace(0.5, 2.0, hidden_size) + w3 = torch.linspace(2.0, 0.5, hidden_size) + for global_expert_id in (3, 7): + module.tmp_pre_quant_scales[global_expert_id] = {"w1": w1.clone(), "w3": w3.clone()} + + def allocate_on_cpu(mod: torch.nn.Module) -> None: + mod.register_parameter( + "fc31_act_scale", + torch.nn.Parameter( + torch.empty(1, mod.hidden_size, dtype=mod.dtype), requires_grad=False + ), + ) + + method._ensure_fc31_act_scale = allocate_on_cpu + method._finalize_pre_quant_scales(module) + + assert not hasattr(module, "tmp_pre_quant_scales") + torch.testing.assert_close( + module.fc31_act_scale.data, + torch.max(w1, w3).to(torch.bfloat16).unsqueeze(0), + rtol=0, + atol=0, + ) + + +@pytest.mark.cpu_only +def test_nvfp4_streaming_plain_checkpoint_leaves_act_scale_unset() -> None: + """A plain NVFP4 streamed load must stay exactly as it was. + + The container is now pre-created for every streamed NVFP4 load, AWQ or + not, so an empty one means "this checkpoint has no pre-quant scale" and + not "the scales went missing". Without this the ``nvidia/Kimi-K3-NVFP4`` + load would raise ``Missing pre_quant_scale data``. + """ + method = NVFP4MegaMoECuteDslMethod.__new__(NVFP4MegaMoECuteDslMethod) + module = _streaming_module_for_pre_quant_scale_test(8) + + # Base implementation, for the reason given in the test above. + NVFP4FusedMoEMethod.prepare_streaming_expert_load(method, module) + method._finalize_pre_quant_scales(module) + + assert module.fc31_act_scale is None + assert not hasattr(module, "tmp_pre_quant_scales") + + +def _streaming_expert_kwargs() -> dict: + """The required ``load_streaming_nvfp4_expert`` arguments, as packed bytes. + + Only the three weights are inspected before the pre-quant-scale checks, so + the scales can be placeholders: the call under test never reaches them. + """ + packed = torch.zeros(2, 2, dtype=torch.uint8) + return { + f"{w}_{kind}": packed if kind == "weight" else None + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + } + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("supplied", ["w1", "w3"]) +def test_nvfp4_streaming_rejects_half_a_pre_quant_scale_pair(supplied: str) -> None: + """One half of the pair cannot produce ``max(w3, w1)``. + + Accepting it would stage an entry the finalizer then rejects with a + message about a missing expert, one layer deep and pointing at the wrong + thing. Refused at the argument instead. + """ + method = NVFP4MegaMoECuteDslMethod.__new__(NVFP4MegaMoECuteDslMethod) + module = SimpleNamespace( + expert_size_per_partition=1, + initial_local_expert_ids=[0], + is_gated_activation=True, + _streamed_expert_slots=set(), + ) + + with pytest.raises(ValueError, match="must be loaded together"): + method.load_streaming_nvfp4_expert( + module, + global_expert_id=0, + local_slot_id=0, + **_streaming_expert_kwargs(), + **{f"{supplied}_pre_quant_scale": torch.ones(4)}, + ) + + +@pytest.mark.cpu_only +def test_nvfp4_streaming_refuses_pre_quant_scale_without_gating() -> None: + """``fc31_act_scale`` is a gated-activation concept, as on the other path. + + ``load_quant_scales`` asserts the same thing: without a gate there is no + ``max(w3, w1)`` to take, so a checkpoint in that shape is refused rather + than loaded with the scale dropped. + """ + method = NVFP4MegaMoECuteDslMethod.__new__(NVFP4MegaMoECuteDslMethod) + module = SimpleNamespace( + expert_size_per_partition=1, + initial_local_expert_ids=[0], + is_gated_activation=False, + _streamed_expert_slots=set(), + ) + + with pytest.raises(ValueError, match="non-gated activations"): + method.load_streaming_nvfp4_expert( + module, + global_expert_id=0, + local_slot_id=0, + **_streaming_expert_kwargs(), + w1_pre_quant_scale=torch.ones(4), + w3_pre_quant_scale=torch.ones(4), + ) + + def test_megamoe_cutedsl_tuning_mode_forces_top_maxt_bucket( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2817,6 +3255,119 @@ def test_megamoe_cutedsl_minimax_m3_swiglu_bias_numerics() -> None: check_accuracy(output, ref_output, rtol=0.1, atol=0.1, percent=0.94) +@pytest.mark.gpu +def test_megamoe_cutedsl_nvfp4_awq_matches_cutlass() -> None: + """MegaMoE and CUTLASS agree on an NVFP4_AWQ layer, and both match the reference. + + The checkpoint carries the per-channel ``w1`` / ``w3`` ``pre_quant_scale`` + an AWQ calibration writes, and its input scale is calibrated on + ``x * pre_quant_scale``, which is what the reference runs FC1 on. The last + check is the control: MegaMoE against the reference on the unscaled input + must fail, or this test could not tell a dropped scale apart. + """ + dtype = torch.bfloat16 + num_experts = 8 + top_k = 2 + hidden_size = 512 + intermediate_size = 512 + seq_len = 16 + + verdict = MegaMoECuteDsl.can_implement( + MoEProblem( + quant=QuantAlgo.NVFP4.value, + dtype_act=dtype, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + num_experts=num_experts, + top_k=top_k, + ), + MoEDeployment( + ep_size=1, + tp_size=1, + parallel_size=1, + use_dp=False, + num_slots=num_experts, + env=collect_moe_environment(), + ), + ) + if not verdict.eligible: + pytest.skip(verdict.detail) + + mapping = Mapping(rank=mpi_rank()) + _ensure_single_proc_dist_for_megamoe(MoeBackendType.MEGAMOE_CUTEDSL, mapping.rank) + with torch.device(f"cuda:{mapping.rank}"): + torch.manual_seed(0) + torch.cuda.manual_seed(0) + AutoTuner.get().setup_distributed_state(mapping) + routing_method = RenormalizeMoeRoutingMethod(top_k=top_k) + x = torch.randn((seq_len, hidden_size), dtype=dtype, device="cuda") + router_logits = torch.randn((seq_len, num_experts), dtype=dtype, device="cuda") + # One vector for every expert, as AWQ writes it, with per-channel + # factors in [1/e, e]: dropping it moves FC1's input far past tolerance. + pre_quant_scale = torch.empty(hidden_size, device="cuda").uniform_(-1, 1).exp().to(dtype) + x_smoothed = x * pre_quant_scale + quantize_util_cls, quant_config, quant_kwargs = get_test_quant_params( + QuantAlgo.NVFP4, x_smoothed, MoeBackendType.MEGAMOE_CUTEDSL + ) + quantize_util = quantize_util_cls( + num_experts=num_experts, + dtype=dtype, + intermediate_size=intermediate_size, + hidden_size=hidden_size, + quant_config=quant_config, + ) + ref_cls = quant_kwargs.pop("ref_cls") + weights = quantize_util.create_weights(**quant_kwargs) + for expert_id in range(num_experts): + weights[f"{expert_id}.w1.pre_quant_scale"] = pre_quant_scale + weights[f"{expert_id}.w3.pre_quant_scale"] = pre_quant_scale + token_selected_experts, token_final_scales = routing_method.apply(router_logits) + + outputs = {} + for backend_type in (MoeBackendType.CUTLASS, MoeBackendType.MEGAMOE_CUTEDSL): + backend = create_test_backend( + backend_type=backend_type, + routing_method=routing_method, + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4_AWQ), + mapping=mapping, + ) + backend.load_weights([weights]) + backend.post_load_weights() + backend.cuda() + assert backend.fc31_act_scale is not None + with torch.inference_mode(): + x_quantized, x_sf = backend.quantize_input(x, post_quant_comm=False) + outputs[backend_type] = run_backend_moe( + backend, + backend_type, + x_quantized, + x_sf, + token_selected_experts, + token_final_scales, + dtype, + router_logits, + ) + + ref_fused_moe = quantize_util.create_ref_module(routing_method, ref_cls=ref_cls) + ref_fused_moe.load_weights([weights]) + ref_fused_moe.cuda() + with torch.inference_mode(): + ref_output = ref_fused_moe(x_smoothed, router_logits) + unscaled_ref_output = ref_fused_moe(x, router_logits) + + cutlass_output = outputs[MoeBackendType.CUTLASS] + megamoe_output = outputs[MoeBackendType.MEGAMOE_CUTEDSL] + ref_fused_moe.check_accuracy(cutlass_output, ref_output) + ref_fused_moe.check_accuracy(megamoe_output, ref_output) + ref_fused_moe.check_accuracy(megamoe_output, cutlass_output) + with pytest.raises(Exception, match="Mismatch percentage"): + ref_fused_moe.check_accuracy(megamoe_output, unscaled_ref_output) + + @pytest.fixture(scope="module") def _run_megamoe_prologue_helper() -> Callable[[int, bool], torch.Tensor]: if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7): diff --git a/tests/unittest/_torch/moe/test_moe_impl.py b/tests/unittest/_torch/moe/test_moe_impl.py index eb4bf631b41b..ba5c2f0510ad 100644 --- a/tests/unittest/_torch/moe/test_moe_impl.py +++ b/tests/unittest/_torch/moe/test_moe_impl.py @@ -1347,6 +1347,7 @@ def admits( TrtllmCutlassMxfp8Impl: quantize_input_mxfp8, TrtllmCutlassW4a8Mxfp4Mxfp8Impl: quantize_input_mxfp8, TrtllmCutlassNvfp4Impl: vars(TrtllmCutlassNvfp4Impl).get("quantize_input"), + TrtllmCutlassW4a16Nvfp4Impl: vars(TrtllmCutlassW4a16Nvfp4Impl).get("quantize_input"), } _CUTLASS_PAD_TO = {TrtllmCutlassW4a16Mxfp4Impl: 128}