From d08149a70db3944ec4c270da4919bde28577efd4 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:35:33 -0700 Subject: [PATCH 1/3] [None][fix] Kimi K3: bound MegaMoE's expert-weight memory at EP8 Two independent leaks in the same lifecycle, found one after the other -- the first fix only moved the OOM, which is how the second was found. 1. Bound the staged expert footprint. MegaMoE never overrode finalize_streamed_expert, so its staged w3_w1 halves accumulated for the whole load instead of being drained per expert the way CUTLASS does. EP8 doubles the rank-local expert count that footprint tracks (112 vs 56 at EP16), which is why this only bites here. Draining was previously blocked by the coverage check, so that accounting moves onto _streamed_expert_slots -- which also closes a documented open item. 2. Bound the mega-format transform's transient. The transform materialized the whole rank-local expert set at once; it is now chunked by slot. Both are pure memory-lifecycle changes; the transform's output is unchanged and test_mega_format_transform_is_slot_blockwise asserts that BITWISE, including an uneven final chunk. Together they are what makes the DEP8 disagg gen worker fit. Note the tests cover the bookkeeping and the transform's bitwise equivalence, but nothing here shows the served model computes correctly at EP8 -- that is what the GSM8K-through-the-proxy run is for. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../_torch/modules/fused_moe/quantization.py | 200 ++++++++++++++---- .../test_lists/test-db/l0_b300.yml | 3 + .../modules/moe/test_kimi_k3_situ_moe.py | 187 ++++++++++++++++ 3 files changed, 349 insertions(+), 41 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 48e972354dc8..9970fc84fecd 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -3818,6 +3818,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 +3852,78 @@ 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. + """ + # 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) def _get_fc2_alpha_input_scale( self, @@ -4243,9 +4322,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 +4467,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 +4758,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..ded1f34b2270 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 @@ -1746,3 +1746,190 @@ 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, + 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) + + 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}" + + 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") From 8c4c38f89c6c56525b8035f86203d0a04e725dca Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Tue, 18 Aug 2026 21:39:50 -0700 Subject: [PATCH 2/3] [None][fix] Kimi K3: drop the MoE TP/EP env overrides Parallelism is configuration, not an environment knob: TLLM_K3_MOE_TP_SIZE / TLLM_K3_MOE_EP_SIZE let the routed-expert split be changed out from under a served config, with no record of it in that config. Review feedback on !10481. _select_moe_tp_ep now has two levels instead of three -- 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 moe_tp=tp_size default is not mistaken for a TP request), else the historical EP-only layout. Behaviour is unchanged for anyone who was not setting the env vars, which nothing in-tree was. test_kimi_k3_moe_split_selection loses its monkeypatch fixture along with the two env cases; the config-driven cases it already covered are the whole surface now. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 33 +++---- .../modules/moe/test_kimi_k3_situ_moe.py | 90 +++++++++---------- 2 files changed, 51 insertions(+), 72 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 86631f8a97ac..7ac22534b569 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,13 +129,13 @@ 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" +_KDA_INDEXED_STATE_POOL_ENABLED = os.environ.get("TLLM_KDA_ENABLE_INDEXED_STATE_POOL", "1") == "1" +# Heuristic ported from SGLang's Blackwell cutoff: +# https://github.com/sgl-project/sglang/blob/e84bbf68efb683c9e2eef4168c5198042544599d/python/sglang/srt/models/kimi_k3.py#L946-L954 +# It has not been tuned for TensorRT-LLM; benchmark and retune it for TRT-LLM's +# projection kernels. Verify intentionally counts B * num_steps because those +# flattened token rows form the projection GEMMs' M dimension. +_KDA_BFA_MULTISTREAM_MAX_ROWS = 128 if TYPE_CHECKING: from transformers import PretrainedConfig @@ -1238,27 +1237,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/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py index ded1f34b2270..f234e0723ffc 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): @@ -1767,9 +1753,7 @@ def test_megamoe_streamed_coverage_survives_per_expert_drain() -> None: Pure bookkeeping, so it needs no GPU and no EP rendezvous. """ - from tensorrt_llm._torch.modules.fused_moe.quantization import ( - NVFP4MegaMoECuteDslMethod, - ) + from tensorrt_llm._torch.modules.fused_moe.quantization import NVFP4MegaMoECuteDslMethod n_slots = 4 method = NVFP4MegaMoECuteDslMethod.__new__(NVFP4MegaMoECuteDslMethod) @@ -1795,10 +1779,10 @@ def _fake_module() -> SimpleNamespace: 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, + "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 @@ -1808,7 +1792,7 @@ def _fake_module() -> SimpleNamespace: 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 + 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. @@ -1817,16 +1801,18 @@ def _fake_module() -> SimpleNamespace: 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) + (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) + (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} + 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) + 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: @@ -1843,6 +1829,8 @@ def test_megamoe_overrides_finalize_streamed_expert() -> None: from tensorrt_llm._torch.modules.fused_moe.quantization import ( NVFP4CutlassFusedMoEMethod, NVFP4FusedMoEMethod, + ) + from tensorrt_llm._torch.modules.fused_moe.quantization import ( NVFP4MegaMoECuteDslMethod as MegaMoE, ) @@ -1852,8 +1840,11 @@ def test_megamoe_overrides_finalize_streamed_expert() -> None: ) assert issubclass(MegaMoE, NVFP4FusedMoEMethod) - for name in ("finalize_streamed_expert", "_resolve_staged_w3_w1_weight", - "_resolve_staged_w3_w1_weight_scale"): + 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}" scale_src = inspect.getsource(MegaMoE._resolve_staged_w3_w1_weight_scale) @@ -1862,7 +1853,8 @@ def test_megamoe_overrides_finalize_streamed_expert() -> None: ) # 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) + NVFP4CutlassFusedMoEMethod._resolve_staged_w3_w1_weight_scale + ) @nvfp4_moe_supported @@ -1897,17 +1889,17 @@ def test_mega_format_transform_is_slot_blockwise() -> None: torch.manual_seed(31) raw_w3_w1 = torch.randint( - 0, 255, (num_slots, expand_intermediate, h_bytes), - dtype=torch.uint8, device="cuda") + 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") + 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") + 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") + 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( @@ -1924,12 +1916,12 @@ def _run(lo: int, hi: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, to 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")): + 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}") + 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") + f"chunk={chunk} {name} differs from the whole-layer transform" + ) From cb66c9f4500e6c0aa4c1f984a0fe29ce984110b3 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Sun, 23 Aug 2026 20:14:51 -0700 Subject: [PATCH 3/3] [None][refactor] deduplicate NVFP4 streamed expert finalization Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 8 -- .../_torch/modules/fused_moe/quantization.py | 76 ++++++++++--------- .../modules/moe/test_kimi_k3_situ_moe.py | 7 ++ 3 files changed, 46 insertions(+), 45 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 7ac22534b569..dad0673ad018 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -129,14 +129,6 @@ os.environ.get("TLLM_K3_DISABLE_MIN_LATENCY_LATENT_PROJ", "0") == "1" ) -_KDA_INDEXED_STATE_POOL_ENABLED = os.environ.get("TLLM_KDA_ENABLE_INDEXED_STATE_POOL", "1") == "1" -# Heuristic ported from SGLang's Blackwell cutoff: -# https://github.com/sgl-project/sglang/blob/e84bbf68efb683c9e2eef4168c5198042544599d/python/sglang/srt/models/kimi_k3.py#L946-L954 -# It has not been tuned for TensorRT-LLM; benchmark and retune it for TRT-LLM's -# projection kernels. Verify intentionally counts B * num_steps because those -# flattened token rows form the projection GEMMs' M dimension. -_KDA_BFA_MULTISTREAM_MAX_ROWS = 128 - if TYPE_CHECKING: from transformers import PretrainedConfig diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 9970fc84fecd..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 @@ -3906,24 +3920,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 _get_fc2_alpha_input_scale( self, 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 f234e0723ffc..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 @@ -1839,6 +1839,7 @@ def test_megamoe_overrides_finalize_streamed_expert() -> None: "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", @@ -1847,6 +1848,12 @@ def test_megamoe_overrides_finalize_streamed_expert() -> None: ): 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"