[None][feat] Dense DWDP for Kimi-K3 context servers and a balanced DWDP expert partition

scripts/generate_trtllm_serve_schemas.py

@@ -145,6 +145,12 @@ def generate_serve_schema() -> dict:
     schema["$defs"]["MoeConfig"]["properties"]["load_balancer"] = {
         "anyOf": [_schema_for(MoeLoadBalancerConfig, schema), {"type": ["string", "null"]}]
     }
+    # Mirror DwdpConfig._validate_expert_partition(): both fields 0 (or absent), or both positive.
+    pair = ("num_experts_per_worker", "num_prefetch_experts")
+    schema["$defs"]["DwdpConfig"]["anyOf"] = [
+        {"properties": {name: {"const": 0} for name in pair}},
+        {"properties": {name: {"minimum": 1} for name in pair}, "required": list(pair)},
+    ]
     # These object/Any annotations are Python API hooks, not arbitrary YAML mappings.
     # Update these restrictions when hooks are added or gain YAML construction support.
     for name in ("batched_logits_processor", "checkpoint_loader", "_mpi_session"):

tensorrt_llm/_torch/models/modeling_kimi_k25.py

@@ -1730,7 +1730,18 @@ def mm_token_ids(self) -> torch.Tensor:
         return self._mm_token_ids
 
     def load_weights(self, weights) -> None:
-        """Load vision + projector + LLM weights from checkpoint."""
+        """Load vision + projector + LLM weights from checkpoint.
+
+        Raises:
+            RuntimeError: Dense DWDP is active on the text model; nothing is loaded.
+        """
+        if getattr(getattr(self, "model", None), "_dense_dwdp_active", False):
+            # The text model's dense weights are views of dense DWDP staging slots that several
+            # layers share; refuse before the vision tower is rebuilt or loaded.
+            raise RuntimeError(
+                f"{type(self).__name__} load_weights is not supported while dense DWDP is "
+                "active: the dense weights live in staging slots shared across layers"
+            )
         # Create vision encoder if it was skipped during meta init.
         # Use the saved VLM config since model_config.pretrained_config
         # has been overwritten with text_config by __init__.

tensorrt_llm/_torch/models/modeling_kimi_linear.py

@@ -121,6 +121,7 @@
 )
 from ..distributed.symm_mem_allgather import SymmetricMemoryAllGather
 from ..model_config import ModelConfig
+from ..modules.dwdp.dense.manager import get_global_dense_dwdp_manager
 from ..modules.gated_mlp import GatedMLP
 from ..modules.kimi_kda import KimiKDALinearAttention
 from ..modules.kimi_kda.kimi_k3_mamba_metadata import KimiK3MambaMetadata
