diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 86631f8a97ac..dad0673ad018 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -45,8 +45,7 @@ (``intermediate / moe_tp_size`` per rank; group-32 MXFP4 packed bytes and scales sliced consistently by the stock TRTLLM-Gen quant-method loaders). The split is EP-only unless the user sets ``moe_tensor_parallel_size`` / -``moe_expert_parallel_size`` explicitly (or the ``TLLM_K3_MOE_TP_SIZE`` / -``TLLM_K3_MOE_EP_SIZE`` env overrides). Routing is computed replicated; the +``moe_expert_parallel_size`` explicitly. Routing is computed replicated; the routed partial sums — EP partials of whole experts, or TP partials over the intermediate shards — are all-reduced in the latent space (before ``routed_expert_norm`` / ``routed_expert_up_proj``, which are @@ -130,14 +129,6 @@ os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" ) -# Routed-expert MoE TP/EP split overrides (read per model init, not import). -# Highest precedence; either one may be set alone, the other is derived from -# tp_size. Without them, an explicit moe_tensor_parallel_size / -# moe_expert_parallel_size pair from the user config is honored, and the -# default stays EP-only (moe_ep == tp_size). -_K3_MOE_TP_ENV = "TLLM_K3_MOE_TP_SIZE" -_K3_MOE_EP_ENV = "TLLM_K3_MOE_EP_SIZE" - if TYPE_CHECKING: from transformers import PretrainedConfig @@ -1238,27 +1229,15 @@ def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: Precedence: - 1. ``TLLM_K3_MOE_TP_SIZE`` / ``TLLM_K3_MOE_EP_SIZE`` env overrides - (either alone; the other is derived from ``tp_size``). - 2. Explicit ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` + 1. Explicit ``moe_tensor_parallel_size`` / ``moe_expert_parallel_size`` from the user config. Detected via ``mapping.moe_tp_ep_user_specified`` so the auto-resolved mapping default (``moe_tp=tp_size, moe_ep=1``) is NOT mistaken for a TP request. - 3. Default: EP-only (``moe_tp=1, moe_ep=tp_size``), the historical + 2. Default: EP-only (``moe_tp=1, moe_ep=tp_size``), the historical K3 layout. """ tp_size = mapping.tp_size - env_tp = os.environ.get(_K3_MOE_TP_ENV) - env_ep = os.environ.get(_K3_MOE_EP_ENV) - if env_tp is not None or env_ep is not None: - moe_tp = int(env_tp) if env_tp is not None else 0 - moe_ep = int(env_ep) if env_ep is not None else 0 - if moe_tp <= 0 and moe_ep > 0: - moe_tp = tp_size // moe_ep - elif moe_ep <= 0 and moe_tp > 0: - moe_ep = tp_size // moe_tp - return moe_tp, moe_ep if getattr(mapping, "moe_tp_ep_user_specified", False): return mapping.moe_tp_size, mapping.moe_ep_size return 1, tp_size diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 48e972354dc8..07463e8550ae 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -19,7 +19,7 @@ import threading from abc import ABC, abstractmethod from enum import Enum, auto -from typing import Dict, List, NamedTuple, Optional, Tuple, Union +from typing import Callable, Dict, List, NamedTuple, Optional, Tuple, Union import torch import torch.nn.functional as F @@ -2398,6 +2398,32 @@ def prepare_streaming_expert_load(self, module: torch.nn.Module) -> None: module.tmp_weight_scale_2 = {} module._streamed_expert_slots = set() + def _finalize_staged_w3_w1_expert( + self, + module: torch.nn.Module, + local_slot_id: int, + weight_resolver: Callable[[Dict[str, torch.Tensor]], None], + weight_scale_resolver: Callable[[Dict[str, torch.Tensor]], None], + ) -> None: + """Drain one expert's staged w3/w1 weights and block scales.""" + # The key must be built exactly as the staging site builds it. The two + # sites do not agree on the accessor -- the weight one uses + # untyped_storage(), the scale one the deprecated storage() -- so + # mirror each rather than assume they return the same address. + for attr, dst_base, resolve in ( + ('tmp_cutlass_w3_w1_weights', lambda: module.w3_w1_weight.data[ + local_slot_id].untyped_storage().data_ptr(), weight_resolver), + ('tmp_cutlass_w3_w1_weight_scales', + lambda: module.w3_w1_weight_scale.data[local_slot_id].storage( + ).data_ptr(), weight_scale_resolver), + ): + staged = getattr(module, attr, None) + if not staged: + continue + entry = staged.pop((dst_base(), local_slot_id), None) + if entry is not None: + resolve(entry) + def finalize_streamed_expert(self, module: torch.nn.Module, local_slot_id: int) -> None: """Resolve any per-expert staging left by ``load_streaming_nvfp4_expert``. @@ -3221,24 +3247,12 @@ def finalize_streamed_expert(self, module: torch.nn.Module, ``dict.pop`` is atomic, so concurrent loader threads draining different slots need no lock. """ - # The key must be built exactly as the staging site builds it. The two - # sites do not agree on the accessor -- the weight one uses - # untyped_storage(), the scale one the deprecated storage() -- so - # mirror each rather than assume they return the same address. - for attr, dst_base, resolve in ( - ('tmp_cutlass_w3_w1_weights', lambda: module.w3_w1_weight.data[ - local_slot_id].untyped_storage().data_ptr(), - self._resolve_staged_w3_w1_weight), - ('tmp_cutlass_w3_w1_weight_scales', - lambda: module.w3_w1_weight_scale.data[local_slot_id].storage( - ).data_ptr(), self._resolve_staged_w3_w1_weight_scale), - ): - staged = getattr(module, attr, None) - if not staged: - continue - entry = staged.pop((dst_base(), local_slot_id), None) - if entry is not None: - resolve(entry) + self._finalize_staged_w3_w1_expert( + module, + local_slot_id, + self._resolve_staged_w3_w1_weight, + self._resolve_staged_w3_w1_weight_scale, + ) def process_weights_after_loading(self, module: torch.nn.Module): # Finalize w3_w1 weights: cat + pad. Streamed loads drain these @@ -3818,6 +3832,22 @@ class NVFP4MegaMoECuteDslMethod(NVFP4FusedMoEMethod): weight_dtype = FUSED_MOE_NVFP4_WEIGHT_DTYPE block_scales_dtype = FUSED_MOE_NVFP4_WEIGHT_BLOCK_SCALE_DTYPE + # How many routed slots the MegaMoE-format transform converts at a time. + # + # The transform materializes several contiguous copies of whatever it is + # handed, so a whole-layer call peaks at a multiple of that layer's routed + # weights on top of the already-allocated destinations. Bounding the call + # bounds the peak; it does not change the result, because the transform is + # per-slot independent (see _build_mega_format_weights). + # + # 16 is a compromise, not a tuned value: small enough that the transient + # stays well under a GiB at EP8 (a whole layer there is ~1.15 GiB, which is + # precisely the allocation that failed in job 489557), large enough that a + # 112-slot layer is 7 calls rather than 112. Lower it if a future topology + # is tighter; there is no correctness constraint on the value, and + # test_mega_format_transform_is_slot_blockwise covers uneven final chunks. + MEGA_FORMAT_SLOT_CHUNK = 16 + def prepare_streaming_expert_load(self, module: torch.nn.Module) -> None: # Like the Cutlass child this backend stages the w3_w1 halves, and it # additionally tracks which w2 rows a partial load covered. Every one of @@ -3836,15 +3866,66 @@ def prepare_streaming_expert_load(self, module: torch.nn.Module) -> None: module._streamed_w2_covered = set() module._streamed_w2_scale_covered = set() - # NOTE: this backend deliberately does NOT override - # finalize_streamed_expert, unlike its Cutlass sibling. Its staged - # dicts are not only staging: _initial_slot_coverage() COUNTS their - # entries to decide which slots a partial load actually populated, so - # draining them per expert would report zero coverage and make a - # complete load look empty. Bounding the staged footprint here means - # moving that accounting off the dicts first (e.g. onto - # _streamed_expert_slots, which already tracks exactly this) rather - # than copying the Cutlass drain over. + def _resolve_staged_w3_w1_weight(self, entry: Dict[str, + torch.Tensor]) -> None: + """Cat + pad one staged expert's w3_w1 halves into its destination.""" + w3 = entry.get('w3') + w1 = entry.get('w1') + dst = entry['dst'] + if w3 is not None and w1 is not None: + cat_weight = torch.cat([w3, w1], dim=0) + cat_weight = self._maybe_padding_shape(cat_weight, dst) + dst.copy_(cat_weight, non_blocking=True) + + def _resolve_staged_w3_w1_weight_scale( + self, entry: Dict[str, torch.Tensor]) -> None: + """Cat + pad one staged expert's w3_w1 block scales. + + Deliberately **no** ``block_scale_interleave``, unlike the Cutlass + sibling's resolver of the same name: MegaMoE's kernel does its own + 16-atom gate/up interleave and ``to_blocked`` swizzle later, in + ``_build_mega_format_weights``. + """ + w3_scale = entry.get('w3') + w1_scale = entry.get('w1') + dst = entry['dst'] + if w3_scale is not None and w1_scale is not None: + cat_scale = torch.cat([w3_scale, w1_scale], dim=0) + cat_scale = self._maybe_padding_shape(cat_scale, dst) + dst.copy_(cat_scale) + + def finalize_streamed_expert(self, module: torch.nn.Module, + local_slot_id: int) -> None: + """Resolve just this slot's staged halves, as soon as it is loaded. + + Without this the staged halves -- a **second copy** of the routed-expert + weights -- accumulate across the whole load, because a loader that + groups its work by shard FILE (Kimi K3 does) finishes a given layer only + when the last of its slots happens to land, which can be arbitrarily + late. That is survivable at EP16 (56 rank-local experts) and is not at + EP8 (112): the DEP8 disagg generation worker OOM'd inside + ``setup_engine`` having had the whole card free beforehand. + + 🔴 This override was previously blocked, and the note that stood here + said why: ``_streamed_coverage`` used to decide which slots a load had + populated by **counting staged dict entries**, so draining them per + expert would have reported zero coverage and turned a complete load into + a spurious "partially covered" error. That accounting now reads + ``_streamed_expert_slots`` for the streamed case -- which records the + same fact and survives draining -- so the two concerns are no longer + entangled. Keep them unentangled: anything that needs to know what a + streamed load covered should ask ``_streamed_expert_slots``, not the + staging dicts. + + ``dict.pop`` is atomic, so concurrent loader threads draining different + slots need no lock. + """ + self._finalize_staged_w3_w1_expert( + module, + local_slot_id, + self._resolve_staged_w3_w1_weight, + self._resolve_staged_w3_w1_weight_scale, + ) def _get_fc2_alpha_input_scale( self, @@ -4243,9 +4324,27 @@ def _streamed_coverage(self, module: torch.nn.Module) -> Dict[str, int]: A w3_w1 stash entry counts only when BOTH halves arrived; the direct-copy w2 paths are tracked via row-pointer sets. Routed slots are told apart from EPLB shared staging by storage base. + + 🔴 **Streamed loads must not be counted from the stashes.** + ``finalize_streamed_expert`` drains each slot's stash entry as soon as + it lands, to keep the staged second copy of the routed weights bounded, + so by the time this runs the stashes are empty and counting them would + report 0/N on a load that was in fact complete. ``_streamed_expert_slots`` + records exactly the same fact -- it is added to at the end of + ``load_streaming_nvfp4_expert``, i.e. only once both halves and w2 have + been handed over -- and it survives draining. + + The stash count remains correct, and is still used, for a + whole-checkpoint load: that path never streams, never drains per slot, + and leaves ``_streamed_expert_slots`` empty. """ + streamed_slots = getattr(module, '_streamed_expert_slots', None) def _stash_covered(stash_name: str, param) -> int: + if streamed_slots: + # Both w3_w1 components are staged in the same + # load_streaming_nvfp4_expert call, so one slot set covers both. + return len(streamed_slots) base = param.data.storage().data_ptr() stash = getattr(module, stash_name, {}) return sum(1 for (b, _idx), e in stash.items() @@ -4370,31 +4469,23 @@ def process_weights_after_loading(self, module: torch.nn.Module): # * module.local_shared_w3_w1_tensors contains cat'd [w3|w1] per shared slot # _maybe_padding_shape is a defensive no-op against future # alignment drift on the grandparent get_weights_shapes. + # Streamed loads drain these per expert in finalize_streamed_expert, so + # this handles whatever is left -- everything for a whole-checkpoint + # load, nothing for a fully streamed one. if hasattr(module, 'tmp_cutlass_w3_w1_weights'): for entry in module.tmp_cutlass_w3_w1_weights.values(): - w3 = entry.get('w3') - w1 = entry.get('w1') - dst = entry['dst'] - if w3 is not None and w1 is not None: - cat_weight = torch.cat([w3, w1], dim=0) - cat_weight = self._maybe_padding_shape(cat_weight, dst) - dst.copy_(cat_weight, non_blocking=True) + self._resolve_staged_w3_w1_weight(entry) delattr(module, 'tmp_cutlass_w3_w1_weights') # ---- Cat raw w3+w1 scales (NO block_scale_interleave) ---- # Same routed + shared cat pattern as weights. MegaMoE's kernel # does its own 16-atom gate/up interleave + to_blocked swizzle # in _build_mega_format_weights below; the Cutlass parent would - # call block_scale_interleave here, which we deliberately skip. + # call block_scale_interleave here, which we deliberately skip -- + # see _resolve_staged_w3_w1_weight_scale, which both paths share. if hasattr(module, 'tmp_cutlass_w3_w1_weight_scales'): for entry in module.tmp_cutlass_w3_w1_weight_scales.values(): - w3_scale = entry.get('w3') - w1_scale = entry.get('w1') - dst = entry['dst'] - if w3_scale is not None and w1_scale is not None: - cat_scale = torch.cat([w3_scale, w1_scale], dim=0) - cat_scale = self._maybe_padding_shape(cat_scale, dst) - dst.copy_(cat_scale) + self._resolve_staged_w3_w1_weight_scale(entry) delattr(module, 'tmp_cutlass_w3_w1_weight_scales') # ---- Build EPLB shared-staging mega buffers ---- @@ -4669,27 +4760,56 @@ def _build_mega_format_weights(self, module: torch.nn.Module): (the transform pipeline itself). """ - # CPU-staged reload sources: upload ONE layer transiently so the + # CPU-staged reload sources: upload ONE slot chunk transiently so the # pack pipeline runs on GPU (see _materialize_source_params). def _on_cuda(t: torch.Tensor) -> torch.Tensor: return t if t.is_cuda else t.cuda() - mega_fc1, mega_fc1_sf, mega_fc2, mega_fc2_sf = ( - self._build_mega_format_buffers( - raw_w3_w1=_on_cuda(module.w3_w1_weight.data), - raw_w3_w1_sf=_on_cuda(module.w3_w1_weight_scale.data), - raw_w2=_on_cuda(module.w2_weight.data), - raw_w2_sf=_on_cuda(module.w2_weight_scale.data), - num_slots=module.expert_size_per_partition, - intermediate=module.intermediate_size_per_partition, - hidden=module.hidden_size, - expand_intermediate=module. - expand_intermediate_size_per_partition, - )) - module.mega_fc1_weight.data.copy_(mega_fc1, non_blocking=True) - module.mega_fc1_weight_sf.data.copy_(mega_fc1_sf, non_blocking=True) - module.mega_fc2_weight.data.copy_(mega_fc2, non_blocking=True) - module.mega_fc2_weight_sf.data.copy_(mega_fc2_sf, non_blocking=True) + # 🔴 Run the transform in slot chunks rather than on the whole layer. + # + # _build_mega_format_buffers materializes several contiguous copies of + # what it is given (gate_part, up_part, the stacked interleave, and the + # final view), so a whole-layer call peaks at a multiple of that layer's + # routed weights ON TOP of the destination buffers, which are already + # allocated. At EP16 that fits; at EP8 a layer's rank-local experts are + # twice as many and it does not -- the DEP8 disagg gen worker OOM'd + # right here (job 489557), asking for 1.15 GiB with ~50 MiB free. + # + # Chunking is exactly equivalent, not an approximation: the transform is + # a pure function documented for "any slot count", and every step in it + # is per-slot (dim-0) independent -- slicing dim 1, viewing with + # num_slots, stacking on an inner axis. Slot i's output depends only on + # slot i's input, which test_mega_format_transform_is_slot_blockwise + # asserts bitwise rather than leaving to inspection. The transform + # itself is deliberately NOT modified: its 16-atom gate/up interleave is + # exactly the kind of layout code where an error is silent. + n_slots = module.expert_size_per_partition + for lo in range(0, n_slots, self.MEGA_FORMAT_SLOT_CHUNK): + hi = min(lo + self.MEGA_FORMAT_SLOT_CHUNK, n_slots) + mega_fc1, mega_fc1_sf, mega_fc2, mega_fc2_sf = ( + self._build_mega_format_buffers( + raw_w3_w1=_on_cuda(module.w3_w1_weight.data[lo:hi]), + raw_w3_w1_sf=_on_cuda( + module.w3_w1_weight_scale.data[lo:hi]), + raw_w2=_on_cuda(module.w2_weight.data[lo:hi]), + raw_w2_sf=_on_cuda(module.w2_weight_scale.data[lo:hi]), + num_slots=hi - lo, + intermediate=module.intermediate_size_per_partition, + hidden=module.hidden_size, + expand_intermediate=module. + expand_intermediate_size_per_partition, + )) + module.mega_fc1_weight.data[lo:hi].copy_(mega_fc1, + non_blocking=True) + module.mega_fc1_weight_sf.data[lo:hi].copy_(mega_fc1_sf, + non_blocking=True) + module.mega_fc2_weight.data[lo:hi].copy_(mega_fc2, + non_blocking=True) + module.mega_fc2_weight_sf.data[lo:hi].copy_(mega_fc2_sf, + non_blocking=True) + # Drop this chunk's intermediates before building the next one; + # otherwise the peak is unchanged and the chunking buys nothing. + del mega_fc1, mega_fc1_sf, mega_fc2, mega_fc2_sf def _build_mega_shared_staging(self, module: torch.nn.Module): """Allocate + populate CPU shared-staging tensors for the four diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 5ec0ade36a7b..057cb2dfb6ca 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -56,6 +56,9 @@ l0_b300: - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_cutlass_situ_bf16_matches_reference - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streamed_experts_match_situ_reference[static_1.0] + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_megamoe_streamed_coverage_survives_per_expert_drain + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_megamoe_overrides_finalize_streamed_expert + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_mega_format_transform_is_slot_blockwise # ------------- MoE: test_moe_backend (by backend) --------------- - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py index 8b564563a92e..e3fca03ca9d3 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py @@ -37,12 +37,7 @@ import tensorrt_llm._torch.models.modeling_kimi_linear as modeling_kimi_linear from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_kimi_linear import ( - _K3_MOE_EP_ENV, - _K3_MOE_TP_ENV, - KimiK3MoEGate, - KimiK3MoERuntime, -) +from tensorrt_llm._torch.models.modeling_kimi_linear import KimiK3MoEGate, KimiK3MoERuntime from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory from tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_deepgemm import ( _MEGA_MOE_SYMM_BUFFER_CACHE, @@ -475,28 +470,19 @@ def test_mapping_records_moe_tp_ep_user_specified(): assert ep.moe_tp_size == 1 and ep.moe_ep_size == 8 -def test_kimi_k3_moe_split_selection(monkeypatch): - monkeypatch.delenv(_K3_MOE_TP_ENV, raising=False) - monkeypatch.delenv(_K3_MOE_EP_ENV, raising=False) - +def test_kimi_k3_moe_split_selection() -> None: # Auto mapping default stays EP-only (the historical K3 layout), even # though the resolved mapping says moe_tp=8. auto = Mapping(world_size=8, tp_size=8) assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (1, 8) - # Explicit pure-TP and hybrid requests are honored. + # Explicit pure-TP and hybrid requests are honored. The MoE split is only + # ever set through the config; there is no env override. tp = Mapping(world_size=8, tp_size=8, moe_tp_size=8, moe_ep_size=1) assert KimiK3MoERuntime._select_moe_tp_ep(tp) == (8, 1) tep = Mapping(world_size=8, tp_size=8, moe_tp_size=4, moe_ep_size=2) assert KimiK3MoERuntime._select_moe_tp_ep(tep) == (4, 2) - # Env override wins; a single side derives the other from tp_size. - monkeypatch.setenv(_K3_MOE_TP_ENV, "4") - assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2) - monkeypatch.delenv(_K3_MOE_TP_ENV) - monkeypatch.setenv(_K3_MOE_EP_ENV, "2") - assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2) - @pytest.mark.parametrize("backend", ["CUTLASS", "TRTLLM", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"]) def test_kimi_k3_routed_config_preserves_explicit_backend(backend): @@ -1746,3 +1732,203 @@ def test_fp8_placeholder_fill_contract(): single.load_checkpoint_pair([pairs[0]]) ref = FP8Linear.from_checkpoint_fp8(pairs[0][0], pairs[0][1], outs[0]) assert _bitwise_equal(single.weight, ref.weight) + + +def test_megamoe_streamed_coverage_survives_per_expert_drain() -> None: + """MegaMoE's coverage check must not be defeated by draining the stashes. + + ``finalize_streamed_expert`` drains each slot's staged w3_w1 halves as soon + as it lands, so the second copy of the routed-expert weights stays bounded + by the experts in flight. That is what makes EP8 fit: at EP16 the unbounded + staging survived on 56 rank-local experts, at EP8 it is 112 and the DEP8 + disagg gen worker OOM'd inside setup_engine with the whole card free + beforehand. + + The drain was previously blocked because ``_streamed_coverage`` decided + which slots a load had populated by COUNTING staged dict entries -- so + draining would report 0/N and turn a complete load into a spurious + "partially covered" error. This asserts the two are now disentangled: + coverage comes from ``_streamed_expert_slots`` for a streamed load, which + records the same fact and survives draining. + + Pure bookkeeping, so it needs no GPU and no EP rendezvous. + """ + from tensorrt_llm._torch.modules.fused_moe.quantization import NVFP4MegaMoECuteDslMethod + + n_slots = 4 + method = NVFP4MegaMoECuteDslMethod.__new__(NVFP4MegaMoECuteDslMethod) + + def _fake_module() -> SimpleNamespace: + m = SimpleNamespace() + m.w3_w1_weight = SimpleNamespace(data=torch.zeros(n_slots, 8, 4, dtype=torch.uint8)) + m.w3_w1_weight_scale = SimpleNamespace(data=torch.zeros(n_slots, 8, 4, dtype=torch.uint8)) + m.w2_weight = SimpleNamespace(data=torch.zeros(n_slots, 4, 4, dtype=torch.uint8)) + m.w2_weight_scale = SimpleNamespace(data=torch.zeros(n_slots, 4, 4, dtype=torch.uint8)) + # w2 is written through directly; its coverage is row-pointer based and + # is unaffected by this change. Mark every row covered. + m._streamed_w2_covered = {m.w2_weight.data[i].data_ptr() for i in range(n_slots)} + m._streamed_w2_scale_covered = { + m.w2_weight_scale.data[i].data_ptr() for i in range(n_slots) + } + return m + + # ---- streamed load, stashes fully drained: must report FULL coverage. + drained = _fake_module() + drained.tmp_cutlass_w3_w1_weights = {} + drained.tmp_cutlass_w3_w1_weight_scales = {} + drained._streamed_expert_slots = set(range(n_slots)) + cov = method._streamed_coverage(drained) + assert cov == { + "w3_w1_weight": n_slots, + "w3_w1_weight_scale": n_slots, + "w2_weight": n_slots, + "w2_weight_scale": n_slots, + }, f"drained streamed load must read as fully covered, got {cov}" + + # ---- a genuinely partial streamed load must still read as partial, so the + # drain does not turn the check into a rubber stamp. + partial = _fake_module() + partial.tmp_cutlass_w3_w1_weights = {} + partial.tmp_cutlass_w3_w1_weight_scales = {} + partial._streamed_expert_slots = {0, 1} + cov = method._streamed_coverage(partial) + assert cov["w3_w1_weight"] == 2 and cov["w3_w1_weight_scale"] == 2, cov + + # ---- whole-checkpoint load (never streams, never drains): the stash count + # is still the source of truth, and a half-staged entry still does not count. + whole = _fake_module() + whole._streamed_expert_slots = set() + wbase = whole.w3_w1_weight.data.storage().data_ptr() + sbase = whole.w3_w1_weight_scale.data.storage().data_ptr() + whole.tmp_cutlass_w3_w1_weights = { + (wbase, i): {"w1": 1, "w3": 1, "dst": None} for i in range(n_slots) + } + whole.tmp_cutlass_w3_w1_weight_scales = { + (sbase, i): {"w1": 1, "w3": 1, "dst": None} for i in range(n_slots - 1) + } + whole.tmp_cutlass_w3_w1_weight_scales[(sbase, n_slots - 1)] = {"w1": 1, "dst": None} + cov = method._streamed_coverage(whole) + assert cov["w3_w1_weight"] == n_slots, cov + assert cov["w3_w1_weight_scale"] == n_slots - 1, ( + "a stash entry missing its w3 half must not count as covered", + cov, + ) + + +def test_megamoe_overrides_finalize_streamed_expert() -> None: + """MegaMoE must actually drain, and must not interleave when it does. + + Two things a reader could get wrong by copying the Cutlass sibling: it is a + SIBLING, not a parent, so nothing is inherited; and its scale resolver must + NOT call block_scale_interleave -- MegaMoE's kernel does its own 16-atom + gate/up interleave and to_blocked swizzle in _build_mega_format_weights, so + interleaving here would apply it twice. + """ + import inspect + + from tensorrt_llm._torch.modules.fused_moe.quantization import ( + NVFP4CutlassFusedMoEMethod, + NVFP4FusedMoEMethod, + ) + from tensorrt_llm._torch.modules.fused_moe.quantization import ( + NVFP4MegaMoECuteDslMethod as MegaMoE, + ) + + assert not issubclass(MegaMoE, NVFP4CutlassFusedMoEMethod), ( + "MegaMoE is a sibling of the Cutlass method, not a child; if this ever " + "changes, re-check which staging behaviour it inherits" + ) + assert issubclass(MegaMoE, NVFP4FusedMoEMethod) + assert "_finalize_staged_w3_w1_expert" in NVFP4FusedMoEMethod.__dict__ + + for name in ( + "finalize_streamed_expert", + "_resolve_staged_w3_w1_weight", + "_resolve_staged_w3_w1_weight_scale", + ): + assert name in MegaMoE.__dict__, f"MegaMoE must define its own {name}" + + for method in (NVFP4CutlassFusedMoEMethod, MegaMoE): + finalize_src = inspect.getsource(method.finalize_streamed_expert) + assert "self._finalize_staged_w3_w1_expert" in finalize_src, ( + f"{method.__name__} must delegate staged draining to the shared helper" + ) + + scale_src = inspect.getsource(MegaMoE._resolve_staged_w3_w1_weight_scale) + assert "_interleave_w3_w1_weight_scale" not in scale_src, ( + "MegaMoE's scale resolver must not interleave; the kernel does that itself" + ) + # The Cutlass sibling's does, which is what makes the distinction load-bearing. + assert "_interleave_w3_w1_weight_scale" in inspect.getsource( + NVFP4CutlassFusedMoEMethod._resolve_staged_w3_w1_weight_scale + ) + + +@nvfp4_moe_supported +def test_mega_format_transform_is_slot_blockwise() -> None: + """Chunking the MegaMoE-format transform must be exactly equivalent. + + ``_build_mega_format_weights`` runs the transform in slot chunks so its + transient stays bounded: the transform materializes several contiguous + copies of whatever it is handed, and at EP8 a whole layer's routed weights + made that peak overflow the card (job 489557 asked for 1.15 GiB with ~50 + MiB free). Chunking only helps if slot i's output depends on slot i's input + and nothing else. + + That property is asserted here BITWISE rather than argued from the code, + because the transform's 16-atom gate/up interleave is exactly the kind of + layout operation whose errors are silent -- shapes stay right, values go + wrong, and only an accuracy run notices. + + Includes an uneven final chunk (a prime slot count against the chunk size), + which is the case a divisible-only test would miss. + """ + from tensorrt_llm._torch.modules.fused_moe.quantization import ( + NVFP4MegaMoECuteDslMethod as MegaMoE, + ) + + method = MegaMoE.__new__(MegaMoE) + + hidden, intermediate = 256, 64 + expand_intermediate = 2 * intermediate + num_slots = 7 # deliberately not a multiple of the chunk size below + h_bytes = hidden // 2 + + torch.manual_seed(31) + raw_w3_w1 = torch.randint( + 0, 255, (num_slots, expand_intermediate, h_bytes), dtype=torch.uint8, device="cuda" + ) + raw_w2 = torch.randint( + 0, 255, (num_slots, hidden, intermediate // 2), dtype=torch.uint8, device="cuda" + ) + raw_w3_w1_sf = torch.randint( + 0, 255, (num_slots, expand_intermediate, hidden // 16), dtype=torch.uint8, device="cuda" + ) + raw_w2_sf = torch.randint( + 0, 255, (num_slots, hidden, intermediate // 16), dtype=torch.uint8, device="cuda" + ) + + def _run(lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return method._build_mega_format_buffers( + raw_w3_w1=raw_w3_w1[lo:hi], + raw_w3_w1_sf=raw_w3_w1_sf[lo:hi], + raw_w2=raw_w2[lo:hi], + raw_w2_sf=raw_w2_sf[lo:hi], + num_slots=hi - lo, + intermediate=intermediate, + hidden=hidden, + expand_intermediate=expand_intermediate, + ) + + whole = _run(0, num_slots) + + for chunk in (1, 3, 16): # 16 > num_slots: the single-call degenerate case + pieces = [_run(lo, min(lo + chunk, num_slots)) for lo in range(0, num_slots, chunk)] + for i, name in enumerate(("mega_fc1", "mega_fc1_sf", "mega_fc2", "mega_fc2_sf")): + stitched = torch.cat([p[i] for p in pieces], dim=0) + assert stitched.shape == whole[i].shape, ( + f"chunk={chunk} {name}: {stitched.shape} vs {whole[i].shape}" + ) + assert _bitwise_equal(stitched, whole[i]), ( + f"chunk={chunk} {name} differs from the whole-layer transform" + )