[None][fix] Kimi K3: bound MegaMoE expert-weight memory at EP8; drop the MoE TP/EP env overrides

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

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