@@ -2070,13 +2071,66 @@ def forward(
         return out
 
 
+# ---------------------------------------------------------------------------
+# Dense DWDP: the tables its setup reads off every decoder layer.
+# ---------------------------------------------------------------------------
+
+# Plain attributes (a tensor, or a tuple of tensors) under the dense roots that dense DWDP may
+# move as weights: the fused KDA projections, decode constants and the low-M GEMM's scale copies.
+# Parameters and buffers need no entry.
+DWDP_DENSE_ATTR_ALLOW: frozenset[str] = frozenset(
+    {
+        "_low_m_weight_scale",
+        "_bfa_proj_weight",
+        "_qkvg_proj_weight",
+        "_gate_b_t",
+        "_packed_conv_weight",
+        "_mtp_conv_weights",
+        "_w_q_t",
+        "_w_k_t",
+        "_w_v_t",
+        "_A_log_f32",
+        "_dt_bias_f32",
+        "_onorm_w_f32",
+    }
+)
+
+# Names that never make dense DWDP move a storage: per-call scratch tensors and the routing bias.
+# A reference under one of these names to a storage that moves through another reference is still
+# repointed with it.
+DWDP_DENSE_EXCLUDE: frozenset[str] = frozenset(
+    {
+        "_o_dense",
+        "_fused_q_placeholder",
+        "softmax_stats_tensor",
+        "temp_softmax_stats_tensor",
+        "e_score_correction_bias",
+    }
+)
+
+_DENSE_DWDP_SEGMENTS = ("attn", "moe")
+# Children of the MoE block that hold its dense weights: the router, the shared expert and the
+# latent projections and norm. The routed experts belong to expert DWDP and are never a dense root.
+_DENSE_DWDP_MOE_ROOTS = (
+    "gate",
+    "shared_experts",
+    "routed_expert_down_proj",
+    "routed_expert_up_proj",
+    "routed_expert_norm",
+)
+
+
 # ---------------------------------------------------------------------------
 # Decoder layer.
 # ---------------------------------------------------------------------------
 
 
 @nccl_window_tensor_scoped_module
 class KimiLinearDecoderLayer(nn.Module):
+    # Read by dense DWDP setup next to ``dense_dwdp_roots``.
+    dense_dwdp_attr_allow: frozenset[str] = DWDP_DENSE_ATTR_ALLOW
+    dense_dwdp_exclude: frozenset[str] = DWDP_DENSE_EXCLUDE
+
     def __init__(
         self,
         model_config: ModelConfig,
@@ -2194,6 +2248,44 @@ def __init__(
         self.self_attention_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype)
         self.mlp_res_proj = nn.Linear(cfg.hidden_size, 1, bias=False, dtype=dtype)
 
+        # The dense DWDP manager exists before the model is built (DwdpManager registers it with
+        # expert DWDP, py_executor_creator when dense DWDP runs standalone); forward() reaches the
+        # dense DWDP runtime through it.
+        self._dense_dwdp = get_global_dense_dwdp_manager()
+
+    def dense_dwdp_roots(self, segment: str) -> list[tuple[str, nn.Module]]:
+        """The modules holding this layer's dense DWDP weights of ``segment``.
+
+        Implements ``modules.dwdp.dense.types.DenseRootsProvider``. Dense DWDP moves the
+        replicated weights reachable from these roots (see ``DWDP_DENSE_ATTR_ALLOW`` and
+        ``DWDP_DENSE_EXCLUDE``) and reads them back only inside the windows ``forward`` opens.
+
+        Args:
+            segment: ``"attn"`` for the attention weights, ``"moe"`` for the MoE-dense weights:
+                the router, the shared expert and the latent projections. The routed experts
+                are never a root.
+
+        Returns:
+            ``(path relative to this layer, module)`` pairs in walk order; [] when the layer's
+            weights stay replicated: a layer without the MoE block (layer 0 of Kimi-K3) or one
+            whose weights were removed (layer-wise benchmarks).
+
+        Raises:
+            ValueError: ``segment`` is neither ``"attn"`` nor ``"moe"``.
+        """
+        if segment not in _DENSE_DWDP_SEGMENTS:
+            raise ValueError(
+                f"Kimi K3 dense DWDP segment must be one of {_DENSE_DWDP_SEGMENTS}; got {segment!r}"
+            )
+        if not self.is_moe or not _has_weights(self):
+            return []
+        if segment == "attn":
+            if self.is_kda:
+                return [("linear_attn", self.linear_attn)]
+            return [("self_attn", self.self_attn)]
+        moe = self.block_sparse_moe
+        return [(f"block_sparse_moe.{name}", getattr(moe, name)) for name in _DENSE_DWDP_MOE_ROOTS]
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -2229,6 +2321,17 @@ def forward(
         def _full_rows(t: torch.Tensor) -> torch.Tensor:
             return t if sliced is None else sliced.all_gather(t, num_tokens)
 
+        # Dense DWDP read windows: every read of this layer's attention weights happens between
+        # H1 (wait_layer) and H2 (attention_done), every read of its MoE-dense weights (router,
+        # shared expert, latent projections) between H1 and H3 (layer_done). Their staging slots
+        # are refilled with a later layer's weights once a window closes, so a read outside its
+        # window sees another layer's bytes.
+        manager = getattr(self, "_dense_dwdp", None)
+        dense = manager.dense if manager is not None else None
+        if dense is not None and not dense.covers(self.layer_idx):
+            dense = None
+        if dense is not None:
+            dense.wait_layer(self.layer_idx)  # H1
         prefix_sum = hidden_states
         valid_block_residual = block_residual[:num_snapshots]
 
@@ -2275,6 +2378,8 @@ def _full_rows(t: torch.Tensor) -> torch.Tensor:
                 hidden_states = sliced.reduce_scatter(hidden_states, num_tokens)
         else:
             hidden_states = self.self_attn(hidden_states, attn_metadata)
+        if dense is not None:
+            dense.attention_done(self.layer_idx)  # H2: no attention weight is read after this.
 
         if prefix_sum is None:
             prefix_sum = hidden_states
@@ -2303,6 +2408,8 @@ def _full_rows(t: torch.Tensor) -> torch.Tensor:
             )
         else:
             hidden_states = self.mlp(hidden_states)
+        if dense is not None:
+            dense.layer_done(self.layer_idx)  # H3: no MoE-dense weight is read after this.
 
         prefix_sum = prefix_sum + hidden_states
         return prefix_sum, num_snapshots
@@ -2394,6 +2501,11 @@ def __init__(self, model_config: ModelConfig):
             key="kimi_k3_aux_capture_mode",
         )
 
+        # The dense DWDP manager exists before the model is built (DwdpManager registers it with
+        # expert DWDP, py_executor_creator when dense DWDP runs standalone); forward() primes dense
+        # DWDP through it.
+        self._dense_dwdp = get_global_dense_dwdp_manager()
+
     def forward(
         self,
         attn_metadata: AttentionMetadata,
@@ -2406,6 +2518,14 @@ def forward(
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
 
+        # Dense DWDP H0: make sure the first two layers' dense weights (and, with expert DWDP,
+        # their experts) are enqueued before any layer runs (the previous forward's end or the
+        # non-overlap executor loop usually already did). After a forward that stopped midway, a
+        # WAR barrier on this stream comes first.
+        manager = getattr(self, "_dense_dwdp", None)
+        if manager is not None and manager.dense is not None:
+            manager.ensure_primed()
+
         if inputs_embeds is None:
             inputs_embeds = self.embed_tokens(input_ids)
         hidden_states = inputs_embeds
@@ -3113,6 +3233,13 @@ def checkpoint_name_plan(
         return name_map, expected_keys, expert_jobs
 
     def load_weights(self, weights: Dict[str, torch.Tensor]) -> None:
+        if getattr(self.model, "_dense_dwdp_active", False):
+            # The dense weights are views of staging slots that several layers share: a load
+            # would write into them, and the next gather overwrites that from the row store.
+            raise RuntimeError(
+                "Kimi K3 load_weights is not supported while dense DWDP is active: the dense "
+                "weights live in staging slots shared across layers"
+            )
         prefix = "language_model." if any(k.startswith("language_model.") for k in weights) else ""
         params = self._trunk_parameters()
         name_map, expected_keys, expert_jobs = self.checkpoint_name_plan(prefix)

tensorrt_llm/_torch/modules/dwdp/__init__.py

@@ -28,20 +28,31 @@
   exchange; WeightManager owns scheduling
 - One group API (allgather + Barrier) served by an MPI sub-communicator, or by
   a TCPStore rendezvous for independently launched workers -- see rendezvous.py
+
+Dense DWDP (``dense/``) shards the replicated dense weights of the MoE layers the
+same way.  Its exports (``DenseDWDP``, ``setup_dense_dwdp``,
+``dense_dwdp_requested``) are resolved by the module-level ``__getattr__`` on
+first use, so importing this package, which every import of a ``dense`` module
+does first, loads no dense module.  Code inside ``dense`` imports from the
+defining submodule, never from this package.
 """
 
+import importlib
+from typing import TYPE_CHECKING
+
 from .page_pool import PagePool, compute_slot_sizes
 from .rendezvous import StoreDwdpGroup, dwdp_global_rank, store_rendezvous_enabled
 from .setup import setup_dwdp, teardown_dwdp
 from .specs import (
     EdgeInfo,
+    ExpertPartition,
     LayerWeightSpecs,
     MnnvlHandleSet,
     PageAlignedLayout,
     PeerRange,
     PeerRanges,
     WeightSpec,
-    compute_peer_ranges,
+    handle_geometry,
     lookup_owner,
 )
 from .transport import DWDPTransport
@@ -57,6 +68,17 @@
 from .weight_buffer import WeightBuffer
 from .weight_manager import DWDPWeightManager
 
+if TYPE_CHECKING:
+    from .dense.config import dense_dwdp_requested
+    from .dense.runtime import DenseDWDP
+    from .dense.setup import setup_dense_dwdp
+
+_DENSE_EXPORTS = {
+    "DenseDWDP": ".dense.runtime",
+    "setup_dense_dwdp": ".dense.setup",
+    "dense_dwdp_requested": ".dense.config",
+}
+
 __all__ = [
     # Data classes
     "WeightSpec",
@@ -66,7 +88,8 @@
     "PageAlignedLayout",
     "PeerRange",
     "PeerRanges",
-    "compute_peer_ranges",
+    "ExpertPartition",
+    "handle_geometry",
     "lookup_owner",
     # VMM utilities
     "align_up",
@@ -90,4 +113,18 @@
     # Setup orchestration
     "setup_dwdp",
     "teardown_dwdp",
+    # Dense DWDP
+    "DenseDWDP",
+    "setup_dense_dwdp",
+    "dense_dwdp_requested",
 ]
+
+
+def __getattr__(name: str) -> object:
+    """Resolve a dense DWDP export on first use and cache it on this module."""
+    module_name = _DENSE_EXPORTS.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(importlib.import_module(module_name, __name__), name)
+    globals()[name] = value
+    return value

tensorrt_llm/_torch/modules/dwdp/dense/__init__.py

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dense DWDP: row-sharded replicated dense weights, gathered two layers ahead of their use.
+
+Every rank of the dense group keeps one row of each MoE decoder layer's attention blob ``A(l)``
+and MoE-dense blob ``D(l)``; the runtime gathers ``L(l+2) = A(l+2) + D(l+2)`` into small staging
+slots on a copy stream.  It runs in one of two modes (``types.DenseMode``): with expert DWDP
+(attached) the dense group is the DWDP group and the copies share the expert DWDP copy stream;
+standalone, on an attention-DP instance without expert DWDP, the dense group is the attention-DP
+(tp) group and the copies run on the dense manager's own stream.  Submodules:
+
+* ``config`` -- ``TRTLLM_DWDP_DENSE*`` switches.
+* ``types`` -- shared records, layout, slots, copy plans, CUDA seams, constants, errors.
+* ``layout`` -- collector, checks, blob layout, checksums.
+* ``repoint`` -- repointing module tensors into slots, rollback, release check, teardown.
+* ``runtime`` -- ``DenseDWDP``: slot rings, events and the per-layer hooks.
+* ``plan`` -- per-layer copy plans and their launcher.
+* ``store`` -- the sharded row store (FABRIC and in-process loopback).
+* ``watchdog`` -- the slot checksum watchdog of ``TRTLLM_DWDP_DENSE_DEBUG=2``.
+* ``setup`` -- setup phases P1-P6 over five fixed collectives.
+* ``manager`` -- ``DenseDwdpManager``: owns the committed runtime in both modes, primes it (H0)
+  and is what the model's hooks find through ``get_global_dense_dwdp_manager``; standalone it
+  also builds the group and the transfer and runs setup.
+* ``gate`` -- the launch gate ``py_executor_creator`` runs before the model is built; it picks
+  the mode.
+
+Import the submodules directly; this package imports none of them.
+"""