diff --git a/scripts/generate_trtllm_serve_schemas.py b/scripts/generate_trtllm_serve_schemas.py index a4e3522d8dd1..8599ab9026a3 100644 --- a/scripts/generate_trtllm_serve_schemas.py +++ b/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"): diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index 7d34e6fc768c..3ea038d8da34 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/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__. diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 5a1b6c41c544..5dbbea992d4c 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/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,6 +2071,55 @@ 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. # --------------------------------------------------------------------------- @@ -2077,6 +2127,10 @@ def forward( @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) diff --git a/tensorrt_llm/_torch/modules/dwdp/__init__.py b/tensorrt_llm/_torch/modules/dwdp/__init__.py index b029c68a35f3..882a15bc2693 100644 --- a/tensorrt_llm/_torch/modules/dwdp/__init__.py +++ b/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 diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/__init__.py b/tensorrt_llm/_torch/modules/dwdp/dense/__init__.py new file mode 100644 index 000000000000..9ccf832f4be5 --- /dev/null +++ b/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. +""" diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/config.py b/tensorrt_llm/_torch/modules/dwdp/dense/config.py new file mode 100644 index 000000000000..bed4e8e240ea --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/config.py @@ -0,0 +1,281 @@ +# 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. +"""Environment switches of dense DWDP; context workers only, no ``llm_args`` field. + +* ``TRTLLM_DWDP_DENSE`` -- turn dense DWDP on; default off. With ``dwdp_config`` it runs with + expert DWDP over the DWDP group; without, standalone over the attention-DP group of an instance + with ``enable_attention_dp`` and ``tp_size > 1`` (``gate.validate_dense_gate``). +* ``TRTLLM_DWDP_DENSE_SEGMENTS`` -- ``attn,moe`` (default) or ``attn`` for attention only. +* ``TRTLLM_DWDP_DENSE_REQUIRED`` -- default on: a failed setup raises on every rank; off: every rank + drops dense DWDP and keeps serving with replicated dense weights. +* ``TRTLLM_DWDP_DENSE_MIN_BYTES`` -- storages smaller than this stay replicated; default 1 MiB. +* ``TRTLLM_DWDP_DENSE_STORE_CHUNK_MB`` -- 0 (default): one store handle; > 0: whole layers per + handle of about this many MiB. +* ``TRTLLM_DWDP_DENSE_STRICT`` -- turn demand fetches, leaks, unknown large attributes, shared + storages and slot watchdog mismatches into errors; default off. It must be the same on every + rank: when the ranks differ, setup fails on every rank, which raises unless + ``TRTLLM_DWDP_DENSE_REQUIRED`` is off. +* ``TRTLLM_DWDP_DENSE_DEBUG`` -- 1: fill the slots with NaN at commit, and name the holders of + the first 8 moved storages that outlive the commit (each a scan of every live Python object); + 2: also run the per-layer slot checksum watchdog (each layer's H1 checks its slots against the + layer's checksum; slow: every slot is read once more and the host waits for the GPU once per + forward); default 0. +* ``TRTLLM_DWDP_DENSE_COPY`` -- the transfer of standalone dense DWDP, on the dense manager's own + copy stream: ``batched`` (default; one ``cudaMemcpyBatchAsync`` per layer on the copy engines), + ``per_slice`` or ``gather`` (the SM fabric gather, which holds SMs next to the MoE all-to-all; + validated on Rubin with ``TRTLLM_DWDP_GATHER_KERNEL=tma`` and + ``TLLM_ROUTING_DISABLE_COOP_LAUNCH=1``). The gate rejects an invalid value standalone. With + expert DWDP the dense copies follow the expert transfer, and the gate warns that any value, + valid or not, is ignored. + +Booleans accept ``1/true/yes/on`` and ``0/false/no/off`` (any case); unset or empty means the +default. A flag that defaults to off is on only for an explicit true value; any other value logs +a warning and keeps the default. Integers that do not parse, or are negative, log a warning and +keep the default. An invalid segment list raises. ``TRTLLM_DWDP_DENSE_COPY`` is read as set and +validated only for standalone dense DWDP (:meth:`DenseConfig.standalone_copy_mode`). + +Standalone dense DWDP is meant for context-only servers. The gate does not check the server role, +which a worker sees as ``TRTLLM_DISAGG_ROLE`` only when ``trtllm-serve`` launches it with a role +(``--server_role``, or a ctx or gen server in a disaggregated config). An aggregated or generation +attention-DP launch with ``TRTLLM_DWDP_DENSE=1`` that meets the gate's other requirements +therefore passes it and pays the per-forward gather of every covered layer on every step. +""" + +from __future__ import annotations + +import os +from collections.abc import Iterable, Mapping +from dataclasses import dataclass + +from tensorrt_llm.logger import logger + +from .types import ( + DEBUG_FILL_BYTE, + DEFAULT_MIN_BYTES, + LOG_PREFIX, + SEGMENT_ORDER, + CopyMode, + DigestSwitches, + SegmentKind, +) + +DENSE_ENV = "TRTLLM_DWDP_DENSE" +DENSE_SEGMENTS_ENV = "TRTLLM_DWDP_DENSE_SEGMENTS" +DENSE_REQUIRED_ENV = "TRTLLM_DWDP_DENSE_REQUIRED" +DENSE_MIN_BYTES_ENV = "TRTLLM_DWDP_DENSE_MIN_BYTES" +DENSE_STORE_CHUNK_MB_ENV = "TRTLLM_DWDP_DENSE_STORE_CHUNK_MB" +DENSE_STRICT_ENV = "TRTLLM_DWDP_DENSE_STRICT" +DENSE_DEBUG_ENV = "TRTLLM_DWDP_DENSE_DEBUG" +DENSE_COPY_ENV = "TRTLLM_DWDP_DENSE_COPY" + +_TRUE_VALUES = ("1", "true", "yes", "on") +_FALSE_VALUES = ("0", "false", "no", "off") + + +def _bool_env(environ: Mapping[str, str], name: str, default: bool) -> bool: + raw = environ.get(name) + if raw is None or not raw.strip(): + return default + value = raw.strip().lower() + if value in _TRUE_VALUES: + return True + if value in _FALSE_VALUES: + return False + logger.warning( + f"{LOG_PREFIX} {name}={raw!r} is not one of 1/true/yes/on or 0/false/no/off; " + f"using {int(default)}" + ) + return default + + +def _int_env(environ: Mapping[str, str], name: str, default: int) -> int: + raw = environ.get(name) + if raw is None or not raw.strip(): + return default + try: + value = int(raw.strip()) + except ValueError: + value = -1 + if value < 0: + logger.warning( + f"{LOG_PREFIX} {name}={raw!r} is not a non-negative integer; using {default}" + ) + return default + return value + + +def parse_segments(value: str) -> tuple[SegmentKind, ...]: + """Parse a ``TRTLLM_DWDP_DENSE_SEGMENTS`` value. + + Args: + value: Comma-separated segment names, any case, surrounding blanks ignored, repeats + allowed; e.g. ``"attn,moe"`` or ``"attn"``. + + Returns: + The segments in canonical order (``attn`` before ``moe``). + + Raises: + ValueError: An unknown name, an empty list, or a list without ``attn``: the attention + segment drives the per-layer prefetch, so it is always enabled. + """ + names = [name.strip().lower() for name in value.split(",") if name.strip()] + known = {kind.value for kind in SEGMENT_ORDER} + unknown = sorted(set(names) - known) + if unknown or SegmentKind.ATTN.value not in names: + raise ValueError( + f"{DENSE_SEGMENTS_ENV}={value!r} must list attn and optionally moe " + f"(comma-separated); unknown: {unknown}" + ) + return tuple(kind for kind in SEGMENT_ORDER if kind.value in names) + + +def parse_copy_mode(value: str) -> CopyMode: + """Parse a ``TRTLLM_DWDP_DENSE_COPY`` value. + + Args: + value: ``batched``, ``per_slice`` or ``gather``, any case, surrounding blanks ignored. + + Returns: + The transfer. + + Raises: + ValueError: Any other value. + """ + try: + return CopyMode(value.strip().lower()) + except ValueError: + raise ValueError( + f"{DENSE_COPY_ENV}={value!r} must be one of batched, per_slice, gather" + ) from None + + +def dense_dwdp_requested(environ: Mapping[str, str] | None = None) -> bool: + """Whether ``TRTLLM_DWDP_DENSE`` asks for dense DWDP; reads no other variable. + + Args: + environ: Environment to read; ``os.environ`` if None. + """ + return _bool_env(os.environ if environ is None else environ, DENSE_ENV, False) + + +@dataclass(frozen=True) +class DenseConfig: + """Dense DWDP switches, normally read once with :meth:`from_env`. + + Attributes: + enabled: ``TRTLLM_DWDP_DENSE``. + segments: ``TRTLLM_DWDP_DENSE_SEGMENTS`` in canonical order; always starts with ``ATTN``. + Strings are accepted and converted. + required: ``TRTLLM_DWDP_DENSE_REQUIRED``. + min_bytes: ``TRTLLM_DWDP_DENSE_MIN_BYTES``. + store_chunk_mb: ``TRTLLM_DWDP_DENSE_STORE_CHUNK_MB``. + strict: ``TRTLLM_DWDP_DENSE_STRICT``. + debug: ``TRTLLM_DWDP_DENSE_DEBUG``. + copy: ``TRTLLM_DWDP_DENSE_COPY`` as set, blanks stripped and case kept; None when unset or + blank. Kept as given and validated only by :meth:`standalone_copy_mode`, so reading + the switches never fails because of it: with expert DWDP it is ignored. + """ + + enabled: bool = False + segments: tuple[SegmentKind, ...] = SEGMENT_ORDER + required: bool = True + min_bytes: int = DEFAULT_MIN_BYTES + store_chunk_mb: int = 0 + strict: bool = False + debug: int = 0 + copy: str | None = None + + def __post_init__(self) -> None: + segments = _canonical_segments(self.segments) + object.__setattr__(self, "segments", segments) + for name in ("min_bytes", "store_chunk_mb", "debug"): + if getattr(self, name) < 0: + raise ValueError( + f"DenseConfig.{name} must be non-negative, got {getattr(self, name)}" + ) + + @classmethod + def from_env(cls, environ: Mapping[str, str] | None = None) -> DenseConfig: + """Read every ``TRTLLM_DWDP_DENSE*`` variable. + + ``TRTLLM_DWDP_DENSE_COPY`` is stored without validation (see ``copy``), so it never makes + this raise. + + Args: + environ: Environment to read; ``os.environ`` if None. + + Raises: + ValueError: ``TRTLLM_DWDP_DENSE_SEGMENTS`` is invalid (see :func:`parse_segments`). + """ + env = os.environ if environ is None else environ + raw_segments = env.get(DENSE_SEGMENTS_ENV) + if raw_segments is None or not raw_segments.strip(): + segments = SEGMENT_ORDER + else: + segments = parse_segments(raw_segments) + copy = (env.get(DENSE_COPY_ENV) or "").strip() + return cls( + enabled=_bool_env(env, DENSE_ENV, False), + segments=segments, + required=_bool_env(env, DENSE_REQUIRED_ENV, True), + min_bytes=_int_env(env, DENSE_MIN_BYTES_ENV, DEFAULT_MIN_BYTES), + store_chunk_mb=_int_env(env, DENSE_STORE_CHUNK_MB_ENV, 0), + strict=_bool_env(env, DENSE_STRICT_ENV, False), + debug=_int_env(env, DENSE_DEBUG_ENV, 0), + copy=copy or None, + ) + + def standalone_copy_mode(self) -> CopyMode: + """The transfer of standalone dense DWDP: ``copy`` parsed, ``BATCHED`` when it is None. + + Only standalone code calls it (the gate and ``DenseDwdpManager.setup``); with expert DWDP + the dense copies follow the expert transfer whatever ``copy`` holds. + + Raises: + ValueError: ``copy`` is not one of batched, per_slice, gather (see + :func:`parse_copy_mode`). + """ + if self.copy is None: + return CopyMode.BATCHED + return parse_copy_mode(self.copy) + + @property + def slot_fill_byte(self) -> int: + """Byte the slots are filled with at commit: 0, or ``DEBUG_FILL_BYTE`` (NaN) if debug.""" + return DEBUG_FILL_BYTE if self.debug >= 1 else 0 + + @property + def slot_watchdog(self) -> bool: + """Whether the per-layer slot checksum watchdog runs (``TRTLLM_DWDP_DENSE_DEBUG >= 2``).""" + return self.debug >= 2 + + def digest_switches(self, num_ranks: int) -> DigestSwitches: + """The switches the layout depends on, for a group of ``num_ranks`` ranks.""" + return DigestSwitches( + segments=self.segments, + min_bytes=self.min_bytes, + store_chunk_mb=self.store_chunk_mb, + num_ranks=num_ranks, + ) + + +def _canonical_segments(segments: Iterable[SegmentKind | str]) -> tuple[SegmentKind, ...]: + try: + kinds = {SegmentKind(kind) for kind in segments} + except ValueError as exc: + raise ValueError(f"DenseConfig.segments has an unknown segment: {exc}") from None + if SegmentKind.ATTN not in kinds: + raise ValueError("DenseConfig.segments must include the attention segment") + return tuple(kind for kind in SEGMENT_ORDER if kind in kinds) diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/gate.py b/tensorrt_llm/_torch/modules/dwdp/dense/gate.py new file mode 100644 index 000000000000..0c19128bbc7e --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/gate.py @@ -0,0 +1,200 @@ +# 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 launch gate: reject what dense DWDP cannot serve before the model exists. + +``py_executor_creator`` calls :func:`validate_dense_gate` next to its DWDP checks, before it +constructs any DWDP or dense DWDP manager and the model, so an unsupported launch fails before the +weights load rather than in setup or in the first forward. The gate also decides how dense DWDP +runs: with expert DWDP (attached) or standalone on an attention-DP instance. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NoReturn + +from tensorrt_llm._utils import mpi_disabled +from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend +from tensorrt_llm.logger import logger + +from .. import fabric_gather, vmm +from .config import DENSE_COPY_ENV, DENSE_ENV, DenseConfig +from .types import LOG_PREFIX, CopyMode, DenseMode + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + from tensorrt_llm.mapping import Mapping + +#: The two launch shapes dense DWDP runs in; every mode rejection names both. +_MODES = ( + "runs either with expert DWDP (dwdp_config with dwdp_size > 1, tp_size=1, " + "disable_overlap_scheduler=True) or standalone on an attention-DP instance (no dwdp_config, " + "enable_attention_dp=True, tp_size > 1)" +) + + +def validate_dense_gate(llm_args: TorchLlmArgs, mapping: Mapping) -> DenseMode: + """Reject launches dense DWDP cannot serve, before the model is built; return the mode. + + Called only when ``dense_dwdp_requested()``. The checks, in this order: + + 1. Valid ``TRTLLM_DWDP_DENSE*`` switches: ``DenseConfig.from_env()`` reads them all, so an + invalid ``TRTLLM_DWDP_DENSE_SEGMENTS`` fails here rather than in setup P1, and the + warnings about values it ignores are logged before the load as well. + ``TRTLLM_DWDP_DENSE_COPY`` is read but not checked here (step 3). + 2. The mode. With ``dwdp_config``, attached: ``dwdp_size > 1`` (the dense rows are sharded + over the DWDP group), ``tp_size == 1`` and the overlap scheduler off, as expert DWDP + requires. Without, standalone: ``enable_attention_dp`` and ``tp_size > 1``; every rank of + the instance then holds the same dense weights, so the rows are sharded over the + attention-DP (tp) group, and the overlap scheduler may be on. Any other launch is + rejected with a message that names both modes, including expert DWDP over an + attention-DP instance (``dwdp_config`` with ``tp_size > 1``), which dense DWDP does + not run with. + 3. ``TRTLLM_DWDP_DENSE_COPY``: standalone it must be valid + (``DenseConfig.standalone_copy_mode``); with expert DWDP the dense copies follow the + expert transfer, so any value, valid or not, is ignored with one warning. + 4. In both modes: + + * ``cuda_graph_config`` and ``torch_compile_config`` None, ``prefill_cuda_graph_backend`` + disabled: the per-layer hooks need eager single-stream execution, since they keep + host-side slot state and order the copies and the weight reads with stream events in + every forward. + * ``sleep_config`` None: the moved weights live in the staging slots and the row store, + outside the memory pools that sleep and wake-up manage. + * ``lora_config`` and ``speculative_config`` None: they read layer weights outside the + layer's read window. + * ``enable_locality_domains`` off: locality domains make ``Linear`` and ``MLA`` keep + ``_locality_domain_*`` shard copies of their weights, which dense DWDP does not move, + so they would stay replicated. + * ``mapping.pp_size`` and ``cp_size`` 1: every rank of the group holds every layer. + * FABRIC peer handles (``vmm.peer_handle_type()``): the row store is shared only through + FABRIC handles. + 5. Standalone only: the MPI backend (``TLLM_DISABLE_MPI`` unset), since the dense group is an + MPI sub-communicator over the attention-DP group; and with ``TRTLLM_DWDP_DENSE_COPY=gather`` + a registered ``trtllm::dwdp_fabric_gather``. A standalone gather launch that passes logs + one warning, which names the settings the gather was validated with. + + Args: + llm_args: The launch configuration. + mapping: The rank mapping, with the DWDP fields filled in. + + Returns: + ``DenseMode.ATTACHED`` with ``dwdp_config``, else ``DenseMode.STANDALONE``. + + Raises: + ValueError: Naming the first violated requirement. + """ + try: + config = DenseConfig.from_env() + except ValueError as exc: + _reject(f"valid {DENSE_ENV}_* switches", str(exc)) + mode = _launch_mode(llm_args, mapping) + copy_mode = None + if mode is DenseMode.STANDALONE: + try: + copy_mode = config.standalone_copy_mode() + except ValueError as exc: + _reject(f"valid {DENSE_ENV}_* switches", str(exc)) + elif config.copy is not None: + logger.warning( + f"{LOG_PREFIX} {DENSE_COPY_ENV}={config.copy!r} is ignored with expert DWDP: the " + f"dense copies run on the expert DWDP copy stream with the expert transfer" + ) + _require_unset("cuda_graph_config", llm_args.cuda_graph_config) + _require_unset("torch_compile_config", llm_args.torch_compile_config) + prefill_graphs = llm_args.prefill_cuda_graph_backend + if prefill_graphs != PrefillCudaGraphBackend.DISABLED: + _reject( + f"prefill_cuda_graph_backend='{PrefillCudaGraphBackend.DISABLED}'", + f"prefill_cuda_graph_backend='{prefill_graphs}'", + ) + _require_unset("sleep_config", llm_args.sleep_config) + _require_unset("lora_config", llm_args.lora_config) + _require_unset("speculative_config", llm_args.speculative_config) + if llm_args.enable_locality_domains: + _reject("enable_locality_domains=False", "enable_locality_domains=True") + for name in ("pp_size", "cp_size"): + size = getattr(mapping, name) + if size != 1: + _reject(f"{name}=1", f"{name}={size}") + handle_type = vmm.peer_handle_type() + if handle_type != vmm.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC: + _reject( + "FABRIC peer handles (CU_MEM_HANDLE_TYPE_FABRIC)", + f"this host shares {handle_type.name} handles", + ) + if mode is DenseMode.STANDALONE: + _check_standalone(copy_mode) + return mode + + +def _launch_mode(llm_args: TorchLlmArgs, mapping: Mapping) -> DenseMode: + """The mode the launch asks for (step 2); rejects a launch that fits neither.""" + dwdp_config = llm_args.dwdp_config + if dwdp_config is not None: + if dwdp_config.dwdp_size <= 1: + _reject_mode(f"dwdp_config with dwdp_size={dwdp_config.dwdp_size}") + if mapping.tp_size != 1: + _reject_mode(f"dwdp_config with tp_size={mapping.tp_size}") + if not llm_args.disable_overlap_scheduler: + _reject_mode("dwdp_config with disable_overlap_scheduler=False") + return DenseMode.ATTACHED + if not mapping.enable_attention_dp: + _reject_mode("no dwdp_config and enable_attention_dp=False") + if mapping.tp_size <= 1: + _reject_mode(f"no dwdp_config and tp_size={mapping.tp_size}") + return DenseMode.STANDALONE + + +def _check_standalone(copy_mode: CopyMode) -> None: + """The standalone-only requirements (step 5), after the common ones.""" + if mpi_disabled(): + _reject( + "the MPI backend (the standalone dense group is an MPI sub-communicator over the " + "attention-DP group)", + "TLLM_DISABLE_MPI=1", + ) + if copy_mode is not CopyMode.GATHER: + return + if fabric_gather.gather_op() is None: + _reject( + f"trtllm::dwdp_fabric_gather for {DENSE_COPY_ENV}=gather", + "this build does not register it", + ) + logger.warning( + f"{LOG_PREFIX} {DENSE_COPY_ENV}=gather without expert DWDP: the gather kernel holds SMs " + f"next to the MoE all-to-all. It was validated on a Rubin DEP16 context server with " + f"TRTLLM_DWDP_GATHER_KERNEL=tma and TLLM_ROUTING_DISABLE_COOP_LAUNCH=1, the latter of " + f"which the setup applies unless it is set: cooperative routing waits for the SMs the " + f"gather holds" + ) + + +def _require_unset(name: str, value: object) -> None: + if value is not None: + _reject(f"{name}=None", f"{name} is set") + + +def _reject(requirement: str, actual: str) -> NoReturn: + raise ValueError( + f"Dense DWDP ({DENSE_ENV}) requires {requirement}, but {actual}. Change the launch " + f"configuration or unset {DENSE_ENV}." + ) + + +def _reject_mode(actual: str) -> NoReturn: + raise ValueError( + f"Dense DWDP ({DENSE_ENV}) {_MODES}, but this launch has {actual}. Change the launch " + f"configuration or unset {DENSE_ENV}." + ) diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/layout.py b/tensorrt_llm/_torch/modules/dwdp/dense/layout.py new file mode 100644 index 000000000000..0bc12ca79719 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/layout.py @@ -0,0 +1,772 @@ +# 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. +"""Collection, checks and blob layout of the dense weights. + +The collector walks the dense roots each decoder layer names, decides which storages move +(anchors), records every reference to them (bindings), places the anchors in one blob per +``(layer, segment)`` and hands the result to :func:`~.types.build_dense_layout` for rows, store +offsets and the digest. Everything here is host logic over tensor metadata; the checksum +reductions run on whatever device the tensors live on (CPU in unit tests), and no function needs +CUDA. +""" + +from __future__ import annotations + +import weakref +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass + +import torch +from torch import nn + +from .config import DenseConfig +from .types import ( + ALLOWED_DTYPES, + CHECKSUM_MULTIPLIER, + POOLED_SLACK_BYTES, + ROW_ALIGN, + ZERO_CHECKSUM, + AnchorRecord, + BindingKind, + Checksum, + DenseLayout, + DenseLayoutError, + DenseRootsProvider, + RowPiece, + SegmentKey, + SegmentKind, + SegmentLayout, + TensorBinding, + align_up, + build_dense_layout, + warn_or_raise, +) + +#: Words one pass of :func:`tensor_checksum` reduces; bounds its int64 temporaries (32 MiB each). +_CHUNK_WORDS = 1 << 22 + +#: Words per partial sum. A word times a 16-bit half of its position hash is below ``2**47`` in +#: magnitude, so a partial sum of ``2**15`` such products stays below ``2**62``: every sum the +#: checksum computes in int64 is exact, and the wrap-around happens only in Python integers. +_PARTIAL_WORDS = 1 << 15 + +_WORD_BYTES = 4 +_HASH_MODULUS = 1 << 32 +_HALF_BITS = 16 +_HALF_MASK = (1 << _HALF_BITS) - 1 +_U64_MASK = (1 << 64) - 1 + +#: Most findings one :class:`~.types.DenseLayoutError` lists. +_MAX_LISTED = 8 + +#: Tensor types a moved reference may have; any other subclass is rejected. +_MOVABLE_TYPES = (torch.Tensor, nn.Parameter) + +_MODULE_KINDS = (BindingKind.PARAMETER, BindingKind.BUFFER) + +_RootList = list[tuple[str, nn.Module]] + + +@dataclass(eq=False) +class _Reference: + """A tensor reference on the dense device met by the walk of one ``(layer, segment)``.""" + + layer: int + module_path: str + kind: BindingKind + name: str + index: int | None + owner: nn.Module + tensor: torch.Tensor + storage: torch.UntypedStorage + storage_key: int + nbytes: int + eligible: bool + unknown: bool + + @property + def path(self) -> str: + return _reference_path(self.layer, self.module_path, self.name, self.index) + + def byte_range(self) -> tuple[int, int]: + start = self.tensor.storage_offset() * self.tensor.element_size() + return start, start + _span_bytes(self.tensor) + + +def collect_dense_layout( + layers: Sequence[tuple[int, DenseRootsProvider]], + config: DenseConfig, + *, + num_ranks: int, + device: torch.device, + attr_allow: frozenset[str], + exclude: frozenset[str], +) -> DenseLayout: + """Collect the anchors and bindings of every covered layer and lay them out (P1, no copies). + + Walk. Layers in ascending index order; per layer the enabled segments in canonical order; per + segment the roots ``layer.dense_dwdp_roots(segment.value)`` in the order returned; per root + ``root.named_modules(prefix=root_path)`` in registration order, visiting a module reached twice + within one ``(layer, segment)`` once. Per module, in this order: ``_parameters`` in insertion + order (``PARAMETER``), ``_buffers`` in insertion order (``BUFFER``), then ``vars(module)`` + sorted by name: tensor values (``ATTRIBUTE``) and non-empty tuples or lists whose elements are + all tensors (``SEQUENCE_ITEM``, one per element). None entries are skipped. Only tensors on + ``device`` with a non-empty storage are references (a ``device`` without an index matches every + index of its type); tensors on another device of the same type whose storage reaches the + threshold are an error, and tensors on other device types are ignored. + + Anchors. A storage becomes an anchor of ``(layer, segment)`` when it is reached through an + eligible reference and ``0 < untyped_storage().nbytes()`` and ``nbytes >= config.min_bytes``. + Parameters and buffers are eligible unless their name is in ``exclude``; attributes and sequence + items are eligible only if their name is in ``attr_allow`` and not in ``exclude``. An attribute + or sequence item named in neither table whose storage reaches the threshold but does not become + an anchor through another reference stays in place with a warning (``config.strict``: + :class:`~.types.DenseLayoutError`). + + Bindings. Every reference under the segment's roots to an anchor storage becomes a + :class:`~.types.TensorBinding` of that anchor, including names in ``exclude`` and attributes + outside ``attr_allow``; bindings are unique by ``(owner, kind, name, index)``. ``original`` + is the referenced tensor (``param.data`` for parameters) and ``source`` a flat ``uint8`` view + of the whole storage. + + Left in place with a warning (``config.strict``: :class:`~.types.DenseLayoutError`): a storage + referenced, through any reference, from more than one ``(layer, segment)``; a pooled storage, + whose ``nbytes`` exceeds the union of its references' byte ranges by more than + ``POOLED_SLACK_BYTES``. + + Errors regardless of ``config.strict`` (:class:`~.types.DenseLayoutError`): a binding whose + owner has ``_weights_transformed`` False; a binding dtype outside ``ALLOWED_DTYPES``; a binding + that is a tensor subclass other than ``nn.Parameter``; a binding whose dry-run repoint fails, + i.e. ``(O_X + byte_offset)`` is not a multiple of its itemsize or + ``byte_offset + span_bytes > nbytes``. Also an error: a layer without ``dense_dwdp_roots``, a + layer with roots for some enabled segments but not all, no covered layer, nothing moved. + + Layout. Per ``(layer, segment)``, anchors in the order their first eligible reference is met, + at offsets from :func:`assign_blob_offsets`; covered layers are those whose roots are non-empty. + Rows, store offsets, chunks, slot specs and the digest come from + :func:`~.types.build_dense_layout` with ``config.digest_switches(num_ranks)``. Checksums are + not computed here (see :func:`compute_checksums`). + + Args: + layers: ``(layer index, decoder layer)`` for every decoder layer, ascending; layers whose + roots are [] (layer 0 of Kimi-K3) are not covered. + config: Dense switches; ``segments``, ``min_bytes``, ``store_chunk_mb`` and ``strict`` + are used. + num_ranks: DWDP group size ``N``. + device: The device whose tensors move (``cuda:`` in production, CPU in tests). + attr_allow: Plain attribute names that may anchor a storage + (the model's ``DWDP_DENSE_ATTR_ALLOW``). + exclude: Names that never anchor a storage (the model's ``DWDP_DENSE_EXCLUDE``). + + Returns: + The layout with anchors, bindings, ``warnings`` and ``digest``; checksums unset. + + Raises: + DenseLayoutError: See above. + ValueError: ``layers`` is not strictly ascending, or ``num_ranks`` is below 2. + """ + switches = config.digest_switches(num_ranks) + device = torch.device(device) + kinds = switches.segments + + walks: dict[SegmentKey, list[_Reference]] = {} + for layer, roots in _covered_roots(layers, kinds): + for kind in kinds: + walks[(layer, kind)] = _walk_segment( + layer, + roots[kind], + device=device, + min_bytes=config.min_bytes, + attr_allow=attr_allow, + exclude=exclude, + ) + + selected, warnings = _select_anchors(walks, min_bytes=config.min_bytes, strict=config.strict) + + blob_bytes: dict[SegmentKey, int] = {} + anchors: dict[SegmentKey, list[AnchorRecord]] = {} + errors: list[str] = [] + for key, groups in selected.items(): + anchors[key], blob_bytes[key] = _place_anchors(key, groups, errors) + if errors: + raise DenseLayoutError(_listing("dense weights cannot move", errors)) + if not any(anchors.values()): + covered = sorted({layer for layer, _ in walks}) + raise DenseLayoutError( + f"nothing to move: no storage of at least {config.min_bytes} bytes under the dense " + f"roots of layers {covered} can move" + ) + try: + return build_dense_layout(blob_bytes, switches, anchors=anchors, warnings=warnings) + except ValueError as exc: + raise DenseLayoutError(f"inconsistent dense layout: {exc}") from exc + + +def assign_blob_offsets(nbytes: Sequence[int]) -> tuple[list[int], int]: + """Place storages of the given sizes back to back in a blob. + + Args: + nbytes: Storage sizes in blob order, each positive. + + Returns: + ``(offsets, blob_bytes)``: offset ``i`` is ``align_up(end of storage i - 1, ROW_ALIGN)`` + (0 for the first); ``blob_bytes`` is the end of the last storage (0 for none). + + Raises: + ValueError: A size is not positive. + """ + offsets: list[int] = [] + end = 0 + for size in nbytes: + if size <= 0: + raise ValueError(f"blob members must have a positive size, got {size}") + offset = align_up(end, ROW_ALIGN) + offsets.append(offset) + end = offset + size + return offsets, end + + +def tensor_checksum(tensor: torch.Tensor, *, word_offset: int = 0) -> Checksum: + """Checksum of a contiguous tensor's bytes. + + The bytes are zero-padded to whole int32 words; word ``i`` of the tensor is word + ``word_offset + i`` of the checksum's origin. The reduction runs on the tensor's device, in + bounded chunks, and the result is copied to the host. + + Args: + tensor: Contiguous tensor of any dtype; callers pass ``t.contiguous()`` for views. + word_offset: Word index of the tensor's first byte relative to the origin (``O_X // 4`` + for an anchor inside its blob). + + Returns: + The two sums, unsigned modulo ``2**64``. + + Raises: + ValueError: ``tensor`` is not contiguous or ``word_offset`` is negative. + """ + if word_offset < 0: + raise ValueError(f"word_offset must be non-negative, got {word_offset}") + if not tensor.is_contiguous(): + raise ValueError("tensor_checksum needs a contiguous tensor; pass t.contiguous()") + flat = tensor.detach().reshape(-1) + if flat.numel() == 0: + return ZERO_CHECKSUM + data = flat.view(torch.uint8) + chunk_bytes = _CHUNK_WORDS * _WORD_BYTES + word_sum = 0 + weighted_sum = 0 + for start in range(0, data.numel(), chunk_bytes): + words, weighted = _chunk_sums( + data[start : start + chunk_bytes], word_offset + start // _WORD_BYTES + ) + word_sum += words + weighted_sum += weighted + return Checksum(word_sum & _U64_MASK, weighted_sum & _U64_MASK) + + +def device_checksum(tensor: torch.Tensor, *, word_offset: int = 0) -> torch.Tensor: + """:func:`tensor_checksum` of a contiguous tensor, left on the tensor's device. + + The same two sums over the same bounded chunks, but nothing is copied to the host: the + reductions are enqueued on the current stream and the host is never synchronized. Device + integer arithmetic wraps modulo ``2**64``, so each sum is the int64 with the 64 bits of the + unsigned sum whatever the reduction order; :func:`checksum_bits` gives the same pair for a + :class:`~.types.Checksum`. + + Args: + tensor: Contiguous tensor of any dtype. + word_offset: Word index of the tensor's first byte relative to the origin. + + Returns: + ``int64 [2]`` on the tensor's device: ``word_sum`` and ``weighted_sum``. + + Raises: + ValueError: ``tensor`` is not contiguous or ``word_offset`` is negative. + """ + if word_offset < 0: + raise ValueError(f"word_offset must be non-negative, got {word_offset}") + if not tensor.is_contiguous(): + raise ValueError("device_checksum needs a contiguous tensor; pass t.contiguous()") + data = tensor.detach().reshape(-1).view(torch.uint8) + sums = torch.zeros(2, dtype=torch.int64, device=data.device) + chunk_bytes = _CHUNK_WORDS * _WORD_BYTES + for start in range(0, data.numel(), chunk_bytes): + words = _words(data[start : start + chunk_bytes]) + first_word = word_offset + start // _WORD_BYTES + # h_i = ((first_word + i) * M) mod 2**32; i * M stays below 2**55, so it is exact, and + # w_i * h_i stays below 2**63 in magnitude. + hashes = torch.arange(words.numel(), dtype=torch.int64, device=words.device) + hashes.mul_(CHECKSUM_MULTIPLIER).add_(first_word * CHECKSUM_MULTIPLIER % _HASH_MODULUS) + hashes.bitwise_and_(_HASH_MODULUS - 1) + sums[0].add_(words.sum(dtype=torch.int64)) + sums[1].add_(hashes.mul_(words).sum()) + return sums + + +def checksum_bits(checksum: Checksum) -> tuple[int, int]: + """The sums of ``checksum`` as the int64 values with the same 64 bits (:func:`device_checksum`).""" + return _signed_64(checksum.word_sum), _signed_64(checksum.weighted_sum) + + +def compute_checksums(layout: DenseLayout) -> None: + """Fill the per-segment and per-binding checksums of a collected layout (P1). + + ``SegmentLayout.checksum`` is the merge of :func:`tensor_checksum` of every anchor's + ``source`` at word offset ``blob_offset // 4``, i.e. the checksum of the zero-padded blob; + ``TensorBinding.checksum`` is :func:`tensor_checksum` of ``original.contiguous()``. + + Raises: + DenseLayoutError: An anchor's ``source`` or a binding's ``original`` was already dropped. + """ + with torch.no_grad(): + for segment in layout.segments.values(): + total = ZERO_CHECKSUM + for anchor in segment.anchors: + if anchor.source is None: + raise DenseLayoutError(f"the source of {anchor.path} was already dropped") + total = total.merge( + tensor_checksum(anchor.source, word_offset=anchor.blob_offset // _WORD_BYTES) + ) + for binding in anchor.bindings: + if binding.original is None: + raise DenseLayoutError( + f"the original of layers.{anchor.layer}.{binding.qualified_name} " + f"was already dropped" + ) + binding.checksum = tensor_checksum(binding.original.contiguous()) + segment.checksum = total + + +def row_pieces(segment: SegmentLayout, rank: int, *, num_ranks: int) -> list[RowPiece]: + """The runs of ``rank``'s row that come from anchors. + + Row ``rank`` is blob bytes ``[rank * S, (rank + 1) * S)``; each anchor overlapping it gives + one piece, in ascending ``row_offset``. Bytes no piece covers are zero. + + Args: + segment: A segment of a collected layout. + rank: Row index. + num_ranks: ``N``. + + Raises: + ValueError: ``rank`` is outside ``[0, num_ranks)``, or ``N`` rows of the segment's width + do not hold its blob. + """ + if not 0 <= rank < num_ranks: + raise ValueError(f"rank {rank} is outside [0, {num_ranks})") + _check_rows(segment, num_ranks) + row_start, row_end = segment.row_range(rank) + pieces = [] + for index, anchor in enumerate(segment.anchors): + start = max(row_start, anchor.blob_offset) + end = min(row_end, anchor.blob_end) + if start < end: + pieces.append( + RowPiece( + anchor_index=index, + anchor_offset=start - anchor.blob_offset, + row_offset=start - row_start, + nbytes=end - start, + ) + ) + return pieces + + +def write_blob(segment: SegmentLayout, out: torch.Tensor, *, num_ranks: int) -> None: + """Write the full blob (all ``N`` rows) of a collected segment into ``out``. + + ``out[:N * S]`` receives each anchor's ``source`` at its ``blob_offset`` and zeros elsewhere, + which is exactly what a gather of ``L(layer)`` leaves in a slot; ``out[N * S:]`` is not + touched. Used by the loopback store, Phase V references and unit tests. + + Args: + segment: A segment of a collected layout (anchors with ``source``). + out: Flat ``uint8`` tensor with at least ``N * S`` elements, on any device. + num_ranks: ``N``. + + Raises: + DenseLayoutError: An anchor's ``source`` was already dropped. + ValueError: ``out`` is not a flat ``uint8`` tensor or is too small, or ``N`` rows of the + segment's width do not hold its blob. + """ + if out.dtype != torch.uint8 or out.dim() != 1: + raise ValueError(f"out must be a flat uint8 tensor, got {out.dtype} {tuple(out.shape)}") + _check_rows(segment, num_ranks) + total = num_ranks * segment.row_bytes + if out.numel() < total: + raise ValueError(f"out holds {out.numel()} bytes; the blob needs {total}") + for anchor in segment.anchors: + if anchor.source is None: + raise DenseLayoutError(f"the source of {anchor.path} was already dropped") + cursor = 0 + for anchor in segment.anchors: + if anchor.blob_offset > cursor: + out[cursor : anchor.blob_offset].zero_() + out[anchor.blob_offset : anchor.blob_end].copy_(anchor.source) + cursor = anchor.blob_end + if total > cursor: + out[cursor:total].zero_() + + +# --------------------------------------------------------------------------- +# Walk +# --------------------------------------------------------------------------- + + +def _covered_roots( + layers: Sequence[tuple[int, DenseRootsProvider]], kinds: Sequence[SegmentKind] +) -> list[tuple[int, dict[SegmentKind, _RootList]]]: + covered = [] + previous: int | None = None + for layer, module in layers: + if previous is not None and layer <= previous: + raise ValueError( + f"decoder layers must be strictly ascending; layer {layer} follows {previous}" + ) + previous = layer + roots_of = getattr(module, "dense_dwdp_roots", None) + if not callable(roots_of): + raise DenseLayoutError( + f"decoder layer {layer} ({type(module).__name__}) does not provide dense_dwdp_roots" + ) + roots = {kind: _checked_roots(layer, kind, roots_of(kind.value)) for kind in kinds} + present = [kind.value for kind in kinds if roots[kind]] + if not present: + continue + if len(present) != len(kinds): + missing = [kind.value for kind in kinds if not roots[kind]] + raise DenseLayoutError( + f"decoder layer {layer} has dense roots for {present} but none for {missing}" + ) + covered.append((layer, roots)) + if not covered: + raise DenseLayoutError("no decoder layer has dense roots, so nothing can move") + return covered + + +def _checked_roots(layer: int, kind: SegmentKind, roots: Sequence[object]) -> _RootList: + checked = [] + for entry in roots: + try: + path, module = entry + except (TypeError, ValueError): + raise DenseLayoutError( + f"decoder layer {layer} returned {entry!r} as a {kind.value} root; " + f"expected a (path, nn.Module) pair" + ) from None + if not isinstance(path, str) or not isinstance(module, nn.Module): + raise DenseLayoutError( + f"decoder layer {layer} returned ({path!r}, {type(module).__name__}) as a " + f"{kind.value} root; expected a (path, nn.Module) pair" + ) + checked.append((path, module)) + return checked + + +def _walk_segment( + layer: int, + roots: _RootList, + *, + device: torch.device, + min_bytes: int, + attr_allow: frozenset[str], + exclude: frozenset[str], +) -> list[_Reference]: + references = [] + visited: set[int] = set() + for root_path, root in roots: + for module_path, module in root.named_modules(prefix=root_path): + if id(module) in visited: + continue + visited.add(id(module)) + for kind, name, index, tensor in _module_tensors(module): + if tensor.device.type != device.type: + continue + storage = tensor.untyped_storage() + nbytes = storage.nbytes() + if device.index is not None and tensor.device.index != device.index: + if 0 < nbytes and nbytes >= min_bytes: + raise DenseLayoutError( + f"{_reference_path(layer, module_path, name, index)} ({nbytes}-byte " + f"storage) is on {tensor.device}, not on the dense device {device}" + ) + continue + if nbytes == 0: + continue + module_kind = kind in _MODULE_KINDS + references.append( + _Reference( + layer=layer, + module_path=module_path, + kind=kind, + name=name, + index=index, + owner=module, + tensor=tensor, + storage=storage, + storage_key=storage.data_ptr(), + nbytes=nbytes, + eligible=name not in exclude and (module_kind or name in attr_allow), + unknown=not module_kind and name not in attr_allow and name not in exclude, + ) + ) + return references + + +def _reference_path(layer: int, module_path: str, name: str, index: int | None) -> str: + suffix = "" if index is None else f"[{index}]" + return f"layers.{layer}.{module_path}.{name}{suffix}" + + +def _module_tensors( + module: nn.Module, +) -> Iterator[tuple[BindingKind, str, int | None, torch.Tensor]]: + for name, param in module._parameters.items(): + if param is not None: + yield BindingKind.PARAMETER, name, None, param + for name, buffer in module._buffers.items(): + if buffer is not None: + yield BindingKind.BUFFER, name, None, buffer + attributes = vars(module) + for name in sorted(attributes): + value = attributes[name] + if isinstance(value, torch.Tensor): + yield BindingKind.ATTRIBUTE, name, None, value + elif ( + isinstance(value, (tuple, list)) + and value + and all(isinstance(item, torch.Tensor) for item in value) + ): + for index, item in enumerate(value): + yield BindingKind.SEQUENCE_ITEM, name, index, item + + +# --------------------------------------------------------------------------- +# Anchor selection and placement +# --------------------------------------------------------------------------- + + +def _select_anchors( + walks: Mapping[SegmentKey, list[_Reference]], *, min_bytes: int, strict: bool +) -> tuple[dict[SegmentKey, list[list[_Reference]]], list[str]]: + holders: dict[int, list[SegmentKey]] = {} + for key, references in walks.items(): + for reference in references: + keys = holders.setdefault(reference.storage_key, []) + if key not in keys: + keys.append(key) + + warnings: list[str] = [] + reported_shared: set[int] = set() + selected: dict[SegmentKey, list[list[_Reference]]] = {} + for key, references in walks.items(): + groups: dict[int, list[_Reference]] = {} + candidates: dict[int, _Reference] = {} + for reference in references: + groups.setdefault(reference.storage_key, []).append(reference) + if reference.eligible and reference.nbytes >= min_bytes: + candidates.setdefault(reference.storage_key, reference) + + chosen = [] + for storage_key, first in candidates.items(): + sharing = holders[storage_key] + if len(sharing) > 1: + if storage_key not in reported_shared: + reported_shared.add(storage_key) + owners = ", ".join(f"layers.{layer}/{kind.value}" for layer, kind in sharing) + warnings.append( + warn_or_raise( + f"storage of {first.path} ({first.nbytes} bytes) is referenced from " + f"{owners}; it stays replicated", + strict=strict, + error=DenseLayoutError, + ) + ) + continue + covered = _covered_bytes(groups[storage_key]) + if first.nbytes - covered > POOLED_SLACK_BYTES: + warnings.append( + warn_or_raise( + f"storage of {first.path} ({first.nbytes} bytes) looks pooled: its " + f"references cover only {covered} bytes; it stays replicated", + strict=strict, + error=DenseLayoutError, + ) + ) + continue + chosen.append(groups[storage_key]) + + anchored = {group[0].storage_key for group in chosen} + for reference in references: + if ( + reference.unknown + and reference.nbytes >= min_bytes + and reference.storage_key not in anchored + ): + warnings.append( + warn_or_raise( + f"{reference.path} ({reference.nbytes}-byte storage) is a plain attribute " + f"in neither the allow nor the exclude table; it stays in place", + strict=strict, + error=DenseLayoutError, + ) + ) + selected[key] = chosen + return selected, warnings + + +def _place_anchors( + key: SegmentKey, groups: list[list[_Reference]], errors: list[str] +) -> tuple[list[AnchorRecord], int]: + layer, kind = key + offsets, blob_bytes = assign_blob_offsets([group[0].nbytes for group in groups]) + anchors = [] + for index, (group, offset) in enumerate(zip(groups, offsets)): + bindings = tuple(_binding(reference) for reference in group) + for reference, binding in zip(group, bindings): + errors.extend(_binding_errors(reference, binding, offset)) + anchors.append( + AnchorRecord( + layer=layer, + segment=kind, + index=index, + nbytes=group[0].nbytes, + blob_offset=offset, + bindings=bindings, + source=_storage_bytes(group[0].storage), + ) + ) + return anchors, blob_bytes + + +def _binding(reference: _Reference) -> TensorBinding: + tensor = reference.tensor + return TensorBinding( + module_path=reference.module_path, + kind=reference.kind, + name=reference.name, + index=reference.index, + dtype=tensor.dtype, + size=tuple(tensor.shape), + stride=tuple(tensor.stride()), + byte_offset=tensor.storage_offset() * tensor.element_size(), + owner_ref=weakref.ref(reference.owner), + original=tensor.data if reference.kind is BindingKind.PARAMETER else tensor, + ) + + +def _binding_errors(reference: _Reference, binding: TensorBinding, blob_offset: int) -> list[str]: + errors = [] + if type(reference.tensor) not in _MOVABLE_TYPES: + errors.append( + f"{reference.path} is a {type(reference.tensor).__name__}; only plain tensors and " + f"nn.Parameter can move" + ) + if binding.dtype not in ALLOWED_DTYPES: + errors.append(f"{reference.path} has dtype {binding.dtype}, which dense DWDP cannot move") + if not getattr(reference.owner, "_weights_transformed", True): + errors.append( + f"{reference.path}: its {type(reference.owner).__name__} has " + f"_weights_transformed=False, so a later weight transform would run on slot memory" + ) + start = blob_offset + binding.byte_offset + if ( + start % binding.dtype.itemsize + or binding.byte_offset + binding.span_bytes > reference.nbytes + ): + errors.append( + f"{reference.path} cannot be repointed: a {binding.span_bytes}-byte view at blob " + f"offset {start} of a {reference.nbytes}-byte storage" + ) + return errors + + +def _storage_bytes(storage: torch.UntypedStorage) -> torch.Tensor: + """Flat ``uint8 [nbytes]`` tensor over the whole of ``storage``.""" + return torch.empty(0, dtype=torch.uint8, device=storage.device).set_(storage) + + +def _span_bytes(tensor: torch.Tensor) -> int: + if tensor.numel() == 0: + return 0 + span = 1 + sum((dim - 1) * step for dim, step in zip(tensor.shape, tensor.stride())) + return span * tensor.element_size() + + +def _covered_bytes(references: Sequence[_Reference]) -> int: + covered = 0 + run_start = run_end = 0 + for start, end in sorted(reference.byte_range() for reference in references): + if start >= end: + continue + if start > run_end: + covered += run_end - run_start + run_start, run_end = start, end + else: + run_end = max(run_end, end) + return covered + run_end - run_start + + +def _check_rows(segment: SegmentLayout, num_ranks: int) -> None: + if num_ranks <= 0 or num_ranks * segment.row_bytes < segment.blob_bytes: + raise ValueError( + f"{num_ranks} rows of {segment.row_bytes} bytes cannot hold the " + f"{segment.blob_bytes}-byte blob of layer {segment.layer} ({segment.segment.value})" + ) + + +def _listing(headline: str, findings: Sequence[str]) -> str: + listed = "; ".join(findings[:_MAX_LISTED]) + more = len(findings) - _MAX_LISTED + return f"{headline}: {listed}" + (f"; and {more} more" if more > 0 else "") + + +# --------------------------------------------------------------------------- +# Checksums +# --------------------------------------------------------------------------- + + +def _signed_64(value: int) -> int: + """The int64 whose 64 bits are those of the unsigned ``value`` (below ``2**64``).""" + return value - (1 << 64) if value >= 1 << 63 else value + + +def _words(data: torch.Tensor) -> torch.Tensor: + """A ``uint8`` run as int32 words, zero-padded to a whole word (copied only if it must be).""" + nbytes = data.numel() + if nbytes % _WORD_BYTES or data.storage_offset() % _WORD_BYTES: + aligned = data.new_zeros(align_up(nbytes, _WORD_BYTES)) + aligned[:nbytes].copy_(data) + data = aligned + return data.view(torch.int32) + + +def _chunk_sums(data: torch.Tensor, first_word: int) -> tuple[int, int]: + """Exact ``(sum w_i, sum w_i * h_i)`` of a ``uint8`` run whose first word is ``first_word``.""" + words = _words(data).to(torch.int64) + count = words.numel() + rows = -(-count // _PARTIAL_WORDS) + if rows * _PARTIAL_WORDS != count: + words = torch.cat((words, words.new_zeros(rows * _PARTIAL_WORDS - count))) + # h_i = ((first_word + i) * M) mod 2**32; positions * M stays below 2**55, so it is exact. + base = first_word * CHECKSUM_MULTIPLIER % _HASH_MODULUS + positions = torch.arange(rows * _PARTIAL_WORDS, dtype=torch.int64, device=words.device) + hashes = ((positions * CHECKSUM_MULTIPLIER + base) % _HASH_MODULUS).view(rows, _PARTIAL_WORDS) + words = words.view(rows, _PARTIAL_WORDS) + low = (words * (hashes & _HALF_MASK)).sum(dim=1) + high = (words * (hashes >> _HALF_BITS)).sum(dim=1) + sums = torch.cat((words.sum().reshape(1), low, high)).tolist() + weighted = sum(sums[1 : 1 + rows]) + (sum(sums[1 + rows :]) << _HALF_BITS) + return sums[0], weighted diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/manager.py b/tensorrt_llm/_torch/modules/dwdp/dense/manager.py new file mode 100644 index 000000000000..06aa3752981f --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/manager.py @@ -0,0 +1,596 @@ +# 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 manager: the owner of the committed dense DWDP runtime, in both modes. + +The model's dense hooks reach dense DWDP only through the manager that +:func:`get_global_dense_dwdp_manager` returns when the model is built: H1-H3 read +:attr:`DenseDwdpManager.dense`, H0 calls :meth:`DenseDwdpManager.ensure_primed`. The manager +does not depend on ``DwdpManager``. + +* Attached (``DenseMode.ATTACHED``, with expert DWDP). ``DwdpManager`` creates and registers the + manager, runs ``setup_dense_dwdp`` over the DWDP group before the expert DWDP setup allocates its + prefetch ring, on a transfer whose copy stream the expert weight manager then adopts, and hands + the committed runtime over with :meth:`DenseDwdpManager.attach`, together with the expert + weight manager. The runtime gathers the next forward's first layers at the end of a forward, + beside the expert prefetch. When they are not in flight, priming enqueues ``L(first), + L(second)`` and then starts the expert prefetch, so its ``E(first), E(second)`` follow them on + the shared copy stream; after an aborted forward H0 runs the dense WAR barrier first, and the + expert prefetch restarts its own sequence. ``DwdpManager.cleanup`` releases dense before the + expert teardown; this manager never releases the weight manager. +* Standalone (``DenseMode.STANDALONE``, an attention-DP instance with ``tp_size > 1`` and no + expert DWDP). The executor creator creates and registers the manager before the model is + built and calls :meth:`DenseDwdpManager.setup` once the weights are loaded. Setup runs over an + MPI sub-communicator of the attention-DP group (:func:`standalone_group`), which it frees when + it ends, with a :class:`StandaloneTransfer`: the manager's own copy stream and + ``TRTLLM_DWDP_DENSE_COPY``. Priming enqueues ``L(first), L(second)`` before every forward. + With the SM gather (``TRTLLM_DWDP_DENSE_COPY=gather``) the manager keeps the gather's resident + SMs out of the persistent compute grids from registration to exit. The final executor exits + the manager at shutdown. + +The module imports no executor or model code. +""" + +from __future__ import annotations + +import traceback +import weakref +from collections.abc import Callable +from typing import TYPE_CHECKING, Protocol + +import torch +from torch import nn + +from tensorrt_llm._torch.cute_dsl_utils import set_sm_reserve +from tensorrt_llm._utils import mpi_comm, mpi_disabled +from tensorrt_llm.logger import logger + +from .. import fabric_gather +from .config import DenseConfig +from .plan import select_copy_mode +from .setup import ( + DenseTransfer, + GroupComm, + LauncherFactory, + StoreFactory, + dense_dwdp_active, + dense_roots_available, + refuse_dense_dwdp_without_roots, + setup_dense_dwdp, +) +from .types import LOG_PREFIX, CopyMode, DenseMode, StreamLike + +if TYPE_CHECKING: + from mpi4py import MPI + + from tensorrt_llm.mapping import Mapping + + from .runtime import DenseDWDP + +_global_dense_dwdp_manager: DenseDwdpManager | None = None + +#: How :meth:`DenseDwdpManager.ensure_primed` names itself when ``DenseDWDP.require_eager`` +#: refuses it. +H0_HOOK = "DenseDwdpManager.ensure_primed (H0)" + +#: The ``Create_group`` tag of :func:`standalone_group` (``"DW"``). MPI pairs the +#: ``MPI_Comm_create_group`` calls on one parent communicator by their tag, and ``MPIDist.tp_comm`` +#: creates its communicator over the same group with the default tag 0, so with a tag of its own +#: the dense group never pairs with a peer's ``tp_comm``: a rank whose ``TRTLLM_DWDP_DENSE`` +#: differs leaves its peers waiting in ``Create_group``, not sharing a communicator with them. At +#: most 32767, the smallest ``MPI_TAG_UB`` the MPI standard allows. +STANDALONE_GROUP_TAG = 0x4457 + + +def set_global_dense_dwdp_manager(manager: DenseDwdpManager | None) -> None: + """Register ``manager`` as the process's dense DWDP manager; None unregisters.""" + global _global_dense_dwdp_manager + _global_dense_dwdp_manager = manager + + +def get_global_dense_dwdp_manager() -> DenseDwdpManager | None: + """The registered dense DWDP manager, or None; the model's hooks capture it when built.""" + return _global_dense_dwdp_manager + + +class ExpertPrefetcher(Protocol): + """The ``DWDPWeightManager`` call attached priming and recovery make.""" + + def start_forward(self) -> None: + """Start the expert prefetch of a forward: its first MoE layers, unless in flight. + + After an aborted forward it restarts the prefetch sequence, synchronizing the device first + if that forward left a MoE layer bound. + """ + ... + + +def new_copy_stream(device_id: int) -> torch.cuda.Stream: + """A new standalone dense copy stream on ``cuda:``. + + Default priority, like the expert DWDP copy stream. The test seam of the standalone + transfer: CPU tests monkeypatch ``manager.new_copy_stream``. + """ + return torch.cuda.Stream(device=torch.device("cuda", device_id)) + + +class StandaloneTransfer: + """The standalone dense transfer (``setup.DenseTransfer``): its own copy stream and mode. + + Only dense copies run on the stream, so its FIFO is ``L(first), L(second), L(3), ...``, the + deadline order. The stream is created by ``stream_factory`` at the first :attr:`copy_stream` + read, which setup makes in P2 inside the phase's local work: constructing the transfer + touches no CUDA, and a failure to create the stream becomes this rank's failed P2 report. + + Args: + mode: The transfer, normally ``DenseConfig.standalone_copy_mode()``. + stream_factory: Creates the copy stream; called at the first :attr:`copy_stream` read. + """ + + def __init__(self, mode: CopyMode, stream_factory: Callable[[], StreamLike]) -> None: + self._mode = CopyMode(mode) + self._stream_factory = stream_factory + self._stream: StreamLike | None = None + + @classmethod + def on_device(cls, mode: CopyMode, device_id: int) -> StandaloneTransfer: + """A transfer whose copy stream is ``new_copy_stream(device_id)``. + + ``new_copy_stream`` is looked up when the stream is created, so a monkeypatched one + applies. + """ + return cls(mode, lambda: new_copy_stream(device_id)) + + @property + def mode(self) -> CopyMode: + """The transfer.""" + return self._mode + + @property + def copy_stream(self) -> StreamLike: + """The copy stream: created on the first read, the same object on every later one.""" + if self._stream is None: + self._stream = self._stream_factory() + return self._stream + + @property + def use_gather(self) -> bool: + """Whether the copies use the SM fabric gather (``CopyMode.GATHER``).""" + return self._mode is CopyMode.GATHER + + @property + def use_batched_prefetch(self) -> bool: + """Whether the copies use one batched copy-engine copy per layer (``CopyMode.BATCHED``).""" + return self._mode is CopyMode.BATCHED + + +def standalone_group(mapping: Mapping) -> MPI.Intracomm: + """The standalone dense group: an MPI sub-communicator over ``mapping.tp_group``. + + Built the way ``MPIDist.tp_comm`` builds its communicator, but a communicator of its own, + created with :data:`STANDALONE_GROUP_TAG`: the dense setup collectives never interleave with + the executor's tp collectives, and the caller frees it when setup ends. Its rank and size are + ``mapping.tp_rank`` and ``mapping.tp_size``, and it has the ``GroupComm`` calls setup makes + plus ``Free``. Creating it is a collective over the tp group: every rank of the group must + call it. + + Raises: + RuntimeError: MPI is disabled (``TLLM_DISABLE_MPI=1``); the gate rejects such a launch + earlier. + """ + if mpi_disabled(): + raise RuntimeError( + "standalone dense DWDP needs the MPI backend: its group is an MPI sub-communicator " + "over the attention-DP group, but TLLM_DISABLE_MPI=1" + ) + comm = mpi_comm() + return comm.Create_group(comm.group.Incl(mapping.tp_group), tag=STANDALONE_GROUP_TAG) + + +class DenseDwdpManager: + """Owner of the committed dense DWDP runtime in either mode; see the module docstring. + + Neither construction nor registration (``__enter__``) runs a collective, so the manager + exists before the model is built, and the model's hooks capture it then; registration touches + CUDA only to check a standalone SM reserve against the device (``set_sm_reserve``). The hooks + are no-ops while :attr:`dense` is None: before :meth:`attach` or :meth:`setup` commits a + runtime, when dense DWDP is off or was abandoned, and after :meth:`cleanup`. + + Args: + mode: ``DenseMode.ATTACHED`` (with expert DWDP; ``DwdpManager`` sets dense DWDP up and + attaches the runtime) or ``DenseMode.STANDALONE`` (:meth:`setup` sets it up). + mapping: Standalone only, and required there: the instance's mapping, with + ``enable_attention_dp`` and ``tp_size > 1``; the dense group is its tp group. + + Raises: + ValueError: A standalone manager without such a mapping, or an attached one with a + mapping. + """ + + def __init__(self, mode: DenseMode, *, mapping: Mapping | None = None) -> None: + mode = DenseMode(mode) + if mode is DenseMode.STANDALONE: + if mapping is None or not mapping.enable_attention_dp or mapping.tp_size <= 1: + raise ValueError( + "a standalone DenseDwdpManager needs the mapping of an attention-DP " + "instance with tp_size > 1 (enable_attention_dp=True), got " + f"{_describe_mapping(mapping)}" + ) + elif mapping is not None: + raise ValueError( + "an attached DenseDwdpManager takes no mapping: with expert DWDP, DwdpManager " + f"sets dense DWDP up over the DWDP group; got {_describe_mapping(mapping)}" + ) + self._mode = mode + self._mapping = mapping + # The committed runtime; None before attach() or setup() commits one, when dense DWDP + # is off or was abandoned, and after cleanup(). + self._dense: DenseDWDP | None = None + # Attached only: the expert weight manager, whose prefetch H0 starts after the dense one. + self._expert: ExpertPrefetcher | None = None + # Standalone only: the model setup() committed the runtime for. + self._model_ref: weakref.ReferenceType[nn.Module] | None = None + # Standalone only: the SMs this manager keeps out of the persistent grids while it is + # registered (set_sm_reserve); 0 when it reserves none. + self._reserved_sms = 0 + self._closed = False + + # ------------------------------------------------------------------ + # Global registration + # ------------------------------------------------------------------ + + def __enter__(self) -> DenseDwdpManager: + """Register this manager globally; refused while any dense DWDP manager is registered. + + Standalone, it first keeps the SMs the SM gather holds (``TRTLLM_DWDP_DENSE_COPY=gather``; + ``fabric_gather.resident_sms``) out of the persistent compute grids (``set_sm_reserve``), + before the model is built, so that kernels sizing their grids at construction (the KDA + prefill) leave them out too. ``TRTLLM_CUTEDSL_SM_RESERVE`` pins the count, and a pinned + count is reserved with every transfer, the copy engines included, as ``DwdpManager`` + reserves it with every expert transfer. The reserve holds until exit, unless a setup that + commits nothing drops it. Attached, ``DwdpManager`` sets the reserve for the expert + transfer, which the dense copies use. + + Raises: + RuntimeError: A dense DWDP manager is registered. + ValueError: Standalone, ``TRTLLM_DWDP_DENSE_COPY`` is invalid (the gate rejects it + earlier), or the reserve leaves no SM of the device. + """ + if get_global_dense_dwdp_manager() is not None: + raise RuntimeError("DenseDwdpManager already registered globally") + if self._mode is DenseMode.STANDALONE: + gather = DenseConfig.from_env().standalone_copy_mode() is CopyMode.GATHER + reserve = fabric_gather.resident_sms(gather=gather) + set_sm_reserve(reserve) + self._reserved_sms = reserve + set_global_dense_dwdp_manager(self) + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> bool: + """:meth:`cleanup`, then unregister this manager (only if it is the registered one) and + drop the SM reserve it set.""" + try: + self.cleanup() + finally: + if get_global_dense_dwdp_manager() is self: + set_global_dense_dwdp_manager(None) + self._release_sm_reserve() + return False + + # ------------------------------------------------------------------ + # State + # ------------------------------------------------------------------ + + @property + def mode(self) -> DenseMode: + """Attached (with expert DWDP) or standalone.""" + return self._mode + + @property + def dense(self) -> DenseDWDP | None: + """The committed dense DWDP runtime, or None when dense DWDP is not active. + + None before :meth:`attach` or :meth:`setup` commits a runtime, when dense DWDP is off or + was abandoned (``TRTLLM_DWDP_DENSE_REQUIRED=0``), and after :meth:`cleanup`. The model's + layer hooks run only while it is not None. + """ + return self._dense + + # ------------------------------------------------------------------ + # Setup + # ------------------------------------------------------------------ + + def attach(self, runtime: DenseDWDP, *, expert: ExpertPrefetcher) -> None: + """Take over a runtime ``DwdpManager.setup`` committed (attached mode). + + Keeps ``expert``, the expert weight manager, whose prefetch :meth:`ensure_primed` starts + after the dense priming; the manager never releases or tears it down. Logs nothing. + + Raises: + RuntimeError: The manager is standalone or closed, or a runtime is already attached. + ValueError: ``runtime`` is None. + """ + if self._mode is not DenseMode.ATTACHED: + raise RuntimeError( + "DenseDwdpManager.attach() is for dense DWDP with expert DWDP; a standalone " + "manager sets dense DWDP up itself" + ) + if self._closed: + raise RuntimeError("DenseDwdpManager.attach() called after cleanup()") + if self._dense is not None: + raise RuntimeError( + "DenseDwdpManager.attach(): a dense DWDP runtime is already attached" + ) + if runtime is None: + raise ValueError("DenseDwdpManager.attach() needs the committed runtime, got None") + self._dense = runtime + self._expert = expert + + def setup( + self, + model: nn.Module, + *, + comm: GroupComm | None = None, + transfer: DenseTransfer | None = None, + config: DenseConfig | None = None, + device_id: int | None = None, + store_factory: StoreFactory | None = None, + launcher_factory: LauncherFactory | None = None, + ) -> DenseDWDP | None: + """Set standalone dense DWDP up on ``model`` over the attention-DP group. + + Called once from ``py_executor_creator`` after the weights are loaded and before the + final KV cache is sized, so that the KV cache gets the memory setup frees, on every rank + of the instance. In order: + + 1. Refused after :meth:`cleanup`. + 2. A runtime is already committed: for the same model it is returned without any + collective; another model is refused. + 3. A model whose dense weights a released runtime moved is refused: its modules hold + zero-size tensors. + 4. A model without dense roots: ``refuse_dense_dwdp_without_roots`` (it raises unless + ``TRTLLM_DWDP_DENSE_REQUIRED`` is off), then None, with no group and no collective. + 5. The switches (``config``, or ``DenseConfig.from_env()``) and the device. + 6. The transfer (``transfer``, or :meth:`StandaloneTransfer.on_device` with + ``config.standalone_copy_mode()``; no stream yet). + 7. One INFO line, so that ranks waiting for a peer in the next step show where they wait. + 8. The group: ``comm``, which the caller owns, or a new :func:`standalone_group`, the + first collective of setup. + 9. ``setup_dense_dwdp`` as rank ``mapping.tp_rank`` of ``mapping.tp_size``, without an + expert layer list. A committed runtime is stored at once; with a gather transfer, + ``fabric_gather.apply_compute_side_defaults()`` follows, still before any GEMM is + tuned. A setup that commits nothing leaves the process environment as it was, and + drops the SM reserve ``__enter__`` set: no gather runs. + 10. A group this call created is freed, whatever happened; a failing ``Free`` is logged, + never raised: the peers would otherwise carry on into the instance's next + collective, and the group carries no traffic after setup. + + Steps 1-7 are host-only and depend only on the model code and the job environment, which + every rank shares, so every rank reaches the group or raises before it alike. Every + fallible local step after it, the creation of the copy stream included, runs inside + ``setup_dense_dwdp``'s phases and becomes a report, so every rank reaches the same + verdict. + + Args: + model: The loaded model. + comm: The dense group; a new :func:`standalone_group` if None. Tests inject one. + transfer: Copy stream and mode; ``StandaloneTransfer.on_device`` if None. + config: Switches; ``DenseConfig.from_env()`` if None. + device_id: CUDA device of this rank; ``torch.cuda.current_device()`` if None. + store_factory: Passed to ``setup_dense_dwdp``. + launcher_factory: Passed to ``setup_dense_dwdp``. + + Returns: + The committed runtime, now :attr:`dense`, or None when dense DWDP is off or was + abandoned with ``TRTLLM_DWDP_DENSE_REQUIRED`` off. + + Raises: + RuntimeError: An attached manager; called after :meth:`cleanup`; a runtime is + committed for another model; or the model's dense weights were already moved. + Each is raised before any work. + ValueError: ``TRTLLM_DWDP_DENSE_SEGMENTS`` or ``TRTLLM_DWDP_DENSE_COPY`` is invalid + (the gate rejects both earlier). + DenseSetupError: Setup failed or was abandoned with ``TRTLLM_DWDP_DENSE_REQUIRED`` + on, or the model has no dense roots and it is on; raised after the group is + freed. + """ + if self._mode is not DenseMode.STANDALONE: + raise RuntimeError( + "DenseDwdpManager.setup() is for standalone dense DWDP; with expert DWDP, " + "DwdpManager.setup sets dense DWDP up and attaches the runtime" + ) + if self._closed: + raise RuntimeError("DenseDwdpManager.setup() called after cleanup()") + if self._dense is not None: + if self._model_ref is None or self._model_ref() is not model: + raise RuntimeError( + "DenseDwdpManager.setup(): dense DWDP is already committed for another " + "model; a manager serves one model" + ) + return self._dense + if dense_dwdp_active(model): + raise RuntimeError( + "DenseDwdpManager.setup(): dense DWDP already moved this model's dense weights " + "and its runtime was released, so the modules hold zero-size tensors; load the " + "model again to set it up" + ) + if not dense_roots_available(model): + refuse_dense_dwdp_without_roots(model) + self._release_sm_reserve(abandoned=True) + return None + if config is None: + config = DenseConfig.from_env() + if device_id is None: + device_id = torch.cuda.current_device() + if transfer is None: + transfer = StandaloneTransfer.on_device(config.standalone_copy_mode(), device_id) + mapping = self._mapping + copy_mode = select_copy_mode( + use_gather=transfer.use_gather, use_batched=transfer.use_batched_prefetch + ) + logger.info( + f"{LOG_PREFIX} standalone setup: rank {mapping.tp_rank} of the {mapping.tp_size}-rank " + f"attention-DP group, {copy_mode.value} copies on the dense copy stream" + ) + group = None + if comm is None: + group = comm = standalone_group(mapping) + try: + dense = setup_dense_dwdp( + model, + transfer, + comm, + dwdp_rank=mapping.tp_rank, + dwdp_size=mapping.tp_size, + device_id=device_id, + moe_layers=None, + config=config, + store_factory=store_factory, + launcher_factory=launcher_factory, + ) + if dense is not None: + # Stored before the group is freed: whatever Free does, the manager owns the + # runtime, so cleanup() (or the creator's rollback) releases it. + self._dense = dense + self._model_ref = weakref.ref(model) + if transfer.use_gather: + fabric_gather.apply_compute_side_defaults() + else: + self._release_sm_reserve(abandoned=True) + finally: + if group is not None: + _free_group(group) + return dense + + # ------------------------------------------------------------------ + # Priming (H0 and the executor loop) + # ------------------------------------------------------------------ + + def prefetch_first_layers(self) -> None: + """Executor-loop priming of the first two covered layers before a forward. + + Enqueues ``L(first), L(second)`` on the copy stream, then marks the slots primed. A no-op + without a runtime, when the slots are already primed (with expert DWDP, the previous + forward's end gathers them), or when a forward is still open: the previous forward + aborted, and the model's :meth:`ensure_primed` recovers on the execution stream. No read + window is open while no forward is, so the stream this runs on needs no fence. With + expert DWDP, ``DwdpManager.start_forward`` calls this before it starts the expert + prefetch, whose first copies then follow these on the shared copy stream. + """ + dense = self._dense + if dense is None or dense.forward_open or dense.primed: + return + self._prime() + + def ensure_primed(self) -> None: + """Model-side priming (H0): called at the model's forward entry while ``dense`` is set. + + Runs on the execution stream, so it also covers the forwards the executor loop does not + start: warmup, autotuning, KV-cache estimation and every forward of the overlap loop. It + refuses a CUDA graph capture and multi-stream execution (``DenseDWDP.require_eager``) + before it records or enqueues anything, so a refused forward leaves dense and expert DWDP + as it found them. Then, if the previous forward aborted between its first and last + covered layer, it records the dense WAR barrier on this stream (every dense consume + event), so the next dense copies wait for every read the aborted forward enqueued. Then + it enqueues ``L(first), L(second)`` unless already primed. Attached, it then starts the + expert prefetch (``ExpertPrefetcher.start_forward``): nothing when the first MoE layers + are in flight, else their copies behind the dense ones; after an aborted forward the + expert sequence restarts, with a device synchronization if that forward left a MoE layer + bound. + + Raises: + RuntimeError: Dense DWDP is not active. + DenseProtocolError: The current stream is capturing a CUDA graph, or multi-stream + execution is on. + """ + dense = self._dense + if dense is None: + raise RuntimeError( + "DenseDwdpManager.ensure_primed() needs an active dense DWDP runtime" + ) + dense.require_eager(H0_HOOK) + if dense.forward_open: + dense.war_barrier() + if not dense.primed: + self._prime() + if self._expert is not None: + self._expert.start_forward() + + def _prime(self) -> None: + """Enqueue the first two covered layers, then mark the dense slots primed. + + Attached, the covered layers are the expert MoE layers and the callers start the expert + prefetch next, so ``L(l)`` goes before ``E(l)`` on the shared copy stream, which runs them + in enqueue order: layer ``l`` reads ``L(l)`` (attention) before ``E(l)`` (experts). + """ + dense = self._dense + for layer in dense.layers[:2]: + dense.prefetch(layer) + dense.mark_primed() + + # ------------------------------------------------------------------ + # Teardown + # ------------------------------------------------------------------ + + def cleanup(self) -> None: + """Release the runtime; idempotent. + + The manager is closed and its state cleared first, then the runtime is released: it + synchronizes the device, repoints the model's dense weights to zero-size tensors and + frees the staging slots and the row store. The model keeps its ``_dense_dwdp_active`` + mark, so weight reloads stay refused. A release error propagates with the state already + cleared. The expert weight manager is never released here. + """ + self._closed = True + dense, self._dense = self._dense, None + self._expert = None + self._model_ref = None + if dense is not None: + dense.release() + + def _release_sm_reserve(self, *, abandoned: bool = False) -> None: + """Drop the SM reserve :meth:`__enter__` set, if any. + + ``abandoned``: setup committed nothing, so no gather runs; kernels that sized their grids + while the model was built keep leaving the SMs out, which costs time only. + """ + reserved, self._reserved_sms = self._reserved_sms, 0 + if not reserved: + return + set_sm_reserve(0) + if abandoned: + logger.warning( + f"{LOG_PREFIX} dense DWDP is off: the persistent grids launched from now on use " + f"all SMs again, while kernels that sized their grids when the model was built " + f"(the KDA prefill) keep leaving {reserved} SMs out" + ) + + +def _free_group(group: MPI.Intracomm) -> None: + """Free the standalone group setup created; a failure is logged, never raised.""" + try: + group.Free() + except Exception: # the peers carry on; a failed free leaks only a communicator handle + logger.error( + f"{LOG_PREFIX} freeing the standalone dense group failed\n{traceback.format_exc()}" + ) + + +def _describe_mapping(mapping: Mapping | None) -> str: + if mapping is None: + return "no mapping" + return ( + f"a mapping with enable_attention_dp={mapping.enable_attention_dp} and " + f"tp_size={mapping.tp_size}" + ) diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/plan.py b/tensorrt_llm/_torch/modules/dwdp/dense/plan.py new file mode 100644 index 000000000000..49249f03b311 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/plan.py @@ -0,0 +1,554 @@ +# 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. +"""Per-layer copy plans of dense DWDP and the launcher that runs them. + +Addressing: rank ``p``'s row of ``(layer, segment)`` is at +``chunk_bases[p][chunk] + chunk_offset`` (``p == rank`` is this rank's own store) and lands in the +slot holding the layer at ``base + p * S``. One plan per layer merges both segments: one row per +peer, each carrying the peer's ``A`` and ``D`` rows plus one slice of this rank's own ``A`` and +``D`` rows, so the whole layer is one gather launch with the same 15-row, 2-blocks-per-row, +cluster-2 geometry as the expert gather. + +The builders are pure host logic over the layout and plain integer addresses (CPU-testable); only +:class:`DenseCopyLauncher` touches CUDA. +""" + +from __future__ import annotations + +import functools +import itertools +import operator +from collections.abc import Callable, Mapping, Sequence +from types import ModuleType +from typing import NamedTuple + +import torch + +from tensorrt_llm.logger import logger + +from .. import fabric_gather +from .types import ( + BATCH_SLICE_BYTES, + LOG_PREFIX, + ROW_ALIGN, + CopyDescriptor, + CopyMode, + DenseLayout, + DenseSetupError, + LayerCopyPlan, + SegmentKind, + SegmentLayout, + align_up, + ceil_div, +) + +#: ``trtllm::dwdp_fabric_gather`` arguments of one layer, as ``fabric_gather.build_plan`` returns +#: them: device ``int64`` destination pointers, source pointers, 16-byte line counts and row starts, +#: plus the blocks per row. +GatherArgs = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int] + +#: One ``cudaMemcpyBatchAsync``: parallel destination, source and size lists of equal length. +BatchedCopies = tuple[list[int], list[int], list[int]] + + +def local_slices(row_bytes: int, num_slices: int) -> list[tuple[int, int]]: + """Cut this rank's row into one slice per peer row. + + With ``q = align_up(ceil(S / num_slices), ROW_ALIGN)``, slice ``k`` is + ``(k * q, max(0, min(q, S - k * q)))``; slices past the row end have length 0 and are skipped + by the plan. + + Args: + row_bytes: ``S``, a multiple of ``ROW_ALIGN``. + num_slices: ``N - 1``, positive. + + Returns: + ``num_slices`` ``(offset, length)`` pairs; the lengths sum to ``S``. + + Raises: + ValueError: ``row_bytes`` is negative or not aligned, or ``num_slices`` is not positive. + """ + if num_slices <= 0: + raise ValueError(f"num_slices must be positive, got {num_slices}") + if row_bytes < 0 or row_bytes % ROW_ALIGN: + raise ValueError( + f"row_bytes must be a non-negative multiple of {ROW_ALIGN}, got {row_bytes}" + ) + width = align_up(ceil_div(row_bytes, num_slices), ROW_ALIGN) + return [(k * width, max(0, min(width, row_bytes - k * width))) for k in range(num_slices)] + + +class _AddressTables(NamedTuple): + """Validated address tables of one rank.""" + + chunk_bases: tuple[tuple[int, ...], ...] + """``[rank][chunk]`` store chunk addresses.""" + + slot_bases: dict[SegmentKind, tuple[int, ...]] + """``[segment][ring slot]`` slot base addresses, for the enabled segments.""" + + +def _checked_address(value: int, name: str) -> int: + """``value`` as an address: an integer, non-negative and ``ROW_ALIGN``-aligned.""" + try: + address = operator.index(value) + except TypeError: + raise ValueError(f"{name} must be an integer address, got {value!r}") from None + if address < 0 or address % ROW_ALIGN: + raise ValueError(f"{name} must be a non-negative multiple of {ROW_ALIGN}, got {address:#x}") + return address + + +def _check_tables( + layout: DenseLayout, + rank: int, + chunk_bases: Sequence[Sequence[int]], + slot_bases: Mapping[SegmentKind, Sequence[int]], +) -> _AddressTables: + """Validate ``rank`` and the address tables against ``layout``. + + Raises: + ValueError: ``rank`` outside ``[0, N)``; a table with the wrong number of ranks, chunks, + segments or ring slots; an address that is not a non-negative multiple of + ``ROW_ALIGN``; or two slots whose ``N * S_max`` byte ranges overlap. + """ + num_ranks = layout.num_ranks + if not 0 <= rank < num_ranks: + raise ValueError(f"rank must be in [0, {num_ranks}), got {rank}") + + num_chunks = len(layout.chunks) + if len(chunk_bases) != num_ranks: + raise ValueError( + f"chunk_bases needs one entry per rank ({num_ranks}), got {len(chunk_bases)}" + ) + chunks = [] + for peer, bases in enumerate(chunk_bases): + if len(bases) != num_chunks: + raise ValueError( + f"chunk_bases[{peer}] needs one address per store chunk ({num_chunks}), " + f"got {len(bases)}" + ) + chunks.append( + tuple( + _checked_address(base, f"chunk_bases[{peer}][{chunk}]") + for chunk, base in enumerate(bases) + ) + ) + + by_kind = {SegmentKind(kind): bases for kind, bases in slot_bases.items()} + if set(by_kind) != set(layout.segment_kinds): + raise ValueError( + f"slot_bases must hold exactly the enabled segments " + f"{[kind.value for kind in layout.segment_kinds]}, got " + f"{sorted(kind.value for kind in by_kind)}" + ) + slots: dict[SegmentKind, tuple[int, ...]] = {} + spans: list[tuple[int, int, str]] = [] + for kind in layout.segment_kinds: + spec = layout.slot_specs[kind] + bases = by_kind[kind] + if len(bases) != spec.depth: + raise ValueError( + f"slot_bases[{kind.value}] needs one base per ring slot ({spec.depth}), " + f"got {len(bases)}" + ) + slots[kind] = tuple( + _checked_address(base, f"slot_bases[{kind.value}][{index}]") + for index, base in enumerate(bases) + ) + spans.extend( + (base, base + spec.valid_bytes, f"{kind.value}[{index}]") + for index, base in enumerate(slots[kind]) + ) + # Sorted by start, two ranges overlap iff some neighbouring pair does. + spans = sorted(span for span in spans if span[0] < span[1]) + for (_, end, first), (start, _, second) in zip(spans, spans[1:]): + if start < end: + raise ValueError(f"slots {first} and {second} overlap") + return _AddressTables(chunk_bases=tuple(chunks), slot_bases=slots) + + +def _check_segment(layout: DenseLayout, segment: SegmentLayout) -> None: + """The row width, chunk and chunk offset of a segment keep every plan address aligned. + + Raises: + ValueError: ``S`` or ``chunk_offset`` is negative or not ``ROW_ALIGN``-aligned, the chunk + does not exist, or ``N * S`` exceeds the segment's slot. + """ + if ( + segment.row_bytes < 0 + or segment.row_bytes % ROW_ALIGN + or segment.chunk_offset < 0 + or segment.chunk_offset % ROW_ALIGN + or not 0 <= segment.chunk < len(layout.chunks) + or layout.num_ranks * segment.row_bytes > layout.slot_specs[segment.segment].valid_bytes + ): + raise ValueError( + f"segment ({segment.layer}, {segment.segment.value}) has S={segment.row_bytes}, " + f"chunk {segment.chunk}, chunk_offset={segment.chunk_offset}; S and the offset must " + f"be non-negative multiples of {ROW_ALIGN}, the chunk must exist and N * S must fit " + f"the slot" + ) + + +def _layer_plan( + layout: DenseLayout, layer: int, rank: int, tables: _AddressTables +) -> LayerCopyPlan: + """:func:`build_layer_plan` of a covered layer over tables :func:`_check_tables` accepted.""" + peers = tuple(peer for peer in range(layout.num_ranks) if peer != rank) + # Per enabled segment: its layout, the base of the slot holding ``layer`` and the local slices. + targets = [] + for segment in layout.layer_segments(layer): + _check_segment(layout, segment) + base = tables.slot_bases[segment.segment][layout.slot_index(segment.segment, layer)] + targets.append((segment, base, local_slices(segment.row_bytes, len(peers)))) + + rows = [] + for k, peer in enumerate(peers): + row = [] + for segment, base, _ in targets: + if segment.row_bytes: + row.append( + CopyDescriptor( + dst=base + peer * segment.row_bytes, + src=tables.chunk_bases[peer][segment.chunk] + segment.chunk_offset, + nbytes=segment.row_bytes, + ) + ) + for segment, base, slices in targets: + offset, length = slices[k] + if length: + row.append( + CopyDescriptor( + dst=base + rank * segment.row_bytes + offset, + src=tables.chunk_bases[rank][segment.chunk] + segment.chunk_offset + offset, + nbytes=length, + ) + ) + rows.append(tuple(row)) + return LayerCopyPlan(layer=layer, rank=rank, peers=peers, rows=tuple(rows)) + + +def build_layer_plan( + layout: DenseLayout, + layer: int, + *, + rank: int, + chunk_bases: Sequence[Sequence[int]], + slot_bases: Mapping[SegmentKind, Sequence[int]], +) -> LayerCopyPlan: + """The gather of ``L(layer)`` on ``rank``. + + Peers ``p_0 < ... < p_{N-2}`` are the ranks other than ``rank``. Row ``k`` holds, in this + order and skipping zero-length ones: + + * ``A`` row of ``p_k``: ``base_A + p_k * S_A`` <- ``chunk_bases[p_k][c_A] + o_A``, ``S_A`` bytes; + * ``D`` row of ``p_k``: ``base_D + p_k * S_D`` <- ``chunk_bases[p_k][c_D] + o_D``, ``S_D`` bytes; + * slice ``k`` of the local ``A`` row: ``base_A + rank * S_A + off_k`` <- + ``chunk_bases[rank][c_A] + o_A + off_k``, ``len_k`` bytes (:func:`local_slices`); + * slice ``k`` of the local ``D`` row, likewise; + + where ``base_X = slot_bases[X][layout.slot_index(X, layer)]`` and ``c_X``/``o_X`` are the + segment's ``chunk``/``chunk_offset``. ``D`` entries exist only with the MoE-dense segment + enabled. Every address and length is a multiple of ``ROW_ALIGN``, the destinations cover + ``[base_X, base_X + N * S_X)`` exactly once per segment, and the plan's shape does not depend + on ``rank``. + + Args: + layout: The layout every rank shares. + layer: Covered layer ``n``. + rank: This rank ``r``. + chunk_bases: ``chunk_bases[p][c]``: address of chunk ``c`` of rank ``p``'s store as mapped + on this rank; ``N`` entries of ``len(layout.chunks)`` addresses, each a multiple of + ``ROW_ALIGN``. + slot_bases: ``slot_bases[segment][i]``: ``SlotBuffer.base`` of ring slot ``i``, for every + enabled segment. The slots' ``N * S_max`` byte ranges must not overlap. + + Returns: + The plan (never None; empty only if every enabled segment of ``layer`` has ``B == 0``). + + Raises: + ValueError: ``layer`` not covered; ``rank`` outside ``[0, N)``; address tables of the + wrong shape or alignment, or with overlapping slots; or a segment of ``layer`` whose + ``S`` or chunk offset is negative or not ``ROW_ALIGN``-aligned, whose chunk does not + exist, or whose ``N * S`` exceeds its slot. + """ + if not layout.covers(layer): + raise ValueError(f"layer {layer} is not covered by the dense layout") + return _layer_plan(layout, layer, rank, _check_tables(layout, rank, chunk_bases, slot_bases)) + + +def build_copy_plans( + layout: DenseLayout, + *, + rank: int, + chunk_bases: Sequence[Sequence[int]], + slot_bases: Mapping[SegmentKind, Sequence[int]], +) -> dict[int, LayerCopyPlan]: + """:func:`build_layer_plan` for every covered layer, keyed by layer. + + Raises: + ValueError: As :func:`build_layer_plan`, for any covered layer. + """ + tables = _check_tables(layout, rank, chunk_bases, slot_bases) + return {layer: _layer_plan(layout, layer, rank, tables) for layer in layout.layers} + + +def gather_launch_args(plan: LayerCopyPlan, device: torch.device) -> GatherArgs: + """``fabric_gather.build_plan`` of the plan's rows, uploaded to ``device``. + + One gather row per plan row (empty rows allowed), so ``N - 1 = 15`` rows give + ``blocks_per_row == 2``. The plan is never silently dropped: alignment that ``build_plan`` + rejects is an error here, not a copy-engine fallback. Build under the stream that later + launches the gather (the copy stream), as the expert plans are. + + Raises: + DenseSetupError: The plan is empty, or ``build_plan`` returned None. + """ + if plan.is_empty: + raise DenseSetupError( + f"the dense copy plan of layer {plan.layer} is empty; nothing to gather" + ) + args = fabric_gather.build_plan([list(row) for row in plan.rows], device) + if args is None: + raise DenseSetupError( + f"fabric_gather.build_plan rejected the dense copy plan of layer {plan.layer}: a " + f"descriptor is not {fabric_gather.LINE_BYTES}-byte aligned (dense DWDP has no " + f"copy-engine fallback)" + ) + return args + + +def batched_copies(plan: LayerCopyPlan, *, slice_bytes: int = BATCH_SLICE_BYTES) -> BatchedCopies: + """The plan as one ``cudaMemcpyBatchAsync``. + + Each descriptor is cut into ``slice_bytes`` pieces; the pieces are then taken round-robin over + the plan rows (one piece from every row that still has one, per round), so consecutive + copies read from different sources. The pieces cover the plan's bytes exactly. + + Raises: + ValueError: ``slice_bytes`` is not positive. + """ + if slice_bytes <= 0: + raise ValueError(f"slice_bytes must be positive, got {slice_bytes}") + row_pieces = [ + [ + (dst + offset, src + offset, min(slice_bytes, nbytes - offset)) + for dst, src, nbytes in row + for offset in range(0, nbytes, slice_bytes) + ] + for row in plan.rows + ] + dst_ptrs: list[int] = [] + src_ptrs: list[int] = [] + sizes: list[int] = [] + for round_pieces in itertools.zip_longest(*row_pieces): + for piece in round_pieces: + if piece is not None: + dst_ptrs.append(piece[0]) + src_ptrs.append(piece[1]) + sizes.append(piece[2]) + return dst_ptrs, src_ptrs, sizes + + +def select_copy_mode(*, use_gather: bool, use_batched: bool) -> CopyMode: + """The dense transfer from a transfer's two flags: gather, else batched, else per-slice. + + With expert DWDP the flags are those of the expert DWDP prefetch (``DwdpManager``'s attached + transfer computes them as the weight manager does), so the dense transfer follows the expert + one; standalone they are ``manager.StandaloneTransfer``'s, which ``TRTLLM_DWDP_DENSE_COPY`` + picks. + """ + if use_gather: + return CopyMode.GATHER + if use_batched: + return CopyMode.BATCHED + return CopyMode.PER_SLICE + + +def _load_cudart() -> ModuleType: + """``cuda.bindings.runtime``; imported on first use, as only the batched transfer needs it.""" + from cuda.bindings import runtime as cudart + + return cudart + + +def _memcpy_batch( + cudart: ModuleType, + attrs: list[object], + device: torch.device, + layer: int, + copies: BatchedCopies, +) -> None: + """One ``cudaMemcpyBatchAsync`` of ``copies`` on the current stream of ``device``.""" + dst_ptrs, src_ptrs, sizes = copies + stream = torch.cuda.current_stream(device) + (err,) = cudart.cudaMemcpyBatchAsync( + dst_ptrs, src_ptrs, sizes, len(dst_ptrs), attrs, [0], 1, stream.cuda_stream + ) + if err != cudart.cudaError_t.cudaSuccess: + raise RuntimeError( + f"cudaMemcpyBatchAsync failed for the dense DWDP copy of layer {layer} " + f"({len(dst_ptrs)} copies): {err}" + ) + + +def _descriptor_views( + plan: LayerCopyPlan, device_index: int +) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]: + """``(dst, src)`` ``uint8`` views of every descriptor, aliasing the plan's addresses. + + Raises: + DenseSetupError: An address cannot be viewed, or its view does not alias it. + """ + from .. import vmm + + views = [] + for descriptor in plan.descriptors: + pair = [] + for address in (descriptor.dst, descriptor.src): + try: + view = vmm.tensor_from_ptr(address, (descriptor.nbytes,), torch.uint8, device_index) + except ValueError as exc: + raise DenseSetupError( + f"layer {plan.layer}: cannot view {descriptor.nbytes} bytes at {address:#x}: " + f"{exc}" + ) from exc + if view.data_ptr() != address: + raise DenseSetupError( + f"layer {plan.layer}: the view of {address:#x} is a copy at " + f"{view.data_ptr():#x}; a per-slice copy needs views that alias the memory" + ) + pair.append(view) + views.append((pair[0], pair[1])) + return tuple(views) + + +def _copy_views(views: tuple[tuple[torch.Tensor, torch.Tensor], ...]) -> None: + """One ``copy_`` per ``(dst, src)`` pair, on the current stream.""" + for dst, src in views: + dst.copy_(src) + + +class DenseCopyLauncher: + """Runs the per-layer plans; the runtime's ``LayerLauncher``. + + Construction prepares everything a launch needs, so a launch only enqueues work on the + current stream (the runtime makes the copy stream current): + + * ``GATHER``: :func:`gather_launch_args` per layer, built on the current stream (construct it + under the copy stream); a launch is one ``trtllm::dwdp_fabric_gather`` with + ``fabric_gather.THREADS``. + * ``BATCHED``: :func:`batched_copies` per layer; a launch is one ``cudaMemcpyBatchAsync`` on + the current stream. + * ``PER_SLICE``: one ``uint8`` destination and source view per descriptor on ``device`` + (``vmm.tensor_from_ptr``); a launch is one ``copy_`` per descriptor. + + Layers whose plan is empty launch nothing. + + Args: + plans: Plan per covered layer, from :func:`build_copy_plans`. + mode: Transfer, from :func:`select_copy_mode`. + device: The dense CUDA device; a CUDA device without an index is the current device. + + Raises: + DenseSetupError: ``GATHER`` without a registered ``trtllm::dwdp_fabric_gather``, or a plan + the chosen transfer cannot run. + ValueError: ``mode`` is not a :class:`~.types.CopyMode`, or a plan is keyed by a layer + other than its own. + """ + + def __init__( + self, plans: Mapping[int, LayerCopyPlan], mode: CopyMode, *, device: torch.device + ) -> None: + self._mode = CopyMode(mode) + device = torch.device(device) + if device.type == "cuda" and device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + self._device = device + self._released = False + for layer, plan in plans.items(): + if plan.layer != layer: + raise ValueError( + f"the plan keyed by layer {layer} is the plan of layer {plan.layer}" + ) + + if self._mode is CopyMode.GATHER: + op = fabric_gather.gather_op() + if op is None: + raise DenseSetupError( + "the dense gather needs trtllm::dwdp_fabric_gather, which this build does " + "not register" + ) + threads = fabric_gather.THREADS + + def prepare(plan: LayerCopyPlan) -> Callable[[], None]: + return functools.partial(op, *gather_launch_args(plan, device), threads) + + elif self._mode is CopyMode.BATCHED: + try: + cudart = _load_cudart() + except ImportError as exc: + raise DenseSetupError( + f"the batched dense copy needs cuda.bindings.runtime: {exc}" + ) from exc + attrs = cudart.cudaMemcpyAttributes() + attrs.srcAccessOrder = cudart.cudaMemcpySrcAccessOrder.cudaMemcpySrcAccessOrderStream + + def prepare(plan: LayerCopyPlan) -> Callable[[], None]: + return functools.partial( + _memcpy_batch, cudart, [attrs], device, plan.layer, batched_copies(plan) + ) + + else: + device_index = device.index or 0 + + def prepare(plan: LayerCopyPlan) -> Callable[[], None]: + return functools.partial(_copy_views, _descriptor_views(plan, device_index)) + + self._launches: dict[int, Callable[[], None] | None] = { + layer: None if plan.is_empty else prepare(plan) for layer, plan in plans.items() + } + logger.debug( + f"{LOG_PREFIX} copy launcher: mode={self._mode.value}, layers={len(plans)}, " + f"descriptors={sum(len(plan.descriptors) for plan in plans.values())}, " + f"bytes={sum(plan.nbytes for plan in plans.values())}" + ) + + @property + def mode(self) -> CopyMode: + """The transfer in use.""" + return self._mode + + def __call__(self, layer: int) -> None: + """Enqueue the copy of ``L(layer)`` on the current stream. + + Raises: + KeyError: ``layer`` has no plan. + RuntimeError: The launcher was released, or ``cudaMemcpyBatchAsync`` failed. + """ + if self._released: + raise RuntimeError("the dense copy launcher was released") + try: + launch = self._launches[layer] + except KeyError: + raise KeyError(f"no dense copy plan for layer {layer}") from None + if launch is not None: + launch() + + def release(self) -> None: + """Drop the prepared launch arguments and views; idempotent.""" + self._launches = {} + self._released = True diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/repoint.py b/tensorrt_llm/_torch/modules/dwdp/dense/repoint.py new file mode 100644 index 000000000000..ffa3890b6ff8 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/repoint.py @@ -0,0 +1,577 @@ +# 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. +"""Repointing module tensors into the staging slots. + +After repointing, every binding of an anchor is a view of the slot that holds its layer: layer +``l`` and the layers that share its slot (``l + 2`` for attention, ``l + 3`` for MoE-dense) see +the same addresses, and the runtime protocol guarantees the slot holds a layer's bytes whenever +that layer reads them. Each binding is rebuilt from its own dtype, size, stride and byte offset, +so aliases stay aliases. The original storages are released only by dropping references; nothing +here resizes a storage. +""" + +from __future__ import annotations + +import enum +import gc +import types +import weakref +from collections.abc import Callable, Mapping, Sequence +from typing import NamedTuple + +import torch +from torch import nn + +from tensorrt_llm.logger import logger + +from .layout import tensor_checksum +from .types import ( + LOG_PREFIX, + RING_DEPTH, + AnchorRecord, + BindingKind, + DenseLayout, + DenseRepointError, + SegmentKind, + SlotBuffer, + TensorBinding, +) + +#: Most leaked anchors or failed bindings one error message lists; also the leaked anchors +#: that ``commit`` warns about one by one and, in debug, scans the heap for. +_MAX_LISTED = 8 + +#: Most holders, and referrers per holder, a leak description names. +_MAX_HOLDERS = 4 + + +class BindingState(enum.Enum): + """Where the modules' dense bindings point.""" + + PENDING = "pending" + """Collected; modules hold the original tensors.""" + + REPOINTED = "repointed" + """Modules hold slot views; the originals are still alive (P5).""" + + ROLLED_BACK = "rolled_back" + """Modules hold the original tensors again.""" + + COMMITTED = "committed" + """Modules hold slot views; the references to the originals are dropped (P6).""" + + DETACHED = "detached" + """Modules hold zero-size tensors (teardown).""" + + +def slot_view(binding: TensorBinding, anchor: AnchorRecord, slot: SlotBuffer) -> torch.Tensor: + """The view of ``slot`` that replaces ``binding``. + + ``slot.storage.view(dtype).as_strided(size, stride, (pad + O_X + byte_offset) // itemsize)``, + where ``pad`` is the slot's, ``O_X`` the anchor's blob offset and the rest the binding's. + + Raises: + DenseRepointError: The byte offset is not a multiple of the itemsize or the view leaves + the slot; the layout's dry run rules both out. + """ + storage = slot.storage + if storage.dtype != torch.uint8 or storage.dim() != 1 or storage.storage_offset() != 0: + raise DenseRepointError( + f"slot {slot.segment.value}[{slot.index}] must be a flat uint8 tensor at storage " + f"offset 0, got {storage.dtype} {tuple(storage.shape)} at {storage.storage_offset()}" + ) + where = f"layers.{anchor.layer}.{binding.qualified_name}" + itemsize = binding.dtype.itemsize + start = slot.pad + anchor.blob_offset + binding.byte_offset + span = binding.span_bytes + if start % itemsize: + raise DenseRepointError( + f"{where} would start at slot byte {start}, not a multiple of its itemsize {itemsize}" + ) + if binding.byte_offset + span > anchor.nbytes or start + span > storage.numel(): + raise DenseRepointError( + f"{where} ({span} bytes at anchor offset {binding.byte_offset}) leaves its " + f"{anchor.nbytes}-byte anchor or the {storage.numel()}-byte slot " + f"{slot.segment.value}[{slot.index}]" + ) + usable = storage.numel() - storage.numel() % itemsize + return ( + storage[:usable] + .view(binding.dtype) + .as_strided(binding.size, binding.stride, start // itemsize) + ) + + +def read_binding(binding: TensorBinding) -> torch.Tensor: + """The tensor the owner exposes for ``binding`` now. + + ``PARAMETER``: ``owner._parameters[name]``; ``BUFFER``: ``owner._buffers[name]``; + ``ATTRIBUTE``: ``vars(owner)[name]``; ``SEQUENCE_ITEM``: ``vars(owner)[name][index]``. + + Raises: + DenseRepointError: The owner is gone or no longer holds the name. + """ + owner = _owner(binding) + if binding.kind is BindingKind.PARAMETER: + value = owner._parameters.get(binding.name) + elif binding.kind is BindingKind.BUFFER: + value = owner._buffers.get(binding.name) + else: + value = vars(owner).get(binding.name) + if binding.kind is BindingKind.SEQUENCE_ITEM: + value = _item(value, binding.index) + if not isinstance(value, torch.Tensor): + raise _missing(owner, binding) + return value + + +def install_binding(binding: TensorBinding, tensor: torch.Tensor) -> None: + """Make the owner expose ``tensor`` for ``binding``. + + ``PARAMETER``: ``param.data = tensor`` (the ``nn.Parameter`` object stays); + ``BUFFER``: ``owner._buffers[name] = tensor`` (stays registered, persistence unchanged); + ``ATTRIBUTE``: ``object.__setattr__(owner, name, tensor)``; + ``SEQUENCE_ITEM``: rebuild the tuple or list with ``tensor`` at ``index`` (same container + type, other elements unchanged) and set it with ``object.__setattr__``. + + Raises: + DenseRepointError: ``tensor.dtype`` differs from ``binding.dtype``, or the owner is gone + or no longer holds the name. + """ + if tensor.dtype != binding.dtype: + raise DenseRepointError( + f"cannot install a {tensor.dtype} tensor for {binding.qualified_name} ({binding.dtype})" + ) + owner = _owner(binding) + name = binding.name + if binding.kind is BindingKind.PARAMETER: + param = owner._parameters.get(name) + if param is None: + raise _missing(owner, binding) + param.data = tensor + elif binding.kind is BindingKind.BUFFER: + if name not in owner._buffers: + raise _missing(owner, binding) + owner._buffers[name] = tensor + elif binding.kind is BindingKind.ATTRIBUTE: + if name not in vars(owner): + raise _missing(owner, binding) + object.__setattr__(owner, name, tensor) + else: + container = vars(owner).get(name) + if _item(container, binding.index) is None: + raise _missing(owner, binding) + items = list(container) + items[binding.index] = tensor + object.__setattr__(owner, name, _rebuild(container, items)) + + +class _Entry(NamedTuple): + """One binding with the anchor it views and the slot that holds its layer.""" + + binding: TensorBinding + anchor: AnchorRecord + slot: SlotBuffer + + +class DenseBindings: + """Every dense binding of a layout and where it points: slots, originals or nothing. + + The state machine (:class:`BindingState`) is ``PENDING -> REPOINTED -> COMMITTED -> DETACHED``, + with ``REPOINTED -> ROLLED_BACK`` (and back to ``REPOINTED``) before the commit point. The + layout keeps the originals alive until :meth:`commit`. + + Args: + layout: A collected layout (anchors with ``source``, bindings with ``original`` and + ``checksum``). + slots: The slot rings, ``slots[segment][i]`` for ``i < RING_DEPTH[segment]``; a binding of + layer ``l`` goes to ``slots[segment][layout.slot_index(segment, l)]``. + + Raises: + ValueError: A ring of an enabled segment is missing, has the wrong depth, or holds a slot + of another segment or position. + DenseRepointError: A binding's ``original`` was already dropped. + """ + + def __init__( + self, layout: DenseLayout, slots: Mapping[SegmentKind, Sequence[SlotBuffer]] + ) -> None: + rings = {SegmentKind(kind): tuple(ring) for kind, ring in slots.items()} + for kind in layout.segment_kinds: + ring = rings.get(kind, ()) + if len(ring) != RING_DEPTH[kind]: + raise ValueError( + f"the {kind.value} ring needs {RING_DEPTH[kind]} slots, got {len(ring)}" + ) + for position, slot in enumerate(ring): + if slot.segment is not kind or slot.index != position: + raise ValueError( + f"slot {position} of the {kind.value} ring is " + f"{slot.segment.value}[{slot.index}]" + ) + entries = [] + for anchor in layout.anchors(): + slot = rings[anchor.segment][layout.slot_index(anchor.segment, anchor.layer)] + for binding in anchor.bindings: + if binding.original is None: + raise DenseRepointError( + f"the original of layers.{anchor.layer}.{binding.qualified_name} was " + f"already dropped" + ) + entries.append(_Entry(binding, anchor, slot)) + self._layout = layout + self._entries = tuple(entries) + self._by_layer = {layer: [] for layer in layout.layers} + for entry in self._entries: + self._by_layer[entry.anchor.layer].append(entry) + self._state = BindingState.PENDING + + @property + def state(self) -> BindingState: + """Where the bindings point.""" + return self._state + + def repoint(self) -> None: + """Install :func:`slot_view` for every binding (P5). + + Valid in ``PENDING`` and ``ROLLED_BACK``; leaves ``REPOINTED``. Reads nothing from the + slots, so it is valid before any gather. If installing fails midway, the bindings already + repointed are restored before the error propagates. + + Raises: + DenseRepointError: Called in another state, or a binding could not be installed. + """ + self._require("repoint", BindingState.PENDING, BindingState.ROLLED_BACK) + views = [slot_view(entry.binding, entry.anchor, entry.slot) for entry in self._entries] + installed = 0 + try: + for entry, view in zip(self._entries, views): + install_binding(entry.binding, view) + installed += 1 + finally: + if installed < len(self._entries): + for entry in reversed(self._entries[:installed]): + install_binding(entry.binding, entry.binding.original) + self._state = BindingState.REPOINTED + + def verify_layer(self, layer: int) -> list[str]: + """P5 self-check of one layer, through the module attributes. + + The caller has gathered ``L(layer)`` into the layer's slots and synchronized. For every + binding of ``layer``, ``read_binding(binding)`` must be the binding's slot view (same + storage, offset, size, stride and dtype) and :func:`~.layout.tensor_checksum` of + ``read_binding(binding).contiguous()`` must equal the P1 ``binding.checksum``. + + Returns: + ``layers..`` of every mismatch; [] if all match. + + Raises: + DenseRepointError: Not ``REPOINTED``, ``layer`` is not covered, a binding has no P1 + checksum, or an owner no longer holds its binding. + """ + self._require("verify_layer", BindingState.REPOINTED) + entries = self._by_layer.get(layer) + if entries is None: + raise DenseRepointError(f"layer {layer} is not covered by dense DWDP") + mismatches = [] + with torch.no_grad(): + for entry in entries: + binding = entry.binding + if binding.checksum is None: + raise DenseRepointError( + f"layers.{layer}.{binding.qualified_name} has no P1 checksum" + ) + tensor = read_binding(binding) + if ( + not _is_slot_view(tensor, entry) + or tensor_checksum(tensor.contiguous()) != binding.checksum + ): + mismatches.append(f"layers.{layer}.{binding.qualified_name}") + return mismatches + + def rollback(self) -> None: + """Make every owner expose its original tensor again (failure before the commit point). + + Valid in ``REPOINTED``; a no-op in ``PENDING`` and ``ROLLED_BACK``. Afterwards every + attribute, parameter and buffer is as collected: same ``nn.Parameter`` objects with their + original ``.data`` views, the original buffer and attribute tensors, containers rebuilt + with the original elements. Every binding is attempted; the state is ``ROLLED_BACK`` even + if some could not be restored. + + Raises: + DenseRepointError: Called after :meth:`commit`, or an owner no longer holds its + binding. + """ + self._require( + "rollback", BindingState.PENDING, BindingState.REPOINTED, BindingState.ROLLED_BACK + ) + if self._state is not BindingState.REPOINTED: + return + failures = self._install_all(lambda entry: entry.binding.original) + self._state = BindingState.ROLLED_BACK + if failures: + raise DenseRepointError(_listing("rollback could not restore every binding", failures)) + + def commit(self, *, strict: bool, describe_referrers: bool = False) -> list[str]: + """Drop the references to the originals and check that their storages die (P6). + + Valid in ``REPOINTED``; leaves ``COMMITTED`` even when it raises. Layer by layer: take a + weak reference to each anchor's storage, call ``AnchorRecord.drop_originals()``, and check + the weak references are dead; after the last layer run ``gc.collect()`` once and check the + survivors again. A surviving storage keeps an old, byte-identical copy alive (memory, not + numerics). The first ``_MAX_LISTED`` survivors (in layer order) get one warning each, + naming the anchor path and its bindings; with ``describe_referrers`` it also summarizes + the holders, which scans every live object (``gc.get_objects`` and + ``gc.get_referrers``), so the scan is bounded to those few. More survivors get one + summary warning with their count, bytes and layers, so a leak that repeats in every layer + neither floods the log nor costs a scan per anchor. + + Args: + strict: Raise if any storage survives. + describe_referrers: Add holder details to the warnings of the first survivors + (debug). + + Returns: + Anchor paths whose storage survived, all of them; [] if none. + + Raises: + DenseRepointError: Called in another state, or ``strict`` and a storage survived. + """ + self._require("commit", BindingState.REPOINTED) + self._state = BindingState.COMMITTED + survivors: list[tuple[AnchorRecord, weakref.ReferenceType[torch.UntypedStorage]]] = [] + for layer in self._layout.layers: + layer_refs = [] + for segment in self._layout.layer_segments(layer): + for anchor in segment.anchors: + storage_ref = _storage_ref(anchor) + anchor.drop_originals() + if storage_ref is not None: + layer_refs.append((anchor, storage_ref)) + survivors.extend( + (anchor, storage_ref) + for anchor, storage_ref in layer_refs + if storage_ref() is not None + ) + if survivors: + gc.collect() + survivors = [(anchor, ref) for anchor, ref in survivors if ref() is not None] + for anchor, storage_ref in survivors[:_MAX_LISTED]: + bindings = ", ".join(binding.qualified_name for binding in anchor.bindings) + message = ( + f"{LOG_PREFIX} storage of {anchor.path} ({anchor.nbytes} bytes) is still alive " + f"after commit: a holder outside its bindings ({bindings}) keeps the old copy" + ) + if describe_referrers: + message += f"; holders: {_describe_holders(storage_ref)}" + logger.warning(message) + if len(survivors) > _MAX_LISTED: + rest = [anchor for anchor, _ in survivors[_MAX_LISTED:]] + kinds = {anchor.segment for anchor in rest} + total_bytes = sum(anchor.nbytes for anchor, _ in survivors) + logger.warning( + f"{LOG_PREFIX} {len(rest)} more moved storages " + f"({sum(anchor.nbytes for anchor in rest)} bytes) are still alive after commit, " + f"in segments {','.join(kind.value for kind in SegmentKind if kind in kinds)} of " + f"layers {_layer_ranges(sorted({anchor.layer for anchor in rest}))}: " + f"{len(survivors)} storages, {total_bytes} bytes in all; the first {_MAX_LISTED} " + f"are named above" + ) + leaked = [anchor.path for anchor, _ in survivors] + if strict and leaked: + raise DenseRepointError(_listing("moved storages are still alive after commit", leaked)) + return leaked + + def detach(self) -> None: + """Point every binding at a zero-size tensor of its dtype on the slot device (teardown). + + From ``COMMITTED``: leaves ``DETACHED``, so a later stray read fails on its shape instead of + reading freed slot memory. A binding whose owner was garbage collected counts as detached: + no module can read it any more, so a model collected before teardown detaches cleanly. + From ``REPOINTED``: rolls back instead (the originals are still alive). A no-op in + ``PENDING``, ``ROLLED_BACK`` and ``DETACHED``; idempotent. Every binding is attempted. + + Raises: + DenseRepointError: A live owner no longer holds its binding. + """ + if self._state is BindingState.REPOINTED: + self.rollback() + return + if self._state is not BindingState.COMMITTED: + return + empties: dict[tuple[torch.dtype, torch.device], torch.Tensor] = {} + + def empty_like(entry: _Entry) -> torch.Tensor: + key = (entry.binding.dtype, entry.slot.storage.device) + if key not in empties: + empties[key] = torch.empty(0, dtype=key[0], device=key[1]) + return empties[key] + + failures = self._install_all(empty_like, skip_collected=True) + self._state = BindingState.DETACHED + if failures: + raise DenseRepointError(_listing("detach could not reach every binding", failures)) + + def _require(self, action: str, *states: BindingState) -> None: + if self._state not in states: + raise DenseRepointError( + f"{action}() is not valid in state {self._state.value}; it needs " + f"{' or '.join(state.value for state in states)}" + ) + + def _install_all( + self, + tensor_for: Callable[[_Entry], torch.Tensor | None], + *, + skip_collected: bool = False, + ) -> list[str]: + """Install ``tensor_for(entry)`` for every binding, last first; returns the failures. + + With ``skip_collected``, a binding whose owner was garbage collected is skipped instead of + reported. + """ + failures = [] + collected = 0 + for entry in reversed(self._entries): + # Held until the install, so a live owner cannot be collected in between. + owner = entry.binding.owner_ref() + if owner is None and skip_collected: + collected += 1 + continue + tensor = tensor_for(entry) + path = f"layers.{entry.anchor.layer}.{entry.binding.qualified_name}" + if tensor is None: + failures.append(f"{path}: no tensor to install") + continue + try: + install_binding(entry.binding, tensor) + except DenseRepointError as exc: + failures.append(f"{path}: {exc}") + if collected: + logger.debug( + f"{LOG_PREFIX} {collected} of {len(self._entries)} bindings belong to " + f"garbage-collected modules; nothing reads them" + ) + return failures + + +def _owner(binding: TensorBinding) -> nn.Module: + owner = binding.owner_ref() + if owner is None: + raise DenseRepointError(f"the module owning {binding.qualified_name} was garbage collected") + return owner + + +def _missing(owner: nn.Module, binding: TensorBinding) -> DenseRepointError: + return DenseRepointError( + f"{type(owner).__name__} at {binding.module_path} no longer holds " + f"{binding.qualified_name} as a {binding.kind.value}" + ) + + +def _item(container: object, index: int | None) -> torch.Tensor | None: + if not isinstance(container, (tuple, list)) or index is None or not 0 <= index < len(container): + return None + item = container[index] + return item if isinstance(item, torch.Tensor) else None + + +def _rebuild(container: tuple | list, items: list[torch.Tensor]) -> tuple | list: + cls = type(container) + if cls is tuple or cls is list: + return cls(items) + make = getattr(cls, "_make", None) + return make(items) if make is not None else cls(items) + + +def _is_slot_view(tensor: torch.Tensor, entry: _Entry) -> bool: + """Whether ``tensor`` is the view :func:`slot_view` makes for ``entry``. + + The placement is compared as storage plus storage offset, not ``data_ptr()``, which is 0 for + a tensor without elements wherever it points. + """ + binding, anchor, slot = entry + start = slot.storage.storage_offset() + slot.pad + anchor.blob_offset + binding.byte_offset + return ( + tensor.dtype == binding.dtype + and tuple(tensor.shape) == binding.size + and tuple(tensor.stride()) == binding.stride + and tensor.untyped_storage().data_ptr() == slot.storage.untyped_storage().data_ptr() + and tensor.storage_offset() * tensor.element_size() == start + ) + + +def _storage_ref(anchor: AnchorRecord) -> weakref.ReferenceType[torch.UntypedStorage] | None: + source = anchor.source + if source is None: + return None + return weakref.ref(source.untyped_storage()) + + +def _describe_holders(storage_ref: weakref.ReferenceType[torch.UntypedStorage]) -> str: + storage = storage_ref() + if storage is None: + return "released while describing" + pointer, device = storage.data_ptr(), storage.device + del storage + holders = [] + for obj in gc.get_objects(): + if not issubclass(type(obj), torch.Tensor) or obj.device != device: + continue + try: + shares = obj.untyped_storage().data_ptr() == pointer + except (RuntimeError, NotImplementedError): + continue + if shares: + holders.append(obj) + parts = [] + for tensor in holders[:_MAX_HOLDERS]: + referrers = [ + _describe_referrer(referrer, tensor) + for referrer in gc.get_referrers(tensor) + if referrer is not holders and not isinstance(referrer, types.FrameType) + ] + parts.append( + f"{type(tensor).__name__}{list(tensor.shape)} held by " + f"{', '.join(referrers[:_MAX_HOLDERS]) or 'no tracked object'}" + ) + holders.clear() + return "; ".join(parts) or "no live tensor (a non-tensor object keeps the storage)" + + +def _describe_referrer(referrer: object, tensor: torch.Tensor) -> str: + if isinstance(referrer, dict): + keys = [key for key, value in referrer.items() if value is tensor] + return f"dict key {keys[0]!r}" if keys else "a dict" + if isinstance(referrer, (list, tuple)): + return f"a {type(referrer).__name__} of {len(referrer)}" + return f"a {type(referrer).__name__}" + + +def _layer_ranges(layers: Sequence[int]) -> str: + """Ascending layers as runs, e.g. ``[1, 2, 3, 5, 7, 8]`` as ``1..3, 5, 7..8``.""" + runs: list[list[int]] = [] + for layer in layers: + if runs and layer == runs[-1][1] + 1: + runs[-1][1] = layer + else: + runs.append([layer, layer]) + return ", ".join(str(first) if first == last else f"{first}..{last}" for first, last in runs) + + +def _listing(headline: str, findings: Sequence[str]) -> str: + listed = "; ".join(findings[:_MAX_LISTED]) + more = len(findings) - _MAX_LISTED + return f"{headline}: {listed}" + (f"; and {more} more" if more > 0 else "") diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py b/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py new file mode 100644 index 000000000000..a971d1aba243 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py @@ -0,0 +1,991 @@ +# 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 runtime: slot rings, events and the per-layer hooks. + +Per covered layer ``l`` (``s_A`` and ``s_D`` are its slots in the attention and MoE-dense rings): + +===== ========================================== ========================================== +Hook Call site Action +===== ========================================== ========================================== +H0 ``KimiLinearModel.forward`` entry ``DenseDwdpManager.ensure_primed()``: + barrier after an abort, then ``L(1), + L(2)`` unless already primed (with expert + DWDP then the expert ``start_forward``) +H1 decoder layer entry compute waits ``ready_A[s_A]`` and + ``ready_D[s_D]`` +H2 after the attention call record ``consume_A[s_A]``, then fill +H3 after the MoE block record ``consume_D[s_D]``, then fill +===== ========================================== ========================================== + +All reads of ``A(l)`` happen between H1 and H2; all reads of ``D(l)`` between H1 and H3. The fill +at H2 and H3 enqueues, in layer order, each next layer whose target slots are free: empty, or +holding a layer whose read window has closed. Within a forward that is ``L(l + 2)`` at H2 of +``l``. The fill never targets a slot that is unread or being read, so it never needs the fence. + +With expert DWDP (``prefetch_next_pass``) the fill goes on into the first two layers of the next +forward as the last layers free their slots, so a completed forward leaves the next one primed, as +the expert prefetch does. For Kimi-K3 (92 layers) the MoE-dense slots of layers 1 and 2 are those +of layers 91 and 92, read until H3, so H3 of layers 91 and 92 enqueue them (H2 does with the +attention segment alone). The copies run on the expert DWDP copy stream, one FIFO. Within a +forward, with two expert prefetch slots, ``L(l)`` is enqueued right before ``E(l)`` (``... E(l+1), +L(l+2), E(l+2), L(l+3) ...``): H2 of ``l - 2`` runs before the MoE of ``l - 2`` finishes and issues +``E(l)``. The next forward's first two layers follow their expert copies when their MoE-dense +slots free only at H3 (``E(1), L(1), E(2), L(2)`` for Kimi-K3); the next forward reads them only +after the gap between executor steps. With more expert prefetch slots each expert copy is issued +earlier, ahead of its layer's dense copy. + +Standalone, the copies run on the dense manager's own stream, ``... L(l+2), L(l+3) ...``, and the +fill stops at the last layer: the next forward starts from the dense manager's priming. + +Between commit and release the runtime never synchronizes the host; every dependency is a stream +wait on an event. The one exception is the slot watchdog of ``TRTLLM_DWDP_DENSE_DEBUG=2`` +(``watchdog.SlotWatchdog``), whose counters H3 of the last covered layer reads once per forward. +""" + +from __future__ import annotations + +import enum +import operator +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import TYPE_CHECKING + +from tensorrt_llm.logger import logger + +from .types import ( + EMPTY_TAG, + LOG_PREFIX, + RING_DEPTH, + SEGMENT_ORDER, + DenseProtocolError, + DeviceOps, + EventLike, + LayerLauncher, + SegmentKind, + SlotBuffer, + StreamLike, + TorchCudaOps, + ring_slot, +) + +if TYPE_CHECKING: + from .repoint import DenseBindings + from .store import DenseShardStore + from .watchdog import Finding, SlotWatchdog + +#: One slot: ``(segment, index in the segment's ring)``. +_SlotKey = tuple[SegmentKind, int] + + +class _Phase(enum.Enum): + """Lifecycle of :class:`DenseDWDP`.""" + + SETUP = "setup" + """Constructed; setup calls (P3-P5) are valid.""" + + ARMED = "armed" + """Committed (P6); the hooks, priming and recovery are valid.""" + + RELEASED = "released" + """Released; only :meth:`DenseDWDP.release` is valid.""" + + +class DenseDWDP: + """Slot rings, events and per-layer hooks of dense DWDP; owned by the dense manager. + + The dense manager (``manager.DenseDwdpManager``) holds the committed runtime in both modes: + with expert DWDP ``DwdpManager`` attaches it, standalone the manager sets it up itself. + + State, per enabled segment and ring slot ``s``: ``ready[seg][s]`` (recorded on the copy stream + after a gather, waited by compute: RAW), ``consume[seg][s]`` (recorded on the compute stream at + the end of the read window, waited by the copy stream: WAR), ``tag[seg][s]`` (layer last + enqueued into the slot, ``EMPTY_TAG`` if none), ``opened[seg][s]`` (whether H1 made compute wait + for the slot's current holder, which opens a read window) and ``consumed[seg][s]`` (whether a + consume record closed the holder's latest read window). One ``fence`` event serves demand + fetches. ``primed`` and ``forward_open`` drive priming, and the fill position (the next + layer to enqueue, counted from the current forward's first covered layer) drives the fill. All + events come from ``ops.new_event()`` and are never shared with the expert DWDP events. + + Invariants, checked on the host: + + * I1 (RAW): at H1 ``tag[seg][s] == l`` for every segment, else a fenced demand fetch of + ``L(l)`` runs first; compute then waits ``ready[seg][s]``, whose latest record follows the + copy of ``L(l)``. + * I2 (WAR): enqueueing ``L(n)`` into slot ``s`` waits ``consume[seg][s]``, whose latest record + follows every read of the slot enqueued so far unless the current holder's read window is + open (``opened`` and not ``consumed``); then the copy also waits the fence (``strict``: + :class:`~.types.DenseProtocolError`). A holder whose window never opened has not been + read, so it needs no fence. + * I3 (window): consume events are recorded only at H2 (attention) and H3 (MoE-dense), after + the window's reads are enqueued, and only to close a window that H1 opened: H2 and H3 raise + :class:`~.types.DenseProtocolError` when the slot holds another layer or its holder's read + window is not open. + * I4 (fill): H2 and H3 enqueue layers in order, each only into slots that are empty or whose + holder's read window has closed, so the fill never takes the fence: ``L(n)`` is enqueued at + H2 of the layer two positions before ``n`` and, with ``prefetch_next_pass``, the first two + layers of the next forward once their slots free up at the end of this one. + + Layers are addressed by position ``p`` among the covered layers: the slot is + ``ring_slot(seg, p)`` in every forward, and the fill reaches ``layers[p + 2]`` at H2 of + position ``p``. For Kimi-K3 (layers 1..92) that is ``(l - 1) % 2``, ``(l - 1) % 3`` and + ``l + 2``; with ``prefetch_next_pass``, H3 of layers 91 and 92 enqueue layers 1 and 2 of the + next forward (H2 of them with the attention segment alone). + + Lifecycle: constructed in setup P2 (records every consume event once on the current stream), + :meth:`set_launcher` in P3, :meth:`fetch_blocking` in P4 and P5, :meth:`commit` in P6. The + hooks, :meth:`prefetch`, :meth:`mark_primed` and :meth:`war_barrier` are valid only after + :meth:`commit`; :meth:`set_launcher` and :meth:`fetch_blocking` only before it. + :meth:`release` is valid at any time; after it, every call named above raises + :class:`~.types.DenseProtocolError`, while the introspection properties and methods keep + reporting the last state (:attr:`slots` is empty). The forward needs eager execution on one + stream: H0 and H1 refuse a CUDA graph capture and multi-stream execution + (:meth:`require_eager`) before they record or enqueue anything. + + With a ``watchdog`` (``TRTLLM_DWDP_DENSE_DEBUG=2``), H1 also enqueues on the compute stream a + check of each slot it waited for against the layer's P1 checksum; H3 of the last layer reads the + mismatch counters, which synchronizes the host once per forward, and logs every + ``(layer, segment)`` whose slot held other bytes (``strict``: raises). + + Args: + layers: Covered layers, ascending; with expert DWDP the same list as the expert DWDP MoE + layers. + segments: Enabled segment kinds in canonical order (``ATTN`` first). + slots: Ring per enabled segment, ``len(slots[seg]) == RING_DEPTH[seg]``; owned by the + runtime once the constructor returns (if it raises, the caller keeps them). + copy_stream: The copy stream: the expert DWDP copy stream with expert DWDP + (``DwdpManager.setup`` creates it for dense setup and hands it to the weight manager), + the dense manager's own stream standalone (``manager.StandaloneTransfer.copy_stream``). + prefetch_next_pass: Whether the fill goes on into the first two layers of the next + forward, which then starts primed: with expert DWDP, whose prefetch does the same, so + those copies are not queued behind its next-pass copies. Standalone, the dense + manager primes every forward. + ops: CUDA seam; ``TorchCudaOps()`` if None. + strict: ``TRTLLM_DWDP_DENSE_STRICT``: demand fetches, fence fallbacks and watchdog + mismatches raise. + store: Row store the launcher reads from; owned and released by the runtime once the + constructor returns. + watchdog: Slot checksum watchdog over the same layers and segments; owned and released by + the runtime once the constructor returns. None (the default) checks nothing. + + Raises: + ValueError: Empty, unsorted, duplicate or negative ``layers``; non-canonical + ``segments``; rings that do not match the enabled segments, or a ring of the wrong + depth or order; no ``copy_stream``; a ``watchdog`` of other layers or segments. + """ + + def __init__( + self, + layers: Sequence[int], + segments: Sequence[SegmentKind], + slots: Mapping[SegmentKind, Sequence[SlotBuffer]], + copy_stream: StreamLike, + *, + prefetch_next_pass: bool, + ops: DeviceOps | None = None, + strict: bool = False, + store: DenseShardStore | None = None, + watchdog: SlotWatchdog | None = None, + ) -> None: + covered = _covered_layers(layers) + kinds = _segment_kinds(segments) + rings = _slot_rings(kinds, slots) + if copy_stream is None: + raise ValueError("dense DWDP needs a copy stream") + if watchdog is not None and (watchdog.layers, watchdog.segments) != (covered, kinds): + raise ValueError( + f"the slot watchdog covers layers {list(watchdog.layers)} and segments " + f"{[kind.value for kind in watchdog.segments]}, the runtime {list(covered)} and " + f"{[kind.value for kind in kinds]}" + ) + + self._ops: DeviceOps = TorchCudaOps() if ops is None else ops + self._layers = covered + self._positions = {layer: position for position, layer in enumerate(covered)} + self._segments = kinds + self._has_moe = SegmentKind.MOE in kinds + self._targets: tuple[tuple[_SlotKey, ...], ...] = tuple( + tuple((kind, ring_slot(kind, position)) for kind in kinds) + for position in range(len(covered)) + ) + self._copy_stream = copy_stream + self._prefetch_next_pass = bool(prefetch_next_pass) + self._strict = bool(strict) + self._slots: dict[SegmentKind, tuple[SlotBuffer, ...]] = rings + self._ready: dict[SegmentKind, tuple[EventLike, ...]] = { + kind: tuple(self._ops.new_event() for _ in rings[kind]) for kind in kinds + } + self._consume: dict[SegmentKind, tuple[EventLike, ...]] = { + kind: tuple(self._ops.new_event() for _ in rings[kind]) for kind in kinds + } + self._fence: EventLike = self._ops.new_event() + self._tag: dict[SegmentKind, list[int]] = { + kind: [EMPTY_TAG] * len(rings[kind]) for kind in kinds + } + self._opened: dict[SegmentKind, list[bool]] = { + kind: [False] * len(rings[kind]) for kind in kinds + } + self._consumed: dict[SegmentKind, list[bool]] = { + kind: [True] * len(rings[kind]) for kind in kinds + } + self._primed = False + self._forward_open = False + # Position of the next layer the fill enqueues, counted from the current forward's first + # covered layer; positions from len(layers) on are the next forward's first layers + # (prefetch_next_pass only). + self._next_fill = 0 + # Forwards that ended before the fill reached the next forward's first layers. + self._unprimed_passes = 0 + self._demand_fetches = 0 + self._war_fences = 0 + self._launcher: LayerLauncher | None = None + self._bindings: DenseBindings | None = None + self._store = store + self._watchdog = watchdog + self._record_consume_events() + # Set last: ``__del__`` releases nothing of an object whose constructor raised. + self._phase = _Phase.SETUP + logger.debug( + f"{LOG_PREFIX} runtime: {len(covered)} layers {covered[0]}..{covered[-1]}, " + f"segments {','.join(kind.value for kind in kinds)}, strict={int(self._strict)}, " + f"next-pass prefetch {'on' if self._prefetch_next_pass else 'off'}, " + f"slot watchdog {'on' if watchdog is not None else 'off'}" + ) + + # ------------------------------------------------------------------ + # Introspection + # ------------------------------------------------------------------ + + @property + def layers(self) -> tuple[int, ...]: + """Covered layers, ascending.""" + return self._layers + + @property + def segments(self) -> tuple[SegmentKind, ...]: + """Enabled segment kinds in canonical order.""" + return self._segments + + @property + def slots(self) -> Mapping[SegmentKind, tuple[SlotBuffer, ...]]: + """The slot rings (read-only); empty after :meth:`release`.""" + return MappingProxyType(self._slots) + + @property + def prefetch_next_pass(self) -> bool: + """Whether the fill goes on into the first two layers of the next forward.""" + return self._prefetch_next_pass + + @property + def primed(self) -> bool: + """Whether ``L(first)`` and ``L(second)`` are enqueued for a forward that has not started. + + Set by :meth:`mark_primed` and, with ``prefetch_next_pass``, by H3 of the last covered + layer once the fill has enqueued them; H1 of the first covered layer clears it. + """ + return self._primed + + @property + def forward_open(self) -> bool: + """Whether a forward passed H1 of the first layer and has not reached H3 of the last.""" + return self._forward_open + + @property + def demand_fetches(self) -> int: + """Fenced demand fetches at H1 so far; 0 after warmup in a healthy run.""" + return self._demand_fetches + + @property + def war_fences(self) -> int: + """Prefetches that waited on the fence because a slot holder's read window was open.""" + return self._war_fences + + @property + def watchdog_checks(self) -> int: + """Slot checks the watchdog enqueued at H1 so far; 0 without a watchdog.""" + return 0 if self._watchdog is None else self._watchdog.checks + + @property + def watchdog_mismatches(self) -> int: + """Slot checks that found other bytes than their layer's, as of the watchdog's latest read. + + The counters are read at H3 of the last covered layer and at :meth:`release`; 0 without a + watchdog. + """ + return 0 if self._watchdog is None else self._watchdog.mismatches + + def covers(self, layer: int) -> bool: + """Whether ``layer`` is covered; the model calls the hooks only for covered layers.""" + return layer in self._positions + + def slot_index(self, segment: SegmentKind, layer: int) -> int: + """Ring slot of ``layer`` in ``segment``. + + Raises: + DenseProtocolError: ``layer`` is not covered, or ``segment`` is not enabled. + ValueError: ``segment`` is not a segment kind. + """ + kind = self._enabled(segment) + return ring_slot(kind, self._position(layer)) + + def slot_tag(self, segment: SegmentKind, index: int) -> int: + """Layer last enqueued into slot ``index`` of ``segment``, or ``EMPTY_TAG``. + + Raises: + DenseProtocolError: ``segment`` is not enabled. + ValueError: ``segment`` is not a segment kind. + IndexError: ``index`` is outside the ring. + """ + kind = self._enabled(segment) + return self._tag[kind][self._ring_index(kind, index)] + + def slot_consumed(self, segment: SegmentKind, index: int) -> bool: + """Whether a consume record closed the latest read window of ``segment``'s slot ``index``. + + False from the copy that installs the holder, and again from each H1 that reads it, until + H2 (attention) or H3 (MoE-dense) records the slot's consume event; True for an empty slot. + + Raises: + DenseProtocolError: ``segment`` is not enabled. + ValueError: ``segment`` is not a segment kind. + IndexError: ``index`` is outside the ring. + """ + kind = self._enabled(segment) + return self._consumed[kind][self._ring_index(kind, index)] + + # ------------------------------------------------------------------ + # Setup (P3-P6) + # ------------------------------------------------------------------ + + def set_launcher(self, launcher: LayerLauncher) -> None: + """Install the copy launcher (P3); the runtime calls it with the copy stream current. + + The runtime drops the launcher at :meth:`release` without calling anything on it. + + Raises: + DenseProtocolError: Called after :meth:`commit`. + TypeError: ``launcher`` is not callable. + """ + self._require_setup("set_launcher") + if not callable(launcher): + raise TypeError(f"the dense DWDP launcher must be callable, got {launcher!r}") + self._launcher = launcher + + def fetch_blocking(self, layer: int) -> None: + """Gather ``L(layer)`` outside the event protocol, then synchronize (P4 Phase V, P5). + + Launches on the copy stream exactly as the runtime does and calls ``ops.synchronize()``; + tags, ``consumed`` flags and events are untouched. + + Raises: + DenseProtocolError: No launcher, ``layer`` not covered, or called after + :meth:`commit`. + """ + self._require_setup("fetch_blocking") + self._position(layer) + launcher = self._launcher + if launcher is None: + raise DenseProtocolError(f"fetch_blocking(L{layer}) before set_launcher") + with self._ops.stream(self._copy_stream): + launcher(layer) + self._ops.synchronize() + + def commit(self, bindings: DenseBindings | None, *, fill_byte: int) -> None: + """Arm the runtime at the end of setup P6, in this fixed order. + + Phase V and the P5 gathers are finished and synchronized by the caller. Then: fill every + slot with ``fill_byte`` (0, or ``DEBUG_FILL_BYTE`` for NaN under debug) -> every tag + ``EMPTY_TAG``, no read window open and every ``consumed`` True -> record every consume + event on the current stream -> ``primed`` and ``forward_open`` False, the fill position + at the first layer. The byte fill is enqueued on the current stream before the consume + records, so the first copy into any slot waits for it. Takes ownership of ``bindings``, + which :meth:`release` detaches (None when nothing is repointed, as in unit tests). + + Raises: + DenseProtocolError: No launcher, or called twice. + ValueError: ``fill_byte`` is not in ``[0, 255]``. + """ + self._require_setup("commit") + if self._launcher is None: + raise DenseProtocolError("commit before set_launcher") + if not 0 <= fill_byte <= 0xFF: + raise ValueError(f"the slot fill byte must be in [0, 255], got {fill_byte}") + for kind in self._segments: + for slot in self._slots[kind]: + slot.storage.fill_(fill_byte) + self._reset_slot_state() + self._record_consume_events() + self._primed = False + self._forward_open = False + self._next_fill = 0 + self._bindings = bindings + self._phase = _Phase.ARMED + logger.debug(f"{LOG_PREFIX} runtime armed; slots filled with {fill_byte:#04x}") + if self._watchdog is not None: + logger.info( + f"{LOG_PREFIX} slot watchdog on (TRTLLM_DWDP_DENSE_DEBUG=2): every H1 checks its " + f"slots against the layer's checksum; the counters are read once per forward" + ) + + # ------------------------------------------------------------------ + # Per-layer hooks (compute stream) + # ------------------------------------------------------------------ + + def require_eager(self, hook: str) -> None: + """Raise unless the current stream runs eagerly with multi-stream execution off. + + The hooks keep host-side slot state and order the copies and the weight reads with + events they record in every forward. Inside a CUDA graph capture those records and + copies would go into the graph instead of running, and the host state would describe + work that never ran; with multi-stream execution, reads would run on streams that no + event covers. H0 (``DenseDwdpManager.ensure_primed``) calls this before it records or + enqueues anything, and H1 (:meth:`wait_layer`) before it waits: a refused hook records, + enqueues and changes nothing, so the next eager forward finds the state the refused hook + found. Valid in every phase. + + Args: + hook: The refused call, for the error message. + + Raises: + DenseProtocolError: The current stream is capturing a CUDA graph, or + ``ops.do_multi_stream()`` is on. + """ + if self._ops.is_current_stream_capturing(): + raise DenseProtocolError( + f"{hook} while the current stream is capturing a CUDA graph; dense DWDP needs " + f"eager execution" + ) + if self._ops.do_multi_stream(): + raise DenseProtocolError( + f"{hook} with multi-stream execution on; dense DWDP covers reads on the compute " + f"stream only" + ) + + def wait_layer(self, layer: int) -> None: + """H1, before the layer's first statement: make compute wait for ``L(layer)``. + + Raises :class:`~.types.DenseProtocolError` if the current stream is capturing a CUDA graph + or ``ops.do_multi_stream()`` is on (:meth:`require_eager`). If any enabled segment's + slot is not tagged ``layer`` (not primed, aborted forward, protocol bug), one fenced + demand fetch of ``L(layer)`` runs first and ``demand_fetches`` grows by one (``strict``: + raises instead, before anything is enqueued); a warning is logged at the 1st, 2nd, 4th, + 8th, ... demand fetch. Then compute waits every enabled segment's ``ready`` event, which + opens the read windows of the slots' holders (``opened`` True, ``consumed`` False until + H2 and H3; a holder that an earlier forward already read gets a new window), and the fill + goes on after ``layer``. For the first layer it sets ``primed`` False and + ``forward_open`` True. Wrapped in the NVTX range ``dwdp.dense.wait/L{layer}``. With a + watchdog, the check of every slot it waited for is then enqueued on the current stream, + inside the read windows, in the NVTX range ``dwdp.dense.watchdog/L{layer}``. + + Raises: + DenseProtocolError: See above; also before :meth:`commit` or for an uncovered layer. + """ + self._require_armed("wait_layer") + position = self._position(layer) + self.require_eager(f"wait_layer(L{layer})") + targets = self._targets[position] + with self._ops.nvtx_range(f"dwdp.dense.wait/L{layer}"): + stale = [(kind, slot) for kind, slot in targets if self._tag[kind][slot] != layer] + if stale: + self._demand_fetch(layer, stale) + compute = self._ops.current_stream() + for kind, slot in targets: + compute.wait_event(self._ready[kind][slot]) + for kind, slot in targets: + self._opened[kind][slot] = True + self._consumed[kind][slot] = False + self._next_fill = max(self._next_fill, position + 1) + if position == 0: + self._primed = False + self._forward_open = True + watchdog = self._watchdog + if watchdog is not None: + with self._ops.nvtx_range(f"dwdp.dense.watchdog/L{layer}"): + for kind, slot in targets: + watchdog.check(layer, kind, self._slots[kind][slot]) + + def attention_done(self, layer: int) -> None: + """H2, after the attention call: release ``A(layer)``, then fill. + + Closes the read window that H1 opened for ``A(layer)``: records ``consume[ATTN][s]`` on + the current stream and marks the slot consumed. Then the fill enqueues, in layer order, + each next layer whose target slots are free (empty, or their holder's read window has + closed), so each copy's wait on a consume event binds to a record already made; in a + forward that runs every hook in order that is ``L(layer + 2)``. + + Raises: + DenseProtocolError: The attention slot is not tagged ``layer``, or its read window is + not open because :meth:`wait_layer` (H1) did not open it since ``L(layer)`` was + enqueued or this hook already closed it (always a protocol bug, regardless of + ``strict``; raised before anything is recorded); before :meth:`commit`; an + uncovered layer. + """ + self._require_armed("attention_done") + position = self._position(layer) + slot = ring_slot(SegmentKind.ATTN, position) + self._close_window(SegmentKind.ATTN, slot, layer, "attention_done") + self._fill() + + def layer_done(self, layer: int) -> None: + """H3, after the MoE block: release ``D(layer)``, then fill. + + With the MoE-dense segment enabled, closes the read window that H1 opened for + ``D(layer)``: records ``consume[MOE][s]`` on the current stream and marks the slot + consumed. Then the fill runs as at H2; with ``prefetch_next_pass``, near the end of a + forward it enqueues the first two layers of the next forward as their slots free up. For + the last layer it sets ``forward_open`` False and, with ``prefetch_next_pass``, ``primed`` + True once both of them are enqueued (if not, ``strict`` raises; otherwise a warning is + logged and the next forward's priming enqueues them); then, with a watchdog, it reads the + watchdog's counters (a host synchronization with the current stream) and logs an error + naming every ``(layer, segment)`` whose slot held other bytes at H1 since the previous + read. + + Raises: + DenseProtocolError: The MoE-dense slot is not tagged ``layer``, or its read window is + not open (as for :meth:`attention_done`); before :meth:`commit`; an uncovered + layer; ``strict`` and the first two layers are not enqueued at the end of the + forward, or the watchdog found mismatches (both raised once the forward is + closed). + """ + self._require_armed("layer_done") + position = self._position(layer) + if self._has_moe: + slot = ring_slot(SegmentKind.MOE, position) + self._close_window(SegmentKind.MOE, slot, layer, "layer_done") + self._fill() + if position == len(self._layers) - 1: + self._forward_open = False + self._next_fill = max(0, self._next_fill - len(self._layers)) + if self._prefetch_next_pass: + self._finish_forward() + if self._watchdog is not None: + self._report_watchdog(self._watchdog.poll(), strict=self._strict) + + # ------------------------------------------------------------------ + # Priming and recovery (driven by the dense manager) + # ------------------------------------------------------------------ + + def prefetch(self, layer: int) -> None: + """Enqueue ``L(layer)`` (every enabled segment, one launch) on the copy stream. + + Used by the dense manager's priming (``L(first), L(second)``, then :meth:`mark_primed`). + Same as ``_prefetch(layer, fence=False)``: a target slot whose holder's read window is + open gets the fence, recorded on the current stream. The model runs H1-H3 of every + covered layer in order, so no window is open while ``forward_open`` is False: every + window closes by H3 of the last layer, and :meth:`war_barrier` resets them all. Priming + then needs no fence, whichever stream it runs on; slots left by an interrupted priming + hold layers that nobody has read and are refilled without one. + + Raises: + DenseProtocolError: Before :meth:`commit`; an uncovered layer; ``strict`` and a target + slot's holder has an open read window. + """ + self._require_armed("prefetch") + self._prefetch(layer, fence=False) + + def mark_primed(self) -> None: + """Record that the dense manager enqueued the first two layers; the fill goes on after them. + + Raises: + DenseProtocolError: Before :meth:`commit`; a forward is open (run + :meth:`war_barrier` first); the slots of the first two covered layers are not + tagged with them, i.e. :meth:`prefetch` did not enqueue them or a forward has + consumed them since. + """ + self._require_armed("mark_primed") + if self._forward_open: + raise DenseProtocolError( + "mark_primed while a forward is open; war_barrier must run before re-priming" + ) + missing = self._missing_first_layers() + if missing: + raise DenseProtocolError( + f"mark_primed but the first layers are not enqueued: {'; '.join(missing)}" + ) + self._primed = True + self._next_fill = min(2, len(self._layers)) + + def war_barrier(self) -> None: + """Recover after a forward aborted between H1 of the first and H3 of the last layer. + + Records every consume event on the current stream (the execution stream, inside the + model's H0), so later copies wait for every read the dead forward enqueued; then every tag + ``EMPTY_TAG``, no read window open, every ``consumed`` True, ``primed`` and + ``forward_open`` False, the fill at the first layer. Copies the dead forward left on the + copy stream still run first (FIFO), and the next ready records follow them. With expert + DWDP the expert prefetch recovers on its own: ``DWDPWeightManager.start_forward`` restarts + its sequence, after a device synchronization if the dead forward left a MoE layer bound. + + Raises: + DenseProtocolError: Before :meth:`commit`. + """ + self._require_armed("war_barrier") + self._record_consume_events() + self._reset_slot_state() + self._primed = False + self._forward_open = False + self._next_fill = 0 + logger.info(f"{LOG_PREFIX} WAR barrier after an aborted forward; slots reset") + + def _prefetch(self, layer: int, *, fence: bool) -> None: + """Enqueue ``L(layer)``: the one place that launches a copy under the event protocol. + + With ``compute = ops.current_stream()`` taken first and the copy stream current: for each + enabled segment's target slot, fall back to the fence if its holder's read window is open + (``war_fences`` grows; ``strict``: raise before anything is enqueued), and make the copy + stream wait its consume event; if fencing, record the fence on ``compute`` and make the + copy stream wait it; call the launcher once; record each target slot's ready event on the + copy stream, set its tag to ``layer`` and mark the new holder neither opened nor consumed. + If the launcher raises, the target slots are tagged ``EMPTY_TAG`` (their bytes are + unknown). Wrapped in the NVTX range ``dwdp.dense.prefetch/L{layer}``. + """ + targets = self._targets[self._position(layer)] + reading = [ + (kind, slot) + for kind, slot in targets + if self._opened[kind][slot] and not self._consumed[kind][slot] + ] + if reading and not fence: + self._fall_back_to_fence(layer, reading) + fence = True + launcher = self._launcher + if launcher is None: + raise DenseProtocolError(f"prefetch(L{layer}) without a launcher") + compute = self._ops.current_stream() + copy_stream = self._copy_stream + with self._ops.nvtx_range(f"dwdp.dense.prefetch/L{layer}"): + with self._ops.stream(copy_stream): + for kind, slot in targets: + copy_stream.wait_event(self._consume[kind][slot]) + if fence: + self._fence.record(compute) + copy_stream.wait_event(self._fence) + try: + launcher(layer) + except BaseException: + for kind, slot in targets: + self._tag[kind][slot] = EMPTY_TAG + raise + for kind, slot in targets: + self._ready[kind][slot].record(copy_stream) + self._tag[kind][slot] = layer + self._opened[kind][slot] = False + self._consumed[kind][slot] = False + + # ------------------------------------------------------------------ + # Teardown + # ------------------------------------------------------------------ + + def release(self) -> None: + """Release everything the runtime owns; idempotent. + + Order: ``ops.synchronize()`` -> the watchdog's last read (its mismatches are logged, never + raised) and its release -> detach the bindings (zero-size tensors) -> drop the launcher + and its plans -> drop the slots -> ``store.release()`` (peer mappings and imports, then the + local store). Every step runs even if an earlier one fails; the first error is raised at + the end and later ones are logged. The dense manager's ``cleanup`` calls it; with expert + DWDP, ``DwdpManager.cleanup`` runs that before the expert teardown. + """ + if self._phase is _Phase.RELEASED: + return + armed = self._phase is _Phase.ARMED + self._phase = _Phase.RELEASED + bindings, self._bindings = self._bindings, None + store, self._store = self._store, None + failures: list[tuple[str, Exception]] = [] + try: + self._ops.synchronize() + except Exception as exc: # every teardown step runs; the first failure is raised below + failures.append(("synchronize", exc)) + watchdog = self._watchdog + if watchdog is not None: + try: + self._report_watchdog(watchdog.poll(), strict=False) + if armed: + logger.info( + f"{LOG_PREFIX} slot watchdog: {watchdog.checks} slot checks, " + f"{watchdog.mismatches} mismatches" + ) + except Exception as exc: # every teardown step runs; the first failure is raised below + failures.append(("read the slot watchdog", exc)) + watchdog.release() + if bindings is not None: + try: + bindings.detach() + except Exception as exc: # every teardown step runs; the first failure is raised below + failures.append(("detach the module bindings", exc)) + self._launcher = None + self._slots = {} + if store is not None: + try: + store.release() + except Exception as exc: # every teardown step runs; the first failure is raised below + failures.append(("release the row store", exc)) + if failures: + for step, exc in failures[1:]: + logger.error(f"{LOG_PREFIX} release: {step} also failed: {exc!r}") + raise failures[0][1] + logger.debug(f"{LOG_PREFIX} runtime released") + + def __del__(self) -> None: + """Best-effort :meth:`release` for owners that miss the explicit teardown. + + Must tolerate an object whose ``__init__`` raised part way, and must never raise, also + at interpreter exit, when this module's globals may already be None: the released check + reaches the phase enum through the phase itself, and the failure log is skipped once the + logger is gone. + """ + phase = getattr(self, "_phase", None) + if phase is None or phase is type(phase).RELEASED: + return + try: + self.release() + except Exception as exc: # a finalizer must not raise + if logger is not None: + logger.debug(f"{LOG_PREFIX} release from __del__ failed: {exc!r}") + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + def _require_setup(self, action: str) -> None: + if self._phase is not _Phase.SETUP: + raise DenseProtocolError(f"dense DWDP {action} called after {self._phase_text()}") + + def _require_armed(self, action: str) -> None: + if self._phase is _Phase.ARMED: + return + if self._phase is _Phase.SETUP: + raise DenseProtocolError(f"dense DWDP {action} called before commit") + raise DenseProtocolError(f"dense DWDP {action} called after release") + + def _phase_text(self) -> str: + return "commit" if self._phase is _Phase.ARMED else "release" + + def _position(self, layer: int) -> int: + position = self._positions.get(layer) + if position is None: + raise DenseProtocolError( + f"layer {layer} is not covered by dense DWDP ({len(self._layers)} covered layers, " + f"{self._layers[0]}..{self._layers[-1]})" + ) + return position + + def _enabled(self, segment: SegmentKind) -> SegmentKind: + kind = SegmentKind(segment) + if kind not in self._tag: + raise DenseProtocolError(f"segment {kind.value} is not enabled") + return kind + + def _ring_index(self, kind: SegmentKind, index: int) -> int: + if not 0 <= index < len(self._tag[kind]): + raise IndexError(f"the {kind.value} ring has {len(self._tag[kind])} slots, got {index}") + return index + + def _holder(self, kind: SegmentKind, slot: int) -> str: + tag = self._tag[kind][slot] + return "nothing" if tag == EMPTY_TAG else f"layer {tag}" + + def _close_window(self, kind: SegmentKind, slot: int, layer: int, hook: str) -> None: + """End the read window of ``layer`` in ``kind``'s slot ``slot``; ``hook`` is the caller. + + Records the slot's consume event on the current stream and marks the slot consumed, after + checking that the slot holds ``layer`` and that the holder's read window is open: opened + by H1 and not closed since. A failed check raises before anything is recorded. + """ + if self._tag[kind][slot] != layer: + raise DenseProtocolError( + f"{hook}(L{layer}): {kind.value} slot {slot} holds {self._holder(kind, slot)}: " + f"the slot was refilled inside the read window of layer {layer}, or H1 " + f"(wait_layer) of layer {layer} did not run" + ) + if not self._opened[kind][slot]: + raise DenseProtocolError( + f"{hook}(L{layer}): H1 (wait_layer) of layer {layer} did not open the read window " + f"of {kind.value} slot {slot} since L({layer}) was enqueued, so the layer's reads " + f"did not wait for its copy" + ) + if self._consumed[kind][slot]: + raise DenseProtocolError( + f"{hook}(L{layer}): the read window of layer {layer} in {kind.value} slot {slot} " + f"is already closed; {hook} ran again without H1 (wait_layer) of layer {layer}" + ) + self._consume[kind][slot].record(self._ops.current_stream()) + self._consumed[kind][slot] = True + + def _slot_free(self, kind: SegmentKind, slot: int) -> bool: + """Whether a copy may go into the slot without the fence. + + True for an empty slot and for a holder whose read window H1 opened and H2 or H3 closed. + """ + return self._tag[kind][slot] == EMPTY_TAG or ( + self._opened[kind][slot] and self._consumed[kind][slot] + ) + + def _fill(self) -> None: + """Enqueue, in layer order, each next layer whose target slots are all free. + + Positions count from the current forward's first covered layer; with + ``prefetch_next_pass`` the fill goes on into the first two positions of the next forward, + which take the slots of positions ``0`` and ``1``. It stops at the first position with a + slot that is unread or being read. + """ + num_layers = len(self._layers) + limit = num_layers + min(2, num_layers) if self._prefetch_next_pass else num_layers + while self._next_fill < limit: + position = self._next_fill % num_layers + if not all(self._slot_free(kind, slot) for kind, slot in self._targets[position]): + return + self._prefetch(self._layers[position], fence=False) + self._next_fill += 1 + + def _missing_first_layers(self) -> list[str]: + """The slots of the first two covered layers that do not hold them, described.""" + return [ + f"L({layer}) {kind.value}[{slot}] holds {self._holder(kind, slot)}" + for position, layer in enumerate(self._layers[:2]) + for kind, slot in self._targets[position] + if self._tag[kind][slot] != layer + ] + + def _finish_forward(self) -> None: + """At the end of a forward: primed once the fill enqueued the next forward's first layers.""" + missing = self._missing_first_layers() + if not missing: + self._primed = True + return + message = ( + f"the forward ended before the next forward's first layers were enqueued: " + f"{'; '.join(missing)}" + ) + if self._strict: + raise DenseProtocolError(message) + self._unprimed_passes += 1 + _warn_sparsely( + f"{message}; the next forward primes them (#{self._unprimed_passes})", + self._unprimed_passes, + ) + + def _report_watchdog(self, findings: Sequence[Finding], *, strict: bool) -> None: + """Log the watchdog's new mismatches as an error; ``strict``: raise them instead.""" + if not findings: + return + total = sum(count for _, _, count in findings) + message = ( + f"slot watchdog: {total} slot checks at H1 found other bytes than their layer's " + f"checksum ({self._watchdog.describe(findings)}; {self._watchdog.mismatches} so far): " + f"a copy or a stream wait of dense DWDP is broken, or its row store changed" + ) + if strict: + raise DenseProtocolError(message) + logger.error(f"{LOG_PREFIX} {message}") + + def _record_consume_events(self) -> None: + stream = self._ops.current_stream() + for kind in self._segments: + for event in self._consume[kind]: + event.record(stream) + + def _reset_slot_state(self) -> None: + for kind in self._segments: + depth = len(self._tag[kind]) + self._tag[kind][:] = [EMPTY_TAG] * depth + self._opened[kind][:] = [False] * depth + self._consumed[kind][:] = [True] * depth + + def _demand_fetch(self, layer: int, stale: Sequence[_SlotKey]) -> None: + held = ", ".join( + f"{kind.value}[{slot}] holds {self._holder(kind, slot)}" for kind, slot in stale + ) + if self._strict: + raise DenseProtocolError( + f"wait_layer(L{layer}): {held}; L({layer}) was not enqueued ahead and " + f"TRTLLM_DWDP_DENSE_STRICT forbids demand fetches" + ) + self._demand_fetches += 1 + _warn_sparsely( + f"wait_layer(L{layer}): {held}; running a fenced demand fetch " + f"(demand fetch #{self._demand_fetches})", + self._demand_fetches, + ) + self._prefetch(layer, fence=True) + + def _fall_back_to_fence(self, layer: int, reading: Sequence[_SlotKey]) -> None: + held = ", ".join( + f"{kind.value}[{slot}] holds {self._holder(kind, slot)}" for kind, slot in reading + ) + if self._strict: + raise DenseProtocolError( + f"prefetch(L{layer}): {held}, whose read window opened at H1 and has no consume " + f"record; TRTLLM_DWDP_DENSE_STRICT forbids the fence fallback" + ) + self._war_fences += 1 + _warn_sparsely( + f"prefetch(L{layer}): {held}, whose read window opened at H1 and has no consume " + f"record; the copy waits for all work on the current stream " + f"(fence fallback #{self._war_fences})", + self._war_fences, + ) + + +def _warn_sparsely(message: str, count: int) -> None: + """Log the 1st, 2nd, 4th, 8th, ... occurrence, so a systematic fault cannot flood the log.""" + if count & (count - 1) == 0: + logger.warning(f"{LOG_PREFIX} {message}") + + +def _covered_layers(layers: Sequence[int]) -> tuple[int, ...]: + covered = tuple(operator.index(layer) for layer in layers) + if not covered: + raise ValueError("dense DWDP needs at least one covered layer") + if any(later <= earlier for earlier, later in zip(covered, covered[1:])): + raise ValueError(f"covered layers must be strictly ascending, got {list(covered)}") + if covered[0] < 0: + raise ValueError(f"covered layers must be non-negative, got {list(covered)}") + return covered + + +def _segment_kinds(segments: Sequence[SegmentKind]) -> tuple[SegmentKind, ...]: + kinds = tuple(SegmentKind(kind) for kind in segments) + canonical = tuple(kind for kind in SEGMENT_ORDER if kind in kinds) + if not kinds or kinds != canonical or kinds[0] is not SegmentKind.ATTN: + raise ValueError( + f"segments must be attn, optionally followed by moe, got " + f"{[kind.value for kind in kinds]}" + ) + return kinds + + +def _slot_rings( + kinds: Sequence[SegmentKind], slots: Mapping[SegmentKind, Sequence[SlotBuffer]] +) -> dict[SegmentKind, tuple[SlotBuffer, ...]]: + rings = {SegmentKind(kind): tuple(ring) for kind, ring in slots.items()} + if set(rings) != set(kinds): + raise ValueError( + f"slot rings must match the enabled segments {[kind.value for kind in kinds]}, " + f"got rings for {sorted(kind.value for kind in rings)}" + ) + for kind in kinds: + ring = rings[kind] + if len(ring) != RING_DEPTH[kind]: + raise ValueError( + f"the {kind.value} ring needs {RING_DEPTH[kind]} slots, got {len(ring)}" + ) + for index, slot in enumerate(ring): + if slot.segment != kind or slot.index != index: + raise ValueError( + f"position {index} of the {kind.value} ring holds slot " + f"{SegmentKind(slot.segment).value}[{slot.index}]" + ) + return {kind: rings[kind] for kind in kinds} diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/setup.py b/tensorrt_llm/_torch/modules/dwdp/dense/setup.py new file mode 100644 index 000000000000..d76c99f3c1cd --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/setup.py @@ -0,0 +1,932 @@ +# 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 setup: phases P1-P6 over five fixed collectives. + +Setup runs once per context worker, after the model is fully loaded and post-processed: from +``DwdpManager.setup`` before the expert DWDP setup (with expert DWDP, so the commit returns the +replicated dense weights to the driver before the expert prefetch ring is allocated), or from +``manager.DenseDwdpManager.setup`` (standalone, over the attention-DP group). Each phase's local +work is wrapped so that a failure becomes ``ok=False`` in the next collective instead of an +exception that would skip it; every rank makes the same collectives in the same order and reaches +the same verdict from the same gathered payloads. +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from typing import NamedTuple, Protocol + +import torch +from torch import nn + +from tensorrt_llm.logger import logger + +from ..setup import _get_decoder_model +from .config import DENSE_ENV, DENSE_REQUIRED_ENV, DenseConfig, dense_dwdp_requested +from .layout import collect_dense_layout, compute_checksums +from .plan import DenseCopyLauncher, build_copy_plans, select_copy_mode +from .repoint import BindingState, DenseBindings +from .runtime import DenseDWDP +from .store import DenseShardStore, FabricShardStore, StorePayload +from .types import ( + LOG_PREFIX, + Checksum, + DenseLayout, + DenseLayoutError, + DenseSetupError, + DeviceOps, + LayerCopyPlan, + LayerLauncher, + SegmentKind, + SlotBuffer, + StreamLike, + TorchCudaOps, + allocate_slot_ring, +) +from .watchdog import SlotWatchdog + +#: Free device memory P1 requires on top of the store and the slots. +SETUP_HEADROOM_BYTES = 2 << 30 + +#: The phase each collective concludes, by collective number. +_COLLECTIVES = { + 1: "P1 collection", + 2: "P2 store and slots", + 3: "P3 peer mapping and copy plans", + 4: "P4 Phase V", + 5: "P5 repoint self-check", +} + +#: Most paths or findings one message lists. +_MAX_LISTED = 8 + + +class GroupComm(Protocol): + """The calls setup makes on the dense group. + + The dense group is the DWDP group with expert DWDP (an mpi4py sub-communicator or + ``StoreDwdpGroup``), or the MPI sub-communicator over the attention-DP group standalone + (``manager.standalone_group``). + """ + + def allgather(self, obj: object) -> list[object]: + """Gather one picklable object per rank, ordered by rank.""" + ... + + def Get_rank(self) -> int: + """This rank within the group.""" + ... + + def Get_size(self) -> int: + """Ranks in the group.""" + ... + + +class DenseTransfer(Protocol): + """Where and how the dense copies run. + + ``DwdpManager``'s attached transfer (with expert DWDP: the copy stream the expert DWDP weight + manager then runs on, and the expert transfer; on a repeated setup the ``DWDPWeightManager`` + itself) and ``manager.StandaloneTransfer`` (standalone: the dense manager's own stream and + ``TRTLLM_DWDP_DENSE_COPY``) implement it. Setup first reads ``copy_stream`` in P2, inside the + phase's local work: an implementation may create the stream on that first read, and a failure + then becomes this rank's failed P2 report. + """ + + @property + def copy_stream(self) -> StreamLike: + """The stream every dense copy runs on; the same object on every read.""" + ... + + @property + def use_gather(self) -> bool: + """Whether the copies use the SM fabric gather (``CopyMode.GATHER``).""" + ... + + @property + def use_batched_prefetch(self) -> bool: + """Whether the copies use one batched copy-engine copy per layer (``CopyMode.BATCHED``).""" + ... + + +class PhaseReport(NamedTuple): + """What one rank contributes to a setup collective. + + Attributes: + ok: The phase's local work succeeded. + error: ``str`` of the local failure; empty when ``ok``. + data: Phase payload (picklable); see :func:`setup_dense_dwdp` for each collective. + """ + + ok: bool + error: str = "" + data: object = None + + +class _CollectData(NamedTuple): + """The data of one rank in collective #1.""" + + requested: bool + """``TRTLLM_DWDP_DENSE`` on this rank.""" + + required: bool + """``TRTLLM_DWDP_DENSE_REQUIRED`` on this rank; True if its switches could not be read.""" + + digest: str + """Layout digest; empty if collection failed.""" + + checksums: tuple[Checksum, ...] + """Per-segment content checksums in store order.""" + + store_bytes: int + """Store bytes this rank would allocate.""" + + slot_bytes: int + """Slot bytes this rank would allocate.""" + + free_bytes: int + """Free device memory before setup allocated anything, with the allocator cache emptied.""" + + strict: bool = False + """``TRTLLM_DWDP_DENSE_STRICT`` on this rank: it decides a raise after the commit point.""" + + +#: Collective #1 data of a rank that did not ask for dense DWDP. +_NOT_REQUESTED = _CollectData(False, True, "", (), 0, 0, 0) + +#: Builds the row store for ``(layout, rank, device)``; tests pass the loopback store. +StoreFactory = Callable[[DenseLayout, int, torch.device], DenseShardStore] + +#: Builds the copy launcher for ``(plans, transfer, device)``; setup calls it with the transfer's +#: copy stream current. +LauncherFactory = Callable[ + [Mapping[int, LayerCopyPlan], DenseTransfer, torch.device], LayerLauncher +] + +#: ``(free, total)`` bytes of the dense device, as ``torch.cuda.mem_get_info`` returns them. +MemGetInfo = Callable[[], tuple[int, int]] + + +def dense_roots_available(model: nn.Module) -> bool: + """Whether the model's decoder layers provide ``dense_dwdp_roots``. + + With expert DWDP, ``DwdpManager.setup`` calls :func:`setup_dense_dwdp` for such models whether + or not dense DWDP is requested on this rank, so collective #1 always runs and catches ranks + whose ``TRTLLM_DWDP_DENSE`` disagrees. Standalone, ``DenseDwdpManager.setup`` checks it + before it creates the dense group. The answer depends only on the model code, so it is the + same on every rank: for a model without roots no rank makes a dense collective (standalone, + no rank builds the group either), and a rank that asks for dense DWDP calls + :func:`refuse_dense_dwdp_without_roots` instead. The decoder is found as the expert DWDP + setup finds it; a model without one has no roots. + """ + try: + decoder = _get_decoder_model(model) + except RuntimeError: + return False + return any(callable(getattr(layer, "dense_dwdp_roots", None)) for layer in decoder.layers) + + +def refuse_dense_dwdp_without_roots(model: nn.Module) -> None: + """Turn down a request for dense DWDP on a model without ``dense_dwdp_roots``. + + ``DwdpManager.setup`` (with expert DWDP) calls it instead of :func:`setup_dense_dwdp` on a + rank whose ``TRTLLM_DWDP_DENSE`` is on when :func:`dense_roots_available` is False, and so + does ``DenseDwdpManager.setup`` (standalone, where every rank asks for dense DWDP) instead of + building the dense group. No rank makes a dense collective for such a model (whether it has + roots is the same on every rank), so this rank decides alone and at once rather than wait in + collective #1 for peers whose switch is off: it raises if its ``TRTLLM_DWDP_DENSE_REQUIRED`` + is on (also when its switches cannot be read), else logs an error and returns, and dense DWDP + stays off with the dense weights replicated. A rank whose ``TRTLLM_DWDP_DENSE`` is off does + not call it and carries on. + + Args: + model: The loaded model, named in the message. + + Raises: + DenseSetupError: ``TRTLLM_DWDP_DENSE_REQUIRED`` is on or unreadable. + """ + try: + required = DenseConfig.from_env().required + except ValueError: + required = True + message = ( + f"{DENSE_ENV} asks for dense DWDP, but the decoder layers of {type(model).__name__} do not " + f"provide dense_dwdp_roots, so it has no dense weights to shard" + ) + if required: + raise DenseSetupError(f"{message}; unset {DENSE_ENV} for this model") + logger.error( + f"{LOG_PREFIX} {message}; {DENSE_REQUIRED_ENV}=0, so dense DWDP is off on this rank and " + f"the dense weights stay replicated" + ) + + +def dense_dwdp_active(model: nn.Module) -> bool: + """Whether dense DWDP moved the dense weights of ``model``. + + True once setup P6 (or a release after the commit point) set ``_dense_dwdp_active = True`` + on the model or its decoder. The mark stays after the runtime is released, when the modules + hold zero-size tensors, so such a model can be neither reloaded nor set up again. The + decoder is found as the expert DWDP setup finds it. + """ + if getattr(model, "_dense_dwdp_active", False) is True: + return True + try: + decoder = _get_decoder_model(model) + except RuntimeError: + return False + return getattr(decoder, "_dense_dwdp_active", False) is True + + +def setup_dense_dwdp( + model: nn.Module, + transfer: DenseTransfer, + comm: GroupComm, + *, + dwdp_rank: int, + dwdp_size: int, + device_id: int, + moe_layers: Sequence[int] | None = None, + config: DenseConfig | None = None, + store_factory: StoreFactory | None = None, + ops: DeviceOps | None = None, + device: torch.device | None = None, + launcher_factory: LauncherFactory | None = None, + mem_get_info: MemGetInfo | None = None, +) -> DenseDWDP | None: + """Shard, verify and repoint the dense weights of ``model``. + + Each collective carries a :class:`PhaseReport` per rank; its local work and verdict: + + 1. P1: read the config, ``layout.collect_dense_layout`` over the decoder layers with the + ``dense_dwdp_attr_allow`` / ``dense_dwdp_exclude`` tables the layers carry, check that the + covered layers are ``moe_layers`` (if given) and that ``free >= store + slots + + SETUP_HEADROOM_BYTES`` (``free`` read after ``torch.cuda.empty_cache()``), then + ``layout.compute_checksums``; no copies. A rank that did not ask for dense DWDP does none + of this. Data ``(requested, required, digest, checksums, store_bytes, slot_bytes, + free_bytes, strict)``. No rank requested: return None and make no further collective. + Some did: every rank raises a configuration error, whatever ``required`` says. All ok + with equal digests, ``strict`` and checksums: continue; else abandon. The checksums catch + a dense tensor that has the same shape on every rank but other bytes, e.g. a + tensor-parallel shard. + 2. P2: allocate the store and the slots, write this rank's rows, ``synchronize``, construct + the runtime on ``transfer.copy_stream`` (the first read of it; it records the consume + events; with ``config.slot_watchdog`` it gets the slot checksum watchdog over the P1 + checksums), export. Data: the exported handles. + 3. P3: import and map the peers, build the plans and, with the copy stream current, the + launcher; install it, ``synchronize``. Doubles as the barrier after every write and + mapping. + 4. P4, Phase V: per layer ``DenseDWDP.fetch_blocking``, then ``torch.equal`` of every + anchor's slot bytes with its ``source``. Data: the mismatching anchor paths. + 5. P5: ``DenseBindings.repoint`` (originals alive), then per layer ``fetch_blocking`` and + ``DenseBindings.verify_layer``. Data: the mismatching binding paths. The commit point: + unless every rank is ok, every rank rolls back and abandons. + + After collectives 2 to 4, any rank not ok means every rank abandons. + + P6 (no collective, host only): ``DenseBindings.commit`` (drop originals, release check), + ``torch.cuda.empty_cache()``, ``DenseDWDP.commit`` with ``config.slot_fill_byte``, set + ``_dense_dwdp_active`` on the model and its decoder, log the commit line with the moved, + store, slot and freed bytes and the leak count. + + Every phase's local work catches ``Exception`` into ``ok=False``; ``BaseException`` propagates. + Abandoning releases everything this rank allocated: the bindings are rolled back, then the + runtime is released (launcher, slots, store) or, if no runtime was built, the store and the + slots; then the launcher. Then every rank raises :class:`~.types.DenseSetupError` if any + rank's ``config.required`` is set (an unreadable config counts as set), else logs an error + and returns None, with the model's original tensors in place (expert DWDP and the replicated + dense weights keep working). A rank whose rollback fails raises regardless. Any exception + that escapes (a broken group, ``BaseException``, a P6 failure) releases this rank's dense + state the same way first; bindings already committed are detached instead of rolled back. + + Args: + model: The loaded model (the VL wrapper or the causal LM); the decoder is found like the + expert setup finds it. + transfer: Copy stream and transfer mode: with expert DWDP, ``DwdpManager``'s attached + transfer, whose copy stream the expert DWDP weight manager then runs on; a + ``manager.StandaloneTransfer`` standalone. + comm: The dense group: the DWDP group, or the attention-DP group standalone. + dwdp_rank: This rank in the dense group (its DWDP rank, or its tp rank standalone). + dwdp_size: Dense group size ``N``. + device_id: CUDA device of this rank. + moe_layers: Expert DWDP MoE layers; the covered dense layers must equal them, and the + runtime prefetches the next forward's first layers at the end of a forward + (``DenseDWDP``'s ``prefetch_next_pass``). None (standalone): no such check; the + covered layers are the decoder layers whose dense roots are not empty, ascending, and + the dense manager primes every forward. + config: Switches; ``DenseConfig.from_env()`` (inside P1) if None. ``config.enabled`` (or + ``TRTLLM_DWDP_DENSE``) says whether this rank asks for dense DWDP. + store_factory: Row store constructor; ``FabricShardStore`` if None. + ops: CUDA seam for the runtime; ``TorchCudaOps(device_id)`` if None. + device: The dense device; ``cuda:`` if None (unit tests pass the CPU). + launcher_factory: Builds the copy launcher from the plans; if None, + ``DenseCopyLauncher(plans, select_copy_mode(use_gather=transfer.use_gather, + use_batched=transfer.use_batched_prefetch), device=device)``. + mem_get_info: ``(free, total)`` device memory; ``torch.cuda.mem_get_info(device_id)`` if + None. + + Returns: + The committed runtime, or None when dense DWDP is off on every rank or was abandoned + with ``required`` off on every rank. + + Raises: + DenseSetupError: Abandoned with ``required`` on, ranks disagree on + ``TRTLLM_DWDP_DENSE``, or this rank could not restore its original dense weights. + DenseRepointError: ``config.strict`` and an original storage survived the commit (this + rank only; its dense state is released first). ``config.strict`` must be the same on + every rank: collective #1 abandons when it differs. + """ + if config is not None: + requested = config.enabled + else: + requested = dense_dwdp_requested() + run = _DenseSetup( + model, + transfer, + comm, + rank=dwdp_rank, + size=dwdp_size, + device_id=device_id, + device=device, + moe_layers=moe_layers, + config=config, + requested=requested, + store_factory=FabricShardStore if store_factory is None else store_factory, + ops=TorchCudaOps(device_id) if ops is None else ops, + launcher_factory=_copy_launcher if launcher_factory is None else launcher_factory, + mem_get_info=_device_mem_get_info(device_id) if mem_get_info is None else mem_get_info, + ) + return run.run() + + +class _DenseSetup: + """One rank's run of :func:`setup_dense_dwdp`: its phases, collectives and resources. + + It holds what this rank allocates until the runtime owns it: the store and the slots until + the runtime constructor returns, the bindings until ``DenseDWDP.commit``, the launcher until + setup ends. Every rank calls :meth:`_exchange` exactly when every other rank does, because + whether collective ``k + 1`` happens depends only on the reports collective ``k`` gathered. + """ + + def __init__( + self, + model: nn.Module, + transfer: DenseTransfer, + comm: GroupComm, + *, + rank: int, + size: int, + device_id: int, + device: torch.device | None, + moe_layers: Sequence[int] | None, + config: DenseConfig | None, + requested: bool, + store_factory: StoreFactory, + ops: DeviceOps, + launcher_factory: LauncherFactory, + mem_get_info: MemGetInfo, + ) -> None: + self._model = model + self._transfer = transfer + self._comm = comm + self._rank = rank + self._size = size + self._device_id = device_id + self._device_override = device + self._moe_layers = moe_layers + self._config = config + self._requested = requested + self._store_factory = store_factory + self._ops = ops + self._launcher_factory = launcher_factory + self._mem_get_info = mem_get_info + self._started = time.monotonic() + # Decided from collective #1; raising is the safe default until then. + self._required = True + self._device: torch.device | None = None + self._decoder: nn.Module | None = None + self._free_before = 0 + self._allocated_before = 0 + self._layout: DenseLayout | None = None + self._store: DenseShardStore | None = None + self._slots: dict[SegmentKind, tuple[SlotBuffer, ...]] | None = None + self._runtime: DenseDWDP | None = None + self._launcher: LayerLauncher | None = None + self._bindings: DenseBindings | None = None + + def run(self) -> DenseDWDP | None: + """Run the phases; whatever escapes releases this rank's dense state first.""" + try: + return self._run() + except BaseException: + self._release() + raise + + def _run(self) -> DenseDWDP | None: + reports = self._exchange(1, self._collect()) + data = [report.data for report in reports] + if not any(item.requested for item in data): + logger.debug(f"{LOG_PREFIX} TRTLLM_DWDP_DENSE is off on every rank; dense DWDP is off") + return None + if not all(item.requested for item in data): + on = [rank for rank, item in enumerate(data) if item.requested] + off = [rank for rank, item in enumerate(data) if not item.requested] + raise DenseSetupError( + f"TRTLLM_DWDP_DENSE differs across the DWDP group: on for {_ranks(on)}, off for " + f"{_ranks(off)}; set it the same on every context worker of the group" + ) + self._required = any(item.required for item in data) + reason = _failures(reports) or _disagreement(self._layout, data) + if reason: + return self._abandon(1, reason) + self._log_agreement(data) + + reports = self._exchange(2, self._allocate()) + if reason := _failures(reports): + return self._abandon(2, reason) + reports = self._exchange(3, self._map([report.data for report in reports])) + if reason := _failures(reports): + return self._abandon(3, reason) + reports = self._exchange(4, self._compare_blobs()) + if reason := _failures(reports): + return self._abandon(4, reason) + reports = self._exchange(5, self._repoint()) + if reason := _failures(reports): + return self._abandon(5, reason) + return self._commit() + + # ------------------------------------------------------------------ + # Collectives and verdicts + # ------------------------------------------------------------------ + + def _exchange(self, number: int, report: PhaseReport) -> list[PhaseReport]: + """Collective ``number``: every rank's report in rank order; malformed ones count as failed.""" + reports = [_checked_report(number, item) for item in self._comm.allgather(report)] + logger.debug( + f"{LOG_PREFIX} collective #{number} ({_COLLECTIVES[number]}): " + f"{sum(report.ok for report in reports)}/{len(reports)} ranks ok" + ) + return reports + + def _abandon(self, number: int, reason: str) -> None: + """Release this rank's dense state, then raise or return as the group's verdict says.""" + intact = self._release() + torch.cuda.empty_cache() + message = ( + f"dense DWDP setup abandoned at collective #{number} ({_COLLECTIVES[number]}): {reason}" + ) + if not intact: + raise DenseSetupError( + f"{message}; this rank could not restore its original dense weights (see the " + f"errors logged above)" + ) + if self._required: + raise DenseSetupError(message) + logger.error( + f"{LOG_PREFIX} {message}; TRTLLM_DWDP_DENSE_REQUIRED=0 on every rank, so dense DWDP " + f"is off and the dense weights stay replicated" + ) + return None + + def _log_agreement(self, data: Sequence[_CollectData]) -> None: + layout = self._layout + lowest = min(range(len(data)), key=lambda rank: data[rank].free_bytes) + logger.info( + f"{LOG_PREFIX} layout agreed by {len(data)} ranks: {len(layout.layers)} layers " + f"{layout.layers[0]}..{layout.layers[-1]}, segments {_segment_names(layout)}, " + f"moved={_gb(layout.moved_bytes)}GB store={_gb(layout.store_bytes)}GB " + f"slots={_gb(layout.slot_bytes)}GB, digest {layout.digest[:16]}; lowest free memory " + f"{_gb(data[lowest].free_bytes)}GB on rank {lowest}" + ) + + # ------------------------------------------------------------------ + # Phases (local work; every failure becomes a report) + # ------------------------------------------------------------------ + + def _collect(self) -> PhaseReport: + """P1, no copies: switches, collection, coverage, memory precheck, checksums.""" + if not self._requested: + return PhaseReport(True, data=_NOT_REQUESTED) + required = True if self._config is None else self._config.required + free_bytes = 0 + try: + config = DenseConfig.from_env() if self._config is None else self._config + self._config = config + required = config.required + self._check_group() + if self._device_override is None: + self._device = torch.device("cuda", self._device_id) + else: + self._device = torch.device(self._device_override) + # Blocks the caching allocator holds count as free: the store needs driver memory, + # and P6 measures freed(mem_get_info) after its own empty_cache. + torch.cuda.empty_cache() + free_bytes, _ = self._mem_get_info() + self._allocated_before = _allocated_bytes(self._device) + self._decoder = _get_decoder_model(self._model) + layers = list(enumerate(self._decoder.layers)) + attr_allow, exclude = _model_tables(layers) + layout = collect_dense_layout( + layers, + config, + num_ranks=self._size, + device=self._device, + attr_allow=attr_allow, + exclude=exclude, + ) + if self._moe_layers is not None: + _check_coverage(layout, self._moe_layers) + _check_free_memory(layout, free_bytes) + compute_checksums(layout) + data = _CollectData( + requested=True, + required=required, + digest=layout.digest, + checksums=layout.checksums(), + store_bytes=layout.store_bytes, + slot_bytes=layout.slot_bytes, + free_bytes=free_bytes, + strict=config.strict, + ) + except Exception as exc: # a local failure must reach collective #1, never skip it + return PhaseReport( + False, _describe(exc), _CollectData(True, required, "", (), 0, 0, free_bytes) + ) + self._layout = layout + self._free_before = free_bytes + return PhaseReport(True, data=data) + + def _allocate(self) -> PhaseReport: + """P2: store, slots, this rank's rows, synchronize, runtime and its watchdog, export.""" + layout = self._layout + try: + store = self._store_factory(layout, self._rank, self._device) + self._store = store + store.allocate() + self._slots = { + kind: allocate_slot_ring(layout.slot_specs[kind], self._device) + for kind in layout.segment_kinds + } + store.write_local_rows() + self._ops.synchronize() + watchdog = None + if self._config.slot_watchdog: + watchdog = SlotWatchdog.from_layout(layout, self._device) + # The first read of the copy stream, which the transfer may create here: a failure + # is this rank's P2 report. + self._runtime = DenseDWDP( + layout.layers, + layout.segment_kinds, + self._slots, + self._transfer.copy_stream, + # With expert DWDP the next forward's first layers are gathered at the end of a + # forward, beside the expert prefetch; standalone, the dense manager primes them. + prefetch_next_pass=self._moe_layers is not None, + ops=self._ops, + strict=self._config.strict, + store=store, + watchdog=watchdog, + ) + payload = store.export() + except Exception as exc: # a local failure must reach collective #2, never skip it + return PhaseReport(False, _describe(exc)) + return PhaseReport(True, data=payload) + + def _map(self, payloads: Sequence[StorePayload]) -> PhaseReport: + """P3: map the peers' rows, build the plans and the launcher, synchronize.""" + try: + store = self._store + store.import_peers(payloads) + slot_bases = {kind: [slot.base for slot in ring] for kind, ring in self._slots.items()} + plans = build_copy_plans( + self._layout, + rank=self._rank, + chunk_bases=store.chunk_bases(), + slot_bases=slot_bases, + ) + with self._ops.stream(self._transfer.copy_stream): + self._launcher = self._launcher_factory(plans, self._transfer, self._device) + self._runtime.set_launcher(self._launcher) + self._ops.synchronize() + except Exception as exc: # a local failure must reach collective #3, never skip it + return PhaseReport(False, _describe(exc)) + return PhaseReport(True) + + def _compare_blobs(self) -> PhaseReport: + """P4, Phase V: gather every layer and compare each anchor's slot bytes with its source.""" + layout = self._layout + mismatches: list[str] = [] + try: + for layer in layout.layers: + self._runtime.fetch_blocking(layer) + for segment in layout.layer_segments(layer): + slot = self._slots[segment.segment][layout.slot_index(segment.segment, layer)] + blob = slot.blob_view(segment.blob_bytes) + for anchor in segment.anchors: + gathered = blob[anchor.blob_offset : anchor.blob_end] + if not torch.equal(gathered, anchor.source): + mismatches.append(anchor.path) + except Exception as exc: # a local failure must reach collective #4, never skip it + return PhaseReport(False, _describe(exc), mismatches) + if mismatches: + return PhaseReport( + False, + f"Phase V: {len(mismatches)} gathered anchors differ from their originals: " + f"{_listing(mismatches)}", + mismatches, + ) + return PhaseReport(True, data=mismatches) + + def _repoint(self) -> PhaseReport: + """P5: repoint every binding, then gather every layer and self-check it through the modules.""" + mismatches: list[str] = [] + try: + bindings = DenseBindings(self._layout, self._slots) + self._bindings = bindings + bindings.repoint() + for layer in self._layout.layers: + self._runtime.fetch_blocking(layer) + mismatches.extend(bindings.verify_layer(layer)) + except Exception as exc: # a local failure must reach collective #5, never skip it + return PhaseReport(False, _describe(exc), mismatches) + if mismatches: + return PhaseReport( + False, + f"P5 self-check: {len(mismatches)} module tensors are not their slot views or " + f"differ from their P1 checksums after repointing: {_listing(mismatches)}", + mismatches, + ) + return PhaseReport(True, data=mismatches) + + def _commit(self) -> DenseDWDP: + """P6, past the commit point, host only: drop the originals and arm the runtime.""" + config, layout = self._config, self._layout + bindings, runtime = self._bindings, self._runtime + leaks = bindings.commit(strict=config.strict, describe_referrers=config.debug >= 1) + torch.cuda.empty_cache() + free_after, _ = self._mem_get_info() + allocated_after = _allocated_bytes(self._device) + runtime.commit(bindings, fill_byte=config.slot_fill_byte) + # The runtime owns the bindings, the slots, the store and the launcher from here on. + self._bindings = self._runtime = self._store = self._slots = self._launcher = None + self._mark_active() + logger.info( + f"{LOG_PREFIX} committed: moved={_gb(layout.moved_bytes)}GB " + f"store={_gb(layout.store_bytes)}GB slots={_gb(layout.slot_bytes)}GB " + f"freed(mem_get_info)={_gb(free_after - self._free_before)}GB " + f"freed(memory_allocated)={_gb(self._allocated_before - allocated_after)}GB " + f"leaks={len(leaks)} layers={len(layout.layers)} segments={_segment_names(layout)} " + f"setup={time.monotonic() - self._started:.1f}s" + ) + return runtime + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _check_group(self) -> None: + """The group must be the one the caller describes; the rows are indexed by its ranks.""" + group = (self._comm.Get_rank(), self._comm.Get_size()) + if group != (self._rank, self._size): + raise DenseSetupError( + f"the DWDP group reports rank {group[0]} of {group[1]}, but dense setup runs as " + f"rank {self._rank} of {self._size}" + ) + + def _mark_active(self) -> None: + """Flag the model and its decoder: their dense tensors are no longer their own weights. + + The weight reload, update and load guards read this flag. + """ + self._model._dense_dwdp_active = True + if self._decoder is not None: + self._decoder._dense_dwdp_active = True + + def _release(self) -> bool: + """Undo this rank's setup work; every step runs and failures are logged, never raised. + + Order: the bindings go back to the originals (to zero-size tensors once committed); then + the runtime is released, which synchronizes, drops the launcher and the slots and + releases the store, or, if no runtime was built, the store is released and the slots + dropped; then the launcher's own ``release`` runs. + + Returns: + Whether the modules hold their original tensors. + """ + intact = True + bindings, self._bindings = self._bindings, None + if bindings is not None: + committed = bindings.state in (BindingState.COMMITTED, BindingState.DETACHED) + if committed: + intact = False + self._mark_active() + try: + if committed: + bindings.detach() + else: + bindings.rollback() + except Exception as exc: # every release step runs; the failure is logged + intact = False + logger.error(f"{LOG_PREFIX} restoring the module tensors failed: {exc!r}") + runtime, self._runtime = self._runtime, None + store, self._store = self._store, None + self._slots = None + try: + if runtime is not None: + runtime.release() + elif store is not None: + store.release() + except Exception as exc: # every release step runs; the failure is logged + logger.error(f"{LOG_PREFIX} releasing the dense store and slots failed: {exc!r}") + launcher, self._launcher = self._launcher, None + release = getattr(launcher, "release", None) + if callable(release): + try: + release() + except Exception as exc: # every release step runs; the failure is logged + logger.error(f"{LOG_PREFIX} releasing the dense copy launcher failed: {exc!r}") + self._layout = None + return intact + + +def _copy_launcher( + plans: Mapping[int, LayerCopyPlan], transfer: DenseTransfer, device: torch.device +) -> LayerLauncher: + """The production launcher in the transfer's mode. + + With expert DWDP that is the transfer the expert DWDP prefetch uses; standalone, the one + ``TRTLLM_DWDP_DENSE_COPY`` picks. + """ + mode = select_copy_mode( + use_gather=transfer.use_gather, use_batched=transfer.use_batched_prefetch + ) + return DenseCopyLauncher(plans, mode, device=device) + + +def _device_mem_get_info(device_id: int) -> MemGetInfo: + def mem_get_info() -> tuple[int, int]: + return torch.cuda.mem_get_info(device_id) + + return mem_get_info + + +def _allocated_bytes(device: torch.device) -> int: + """Bytes the caching allocator hands out on ``device``; 0 for a device that is not CUDA.""" + return torch.cuda.memory_allocated(device) if device.type == "cuda" else 0 + + +def _model_tables(layers: Sequence[tuple[int, nn.Module]]) -> tuple[frozenset[str], frozenset[str]]: + """The model's allow and exclude tables, read from the decoder layers. + + Raises: + DenseLayoutError: Decoder layers carry different tables; the collector applies one pair + to every layer. + """ + tables: dict[tuple[frozenset[str], frozenset[str]], list[int]] = {} + for index, layer in layers: + pair = ( + frozenset(getattr(layer, "dense_dwdp_attr_allow", frozenset())), + frozenset(getattr(layer, "dense_dwdp_exclude", frozenset())), + ) + tables.setdefault(pair, []).append(index) + if len(tables) > 1: + groups = "; ".join(f"layers {indices[:_MAX_LISTED]}" for indices in tables.values()) + raise DenseLayoutError( + f"the decoder layers carry {len(tables)} different dense_dwdp_attr_allow / " + f"dense_dwdp_exclude tables ({groups}); dense DWDP applies one pair to every layer" + ) + if not tables: + return frozenset(), frozenset() + (pair,) = tables + return pair + + +def _check_coverage(layout: DenseLayout, moe_layers: Sequence[int]) -> None: + """The covered layers must be the expert DWDP MoE layers. + + The dense priming and the expert prefetch start with the same first two layers, and each + layer's dense copy ``L(l)`` pairs with its expert copy ``E(l)`` on the shared copy stream. + + Raises: + DenseLayoutError: They differ, or ``moe_layers`` repeats a layer. + """ + expected = sorted(moe_layers) + if len(set(expected)) != len(expected): + raise DenseLayoutError(f"the expert DWDP MoE layers repeat a layer: {expected}") + if list(layout.layers) != expected: + dense_only = sorted(set(layout.layers) - set(expected)) + expert_only = sorted(set(expected) - set(layout.layers)) + raise DenseLayoutError( + f"dense DWDP covers {len(layout.layers)} layers and expert DWDP manages " + f"{len(expected)}; they must be the same layers: only dense " + f"{dense_only[:_MAX_LISTED]}, only expert {expert_only[:_MAX_LISTED]}" + ) + + +def _check_free_memory(layout: DenseLayout, free_bytes: int) -> None: + """P1 precheck: ``free >= store + slots + SETUP_HEADROOM_BYTES``. + + Raises: + DenseSetupError: Too little free device memory. + """ + needed = layout.store_bytes + layout.slot_bytes + SETUP_HEADROOM_BYTES + if free_bytes < needed: + raise DenseSetupError( + f"{_gb(free_bytes)} GB of device memory is free, but the store " + f"({_gb(layout.store_bytes)} GB), the slots ({_gb(layout.slot_bytes)} GB) and the " + f"setup headroom ({_gb(SETUP_HEADROOM_BYTES)} GB) need {_gb(needed)} GB" + ) + + +def _checked_report(number: int, item: object) -> PhaseReport: + """``item`` as a report of collective ``number``; anything malformed is a failed report.""" + if not isinstance(item, PhaseReport): + item = PhaseReport(False, f"sent a {type(item).__name__} instead of a PhaseReport") + if number == 1 and not isinstance(item.data, _CollectData): + return PhaseReport( + False, + item.error or "sent no collective #1 data", + _CollectData(True, True, "", (), 0, 0, 0), + ) + return item + + +def _failures(reports: Sequence[PhaseReport]) -> str: + """The failed ranks grouped by error; empty if every rank is ok.""" + by_error: dict[str, list[int]] = {} + for rank, report in enumerate(reports): + if not report.ok: + by_error.setdefault(report.error or "failed without a message", []).append(rank) + return "; ".join(f"{_ranks(ranks)}: {error}" for error, ranks in by_error.items()) + + +def _disagreement(layout: DenseLayout | None, data: Sequence[_CollectData]) -> str: + """Why ok ranks cannot share rows: different layouts, or the same layout with other bytes. + + Nor can ranks whose ``TRTLLM_DWDP_DENSE_STRICT`` differs go on together: a leak found after + the commit point would raise on some of them only. + """ + digests = _group_ranks(item.digest for item in data) + if len(digests) > 1: + listed = "; ".join(f"{digest[:16]} on {_ranks(ranks)}" for digest, ranks in digests.items()) + return ( + f"the dense layouts differ across ranks (digest {listed}): the ranks' models or " + f"TRTLLM_DWDP_DENSE_* switches differ" + ) + stricts = _group_ranks(item.strict for item in data) + if len(stricts) > 1: + listed = " vs ".join(f"{int(value)} on {_ranks(ranks)}" for value, ranks in stricts.items()) + return f"TRTLLM_DWDP_DENSE_STRICT differs across ranks ({listed})" + findings = [] + for index, (layer, kind) in enumerate(layout.segments): + holders = _group_ranks(item.checksums[index] for item in data) + if len(holders) > 1: + findings.append( + f"layers.{layer} {kind.value}: " + + " vs ".join(_ranks(ranks) for ranks in holders.values()) + ) + if findings: + return f"the dense weights differ across ranks: {_listing(findings)}" + return "" + + +def _group_ranks(values: Iterable[object]) -> dict[object, list[int]]: + groups: dict[object, list[int]] = {} + for rank, value in enumerate(values): + groups.setdefault(value, []).append(rank) + return groups + + +def _ranks(ranks: Sequence[int]) -> str: + return f"rank {ranks[0]}" if len(ranks) == 1 else f"ranks {list(ranks)}" + + +def _describe(exc: Exception) -> str: + return f"{type(exc).__name__}: {exc}" + + +def _listing(findings: Sequence[str]) -> str: + listed = "; ".join(findings[:_MAX_LISTED]) + more = len(findings) - _MAX_LISTED + return listed + (f"; and {more} more" if more > 0 else "") + + +def _segment_names(layout: DenseLayout) -> str: + return ",".join(kind.value for kind in layout.segment_kinds) + + +def _gb(nbytes: int) -> str: + return f"{nbytes / 1e9:.2f}" diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/store.py b/tensorrt_llm/_torch/modules/dwdp/dense/store.py new file mode 100644 index 000000000000..77e00c6c213a --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/store.py @@ -0,0 +1,611 @@ +# 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. +"""The sharded row store of dense DWDP. + +Every rank holds its row of every ``(layer, segment)`` blob in one handle per store chunk (one +handle unless ``TRTLLM_DWDP_DENSE_STORE_CHUNK_MB`` is set); every rank maps every peer's handles +read-only. The store is written once in setup P2 and never changes afterwards; the runtime reads +it only through the copy plans. + +* :class:`FabricShardStore` -- FABRIC handles exported over the DWDP group (aarch64 / NVL72). +* :class:`LoopbackShardStore` -- all ``N`` ranks' rows in this process, for single-process tests + of the whole path (collector, rows, gather, repoint). + +Addressing: rank ``p``'s row of ``(layer, segment)`` starts at +``chunk_bases()[p][segment.chunk] + segment.chunk_offset`` and is ``segment.row_bytes`` long; +``chunk_bases()[rank]`` is this rank's own, writable mapping. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable, Sequence +from dataclasses import dataclass + +import torch + +from tensorrt_llm.logger import logger + +from .. import transport, vmm +from .config import DENSE_STORE_CHUNK_MB_ENV +from .layout import row_pieces +from .types import LOG_PREFIX, ROW_ALIGN, DenseLayout, DenseSetupError + +#: What one rank contributes to collective #2: one exported handle per store chunk, or None for +#: the loopback store. Must be picklable. +StorePayload = list[bytes] | None + +_FABRIC = vmm.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + + +class DenseShardStore(ABC): + """This rank's rows plus a read-only view of every peer's rows. + + Call order (setup): :meth:`allocate` and :meth:`write_local_rows` in P2, then + ``synchronize``, then :meth:`export` (collective #2), then :meth:`import_peers` in P3, after + which :meth:`chunk_bases` is complete. :meth:`release` is valid at any point and idempotent. + + Each step runs at most once: calling :meth:`allocate` or :meth:`import_peers` a second time + (also after one that raised part way) raises ``RuntimeError``, and so does every call except + :meth:`release` after :meth:`release`. Whatever a failed step created is freed by + :meth:`release`. + + Args: + layout: The shared layout (chunks and rows); the collected anchors supply the bytes. + rank: This rank. + device: The dense CUDA device (CPU is allowed for the loopback store in unit tests). + + Raises: + ValueError: ``rank`` is outside ``[0, N)``, or a segment's rows do not fit its chunk at a + ``ROW_ALIGN``-aligned offset. + """ + + def __init__(self, layout: DenseLayout, rank: int, device: torch.device) -> None: + num_ranks = layout.num_ranks + if not 0 <= rank < num_ranks: + raise ValueError(f"rank must be in [0, {num_ranks}), got {rank}") + for segment in layout.segments.values(): + if ( + not 0 <= segment.chunk < len(layout.chunks) + or segment.chunk_offset < 0 + or segment.chunk_offset % ROW_ALIGN + or segment.chunk_offset + segment.row_bytes > layout.chunks[segment.chunk].nbytes + ): + raise ValueError( + f"the rows of ({segment.layer}, {segment.segment.value}) do not fit store " + f"chunk {segment.chunk} at offset {segment.chunk_offset} " + f"(S={segment.row_bytes}); chunk offsets must be {ROW_ALIGN}-aligned" + ) + self._layout = layout + self._rank = rank + self._num_ranks = num_ranks + self._device = torch.device(device) + self._allocate_started = False + self._allocated = False + self._import_started = False + self._bases: tuple[tuple[int, ...], ...] | None = None + # Set last: ``__del__`` releases nothing of an object whose constructor raised. + self._released = False + + @abstractmethod + def allocate(self) -> None: + """P2: allocate and map the local chunks (``StoreChunk.handle_bytes`` each), zero-filled. + + Raises: + RuntimeError: The allocation or mapping was refused (e.g. a FABRIC handle that large). + """ + + @abstractmethod + def write_local_rows(self) -> None: + """P2: copy this rank's row of every ``(layer, segment)`` into its local chunk. + + Bytes come from ``layout.row_pieces`` over the anchors' ``source``; everything else stays + zero. Asynchronous on the current stream; the caller synchronizes before :meth:`export`. + """ + + @abstractmethod + def export(self) -> StorePayload: + """P2: the payload peers need to import this rank's chunks (collective #2).""" + + @abstractmethod + def import_peers(self, payloads: Sequence[StorePayload]) -> None: + """P3: import and map every peer's chunks read-only. + + Args: + payloads: Collective #2 results, indexed by rank (this rank's entry included). + """ + + @abstractmethod + def chunk_bases(self) -> tuple[tuple[int, ...], ...]: + """``bases[p][c]``: address of rank ``p``'s chunk ``c`` as mapped here; ``ROW_ALIGN``-aligned. + + Raises: + RuntimeError: Before :meth:`import_peers`. + """ + + @abstractmethod + def release(self) -> None: + """Unmap and free the peer imports, then the local chunks; idempotent.""" + + def __del__(self) -> None: + """Best-effort :meth:`release` for owners that miss the explicit teardown. + + Must tolerate an object whose ``__init__`` raised part way, and must never raise. + """ + if getattr(self, "_released", True): + return + try: + self.release() + except Exception as exc: # a finalizer must not raise + logger.debug(f"{LOG_PREFIX} store release from __del__ failed: {exc!r}") + + # ------------------------------------------------------------------ + # Shared helpers + # ------------------------------------------------------------------ + + def _require_open(self, action: str) -> None: + if self._released: + raise RuntimeError(f"dense DWDP store: {action} called after release") + + def _begin_allocate(self) -> None: + """Claim the single :meth:`allocate` of this store.""" + self._require_open("allocate") + if self._allocate_started: + raise RuntimeError("dense DWDP store: allocate already ran") + self._allocate_started = True + + def _require_allocated(self, action: str) -> None: + self._require_open(action) + if not self._allocated: + raise RuntimeError(f"dense DWDP store: {action} called before a successful allocate") + + def _begin_import(self, payloads: Sequence[StorePayload]) -> list[StorePayload]: + """Claim the single :meth:`import_peers` and return the payloads as a list of ``N``.""" + self._require_allocated("import_peers") + if self._import_started: + raise RuntimeError("dense DWDP store: import_peers already ran") + self._import_started = True + gathered = list(payloads) + if len(gathered) != self._num_ranks: + raise ValueError( + f"import_peers needs one payload per rank ({self._num_ranks}), got {len(gathered)}" + ) + return gathered + + def _checked_bases(self) -> tuple[tuple[int, ...], ...]: + self._require_open("chunk_bases") + if self._bases is None: + raise RuntimeError("dense DWDP store: chunk_bases called before import_peers") + return self._bases + + def _check_sources(self) -> None: + """Every anchor still holds its flat ``uint8`` source of the recorded size. + + Raises: + DenseSetupError: A source was dropped (after commit) or does not match its anchor. + """ + for segment in self._layout.segments.values(): + for anchor in segment.anchors: + source = anchor.source + if source is None: + raise DenseSetupError( + f"cannot write the store rows: the source of {anchor.path} was dropped" + ) + if ( + source.dtype != torch.uint8 + or source.dim() != 1 + or source.numel() != anchor.nbytes + ): + raise DenseSetupError( + f"the source of {anchor.path} is {source.dtype} {tuple(source.shape)}; " + f"expected a flat uint8 tensor of {anchor.nbytes} bytes" + ) + + def _write_rows(self, row: int, chunks: Sequence[torch.Tensor]) -> None: + """Copy row ``row`` of every ``(layer, segment)`` into ``chunks[c]`` at its chunk offset. + + Args: + row: The row (rank) to write. + chunks: Flat ``uint8`` tensors over the chunks of the store holding row ``row``. + """ + for segment in self._layout.segments.values(): + target = chunks[segment.chunk] + for piece in row_pieces(segment, row, num_ranks=self._num_ranks): + source = segment.anchors[piece.anchor_index].source + start = segment.chunk_offset + piece.row_offset + target[start : start + piece.nbytes].copy_( + source[piece.anchor_offset : piece.anchor_offset + piece.nbytes] + ) + + def _synchronize_current_stream(self) -> None: + if self._device.type == "cuda": + torch.cuda.current_stream(self._device).synchronize() + + +@dataclass(eq=False) +class _Mapping: + """One VMM handle and its mapping; the fields record how far its setup got.""" + + handle: int + size: int + va: int = 0 + mapped: bool = False + + +def _handle_size(handle_bytes: int, granularity: int) -> int: + """Physical size of a chunk's handle: whole granules, at least one (an empty chunk too).""" + return max(vmm.align_up(handle_bytes, granularity), granularity) + + +def _layer_span(layers: Sequence[int]) -> str: + """How a store chunk's layers are named in messages.""" + if not layers: + return "no layers" + if len(layers) == 1: + return f"layer {layers[0]}" + return f"layers {layers[0]}..{layers[-1]}" + + +def _describe_payload(payload: object) -> str: + if isinstance(payload, (list, tuple)): + kinds = sorted({type(blob).__name__ for blob in payload}) + return f"a {type(payload).__name__} of {len(payload)} ({', '.join(kinds) or 'nothing'})" + return type(payload).__name__ + + +def _set_read_only_access(va: int, size: int, device_id: int) -> None: + """``cuMemSetAccess`` with ``PROT_READ`` for ``device_id`` over ``[va, va + size)``.""" + desc = vmm.get_access_desc(device_id) + desc.flags = vmm.cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READ + vmm.check_cu_result(vmm.cuda.cuMemSetAccess(va, size, [desc], 1)) + + +class FabricShardStore(DenseShardStore): + """Store backed by one FABRIC handle per chunk. + + Each chunk is a ``vmm.create_fabric_handle`` allocation, mapped locally for the P2 writes and + as the copy source of the local row slices, exported like the expert transport's FABRIC + handles and imported and mapped read-only on every peer. + + A chunk's handle is ``StoreChunk.handle_bytes`` rounded up to whole allocation granules, and + at least one granule: a chunk without rows still gets a handle, so every chunk has a mapped + base. Each rank therefore allocates one granule more than ``DenseLayout.store_bytes`` counts + for every chunk without rows and, on a device whose granularity exceeds ``STORE_ALIGN``, up to + a granule more for every other chunk; :meth:`allocate` logs the difference. + + The local mapping is read-write; peer mappings grant ``PROT_READ`` only, so a stray write into + a peer's rows faults instead of corrupting them. Every rank maps each peer handle whole, with + the size it computes from the shared layout. Run it on the dense device's thread with that + device's context current, as setup does. + + Raises: + ValueError: ``device`` is not a CUDA device (and the base class checks). + """ + + def __init__(self, layout: DenseLayout, rank: int, device: torch.device) -> None: + device = torch.device(device) + if device.type != "cuda": + raise ValueError(f"the FABRIC store needs a CUDA device, got {device}") + if device.index is None: + device = torch.device("cuda", torch.cuda.current_device()) + super().__init__(layout, rank, device) + self._device_id: int = device.index + self._local: list[_Mapping] = [] + self._imports: list[_Mapping] = [] + self._views: list[torch.Tensor] = [] + self._exported: list[bytes] | None = None + + def allocate(self) -> None: + """See :meth:`DenseShardStore.allocate`. + + Each chunk's handle is ``max(align_up(handle_bytes, granularity), granularity)`` bytes + (see the class docstring). The zero fill runs on the current stream, which is + synchronized before returning. + + Raises: + RuntimeError: The platform's peer handles are not FABRIC, a CUDA VMM call failed, or + a local view does not alias its mapping. A refused handle names its chunk, its + size and ``TRTLLM_DWDP_DENSE_STORE_CHUNK_MB``, which splits the store into + smaller handles. + """ + self._begin_allocate() + handle_type = vmm.peer_handle_type() + if handle_type != _FABRIC: + raise RuntimeError( + f"the dense DWDP store needs FABRIC peer handles; this platform uses " + f"{handle_type.name}" + ) + granularity = vmm.get_allocation_granularity(self._device_id) + chunks = self._layout.chunks + for chunk in chunks: + size = _handle_size(chunk.handle_bytes, granularity) + try: + handle = vmm.create_fabric_handle(size, self._device_id) + except RuntimeError as exc: + raise RuntimeError( + f"cuMemCreate of the FABRIC handle of store chunk {chunk.index} of " + f"{len(chunks)} ({size} bytes, {_layer_span(chunk.layers)}) on " + f"cuda:{self._device_id} failed: {exc}. {DENSE_STORE_CHUNK_MB_ENV} (now " + f"{self._layout.switches.store_chunk_mb}) splits the store into handles of " + f"whole layers of about that many MiB, e.g. 1024" + ) from exc + record = _Mapping(handle=handle, size=size) + self._local.append(record) + record.va = vmm.reserve_va(size, granularity) + vmm.map_handle(record.va, size, record.handle) + record.mapped = True + vmm.set_access(record.va, size, self._device_id) + view = vmm.tensor_from_ptr(record.va, (size,), torch.uint8, self._device_id) + if view.data_ptr() != record.va: + raise RuntimeError( + f"the view of store chunk {chunk.index} at {record.va:#x} is a copy at " + f"{view.data_ptr():#x}; the store needs views that alias its mapping" + ) + self._views.append(view) + view.zero_() + self._synchronize_current_stream() + self._allocated = True + allocated = sum(record.size for record in self._local) + message = ( + f"{LOG_PREFIX} FABRIC store: rank {self._rank} allocated {allocated} bytes in " + f"{len(self._local)} handle(s) on cuda:{self._device_id}" + ) + if allocated == self._layout.store_bytes: + logger.debug(message) + else: + logger.info( + f"{message}, {allocated - self._layout.store_bytes} more than the layout's " + f"store_bytes ({self._layout.store_bytes}): every handle is whole granules of " + f"{granularity} bytes, at least one" + ) + + def write_local_rows(self) -> None: + """See :meth:`DenseShardStore.write_local_rows`. + + Raises: + RuntimeError: Before a successful :meth:`allocate`, or after :meth:`release`. + DenseSetupError: An anchor's source was dropped or does not match its anchor. + """ + self._require_allocated("write_local_rows") + self._check_sources() + self._write_rows(self._rank, self._views) + + def export(self) -> StorePayload: + """See :meth:`DenseShardStore.export`. + + Returns: + One exported FABRIC handle (``bytes``) per store chunk, in chunk order. + + Raises: + RuntimeError: Before a successful :meth:`allocate`, after :meth:`release`, or the + export failed. + """ + self._require_allocated("export") + if self._exported is None: + self._exported = [ + transport._export_handle(record.handle, _FABRIC) for record in self._local + ] + return list(self._exported) + + def import_peers(self, payloads: Sequence[StorePayload]) -> None: + """See :meth:`DenseShardStore.import_peers`. + + Peers are imported in ascending rank order, chunks in chunk order; this rank's entry + must be what :meth:`export` returned and its local mappings serve as its bases. + + Raises: + ValueError: Not one payload per rank, a payload that is not one exported handle per + chunk, or this rank's entry differs from its export (payloads out of rank order). + RuntimeError: Before :meth:`export`, on a second call, after :meth:`release`, a CUDA + VMM call failed, or a mapping is not ``ROW_ALIGN``-aligned. + """ + self._require_allocated("import_peers") + if self._exported is None: + raise RuntimeError("dense DWDP store: import_peers called before export") + gathered = self._begin_import(payloads) + num_chunks = len(self._local) + for peer, payload in enumerate(gathered): + if ( + not isinstance(payload, (list, tuple)) + or len(payload) != num_chunks + or not all(isinstance(blob, bytes) and blob for blob in payload) + ): + raise ValueError( + f"rank {peer}'s store payload must be {num_chunks} exported FABRIC handle(s) " + f"(non-empty bytes), got {_describe_payload(payload)}" + ) + if list(gathered[self._rank]) != self._exported: + raise ValueError( + f"payload {self._rank} is not this rank's export; the payloads must be indexed by " + f"rank" + ) + granularity = vmm.get_allocation_granularity(self._device_id) + bases = [] + for peer, payload in enumerate(gathered): + if peer == self._rank: + bases.append(tuple(record.va for record in self._local)) + continue + row = [] + for local, blob in zip(self._local, payload): + record = _Mapping(handle=transport._import_handle(blob, _FABRIC), size=local.size) + self._imports.append(record) + record.va = vmm.reserve_va(record.size, granularity) + vmm.map_handle(record.va, record.size, record.handle) + record.mapped = True + _set_read_only_access(record.va, record.size, self._device_id) + row.append(record.va) + bases.append(tuple(row)) + unaligned = [hex(base) for row in bases for base in row if base % ROW_ALIGN] + if unaligned: + raise RuntimeError(f"store mappings {unaligned[:4]} are not {ROW_ALIGN}-aligned") + self._bases = tuple(bases) + logger.debug( + f"{LOG_PREFIX} FABRIC store: rank {self._rank} mapped {len(self._imports)} peer " + f"chunk(s) read-only" + ) + + def chunk_bases(self) -> tuple[tuple[int, ...], ...]: + """See :meth:`DenseShardStore.chunk_bases`. + + Raises: + RuntimeError: Before a successful :meth:`import_peers`, or after :meth:`release`. + """ + return self._checked_bases() + + def release(self) -> None: + """See :meth:`DenseShardStore.release`. + + Order (as the expert transport): synchronize the device if anything was created, then + unmap and free the VA of every peer mapping, release the imported handles, unmap and free + the local mappings, and release the local handles. Each step covers exactly what + :meth:`allocate` and :meth:`import_peers` created, also after one of them raised part + way. Every step runs even if an earlier one fails; the first error is raised at the end + and later ones are logged. + + Raises: + RuntimeError: The first failed step's CUDA error; the store is released as far as + possible. + """ + if self._released: + return + self._released = True + self._views = [] + self._bases = None + imports, self._imports = self._imports, [] + local, self._local = self._local, [] + failures: list[tuple[str, Exception]] = [] + + def attempt(step: str, call: Callable[..., object], *args: int) -> None: + try: + call(*args) + except Exception as exc: # every teardown step runs; the first failure is raised below + failures.append((step, exc)) + + if imports or local: + attempt("synchronize", torch.cuda.synchronize, self._device_id) + for records, kind in ((imports, "peer"), (local, "local")): + for record in records: + if record.mapped: + attempt(f"unmap {kind} VA {record.va:#x}", vmm.unmap_va, record.va, record.size) + if record.va: + attempt(f"free {kind} VA {record.va:#x}", vmm.free_va, record.va, record.size) + for record in records: + attempt(f"release {kind} handle {record.handle}", vmm.release_handle, record.handle) + if failures: + for step, exc in failures[1:]: + logger.error(f"{LOG_PREFIX} FABRIC store release: {step} also failed: {exc!r}") + raise failures[0][1] + logger.debug( + f"{LOG_PREFIX} FABRIC store released: rank {self._rank}, {len(imports)} peer and " + f"{len(local)} local handle(s)" + ) + + +class LoopbackShardStore(DenseShardStore): + """All ``N`` ranks' rows in this process, as ``ROW_ALIGN``-aligned torch buffers. + + Every rank's weights are identical, so :meth:`write_local_rows` writes the rows of every rank + from the local anchors; :meth:`export` returns None and :meth:`import_peers` only checks the + payloads. Used by the CPU setup tests. + + Rank ``p``'s chunk ``c`` is its own zero-filled allocation of ``StoreChunk.handle_bytes`` + bytes (plus ``ROW_ALIGN`` to align its base), so every rank's bases are distinct. + """ + + def __init__(self, layout: DenseLayout, rank: int, device: torch.device) -> None: + super().__init__(layout, rank, device) + self._chunks: list[list[torch.Tensor]] = [] + self._addresses: list[tuple[int, ...]] = [] + + def allocate(self) -> None: + """See :meth:`DenseShardStore.allocate`; allocates the chunks of all ``N`` ranks. + + On CUDA the zero fill runs on the current stream, which is synchronized before + returning. + """ + self._begin_allocate() + for _ in range(self._num_ranks): + chunks = [] + addresses = [] + for chunk in self._layout.chunks: + storage = torch.zeros( + chunk.handle_bytes + ROW_ALIGN, dtype=torch.uint8, device=self._device + ) + pad = (-storage.data_ptr()) % ROW_ALIGN + chunks.append(storage[pad : pad + chunk.handle_bytes]) + addresses.append(storage.data_ptr() + pad) + self._chunks.append(chunks) + self._addresses.append(tuple(addresses)) + self._synchronize_current_stream() + self._allocated = True + + def write_local_rows(self) -> None: + """See :meth:`DenseShardStore.write_local_rows`; writes all ``N`` ranks' rows. + + Raises: + RuntimeError: Before a successful :meth:`allocate`, or after :meth:`release`. + DenseSetupError: An anchor's source was dropped or does not match its anchor. + """ + self._require_allocated("write_local_rows") + self._check_sources() + for row, chunks in enumerate(self._chunks): + self._write_rows(row, chunks) + + def export(self) -> StorePayload: + """See :meth:`DenseShardStore.export`; always None. + + Raises: + RuntimeError: Before a successful :meth:`allocate`, or after :meth:`release`. + """ + self._require_allocated("export") + return None + + def import_peers(self, payloads: Sequence[StorePayload]) -> None: + """See :meth:`DenseShardStore.import_peers`; every payload must be None. + + Raises: + ValueError: Not one payload per rank, or a payload that is not None (a peer with + another kind of store). + RuntimeError: Before a successful :meth:`allocate`, on a second call, or after + :meth:`release`. + """ + gathered = self._begin_import(payloads) + foreign = [peer for peer, payload in enumerate(gathered) if payload is not None] + if foreign: + raise ValueError( + f"the loopback store expects a None payload from every rank; ranks {foreign} " + f"sent store handles" + ) + self._bases = tuple(self._addresses) + + def chunk_bases(self) -> tuple[tuple[int, ...], ...]: + """See :meth:`DenseShardStore.chunk_bases`. + + Raises: + RuntimeError: Before a successful :meth:`import_peers`, or after :meth:`release`. + """ + return self._checked_bases() + + def release(self) -> None: + """See :meth:`DenseShardStore.release`; drops the buffers after a device synchronize.""" + if self._released: + return + self._released = True + self._bases = None + chunks, self._chunks = self._chunks, [] + self._addresses = [] + if chunks and self._device.type == "cuda": + torch.cuda.synchronize(self._device) + del chunks diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/types.py b/tensorrt_llm/_torch/modules/dwdp/dense/types.py new file mode 100644 index 000000000000..a3be8e3db14d --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/types.py @@ -0,0 +1,1134 @@ +# 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. +"""Shared data model of dense DWDP. + +Dense DWDP keeps 1/N of each MoE decoder layer's replicated dense weights on every rank of an +N-rank DWDP group and gathers the layer back into a small staging slot two layers ahead of its use. +Each layer contributes one blob per segment kind: attention ``A(l)`` and MoE-dense ``D(l)``. A +blob is the concatenation of the layer's moved storages ("anchors"), viewed as ``uint8[N, S]``; +rank ``r`` stores row ``r``. + +This module holds what more than one part of the implementation reads or writes: + +* segment kinds and their slot rings (``A``: depth 2, ``D``: depth 3), copy modes and run modes, +* the per-tensor (:class:`TensorBinding`) and per-storage (:class:`AnchorRecord`) records, +* the per-(layer, segment) row layout, the store offsets and chunks (:class:`DenseLayout`), +* slot specifications and slot buffers, +* copy descriptors and per-layer copy plans, +* the CUDA seams the runtime calls, with their torch implementation (:class:`TorchCudaOps`), +* alignment constants, checksums and exceptions. + +Nothing here calls CUDA at import time; only :class:`TorchCudaOps` and +:func:`allocate_slot_ring` with a CUDA device do, when called. +""" + +from __future__ import annotations + +import enum +import hashlib +import weakref +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager +from dataclasses import dataclass, field +from types import MappingProxyType +from typing import NamedTuple, Protocol + +import torch +from torch import nn + +from tensorrt_llm.logger import logger + +from ...multi_stream_utils import do_multi_stream + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +#: Alignment of anchor offsets inside a blob, of row widths ``S``, of the local-row slices of a +#: copy plan and of slot bases. Covers the gather's 16-byte lines, every itemsize of a moved +#: dtype, the CuTe DSL operand alignment and the caching allocator's 512 bytes. +ROW_ALIGN = 4096 + +#: Size granularity of each FABRIC store handle. +STORE_ALIGN = 2 << 20 + +#: Slice size of the batched copy-engine transfer (one ``cudaMemcpyBatchAsync`` per layer). +BATCH_SLICE_BYTES = 2 << 20 + +#: A candidate storage larger than the union of its views' byte ranges by more than this is +#: treated as pooled and stays replicated. +POOLED_SLACK_BYTES = 1 << 20 + +#: Default of ``TRTLLM_DWDP_DENSE_MIN_BYTES``: smaller storages stay replicated. +DEFAULT_MIN_BYTES = 1 << 20 + +#: Tag of a slot that holds no layer. +EMPTY_TAG = -1 + +#: Slot fill byte at commit under ``TRTLLM_DWDP_DENSE_DEBUG >= 1``: NaN in FP8 E4M3, BF16 and FP32. +DEBUG_FILL_BYTE = 0xFF + +#: Odd multiplier of the word-position hash of the weighted checksum. +CHECKSUM_MULTIPLIER = 2654435761 + +#: Dtypes a moved tensor may have. +ALLOWED_DTYPES = frozenset( + { + torch.float8_e4m3fn, + torch.uint8, + torch.bfloat16, + torch.float16, + torch.float32, + torch.int32, + } +) + +#: Prefix of every dense DWDP log line. +LOG_PREFIX = "[DWDP Dense]" + +_U64_MASK = (1 << 64) - 1 + +# --------------------------------------------------------------------------- +# Segment kinds, slot rings, binding kinds, copy modes +# --------------------------------------------------------------------------- + + +class SegmentKind(str, enum.Enum): + """The part of a decoder layer a blob holds. + + The values are the strings the model's ``dense_dwdp_roots(segment)`` and + ``TRTLLM_DWDP_DENSE_SEGMENTS`` use. + """ + + ATTN = "attn" + """Attention weights ``A(l)``, read between H1 ``wait_layer`` and H2 ``attention_done``.""" + + MOE = "moe" + """MoE-dense weights ``D(l)`` (router, shared expert, latent projections), read between H1 and + H3 ``layer_done``.""" + + +#: Canonical segment order: blob order inside a layer, store order, digest order. +SEGMENT_ORDER: tuple[SegmentKind, ...] = (SegmentKind.ATTN, SegmentKind.MOE) + +#: Slot ring depth per segment. ``A(l)`` is consumed at H2(l) and its slot is refilled with +#: ``A(l+2)`` right after, so two slots suffice; ``D(l)`` is consumed only at H3(l) while ``D(l+2)`` +#: is enqueued at H2(l), so it needs three. +RING_DEPTH: Mapping[SegmentKind, int] = MappingProxyType({SegmentKind.ATTN: 2, SegmentKind.MOE: 3}) + + +def ring_slot(segment: SegmentKind, position: int) -> int: + """Slot of a layer in ``segment``'s ring. + + Args: + segment: Segment kind. + position: Index of the layer among the covered layers (ascending). For Kimi-K3 the + covered layers are 1..92, so ``position = l - 1``: layer ``l`` takes slot + ``(l - 1) % 2`` of the attention ring and slot ``(l - 1) % 3`` of the MoE-dense ring. + + Returns: + ``position % RING_DEPTH[segment]``. + + Raises: + ValueError: ``position`` is negative. + """ + if position < 0: + raise ValueError(f"layer position must be non-negative, got {position}") + return position % RING_DEPTH[segment] + + +class BindingKind(str, enum.Enum): + """How a module holds a tensor; decides how repointing installs the slot view.""" + + PARAMETER = "param" + """An ``nn.Parameter`` in ``module._parameters``; repointed with ``param.data = view``.""" + + BUFFER = "buffer" + """A registered buffer; repointed with ``module._buffers[name] = view``.""" + + ATTRIBUTE = "attr" + """A plain tensor attribute; repointed with ``object.__setattr__(module, name, view)``.""" + + SEQUENCE_ITEM = "tuple" + """Element ``index`` of a tuple or list attribute whose elements are all tensors; repointed by + rebuilding the container with the view in place and setting it again.""" + + +class CopyMode(str, enum.Enum): + """Transfer used to gather a layer into its slots. + + With expert DWDP it follows the expert DWDP transfer; standalone, ``TRTLLM_DWDP_DENSE_COPY`` + picks it. + """ + + GATHER = "gather" + """One ``trtllm::dwdp_fabric_gather`` launch per layer (expert ``use_gather``).""" + + BATCHED = "batched" + """One ``cudaMemcpyBatchAsync`` of ``BATCH_SLICE_BYTES`` slices per layer + (expert ``use_batched_prefetch``).""" + + PER_SLICE = "per_slice" + """One ``Tensor.copy_`` per copy descriptor.""" + + +class DenseMode(str, enum.Enum): + """How dense DWDP runs; ``gate.validate_dense_gate`` decides it from the launch.""" + + ATTACHED = "attached" + """With expert DWDP (``dwdp_config``): the DWDP group, the expert DWDP copy stream and its + transfer.""" + + STANDALONE = "standalone" + """On an attention-DP instance with ``tp_size > 1``, without expert DWDP: the tp group, the + dense manager's own copy stream and ``TRTLLM_DWDP_DENSE_COPY``.""" + + +#: Key of one blob: ``(layer index, segment kind)``. +SegmentKey = tuple[int, SegmentKind] + +# --------------------------------------------------------------------------- +# Arithmetic +# --------------------------------------------------------------------------- + + +def align_up(value: int, alignment: int) -> int: + """Round a non-negative ``value`` up to a multiple of the power-of-two ``alignment``. + + Raises: + ValueError: ``alignment`` is not a positive power of two, or ``value`` is negative. + """ + if alignment <= 0 or alignment & (alignment - 1): + raise ValueError(f"alignment must be a positive power of 2, got {alignment}") + if value < 0: + raise ValueError(f"cannot align a negative value ({value})") + return (value + alignment - 1) // alignment * alignment + + +def ceil_div(numerator: int, denominator: int) -> int: + """``ceil(numerator / denominator)`` for a non-negative numerator and positive denominator.""" + if denominator <= 0 or numerator < 0: + raise ValueError( + f"ceil_div needs numerator >= 0 and denominator > 0, got {numerator}, {denominator}" + ) + return -(-numerator // denominator) + + +def row_bytes_for(blob_bytes: int, num_ranks: int) -> int: + """Row width of a blob: ``S = align_up(ceil(B / N), ROW_ALIGN)``.""" + return align_up(ceil_div(blob_bytes, num_ranks), ROW_ALIGN) + + +# --------------------------------------------------------------------------- +# Checksums +# --------------------------------------------------------------------------- + + +class Checksum(NamedTuple): + """Two wrap-around sums over a byte range read as int32 words. + + With ``w_i`` the i-th word read as a signed little-endian int32 and ``h_i = (i * + CHECKSUM_MULTIPLIER) mod 2**32``: ``word_sum = sum(w_i)`` and ``weighted_sum = sum(w_i * h_i)``, + both modulo ``2**64`` and stored as unsigned ints. ``i`` counts words from an origin the caller + names (the blob start for per-segment sums, the tensor start for per-tensor sums), so sums of + disjoint ranges combine with :meth:`merge` and do not depend on the reduction order. A range + that does not end on a word boundary is zero-padded to one. + """ + + word_sum: int + weighted_sum: int + + def merge(self, other: Checksum) -> Checksum: + """Checksum of the union of two disjoint ranges with the same word origin.""" + return Checksum( + (self.word_sum + other.word_sum) & _U64_MASK, + (self.weighted_sum + other.weighted_sum) & _U64_MASK, + ) + + +#: Checksum of an empty (or all-zero) range. +ZERO_CHECKSUM = Checksum(0, 0) + +# --------------------------------------------------------------------------- +# Collection records +# --------------------------------------------------------------------------- + + +@dataclass(eq=False) +class TensorBinding: + """One reference through which a decoder layer reaches an anchor storage. + + Every reference under the layer's roots for the anchor's segment is recorded, including names + in the model's exclude table, so repointing preserves every alias. A reference is unique by + ``(owner, kind, name, index)``. + + Attributes: + module_path: Path of the owning module relative to the decoder layer, as produced by + ``root.named_modules(prefix=root_name)``; e.g. ``"linear_attn.qkvg_proj"``. + kind: How the owner holds the tensor. + name: ``_parameters`` / ``_buffers`` key or attribute name on the owner. + index: Element index for ``SEQUENCE_ITEM``; None for the other kinds. + dtype: Dtype of the view the owner exposes. + size: Size of that view. + stride: Stride of that view, in elements. + byte_offset: ``storage_offset() * element_size()``: where the view starts inside the anchor + storage. + owner_ref: Weak reference to the owning module. + original: The tensor the owner exposed at collection (for ``PARAMETER``, ``param.data``). + Held until commit for rollback and for the P2 to P5 checks; None after commit. + checksum: Per-tensor checksum of ``original`` (made contiguous), from P1, for the P5 + self-check through the module attribute. + """ + + module_path: str + kind: BindingKind + name: str + index: int | None + dtype: torch.dtype + size: tuple[int, ...] + stride: tuple[int, ...] + byte_offset: int + owner_ref: weakref.ReferenceType[nn.Module] + original: torch.Tensor | None = None + checksum: Checksum | None = None + + @property + def qualified_name(self) -> str: + """``module_path.name`` plus ``[index]`` for sequence items; unique within a layer.""" + suffix = "" if self.index is None else f"[{self.index}]" + return f"{self.module_path}.{self.name}{suffix}" + + @property + def span_bytes(self) -> int: + """Bytes from the view start to one past its last element; 0 for an empty view.""" + if any(dim == 0 for dim in self.size): + return 0 + span = 1 + sum((dim - 1) * step for dim, step in zip(self.size, self.stride)) + return span * self.dtype.itemsize + + def owner(self) -> nn.Module: + """The owning module. + + Raises: + DenseDWDPError: The owner was garbage collected. + """ + module = self.owner_ref() + if module is None: + raise DenseDWDPError(f"the module owning {self.qualified_name} was garbage collected") + return module + + def digest_tuple(self) -> tuple[object, ...]: + """``(path, kind, name, index, dtype, size, stride, byte_offset)`` in primitive types.""" + return ( + self.module_path, + self.kind.value, + self.name, + self.index, + str(self.dtype), + tuple(self.size), + tuple(self.stride), + self.byte_offset, + ) + + +@dataclass(eq=False) +class AnchorRecord: + """One moved storage and every binding of it. + + Attributes: + layer: Decoder layer index. + segment: Segment kind; an anchor belongs to exactly one ``(layer, segment)``. + index: Position among the segment's anchors, in first-encounter order of the walk. + nbytes: ``untyped_storage().nbytes()``; the whole storage moves, including bytes no view + covers (e.g. the pad rows of ``_bfa_proj_weight``). + blob_offset: ``O_X``, the anchor's byte offset in its blob; a multiple of ``ROW_ALIGN``. + bindings: Every reference to the storage under the segment's roots, in walk order. + source: Flat ``uint8 [nbytes]`` view of the original storage; the bytes P2 writes into the + store rows and P4 compares against. None after commit. + """ + + layer: int + segment: SegmentKind + index: int + nbytes: int + blob_offset: int + bindings: tuple[TensorBinding, ...] + source: torch.Tensor | None = None + + @property + def blob_end(self) -> int: + """One past the anchor's last blob byte.""" + return self.blob_offset + self.nbytes + + @property + def path(self) -> str: + """``layers..``: how logs and leak reports name the anchor.""" + first = self.bindings[0].qualified_name if self.bindings else "" + return f"layers.{self.layer}.{first}" + + def drop_originals(self) -> None: + """Drop this record's strong references to the original storage (``source``, originals).""" + self.source = None + for binding in self.bindings: + binding.original = None + + def digest_tuple(self) -> tuple[object, ...]: + """``(idx, nbytes, O_X, bindings)`` in primitive types.""" + return ( + self.index, + self.nbytes, + self.blob_offset, + tuple(binding.digest_tuple() for binding in self.bindings), + ) + + +class RowPiece(NamedTuple): + """One contiguous run of a rank's row that comes from one anchor. + + The rest of the row is zero. + """ + + anchor_index: int + """Index into ``SegmentLayout.anchors``.""" + + anchor_offset: int + """Byte offset inside the anchor storage (and ``AnchorRecord.source``).""" + + row_offset: int + """Byte offset inside the row, ``0 <= row_offset < S``.""" + + nbytes: int + """Length of the run.""" + + +# --------------------------------------------------------------------------- +# Layout +# --------------------------------------------------------------------------- + + +@dataclass(eq=False) +class SegmentLayout: + """Blob, rows and store placement of one ``(layer, segment)``. + + Attributes: + layer: Decoder layer index. + segment: Segment kind. + blob_bytes: ``B``: end of the last anchor; 0 if nothing of this segment moves. + row_bytes: ``S = align_up(ceil(B / N), ROW_ALIGN)``; row ``p`` is blob bytes + ``[p * S, (p + 1) * S)``, zero-padded past ``B``. + store_offset: ``off``: offset of the rows in the logical store, the running sum of ``S`` + over ``(layer ascending, attn before moe)``; identical on every rank. + chunk: Index of the store chunk holding the rows. + chunk_offset: ``store_offset`` relative to the start of that chunk; rank ``p``'s row is at + ``chunk_bases[p][chunk] + chunk_offset``. + anchors: The anchors in blob order; empty for layouts built from blob sizes alone. + checksum: Per-segment content checksum over the blob (word origin: blob start), from P1. + """ + + layer: int + segment: SegmentKind + blob_bytes: int + row_bytes: int + store_offset: int + chunk: int + chunk_offset: int + anchors: tuple[AnchorRecord, ...] = () + checksum: Checksum | None = None + + @property + def key(self) -> SegmentKey: + """``(layer, segment)``.""" + return (self.layer, self.segment) + + @property + def moved_bytes(self) -> int: + """Sum of the anchors' storage sizes.""" + return sum(anchor.nbytes for anchor in self.anchors) + + def row_range(self, rank: int) -> tuple[int, int]: + """Blob byte range ``[rank * S, (rank + 1) * S)`` of ``rank``'s row.""" + return rank * self.row_bytes, (rank + 1) * self.row_bytes + + def digest_tuple(self) -> tuple[object, ...]: + """``(l, seg, B, S, off, chunk, anchors)`` in primitive types.""" + return ( + self.layer, + self.segment.value, + self.blob_bytes, + self.row_bytes, + self.store_offset, + self.chunk, + tuple(anchor.digest_tuple() for anchor in self.anchors), + ) + + +@dataclass(frozen=True) +class StoreChunk: + """A run of whole layers whose rows share one store handle. + + ``TRTLLM_DWDP_DENSE_STORE_CHUNK_MB=0`` gives a single chunk. + + Attributes: + index: Chunk index, ascending with the layers. + layers: The layers whose rows (all segments) live in this chunk. + offset: Logical store offset of the chunk's first row. + nbytes: Sum of the chunk's row widths. + """ + + index: int + layers: tuple[int, ...] + offset: int + nbytes: int + + @property + def handle_bytes(self) -> int: + """Size of the chunk's handle on every rank: ``align_up(nbytes, STORE_ALIGN)``.""" + return align_up(self.nbytes, STORE_ALIGN) + + +@dataclass(frozen=True) +class SlotSpec: + """Shape of one segment's slot ring. + + Attributes: + segment: Segment kind. + depth: Ring depth, ``RING_DEPTH[segment]``. + row_bytes: ``S_max``: the largest row width of the segment over the covered layers. + num_ranks: ``N``. + """ + + segment: SegmentKind + depth: int + row_bytes: int + num_ranks: int + + @property + def valid_bytes(self) -> int: + """``N * S_max``: the largest blob a slot holds, starting at its aligned base.""" + return self.num_ranks * self.row_bytes + + @property + def alloc_bytes(self) -> int: + """Allocation size of one slot: ``valid_bytes`` plus ``ROW_ALIGN`` for base alignment.""" + return self.valid_bytes + ROW_ALIGN + + +@dataclass(frozen=True) +class DigestSwitches: + """The switches the layout depends on; part of the cross-rank digest. + + Attributes: + segments: Enabled segment kinds in canonical order; always starts with ``ATTN``. + min_bytes: ``TRTLLM_DWDP_DENSE_MIN_BYTES``. + store_chunk_mb: ``TRTLLM_DWDP_DENSE_STORE_CHUNK_MB``. + num_ranks: ``N``, the DWDP group size; at least 2. + row_align: ``ROW_ALIGN``. + store_align: ``STORE_ALIGN``. + ring_depths: Ring depths of ``SEGMENT_ORDER``. + """ + + segments: tuple[SegmentKind, ...] + min_bytes: int + store_chunk_mb: int + num_ranks: int + row_align: int = ROW_ALIGN + store_align: int = STORE_ALIGN + ring_depths: tuple[int, ...] = tuple(RING_DEPTH[kind] for kind in SEGMENT_ORDER) + + def __post_init__(self) -> None: + canonical = tuple(kind for kind in SEGMENT_ORDER if kind in self.segments) + if not self.segments or tuple(self.segments) != canonical: + raise ValueError( + f"segments must be a non-empty subset of {[k.value for k in SEGMENT_ORDER]} " + f"in that order, got {self.segments}" + ) + if self.segments[0] is not SegmentKind.ATTN: + raise ValueError("dense DWDP always includes the attention segment") + if self.num_ranks < 2: + raise ValueError(f"dense DWDP needs at least 2 ranks, got {self.num_ranks}") + if self.min_bytes < 0 or self.store_chunk_mb < 0: + raise ValueError( + f"min_bytes and store_chunk_mb must be non-negative, got " + f"{self.min_bytes} and {self.store_chunk_mb}" + ) + + def as_tuple(self) -> tuple[object, ...]: + """The switch tuple in primitive types, in field order.""" + return ( + tuple(kind.value for kind in self.segments), + self.min_bytes, + self.store_chunk_mb, + self.num_ranks, + self.row_align, + self.store_align, + tuple(self.ring_depths), + ) + + +@dataclass(eq=False) +class DenseLayout: + """Everything every rank must agree on, plus this rank's collection records. + + Built by :func:`build_dense_layout`; the collector (``layout.collect_dense_layout``) adds the + anchors, tests can build it from blob sizes alone. The plan, store, slots and repointing all + read it; only the checksum fields and the anchors' originals change after it is built. + + Attributes: + switches: Layout switches (segments, threshold, chunking, N, alignments, ring depths). + layers: Covered decoder layers, ascending. Every covered layer has an entry for every + enabled segment. + segments: ``(layer, segment) -> SegmentLayout`` in store order. + chunks: Store chunks in store order. + slot_specs: Ring shape per enabled segment. + digest: sha256 hex digest of :meth:`digest_input`; equal on two ranks iff their layouts + are. + warnings: Non-fatal findings of the collector (storages left in place), in walk order. + """ + + switches: DigestSwitches + layers: tuple[int, ...] + segments: dict[SegmentKey, SegmentLayout] + chunks: tuple[StoreChunk, ...] + slot_specs: dict[SegmentKind, SlotSpec] + digest: str = "" + warnings: list[str] = field(default_factory=list) + + def __post_init__(self) -> None: + self._positions = {layer: position for position, layer in enumerate(self.layers)} + + @property + def num_ranks(self) -> int: + """``N``.""" + return self.switches.num_ranks + + @property + def segment_kinds(self) -> tuple[SegmentKind, ...]: + """Enabled segment kinds in canonical order.""" + return self.switches.segments + + @property + def moved_bytes(self) -> int: + """Bytes of replicated storage this layout moves.""" + return sum(segment.moved_bytes for segment in self.segments.values()) + + @property + def store_bytes(self) -> int: + """Bytes of store handles each rank allocates.""" + return sum(chunk.handle_bytes for chunk in self.chunks) + + @property + def slot_bytes(self) -> int: + """Bytes of slot allocations each rank makes.""" + return sum(spec.depth * spec.alloc_bytes for spec in self.slot_specs.values()) + + def covers(self, layer: int) -> bool: + """Whether ``layer`` is a covered layer.""" + return layer in self._positions + + def position(self, layer: int) -> int: + """Index of ``layer`` among the covered layers. + + Raises: + KeyError: ``layer`` is not covered. + """ + try: + return self._positions[layer] + except KeyError: + raise KeyError(f"layer {layer} is not covered by dense DWDP") from None + + def slot_index(self, segment: SegmentKind, layer: int) -> int: + """Ring slot holding ``layer``'s ``segment`` blob.""" + return ring_slot(segment, self.position(layer)) + + def segment(self, layer: int, segment: SegmentKind) -> SegmentLayout: + """The ``(layer, segment)`` entry. + + Raises: + KeyError: No such entry. + """ + return self.segments[(layer, segment)] + + def layer_segments(self, layer: int) -> tuple[SegmentLayout, ...]: + """``layer``'s entries in canonical segment order.""" + return tuple(self.segments[(layer, kind)] for kind in self.segment_kinds) + + def anchors(self) -> Iterator[AnchorRecord]: + """Every anchor, in store order.""" + for segment in self.segments.values(): + yield from segment.anchors + + def checksums(self) -> tuple[Checksum, ...]: + """Per-segment content checksums in store order: the #1 payload. + + Raises: + DenseLayoutError: Checksums were not computed yet. + """ + missing = [key for key, segment in self.segments.items() if segment.checksum is None] + if missing: + raise DenseLayoutError(f"segment checksums not computed for {missing[:4]}...") + return tuple(segment.checksum for segment in self.segments.values()) + + def digest_input(self) -> tuple[object, ...]: + """What the digest covers: switches, every segment with its anchors and bindings, chunks.""" + return ( + self.switches.as_tuple(), + tuple(segment.digest_tuple() for segment in self.segments.values()), + tuple((chunk.index, chunk.layers, chunk.offset, chunk.nbytes) for chunk in self.chunks), + ) + + +def compute_layout_digest(layout: DenseLayout) -> str: + """The sha256 hex digest of ``repr(layout.digest_input())``; stable, as it holds primitives.""" + return hashlib.sha256(repr(layout.digest_input()).encode("utf-8")).hexdigest() + + +def build_dense_layout( + blob_bytes: Mapping[SegmentKey, int], + switches: DigestSwitches, + *, + anchors: Mapping[SegmentKey, Sequence[AnchorRecord]] | None = None, + warnings: Sequence[str] = (), +) -> DenseLayout: + """Compute rows, store offsets, chunks, slot specs and the digest. + + ``S`` and ``off`` are as :class:`SegmentLayout` defines them. Chunks hold whole layers: with + ``store_chunk_mb == 0`` there is one chunk; otherwise layers are added in ascending order and + a new chunk starts when the next layer's rows would push a non-empty chunk past + ``store_chunk_mb`` MiB (a single larger layer gets a chunk of its own). + + Args: + blob_bytes: ``B`` for every ``(covered layer, enabled segment)``; the covered layers are + the layers that appear in the keys. + switches: Layout switches. + anchors: The collector's anchors per key, in blob order. When given, every key's ``B`` + must equal its last anchor's end (0 without anchors) and the offsets must be + ``ROW_ALIGN``-aligned, ascending and non-overlapping. Omit it to build a layout from + blob sizes alone (copy-plan and runtime tests). + warnings: Collector findings to carry on the layout. + + Returns: + The layout, with ``digest`` set and checksums unset. + + Raises: + ValueError: Keys do not cover every enabled segment of every covered layer, a size is + negative, or the anchors are inconsistent with ``blob_bytes``. + """ + kinds = switches.segments + sizes = {(layer, SegmentKind(kind)): size for (layer, kind), size in blob_bytes.items()} + layers = tuple(sorted({layer for layer, _ in sizes})) + if not layers: + raise ValueError("a dense layout needs at least one covered layer") + expected = {(layer, kind) for layer in layers for kind in kinds} + if set(sizes) != expected: + raise ValueError( + f"blob sizes must cover every enabled segment of every covered layer; " + f"missing {sorted(expected - set(sizes))[:4]}, " + f"unexpected {sorted(set(sizes) - expected)[:4]}" + ) + anchor_map = { + (layer, SegmentKind(kind)): tuple(value) for (layer, kind), value in (anchors or {}).items() + } + unknown = set(anchor_map) - expected + if unknown: + raise ValueError(f"anchors for keys outside the layout: {sorted(unknown)[:4]}") + + num_ranks = switches.num_ranks + rows: list[tuple[SegmentKey, int, int, int, tuple[AnchorRecord, ...]]] = [] + layer_bytes: dict[int, int] = {} + offset = 0 + for layer in layers: + for kind in kinds: + key = (layer, kind) + size = sizes[key] + if size < 0: + raise ValueError(f"blob size of {key} is negative ({size})") + segment_anchors = anchor_map.get(key, ()) + if anchors is not None: + _check_anchor_placement(key, size, segment_anchors) + row = row_bytes_for(size, num_ranks) + rows.append((key, size, row, offset, segment_anchors)) + layer_bytes[layer] = layer_bytes.get(layer, 0) + row + offset += row + + chunk_layers = _assign_chunks(layers, layer_bytes, switches.store_chunk_mb << 20) + chunk_of_layer = {layer: index for index, group in enumerate(chunk_layers) for layer in group} + first_offset = {} + for key, _, _, row_offset, _ in rows: + first_offset.setdefault(chunk_of_layer[key[0]], row_offset) + chunks = tuple( + StoreChunk( + index=index, + layers=tuple(group), + offset=first_offset[index], + nbytes=sum(layer_bytes[layer] for layer in group), + ) + for index, group in enumerate(chunk_layers) + ) + + segments: dict[SegmentKey, SegmentLayout] = {} + for key, size, row, row_offset, segment_anchors in rows: + chunk = chunk_of_layer[key[0]] + segments[key] = SegmentLayout( + layer=key[0], + segment=key[1], + blob_bytes=size, + row_bytes=row, + store_offset=row_offset, + chunk=chunk, + chunk_offset=row_offset - chunks[chunk].offset, + anchors=segment_anchors, + ) + + slot_specs = { + kind: SlotSpec( + segment=kind, + depth=RING_DEPTH[kind], + row_bytes=max(segments[(layer, kind)].row_bytes for layer in layers), + num_ranks=num_ranks, + ) + for kind in kinds + } + layout = DenseLayout( + switches=switches, + layers=layers, + segments=segments, + chunks=chunks, + slot_specs=slot_specs, + warnings=list(warnings), + ) + layout.digest = compute_layout_digest(layout) + return layout + + +def _check_anchor_placement( + key: SegmentKey, blob_bytes: int, anchors: Sequence[AnchorRecord] +) -> None: + end = 0 + for position, anchor in enumerate(anchors): + if (anchor.layer, anchor.segment) != key or anchor.index != position: + raise ValueError( + f"anchor {anchor.path} is recorded as ({anchor.layer}, {anchor.segment.value}) " + f"#{anchor.index} but placed in {key} at #{position}" + ) + if anchor.blob_offset % ROW_ALIGN or anchor.blob_offset < end or anchor.nbytes <= 0: + raise ValueError( + f"anchor {anchor.path} has offset {anchor.blob_offset} and {anchor.nbytes} bytes; " + f"offsets must be {ROW_ALIGN}-aligned, ascending, non-overlapping and sizes positive" + ) + end = anchor.blob_end + if end != blob_bytes: + raise ValueError(f"blob size of {key} is {blob_bytes} but its anchors end at {end}") + + +def _assign_chunks( + layers: Sequence[int], layer_bytes: Mapping[int, int], limit: int +) -> list[list[int]]: + groups: list[list[int]] = [[]] + used = 0 + for layer in layers: + size = layer_bytes[layer] + if limit > 0 and used > 0 and used + size > limit: + groups.append([]) + used = 0 + groups[-1].append(layer) + used += size + return groups + + +# --------------------------------------------------------------------------- +# Slots +# --------------------------------------------------------------------------- + + +@dataclass(eq=False) +class SlotBuffer: + """One staging slot of a segment ring. + + Blob byte ``b`` of the layer the slot currently holds lives at ``storage[pad + b]``; rows use + that layer's own ``S``. Module views made by repointing are ``as_strided`` views of + ``storage``, so the views keep the allocation alive. + + Attributes: + segment: Ring the slot belongs to. + index: Position in the ring, ``0 .. RING_DEPTH[segment] - 1``. + storage: Flat ``uint8 [SlotSpec.alloc_bytes]`` tensor with storage offset 0. + pad: ``(-storage.data_ptr()) % ROW_ALIGN``. + """ + + segment: SegmentKind + index: int + storage: torch.Tensor + pad: int + + @property + def base(self) -> int: + """``ROW_ALIGN``-aligned device address of blob byte 0.""" + return self.storage.data_ptr() + self.pad + + @property + def capacity(self) -> int: + """Bytes available from ``base`` to the end of the allocation.""" + return self.storage.numel() - self.pad + + def blob_view(self, nbytes: int) -> torch.Tensor: + """``uint8 [nbytes]`` view of the first ``nbytes`` blob bytes. + + Raises: + ValueError: ``nbytes`` exceeds ``capacity``. + """ + if nbytes < 0 or nbytes > self.capacity: + raise ValueError( + f"slot {self.segment.value}[{self.index}] holds {self.capacity} bytes, asked {nbytes}" + ) + return self.storage[self.pad : self.pad + nbytes] + + +def allocate_slot_ring(spec: SlotSpec, device: torch.device | str) -> tuple[SlotBuffer, ...]: + """Allocate ``spec.depth`` zero-filled slots of ``spec.alloc_bytes`` bytes on ``device``. + + Slots are ordinary torch allocations (not cuMem), so the module views repointed into them can + never dangle. CPU devices work for unit tests. + """ + slots = [] + for index in range(spec.depth): + storage = torch.zeros(spec.alloc_bytes, dtype=torch.uint8, device=device) + slots.append( + SlotBuffer( + segment=spec.segment, + index=index, + storage=storage, + pad=(-storage.data_ptr()) % ROW_ALIGN, + ) + ) + return tuple(slots) + + +# --------------------------------------------------------------------------- +# Copy plans +# --------------------------------------------------------------------------- + + +class CopyDescriptor(NamedTuple): + """One copy of ``nbytes`` from device address ``src`` to ``dst``. + + A plain tuple, so lists of descriptors are valid ``fabric_gather.build_plan`` groups. + """ + + dst: int + src: int + nbytes: int + + +@dataclass(frozen=True) +class LayerCopyPlan: + """Everything one gather of ``L(layer)`` copies on one rank. + + Row ``k`` serves peer ``peers[k]`` (ascending, this rank excluded) and holds, in this order and + skipping empty ones: the ``A`` row of that peer, the ``D`` row of that peer, slice ``k`` of + this rank's ``A`` row and slice ``k`` of this rank's ``D`` row. Together the rows write every + byte of ``[base_A, base_A + N * S_A)`` and ``[base_D, base_D + N * S_D)`` exactly once. + + Attributes: + layer: Target layer ``n``. + rank: This rank ``r``. + peers: ``p_0 < ... < p_{N-2}``. + rows: One descriptor tuple per peer, aligned with ``peers``. + """ + + layer: int + rank: int + peers: tuple[int, ...] + rows: tuple[tuple[CopyDescriptor, ...], ...] + + @property + def descriptors(self) -> tuple[CopyDescriptor, ...]: + """All descriptors in row order.""" + return tuple(descriptor for row in self.rows for descriptor in row) + + @property + def nbytes(self) -> int: + """Bytes the plan copies.""" + return sum(descriptor.nbytes for descriptor in self.descriptors) + + @property + def is_empty(self) -> bool: + """True if the plan copies nothing (every enabled segment of the layer has ``B == 0``).""" + return self.nbytes == 0 + + +#: Enqueues the copy of ``L(layer)`` (all enabled segments) on the current CUDA stream. The +#: runtime makes the copy stream current around the call. +LayerLauncher = Callable[[int], None] + +# --------------------------------------------------------------------------- +# CUDA seams of the runtime +# --------------------------------------------------------------------------- + + +class EventLike(Protocol): + """The part of ``torch.cuda.Event`` the runtime uses.""" + + def record(self, stream: StreamLike | None = None) -> None: + """Record the event on ``stream`` (the current stream if None).""" + ... + + +class StreamLike(Protocol): + """The part of ``torch.cuda.Stream`` the runtime uses.""" + + def wait_event(self, event: EventLike) -> None: + """Make future work on this stream wait for the event's most recent record.""" + ... + + +class DeviceOps(Protocol): + """Every CUDA call the dense runtime makes; unit tests pass a recording fake.""" + + def new_event(self) -> EventLike: + """A new non-timing event.""" + ... + + def current_stream(self) -> StreamLike: + """The current stream of the dense device.""" + ... + + def stream(self, stream: StreamLike) -> AbstractContextManager[object]: + """Context that makes ``stream`` current (``torch.cuda.stream``).""" + ... + + def synchronize(self) -> None: + """Wait for all work on the dense device.""" + ... + + def is_current_stream_capturing(self) -> bool: + """Whether the current stream is capturing a CUDA graph.""" + ... + + def do_multi_stream(self) -> bool: + """Whether modules may fork work onto auxiliary streams (``multi_stream_utils``).""" + ... + + def nvtx_range(self, message: str) -> AbstractContextManager[object]: + """Context that wraps an NVTX range named ``message``.""" + ... + + +class TorchCudaOps: + """:class:`DeviceOps` backed by ``torch.cuda``. + + Args: + device: The dense device; None means the current device. + """ + + def __init__(self, device: torch.device | int | None = None) -> None: + self._device = device + + def new_event(self) -> torch.cuda.Event: + """A new non-timing ``torch.cuda.Event``.""" + return torch.cuda.Event() + + def current_stream(self) -> torch.cuda.Stream: + """``torch.cuda.current_stream(device)``.""" + return torch.cuda.current_stream(self._device) + + def stream(self, stream: StreamLike) -> AbstractContextManager[object]: + """``torch.cuda.stream(stream)``.""" + return torch.cuda.stream(stream) + + def synchronize(self) -> None: + """``torch.cuda.synchronize(device)``.""" + torch.cuda.synchronize(self._device) + + def is_current_stream_capturing(self) -> bool: + """``torch.cuda.is_current_stream_capturing()``.""" + return torch.cuda.is_current_stream_capturing() + + def do_multi_stream(self) -> bool: + """``multi_stream_utils.do_multi_stream()``.""" + return do_multi_stream() + + def nvtx_range(self, message: str) -> AbstractContextManager[object]: + """``torch.cuda.nvtx.range(message)``.""" + return torch.cuda.nvtx.range(message) + + +# --------------------------------------------------------------------------- +# Model seam +# --------------------------------------------------------------------------- + + +class DenseRootsProvider(Protocol): + """A decoder layer that names its dense roots. + + ``modules/dwdp`` does not import model code; a layer opts in by implementing this method. + """ + + def dense_dwdp_roots(self, segment: str) -> list[tuple[str, nn.Module]]: + """The ``(path, module)`` roots of ``segment``. + + Args: + segment: ``"attn"`` or ``"moe"``. + + Returns: + The roots in walk order; [] if the layer stays replicated. + """ + ... + + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +class DenseDWDPError(RuntimeError): + """Base class of dense DWDP errors.""" + + +class DenseLayoutError(DenseDWDPError): + """Collection, its checks or the layout rejected the model.""" + + +class DenseRepointError(DenseDWDPError): + """Repointing, the self-check, rollback or the release check failed.""" + + +class DenseProtocolError(DenseDWDPError): + """A runtime hook ran out of protocol.""" + + +class DenseSetupError(DenseDWDPError): + """A setup phase failed, or the group decided to abandon dense DWDP.""" + + +def warn_or_raise( + message: str, *, strict: bool, error: type[DenseDWDPError] = DenseDWDPError +) -> str: + """Report a finding that ``TRTLLM_DWDP_DENSE_STRICT`` turns into an error. + + Args: + message: What happened, without the log prefix. + strict: Raise instead of warning. + error: Exception type to raise under ``strict``. + + Returns: + ``message``, after logging it as a warning. + + Raises: + DenseDWDPError: ``error(message)`` when ``strict``. + """ + if strict: + raise error(message) + logger.warning(f"{LOG_PREFIX} {message}") + return message diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/watchdog.py b/tensorrt_llm/_torch/modules/dwdp/dense/watchdog.py new file mode 100644 index 000000000000..ed2078a0b9a9 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/watchdog.py @@ -0,0 +1,200 @@ +# 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. +"""Slot checksum watchdog of dense DWDP: ``TRTLLM_DWDP_DENSE_DEBUG=2``. + +At H1 of every covered layer, once compute waits for the layer's copies, ``DenseDWDP.wait_layer`` +has the watchdog check each slot the layer reads: on the compute stream, the checksum of the slot's +blob is compared with the P1 content checksum of the layer's segment, and a mismatch adds one to +that ``(layer, segment)``'s counter on the device. Nothing reaches the host there. +``DenseDWDP.layer_done`` reads the counters once per forward, at the last covered layer, and +``DenseDWDP.release`` once more; each read synchronizes the host with the current stream. Every +check reads its whole blob once more, so the watchdog is a tool for validation runs. + +A correct gather leaves exactly the zero-padded blob in a slot: the row store is zero outside the +anchors, and Phase V proved that the gathered anchors equal the originals the P1 checksums were +computed from. A mismatch therefore means the slot did not hold its layer's bytes when compute +reached the layer: a missing or misbound stream wait, a copy that wrote other bytes, or a row +store that changed after setup. +""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence + +import torch + +from .layout import checksum_bits, device_checksum +from .types import ( + Checksum, + DenseLayout, + DenseProtocolError, + SegmentKey, + SegmentKind, + SlotBuffer, + align_up, +) + +#: Bytes per checksum word: a blob is checked up to its end rounded up to a word, which the zero +#: padding of its last row covers. +_WORD_BYTES = 4 + +#: Most ``(layer, segment)`` pairs one report names. +_MAX_LISTED = 8 + +#: ``(layer, segment, mismatches)``: a pair whose counter grew since the previous read. +Finding = tuple[int, SegmentKind, int] + + +class SlotWatchdog: + """Device-side checks of the slots at H1 against the layers' P1 content checksums. + + Args: + layers: Covered layers, ascending. + segments: Enabled segment kinds in canonical order. + expected: ``(layer, segment) -> (blob_bytes, checksum)`` for every covered layer and + enabled segment: the blob size ``B`` and the P1 content checksum of the blob, whose + word origin is the blob start. + device: Device of the slots. + + Raises: + ValueError: ``expected`` does not hold exactly the covered ``(layer, segment)`` pairs, or + a blob size is negative. + """ + + def __init__( + self, + layers: Sequence[int], + segments: Sequence[SegmentKind], + expected: Mapping[SegmentKey, tuple[int, Checksum]], + device: torch.device | str, + ) -> None: + self._layers = tuple(layers) + self._segments = tuple(SegmentKind(kind) for kind in segments) + keys = tuple((layer, kind) for layer in self._layers for kind in self._segments) + given = {(layer, SegmentKind(kind)): value for (layer, kind), value in expected.items()} + if set(given) != set(keys): + missing = sorted(f"{layer}/{kind.value}" for layer, kind in set(keys) - set(given)) + extra = sorted(f"{layer}/{kind.value}" for layer, kind in set(given) - set(keys)) + raise ValueError( + f"the slot watchdog needs a checksum for every covered (layer, segment): missing " + f"{missing[:_MAX_LISTED]}, not covered {extra[:_MAX_LISTED]}" + ) + check_bytes = [] + bits = [] + for layer, kind in keys: + blob_bytes, checksum = given[(layer, kind)] + if blob_bytes < 0: + raise ValueError(f"layer {layer} {kind.value}: negative blob size {blob_bytes}") + check_bytes.append(align_up(blob_bytes, _WORD_BYTES)) + bits.append(checksum_bits(checksum)) + self._keys = keys + self._indices = {key: index for index, key in enumerate(keys)} + self._check_bytes = tuple(check_bytes) + self._expected: torch.Tensor | None = torch.tensor(bits, dtype=torch.int64, device=device) + self._counts: torch.Tensor | None = torch.zeros(len(keys), dtype=torch.int64, device=device) + self._seen = [0] * len(keys) + self._checks = 0 + self._mismatches = 0 + + @classmethod + def from_layout(cls, layout: DenseLayout, device: torch.device | str) -> SlotWatchdog: + """The watchdog of a collected layout whose P1 checksums are computed. + + Raises: + DenseLayoutError: The segment checksums were not computed. + """ + checksums = layout.checksums() + expected = { + key: (segment.blob_bytes, checksum) + for (key, segment), checksum in zip(layout.segments.items(), checksums) + } + return cls(layout.layers, layout.segment_kinds, expected, device) + + @property + def layers(self) -> tuple[int, ...]: + """Covered layers, ascending.""" + return self._layers + + @property + def segments(self) -> tuple[SegmentKind, ...]: + """Enabled segment kinds in canonical order.""" + return self._segments + + @property + def checks(self) -> int: + """Slot checks enqueued so far.""" + return self._checks + + @property + def mismatches(self) -> int: + """Checks that found other bytes than their layer's, as of the latest :meth:`poll`.""" + return self._mismatches + + def check(self, layer: int, segment: SegmentKind, slot: SlotBuffer) -> None: + """Enqueue the check of ``slot``, which must hold ``layer``'s ``segment`` blob. + + Runs on the current stream without synchronizing the host; a mismatch counts on the + device until the next :meth:`poll`. + + Raises: + DenseProtocolError: ``(layer, segment)`` is not covered, or called after + :meth:`release`. + """ + index = self._indices.get((layer, SegmentKind(segment))) + if index is None: + raise DenseProtocolError( + f"the slot watchdog has no checksum for layer {layer} {SegmentKind(segment).value}" + ) + counts, expected = self._counts, self._expected + if counts is None or expected is None: + raise DenseProtocolError("slot watchdog check after release") + self._checks += 1 + nbytes = self._check_bytes[index] + if nbytes: + sums = device_checksum(slot.blob_view(nbytes)) + counts[index].add_(sums.ne(expected[index]).any()) + + def poll(self) -> list[Finding]: + """Read the counters; synchronizes the host with the current stream. + + Returns: + ``(layer, segment, mismatches)`` for every pair whose counter grew since the + previous read, in layer order; empty after :meth:`release`. + """ + counts = self._counts + if counts is None: + return [] + findings = [] + for index, count in enumerate(counts.tolist()): + new = count - self._seen[index] + if new: + self._seen[index] = count + self._mismatches += new + layer, kind = self._keys[index] + findings.append((layer, kind, new)) + return findings + + def release(self) -> None: + """Drop the device tensors; :attr:`checks` and :attr:`mismatches` keep their values.""" + self._counts = self._expected = None + + @staticmethod + def describe(findings: Sequence[Finding]) -> str: + """One line naming the ``(layer, segment)`` pairs of ``findings`` and their counts.""" + listed = ", ".join( + f"layers.{layer} {kind.value} x{count}" for layer, kind, count in findings[:_MAX_LISTED] + ) + more = len(findings) - _MAX_LISTED + return listed + (f" and {more} more" if more > 0 else "") diff --git a/tensorrt_llm/_torch/modules/dwdp/fabric_gather.py b/tensorrt_llm/_torch/modules/dwdp/fabric_gather.py index 43311d9189d8..d8aff2a22a25 100644 --- a/tensorrt_llm/_torch/modules/dwdp/fabric_gather.py +++ b/tensorrt_llm/_torch/modules/dwdp/fabric_gather.py @@ -109,17 +109,22 @@ def available() -> bool: return is_enabled() and gather_op() is not None -def resident_sms() -> int: - """SMs the persistent compute grids leave to the gather; read once, by ``DwdpManager``. +def resident_sms(gather: Optional[bool] = None) -> int: + """SMs the persistent compute grids leave to the gather; read once, by the DWDP managers. ``TRTLLM_CUTEDSL_SM_RESERVE`` pins the count. Otherwise it is the gather's block count where the gather is the transport and one block of ``THREADS >= 512`` fills an SM; smaller blocks share SMs with the compute kernels and reserve none. + + Args: + gather: Whether the gather is the transport; None means :func:`available`, the expert + DWDP choice (``DwdpManager``). Standalone dense DWDP passes its own choice. """ pinned = _int_env(_SM_RESERVE_ENV, -1) if pinned >= 0: return pinned - if not available() or THREADS < 512: + uses_gather = available() if gather is None else gather + if not uses_gather or THREADS < 512: return 0 return total_blocks() diff --git a/tensorrt_llm/_torch/modules/dwdp/page_pool.py b/tensorrt_llm/_torch/modules/dwdp/page_pool.py index 84c09788a68c..2d10f56d5a7a 100644 --- a/tensorrt_llm/_torch/modules/dwdp/page_pool.py +++ b/tensorrt_llm/_torch/modules/dwdp/page_pool.py @@ -15,12 +15,13 @@ """Page pool for the DWDP remote prefetch slots. -The page pool manages per-page fabric handles for the remote regions of the +The page pool manages per-page handles for the remote regions of the composite buffer. One pool per prefetch slot (``DwdpConfig.num_prefetch_slots``, two by default). -Each pool is a list of page-sized fabric handles that can be mapped at different -VA positions for different layers. Pages are reused across layers in the same slot: +Each pool is a list of page-sized local (``CU_MEM_HANDLE_TYPE_NONE``) handles +that can be mapped at different VA positions for different layers. Pages are +reused across layers in the same slot: - Layer 3 (buf_idx=0): pool[0] pages 0..K for pre/post regions - Layer 5 (buf_idx=0): SAME pool[0] pages (slot reuse) - Layer 4 (buf_idx=1): pool[1] pages 0..M @@ -28,10 +29,9 @@ This naturally handles heterogeneous layers: each layer maps exactly the pages it needs, even if remote sizes differ across layers. -Verified on GB200: - - Per-page fabric handles composited into one VA: PASS - - Same page handles reused across layers: PASS - - Non-fabric handles cannot be mapped: FAIL (not supported on GB200) +The pages are never exported: they are mapped next to the fabric handle of a +weight's local experts in one VA reservation, and one set_access covers the +whole reservation. """ from __future__ import annotations @@ -51,10 +51,10 @@ class PagePool: - """Pool of per-page fabric handles for the remote prefetch slots. + """Pool of per-page local handles for the remote prefetch slots. - One pool per prefetch slot. Each pool is a list - of page-sized fabric handles. Different layers map these pages at + One pool per prefetch slot. Each pool is a list of page-sized + ``CU_MEM_HANDLE_TYPE_NONE`` handles. Different layers map these pages at different VA positions within their composite VA. Pages are reused across layers in the same slot: @@ -65,9 +65,8 @@ class PagePool: Pool page size vs VMM granularity: Each pool handle is ``page_size`` bytes (default 8 * granularity = 16MB on GB200). Larger pool pages reduce the number of cuMemMap - calls by 8x, avoiding the CUDA driver internal tracking limit - (~130K mappings per process) that causes OOM when 4 processes share - a GB200 node with 2MB pages. + calls and live mappings by 8x. With 2MB pages, 4 processes on a + GB200 node run out of memory at ~130K live mappings per process. Attributes: device_id: CUDA device ordinal. @@ -156,15 +155,16 @@ def __init__( handle = create_local_handle(self._page_size, device_id) handles.append(handle) except Exception as e: - # Cleanup on failure - for h in handles: - release_handle(h) - for prev_handles in self._page_handles: - for h in prev_handles: - release_handle(h) + # Release every page made so far exactly once, here, as release() does: a + # page that fails to release is logged, so the others are still released and + # this allocation error stays the one raised; release() also marks the pool + # released, so __del__ does not release the handles again once the driver + # may have reused their values. + self._page_handles.append(handles) + self.release() raise RuntimeError( f"Failed to allocate page {page_idx} in slot {slot_idx}: {e}" - ) + ) from e self._page_handles.append(handles) logger.debug( f"[PagePool] Allocated {num_pages} pages ({self._page_size} B each) " diff --git a/tensorrt_llm/_torch/modules/dwdp/setup.py b/tensorrt_llm/_torch/modules/dwdp/setup.py index b8203c2e6e0b..8918c9ec3509 100644 --- a/tensorrt_llm/_torch/modules/dwdp/setup.py +++ b/tensorrt_llm/_torch/modules/dwdp/setup.py @@ -22,7 +22,9 @@ ``layer_indices`` (the Single Source of Truth coming from ``DwdpManager._registered_layers``). 2. Builds per-layer WeightSpecs from the collected parameter shapes. -3. Computes each rank's local expert range. +3. Computes the expert partition (``ExpertPartition``), the only source of + every rank's expert range, and agrees on it across the DWDP group + (``_agree_on_partition``). 4. Allocates MNNVL fabric handles via ``DWDPTransport.create()``, copying local weights into fabric memory and freeing the originals. Handles are exchanged via the DWDP MPI communicator. @@ -30,10 +32,21 @@ + page-pool-backed remote prefetch slots). 6. Fills page-alignment edge bytes at MNNVL region boundaries. 7. Creates a DWDPWeightManager for runtime P2P prefetch scheduling. -8. Patches each MoE backend so it sees all ``num_experts`` (ep_size=1) and - allgathers small EP-sharded parameters (e.g. e_score_correction_bias, - fc31_alpha, fc2_alpha) via MPI. Large scale params - (w3_w1_weight_scale, w2_weight_scale) go through Transport (MNNVL + P2P). +8. Exchanges the small per-expert parameters (e.g. fc31_alpha, fc2_alpha, + an EP-sharded e_score_correction_bias) in one collective + (``_exchange_expert_scalars``), which is also the group's vote on steps + 5-7, then patches each MoE backend so it sees all ``num_experts`` + (ep_size=1). Large scale params (w3_w1_weight_scale, w2_weight_scale) + go through Transport (MNNVL + P2P). + +From the handshake to the vote, every check of the setup is made from data +the whole group gathered, the same way on every rank: the handshake, each +transport handle allgather, the transport's import status and the vote. A +failure on one rank in that span thus raises on every rank after the same +collective, instead of leaving its peers blocked in the next one. The +backend fixup after the vote is local: it installs the exchanged tensors (a +few KiB of device memory per MoE layer), and a failure there raises on that +rank only. The function is idempotent: calling it a second time on the same model is a no-op (returns the existing manager stored on the model). @@ -41,16 +54,19 @@ from __future__ import annotations +import itertools import traceback -from typing import Dict, List, Optional, Sequence, Tuple +from collections import Counter +from typing import Dict, List, NamedTuple, Optional, Sequence, Tuple import torch from torch import nn from tensorrt_llm.logger import logger -from .specs import LayerWeightSpecs, PeerRanges, WeightSpec, lookup_owner -from .transport import DWDPTransport +from .specs import ExpertPartition, LayerWeightSpecs, PeerRanges, WeightSpec, lookup_owner +from .transport import DWDPTransport, _allgather_keeping, _describe_error +from .vmm import get_allocation_granularity from .weight_buffer import WeightBuffer from .weight_manager import DWDPWeightManager @@ -68,6 +84,12 @@ "w2_weight_scale", ) +# Name keywords of the small parameters of an experts module that the fixup +# exchanges (``_exchange_expert_scalars``): fc31_alpha, fc2_alpha, the +# activation input scales, ... ; names in _EXPERT_WEIGHT_NAMES go through +# the transport instead. +_SCALE_KEYWORDS = ("scale", "scaling_factor", "dequant", "alpha") + def _release_during_rollback(resource, what: str) -> None: """Release a partially built resource without masking the in-flight error. @@ -85,6 +107,54 @@ def _release_during_rollback(resource, what: str) -> None: ) +def _roll_back_setup( + transport: DWDPTransport, + weight_buffer: Optional[WeightBuffer], + weight_manager: Optional[DWDPWeightManager], +) -> None: + """Release what ``setup_dwdp`` built after the transport, without masking the error in flight. + + The weight manager owns the buffer and the transport once it holds the transport. + """ + if weight_manager is not None and weight_manager._transport is transport: + _release_during_rollback(weight_manager, "weight manager") + return + if weight_buffer is not None: + _release_during_rollback(weight_buffer, "weight buffer") + _release_during_rollback(transport, "transport") + + +def _refuse_reserved_storages( + local_params: Dict[Tuple[int, str], torch.Tensor], reserved_storages: frozenset[int] +) -> None: + """Refuse expert weights whose storage another owner holds, before the transport takes them. + + The transport copies each weight into its own handle and then frees the weight's storage. + ``DwdpManager.setup`` reserves the slots of the committed attached dense DWDP runtime: a + weight there is a view dense DWDP repointed into a slot, which happens only when a dense + root reaches a routed-expert weight. The transport would then take the slot's bytes for the + experts and free the slot under the dense runtime. + + Raises: + RuntimeError: Expert weights live in reserved storages; the message lists them. + """ + if not reserved_storages: + return + taken = [ + f"layers.{layer}.{name}" + for (layer, name), tensor in sorted(local_params.items(), key=lambda item: item[0]) + if tensor.untyped_storage().data_ptr() in reserved_storages + ] + if taken: + listed = ", ".join(taken[:8]) + (f" and {len(taken) - 8} more" if len(taken) > 8 else "") + raise RuntimeError( + f"[DWDP Setup] Expert weights in reserved storage, which the transport would free: " + f"{listed}. They are views of the committed dense DWDP slots: dense DWDP moved them, " + f"so the decoder layers' dense_dwdp_roots (or dense_dwdp_attr_allow) reach the routed " + f"experts, which expert DWDP owns." + ) + + def setup_dwdp( model: nn.Module, mapping, @@ -96,6 +166,8 @@ def setup_dwdp( num_prefetch_slots: int, contention_opt: bool = False, repeated_layer_indices: Sequence[int] = (), + copy_stream: Optional[torch.cuda.Stream] = None, + reserved_storages: frozenset[int] = frozenset(), ) -> Optional[DWDPWeightManager]: """Set up DWDP for a model after weight loading. @@ -111,19 +183,18 @@ def setup_dwdp( layer_indices: MoE layer indices collected by ``DwdpManager._registered_layers`` (SSOT). Each index must correspond to an MoE layer in the model. - num_experts_per_worker: Range size — number of experts each rank - owns (``DwdpConfig.num_experts_per_worker``). Used as the local - expert range length: ``[start, start + num_experts_per_worker)``. - This is also the *uniform storage chunk size* installed on every - fused MoE backend by ``_init_dwdp_expert_layout``; the tail - rank may have a few padding slots past ``num_experts``. - num_prefetch_experts: Stride between consecutive ranks - (``DwdpConfig.num_prefetch_experts``). Used as the rank-to-rank - offset: rank ``r`` owns ``[r * num_prefetch_experts, ...)``. - For uniform integer-division partitioning these two values are - equal (e.g. dwdp=4: 64/64, dwdp=8: 32/32). Setting - ``num_prefetch_experts < num_experts_per_worker`` enables - redundancy: adjacent ranks' ranges overlap by + num_experts_per_worker: ``DwdpConfig.num_experts_per_worker``: the + experts each rank stores in a strided expert partition, or 0 + (with ``num_prefetch_experts``) for the balanced partition. + num_prefetch_experts: ``DwdpConfig.num_prefetch_experts``: the + stride between consecutive ranks' first experts in a strided + partition, or 0 for the balanced one. Both go to + ``ExpertPartition.from_config`` with the model's expert count, + as in ``DwdpManager.expert_partition``, which sized the MoE + backends before their weights were created; the handshake checks + that every rank computed the same partition and loaded the rows + of its range. Under redundancy (``num_prefetch_experts < + num_experts_per_worker``) adjacent ranks' ranges overlap by ``num_experts_per_worker - num_prefetch_experts`` experts and ``lookup_owner`` resolves the lowest-rank owner. num_prefetch_slots: Number of prefetch buffer slots, each holding the @@ -132,11 +203,30 @@ def setup_dwdp( ``cudaMemcpyBatchAsync`` per layer (``DwdpConfig.contention_opt``). repeated_layer_indices: Layers of ``layer_indices`` that run several times in a row within a forward pass (``DwdpManager.add_repeated_layer``). + copy_stream: The weight manager's copy stream (``DWDPWeightManager`` + creates one if None); ``DwdpManager.setup`` passes the stream + attached dense DWDP was set up with. A cached manager keeps its own. + reserved_storages: ``untyped_storage().data_ptr()`` of storages that + another owner holds, which the transport must not take over: + ``DwdpManager.setup`` passes the slots of the committed attached + dense DWDP runtime. An expert weight in one of them raises a + RuntimeError before the transport starts. Returns: A ready-to-use DWDPWeightManager if DWDP is enabled, ``None`` otherwise. The manager is also stored on ``model.dwdp_weight_manager`` so that a second call returns the cached instance (idempotent). + + Raises: + RuntimeError: The DWDP group disagrees on the expert partition, a + rank failed in the transport or in a later step before the vote, + or the ranks' small expert parameters do not fit together. + Raised on every rank after the same collective: a rank that + failed raises its own error, the others an error naming it. A + failure in the backend fixup after the vote raises on that rank + only. + ValueError: The config gives no valid partition of the model's + experts (raised on every rank alike). """ # --- Guard: DWDP not enabled --- if not mapping.dwdp_enabled: @@ -171,6 +261,7 @@ def setup_dwdp( # 1. Collect MoE params (SSOT: iterate passed layer_indices only) local_params, weight_names = collect_moe_params(model, layer_indices) + _refuse_reserved_storages(local_params, reserved_storages) logger.info( f"[DWDP Setup] Collected params from {len(layer_indices)} MoE layers, " f"weight_names={weight_names}" @@ -178,11 +269,9 @@ def setup_dwdp( # Query the model's *actual* gate-side num_experts from the first MoE # layer. This is the routing-side expert count (e.g., 72 for DSv3-Lite, - # 256 for DeepSeek-V3 / R1). It is distinct from the per-rank storage - # extent (``num_experts_per_worker * dwdp_size``) which may be larger - # under non-uniform / redundancy Mode B partitioning. All Phase 1+2 - # plumbing — composite VA size, peer_ranges cap, validation — uses - # this gate-side count, not storage extent. + # 256 for DeepSeek-V3 / R1, 896 for Kimi-K3), the one the expert + # partition splits. All Phase 1+2 plumbing — composite VA size, the + # partition, validation — uses this gate-side count. num_experts_total = _get_model_num_experts(model, layer_indices) # 2. Build weight specs (full_shape uses the gate-side num_experts) @@ -194,21 +283,28 @@ def setup_dwdp( f"num_experts={num_experts_total}" ) - # 3. Validate partition config against the gate-side expert count. - first_spec = _get_first_spec(layer_weight_specs) - _validate_partition_config( + # 3. The expert partition, from the function _init_dwdp_expert_layout + # sized the backends with, and the group's verdict on it before any + # handle exists: one allgather after which every rank goes on, or + # every rank raises. + partition = ExpertPartition.from_config( + num_experts=num_experts_total, + dwdp_size=mapping.dwdp_size, num_experts_per_worker=num_experts_per_worker, num_prefetch_experts=num_prefetch_experts, - num_experts_total=num_experts_total, - dwdp_size=mapping.dwdp_size, - loaded_local_experts=first_spec.local_experts, ) - local_start = mapping.dwdp_rank * num_prefetch_experts - local_end = local_start + num_experts_per_worker + _agree_on_partition( + comm, + partition, + dwdp_rank=mapping.dwdp_rank, + granularity=get_allocation_granularity(device_id), + layer_weight_specs=layer_weight_specs, + ) + local_start, local_end = partition.range_of(mapping.dwdp_rank) logger.info( - f"[DWDP Setup] Expert range: local=[{local_start}, {local_end}), " - f"total={num_experts_total}, " - f"size={num_experts_per_worker}, stride={num_prefetch_experts}" + f"[DWDP Setup] Expert range: local=[{local_start}, {local_end}) " + f"({local_end - local_start} experts), total={num_experts_total}, " + f"partition {partition.describe()}" ) # 4. Transport: MNNVL alloc -> copy -> free originals -> exchange handles @@ -218,24 +314,29 @@ def setup_dwdp( local_params=local_params, comm=comm, dwdp_rank=mapping.dwdp_rank, - dwdp_size=mapping.dwdp_size, device_id=device_id, - local_start=local_start, - local_end=local_end, - num_experts_per_worker=num_experts_per_worker, - num_prefetch_experts=num_prefetch_experts, + partition=partition, ) logger.info("[DWDP Setup] Transport created successfully.") - free_mem, total_mem = torch.cuda.mem_get_info(device_id) - logger.info( - f"[DWDP Setup] GPU {device_id} memory after transport: " - f"free={free_mem/(1024**3):.2f}GB / total={total_mem/(1024**3):.2f}GB " - f"(used={((total_mem-free_mem)/(1024**3)):.2f}GB)" - ) - # 5. WeightBuffer: composite VA layout - logger.info("[DWDP Setup] Creating weight buffer (composite VA layout)...") + # 5-7. Local work: an error here (e.g. this rank's prefetch ring does not + # fit) is only recorded, and raised after the vote in step 8, which every + # rank enters; the peers then raise there too, instead of waiting for + # this rank in a later collective. + weight_buffer = None + weight_manager = None + local_error = None try: + free_mem, total_mem = torch.cuda.mem_get_info(device_id) + logger.info( + f"[DWDP Setup] GPU {device_id} memory after transport: " + f"free={free_mem / (1024**3):.2f}GB / total={total_mem / (1024**3):.2f}GB " + f"(used={((total_mem - free_mem) / (1024**3)):.2f}GB)" + ) + peer_ranges = partition.peer_ranges + + # 5. WeightBuffer: composite VA layout + logger.info("[DWDP Setup] Creating weight buffer (composite VA layout)...") weight_buffer = WeightBuffer.create( layer_weight_specs=layer_weight_specs, handles=transport.get_handle_set(), @@ -245,17 +346,19 @@ def setup_dwdp( device_id=device_id, num_slots=num_prefetch_slots, ) - except BaseException: - _release_during_rollback(transport, "transport") - raise - logger.info( - f"[DWDP Setup] Weight buffer created successfully ({num_prefetch_slots} prefetch slots)." - ) + logger.info( + f"[DWDP Setup] Weight buffer created successfully ({num_prefetch_slots} prefetch " + f"slots)." + ) + free_mem, total_mem = torch.cuda.mem_get_info(device_id) + logger.info( + f"[DWDP Setup] GPU {device_id} memory after weight buffer: " + f"free={free_mem / (1024**3):.2f}GB / total={total_mem / (1024**3):.2f}GB " + f"(used={((total_mem - free_mem) / (1024**3)):.2f}GB)" + ) - # 6. Fill edge bytes - logger.info("[DWDP Setup] Filling page-alignment edge bytes...") - peer_ranges = transport.get_peer_ranges() - try: + # 6. Fill edge bytes + logger.info("[DWDP Setup] Filling page-alignment edge bytes...") fill_edge_bytes( weight_buffer=weight_buffer, peer_views=transport.get_peer_views(), @@ -263,15 +366,10 @@ def setup_dwdp( local_end=local_end, peer_ranges=peer_ranges, ) - except BaseException: - _release_during_rollback(weight_buffer, "weight buffer") - _release_during_rollback(transport, "transport") - raise - logger.info("[DWDP Setup] Edge bytes filled.") + logger.info("[DWDP Setup] Edge bytes filled.") - # 7. WeightManager - logger.info("[DWDP Setup] Creating weight manager...") - try: + # 7. WeightManager + logger.info("[DWDP Setup] Creating weight manager...") weight_manager = DWDPWeightManager( weight_buffer=weight_buffer, peer_views=transport.get_peer_views(), @@ -282,31 +380,38 @@ def setup_dwdp( dwdp_size=mapping.dwdp_size, contention_opt=contention_opt, repeated_layer_indices=repeated_layer_indices, + copy_stream=copy_stream, + ) + # Transfer resource ownership before backend fixup so every later failure + # can be rolled back through one idempotent manager release. + weight_manager._transport = transport + logger.info("[DWDP Setup] Weight manager created.") + except Exception as exc: # noqa: BLE001 - any failure is this rank's vote below + local_error = exc + logger.error( + f"[DWDP Setup] Rank {mapping.dwdp_rank} failed after the transport with {exc!r}; " + f"reporting it to the DWDP group" ) except BaseException: - _release_during_rollback(weight_buffer, "weight buffer") - _release_during_rollback(transport, "transport") + _roll_back_setup(transport, weight_buffer, weight_manager) raise - # Transfer resource ownership before backend fixup so every later failure - # can be rolled back through one idempotent manager release. - weight_manager._transport = transport - logger.info("[DWDP Setup] Weight manager created.") - # 8. Patch MoE backends + # 8. The vote and the backend fixup: _exchange_expert_scalars is the + # fixup's only collective, entered by every rank; it raises on every + # rank if any rank failed, and otherwise returns the full small + # per-expert tensors, which fixup_moe_backends installs locally. logger.info("[DWDP Setup] Patching MoE backends...") try: - fixup_moe_backends( + expert_scalars = _exchange_expert_scalars( + comm, model, layer_indices, - num_experts_total, - comm, - mapping.dwdp_rank, - mapping.dwdp_size, - num_experts_per_worker=num_experts_per_worker, - peer_ranges=peer_ranges, + partition, + error=local_error, ) + fixup_moe_backends(model, layer_indices, num_experts_total, expert_scalars) except BaseException: - _release_during_rollback(weight_manager, "weight manager") + _roll_back_setup(transport, weight_buffer, weight_manager) raise logger.info("[DWDP Setup] MoE backends patched.") @@ -447,16 +552,16 @@ def build_weight_specs( ) -> LayerWeightSpecs: """Convert collected parameters to ``LayerWeightSpecs``. - Each local parameter has shape ``(experts_per_rank, ...)`` where - ``experts_per_rank`` is the per-rank storage size - (``num_experts_per_worker`` from ``DwdpConfig`` after - ``_init_dwdp_expert_layout`` runs). The composite VA's ``full_shape`` - uses ``num_experts_total`` (gate-side, e.g. 72 for DSv3-Lite) so - every expert id in ``[0, num_experts_total)`` is addressable — and - *only* those. Under Mode B redundancy ``experts_per_rank * - dwdp_size`` may exceed ``num_experts_total`` (overlap duplicates), - but those duplicates live within each peer's MNNVL handle, not in - the composite VA. + Each local parameter has shape ``(local_experts, ...)`` where + ``local_experts`` is the number of experts this rank stores: the length + of its range in the expert partition, which ``_init_dwdp_expert_layout`` + sized the backend with (ranks may differ, e.g. 75 and 74). The + composite VA's ``full_shape`` uses ``num_experts_total`` (gate-side, + e.g. 72 for DSv3-Lite) so every expert id in ``[0, num_experts_total)`` + is addressable — and *only* those. Under redundancy the ranks' counts + add up to more than ``num_experts_total`` (overlap duplicates), but + those duplicates live within each peer's MNNVL handle, not in the + composite VA. Args: local_params: Dict from ``collect_moe_params``. @@ -536,6 +641,126 @@ def _get_model_num_experts(model: nn.Module, layer_indices: List[int]) -> int: return n +# --------------------------------------------------------------------------- +# Helper: agree on the expert partition (handshake) +# --------------------------------------------------------------------------- + + +def _agree_on_partition( + comm, + partition: ExpertPartition, + *, + dwdp_rank: int, + granularity: int, + layer_weight_specs: LayerWeightSpecs, +) -> None: + """One allgather after which every rank of the DWDP group goes on, or every rank raises. + + Each rank reports its DWDP rank, the partition it computed (expert count, every rank's + range, kind), its VMM granularity and, for every (layer, weight) expert chunk it loaded, + the rows, the shape of one expert and the dtype. The verdict is a function of the gathered + reports only, so it is the same on every rank: the transport after it is entered by every + rank or by none. The transport then maps every peer's handle with the peer's range and its + own granularity and bytes per expert, which this makes the same on every rank. + Rejected: a rank that sent no handshake (it runs another TensorRT-LLM tree); a rank whose + partition (config or model drift), granularity, or loaded (layer, weight) chunks, expert + shapes or dtypes differ from the most common report's, so that the ranks named are the + ones that differ from the others, rank 0 included; a chunk whose rows are not its rank's + expert count (``_init_dwdp_expert_layout`` sized the backend from another partition); and + a report out of rank order. + + Raises: + RuntimeError: On every rank, naming each offending rank and item (the first 8). + """ + chunks = tuple( + (layer_idx, name, tuple(spec.chunk_shape[1:]), str(spec.dtype), spec.local_experts) + for layer_idx in sorted(layer_weight_specs) + for name, spec in sorted(layer_weight_specs[layer_idx].items()) + ) + report = ( + dwdp_rank, + partition.num_experts, + partition.ranges, + partition.kind, + granularity, + chunks, + ) + reports = comm.allgather(report) + malformed = [ + rank + for rank, rank_report in enumerate(reports) + if not (isinstance(rank_report, tuple) and len(rank_report) == len(report)) + ] + if malformed: + raise RuntimeError( + f"[DWDP Setup] DWDP ranks {malformed} sent no expert partition handshake: every rank " + f"of the DWDP group must run the same TensorRT-LLM tree." + ) + + problems = [] + if len(reports) != partition.dwdp_size: + problems.append(f"{len(reports)} ranks reported, the partition has {partition.dwdp_size}") + # The reference is the most common report, rank and rows aside (on a tie, the lowest rank's), + # so that the ranks named are the ones that differ from the others, rank 0 included. + keys = [ + (num_experts, ranges, kind, rank_granularity, tuple(chunk[:-1] for chunk in rank_chunks)) + for _, num_experts, ranges, kind, rank_granularity, rank_chunks in reports + ] + ref_rank = keys.index(Counter(keys).most_common(1)[0][0]) + _, ref_experts, ref_ranges, ref_kind, ref_granularity, ref_chunks = reports[ref_rank] + # Every rank must load the same (layer, weight) chunks, of the same expert shape and dtype. + ref_layout = [chunk[:-1] for chunk in ref_chunks] + for rank, rank_report in enumerate(reports): + reported_rank, num_experts, ranges, kind, rank_granularity, rank_chunks = rank_report + if reported_rank != rank: + problems.append(f"position {rank} holds the report of DWDP rank {reported_rank}") + continue + if (num_experts, ranges, kind) != (ref_experts, ref_ranges, ref_kind): + problems.append( + f"rank {rank} computed the {kind} partition {list(ranges)} of {num_experts} " + f"experts, rank {ref_rank} the {ref_kind} partition {list(ref_ranges)} of " + f"{ref_experts}" + ) + continue + if rank_granularity != ref_granularity: + problems.append( + f"rank {rank} maps with a granularity of {rank_granularity} bytes, rank " + f"{ref_rank} with {ref_granularity}" + ) + layout = [chunk[:-1] for chunk in rank_chunks] + if layout != ref_layout: + mine, ref = next( + (mine, ref) + for mine, ref in itertools.zip_longest(layout, ref_layout) + if mine != ref + ) + problems.append( + f"rank {rank} loaded other expert chunks than rank {ref_rank}: (layer, weight, " + f"expert shape, dtype) {mine}, rank {ref_rank} {ref}" + ) + continue + if rank >= len(ranges): + continue + start, end = ranges[rank] + wrong = [(layer_idx, name, n) for layer_idx, name, *_, n in rank_chunks if n != end - start] + if wrong: + layer_idx, name, n = wrong[0] + problems.append( + f"rank {rank} stores experts [{start}, {end}) ({end - start} experts) but " + f"loaded {len(wrong)} expert chunks with other row counts, e.g. layer " + f"{layer_idx} {name} with {n} rows" + ) + if problems: + raise RuntimeError( + f"[DWDP Setup] The DWDP group does not agree on the expert partition: " + f"{'; '.join(problems[:8])}" + + (f"; and {len(problems) - 8} more" if len(problems) > 8 else "") + + ". Every rank must run the same DwdpConfig and model; " + "_init_dwdp_expert_layout sizes the MoE backends from " + "DwdpManager.expert_partition()." + ) + + # --------------------------------------------------------------------------- # Helper: fill edge bytes # --------------------------------------------------------------------------- @@ -551,9 +776,12 @@ def fill_edge_bytes( """Fill page-alignment edge bytes at MNNVL region boundaries. When local expert data does not align to page boundaries, the MNNVL pages - contain edge regions that must be filled with the correct peer data. This - happens at most twice per (layer, weight): a *leading edge* before - ``local_start`` and a *trailing edge* after ``local_end``. + contain edge regions that hold peer data: per (layer, weight) at most one + *leading edge* ``[page_start, local_start * expert_bytes)`` and one + *trailing edge* ``[local_end * expert_bytes, page_end)``. With experts + smaller than the granularity an edge holds bytes of several experts; only + the experts next to the local range are filled here. Every prefetch + rewrites the remote experts, edges included, before any read. The data is read from peer MNNVL tensor views and written into the local MNNVL handle's mapped VA. This is a one-time setup operation. @@ -563,10 +791,11 @@ def fill_edge_bytes( peer_views: ``{(peer_rank, layer_idx, name): tensor}`` from Transport. local_start: First local expert index (inclusive). local_end: Last local expert index (exclusive). - peer_ranges: Per-peer ``(local_start, local_end_capped)`` tuples. + peer_ranges: Every rank's ``(start, end)`` in the expert partition. ``lookup_owner`` resolves the owner of ``prev_expert`` / - ``next_expert`` here, so non-uniform / redundancy partitions - pick the right peer. + ``next_expert`` here, and each peer view has that peer's rows, + so any partition (balanced, capped, redundant) picks the right + peer and row. """ if local_start >= local_end: return @@ -585,12 +814,11 @@ def fill_edge_bytes( full_tensor = weight_buffer.get_full_tensor(layer_idx, name) # Leading edge: bytes in [page_start, local_start_bytes). - # These belong to experts just before local_start. + # These belong to the experts before local_start (possibly several). if edge_info.leading_edge > 0 and local_start > 0: - # The leading edge is within the expert just before - # local_start. Use lookup_owner so non-uniform / redundant - # partitions pick the lowest-rank peer that owns - # ``prev_expert``. + # Fill the expert just before local_start. Use lookup_owner so + # non-uniform / redundant partitions pick the lowest-rank peer + # that owns ``prev_expert``. prev_expert = local_start - 1 peer_rank = lookup_owner(prev_expert, peer_ranges) peer_chunk_start, _ = peer_ranges[peer_rank] @@ -599,9 +827,8 @@ def fill_edge_bytes( peer_key = (peer_rank, layer_idx, name) if peer_key in peer_views: src = peer_views[peer_key] - # The edge bytes correspond to the tail portion of - # prev_expert's data that spills into our MNNVL page. - # Copy the full expert slice (the composite VA's + # prev_expert's data, or its tail, lies in our MNNVL + # page. Copy the full expert slice (the composite VA's # expert dimension is page-aligned). full_tensor[prev_expert].copy_(src[peer_local_offset]) logger.debug( @@ -610,7 +837,7 @@ def fill_edge_bytes( ) # Trailing edge: bytes in [local_end_bytes, page_end). - # These belong to experts just after local_end. + # These belong to the experts from local_end on (possibly several). if edge_info.trailing_edge > 0 and local_end < full_tensor.shape[0]: next_expert = local_end peer_rank = lookup_owner(next_expert, peer_ranges) @@ -630,20 +857,274 @@ def fill_edge_bytes( # --------------------------------------------------------------------------- -# Helper: fixup MoE backends +# Helper: exchange the small expert parameters (the vote), fixup MoE backends # --------------------------------------------------------------------------- +class _ParamReport(NamedTuple): + """One small parameter of one layer as a rank reports it to ``_exchange_expert_scalars``. + + Attributes: + shape: The parameter's shape on the rank. + data: A CPU copy; None when the first dimension is the model's expert count, since a + full-sized parameter needs no exchange. + """ + + shape: Tuple[int, ...] + data: Optional[torch.Tensor] + + +class _LayerScalars(NamedTuple): + """What ``fixup_moe_backends`` installs on one MoE layer after ``_exchange_expert_scalars``. + + Attributes: + params: ``[num_experts, ...]`` CPU tensors of the experts module's per-expert small + parameters, by name, put together from every rank's shard. + bias: The ``[num_experts]`` CPU gate ``e_score_correction_bias`` put together from every + rank's shard, or None when it is full-sized on every rank already, or absent. + """ + + params: Dict[str, torch.Tensor] + bias: Optional[torch.Tensor] + + +def _report_param(tensor: torch.Tensor, num_experts: int) -> _ParamReport: + """``tensor`` as this rank reports it: its shape, and a CPU copy unless it is full-sized.""" + shape = tuple(tensor.shape) + if shape and shape[0] == num_experts: + return _ParamReport(shape, None) + return _ParamReport(shape, tensor.detach().cpu().contiguous()) + + +def _collect_expert_scalars( + model: nn.Module, layer_indices: List[int], num_experts: int +) -> Dict[int, Tuple[Dict[str, _ParamReport], Optional[_ParamReport]]]: + """This rank's small parameters of every MoE layer, as ``_exchange_expert_scalars`` sends them. + + Per layer: the direct parameters of the experts module whose names contain one of + ``_SCALE_KEYWORDS`` and are not in ``_EXPERT_WEIGHT_NAMES`` (the transport moves those), + and the gate's ``e_score_correction_bias`` (None without one). + """ + layers = _get_decoder_model(model).layers + collected = {} + for layer_idx in layer_indices: + moe_module, experts_module = _get_moe_and_experts(layers[layer_idx]) + params = {} + if experts_module is not None: + for pname, param in experts_module.named_parameters(recurse=False): + if pname in _EXPERT_WEIGHT_NAMES or not any(kw in pname for kw in _SCALE_KEYWORDS): + continue + params[pname] = _report_param(param.data, num_experts) + bias = None + gate_module = _get_gate_module(moe_module) if moe_module is not None else None + bias_param = getattr(gate_module, "e_score_correction_bias", None) + if bias_param is not None: + bias_data = bias_param.data if isinstance(bias_param, nn.Parameter) else bias_param + bias = _report_param(bias_data, num_experts) + collected[layer_idx] = (params, bias) + return collected + + +def _bitwise_equal(a: Optional[torch.Tensor], b: Optional[torch.Tensor]) -> bool: + """Same dtype, shape and bytes (NaN payloads included).""" + if a is None or b is None or a.dtype != b.dtype or a.shape != b.shape: + return False + return torch.equal(a.reshape(-1).view(torch.uint8), b.reshape(-1).view(torch.uint8)) + + +def _merge_shards( + what: str, + reports: List[_ParamReport], + partition: ExpertPartition, + problems: List[str], + *, + allow_per_layer: bool, +) -> Tuple[str, Optional[torch.Tensor]]: + """Classify one parameter from every rank's report and put a per-expert one together. + + Returns: + ``(category, full)``: ``"per expert"`` with the ``[num_experts, ...]`` tensor, + ``"full"`` (full-sized on every rank, nothing to install), ``"per layer"`` (the same + shape on every rank and not per expert; it must be bitwise equal on every rank, nothing + to install) or ``"invalid"``; every inconsistency is appended to ``problems``. + """ + num_experts = partition.num_experts + shapes = [report.shape for report in reports] + counts = [partition.count_of(rank) for rank in range(len(reports))] + trailing = {shape[1:] for shape in shapes} + if len(trailing) == 1 and all(shape and shape[0] == num_experts for shape in shapes): + return "full", None + if len(trailing) == 1 and all( + shape and shape[0] == count for shape, count in zip(shapes, counts) + ): + shards = [report.data for report in reports] + if len({shard.dtype for shard in shards}) != 1: + problems.append(f"{what}: dtypes {[shard.dtype for shard in shards]} per rank") + return "invalid", None + full = _scatter_shards_to_full( + shards=shards, + peer_ranges=partition.peer_ranges, + num_experts_total=num_experts, + ref=shards[0], + ) + # Overlapping ranges store an expert on more than one rank: every copy must agree, + # i.e. each rank's shard is what was put together over its range. + differing = set() + for rank, shard in enumerate(shards): + start, end = partition.range_of(rank) + if shard.numel() == 0: + continue + own = shard.reshape(end - start, -1).view(torch.uint8) + merged = full[start:end].reshape(end - start, -1).view(torch.uint8) + rows = (own != merged).any(dim=1).nonzero().flatten().tolist() + differing.update(start + row for row in rows) + if differing: + problems.append( + f"{what}: the ranks storing experts {sorted(differing)[:8]} hold different " + f"values for them" + ) + return "per expert", full + if allow_per_layer and len(set(shapes)) == 1: + differing = [ + rank + for rank, report in enumerate(reports) + if not _bitwise_equal(report.data, reports[0].data) + ] + if differing: + problem = f"{what} is per layer, but ranks {differing} hold other values than rank 0" + if not shapes[0]: + # The activation input scales. + problem += ( + f" ({[report.data.item() for report in reports]}); a rank reduces it over the " + f"experts it loads (quantization's load_streaming_nvfp4_expert), so all " + f"experts of the checkpoint must have the same one" + ) + problems.append(problem) + return "per layer", None + problems.append( + f"{what}: shapes {shapes} per rank; a per-expert parameter has the rank's expert count " + f"{counts} or {num_experts} rows on every rank" + ) + return "invalid", None + + +def _exchange_expert_scalars( + comm, + model: nn.Module, + layer_indices: List[int], + partition: ExpertPartition, + *, + error: Optional[Exception] = None, +) -> Dict[int, _LayerScalars]: + """The fixup's one collective, which is also the group's vote on the setup steps before it. + + Every rank enters it, also after a failure, with its status (``error``, None on success) + and its small expert parameters (``_collect_expert_scalars``) of every MoE layer. The + verdict is a function of the gathered reports only, so it is the same on every rank: + + * A rank failed: every rank raises, the failed rank its own error and the others a + RuntimeError naming it, and each rolls its setup back. + * Each layer's parameter names must match across ranks. A parameter with ``count_p`` rows + on every rank ``p`` (``fc31_alpha``, ``fc2_alpha``) is per expert: the shards are put + together into ``[num_experts, ...]`` along the partition (``_scatter_shards_to_full``), + and every rank storing an expert (overlapping ranges) must hold the same values for it. + One with ``num_experts`` rows on every rank is full-sized already. A zero-dimensional + one (``fc31_input_scale``, ``fc2_input_scale``), or any other one of the same shape on + every rank, is per layer and must be bitwise equal on every rank. Anything else raises. + * The gate's ``e_score_correction_bias``: full-sized on every rank (the gate loads it + whole under ``moe_ep_size = 1``) stays; ``count_p`` entries on every rank are put + together like a per-expert parameter; anything else raises. + + Returns: + ``_LayerScalars`` per layer index, for ``fixup_moe_backends``. + + Raises: + RuntimeError: As above, on every rank after the same collective. + """ + collected = {} + if error is None: + try: + collected = _collect_expert_scalars(model, layer_indices, partition.num_experts) + except Exception as exc: # noqa: BLE001 - reported to the group below + error = exc + status = None if error is None else _describe_error(error) + reports = _allgather_keeping(comm, (status, collected), error) + if error is not None: + raise error + failed = [ + f"rank {rank}: {rank_status}" + for rank, (rank_status, _) in enumerate(reports) + if rank_status is not None + ] + if failed: + raise RuntimeError( + f"[DWDP Setup] DWDP setup failed on {'; '.join(failed)}. Every rank of the DWDP " + f"group raises this and releases its DWDP resources." + ) + + problems: List[str] = [] + categories: Dict[str, set[str]] = {} + expert_scalars: Dict[int, _LayerScalars] = {} + for layer_idx in layer_indices: + per_rank = [layers.get(layer_idx) for _, layers in reports] + if any(entry is None for entry in per_rank): + problems.append(f"layer {layer_idx} is missing on some ranks") + continue + names = sorted(per_rank[0][0]) + if any(sorted(params) != names for params, _ in per_rank): + problems.append( + f"layer {layer_idx}: the ranks have other small expert parameters " + f"({[sorted(params) for params, _ in per_rank]})" + ) + continue + params = {} + for pname in names: + category, full = _merge_shards( + f"layer {layer_idx} {pname}", + [rank_params[pname] for rank_params, _ in per_rank], + partition, + problems, + allow_per_layer=True, + ) + categories.setdefault(category, set()).add(pname) + if full is not None: + params[pname] = full + biases = [bias for _, bias in per_rank] + bias = None + if all(report is None for report in biases): + categories.setdefault("absent", set()).add("e_score_correction_bias") + elif any(report is None for report in biases): + problems.append(f"layer {layer_idx}: only some ranks have an e_score_correction_bias") + else: + category, bias = _merge_shards( + f"layer {layer_idx} e_score_correction_bias", + biases, + partition, + problems, + allow_per_layer=False, + ) + categories.setdefault(category, set()).add("e_score_correction_bias") + expert_scalars[layer_idx] = _LayerScalars(params, bias) + if problems: + raise RuntimeError( + f"[DWDP Setup] The DWDP group's small expert parameters do not fit together: " + f"{'; '.join(problems[:8])}" + + (f"; and {len(problems) - 8} more" if len(problems) > 8 else "") + + ". Every rank of the DWDP group raises this." + ) + summary = "; ".join(f"{kind}: {sorted(names)}" for kind, names in sorted(categories.items())) + logger.info( + f"[DWDP Setup] Exchanged the small expert parameters of {len(layer_indices)} MoE layers " + f"in one allgather (partition {partition.describe()}): {summary}" + ) + return expert_scalars + + def fixup_moe_backends( model: nn.Module, moe_layer_indices: List[int], num_experts_total: int, - comm, - dwdp_rank: int, - dwdp_size: int, - *, - num_experts_per_worker: int, - peer_ranges: PeerRanges, + expert_scalars: Dict[int, _LayerScalars], ) -> None: """Patch each MoE backend to see all experts (ep_size=1) after DWDP setup. @@ -660,28 +1141,18 @@ def fixup_moe_backends( - ``initial_global_assignments = list(range(num_experts))``. - ``num_slots = num_experts``. - Additionally, small EP-sharded scale parameters - (``fc31_alpha`` / ``fc2_alpha``) are allgathered to full size. The - gate's ``e_score_correction_bias`` is checked too, but on the - Phase-1 ``moe_ep_size = 1`` setup it is already loaded full and the - helper short-circuits. + Additionally, it installs the full-sized small parameters that + ``_exchange_expert_scalars`` put together from every rank's shard + (``fc31_alpha`` / ``fc2_alpha``; the gate's ``e_score_correction_bias`` + only if it was loaded per expert: on the ``moe_ep_size = 1`` setup the + gate loads it whole) and rebuilds the quant scales. Local only: the + exchange was the setup's last collective. Args: model: Top-level CausalLM model. moe_layer_indices: MoE layer indices to patch. num_experts_total: Total number of experts across all ranks. - comm: DWDP MPI communicator used for small tensor allgathers. - dwdp_rank: This rank's DWDP rank. - dwdp_size: Total DWDP ranks. - num_experts_per_worker: Storage chunk size — uniform across ranks. - Used as the EP shard size for scale-param detection (the scale - params loaded by the fused MoE backend match - ``self.expert_size_per_partition`` set by - ``_init_dwdp_expert_layout``). - peer_ranges: Per-peer ``(local_start, local_end_capped)`` tuples; - used to scatter each peer's allgathered shard back into the - correct slice of the reconstructed full-size tensor (handles - tail-rank padding and redundant overlap). + expert_scalars: ``_exchange_expert_scalars``'s result, per layer index. """ decoder_model = _get_decoder_model(model) layers = decoder_model.layers @@ -746,253 +1217,40 @@ def fixup_moe_backends( f"slots=[{old_slot_start},{old_slot_end})->[0,{num_experts_total})" ) - # --- Allgather e_score_correction_bias --- - gate_module = _get_gate_module(moe_module) - if gate_module is not None: - _allgather_e_score_correction_bias( - gate_module, layer_idx, dwdp_rank, dwdp_size, - num_experts_total, comm, - num_experts_per_worker=num_experts_per_worker, - peer_ranges=peer_ranges, - ) - - # --- Allgather small expert scale/dequant parameters via MPI --- - # Large scale params (w3_w1_weight_scale, w2_weight_scale) are already - # handled by Transport (MNNVL + P2P) — they are in _EXPERT_WEIGHT_NAMES. - # Small params (fc31_alpha, fc2_alpha) are allgathered here. - # ``num_experts_per_worker`` (uniform storage chunk size across ranks) - # is the right shape match for these EP-sharded scale params under the - # Phase-1 layout, even when the partition is non-uniform or has - # redundancy. - _allgather_expert_scales( - experts_module, layer_idx, dwdp_rank, dwdp_size, comm, - experts_per_rank=num_experts_per_worker, - num_experts_total=num_experts_total, - peer_ranges=peer_ranges, - ) + layer_scalars = expert_scalars.get(layer_idx) + if layer_scalars is not None: + # --- The gate's e_score_correction_bias, if it was per expert --- + gate_module = _get_gate_module(moe_module) + if layer_scalars.bias is not None and gate_module is not None: + bias_param = gate_module.e_score_correction_bias + bias_data = bias_param.data if isinstance(bias_param, nn.Parameter) else bias_param + full_bias = layer_scalars.bias.to(bias_data.device) + if isinstance(bias_param, nn.Parameter): + gate_module.e_score_correction_bias = nn.Parameter( + full_bias, requires_grad=False + ) + else: + bias_param.data = full_bias + + # --- Small per-expert scale/dequant parameters --- + # Large scale params (w3_w1_weight_scale, w2_weight_scale) are + # handled by Transport (MNNVL + P2P) — they are in + # _EXPERT_WEIGHT_NAMES. The small ones (fc31_alpha, fc2_alpha) + # replace the rank's shard with the full-sized tensor here. + for pname, full in layer_scalars.params.items(): + device = getattr(experts_module, pname).device + setattr( + experts_module, + pname, + nn.Parameter(full.to(device), requires_grad=False), + ) + logger.debug( + f"[DWDP Setup] Layer {layer_idx}: installed the full {pname} " + f"({full.shape[0]} experts)" + ) _rebuild_quant_scales(experts_module, layer_idx) -def _allgather_e_score_correction_bias( - gate_module: nn.Module, - layer_idx: int, - dwdp_rank: int, - dwdp_size: int, - num_experts_total: int, - comm, - *, - num_experts_per_worker: Optional[int] = None, - peer_ranges: Optional[PeerRanges] = None, -) -> None: - """Allgather the EP-sharded e_score_correction_bias to get the full vector. - - Each rank has a shard of size ``num_experts_per_worker`` (the uniform - storage chunk set by ``_init_dwdp_expert_layout``). After allgather - each rank holds the full ``(num_experts,)`` bias vector reconstructed - via ``peer_ranges`` (so non-uniform tail padding and redundant - overlapping ranges are placed at the correct global expert indices). - - If the bias is already full-sized (not EP-sharded), this is a no-op. - The gate's e_score_correction_bias covers ALL experts for routing and - is typically loaded full when ``moe_ep_size = 1``. - - Args: - gate_module: The DeepseekV3Gate (or compatible) module with - ``e_score_correction_bias`` parameter. - layer_idx: Layer index (for log messages only). - dwdp_rank: This rank's DWDP rank. - dwdp_size: Total DWDP ranks. - num_experts_total: Total number of experts across all ranks. - comm: DWDP MPI communicator. - num_experts_per_worker: Uniform storage shard size used to detect - EP-sharded bias. Defaults to ``num_experts_total // dwdp_size`` - (the uniform-partition shard size). - peer_ranges: ``(local_start, local_end_capped)`` per peer; used - to scatter shards back to the right global expert indices. - When ``None`` the function falls back to ``torch.cat`` and - truncates to ``num_experts_total`` (uniform / non-divisible - tail-padding case). Required for redundancy. - """ - bias_param = getattr(gate_module, "e_score_correction_bias", None) - if bias_param is None: - logger.debug( - f"[DWDP Setup] Layer {layer_idx}: no e_score_correction_bias found" - ) - return - - bias_data = bias_param.data if isinstance(bias_param, nn.Parameter) else bias_param - bias_size = bias_data.shape[0] - expected_shard_size = ( - num_experts_per_worker if num_experts_per_worker is not None - else num_experts_total // dwdp_size - ) - - logger.info( - f"[DWDP Setup] Layer {layer_idx}: e_score_correction_bias " - f"type={type(bias_param).__name__}, shape={bias_data.shape}, " - f"device={bias_data.device}, " - f"num_experts_total={num_experts_total}, " - f"expected_shard={expected_shard_size}" - ) - - # The gate's e_score_correction_bias covers ALL experts for routing. - # It may or may not be EP-sharded depending on the model. - # Only allgather if the bias is actually sharded (size matches shard). - if bias_size == num_experts_total: - logger.info( - f"[DWDP Setup] Layer {layer_idx}: e_score_correction_bias " - f"already full-sized ({bias_size}), skipping allgather" - ) - return - if bias_size != expected_shard_size: - logger.warning( - f"[DWDP Setup] Layer {layer_idx}: e_score_correction_bias " - f"size={bias_size} doesn't match expected shard size " - f"({expected_shard_size}) or full size ({num_experts_total}), " - f"skipping allgather" - ) - return - - # Allgather local shards via MPI. Move to CPU first so the allgather - # doesn't require CUDA-aware MPI. - local_shard = bias_data.cpu().contiguous() - all_shards = comm.allgather(local_shard) - if peer_ranges is not None: - full_bias = _scatter_shards_to_full( - shards=all_shards, - peer_ranges=peer_ranges, - num_experts_total=num_experts_total, - ref=bias_data, - ) - else: - # Legacy uniform path: cat then truncate to handle non-divisible - # tail padding. Sufficient for the no-redundancy uniform case. - full_bias = torch.cat(all_shards, dim=0).to(bias_data.device) - if full_bias.shape[0] > num_experts_total: - full_bias = full_bias[:num_experts_total].contiguous() - - if isinstance(bias_param, nn.Parameter): - # Re-create parameter with the full size - gate_module.e_score_correction_bias = nn.Parameter( - full_bias, requires_grad=False - ) - else: - bias_param.data = full_bias - - logger.info( - f"[DWDP Setup] Layer {layer_idx}: allgathered e_score_correction_bias " - f"({expected_shard_size} -> {full_bias.shape[0]})" - ) - - -def _allgather_expert_scales( - experts_module: nn.Module, - layer_idx: int, - dwdp_rank: int, - dwdp_size: int, - comm, - experts_per_rank: Optional[int] = None, - *, - num_experts_total: Optional[int] = None, - peer_ranges: Optional[PeerRanges] = None, -) -> None: - """Allgather small EP-sharded quantization scale/dequant parameters via MPI. - - Only handles SMALL scale params (e.g. ``fc31_alpha``, ``fc2_alpha``) - that are NOT in ``_EXPERT_WEIGHT_NAMES``. Large scale params - (``w3_w1_weight_scale``, ``w2_weight_scale``) are handled by DWDP - Transport (MNNVL + P2P). - - Heuristic: any ``nn.Parameter`` on ``experts_module`` whose name contains - ``"scale"``, ``"scaling_factor"``, ``"dequant"``, or ``"alpha"``, whose - ``shape[0] == experts_per_rank``, and which is NOT in - ``_EXPERT_WEIGHT_NAMES``, is treated as a small EP-sharded tensor. - - When ``peer_ranges`` is provided each shard is scattered into the - correct global expert slice — this is required for non-uniform / - redundant partitions, where naive ``torch.cat`` would mis-place - overlapping ranges or include tail-rank padding. When ``peer_ranges`` - is ``None`` the function falls back to the simple uniform-partition - concat (legacy behavior). - - Args: - experts_module: The MoE backend module (e.g. CutlassFusedMoE). - layer_idx: Layer index (for log messages only). - dwdp_rank: This rank's DWDP rank. - dwdp_size: Total DWDP ranks. - comm: DWDP MPI communicator. - experts_per_rank: Number of experts per DWDP rank — i.e., the - uniform storage chunk size. If None, inferred from - ``w3_w1_weight.shape[0]`` (only valid before DWDP weight - buffer replaces the original tensor). - num_experts_total: Global expert count. Required when - ``peer_ranges`` is provided (for the destination tensor size). - peer_ranges: Per-peer ``(local_start, local_end_capped)``; when - provided enables scatter-based reconstruction that handles - non-uniform / redundant partitions. - """ - _SCALE_KEYWORDS = ("scale", "scaling_factor", "dequant", "alpha") - - # Discover EP-sharded scale parameters. - if experts_per_rank is None: - # Fallback: infer from main weight (only valid before DWDP buffer swap). - ref_param = getattr(experts_module, "w3_w1_weight", None) - if ref_param is None: - return - ref_data = ref_param.data if isinstance(ref_param, nn.Parameter) else ref_param - experts_per_rank = ref_data.shape[0] - if experts_per_rank == 0: - return - - for pname, param in list(experts_module.named_parameters()): - # Only look at direct parameters (no '.' in name — skip sub-modules). - if "." in pname: - continue - if not any(kw in pname for kw in _SCALE_KEYWORDS): - continue - # Large scale params in _EXPERT_WEIGHT_NAMES are handled by Transport - # (MNNVL + P2P), not MPI allgather. - if pname in _EXPERT_WEIGHT_NAMES: - continue - - pdata = param.data - if pdata.ndim == 0 or pdata.shape[0] != experts_per_rank: - # Scalar or non-EP-sharded — skip. - continue - - local_shard = pdata.cpu().contiguous() - logger.info( - f"[DWDP Setup] Layer {layer_idx}: allgathering {pname} " - f"shape={local_shard.shape} dtype={local_shard.dtype} via MPI" - ) - - all_shards = comm.allgather(local_shard) - if peer_ranges is not None: - assert num_experts_total is not None, ( - "num_experts_total must be provided alongside peer_ranges" - ) - full_tensor = _scatter_shards_to_full( - shards=all_shards, - peer_ranges=peer_ranges, - num_experts_total=num_experts_total, - ref=pdata, - ) - else: - full_tensor = torch.cat(all_shards, dim=0).to(pdata.device) - - # Replace parameter with full-sized version. - setattr( - experts_module, - pname, - nn.Parameter(full_tensor, requires_grad=False), - ) - - logger.info( - f"[DWDP Setup] Layer {layer_idx}: allgathered {pname} " - f"({experts_per_rank} -> {full_tensor.shape[0]})" - ) - - def _scatter_shards_to_full( *, shards: List[torch.Tensor], @@ -1003,15 +1261,16 @@ def _scatter_shards_to_full( """Reconstruct a full ``(num_experts_total, *trailing)`` tensor from per-peer shards using ``peer_ranges`` as the index map. - Each peer's shard has shape ``(num_experts_per_worker, *trailing)`` - (uniform across ranks). The valid prefix - ``shard[:end_capped - start]`` is placed at ``full[start:end_capped]``. - For redundancy (overlapping ranges) writes overwrite — values agree - across peers so the result is deterministic. + Each peer's shard has shape ``(count, *trailing)``, where ``count`` is + at least ``end - start`` of the peer's range; counts may differ between + peers (balanced partition, capped last rank). The prefix + ``shard[:end - start]`` is placed at ``full[start:end]``. For redundancy + (overlapping ranges) writes overwrite — ``_exchange_expert_scalars`` + checks that the values agree, so the result is deterministic. Args: shards: Per-peer tensors as returned by ``comm.allgather``. - peer_ranges: Per-peer ``(local_start, local_end_capped)``. + peer_ranges: Per-peer ``(start, end)``. num_experts_total: Global expert count (= full tensor's first dim). ref: Reference tensor for dtype + device. """ @@ -1045,9 +1304,10 @@ def _rebuild_quant_scales(experts_module: nn.Module, layer_idx: int) -> None: still exist. At runtime, wait_and_bind() does param.data = composite_VA, which swaps in the full tensor. The NamedTuple references the same Parameter, so it sees the updated data. - - fc31_alpha / fc2_alpha: allgathered via MPI. _allgather_expert_scales - replaces these with new full-sized Parameter objects via setattr(), so the - old NamedTuple references are stale and must be rebuilt. + - fc31_alpha / fc2_alpha: put together from every rank's shard by + _exchange_expert_scalars; fixup_moe_backends replaces these with new + full-sized Parameter objects via setattr(), so the old NamedTuple + references are stale and must be rebuilt. This function reconstructs the NamedTuple from the module's current Parameter state so all fields point to the live objects. @@ -1067,10 +1327,10 @@ def _rebuild_quant_scales(experts_module: nn.Module, layer_idx: int) -> None: # FusedMoEQuantScalesNVFP4: # fc1_act_global = module.fc31_input_scale (scalar, not EP-sharded) # fc1_weight_block = module.w3_w1_weight_scale (Transport: empty now, filled at runtime) - # fc1_global = module.fc31_alpha (MPI allgathered: full-sized) + # fc1_global = module.fc31_alpha (exchanged: full-sized) # fc2_act_global = module.fc2_input_scale (scalar, not EP-sharded) # fc2_weight_block = module.w2_weight_scale (Transport: empty now, filled at runtime) - # fc2_global = module.fc2_alpha (MPI allgathered: full-sized) + # fc2_global = module.fc2_alpha (exchanged: full-sized) new_qs = qs_type( fc1_act_global=experts_module.fc31_input_scale, fc1_weight_block=experts_module.w3_w1_weight_scale, @@ -1082,7 +1342,7 @@ def _rebuild_quant_scales(experts_module: nn.Module, layer_idx: int) -> None: experts_module.quant_scales = new_qs logger.info( f"[DWDP Setup] Layer {layer_idx}: rebuilt FusedMoEQuantScalesNVFP4 " - f"(weight_scale params via Transport, alpha params via MPI allgather)" + f"(weight_scale params via Transport, alpha params via the exchange)" ) else: logger.debug( @@ -1198,120 +1458,3 @@ def _get_gate_module(moe_module: nn.Module) -> Optional[nn.Module]: if gate is not None and hasattr(gate, "e_score_correction_bias"): return gate return None - - -def _get_first_spec(layer_weight_specs: LayerWeightSpecs) -> WeightSpec: - """Get the first WeightSpec from layer_weight_specs. - - Args: - layer_weight_specs: The weight specs dict. - - Returns: - First WeightSpec found. - - Raises: - ValueError: If specs are empty. - """ - for layer_idx, weight_specs in layer_weight_specs.items(): - for name, spec in weight_specs.items(): - return spec - raise ValueError("layer_weight_specs is empty") - - -def _validate_partition_config( - *, - num_experts_per_worker: int, - num_prefetch_experts: int, - num_experts_total: int, - dwdp_size: int, - loaded_local_experts: int, -) -> None: - """Validate that the DwdpConfig partition is internally consistent. - - Checks (in order, each independent): - - 1. Both ``num_experts_per_worker`` and ``num_prefetch_experts`` are - positive integers. - 2. ``num_prefetch_experts <= num_experts_per_worker`` — stride must - not exceed range size, otherwise consecutive ranks leave a gap - (some experts owned by no one). Equality is the no-redundancy - case; ``stride < size`` is intentional redundancy where adjacent - ranks' ranges overlap by ``size - stride`` experts. - 3. Strict equality ``(dwdp_size - 1) * stride + size == num_experts``. - The last rank's storage must end exactly at ``num_experts``: - - * ``< num_experts``: insufficient coverage — some experts owned - by no rank. - * ``> num_experts``: tail padding — last rank's storage extends - past the final expert, which the GB200 cuMemMap-with-fabric- - handle ABI cannot map into a ``num_experts``-sized composite - VA (partial mapping returns ``CUDA_ERROR_NOT_SUPPORTED``). - Pick a (size, stride) pair satisfying equality instead, e.g. - dwdp=3 + 256 experts → ``size=86, stride=85``; - dwdp=5 + 256 → ``52, 51``; - dwdp=7 + 256 → ``40, 36`` (or any other ``a*stride + b*size`` - summing to ``num_experts``). - 4. ``loaded_local_experts == num_experts_per_worker`` — the chunk - shape produced by the fused MoE weight loader must match what - DwdpConfig declares, so DWDP does not redirect prefetch to a - tensor of the wrong size. ``_init_dwdp_expert_layout`` arranges - this by overriding ``expert_size_per_partition`` to - ``num_experts_per_worker`` before ``create_weights`` runs. - - Together these guarantee that every expert in - ``[0, num_experts_total)`` is owned by at least one rank, that the - runtime range computation matches the stored weights, and that - every storage slot maps to a valid (in-range) expert id. - """ - if num_experts_per_worker <= 0: - raise ValueError( - f"[DWDP Setup] num_experts_per_worker must be positive, " - f"got {num_experts_per_worker}" - ) - if num_prefetch_experts <= 0: - raise ValueError( - f"[DWDP Setup] num_prefetch_experts must be positive, " - f"got {num_prefetch_experts}" - ) - if num_prefetch_experts > num_experts_per_worker: - raise ValueError( - f"[DWDP Setup] num_prefetch_experts ({num_prefetch_experts}) " - f"must be <= num_experts_per_worker ({num_experts_per_worker}); " - f"a stride larger than the range size leaves expert gaps " - f"between consecutive ranks." - ) - coverage = (dwdp_size - 1) * num_prefetch_experts + num_experts_per_worker - if coverage < num_experts_total: - raise ValueError( - f"[DWDP Setup] DWDP partition does not cover all experts: " - f"(dwdp_size - 1) * num_prefetch_experts + num_experts_per_worker " - f"= ({dwdp_size} - 1) * {num_prefetch_experts} + " - f"{num_experts_per_worker} = {coverage} < num_experts_total " - f"= {num_experts_total}. Increase num_prefetch_experts or " - f"num_experts_per_worker so coverage >= num_experts." - ) - if coverage > num_experts_total: - raise ValueError( - f"[DWDP Setup] DWDP partition exceeds num_experts: " - f"(dwdp_size - 1) * num_prefetch_experts + num_experts_per_worker " - f"= ({dwdp_size} - 1) * {num_prefetch_experts} + " - f"{num_experts_per_worker} = {coverage} > num_experts_total " - f"= {num_experts_total}. The last rank's storage extends " - f"past the final expert; GB200 fabric handles cannot be " - f"partially mapped into the composite VA, so tail padding " - f"is unsupported. Pick a (size, stride) pair with " - f"(dwdp_size - 1) * stride + size == num_experts (Mode B " - f"overlap), e.g. dwdp=3 + 256 -> size=86 stride=85; " - f"dwdp=5 + 256 -> 52 / 51; dwdp=7 + 256 -> 40 / 36." - ) - if loaded_local_experts != num_experts_per_worker: - raise ValueError( - f"[DWDP Setup] DwdpConfig.num_experts_per_worker=" - f"{num_experts_per_worker} does not match the fused MoE chunk " - f"shape ({loaded_local_experts}). _init_dwdp_expert_layout is " - f"expected to set expert_size_per_partition = " - f"num_experts_per_worker before create_weights() runs; check " - f"that get_global_dwdp_manager() is set during model " - f"construction (it is set by py_executor_creator before " - f"model loading)." - ) diff --git a/tensorrt_llm/_torch/modules/dwdp/specs.py b/tensorrt_llm/_torch/modules/dwdp/specs.py index dd56862dc64c..03fd509acfe4 100644 --- a/tensorrt_llm/_torch/modules/dwdp/specs.py +++ b/tensorrt_llm/_torch/modules/dwdp/specs.py @@ -19,60 +19,33 @@ from dataclasses import dataclass from functools import cached_property -from typing import Dict, List, Tuple +from typing import Dict, List, Literal, Sequence, Tuple import torch -# (local_start, local_end_capped) for one peer DWDP rank. ``local_start`` -# is ``peer_rank * num_prefetch_experts`` (the rank-to-rank stride). -# ``local_end_capped`` is the *valid* upper bound — i.e., -# ``min(local_start + num_experts_per_worker, num_experts_total)``; the -# tail rank's storage may extend past ``num_experts_total`` but the cap -# truncates the valid range. Adjacent ranks' valid ranges may overlap -# when ``num_prefetch_experts < num_experts_per_worker`` (redundancy mode). +# ``(start, end)``: one DWDP rank stores the routed experts ``[start, end)`` in rows +# ``0 .. end - start - 1`` of its expert weights. ``ExpertPartition.ranges`` holds one per rank; +# adjacent ranges overlap only in a redundant partition. PeerRange = Tuple[int, int] # Indexed by peer DWDP rank; ``len(PeerRanges) == dwdp_size``. PeerRanges = List[PeerRange] - -def compute_peer_ranges( - *, - dwdp_size: int, - num_experts_per_worker: int, - num_prefetch_experts: int, - num_experts_total: int, -) -> PeerRanges: - """Compute every peer rank's valid expert range. - - Each peer's storage holds ``num_experts_per_worker`` slots starting - at ``peer_rank * num_prefetch_experts``. The valid range is the - storage range capped at ``num_experts_total`` (the model's total - expert count), which truncates any padding on the tail rank. - - All ranks deterministically compute the same ``PeerRanges`` from the - shared ``DwdpConfig``, so no allgather is required. - """ - ranges: PeerRanges = [] - for peer_rank in range(dwdp_size): - start = peer_rank * num_prefetch_experts - end_capped = min(start + num_experts_per_worker, num_experts_total) - ranges.append((start, end_capped)) - return ranges +# How an ``ExpertPartition`` was made; see its docstring. +PartitionKind = Literal["balanced", "uniform", "redundant", "capped", "custom"] def lookup_owner(expert_id: int, peer_ranges: PeerRanges) -> int: - """Return the lowest peer rank whose valid range contains ``expert_id``. + """Return the lowest peer rank whose range contains ``expert_id``. - For uniform partition this is equivalent to - ``expert_id // (num_experts_total // dwdp_size)``. For redundancy + Without overlap that is the one rank storing it. For redundancy (overlapping ranges) the first-match policy picks the lowest-rank owner, matching the IPC implementation's behavior so reads of the same expert id are deterministic across peers. Raises: - ValueError: If no peer's valid range contains ``expert_id``. - This is impossible when the partition passes the coverage - check ``(dwdp_size - 1) * stride + size >= num_experts``. + ValueError: If no peer's range contains ``expert_id``. The ranges + of an ``ExpertPartition`` cover every expert, so this never + happens for them. """ for peer_rank, (start, end) in enumerate(peer_ranges): if start <= expert_id < end: @@ -80,6 +53,235 @@ def lookup_owner(expert_id: int, peer_ranges: PeerRanges) -> int: raise ValueError(f"expert_id={expert_id} not owned by any peer in peer_ranges={peer_ranges}") +_BALANCED_HINT = ( + "Leave num_experts_per_worker and num_prefetch_experts both at 0 for the balanced partition, " + "which stores every expert once on any dwdp_size." +) + + +@dataclass(frozen=True) +class ExpertPartition: + """Which routed experts each DWDP rank stores: one contiguous range per rank, every MoE layer. + + ``ranges[r] == (start, end)``: rank ``r`` stores experts ``[start, end)`` in rows + ``0 .. end - start - 1`` of its expert weights, allocates exactly those rows and exports them + in its handles. Every rank computes the same partition (``from_config``) from + ``DwdpConfig``, ``dwdp_size`` and the model's expert count; it is the only source of any + rank's range for ``DwdpManager``, ``MoE._init_dwdp_expert_layout``, ``setup_dwdp`` and the + transport. The ranges ascend, the first starts at 0, the last ends at ``num_experts`` and + none starts after the previous one ends, so every expert has an owner: ``lookup_owner`` + picks the lowest rank storing it. + + Kinds: + + * ``balanced`` (both config fields 0): with ``divmod(num_experts, dwdp_size) == (base, + rem)``, ranks below ``rem`` store ``base + 1`` experts and the others ``base``, as + ``_compute_ep_partition`` splits EP ranks. No expert is stored twice. + * ``uniform``, ``redundant`` and ``capped`` (``strided``, both fields set): rank ``r`` stores + ``num_experts_per_worker`` experts from ``r * num_prefetch_experts``, capped at + ``num_experts``. + * ``custom`` (``from_counts``): contiguous ranges of given counts, for tests and tools. + + Attributes: + num_experts: Routed experts of every MoE layer. + ranges: ``(start, end)`` per DWDP rank. + kind: How the ranges were made. + """ + + num_experts: int + ranges: Tuple[PeerRange, ...] + kind: PartitionKind + + def __post_init__(self) -> None: + if self.num_experts <= 0: + raise ValueError(f"[DWDP] num_experts must be positive, got {self.num_experts}") + if not self.ranges: + raise ValueError("[DWDP] an expert partition needs at least one rank") + prev_start, prev_end = 0, 0 + for rank, (start, end) in enumerate(self.ranges): + if not 0 <= start < end <= self.num_experts: + raise ValueError( + f"[DWDP] rank {rank} stores experts [{start}, {end}), which is empty or not " + f"within [0, {self.num_experts})" + ) + if rank == 0 and start != 0: + raise ValueError( + f"[DWDP] rank 0 stores experts [{start}, {end}); experts [0, {start}) have " + f"no owner" + ) + if rank > 0 and (not prev_start < start <= prev_end or end <= prev_end): + raise ValueError( + f"[DWDP] rank {rank} stores experts [{start}, {end}) after rank {rank - 1}'s " + f"[{prev_start}, {prev_end}); the ranges must ascend without a gap" + ) + prev_start, prev_end = start, end + if prev_end != self.num_experts: + raise ValueError( + f"[DWDP] the last rank's range ends at {prev_end}; experts " + f"[{prev_end}, {self.num_experts}) have no owner" + ) + + @classmethod + def balanced(cls, *, num_experts: int, dwdp_size: int) -> ExpertPartition: + """Floor or ceil of ``num_experts / dwdp_size`` per rank, the larger ones on the low ranks. + + Raises: + ValueError: ``num_experts < dwdp_size``: a rank would store no expert. + """ + if dwdp_size <= 0 or num_experts < dwdp_size: + raise ValueError( + f"[DWDP] cannot split {num_experts} routed experts over dwdp_size={dwdp_size} " + f"ranks: every rank must store at least one expert" + ) + base, rem = divmod(num_experts, dwdp_size) + ranges = tuple( + (rank * base + min(rank, rem), (rank + 1) * base + min(rank + 1, rem)) + for rank in range(dwdp_size) + ) + return cls(num_experts=num_experts, ranges=ranges, kind="balanced") + + @classmethod + def strided( + cls, *, num_experts: int, dwdp_size: int, size: int, stride: int + ) -> ExpertPartition: + """``size`` experts per rank, ``stride`` apart: the two strided ``DwdpConfig`` fields. + + ``size`` is ``num_experts_per_worker`` and ``stride`` is ``num_prefetch_experts``. Rank + ``r`` stores ``[r * stride, min(r * stride + size, num_experts))``. Accepted: + + * ``uniform``: ``stride == size`` and ``dwdp_size * size == num_experts``. + * ``redundant``: ``stride < size`` and ``(dwdp_size - 1) * stride + size == + num_experts``; adjacent ranks both store ``size - stride`` experts, which the lower + rank serves. + * ``capped``: ``stride == size`` and ``(dwdp_size - 1) * size < num_experts < dwdp_size * + size``; the last rank stores the ``num_experts - (dwdp_size - 1) * size`` experts left + and allocates only those rows. + + Raises: + ValueError: Any other combination: a value that is not positive, a stride larger + than the size (experts nobody stores), too few experts in total, overlapping + ranges past ``num_experts``, or a last rank left without experts. + """ + if size <= 0 or stride <= 0: + raise ValueError( + f"[DWDP] num_experts_per_worker ({size}) and num_prefetch_experts ({stride}) " + f"must both be positive in a strided expert partition. {_BALANCED_HINT}" + ) + if stride > size: + raise ValueError( + f"[DWDP] num_prefetch_experts ({stride}) must be <= num_experts_per_worker " + f"({size}): a stride larger than the range size leaves experts nobody stores. " + f"{_BALANCED_HINT}" + ) + coverage = (dwdp_size - 1) * stride + size + if coverage == num_experts: + kind = "uniform" if stride == size else "redundant" + elif stride == size and (dwdp_size - 1) * size < num_experts < coverage: + kind = "capped" + else: + if coverage < num_experts: + problem = f"so experts [{coverage}, {num_experts}) have no owner" + elif stride < size: + problem = f"past the {num_experts} routed experts, with overlapping ranges" + else: + problem = f"and the last rank would store none of the {num_experts} routed experts" + raise ValueError( + f"[DWDP] num_experts_per_worker={size}, num_prefetch_experts={stride} on " + f"dwdp_size={dwdp_size}: the last rank's range ends at (dwdp_size - 1) * " + f"num_prefetch_experts + num_experts_per_worker = {coverage}, {problem}. A strided " + f"partition must end at {num_experts}, or use equal values and cut the last " + f"range short at {num_experts}. {_BALANCED_HINT}" + ) + ranges = tuple( + (rank * stride, min(rank * stride + size, num_experts)) for rank in range(dwdp_size) + ) + return cls(num_experts=num_experts, ranges=ranges, kind=kind) + + @classmethod + def from_counts(cls, counts: Sequence[int]) -> ExpertPartition: + """Contiguous ranges of ``counts[r]`` experts each, from expert 0 (tests and tools).""" + ranges = [] + start = 0 + for count in counts: + ranges.append((start, start + count)) + start += count + return cls(num_experts=start, ranges=tuple(ranges), kind="custom") + + @classmethod + def from_config( + cls, + *, + num_experts: int, + dwdp_size: int, + num_experts_per_worker: int, + num_prefetch_experts: int, + ) -> ExpertPartition: + """The partition ``DwdpConfig`` asks for, for a model with ``num_experts`` routed experts. + + Both fields 0: ``balanced``; otherwise ``strided``, which needs both positive. + + Raises: + ValueError: The config gives no valid partition of ``num_experts`` experts over + ``dwdp_size`` ranks. + """ + if num_experts_per_worker == 0 and num_prefetch_experts == 0: + return cls.balanced(num_experts=num_experts, dwdp_size=dwdp_size) + return cls.strided( + num_experts=num_experts, + dwdp_size=dwdp_size, + size=num_experts_per_worker, + stride=num_prefetch_experts, + ) + + @property + def dwdp_size(self) -> int: + """Number of DWDP ranks.""" + return len(self.ranges) + + @property + def peer_ranges(self) -> PeerRanges: + """``ranges`` as the ``PeerRanges`` list ``lookup_owner`` and the copy plans take.""" + return list(self.ranges) + + def range_of(self, rank: int) -> PeerRange: + """``(start, end)`` of the experts ``rank`` stores.""" + return self.ranges[rank] + + def count_of(self, rank: int) -> int: + """Number of experts ``rank`` stores: the rows of its expert weights.""" + start, end = self.ranges[rank] + return end - start + + def describe(self) -> str: + """The kind and the per-rank counts, e.g. ``balanced, 8 x 75 + 4 x 74 of 896``. + + Overlapping ranges add the extra copies: ``redundant, 12 x 82 of 896, 88 duplicate + copies``. + """ + counts = [end - start for start, end in self.ranges] + groups = " + ".join( + f"{counts.count(count)} x {count}" for count in sorted(set(counts), reverse=True) + ) + duplicates = sum(counts) - self.num_experts + extra = f", {duplicates} duplicate copies" if duplicates else "" + return f"{self.kind}, {groups} of {self.num_experts}{extra}" + + +def handle_geometry(expert_bytes: int, start: int, end: int, granularity: int) -> Tuple[int, int]: + """``(phys_size, data_offset)`` of the handle that holds experts ``[start, end)`` of one weight. + + The handle spans the ``granularity`` pages the experts' bytes touch, and they start + ``data_offset`` bytes into it. A VMM handle maps only whole (offset 0 and its full size, + whatever its type), and an imported handle carries no size: the exporter allocates exactly + ``phys_size`` and every importer maps exactly ``phys_size``, computed from the exporter's + range. ``PageAlignedLayout.compute`` gives the same ``mnnvl_size`` and ``data_offset``. + """ + start_bytes = start * expert_bytes + page_start = _align_down(start_bytes, granularity) + page_end = _align_up(end * expert_bytes, granularity) + return page_end - page_start, start_bytes - page_start + + @dataclass(frozen=True) class WeightSpec: """Shape and dtype of one expert weight parameter. @@ -232,11 +434,12 @@ class PageAlignedLayout: that each pool-backed region is an exact multiple of pool page handles, avoiding cuMemMap overlap between pool pages and the MNNVL handle. - When ``pool_granularity > granularity``, a padding region appears between - the pool-backed pre-region and the MNNVL region (``pre_padding`` bytes) and - between the MNNVL region and the pool-backed post-region (``post_padding`` - bytes). The tensor view must start at ``va_base + pre_padding`` so that - expert indices line up with the correct physical memory. + When ``pool_granularity > granularity``, ``pre_size`` is ``page_start`` + rounded up: the first ``pre_padding`` bytes of the reservation (before + expert 0) are unused pool-backed space, and the tensor view must start at + ``va_base + pre_padding`` so that expert indices line up with the correct + physical memory. Likewise the last ``post_padding`` bytes of the + reservation are unused. The MNNVL region abuts both pool-backed regions. Attributes: expert_bytes: Bytes per single expert. @@ -253,8 +456,8 @@ class PageAlignedLayout: mnnvl_size: VA bytes for MNNVL region (multiple of granularity). post_size: VA bytes for pool-backed post-region (multiple of pool_granularity). - pre_padding: Extra bytes at end of pre-region due to pool_granularity - rounding (pre_size - page_start). + pre_padding: Unused bytes at the start of the pre-region (before + expert 0) due to pool_granularity rounding (pre_size - page_start). post_padding: Extra bytes at end of post-region due to pool_granularity rounding. data_offset: local_start_bytes - page_start (same as leading_edge). @@ -349,8 +552,9 @@ def compute( mnnvl_size = page_end - page_start # Validate: mnnvl_size must not exceed the handle's physical size. - # On GB200, cuMemMap with fabric handles requires size == phys_size - # (partial mapping returns CUDA_ERROR_NOT_SUPPORTED). + # cuMemMap maps a handle only whole (offset 0, size == phys_size; any + # handle type, else CUDA_ERROR_NOT_SUPPORTED), and the transport + # allocates it with handle_geometry(), whose size is mnnvl_size. if mnnvl_size > handle_phys_size: raise ValueError( f"mnnvl_size ({mnnvl_size}) exceeds handle_phys_size " diff --git a/tensorrt_llm/_torch/modules/dwdp/transport.py b/tensorrt_llm/_torch/modules/dwdp/transport.py index c0a797562720..986b378d7ea0 100644 --- a/tensorrt_llm/_torch/modules/dwdp/transport.py +++ b/tensorrt_llm/_torch/modules/dwdp/transport.py @@ -31,15 +31,17 @@ After copying data into the MNNVL handle the original parameter is freed immediately to avoid OOM on memory-constrained systems. - GB200 constraint: all handles use CU_MEM_HANDLE_TYPE_FABRIC. -- Peer views are immutable: after setup, MNNVL handles are never written again. - P2P reads are safe across processes without synchronization. +- Peer views are immutable: after setup the bytes a peer view covers (the + owner's own experts) are never written; the owner's prefetches rewrite only + the page-alignment edge bytes outside every view, so P2P reads are safe + across processes without synchronization. """ from __future__ import annotations import ctypes import os -from typing import Dict, List, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import torch @@ -54,10 +56,8 @@ logger.debug("[DWDP transport] Falling back to legacy `cuda` bindings") -from .specs import LayerWeightSpecs, MnnvlHandleSet, PeerRanges, compute_peer_ranges +from .specs import ExpertPartition, LayerWeightSpecs, MnnvlHandleSet, WeightSpec, handle_geometry from .vmm import ( - align_down, - align_up, check_cu_result, create_fabric_handle, free_va, @@ -138,6 +138,126 @@ def _import_handle(payload: Union[bytes, int], htype: cuda.CUmemAllocationHandle return int(imported_handle) +def _describe_error(exc: BaseException) -> str: + """A failure as the status a rank reports to its peers: picklable, one line.""" + return f"{type(exc).__name__}: {exc}" + + +def _allgather_keeping(comm, obj: object, error: Optional[Exception]) -> list: + """``comm.allgather(obj)`` that raises this rank's deferred ``error`` if the collective fails. + + A rank that failed enters the verdict collective anyway, with its failure as its status, + and raises ``error`` after it. If the collective fails as well (a peer died, the store + timed out), the rank raises ``error`` rather than the collective's failure, which becomes + ``error.__context__``: the rank's own failure is not lost. ``KeyboardInterrupt`` and + ``SystemExit`` pass through unchanged. + """ + try: + return comm.allgather(obj) + except Exception: + if error is None: + raise + raise error + + +def _export_local_chunk( + local_param: torch.Tensor, + spec: WeightSpec, + phys_size: int, + data_offset: int, + granularity: int, + device_id: int, + htype: cuda.CUmemAllocationHandleType, + local_handles: List[int], + local_export_fds: List[int], +) -> Tuple[int, Union[bytes, int]]: + """Move one local expert chunk into a new handle of ``phys_size`` bytes and export it. + + The handle and, on POSIX_FD, the exported FD are appended to ``local_handles`` and + ``local_export_fds`` as soon as they exist, so the caller's rollback releases them. + + Returns: + ``(handle, payload)``: the handle and its export (FABRIC bytes or a POSIX FD). + """ + # (a) Create the fabric handle + handle = create_fabric_handle(phys_size, device_id) + local_handles.append(handle) + + # (b) Reserve temp VA, map, set access, create tensor view + temp_va = reserve_va(phys_size, granularity) + mapped = False + try: + map_handle(temp_va, phys_size, handle, 0) + mapped = True + set_access(temp_va, phys_size, device_id) + + # Create a flat tensor view at data_offset within the + # mapped VA, sized exactly to the chunk. + data_ptr = temp_va + data_offset + handle_tensor = tensor_from_ptr( + ptr=data_ptr, + shape=spec.chunk_shape, + dtype=spec.dtype, + device_id=device_id, + ) + + # (c) Copy local param into the handle memory + handle_tensor.copy_(local_param) + torch.cuda.synchronize(device_id) + + # (d) Free the original param to reclaim memory + local_param.untyped_storage().resize_(0) + torch.cuda.empty_cache() + + # (e) Unmap temp VA (handle stays alive) + unmap_va(temp_va, phys_size) + mapped = False + finally: + if mapped: + try: + unmap_va(temp_va, phys_size) + except Exception as e: + logger.warning( + f"[DWDPTransport] temp VA unmap on error path failed (ignored): {e!r}" + ) + free_va(temp_va, phys_size) + + # (f) Export for the peers + exported = _export_handle(handle, htype) + if htype == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR: + # Track local FD so it can be closed at release time. + # It must remain open through Phase 2 so peers can + # pidfd_getfd it. + local_export_fds.append(int(exported)) + return handle, exported + + +def _check_peer_exports( + reports: List[Tuple[Optional[str], Union[bytes, int, None]]], *, layer_idx: int, name: str +) -> None: + """The group's verdict on one Phase-1 allgather, before any rank imports a handle of it. + + ``reports[p]`` is rank ``p``'s ``(status, payload)``: its error (None on success) and its + export. Every rank checks the same list, so either every rank goes on to import or every + rank raises the same error. The handles need no geometry check: every rank computes each + handle's size and data offset with ``handle_geometry`` from the partition, the granularity + and the expert shape and dtype that ``setup_dwdp``'s handshake agreed on. + + Raises: + RuntimeError: Listing every rank whose export failed. + """ + problems = [ + f"rank {peer_rank} failed to export its handle: {status}" + for peer_rank, (status, _payload) in enumerate(reports) + if status is not None + ] + if problems: + raise RuntimeError( + f"[DWDPTransport] Layer {layer_idx} {name}: {'; '.join(problems)}. No peer handle " + f"was imported; every rank of the DWDP group raises this." + ) + + class DWDPTransport: """MNNVL handle allocation and cross-process exchange for DWDP. @@ -164,7 +284,6 @@ class DWDPTransport: "_granularity", "_handle_set", "_peer_views", - "_peer_ranges", "_local_handles", "_imported_handles", "_peer_va_mappings", @@ -180,7 +299,6 @@ def __init__( granularity: int, handle_set: MnnvlHandleSet, peer_views: Dict[Tuple[int, int, str], torch.Tensor], - peer_ranges: PeerRanges, local_handles: List[int], imported_handles: List[int], peer_va_mappings: List[Tuple[int, int]], @@ -195,11 +313,6 @@ def __init__( granularity: VMM page granularity in bytes. handle_set: Local MNNVL handles for WeightBuffer. peer_views: Immutable tensor views into peer MNNVL memory. - peer_ranges: Per-peer ``(local_start, local_end_capped)`` tuples - indexed by DWDP rank. The capped end truncates the tail - rank's padding; consumers (lookup_owner, fill_edge_bytes, - weight_manager) read these to find the actual owner of a - given expert id. local_handles: Raw local handle integers (for cleanup). imported_handles: Raw imported handle integers (for cleanup). peer_va_mappings: List of (va, size) for peer VA regions (for cleanup). @@ -210,7 +323,6 @@ def __init__( self._granularity = granularity self._handle_set = handle_set self._peer_views = peer_views - self._peer_ranges = peer_ranges self._local_handles = local_handles self._imported_handles = imported_handles self._peer_va_mappings = peer_va_mappings @@ -224,36 +336,38 @@ def create( local_params: Dict[Tuple[int, str], torch.Tensor], comm, dwdp_rank: int, - dwdp_size: int, device_id: int, - local_start: int, - local_end: int, - num_experts_per_worker: int, - num_prefetch_experts: int, + partition: ExpertPartition, ) -> DWDPTransport: """Allocate MNNVL handles, populate them, and exchange with all peers. - This is the main entry point. Each of the ``dwdp_size`` processes calls - this method concurrently. The method: + This is the main entry point. Each of the ``partition.dwdp_size`` + processes calls this method concurrently. The method: 1. For each (layer_idx, weight_name) — one at a time to limit peak memory usage — every rank in the DWDP communicator: - a. Allocates a fabric handle large enough for the local chunk. + a. Allocates a fabric handle of ``handle_geometry`` of its range. b. Maps the handle to a temporary VA and copies the local parameter tensor into it at the correct ``data_offset``. c. Frees the original parameter tensor immediately. d. Unmaps the temporary VA (the physical handle persists). - e. Exports the handle to bytes and shares it with every peer via - a per-pair ``comm.allgather(handle_bytes)``. The allgather is - itself a synchronization point so no explicit barrier is needed - between Phase 1 iterations. + e. Exports the handle and shares it with every peer via a + per-pair ``comm.allgather((status, payload))``, which a rank + also enters after a-e failed, with the error as its status. + If any rank failed, every rank raises after this allgather + (``_check_peer_exports``), before any import. The allgather + is itself a synchronization point so no explicit barrier is + needed between Phase 1 iterations. 2. For each peer rank (skipping self), for each (layer_idx, weight_name): - a. Retrieves the peer's exported handle bytes from the Phase 1 + a. Retrieves the peer's exported payload from the Phase 1 allgather cache. - b. Imports the fabric handle. - c. Maps it to a VA and creates a tensor view for P2P reads. - 3. ``comm.Barrier()`` after all imports — ensures no rank starts - consuming peer views before every rank has finished importing. + b. Imports the handle. + c. Maps it whole, with the size of the peer's range, and creates + a tensor view of the peer's rows for P2P reads. + 3. ``comm.allgather(status)`` after all imports — ensures no rank + starts consuming peer views before every rank has finished + importing, and makes an import failure on any rank raise on + every rank. Args: layer_weight_specs: Per-layer weight specifications. Keys are layer @@ -262,19 +376,21 @@ def create( Each tensor has shape == spec.chunk_shape and is on the correct CUDA device. These tensors are consumed (freed) during the call. comm: mpi4py communicator scoped to the DWDP group (created by - DwdpManager). Must have exactly ``dwdp_size`` ranks. - dwdp_rank: This process's DWDP rank (0..dwdp_size-1). - dwdp_size: Total number of DWDP ranks in ``comm``. + DwdpManager). Must have exactly ``partition.dwdp_size`` ranks. + dwdp_rank: This process's DWDP rank (0..partition.dwdp_size-1). device_id: CUDA device ordinal. - local_start: First local expert index (inclusive). - local_end: Last local expert index (exclusive). + partition: The expert partition the DWDP group agreed on in + ``setup_dwdp``. A rank's range sizes the handle it exports, + which every peer maps whole, and the rows of the peers' + views of it. Returns: Ready-to-use DWDPTransport with handles and peer views. Raises: ValueError: If local_params keys do not match layer_weight_specs. - RuntimeError: If any CUDA operation fails. + RuntimeError: If any CUDA operation fails, on any rank: every rank + raises after the same allgather. """ granularity = get_allocation_granularity(device_id) htype = peer_handle_type() @@ -283,13 +399,13 @@ def create( ) # POSIX_FD path: dup peer FDs into our fd table via pidfd_getfd. - # We open one pidfd per peer ≠ self once up front and close them - # after Phase 3 Barrier. ``peer_pidfds`` is empty on the FABRIC - # path (aarch64 / GB200). + # Phase 2 opens one pidfd per peer ≠ self before importing its + # handles; they are closed after the Phase 3 allgather. + # ``peer_pidfds`` is empty on the FABRIC path (aarch64 / GB200). peer_pidfds: Dict[int, int] = {} # FDs returned by ``cuMemExportToShareableHandle`` on this rank. # Must outlive every peer's ``pidfd_getfd`` (Phase 2) but are - # closed right after the Phase 3 Barrier — at that point every + # closed right after the Phase 3 allgather — at that point every # peer has finished dup-ing them and the FDs no longer back # anything we need. ``cuMemImportFromShareableHandle`` dups the # peer-side FD internally, so we do NOT have to track imported @@ -299,23 +415,9 @@ def create( # the FABRIC path (aarch64 / GB200). local_export_fds: List[int] = [] - # Pick any spec to read the model's total expert count. All layers - # share the same ``num_experts`` (the gate-side global expert table). - first_spec = next(iter(next(iter(layer_weight_specs.values())).values())) - num_experts_total = first_spec.num_experts - - # Compute every peer's valid expert range deterministically from the - # shared DwdpConfig. Used by lookup_owner downstream - # (fill_edge_bytes, weight_manager.prefetch_layer, fixup-side - # allgathers) so the same first-match owner policy is applied - # everywhere — including non-uniform tail padding and redundancy - # where adjacent ranges overlap. - peer_ranges = compute_peer_ranges( - dwdp_size=dwdp_size, - num_experts_per_worker=num_experts_per_worker, - num_prefetch_experts=num_prefetch_experts, - num_experts_total=num_experts_total, - ) + dwdp_size = partition.dwdp_size + local_start, local_end = partition.range_of(dwdp_rank) + peer_ranges = partition.peer_ranges # Validate that every (layer_idx, name) in specs has a matching param expected_keys = set() @@ -341,16 +443,11 @@ def create( all_exports: Dict[Tuple[int, str], List[Union[bytes, int]]] = {} try: - # On POSIX_FD: open one pidfd per peer ≠ self before Phase 1. - if is_posix_fd: - all_pids = comm.allgather(os.getpid()) - for peer_rank, peer_pid in enumerate(all_pids): - if peer_rank == dwdp_rank: - continue - peer_pidfds[peer_rank] = _pidfd_open(peer_pid) + # On POSIX_FD: every rank's pid, for the pidfds Phase 2 opens. + all_pids = comm.allgather(os.getpid()) if is_posix_fd else [] # ---------------------------------------------------------- - # Phase 1: Allocate local handles and allgather bytes + # Phase 1: Allocate local handles and allgather exports # ---------------------------------------------------------- # Every rank iterates pairs in the same sorted order so that # per-pair comm.allgather calls match across ranks. @@ -359,76 +456,45 @@ def create( for name in sorted(weight_specs.keys()): spec = weight_specs[name] key = (layer_idx, name) - local_param = local_params[key] - - # Compute physical handle size using the SAME formula - # as PageAlignedLayout.compute() to guarantee that - # phys_size == layout.mnnvl_size. On GB200, fabric - # handles require cuMemMap with size == phys_size - # (partial mapping returns CUDA_ERROR_NOT_SUPPORTED). - local_start_bytes = local_start * spec.expert_bytes - local_end_bytes = local_end * spec.expert_bytes - page_start = align_down(local_start_bytes, granularity) - page_end = align_up(local_end_bytes, granularity) - phys_size = page_end - page_start - data_offset = local_start_bytes - page_start - - # (a) Create the fabric handle - handle = create_fabric_handle(phys_size, device_id) - local_handles.append(handle) - handles[key] = handle - sizes[key] = phys_size - - # (b) Reserve temp VA, map, set access, create tensor view - temp_va = reserve_va(phys_size, granularity) - mapped = False + + # The handle spans the pages of this rank's range: the + # handle_geometry() every importer computes for this + # rank, and PageAlignedLayout.mnnvl_size, since a handle + # maps only whole. + phys_size, data_offset = handle_geometry( + spec.expert_bytes, local_start, local_end, granularity + ) + export_error = None try: - map_handle(temp_va, phys_size, handle, 0) - mapped = True - set_access(temp_va, phys_size, device_id) - - # Create a flat tensor view at data_offset within the - # mapped VA, sized exactly to the chunk. - data_ptr = temp_va + data_offset - handle_tensor = tensor_from_ptr( - ptr=data_ptr, - shape=spec.chunk_shape, - dtype=spec.dtype, - device_id=device_id, + handle, exported = _export_local_chunk( + local_params[key], + spec, + phys_size, + data_offset, + granularity, + device_id, + htype, + local_handles, + local_export_fds, ) - - # (c) Copy local param into the handle memory - handle_tensor.copy_(local_param) - torch.cuda.synchronize(device_id) - - # (d) Free the original param to reclaim memory - local_param.untyped_storage().resize_(0) - torch.cuda.empty_cache() - - # (e) Unmap temp VA (handle stays alive) - unmap_va(temp_va, phys_size) - mapped = False - finally: - if mapped: - try: - unmap_va(temp_va, phys_size) - except Exception as e: - logger.warning( - f"[DWDPTransport] temp VA unmap on error " - f"path failed (ignored): {e!r}" - ) - free_va(temp_va, phys_size) - - # (f) Export and allgather to all peers - exported = _export_handle(handle, htype) - if is_posix_fd: - # Track local FD so it can be closed at release time. - # It must remain open through Phase 2 so peers can - # pidfd_getfd it. - local_export_fds.append(int(exported)) - # comm.allgather implicitly synchronizes all ranks on this + except Exception as exc: # reported to the peers below + export_error = exc + report = (_describe_error(exc), None) + else: + handles[key] = handle + sizes[key] = phys_size + report = (None, exported) + + # Entered by every rank, also after a local failure, so + # the verdict below is made from the same reports on every + # rank, before anything is imported. comm.allgather + # implicitly synchronizes all ranks on this # (layer_idx, name) — no explicit barrier required. - all_exports[key] = comm.allgather(exported) + reports = _allgather_keeping(comm, report, export_error) + if export_error is not None: + raise export_error + _check_peer_exports(reports, layer_idx=layer_idx, name=name) + all_exports[key] = [payload for _status, payload in reports] logger.debug( f"[DWDPTransport] Rank {dwdp_rank}: exported + allgathered " @@ -443,81 +509,125 @@ def create( # ---------------------------------------------------------- # Phase 2: Import peer handles and create tensor views # ---------------------------------------------------------- - for peer_rank in range(dwdp_size): - if peer_rank == dwdp_rank: - continue - - # Compute peer's expert range using the same config-driven - # formula every rank uses for itself (size = num_experts_per_worker, - # stride = num_prefetch_experts). All ranks share the same - # DwdpConfig, so peer ranges are deterministic — no allgather needed. - for layer_idx in sorted(layer_weight_specs.keys()): - weight_specs = layer_weight_specs[layer_idx] - for name in sorted(weight_specs.keys()): - spec = weight_specs[name] - - peer_local_start = peer_rank * num_prefetch_experts - peer_local_end = peer_local_start + num_experts_per_worker - - peer_start_bytes = peer_local_start * spec.expert_bytes - peer_end_bytes = peer_local_end * spec.expert_bytes - peer_page_start = align_down(peer_start_bytes, granularity) - peer_page_end = align_up(peer_end_bytes, granularity) - phys_size = peer_page_end - peer_page_start - peer_data_offset = peer_start_bytes - peer_page_start - - # (a) Retrieve peer handle payload from Phase 1 cache. - peer_payload = all_exports[(layer_idx, name)][peer_rank] - - # (b) Import the handle. On POSIX_FD we first dup - # the peer's FD into our fd table via pidfd_getfd; - # cuMemImportFromShareableHandle then accepts that - # local FD and dups it internally, so we close - # ours immediately after the import returns. On - # FABRIC the bytes blob is consumed directly. - if is_posix_fd: - local_fd = _pidfd_getfd(peer_pidfds[peer_rank], int(peer_payload)) + # Local work only: a failure becomes this rank's Phase 3 status. + import_error = None + try: + for peer_rank in range(dwdp_size): + if peer_rank == dwdp_rank: + continue + if is_posix_fd: + # The pidfd through which (b) dups the peer's FDs. + peer_pidfds[peer_rank] = _pidfd_open(all_pids[peer_rank]) + + # The peer's range in the agreed partition: its handle is + # handle_geometry() of that range, the size it allocated, + # and its view has that range's rows, which may differ + # from ours under the balanced partition or for a capped + # last rank. + peer_local_start, peer_local_end = peer_ranges[peer_rank] + for layer_idx in sorted(layer_weight_specs.keys()): + weight_specs = layer_weight_specs[layer_idx] + for name in sorted(weight_specs.keys()): + spec = weight_specs[name] + + phys_size, peer_data_offset = handle_geometry( + spec.expert_bytes, peer_local_start, peer_local_end, granularity + ) + + # (a) Retrieve peer handle payload from Phase 1 cache. + peer_payload = all_exports[(layer_idx, name)][peer_rank] + + # (b) Import the handle. On POSIX_FD we first dup + # the peer's FD into our fd table via pidfd_getfd; + # cuMemImportFromShareableHandle then accepts that + # local FD and dups it internally, so we close + # ours immediately after the import returns. On + # FABRIC the bytes blob is consumed directly. + if is_posix_fd: + local_fd = _pidfd_getfd(peer_pidfds[peer_rank], int(peer_payload)) + try: + imported_handle = _import_handle(local_fd, htype) + finally: + try: + os.close(local_fd) + except OSError as e: + logger.warning( + f"[DWDPTransport] Failed to close " + f"imported FD {local_fd}: {e}" + ) + else: + imported_handle = _import_handle(peer_payload, htype) + imported_handles.append(imported_handle) + + # (c) Reserve VA, map the whole handle, set access. On a + # failure the VA is unmapped and freed here, best effort, and + # the error is re-raised: the rollback releases only the + # complete mappings in peer_va_mappings. + peer_va = reserve_va(phys_size, granularity) + mapped = False try: - imported_handle = _import_handle(local_fd, htype) - finally: + map_handle(peer_va, phys_size, imported_handle, 0) + mapped = True + set_access(peer_va, phys_size, device_id) + except BaseException: + if mapped: + try: + unmap_va(peer_va, phys_size) + except Exception as e: + logger.warning( + f"[DWDPTransport] peer VA unmap on error path failed " + f"(ignored): {e!r}" + ) try: - os.close(local_fd) - except OSError as e: + free_va(peer_va, phys_size) + except Exception as e: logger.warning( - f"[DWDPTransport] Failed to close " - f"imported FD {local_fd}: {e}" + f"[DWDPTransport] peer VA free on error path failed " + f"(ignored): {e!r}" ) - else: - imported_handle = _import_handle(peer_payload, htype) - imported_handles.append(imported_handle) - - # (c) Reserve VA, map, set access - peer_va = reserve_va(phys_size, granularity) - map_handle(peer_va, phys_size, imported_handle, 0) - set_access(peer_va, phys_size, device_id) - peer_va_mappings.append((peer_va, phys_size)) - - peer_data_ptr = peer_va + peer_data_offset - peer_tensor = tensor_from_ptr( - ptr=peer_data_ptr, - shape=spec.chunk_shape, - dtype=spec.dtype, - device_id=device_id, - ) - - view_key = (peer_rank, layer_idx, name) - peer_views[view_key] = peer_tensor - - logger.debug( - f"[DWDPTransport] Rank {dwdp_rank}: imported " - f"peer={peer_rank} layer={layer_idx} name={name} " - f"peer_data_offset={peer_data_offset}" - ) + raise + peer_va_mappings.append((peer_va, phys_size)) + + peer_data_ptr = peer_va + peer_data_offset + peer_tensor = tensor_from_ptr( + ptr=peer_data_ptr, + shape=(peer_local_end - peer_local_start,) + + tuple(spec.chunk_shape[1:]), + dtype=spec.dtype, + device_id=device_id, + ) + + view_key = (peer_rank, layer_idx, name) + peer_views[view_key] = peer_tensor + + logger.debug( + f"[DWDPTransport] Rank {dwdp_rank}: imported " + f"peer={peer_rank} layer={layer_idx} name={name} " + f"peer_data_offset={peer_data_offset}" + ) + except Exception as exc: # reported to the peers in Phase 3 + import_error = exc # ---------------------------------------------------------- - # Phase 3: Barrier — all handles imported + # Phase 3: allgather(status) — all handles imported # ---------------------------------------------------------- - comm.Barrier() + # Every rank enters this allgather, also after a failed import, + # so it synchronizes the group and an import failure on any rank + # raises on every rank here. + import_status = None if import_error is None else _describe_error(import_error) + statuses = _allgather_keeping(comm, import_status, import_error) + if import_error is not None: + raise import_error + failed = [ + f"rank {rank}: {status}" + for rank, status in enumerate(statuses) + if status is not None + ] + if failed: + raise RuntimeError( + f"[DWDPTransport] Importing the peer handles failed on " + f"{'; '.join(failed)}. Every rank of the DWDP group raises this." + ) # POSIX_FD: close the per-peer pidfds — they're no longer # needed once every imported handle has been created. @@ -530,9 +640,9 @@ def create( # POSIX_FD: close our exported FDs. They had to stay open # through Phase 2 so peers could ``pidfd_getfd`` them, but - # the Barrier above guarantees every peer has finished doing - # so. Closing here (instead of at ``release()``) caps peak - # fd usage during setup at ~num_handles per rank. + # the Phase 3 allgather above guarantees every peer has + # finished doing so. Closing here (instead of at ``release()``) + # caps peak fd usage during setup at ~num_handles per rank. for fd in local_export_fds: try: os.close(fd) @@ -555,7 +665,6 @@ def create( granularity=granularity, handle_set=handle_set, peer_views=peer_views, - peer_ranges=peer_ranges, local_handles=local_handles, imported_handles=imported_handles, peer_va_mappings=peer_va_mappings, @@ -605,9 +714,11 @@ def get_handle_set(self) -> MnnvlHandleSet: def get_peer_views(self) -> Dict[Tuple[int, int, str], torch.Tensor]: """Return immutable P2P tensor views into every peer's MNNVL memory. - Keys are (peer_rank, layer_idx, weight_name). Each tensor has - shape == spec.chunk_shape and is read-only (writes would corrupt - the peer's data). + Keys are (peer_rank, layer_idx, weight_name). Each tensor holds the + peer's experts: ``end - start`` rows of its range in the expert + partition, which may differ from this rank's row count, and + ``spec.chunk_shape[1:]``. It is read-only (writes would corrupt the + peer's data). Returns: Dict mapping (peer_rank, layer_idx, name) to tensor views. @@ -619,22 +730,6 @@ def get_peer_views(self) -> Dict[Tuple[int, int, str], torch.Tensor]: raise RuntimeError("DWDPTransport has been released") return self._peer_views - def get_peer_ranges(self) -> PeerRanges: - """Return per-peer ``(local_start, local_end_capped)`` tuples. - - Used by ``fill_edge_bytes`` and the runtime ``WeightManager`` to - resolve the owner of any given expert id under non-uniform - partition or redundancy. The end is capped at ``num_experts`` - (i.e., reflects the *valid* expert range — not the storage range - which may include tail-padding slots). - - Raises: - RuntimeError: If transport has been released. - """ - if self._released: - raise RuntimeError("DWDPTransport has been released") - return self._peer_ranges - @property def dwdp_rank(self) -> int: """This process's DWDP rank.""" @@ -719,7 +814,7 @@ def _cleanup_resources( The POSIX_FD path releases its FDs eagerly during ``create()`` — imported FDs are closed inline right after ``cuMemImportFromShareableHandle`` (which dups them internally), and - local export FDs are closed right after the Phase 3 Barrier — so by + local export FDs are closed right after the Phase 3 allgather — so by the time normal-path ``release()`` runs ``local_export_fds`` is already empty. This function still iterates it because the error-path caller invokes us with whatever FDs were tracked at the diff --git a/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py b/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py index 42ac02afb167..74b8c3039b25 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py @@ -44,6 +44,7 @@ from .page_pool import PagePool, compute_slot_sizes from .specs import EdgeInfo, LayerWeightSpecs, MnnvlHandleSet, PageAlignedLayout from .vmm import ( + align_up, free_va, get_allocation_granularity, map_handle, @@ -53,6 +54,99 @@ unmap_va, ) +_GIB = 1 << 30 +_MIB = 1 << 20 + + +class PrefetchRingMemoryError(RuntimeError): + """The expert prefetch ring does not fit in the free device memory. + + The ring is the page pool of the remote prefetch slots (``DwdpConfig.num_prefetch_slots``, + two by default), each holding the remote experts of one MoE layer. ``WeightBuffer.create`` + raises this before it allocates any page. ``str()`` is the message; the attributes carry its + numbers. The message is the only positional argument, so the error pickles, as the + executor's report of a failed setup requires. + + Attributes: + needed_bytes: Bytes of every page of every slot. + free_bytes: Free device memory, read after emptying the caching allocator. + total_bytes: Device memory. + device_id: CUDA device ordinal. + slot_pages: Pages per slot. + page_size: Bytes per page. + """ + + def __init__( + self, + message: str, + *, + needed_bytes: int = 0, + free_bytes: int = 0, + total_bytes: int = 0, + device_id: int = -1, + slot_pages: Tuple[int, ...] = (), + page_size: int = 0, + ) -> None: + super().__init__(message) + self.needed_bytes = needed_bytes + self.free_bytes = free_bytes + self.total_bytes = total_bytes + self.device_id = device_id + self.slot_pages = tuple(slot_pages) + self.page_size = page_size + + +def _describe_ring(slot_pages: Tuple[int, ...], page_size: int) -> str: + """The ring's size and pages, e.g. ``27.22 GiB (871 + 871 pages of 16 MiB, ...)``.""" + per_slot = " + ".join(str(count) for count in slot_pages) + return ( + f"{sum(slot_pages) * page_size / _GIB:.2f} GiB ({per_slot} pages of " + f"{page_size / _MIB:g} MiB, the remote experts of {len(slot_pages)} MoE layers)" + ) + + +def _check_ring_fits(slot_sizes: List[int], page_size: int, device_id: int) -> Tuple[int, ...]: + """Refuse a page pool larger than the free device memory, before any ``cuMemCreate``. + + The pool is driver memory that blocks held by the PyTorch caching allocator cannot serve, so + those go back to the driver first (the transport ends with that too: a no-op in practice). + The pool takes whole pages per slot, as ``PagePool`` counts them. With exactly enough free + memory the pool is allocated as before; with less, every rank would otherwise die in + ``cuMemCreate`` partway through it. + + Returns: + The pages per slot. + + Raises: + PrefetchRingMemoryError: The pages of all slots need more than the free memory. + """ + slot_pages = tuple(align_up(size, page_size) // page_size for size in slot_sizes) + needed = sum(slot_pages) * page_size + torch.cuda.empty_cache() + free, total = torch.cuda.mem_get_info(device_id) + if free >= needed: + return slot_pages + # Two slots are the fewest the prefetch runs with. + fewer_slots = ( + ", fewer prefetch slots (DwdpConfig.num_prefetch_slots)" if len(slot_pages) > 2 else "" + ) + raise PrefetchRingMemoryError( + f"[DWDP Setup] The expert prefetch ring does not fit on GPU {device_id}: its " + f"{len(slot_pages)} slots need {_describe_ring(slot_pages, page_size)}, but " + f"{free / _GIB:.2f} GiB of {total / _GIB:.2f} GiB is free. Every rank holds its resident " + f"experts, the weights it replicates and this ring at once: use a larger dwdp_size (fewer " + f"resident experts per rank){fewer_slots} or leave more memory free before the setup. On " + f"a model with dense DWDP roots (Kimi-K3), dense DWDP (TRTLLM_DWDP_DENSE=1) releases the " + f"replicated dense weights before the ring is allocated, unless it is off or was abandoned " + f"(see the [DWDP Dense] lines above).", + needed_bytes=needed, + free_bytes=free, + total_bytes=total, + device_id=device_id, + slot_pages=slot_pages, + page_size=page_size, + ) + @dataclass class LayerBufferState: @@ -190,7 +284,8 @@ def create( Steps: 1. Compute page-aligned layout for each (layer, weight) pair (num_experts comes from WeightSpec.num_experts per layer) - 2. Allocate page pool for the remote prefetch slots + 2. Check that the page pool fits in the free device memory, then + allocate it for the remote prefetch slots 3. Create per-layer composite VA mappings 4. Create tensor views @@ -210,6 +305,8 @@ def create( Raises: ValueError: If parameters are inconsistent. + PrefetchRingMemoryError: The page pool does not fit in the free + device memory; raised before any page is allocated. """ # Validate inputs cls._validate_inputs(layer_weight_specs, handles, local_start, local_end, dwdp_size) @@ -234,6 +331,7 @@ def create( # Because layouts already have pre/post sizes aligned to # pool_page_size, the slot sizes are automatically aligned too. slot_sizes = compute_slot_sizes(layouts, buffer._slot_of_layer, num_slots) + slot_pages = _check_ring_fits(slot_sizes, buffer._pool_page_size, device_id) buffer._page_pool = PagePool.create( slot_sizes, device_id, page_size=buffer._pool_page_size @@ -245,7 +343,8 @@ def create( logger.info( f"[WeightBuffer] Created for {len(buffer._moe_layer_indices)} layers, " - f"local experts [{local_start}, {local_end}), DWDP size {dwdp_size}" + f"local experts [{local_start}, {local_end}), DWDP size {dwdp_size}; prefetch " + f"ring {_describe_ring(slot_pages, buffer._pool_page_size)}" ) return buffer @@ -320,20 +419,16 @@ def _setup_layer( [ pre_region | mnnvl_region | post_region ] (page pool) (MNNVL handle) (page pool) - - pre_region : page-pool fabric pages for remote experts [0, local_start) + - pre_region : page-pool pages for remote experts [0, local_start) - mnnvl_region: zero-copy mapping of the MNNVL fabric handle (contains local experts [local_start, local_end)) - - post_region : page-pool fabric pages for remote experts [local_end, N) - - Routing table budget (GB200): - ~928 total entries - - 696 used by Transport (3 peers × 4 weights × 58 layers, persistent) - - 232 used here (1 entry per weight per layer, persistent for lifetime) - = exactly 928 = 0 remaining. + - post_region : page-pool pages for remote experts [local_end, N) - We achieve 1 routing entry per weight by calling set_access ONCE on the - full composite VA (va_base, total_size). Calling it three times - (once per sub-region) would consume 696 entries here, exhausting the table. + The page-pool pages are local ``CU_MEM_HANDLE_TYPE_NONE`` handles + (``PagePool``); only the MNNVL region is a fabric handle. The three + regions are mapped back to back into one reservation, and one + set_access call covers the full composite VA (va_base, total_size) + instead of one call per region. """ weight_specs = self._layer_weight_specs[layer_idx] buf_slot = self.buffer_index_for_layer(layer_idx) @@ -391,10 +486,8 @@ def _setup_layer( page_pool_offset += layout.post_pages # --- Single set_access for the ENTIRE composite VA --- - # One cuMemSetAccess call = 1 routing table entry consumed. - # This is critical: calling it per sub-region would consume 3 - # entries per weight per layer (696 total) and exhaust the - # routing table on top of Transport's 696 persistent entries. + # One call covers the whole reservation instead of one call + # per region. set_access(va_base, layout.total_size, self._device_id) # --- Create full [num_experts, ...] tensor view --- @@ -405,7 +498,7 @@ def _setup_layer( # # When pool_granularity > granularity, pre_size may be larger # than page_start due to rounding. The difference is - # pre_padding bytes of unused pool-backed space at the end of + # pre_padding bytes of unused pool-backed space at the start of # the pre-region. The tensor view must start at # va_base + pre_padding so that expert index 0 maps to the # correct physical memory: @@ -561,25 +654,6 @@ def get_data_offset(self, layer_idx: int, name: str) -> int: """ return self.get_layout(layer_idx, name).data_offset - def compute_peer_data_offset(self, layer_idx: int, name: str, peer_local_start: int) -> int: - """Compute data_offset for a peer rank given their local_start. - - layer_idx is required because different layers may have - different expert_bytes, affecting page alignment. - - Args: - layer_idx: Layer index. - name: Weight name. - peer_local_start: The peer rank's local_start expert index. - - Returns: - Data offset in bytes for the peer. - """ - layout = self.get_layout(layer_idx, name) - peer_start_bytes = peer_local_start * layout.expert_bytes - peer_page_start = (peer_start_bytes // self._granularity) * self._granularity - return peer_start_bytes - peer_page_start - # --- Layout queries (diagnostics, tests) --- def get_layout(self, layer_idx: int, name: str) -> PageAlignedLayout: diff --git a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py index a918e2e2c7af..3ce78a52474a 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py @@ -129,10 +129,11 @@ def __init__( dwdp_size: int, contention_opt: bool = False, repeated_layer_indices: Sequence[int] = (), + copy_stream: Optional[torch.cuda.Stream] = None, ) -> None: """Initialize the DWDPWeightManager. - Creates a dedicated copy stream and pre-allocates the CUDA events of + Takes or creates the copy stream and pre-allocates the CUDA events of every buffer slot. The consume_events are recorded immediately so the first prefetch does not stall waiting for a never-recorded event. @@ -141,11 +142,12 @@ def __init__( peer_views: Mapping of (peer_rank, layer_idx, weight_name) to tensor views into peer's MNNVL handle. These are IMMUTABLE — the source GPU never writes to them during inference. - peer_ranges: Per-peer ``(local_start, local_end_capped)`` tuples - indexed by DWDP rank. ``prefetch_layer`` resolves the owner of - a remote expert id with ``lookup_owner(expert_id, peer_ranges)``, - which handles non-uniform partition (tail-rank padding) and - redundancy (overlapping ranges) uniformly. + peer_ranges: Every rank's ``(start, end)`` in the expert partition, + indexed by DWDP rank (``ExpertPartition.peer_ranges``); + peer ``p``'s views have ``end - start`` rows. ``prefetch_layer`` + resolves the owner of a remote expert id with + ``lookup_owner(expert_id, peer_ranges)``, the lowest-rank owner + under redundancy (overlapping ranges). moe_layer_indices: Sorted list of decoder layer indices that are MoE layers (e.g., layers 3..60 in a typical MoE model where the first few layers are dense). @@ -162,10 +164,15 @@ def __init__( repeated_layer_indices: MoE layers that run several times in a row within a forward pass; their repeats bind the copy still in the slot (see ``PrefetchSchedule``). + copy_stream: The stream every prefetch copy runs on; a new stream + on the buffer's device if None. With attached dense DWDP, + ``DwdpManager.setup`` passes the stream the dense copies + already run on. Raises: ValueError: If dwdp_rank is out of range, dwdp_size is invalid, - there are no MoE layers or a repeated layer is not one of them. + there are no MoE layers, a repeated layer is not one of them, + or copy_stream is on another device than the weight buffer. """ if dwdp_size <= 0: raise ValueError(f"dwdp_size must be positive, got {dwdp_size}") @@ -181,16 +188,23 @@ def __init__( self._dwdp_rank = dwdp_rank self._dwdp_size = dwdp_size - # Storage chunk size of every peer's MNNVL handle. Uniform across - # ranks (Phase 1 made ``num_experts_per_worker`` the storage size for - # every rank) — the variability lives in the *valid* sub-range - # tracked by ``_peer_ranges``. + # The number of experts this rank stores. Peers may store another + # number (e.g. one more or one fewer under the balanced partition): + # each peer's range is in ``_peer_ranges``, and its view has its rows. self._experts_per_rank = weight_buffer.local_end - weight_buffer.local_start # Dedicated CUDA stream for P2P copy operations. # Using a separate stream allows copy/compute overlap. device = torch.device("cuda", weight_buffer.device_id) - self._copy_stream = torch.cuda.Stream(device=device) + if copy_stream is None: + copy_stream = torch.cuda.Stream(device=device) + elif copy_stream.device != device: + raise ValueError( + f"copy_stream is on {copy_stream.device}, the weight buffer on {device}" + ) + # Attached dense DWDP, set up before this manager exists, already runs on a given + # stream: one copy stream, one FIFO, for the dense and the expert copies. + self._copy_stream = copy_stream num_slots = weight_buffer.num_slots self._layer_order = {layer: order for order, layer in enumerate(self._moe_layer_indices)} @@ -255,7 +269,8 @@ def __init__( f"[DWDPWeightManager] Initialized: rank={dwdp_rank}/{dwdp_size}, " f"moe_layers={len(self._moe_layer_indices)}, " f"weights={self._weight_names}, " - f"experts_per_rank={self._experts_per_rank}, " + f"local_experts={self._experts_per_rank} " + f"[{weight_buffer.local_start}, {weight_buffer.local_end}), " f"contention_opt={contention_opt}, " f"gather={self._use_gather}" ) @@ -275,6 +290,21 @@ def dwdp_size(self) -> int: """Total number of DWDP ranks.""" return self._dwdp_size + @property + def copy_stream(self) -> torch.cuda.Stream: + """The CUDA stream every DWDP prefetch copy runs on.""" + return self._copy_stream + + @property + def use_gather(self) -> bool: + """Whether prefetches use the SM fabric gather (``TRTLLM_DWDP_PREFETCH_GATHER``).""" + return self._use_gather + + @property + def use_batched_prefetch(self) -> bool: + """Whether prefetches use one batched copy-engine copy per layer (``contention_opt``).""" + return self._use_batched_prefetch + def is_moe_layer(self, layer_idx: int) -> bool: """Check if a given layer index is a MoE layer. diff --git a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index 999626e8e30f..647ef92562b8 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -103,6 +103,30 @@ protocol. `ConfigurableMoE` owns wait/bind/record placement so those calls happen once per layer rather than once per scheduler chunk. Schedulers must not perform DWDP lifecycle calls. +Dense DWDP (`TRTLLM_DWDP_DENSE`, which row-shards the replicated dense weights +of the MoE decoder layers the same way) is owned by `DenseDwdpManager` +(`modules/dwdp/dense/manager.py`), which the model's decoder-layer hooks reach +through `get_global_dense_dwdp_manager()`. With expert DWDP, `DwdpManager` +builds the dense manager and drives it: dense DWDP shards over the DWDP group, +and the hooks enqueue each layer's dense weights on the same copy stream two +layers ahead, and those of the first two layers of the next pass at the end of +a pass, as the expert prefetch does. `DwdpManager.start_forward` enqueues the +dense weights of the first two MoE layers before it starts the expert +prefetch, unless the previous pass already enqueued them; after an aborted +forward it leaves both to the model's forward entry +(`DenseDwdpManager.ensure_primed`). That entry does the same in every forward, +also those the executor loop does not start (warmup, KV-cache estimation); +after an aborted forward it first resets the dense slots, then the expert +prefetch restarts its sequence. Without `dwdp_config`, dense DWDP runs +standalone on an attention-DP instance (`tp_size > 1`): the executor creator +builds the dense manager, which shards over the attention-DP group on its own +copy stream. The non-overlap executor loop primes it before each forward, and +the model's forward entry primes any forward the loop did not, so the overlap +scheduler, whose loop does not prime, is allowed there. In both modes +`ConfigurableMoE` is unaware of dense DWDP: with expert DWDP it keeps the +wait/bind/finish protocol above unchanged; standalone it runs without DWDP +(expert-parallel experts, no wait/bind/finish). + The composite mapping exposes the full expert table through the backend's ordinary weight attributes, so DWDP returns `None` from `_create_comm_strategy_auto` and continues through the `EXTERNAL_COMM` diff --git a/tensorrt_llm/_torch/moe/fused_moe/interface.py b/tensorrt_llm/_torch/moe/fused_moe/interface.py index 3c14d8a174b6..fc37c57f2a7b 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/interface.py +++ b/tensorrt_llm/_torch/moe/fused_moe/interface.py @@ -379,50 +379,45 @@ def __init__( # Override expert layout when DWDP is enabled. This must run before # create_weights() (which is invoked later by ConfigurableMoE) so the - # fused MoE backend allocates ``num_experts_per_worker`` slots per - # rank — not ``num_experts // ep_size``. fixup_moe_backends() at - # setup_dwdp() time later promotes these to the full composite-VA - # view (``ep_size = 1``, ``slot_start = 0``, - # ``slot_end = num_experts``); the earlier override is still - # required because ``mapping.moe_ep_size`` is now 1, so the - # un-overridden default would be ``expert_size_per_partition = - # num_experts`` (each rank allocating storage for every expert). + # fused MoE backend allocates the rows of this rank's experts only: + # ``mapping.moe_ep_size`` is 1 under DWDP, so the default layout would + # allocate every expert on every rank. An invalid DwdpConfig raises + # here, on every rank alike, before any weight exists. + # fixup_moe_backends() at setup_dwdp() time later promotes the layout + # to the full composite-VA view (``ep_size = 1``, ``slot_start = 0``, + # ``slot_end = num_experts``). self._init_dwdp_expert_layout() self._init_perfect_router() def _init_dwdp_expert_layout(self): """Override expert layout when DWDP is enabled. - Plumbs ``num_experts_per_worker`` (storage size) and - ``start_expert_id`` (storage start) from the active - ``DwdpManager``. This is a no-op when DWDP is not enabled - (``get_global_dwdp_manager()`` returns ``None``). - - For the uniform partition case (``num_prefetch_experts == - num_experts_per_worker == num_experts // dwdp_size``) this is - mathematically equivalent to the legacy ``ep_size = dwdp_size`` - layout. It additionally enables: - - * Non-uniform partition (``dwdp_size`` does not divide - ``num_experts``): user picks ``num_prefetch_experts < num_experts_per_worker`` - so that ``(dwdp_size - 1) * num_prefetch_experts + num_experts_per_worker - == num_experts`` exactly. Adjacent ranks' valid ranges - overlap by ``num_experts_per_worker - num_prefetch_experts`` - experts; ``_validate_partition_config`` rejects - configurations whose last-rank end exceeds ``num_experts`` - because the GB200 cuMemMap-with-fabric-handle ABI requires - ``mnnvl_size == handle_phys_size`` (no partial mapping), so - tail-padded storage cannot be partially mapped into a - ``num_experts``-sized composite VA. - * Redundancy (``num_prefetch_experts < num_experts_per_worker`` - with ``(dwdp_size - 1) * stride + size == num_experts``): - consecutive ranks' ranges overlap; peer-side - ``lookup_owner`` (Phase 2) picks the lowest-rank owner so - reads of shared experts are deterministic. + Sizes this backend from this rank's range ``[start, end)`` of the + ``ExpertPartition`` the active ``DwdpManager`` computes for + ``num_experts`` (``DwdpManager.expert_partition``): it allocates + ``end - start`` rows, and the weight loaders put experts + ``start .. end - 1`` into them. Ranks may store different counts: + the balanced partition (``num_experts_per_worker == + num_prefetch_experts == 0``) stores the floor or the ceil of + ``num_experts / dwdp_size`` experts per rank, the larger counts on + the lower ranks, and a capped strided partition cuts the last range + short at ``num_experts``. No rank allocates a row for an expert it + does not store. ``setup_dwdp`` checks every rank's rows against the + same partition across the DWDP group. This is a no-op when DWDP is + not enabled (``get_global_dwdp_manager()`` returns ``None``). + + In a strided partition with ``num_prefetch_experts < + num_experts_per_worker`` adjacent ranges overlap (redundancy), and + ``lookup_owner`` picks the lowest-rank owner so reads of shared + experts are deterministic. DWDP and MoE EPLB are mutually exclusive — DWDP swaps ``param.data`` to a composite-VA tensor at runtime, which the EPLB rebalancer would clobber. + + Raises: + ValueError: The DwdpConfig gives no valid partition of + ``num_experts`` experts over the DWDP group. """ from tensorrt_llm._torch.pyexecutor.dwdp import get_global_dwdp_manager @@ -433,28 +428,17 @@ def _init_dwdp_expert_layout(self): "DWDP and EPLB (MoE load balancer) cannot be used together. " "Disable one of dwdp_config or moe_load_balancer.") + partition = dwdp_manager.expert_partition(self.num_experts) + start, end = partition.range_of(dwdp_manager.dwdp_rank) self.num_slots = self.num_experts - self.expert_size_per_partition = dwdp_manager.num_experts_per_worker - dwdp_size = dwdp_manager.dwdp_size - # Routing-side expert assignment: distribute ``num_experts`` round-robin - # across DWDP ranks. Independent of the storage layout — the gate uses - # this to map expert ids to ranks, while DWDP composite VA handles the - # actual weight access. - self.initial_global_assignments = [ - (ep_rank * self.num_experts // dwdp_size + local_slot_id) % - self.num_experts for ep_rank in range(dwdp_size) - for local_slot_id in range(self.expert_size_per_partition) - ] - # Storage range: ``[start, start + size)``. Phase 2's strict - # validation guarantees ``slot_end <= num_experts``, so every - # storage slot maps to a valid expert id. ``len(initial_local_expert_ids) - # == expert_size_per_partition`` is preserved, which is required - # by the per-slot weight-scale copy in - # ``quantization.load_quant_scales``. - self.slot_start = dwdp_manager.start_expert_id - self.slot_end = self.slot_start + self.expert_size_per_partition - self.initial_local_expert_ids = list( - range(self.slot_start, self.slot_end)) + self.expert_size_per_partition = end - start + # Only a load balancer reads it, and DWDP excludes one (above). + self.initial_global_assignments = list(range(self.num_experts)) + # Storage range: row ``i`` holds expert ``start + i``. The loaders + # enumerate ``initial_local_expert_ids`` to pick the experts to load. + self.slot_start = start + self.slot_end = end + self.initial_local_expert_ids = list(range(start, end)) def _get_perfect_router_dtype(self) -> torch.dtype: if self.routing_method.routing_method_type in ( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index d60b199ffbba..38af4aff44d7 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -50,6 +50,7 @@ from ..hostfunc import set_low_latency_dispatch from ..model_config import ModelConfig, _get_kv_cache_layer_specs from ..models.modeling_multimodal_mixin import MultimodalModelMixin +from ..modules.dwdp.dense.manager import DenseDwdpManager from ..nccl_window_tensor_scope import discard_nccl_window_tensor_outputs from ..speculative import (draft_prompt_lookahead, get_num_extra_kv_tokens, get_num_spec_layers, get_spec_decoder, @@ -4044,6 +4045,7 @@ def create_py_executor_instance( virtual_memory_pools: Optional[dict] = None, execution_stream: Optional[torch.cuda.Stream] = None, dwdp_manager: Optional[DwdpManager] = None, + dense_dwdp_manager: Optional[DenseDwdpManager] = None, max_num_sequences: Optional[int] = None, ) -> PyExecutor: set_low_latency_dispatch( @@ -4454,6 +4456,7 @@ def create_py_executor_instance( execution_stream=execution_stream, waiting_queue_policy=waiting_queue_policy, dwdp_manager=dwdp_manager, + dense_dwdp_manager=dense_dwdp_manager, enable_kv_pool_rebalance=llm_args.kv_cache_config. enable_kv_pool_rebalance, kv_pool_rebalance_check_interval=llm_args.kv_cache_config. diff --git a/tensorrt_llm/_torch/pyexecutor/dwdp.py b/tensorrt_llm/_torch/pyexecutor/dwdp.py index 0497b1299119..c9ff0ad50995 100644 --- a/tensorrt_llm/_torch/pyexecutor/dwdp.py +++ b/tensorrt_llm/_torch/pyexecutor/dwdp.py @@ -14,7 +14,7 @@ # limitations under the License. """DwdpManager facade for VA-based DWDP. -This module owns three concerns and delegates the rest: +This module owns four concerns and delegates the rest: * Lifecycle / global singleton (``__enter__`` / ``__exit__`` + ``set_global_dwdp_manager`` / ``get_global_dwdp_manager``) * DWDP group creation (``_create_dwdp_comm``): an MPI sub-communicator, or @@ -23,16 +23,27 @@ plus runtime entry points that forward to ``DWDPWeightManager`` (``start_forward``, ``wait_and_bind``, ``finish_layer``, ``invalidate_prefetch``) + * Attached dense DWDP: the ``DenseDwdpManager`` created and registered + with this manager, through which the model's layer hooks reach the + ``DenseDWDP`` runtime. ``setup(model)`` commits the runtime over the DWDP + group before the expert setup allocates its prefetch ring, so the + replicated dense weights are freed first, then hands it to the dense + manager together with the weight manager, which runs on the same copy + stream; each forward starts the dense prefetch ahead of the expert + prefetch (``start_forward``, ``DenseDwdpManager.ensure_primed``), and + ``cleanup`` releases dense DWDP before the expert teardown The heavy lifting — MNNVL handle allocation, composite VA layout, prefetch scheduling — lives in ``tensorrt_llm._torch.modules.dwdp`` -and is wired up by ``setup_dwdp()`` during ``setup(model)``. +and is wired up by ``setup_dwdp()`` and ``setup_dense_dwdp()`` during +``setup(model)``. """ from __future__ import annotations +import traceback import weakref -from typing import List, Optional +from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn @@ -41,14 +52,29 @@ from tensorrt_llm._torch.cute_dsl_utils import set_sm_reserve from tensorrt_llm._torch.distributed import MPIDist from tensorrt_llm._torch.modules.dwdp import ( + DenseDWDP, DWDPWeightManager, + ExpertPartition, StoreDwdpGroup, + dense_dwdp_requested, dwdp_global_rank, fabric_gather, + setup_dense_dwdp, setup_dwdp, store_rendezvous_enabled, teardown_dwdp, ) +from tensorrt_llm._torch.modules.dwdp.dense.manager import ( + DenseDwdpManager, + get_global_dense_dwdp_manager, +) +from tensorrt_llm._torch.modules.dwdp.dense.setup import ( + dense_dwdp_active, + dense_roots_available, + refuse_dense_dwdp_without_roots, +) +from tensorrt_llm._torch.modules.dwdp.dense.types import DenseMode +from tensorrt_llm._torch.modules.dwdp.setup import _release_during_rollback from tensorrt_llm.llmapi.llm_args import DwdpConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -65,20 +91,135 @@ def get_global_dwdp_manager() -> Optional["DwdpManager"]: return _global_dwdp_manager +def new_expert_copy_stream(device_id: int) -> torch.cuda.Stream: + """The expert DWDP copy stream on ``cuda:``, made as ``DWDPWeightManager`` makes it. + + ``DwdpManager.setup`` creates it for attached dense DWDP, which runs before the weight manager + exists, and hands it to ``setup_dwdp``. The test seam: tests monkeypatch it. + """ + return torch.cuda.Stream(device=torch.device("cuda", device_id)) + + +class _AttachedTransfer: + """The transfer of attached dense DWDP (``DenseTransfer``), before the weight manager exists. + + Dense DWDP commits before ``setup_dwdp`` allocates the expert prefetch ring, so it cannot take + the expert weight manager as its transfer. This stands in with what the weight manager will + use: + + * ``copy_stream``: made by ``new_expert_copy_stream`` at the first read, which dense setup + makes in P2, inside the phase's local work (a failure is this rank's failed P2 report); + ``DwdpManager.setup`` then hands it to ``setup_dwdp``. + * ``use_gather`` and ``use_batched_prefetch``: computed when read, as + ``DWDPWeightManager.__init__`` computes them. + + Construction touches no CUDA, probes nothing and logs nothing. + """ + + __slots__ = ("_device_id", "_contention_opt", "_stream") + + def __init__(self, device_id: int, *, contention_opt: bool) -> None: + self._device_id = device_id + self._contention_opt = contention_opt + self._stream: Optional[torch.cuda.Stream] = None + + @property + def copy_stream(self) -> torch.cuda.Stream: + """The copy stream: created on the first read, the same object on every later one.""" + if self._stream is None: + self._stream = new_expert_copy_stream(self._device_id) + return self._stream + + @property + def created_stream(self) -> Optional[torch.cuda.Stream]: + """The copy stream if a read created it, else None.""" + return self._stream + + @property + def use_gather(self) -> bool: + """Whether the expert prefetch uses the SM fabric gather (``TRTLLM_DWDP_PREFETCH_GATHER``).""" + return fabric_gather.available() + + @property + def use_batched_prefetch(self) -> bool: + """Whether the expert prefetch uses one batched copy per layer (``contention_opt``).""" + return self._contention_opt + + +def _check_shared_transfer( + transfer: _AttachedTransfer | DWDPWeightManager, weight_manager: DWDPWeightManager +) -> None: + """The weight manager must run on the copy stream and transfer dense DWDP was set up with. + + Dense and expert copies share one copy stream, which runs them in the order they are enqueued, + and the dense copies follow the expert transfer. On a repeated setup dense DWDP ran on the + weight manager itself, which the expert setup must then have returned. Nothing to check when + dense setup never read the attached transfer's stream. + + Raises: + RuntimeError: The weight manager runs on another stream or with other transfer flags. + """ + if transfer is weight_manager: + return + if isinstance(transfer, _AttachedTransfer): + stream = transfer.created_stream + if stream is None: + return + else: + stream = transfer.copy_stream + dense_flags = (transfer.use_gather, transfer.use_batched_prefetch) + expert_flags = (weight_manager.use_gather, weight_manager.use_batched_prefetch) + if weight_manager.copy_stream is not stream or expert_flags != dense_flags: + raise RuntimeError( + "DwdpManager.setup(): the expert weight manager does not run on the copy stream and " + "transfer dense DWDP was set up with; the dense and the expert copies must share one " + f"copy stream (expert: {weight_manager.copy_stream}, (use_gather, " + f"use_batched_prefetch)={expert_flags}; dense: {stream}, {dense_flags})" + ) + + +def _dense_slot_storages(dense: DenseDWDP | None) -> frozenset[int]: + """The storages of a committed dense runtime's slots, which the expert transport must not take. + + After the commit the modules' dense tensors are views of these slots. Empty without a + runtime. + """ + if dense is None: + return frozenset() + return frozenset( + slot.storage.untyped_storage().data_ptr() for ring in dense.slots.values() for slot in ring + ) + + class DwdpManager: """Lifecycle facade that plumbs the VA-based DWDP pipeline. Construction creates the DWDP group: an MPI sub-communicator, or a ``StoreDwdpGroup`` when the context workers are launched independently - (see ``modules.dwdp.rendezvous``). ``add_layer`` is the - Single Source of Truth for MoE layer indices — ``configurable_moe`` - registers each MoE layer during model construction, and ``setup(model)`` - passes that list to ``setup_dwdp()`` so ``collect_moe_params`` never has - to re-discover layers by walking the model tree. + (see ``modules.dwdp.rendezvous``), and the attached dense DWDP manager + (``dense_manager``); ``__enter__`` registers both globally. + ``add_layer`` is the Single Source of Truth for MoE layer indices — + ``configurable_moe`` registers each MoE layer during model construction, + and ``setup(model)`` passes that list to ``setup_dwdp()`` so + ``collect_moe_params`` never has to re-discover layers by walking the + model tree. After ``setup()`` returns, the runtime methods (``start_forward``, ``wait_and_bind``, ``finish_layer``) forward to the ``DWDPWeightManager`` instance produced by ``setup_dwdp``. + + When dense DWDP is committed, the dense manager owns its runtime + (``dense``) and holds the weight manager: the model's decoder-layer hooks + drive the runtime and the model's forward calls + ``DenseDwdpManager.ensure_primed``. ``L`` is a layer's dense gather and + ``E`` its expert prefetch, both on the weight manager's copy stream. A + completed forward leaves ``L`` and ``E`` of the first two layers enqueued + for the next one; otherwise ``start_forward`` or ``ensure_primed`` + enqueues ``L(first), L(second)``, then the expert prefetch its + ``E(first), E(second)``. Without a committed runtime the dense manager + does nothing. ``setup`` commits dense DWDP before the expert setup + allocates its prefetch ring, on the copy stream the weight manager then + adopts. """ def __init__( @@ -89,6 +230,10 @@ def __init__( ) -> None: if get_global_dwdp_manager() is not None: raise RuntimeError("DwdpManager already registered globally") + # __enter__ registers the attached dense DWDP manager as well; refuse a + # registered one here, before the DWDP group exists. + if get_global_dense_dwdp_manager() is not None: + raise RuntimeError("DenseDwdpManager already registered globally") # Independently launched workers each have their own MPI world, so the # group rendezvous over a TCPStore instead of COMM_WORLD. self._store_rendezvous = store_rendezvous_enabled() @@ -145,18 +290,10 @@ def __init__( self.num_groups = config.num_groups self.num_prefetch_experts = config.num_prefetch_experts - # Per-rank expert range. ``num_prefetch_experts`` is the rank-to-rank - # stride and ``num_experts_per_worker`` is the (uniform across ranks) - # storage chunk size. When the user requests redundancy - # (``num_prefetch_experts < num_experts_per_worker``) consecutive - # ranks' ranges overlap; when the partition is non-uniform (e.g. - # dwdp_size = 3 over 256 experts) the tail rank's - # ``[start, end)`` runs past ``num_experts`` and is later capped at - # weight-load time. ``_init_dwdp_expert_layout`` reads - # ``start_expert_id`` to install the correct ``slot_start`` / - # ``slot_end`` on every fused MoE backend before ``create_weights``. - self.start_expert_id = self.dwdp_rank * self.num_prefetch_experts - self.end_expert_id = self.start_expert_id + self.num_experts_per_worker + # One expert partition per expert count: ``expert_partition`` computes + # it when a MoE layer of that count is built. DwdpConfig accepts the + # two partition fields only when both are 0 or both are positive. + self._expert_partitions: Dict[int, ExpertPartition] = {} # SMs the expert gather holds for the whole step; persistent grids leave # them out (``set_sm_reserve`` on ``__enter__``, before the model is built). @@ -172,6 +309,10 @@ def __init__( # Set by setup(model); None until then self._weight_manager: Optional[DWDPWeightManager] = None self._model_ref: Optional[weakref.ReferenceType[nn.Module]] = None + # Owns the committed dense DWDP runtime. Created whether or not + # TRTLLM_DWDP_DENSE is set: the model's layers capture it when they are + # built, and without a runtime it does nothing. + self._dense_manager = DenseDwdpManager(DenseMode.ATTACHED) # ------------------------------------------------------------------ # DWDP group @@ -204,10 +345,28 @@ def _create_dwdp_comm(self): # ------------------------------------------------------------------ def __enter__(self) -> "DwdpManager": + """Register this manager and its dense DWDP manager globally (both or neither), then + keep the gather's resident SMs out of the persistent grids.""" if get_global_dwdp_manager() is not None: raise RuntimeError("DwdpManager already registered globally") set_global_dwdp_manager(self) - set_sm_reserve(self.resident_sms) + try: + self._dense_manager.__enter__() + set_sm_reserve(self.resident_sms) + except BaseException: + # Not entered, so no __exit__ follows: unregister both and free the DWDP group + # __init__ created. Nothing else exists before setup(). + try: + self.cleanup() + except BaseException: + logger.error( + "[DwdpManager] Failed to free the DWDP group after a failed registration\n" + f"{traceback.format_exc()}" + ) + finally: + self._dense_manager.__exit__(None, None, None) + set_global_dwdp_manager(None) + raise return self def __exit__(self, exc_type, exc_val, exc_tb) -> bool: @@ -216,25 +375,64 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> bool: finally: set_sm_reserve(0) set_global_dwdp_manager(None) + self._dense_manager.__exit__(None, None, None) return False def is_enabled(self) -> bool: return self.dwdp_size > 1 + @property + def dense_manager(self) -> DenseDwdpManager: + """The attached dense DWDP manager, registered globally while this manager is entered.""" + return self._dense_manager + + @property + def dense(self) -> DenseDWDP | None: + """The committed dense DWDP runtime, or None when dense DWDP is not active. + + The dense manager's runtime: None when dense DWDP is off or was + abandoned, before ``setup(model)`` commits it and after ``cleanup()``. + The model's layer hooks run only while it is not None. + """ + return self._dense_manager.dense + def cleanup(self) -> None: - """Release the DWDP group (MPI sub-communicator or TCPStore) and reset - the weight manager. Idempotent.""" + """Release dense DWDP, tear down the weight manager, free the DWDP group. + + Dense DWDP goes first (``DenseDwdpManager.cleanup``, which also drops + its reference to the weight manager): the release synchronizes the + device, repoints the model's dense weights to zero-size tensors and + frees the staging slots and the row store. The model keeps its + ``_dense_dwdp_active`` mark, so weight reloads stay refused. If the + model was garbage collected first, the bindings of its modules count + as detached, since nothing can read them any more. Every step runs even + when an earlier one fails; the first error is raised at the end and + later ones are logged. Idempotent. + """ cleanup_error = None cleanup_traceback = None try: model = self._model_ref() if self._model_ref is not None else None + try: + self._dense_manager.cleanup() # a no-op without a committed runtime + except BaseException as exc: # the expert teardown below must still run + cleanup_error, cleanup_traceback = exc, exc.__traceback__ + logger.error( + "[DwdpManager] Dense DWDP release failed; tearing down expert " + f"DWDP anyway\n{traceback.format_exc()}" + ) if model is not None and self._weight_manager is not None: teardown_dwdp(model, expected_manager=self._weight_manager) elif self._weight_manager is not None: self._weight_manager.release() except BaseException as exc: - cleanup_error = exc - cleanup_traceback = exc.__traceback__ + if cleanup_error is None: + cleanup_error, cleanup_traceback = exc, exc.__traceback__ + else: + logger.error( + "[DwdpManager] Expert DWDP teardown failed after the dense DWDP " + f"release error\n{traceback.format_exc()}" + ) finally: self._weight_manager = None self._model_ref = None @@ -244,13 +442,57 @@ def cleanup(self) -> None: except BaseException: if cleanup_error is None: raise - logger.exception("Failed to free DWDP group after runtime cleanup error") + logger.error( + "[DwdpManager] Failed to free DWDP group after runtime cleanup " + f"error\n{traceback.format_exc()}" + ) finally: self.dwdp_group = None if cleanup_error is not None: raise cleanup_error.with_traceback(cleanup_traceback) + # ------------------------------------------------------------------ + # Expert partition + # ------------------------------------------------------------------ + + def expert_partition(self, num_experts: int) -> ExpertPartition: + """The expert partition of the DWDP group for MoE layers of ``num_experts`` routed experts. + + Computed from ``DwdpConfig`` with ``ExpertPartition.from_config``, which ``setup_dwdp`` + repeats and checks across the group; pure (no collective) and cached per expert count. + ``MoE._init_dwdp_expert_layout`` sizes every MoE backend from it before + ``create_weights``. + + Raises: + ValueError: The config gives no valid partition of ``num_experts`` experts, raised + the same way on every rank while the model is built. + """ + partition = self._expert_partitions.get(num_experts) + if partition is not None: + return partition + partition = ExpertPartition.from_config( + num_experts=num_experts, + dwdp_size=self.dwdp_size, + num_experts_per_worker=self.num_experts_per_worker, + num_prefetch_experts=self.num_prefetch_experts, + ) + self._expert_partitions[num_experts] = partition + start, end = partition.range_of(self.dwdp_rank) + logger.info( + f"[DwdpManager] DWDP rank {self.dwdp_rank} of {self.dwdp_size} stores routed experts " + f"[{start}, {end}): {end - start} of {num_experts} ({partition.describe()})" + ) + return partition + + def local_expert_range(self, num_experts: int) -> Tuple[int, int]: + """``(start, end)`` of the routed experts this rank stores, of ``num_experts`` per layer. + + Raises: + ValueError: As ``expert_partition``. + """ + return self.expert_partition(num_experts).range_of(self.dwdp_rank) + # ------------------------------------------------------------------ # Layer registration (SSOT) + setup # ------------------------------------------------------------------ @@ -273,31 +515,149 @@ def add_repeated_layer(self, layer_idx: int) -> None: self._repeated_layers.append(layer_idx) def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: - """Run the VA initialization pipeline on ``model``. - - Delegates to ``setup_dwdp`` (transport + weight buffer + weight - manager + MoE backend fixup). Stores the returned ``DWDPWeightManager`` - for subsequent runtime calls. - - Called once from ``py_executor_creator`` after weight loading. + """Set dense DWDP, then the VA pipeline of expert DWDP, up on ``model``. + + First, if the model's decoder layers provide dense roots, runs + ``setup_dense_dwdp`` on every rank, also on a rank that did not ask + for dense DWDP, so the first dense collective gives every rank of the + group the same verdict when ranks disagree on ``TRTLLM_DWDP_DENSE``. + Dense setup errors propagate before any expert setup work. A commit + returns the replicated dense weights to the driver before the expert + setup allocates its prefetch ring, so the setup peak is the larger of + the replicated dense weights with the dense store and slots, and the + dense shard with the ring, not their sum. Dense setup runs on an + attached transfer: a copy stream that the transfer creates when dense + setup first reads it, and the transfer flags computed as the weight + manager computes them. + + Then ``setup_dwdp`` (transport + weight buffer + weight manager + MoE + backend fixup) runs with that copy stream, if dense setup created + it, so the dense and the expert copies share one copy stream, and the + returned ``DWDPWeightManager`` is stored for subsequent runtime calls. + It gets the committed runtime's slots as reserved storages, so that + an expert weight dense DWDP moved is refused before the transport + would free a slot. The committed runtime, if any, goes to the dense + manager together with the weight manager (``DenseDwdpManager.attach``) + and becomes ``dense``. If anything fails after the commit and before + the dense manager took the runtime over, the runtime is released and + the error propagates. Whether a model has roots depends only on its + code, so on a model without them no rank makes a dense collective: a + rank that asks for dense DWDP there decides alone, at once, after the + expert setup (``refuse_dense_dwdp_without_roots``: it raises unless + its ``TRTLLM_DWDP_DENSE_REQUIRED`` is off). + + Called once from ``py_executor_creator`` after weight loading. Once + dense DWDP is committed, a later call for the same model returns the + weight manager and keeps the runtime, making no collective; every + rank committed together, so every rank returns early alike. A later + call for the same model after dense DWDP stayed off runs dense setup + on the weight manager, whose stream and flags it then uses; the + expert setup must return that same weight manager. + + Raises: + RuntimeError: Called after ``cleanup()``; dense DWDP is committed + for another model; or the model's dense weights were moved + by a dense DWDP runtime this manager does not hold (it was + released, so the modules hold zero-size tensors). Each is + raised before any expert or dense setup work. Also raised + when the weight manager does not run on the copy stream and + transfer dense DWDP was set up with, and when dense DWDP + moved an expert weight the transport takes over. The expert + setup raises it on every rank after the same collective when + the DWDP group disagrees on the expert partition, when another + rank failed in the transport or in a later step before the + vote, or when the ranks' small expert parameters do not fit + together; a failure in the backend fixup after the vote raises + on that rank only. + DenseSetupError: Dense DWDP setup failed, or was requested on a + model without dense roots, with ``TRTLLM_DWDP_DENSE_REQUIRED`` + on. + PrefetchRingMemoryError: This rank's expert prefetch ring does not + fit in the free device memory (the other ranks raise a + RuntimeError naming this rank). """ + if self.dwdp_group is None: + raise RuntimeError("DwdpManager.setup() called after cleanup()") + if self._dense_manager.dense is not None: + if self._model_ref is None or self._model_ref() is not model: + raise RuntimeError( + "DwdpManager.setup(): dense DWDP is already committed for another " + "model; a manager serves one model" + ) + return self._weight_manager + if dense_dwdp_active(model): + raise RuntimeError( + "DwdpManager.setup(): dense DWDP already moved this model's dense " + "weights and its runtime was released, so the modules hold zero-size " + "tensors; load the model again to set it up" + ) layer_indices = sorted(self._registered_layers) device_id = torch.cuda.current_device() - self._weight_manager = setup_dwdp( - model=model, - mapping=self._mapping, - device_id=device_id, - comm=self.dwdp_group, - layer_indices=layer_indices, - num_experts_per_worker=self.num_experts_per_worker, - num_prefetch_experts=self.num_prefetch_experts, - contention_opt=self.config.contention_opt, - num_prefetch_slots=self.config.num_prefetch_slots, - repeated_layer_indices=[ - layer for layer in self._repeated_layers if layer in self._registered_layers - ], + roots = dense_roots_available(model) + # A repeated call for the model whose expert setup is done (dense DWDP stayed off): the + # expert setup returns its cached weight manager, so dense setup runs on that one. + expert_done = ( + self._weight_manager is not None + and self._model_ref is not None + and self._model_ref() is model + ) + attached = ( + None + if expert_done + else _AttachedTransfer(device_id, contention_opt=self.config.contention_opt) ) - self._model_ref = weakref.ref(model) + transfer = self._weight_manager if attached is None else attached + dense = None + if roots: # on every rank, requested or not: collective #1 gives one verdict + # Collective #1 is the group's first collective of the setup, where a rank waits for + # peers still loading: one line, so that a rank waiting there shows where it waits. + logger.info( + f"[DwdpManager] DWDP rank {self.dwdp_rank} of {self.dwdp_size}: dense DWDP setup " + f"first (its collective #1 waits for every rank of the DWDP group), then expert " + f"DWDP setup" + ) + dense = setup_dense_dwdp( + model, + transfer, + self.dwdp_group, + dwdp_rank=self.dwdp_rank, + dwdp_size=self.dwdp_size, + device_id=device_id, + moe_layers=layer_indices, + ) + try: + self._weight_manager = setup_dwdp( + model=model, + mapping=self._mapping, + device_id=device_id, + comm=self.dwdp_group, + layer_indices=layer_indices, + num_experts_per_worker=self.num_experts_per_worker, + num_prefetch_experts=self.num_prefetch_experts, + contention_opt=self.config.contention_opt, + num_prefetch_slots=self.config.num_prefetch_slots, + repeated_layer_indices=[ + layer for layer in self._repeated_layers if layer in self._registered_layers + ], + copy_stream=None if attached is None else attached.created_stream, + reserved_storages=_dense_slot_storages(dense), + ) + self._model_ref = weakref.ref(model) + if dense is not None: + _check_shared_transfer(transfer, self._weight_manager) + self._dense_manager.attach(dense, expert=self._weight_manager) + except BaseException: + if dense is not None and self._dense_manager.dense is not dense: + _release_during_rollback(dense, "dense DWDP runtime") + raise + if not roots and dense_dwdp_requested(): + refuse_dense_dwdp_without_roots(model) + if self.dense is not None and self.config.num_prefetch_slots > 2: + logger.info( + f"[DwdpManager] dense DWDP with {self.config.num_prefetch_slots} expert prefetch " + f"slots: each expert copy runs ahead of its layer's dense copy on the shared copy " + f"stream" + ) if self._store_rendezvous: # Lets DWDP rank 0, which serves the store, release it once every group is set up. self.dwdp_group.mark_setup_done() @@ -308,15 +668,47 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: # ------------------------------------------------------------------ def start_forward(self) -> None: - """Start a forward pass: prefetch its first MoE layers. - - Enqueues nothing when the previous forward pass already prefetched - them (its last MoE layers gather the first layers of the next pass). + """Start a forward pass: prefetch its first layers. + + Prefetches the experts of the first MoE layers unless the previous + forward pass already did (its last MoE layers gather the first layers + of the next pass). With dense DWDP, the dense manager first enqueues + the dense weights of the first two layers, unless they are already + enqueued, so the copies this enqueues put each layer's dense copy + ahead of its expert copy on the shared copy stream. While a forward is + still open, which means the previous forward aborted inside the covered + layers, this enqueues nothing: the model's forward entry + (``DenseDwdpManager.ensure_primed``) recovers on the execution stream, + dense first. """ if self._weight_manager is None: raise RuntimeError("DwdpManager.setup() has not been called yet") + dense = self._dense_manager.dense + if dense is not None and dense.forward_open: + return + self._dense_manager.prefetch_first_layers() # a no-op without a committed runtime self._weight_manager.start_forward() + def ensure_primed(self) -> None: + """Model-side priming (H0) through ``DenseDwdpManager.ensure_primed``. + + The model reaches the dense manager directly; this delegate serves + callers that hold the ``DwdpManager``. The dense manager refuses a CUDA + graph capture and multi-stream execution before it records or enqueues + anything; after an aborted forward it records the dense WAR barrier on + the current stream; it enqueues ``L(first), L(second)`` unless already + primed, then starts the expert prefetch (``start_forward`` of the + weight manager). + + Raises: + RuntimeError: Dense DWDP is not active. + DenseProtocolError: The current stream is capturing a CUDA graph, + or multi-stream execution is on. + """ + if self._dense_manager.dense is None or self._weight_manager is None: + raise RuntimeError("DwdpManager.ensure_primed() needs an active dense DWDP runtime") + self._dense_manager.ensure_primed() + def wait_and_bind(self, backend_module: nn.Module, layer_idx: int) -> None: """Wait for prefetch and bind full-tensor views onto backend weights.""" if self._weight_manager is None: diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index f967e92d7069..b518dab51ff5 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -724,6 +724,32 @@ def _check_compile_wrapper_before_reload(model: torch.nn.Module) -> None: _assert_no_compile_wrapper_in_paths(model) +def _reject_dense_dwdp_update(model: torch.nn.Module, action: str) -> None: + """Refuse to rewrite the weights of a model that dense DWDP sharded. + + Dense DWDP repoints the dense weights of the MoE decoder layers into staging + slots that several layers share and that the next gather overwrites, so + weights written through the modules would be lost. Its setup marks the + model and its decoder with ``_dense_dwdp_active``; the mark stays after the + dense DWDP teardown, which leaves zero-size tensors behind. + + Args: + model: The model about to be updated. + action: What the caller was about to do, for the error message. + + Raises: + RuntimeError: The model or its decoder (``model.model``) is marked. + """ + decoder = getattr(model, "model", None) + if getattr(model, "_dense_dwdp_active", False) or getattr( + decoder, "_dense_dwdp_active", False): + raise RuntimeError( + f"Cannot {action}: dense DWDP (TRTLLM_DWDP_DENSE) moved this " + "model's dense weights into staging slots shared across layers, so " + "they cannot be rewritten in place. Launch without " + "TRTLLM_DWDP_DENSE to update weights.") + + class ModelLoaderMetricNames(Enum): TOTAL_MODEL_LOADING_SECONDS = "total_model_loading_seconds" CHECKPOINT_PREPARATION_SECONDS = "checkpoint_preparation_seconds" @@ -1960,7 +1986,12 @@ def reload(self, Returns: None. + + Raises: + RuntimeError: Dense DWDP sharded the model's dense weights, or the + weight mapper was not initialized. """ + _reject_dense_dwdp_update(model, "reload weights") if self.weight_mapper is None: raise RuntimeError( "Cannot reload weights: weight_mapper was not initialized. " @@ -1977,10 +2008,16 @@ def reload(self, torch.cuda.current_stream().synchronize() def begin_update_weights(self) -> None: - """Start an incremental update session on the persistent mapper.""" + """Start an incremental update session on the persistent mapper. + + Raises: + RuntimeError: The weight mapper was not initialized, or dense DWDP + sharded the dense weights of the model it is bound to. + """ if self.weight_mapper is None: raise RuntimeError( "Cannot update weights: weight_mapper was not initialized") + _reject_dense_dwdp_update(self.weight_mapper.model, "update weights") self.weight_mapper.begin_update_weights() def finalize_update_weights(self) -> None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a0e323196a53..1dce7afd26cf 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -75,6 +75,7 @@ maybe_prefetch_mm_encoder_for_next_iter from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.decoder_layer import DecoderLayer +from ..modules.dwdp.dense.manager import DenseDwdpManager from ..moe.expert_statistic import ExpertStatistic from ..speculative.drafter import Drafter from ..speculative.spec_sampler_base import SpecSampler @@ -206,6 +207,53 @@ def _fill_stall_timeout_sec() -> float: return value +# Environment variable that lets a requested shutdown stop the executor loop +# while requests are still active or queued, instead of draining them first. +# "1" enables it; any other value, or leaving it unset, keeps the drain. +# Dropped requests get no final response, and a request still receiving KV +# has its cache torn down with the pool instead of waited for, so it is meant +# for context-only servers whose clients are gone by the time they shut down. +SHUTDOWN_DROP_INFLIGHT_ENV_VAR_NAME = "TRTLLM_EXECUTOR_SHUTDOWN_DROP_INFLIGHT" + + +def _shutdown_drop_inflight_enabled(pp_size: int, + disable_overlap_scheduler: bool) -> bool: + """Resolve TRTLLM_EXECUTOR_SHUTDOWN_DROP_INFLIGHT for one executor. + + Only the pp_size == 1 loop without the overlap scheduler may stop with + requests in flight. A pipeline-parallel loop can still have microbatches + travelling between ranks, whose send handles the loop cleanup waits on, + and the overlap loop still owes the previous batch its processing. The + effective value is logged whenever the variable is set. + + Args: + pp_size: Pipeline-parallel size of the executor. + disable_overlap_scheduler: Whether the executor runs the loop without + the overlap scheduler. + + Returns: + True when a requested shutdown may drop the requests that are still + active or queued. + """ + raw = os.environ.get(SHUTDOWN_DROP_INFLIGHT_ENV_VAR_NAME) + if raw is None: + return False + enabled = raw == "1" and pp_size == 1 and disable_overlap_scheduler + if enabled: + state = "on" + elif raw != "1": + state = "off (only '1' enables it)" + else: + state = (f"off (needs pp_size == 1 and disable_overlap_scheduler; " + f"pp_size={pp_size}, disable_overlap_scheduler=" + f"{disable_overlap_scheduler})") + # Any value but "1" is quoted, so stray whitespace or case shows. + value = raw if raw == "1" else repr(raw) + logger.info(f"{SHUTDOWN_DROP_INFLIGHT_ENV_VAR_NAME}={value}: dropping " + f"in-flight requests at shutdown is {state}") + return enabled + + # Environment variable to control which ranks print step logging. # Format: comma-separated rank IDs, e.g. "0,1,3", or "all" for all ranks. # Default: "0" (only rank 0 prints, matching existing behavior). @@ -488,6 +536,7 @@ def __init__( waiting_queue_policy: WaitingQueuePolicy = WaitingQueuePolicy.FCFS, adp_router: Optional[ADPRouter] = None, dwdp_manager: Optional[DwdpManager] = None, + dense_dwdp_manager: Optional[DenseDwdpManager] = None, enable_kv_pool_rebalance: bool = False, kv_pool_rebalance_check_interval: int = KV_POOL_REBALANCE_CHECK_INTERVAL, enable_adp_dummy_fixes: bool = False, @@ -998,6 +1047,11 @@ def __init__( self.kv_cache_manager.snapshot_warmup_baseline() self.is_shutdown = False + # Whether a requested shutdown may drop the requests still in flight. + # Read through getattr, so executors assembled without __init__ (unit + # tests) keep draining. + self._drop_inflight_on_shutdown = _shutdown_drop_inflight_enabled( + self.dist.pp_size, self.disable_overlap_scheduler) # Set at the executor loops' normal-exit `break` sites, and ONLY # there. It answers exactly one question for _event_loop_wrapper: # "did event_loop() reach its own termination?" -- which decides @@ -1006,6 +1060,8 @@ def __init__( # rank-locally on a fatal error (e.g. a CUDA illegal address on this # rank alone) while peers are told nothing and the loop keeps running # collectives, so a crash after that point still strands them. + # _executor_loop also reads it once its loop ends, so that only the + # normal exit releases what a dropping shutdown left behind. self._event_loop_completed = False self._fatal_error: Optional[BaseException] = None self._kv_connector_failed = False @@ -1190,10 +1246,8 @@ def on_detected(): if is_trace_enabled("TLLM_TRACE_EXECUTOR_LOOP"): self.event_loop = trace_func(self.event_loop) - if dwdp_manager is not None and not self.disable_overlap_scheduler: - raise ValueError( - "DWDP requires disable_overlap_scheduler=True. " - "Overlap scheduler is not yet supported with DWDP.") + self._validate_dwdp_managers(dwdp_manager, dense_dwdp_manager, + self.disable_overlap_scheduler) if self.drafter is not None: if self.event_loop.__name__ == self._executor_loop_pp.__name__: @@ -1233,6 +1287,7 @@ def on_detected(): self._maybe_init_kv_connector_manager() self.dwdp_manager = dwdp_manager + self.dense_dwdp_manager = dense_dwdp_manager self._validate_moe_a2a_checkpoint_support() @@ -1257,6 +1312,32 @@ def _kv_connector_moves_kv(self) -> bool: connector = getattr(self, "kv_connector_manager", None) return connector is not None and not connector.capacity_only + @staticmethod + def _validate_dwdp_managers(dwdp_manager: Optional[DwdpManager], + dense_dwdp_manager: Optional[DenseDwdpManager], + disable_overlap_scheduler: bool) -> None: + """Reject DWDP managers this executor cannot drive. + + Expert DWDP (``dwdp_manager``) runs only with the overlap scheduler + off. A standalone dense DWDP manager runs with either loop: the + non-overlap loop primes it before each forward, and in the overlap loop + the model's forward entry (H0) primes it. With expert DWDP, + ``DwdpManager`` drives dense DWDP itself, so only one of the two + managers may be given. + + Raises: + ValueError: ``dwdp_manager`` with the overlap scheduler on, or both + managers given. + """ + if dwdp_manager is not None and not disable_overlap_scheduler: + raise ValueError( + "DWDP requires disable_overlap_scheduler=True. " + "Overlap scheduler is not yet supported with DWDP.") + if dwdp_manager is not None and dense_dwdp_manager is not None: + raise ValueError( + "pass dense_dwdp_manager only without DWDP: with dwdp_manager, " + "DwdpManager drives dense DWDP") + def _maybe_init_kv_connector_manager(self) -> None: if self.kv_connector_manager is not None: # A connector that moves no KV cannot contend with the transceiver @@ -2167,6 +2248,9 @@ def _shutdown(self, *, release_model_owned_graphs: bool) -> None: if self.dwdp_manager is not None: self.dwdp_manager.__exit__(None, None, None) self.dwdp_manager = None + if self.dense_dwdp_manager is not None: + self.dense_dwdp_manager.__exit__(None, None, None) + self.dense_dwdp_manager = None def can_enqueue_requests(self) -> bool: """ @@ -2320,8 +2404,48 @@ def set_gather_responses(self, gather_all_responses): @property def should_stop_processing(self) -> bool: - return self.is_shutdown and len(self.active_requests) == 0 and \ - len(self.waiting_queue) == 0 and not self._has_pending_connector_transfers() + if not self.is_shutdown: + return False + if len(self.active_requests) == 0 and len(self.waiting_queue) == 0: + return not self._has_pending_connector_transfers() + # Requests are still active or queued. They are drained unless + # TRTLLM_EXECUTOR_SHUTDOWN_DROP_INFLIGHT is in effect, and even then a + # pending connector transfer is waited for. The loop hands what they + # hold back through _release_requests_dropped_at_shutdown. + if (not getattr(self, "_drop_inflight_on_shutdown", False) + or self._has_pending_connector_transfers()): + return False + logger.warning_once( + f"shutdown requested: dropping {len(self.active_requests)} active " + f"and {len(self.waiting_queue)} queued requests instead of " + f"draining them ({SHUTDOWN_DROP_INFLIGHT_ENV_VAR_NAME}=1)", + key="executor_shutdown_drop_inflight") + return True + + def _release_requests_dropped_at_shutdown(self) -> None: + """Release what the requests dropped by a non-draining shutdown hold. + + Runs once the loop has stopped. A no-op unless + TRTLLM_EXECUTOR_SHUTDOWN_DROP_INFLIGHT is in effect, since a draining + shutdown only stops with no request left. Queued items hold no + executor resources and are discarded. Each active request goes through + ``_terminate_request``, as a finished or cancelled one does, so its KV + cache, sequence slot and per-request state are back with their + managers before ``shutdown`` tears those down. A request whose KV cache + the transfer fabric may still access keeps it, as it would on + cancellation; the KV cache manager closes that cache at its own + shutdown, as it does for a context still in transfer. No dropped + request gets a final response. + """ + if not getattr(self, "_drop_inflight_on_shutdown", False): + return + while self.waiting_queue: + self.waiting_queue.pop_request() + dropped_requests = list(self.active_requests) + self.active_requests.clear() + for request in dropped_requests: + if not self._is_request_in_transmission(request): + self._terminate_request(request) def _has_pending_connector_transfers(self) -> bool: connector = getattr(self, "kv_connector_manager", None) @@ -4280,6 +4404,17 @@ def _finalize_adp_dummy_allocation(self, can_queue: bool) -> None: if spec_resource_manager is not None: spec_resource_manager.free_resources(dummy_request) self.kv_cache_manager.free_resources(dummy_request) + # The V2 scheduler also admits the draft pool's mirror while it + # schedules the dummy (KVCacheV2Scheduler._admit_unpaired_draft), so + # that pool holds an empty cache under the dummy ID too. Left behind, + # it outlives this dummy, and the deferred previous batch -- which + # still holds the earlier, finished dummy with the same ID -- resizes + # it to that dummy's history in update_resources ("History length + # cannot exceed capacity"). + draft_kv_cache_manager = getattr(self, "draft_kv_cache_manager", None) + if (getattr(self, "_is_kv_manager_v2", False) + and draft_kv_cache_manager is not None): + draft_kv_cache_manager.free_resources(dummy_request) self.active_requests.remove(dummy_request) def _revert_ctx_alloc(self, dropped_context_requests): @@ -5229,6 +5364,24 @@ def _step_scope(self, return torch.profiler.record_function( self._build_step_scope_label(scheduled_batch, phase)) + def _prime_dwdp(self) -> None: + """Start DWDP for the next forward of the non-overlap loop. + + Expert DWDP starts its prefetch (``DwdpManager.start_forward``): the + first MoE layers, unless the previous forward already prefetched them, + each after the layer's dense weights when dense DWDP is attached. A + standalone dense DWDP manager enqueues the dense weights of the first + two layers. Dense DWDP skips what is primed already and leaves the + recovery after an aborted forward to the model's forward entry (H0). + The overlap and pipeline-parallel loops do not call this: expert DWDP + refuses the overlap scheduler, and a standalone dense DWDP manager is + primed by H0 in every forward there. + """ + if self.dwdp_manager is not None: + self.dwdp_manager.start_forward() + elif self.dense_dwdp_manager is not None: + self.dense_dwdp_manager.prefetch_first_layers() + def _executor_loop(self) -> None: torch.cuda.set_device(self.device_id) # ensure the context is created, otherwise, some MPI calls will fail. @@ -5391,8 +5544,7 @@ def _executor_loop(self) -> None: with self.perf_manager.record_perf_events( gpu_forward_start, gpu_forward_end) as fwd_timing: - if self.dwdp_manager is not None: - self.dwdp_manager.start_forward() + self._prime_dwdp() batch_outputs = self._forward_step(scheduled_batch) # None: _forward_step hit an exception and @@ -5527,6 +5679,12 @@ def _executor_loop(self) -> None: self.iter_counter += 1 + # Rank-local, after the loop's last collective and once the loop + # has marked itself completed, so an error here is a teardown error + # rather than a stranded peer. + if getattr(self, "_event_loop_completed", False): + self._release_requests_dropped_at_shutdown() + def _prepare_draft_requests(self): try: # Set draft tokens here to make the KV cache manager diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f0e465944c5a..7dbec12299d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -18,7 +18,12 @@ import tensorrt_llm from tensorrt_llm._startup import _StartupTimer -from tensorrt_llm._torch.modules.dwdp import dwdp_global_rank +from tensorrt_llm._torch.modules.dwdp import (dense_dwdp_requested, + dwdp_global_rank) +from tensorrt_llm._torch.modules.dwdp.dense.gate import validate_dense_gate +from tensorrt_llm._torch.modules.dwdp.dense.manager import ( + DenseDwdpManager, get_global_dense_dwdp_manager) +from tensorrt_llm._torch.modules.dwdp.dense.types import DenseMode from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, ResourceManagerType) from tensorrt_llm._utils import get_sm_version @@ -207,6 +212,28 @@ def observe_creation_stage(self, current_stage: ExecutorMemoryType): )) +def _check_dense_dwdp_launch(llm_args: TorchLlmArgs, + mapping: Mapping) -> Optional[DenseMode]: + """Run the dense DWDP launch gate if ``TRTLLM_DWDP_DENSE`` asks for it. + + It runs before the DWDP manager, the standalone dense DWDP manager and the + model are built, so a launch that dense DWDP cannot serve fails before the + weights load. + + Returns: + The mode dense DWDP runs in: attached to expert DWDP (``dwdp_config``) + or standalone on an attention-DP instance. None when dense DWDP is not + requested. + + Raises: + ValueError: Dense DWDP is requested and the launch violates one of its + requirements (see ``validate_dense_gate``). + """ + if dense_dwdp_requested(): + return validate_dense_gate(llm_args, mapping) + return None + + def _flashinfer_one_engine_spec_supported(attn_backend: str, spec_config) -> bool: """Whether this speculation is qualified on the FlashInfer target backend. @@ -733,6 +760,11 @@ def _create_py_executor( ) logger.info("ATTENTION RUNTIME FEATURES: ", attn_runtime_features) + # The dense DWDP gate runs outside the DWDP block below: it picks the dense + # DWDP mode, attached to expert DWDP or standalone without dwdp_config, and + # rejects a launch that neither mode serves. + dense_mode = _check_dense_dwdp_launch(llm_args, mapping) + # Initialize DWDP Manager. # # DWDP needs every global DWDP rank to be a complete, unsharded model replica @@ -774,6 +806,19 @@ def _create_py_executor( dwdp_manager.__enter__() logger.info(f"Dwdp Manager initialized. Config: {llm_args.dwdp_config}") + # Standalone dense DWDP shards the dense weights over the attention-DP + # group. Its manager is registered before the model is built, because the + # model's layers capture it then; with expert DWDP, DwdpManager creates and + # registers the dense manager itself. + dense_dwdp_manager: Optional[DenseDwdpManager] = None + if dense_mode is DenseMode.STANDALONE: + dense_dwdp_manager = DenseDwdpManager(DenseMode.STANDALONE, + mapping=mapping) + dense_dwdp_manager.__enter__() + logger.info( + "Dense DWDP manager initialized (standalone): attention-DP group " + f"of {mapping.tp_size} ranks") + _startup_timer.mark_initialization("configuration_and_distributed_init") mem_monitor = _ExecutorMemoryMonitor() @@ -1129,6 +1174,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): resources = {} estimating_kv_cache = False kv_cache_creator = None + # Whether DWDP was set up before the KV caches (below). + dwdp_set_up = False # Create the execution stream for model forward operations # for proper synchronization with KVCacheTransferManager's onboard/offload operations. @@ -1185,15 +1232,32 @@ def allocation_scope(current_stage: ExecutorMemoryType): estimating_kv_cache = kv_cache_creator.try_prepare_estimation() + # Without the estimation, build_managers sizes the final KV cache from + # the device memory free now, so dense DWDP is set up first, standalone + # or attached with the expert setup: the replicated dense weights it + # frees then go to the KV cache. With the estimation, it is set up + # below, before the estimation run measures the memory. Expert DWDP + # without a dense request keeps its place after the KV cache. + if dense_mode is not None and not estimating_kv_cache: + manager = (dwdp_manager + if dwdp_manager is not None else dense_dwdp_manager) + manager.setup(model_engine.model) + dwdp_set_up = True + with allocation_scope( ExecutorMemoryType.INIT_KV_CACHE if estimating_kv_cache else ExecutorMemoryType.KV_CACHE): kv_cache_creator.build_managers(resources, estimating_kv_cache) # DWDP setup: MNNVL handle exchange + composite VA weight buffer + - # weight manager + MoE backend fixup (single entry point). - if dwdp_manager is not None: - dwdp_manager.setup(model_engine.model) + # weight manager + MoE backend fixup (single entry point). Dense DWDP is set + # up there with expert DWDP, else here when it runs standalone. A launch + # with dense DWDP and without the estimation was set up above instead. + if not dwdp_set_up: + if dwdp_manager is not None: + dwdp_manager.setup(model_engine.model) + elif dense_dwdp_manager is not None: + dense_dwdp_manager.setup(model_engine.model) # Resource managers for speculative decoding # For user-specified drafters, use extra_resource_managers in PyTorchBackend config @@ -1246,6 +1310,12 @@ def allocation_scope(current_stage: ExecutorMemoryType): cache_transceiver_config=cache_transceiver_config, virtual_memory_pools=vm_pools if not estimating_kv_cache else None, execution_stream=execution_stream, + # Without the KV-cache estimation this is the final executor, which + # primes the DWDP manager or the standalone dense DWDP manager and + # exits it at shutdown. + dwdp_manager=dwdp_manager if not estimating_kv_cache else None, + dense_dwdp_manager=dense_dwdp_manager + if not estimating_kv_cache else None, max_num_sequences=max_num_seq_slots, ) @@ -1321,6 +1391,7 @@ def allocation_scope(current_stage: ExecutorMemoryType): virtual_memory_pools=vm_pools, execution_stream=execution_stream, dwdp_manager=dwdp_manager, + dense_dwdp_manager=dense_dwdp_manager, max_num_sequences=max_num_seq_slots, ) @@ -1360,8 +1431,15 @@ def create_py_executor( profiling_stage_data: Optional[dict] = None, resource_governor_queue=None, ) -> PyExecutor: - """Create a PyExecutor and roll back a partially initialized DWDP runtime.""" + """Create a PyExecutor and roll back a partially initialized DWDP runtime. + + On an error, a DWDP manager and a dense DWDP manager registered during the + construction are exited, in that order; managers registered before it are + left alone. Exiting the DWDP manager also exits the dense manager it + registered, so the second rollback finds a standalone one only. + """ previous_dwdp_manager = get_global_dwdp_manager() + previous_dense_manager = get_global_dense_dwdp_manager() try: with monitor_executor_initialization(), _StartupTimer( "executor_creation") as startup_timer: @@ -1383,4 +1461,13 @@ def create_py_executor( logger.error( "Failed to roll back DWDP after PyExecutor construction error\n" f"{traceback.format_exc()}") + current_dense_manager = get_global_dense_dwdp_manager() + if (current_dense_manager is not None + and current_dense_manager is not previous_dense_manager): + try: + current_dense_manager.__exit__(None, None, None) + except BaseException: + logger.error( + "Failed to roll back dense DWDP after PyExecutor construction " + f"error\n{traceback.format_exc()}") raise diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 08d018dcfb44..710756d7bdf7 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -77,6 +77,7 @@ from tensorrt_llm.serve.multi_frontend import (LAUNCHER_UDS_ENV, LAUNCHER_UDS_NAME, MultiFrontendServing) +from tensorrt_llm.serve.openai_server import _graceful_shutdown_timeout_from_env from tensorrt_llm.serve.tool_parser import ToolParserFactory from tensorrt_llm.serve.tool_parser.tool_parser_factory import ( MODEL_TYPE_TO_TOOL_PARSER, resolve_auto_tool_parser) @@ -579,6 +580,11 @@ def launch_server( enable_runtime_control_endpoints: bool = False, runtime_control_api_key: Optional[str] = None): + # OpenAIServer reads the graceful-shutdown timeout when it starts uvicorn, + # after the model has loaded; read it here too, so that a malformed value + # fails before the model loads. + _graceful_shutdown_timeout_from_env() + # Publish an explicit serving role before constructing/spawning the model. # Remote-tail context workers use it to omit generation-only weights. if server_role is not None: @@ -890,6 +896,10 @@ def launch_visual_gen_server( visual_gen_args: Optional validated VisualGenArgs for model configuration. metadata_server_cfg: Optional metadata server configuration. """ + # See launch_server: read the graceful-shutdown timeout before any rank + # builds VisualGen, so that a malformed value fails before the model loads. + _graceful_shutdown_timeout_from_env() + # Only global rank 0 serves HTTP; in external multi-rank launch ranks > 0 # must become workers before binding, else every rank on a multi-GPU node # races the same port and all but one die EADDRINUSE. VisualGen() on a diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 514aa4e26f05..9c89cba49e4e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5716,10 +5716,22 @@ class DwdpConfig(StrictBaseModel): default=1, description= "The number of DWDP groups. Total workers = num_groups * dwdp_size.") - num_experts_per_worker: int = Field( - default=0, description="The number of experts per worker.") - num_prefetch_experts: int = Field( - default=0, description="The number of prefetch experts per worker.") + num_experts_per_worker: NonNegativeInt = Field( + default=0, + description= + "The number of routed experts each worker stores in a strided expert " + "partition, where worker r stores the experts from " + "r * num_prefetch_experts. Leave both fields at 0 (the default) for " + "the balanced partition: every worker stores floor or ceil of " + "num_experts / dwdp_size contiguous experts, and no expert is stored " + "twice.") + num_prefetch_experts: NonNegativeInt = Field( + default=0, + description= + "The stride between the first experts of consecutive workers in a " + "strided expert partition (not a number of prefetched experts). Equal " + "to num_experts_per_worker for disjoint ranges, smaller for ranges " + "that overlap. Leave both fields at 0 for the balanced partition.") contention_opt: bool = Field( default=False, description= @@ -5735,6 +5747,22 @@ class DwdpConfig(StrictBaseModel): "further ahead of the MoE compute, at the cost of device memory that " "is then not available to the KV cache.") + @model_validator(mode="after") + def _validate_expert_partition(self) -> 'DwdpConfig': + """Both partition fields 0 (balanced) or both positive (strided). + + Whether a strided pair fits the model's routed experts is checked + when the model is built, where their count is known. + """ + size = self.num_experts_per_worker + stride = self.num_prefetch_experts + if (size == 0) != (stride == 0): + raise ValueError( + f"DwdpConfig: num_experts_per_worker={size} and " + f"num_prefetch_experts={stride}: set both (a strided expert " + f"partition) or leave both at 0 (the balanced partition).") + return self + class BaseLlmArgs(StrictBaseModel): """Base class for the LLM arguments. It contains all the common arguments.""" diff --git a/tensorrt_llm/mapping.py b/tensorrt_llm/mapping.py index 8a38434438b3..111106064ad9 100644 --- a/tensorrt_llm/mapping.py +++ b/tensorrt_llm/mapping.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -99,10 +99,11 @@ def __init__( # see a single (full) expert table; the DWDP-specific layout # (``slot_start`` / ``slot_end`` / ``expert_size_per_partition``) is # installed by ``_init_dwdp_expert_layout`` at the end of - # ``MoE.__init__`` using ``DwdpManager.start_expert_id`` and - # ``num_experts_per_worker``. This is what unlocks dwdp_size values - # that do not divide ``num_experts`` evenly (e.g. dwdp_size = 3, 5) - # and the IPC-era redundancy mode where adjacent peer ranges overlap. + # ``MoE.__init__`` from this rank's range of the DWDP group's + # ``ExpertPartition`` (``DwdpManager.expert_partition``): the balanced + # partition of any dwdp_size (e.g. 8 x 75 + 4 x 74 experts over + # dwdp_size = 12), a strided one with a capped last range, or a + # redundant one whose adjacent ranges overlap. # Record whether the caller explicitly requested a MoE TP/EP split # (the -1 sentinels mean "auto"). Model code (e.g. Kimi K3) uses this # to distinguish an explicit ``moe_tensor_parallel_size`` / diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index b2ca2567a6fe..57daa99b894e 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -1836,6 +1836,32 @@ }, "DwdpConfig": { "additionalProperties": false, + "anyOf": [ + { + "properties": { + "num_experts_per_worker": { + "const": 0 + }, + "num_prefetch_experts": { + "const": 0 + } + } + }, + { + "properties": { + "num_experts_per_worker": { + "minimum": 1 + }, + "num_prefetch_experts": { + "minimum": 1 + } + }, + "required": [ + "num_experts_per_worker", + "num_prefetch_experts" + ] + } + ], "description": "Configuration for Distributed Weight Data Parallelism (DWDP).\n\nDWDP accelerates the context (prefill) phase of disaggregated MoE serving\nby combining data parallelism with NVLink-based expert weight sharing.\nEach worker holds a subset of experts locally and asynchronously prefetches\nthe remaining experts from peer workers via CUDA VMM + MNNVL fabric\nhandles (composite virtual-address layout with zero-copy local mapping and\nP2P remote prefetch slots), enabling fully asynchronous execution\nacross ranks without synchronization barriers.\n\nCurrently supported with the CuteDSL MoE backend and NVFP4 quantization\non MNNVL-connected multi-GPU systems (e.g., GB200).", "properties": { "contention_opt": { @@ -1852,7 +1878,8 @@ }, "num_experts_per_worker": { "default": 0, - "description": "The number of experts per worker.", + "description": "The number of routed experts each worker stores in a strided expert partition, where worker r stores the experts from r * num_prefetch_experts. Leave both fields at 0 (the default) for the balanced partition: every worker stores floor or ceil of num_experts / dwdp_size contiguous experts, and no expert is stored twice.", + "minimum": 0, "title": "Num Experts Per Worker", "type": "integer" }, @@ -1864,7 +1891,8 @@ }, "num_prefetch_experts": { "default": 0, - "description": "The number of prefetch experts per worker.", + "description": "The stride between the first experts of consecutive workers in a strided expert partition (not a number of prefetched experts). Equal to num_experts_per_worker for disjoint ranges, smaller for ranges that overlap. Leave both fields at 0 for the balanced partition.", + "minimum": 0, "title": "Num Prefetch Experts", "type": "integer" }, diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 704cffba5538..7d7c5e38b70c 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -3400,6 +3400,32 @@ }, "DwdpConfig": { "additionalProperties": false, + "anyOf": [ + { + "properties": { + "num_experts_per_worker": { + "const": 0 + }, + "num_prefetch_experts": { + "const": 0 + } + } + }, + { + "properties": { + "num_experts_per_worker": { + "minimum": 1 + }, + "num_prefetch_experts": { + "minimum": 1 + } + }, + "required": [ + "num_experts_per_worker", + "num_prefetch_experts" + ] + } + ], "description": "Configuration for Distributed Weight Data Parallelism (DWDP).\n\nDWDP accelerates the context (prefill) phase of disaggregated MoE serving\nby combining data parallelism with NVLink-based expert weight sharing.\nEach worker holds a subset of experts locally and asynchronously prefetches\nthe remaining experts from peer workers via CUDA VMM + MNNVL fabric\nhandles (composite virtual-address layout with zero-copy local mapping and\nP2P remote prefetch slots), enabling fully asynchronous execution\nacross ranks without synchronization barriers.\n\nCurrently supported with the CuteDSL MoE backend and NVFP4 quantization\non MNNVL-connected multi-GPU systems (e.g., GB200).", "properties": { "contention_opt": { @@ -3416,7 +3442,8 @@ }, "num_experts_per_worker": { "default": 0, - "description": "The number of experts per worker.", + "description": "The number of routed experts each worker stores in a strided expert partition, where worker r stores the experts from r * num_prefetch_experts. Leave both fields at 0 (the default) for the balanced partition: every worker stores floor or ceil of num_experts / dwdp_size contiguous experts, and no expert is stored twice.", + "minimum": 0, "title": "Num Experts Per Worker", "type": "integer" }, @@ -3428,7 +3455,8 @@ }, "num_prefetch_experts": { "default": 0, - "description": "The number of prefetch experts per worker.", + "description": "The stride between the first experts of consecutive workers in a strided expert partition (not a number of prefetched experts). Equal to num_experts_per_worker for disjoint ranges, smaller for ranges that overlap. Leave both fields at 0 for the balanced partition.", + "minimum": 0, "title": "Num Prefetch Experts", "type": "integer" }, diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index b01445d2a086..ed045d053764 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -213,6 +213,38 @@ async def route_handler(request: Request): TIMEOUT_KEEP_ALIVE = 5 # seconds. +GRACEFUL_SHUTDOWN_TIMEOUT_ENV = "TRTLLM_SERVE_GRACEFUL_SHUTDOWN_TIMEOUT_S" + + +def _graceful_shutdown_timeout_from_env() -> int | None: + """Return the uvicorn graceful-shutdown bound requested by the environment. + + A positive integer N in ``TRTLLM_SERVE_GRACEFUL_SHUTDOWN_TIMEOUT_S`` makes + uvicorn wait at most N seconds, after a shutdown signal, for open + connections and running request handlers. Handlers still running then are + cancelled, and the application lifespan shutdown runs as usual. When the + variable is unset or empty, uvicorn waits for every handler to finish. + + Returns: + The bound in seconds, or None when the variable is unset or empty. + + Raises: + ValueError: The variable is set to anything but a positive integer. + """ + raw = os.environ.get(GRACEFUL_SHUTDOWN_TIMEOUT_ENV) + if not raw: + return None + try: + seconds = int(raw) + except ValueError: + seconds = 0 + if seconds <= 0: + raise ValueError( + f"{GRACEFUL_SHUTDOWN_TIMEOUT_ENV}={raw!r} is invalid: expected a " + "positive integer number of seconds") + return seconds + + _warned_unresolvable_thinking = False @@ -4329,11 +4361,21 @@ async def __call__(self, self.binding_addr = f"http://{host}:{port}" self.host = host self.port = port + # Unless GRACEFUL_SHUTDOWN_TIMEOUT_ENV is set, uvicorn keeps its default + # and waits for every request handler during graceful shutdown. + config_kwargs = {} + shutdown_timeout = _graceful_shutdown_timeout_from_env() + if shutdown_timeout is not None: + config_kwargs["timeout_graceful_shutdown"] = shutdown_timeout + logger.info(f"{GRACEFUL_SHUTDOWN_TIMEOUT_ENV}={shutdown_timeout}: " + "uvicorn graceful shutdown cancels request handlers " + f"still running after {shutdown_timeout}s") config = uvicorn.Config(self.app, host=host, port=port, log_level="info", - timeout_keep_alive=TIMEOUT_KEEP_ALIVE) + timeout_keep_alive=TIMEOUT_KEEP_ALIVE, + **config_kwargs) server = create_uvicorn_server(config) async def _register_after_serving(): diff --git a/tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py b/tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py index 662e60e94d77..b12142890f16 100644 --- a/tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_dwdp_disaggregated_serving.py @@ -423,10 +423,10 @@ def test_dwdp_accuracy_mode_b_overlap(self): * peer ranges ``[(0, 40), (32, 72)]`` have a non-empty 8-expert overlap, exercising ``lookup_owner``'s first-match policy in fill_edge_bytes / weight_manager / scatter; - * ``_scatter_shards_to_full`` writes overlapping experts - (32..39) twice — once from rank 0 (lowest, wins) and once - from rank 1 — and values must agree because every rank that - owns expert ``e`` loaded ``e`` from the same checkpoint. + * ``_exchange_expert_scalars`` puts the overlapping experts + (32..39) together from both ranks' shards and requires the two + copies to be equal, which holds because every rank that owns + expert ``e`` loaded ``e`` from the same checkpoint. Passing GSM8K at the standard threshold is sufficient evidence that the runtime fetches every remote expert from the right @@ -477,9 +477,9 @@ def test_dwdp_accuracy_mode_b_overlap(self): "dwdp_size": 2, "num_groups": 1, # Mode B redundancy: stride=32 < size=40. Strict - # equality 1*32 + 40 == 72 is enforced by - # _validate_partition_config; 8-expert overlap between - # rank 0's range [0, 40) and rank 1's range [32, 72). + # equality 1*32 + 40 == 72 is required by + # ExpertPartition.strided (redundant); 8-expert overlap + # between rank 0's range [0, 40) and rank 1's range [32, 72). "num_experts_per_worker": 40, "num_prefetch_experts": 32, }, diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 1bce3cdda018..d2f78a09d80b 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -42,6 +42,7 @@ l0_a10: - unittest/_torch/executor/kv_cache/test_kv_cache_v2_capacity_only.py - unittest/_torch/executor/kv_cache/sharing - unittest/_torch/executor/kv_cache/test_kv_cache_v2_branch_snapshot.py + - unittest/_torch/executor/kv_cache/test_kv_cache_v2_adp_dummy_draft_rollback.py - unittest/_torch/executor/test_kv_cache_layout.py - unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py - unittest/_torch/executor/test_error_classification.py diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 31527968a7dd..6cced3e3e13b 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -250,6 +250,8 @@ l0_b200: # DWDP expert prefetch on real streams and copies; the peer-GPU gather # cases skip below 2 GPUs. - unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py + # Attached dense DWDP beside the expert prefetch on real streams and copies. + - unittest/_torch/modules/dwdp/test_dense_attached_gpu.py # ------------- MoE components tests --------------- - unittest/_torch/moe/fused_moe/test_cutedsl_fc2_tuning.py - unittest/_torch/moe/test_cutedsl_fc1_weight_layout.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index b3e075c3ce14..66c696529a1b 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -240,6 +240,7 @@ l0_cpu: - unittest/llmapi/apps/test_openai_chat_disagg_multimodal.py - unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py - unittest/llmapi/apps/test_openai_protocol_chat_store.py + - unittest/llmapi/apps/test_openai_server_graceful_shutdown.py - unittest/llmapi/apps/test_openai_server_iteration_stats.py - unittest/llmapi/apps/test_reasoning_prompt_resolution.py - unittest/llmapi/apps/test_render_callers.py diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml index cd7500c03d99..51d23d218627 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml @@ -58,6 +58,8 @@ l0_gb300_multi_gpus: # DWDP expert prefetch on real streams and copies, incl. the gather from # peer GPUs. - unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py + # Attached dense DWDP beside the expert prefetch on real streams and copies. + - unittest/_torch/modules/dwdp/test_dense_attached_gpu.py # ------------- MoE components tests --------------- - unittest/_torch/moe/test_moe_load_balancer.py - unittest/_torch/moe/test_moe_routing.py diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_adp_dummy_draft_rollback.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_adp_dummy_draft_rollback.py new file mode 100644 index 000000000000..5b870a19a971 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_adp_dummy_draft_rollback.py @@ -0,0 +1,217 @@ +# 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. +"""A skipped iteration must not strand the attention-DP dummy's draft KV mirror. + +``KVCacheV2Scheduler`` admits an unpaired draft pool's mirror while it schedules +a generation request (``_admit_unpaired_draft`` -> ``admit_mirror``). For the +attention-DP pad dummy that is a fresh, empty draft cache under the reused +``ATTENTION_DP_DUMMY_REQUEST_ID``. When a peer rank schedules nothing, the +iteration is not queued and ``PyExecutor._finalize_adp_dummy_allocation`` rolls +the dummy back before ``prepare_resources``. If the draft mirror survives that +rollback, the deferred previous batch -- which still holds the earlier, +finished dummy with the same ID -- is processed on the next queued iteration, +and the draft pool's ``update_resources`` resizes the stranded empty cache to +that dummy's history: ``ValueError: History length cannot exceed capacity``. +That is the crash of a disaggregated DSpark generation server (separate draft +KV cache, overlap scheduler, attention DP) one iteration after a skipped one. +""" + +import gc +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +import tensorrt_llm.bindings +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import ( + ATTENTION_DP_DUMMY_REQUEST_ID, + LlmRequestState, +) +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManager, ResourceManagerType +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler +from tensorrt_llm.llmapi.llm_args import ( + CapacitySchedulerPolicy, + DFlashDecodingConfig, + KvCacheConfig, +) +from tensorrt_llm.mapping import Mapping + +DataType = tensorrt_llm.bindings.DataType +CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType + +DRAFT_LEN = 5 + + +def _executor(**kwargs): + """The attributes ``_finalize_adp_dummy_allocation`` and ``_revert_gen_alloc`` read.""" + defaults = dict( + _pending_adp_dummy_request=None, + _is_kv_manager_v2=True, + enable_joint_kv_cache_reuse=False, + active_requests=[], + ) + defaults.update(kwargs) + return SimpleNamespace(**defaults) + + +@pytest.mark.cpu_only +def test_rollback_releases_the_dummy_draft_mirror(): + dummy = SimpleNamespace( + py_request_id=ATTENTION_DP_DUMMY_REQUEST_ID, + py_seq_slot=None, + state=LlmRequestState.GENERATION_IN_PROGRESS, + ) + resource_manager = Mock() + resource_manager.get_resource_manager.return_value = None + executor = _executor( + _pending_adp_dummy_request=dummy, + kv_cache_manager=Mock(), + draft_kv_cache_manager=Mock(), + resource_manager=resource_manager, + active_requests=[dummy], + ) + + PyExecutor._finalize_adp_dummy_allocation(executor, False) + + assert dummy.state == LlmRequestState.GENERATION_COMPLETE + assert executor.active_requests == [] + executor.kv_cache_manager.free_resources.assert_called_once_with(dummy) + executor.draft_kv_cache_manager.free_resources.assert_called_once_with(dummy) + + +@pytest.mark.cpu_only +def test_queued_dummy_keeps_its_draft_mirror(): + dummy = SimpleNamespace(py_request_id=ATTENTION_DP_DUMMY_REQUEST_ID, py_seq_slot=None) + executor = _executor( + _pending_adp_dummy_request=dummy, + kv_cache_manager=Mock(), + draft_kv_cache_manager=Mock(), + resource_manager=Mock(), + active_requests=[dummy], + ) + + PyExecutor._finalize_adp_dummy_allocation(executor, True) + + assert executor.active_requests == [dummy] + executor.kv_cache_manager.free_resources.assert_not_called() + executor.draft_kv_cache_manager.free_resources.assert_not_called() + + +def _manager(*, is_draft: bool) -> KVCacheManagerV2: + return KVCacheManagerV2( + KvCacheConfig(enable_block_reuse=False, max_gpu_total_bytes=4 << 20), + CacheType.SELF, + num_layers=1, + num_kv_heads=2, + head_dim=64, + tokens_per_block=32, + max_seq_len=4096, + max_batch_size=4, + max_num_tokens=128, + mapping=Mapping(), + dtype=DataType.HALF, + spec_config=DFlashDecodingConfig(max_draft_len=DRAFT_LEN, speculative_model="draft"), + is_draft=is_draft, + is_disagg=True, + ) + + +@pytest.fixture +def pools(): + torch.cuda.init() + target = _manager(is_draft=False) + draft = _manager(is_draft=True) + yield target, draft + draft.shutdown() + target.shutdown() + gc.collect() + torch.cuda.empty_cache() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="allocates real KV cache pools") +def test_skipped_iteration_does_not_strand_the_dummy_draft_mirror(pools): + """Replay the overlap loop on one attention-DP rank with real KV pools. + + J: the rank is idle, pads a dummy, every rank queues, the dummy runs and is + terminated; batch J becomes the previous batch. + K: the rank pads a new dummy, but a peer scheduled nothing, so can_queue is + False: the dummy is rolled back, and batch J stays deferred because this + rank's own batch was not empty. + K+1: the rank has real work and pads no dummy, so the deferred batch J is + processed (_process_previous_batch skips the V2 target manager there). + """ + target, draft = pools + resources = ResourceManager( + { + ResourceManagerType.KV_CACHE_MANAGER: target, + ResourceManagerType.DRAFT_KV_CACHE_MANAGER: draft, + } + ) + scheduler = KVCacheV2Scheduler( + max_batch_size=4, + max_num_tokens=128, + kv_cache_manager=target, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + draft_kv_cache_manager=draft, + ) + # Without block reuse the draft pool is unpaired, as on the DSpark + # generation server; the scheduler admits its mirror on its own. + assert scheduler._joint_draft_manager is None + executor = _executor( + kv_cache_manager=target, + draft_kv_cache_manager=draft, + resource_manager=resources, + ) + + def pad_and_schedule(): + # _pad_attention_dp_dummy_request, then _schedule(). + dummy = target.add_dummy_requests( + [ATTENTION_DP_DUMMY_REQUEST_ID], is_gen=True, max_num_draft_tokens=DRAFT_LEN + )[0] + dummy.is_attention_dp_dummy = True + executor.active_requests.append(dummy) + executor._pending_adp_dummy_request = dummy + output = scheduler.schedule_request(executor.active_requests, set()) + assert output.generation_requests == [dummy] + batch = ScheduledRequests() + batch.generation_requests = list(output.generation_requests) + return dummy, batch + + # Iteration J. + dummy_j, batch_j = pad_and_schedule() + resources.prepare_resources(batch_j) + PyExecutor._finalize_adp_dummy_allocation(executor, True) + # _update_request_states -> _terminate_attention_dp_dummy_requests. + dummy_j.state = LlmRequestState.GENERATION_COMPLETE + resources.free_resources(dummy_j) + executor.active_requests.remove(dummy_j) + assert ATTENTION_DP_DUMMY_REQUEST_ID not in draft.kv_cache_map + + # Iteration K: not queued. + _, batch_k = pad_and_schedule() + assert ATTENTION_DP_DUMMY_REQUEST_ID in draft.kv_cache_map + PyExecutor._revert_gen_alloc(executor, batch_k) + PyExecutor._finalize_adp_dummy_allocation(executor, False) + assert ATTENTION_DP_DUMMY_REQUEST_ID not in target.kv_cache_map + + # Iteration K+1: _process_previous_batch on the deferred batch J. + resources.update_resources(batch_j, skip_kv_cache_manager=True) + + assert ATTENTION_DP_DUMMY_REQUEST_ID not in draft.kv_cache_map + assert executor.active_requests == [] diff --git a/tests/unittest/_torch/executor/test_executor_shutdown_drop_inflight.py b/tests/unittest/_torch/executor/test_executor_shutdown_drop_inflight.py new file mode 100644 index 000000000000..287eca9c2b3e --- /dev/null +++ b/tests/unittest/_torch/executor/test_executor_shutdown_drop_inflight.py @@ -0,0 +1,230 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""A requested shutdown that drops in-flight requests instead of draining them. + +``TRTLLM_EXECUTOR_SHUTDOWN_DROP_INFLIGHT=1`` lets the executor loop stop while +requests are still active or queued. Without it every case keeps the drain: +the loop stops only once nothing is left and no connector transfer is pending. +""" + +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock + +import pytest + +from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, _shutdown_drop_inflight_enabled +from tensorrt_llm._torch.pyexecutor.scheduler import create_waiting_queue +from tensorrt_llm.logger import logger + +pytestmark = pytest.mark.cpu_only + +# The name deployments set. It is spelled out rather than imported, so renaming the +# executor's variable fails these tests instead of silently leaving the drain on. +_ENV = "TRTLLM_EXECUTOR_SHUTDOWN_DROP_INFLIGHT" + +# Leaves `_drop_inflight_on_shutdown` unset, as on executors assembled without `__init__`. +_UNSET = object() + + +class _KeptDraining(Exception): + """Raised by the first step after the stop check: the loop kept going.""" + + +def _executor(*, drop_inflight=_UNSET, active=0, queued=0, connector=None) -> PyExecutor: + """An executor asked to shut down, assembled without ``__init__``. + + ``connector`` is None (no connector), "idle", "load" (a prefix load in flight) + or "transfer" (a save in flight). + """ + executor = object.__new__(PyExecutor) + if drop_inflight is not _UNSET: + executor._drop_inflight_on_shutdown = drop_inflight + executor.is_shutdown = True + executor.active_requests = [ + SimpleNamespace(py_request_id=i, state=LlmRequestState.CONTEXT_INIT) for i in range(active) + ] + executor.waiting_queue = create_waiting_queue() + executor.waiting_queue.add_requests(RequestQueueItem(100 + i) for i in range(queued)) + executor.enable_attention_dp = False + executor.kv_connector_manager = None + if connector is not None: + executor.kv_connector_manager = Mock( + **{"has_pending_loads.return_value": connector == "load"} + ) + executor.async_transfer_manager = Mock( + **{"has_any_inflight_requests.return_value": connector == "transfer"} + ) + # The termination path, down to the resource managers. + executor.kv_cache_transceiver = None + executor._disagg_pp_termination_handler = None + executor._is_kv_manager_v2 = False + executor.sampler = Mock() + executor.resource_manager = Mock() + executor._prefetched_request_ids = set() + executor._disagg_coordinator = Mock() + executor.model_engine = None + executor.draft_model_engine = None + executor.gather_all_responses = False + executor.dist = SimpleNamespace(rank=0) + executor.result_wait_queues = {} + return executor + + +@pytest.fixture +def log_calls(monkeypatch) -> Mock: + """Capture what the TRT-LLM logger emits, with every log-once key unused.""" + calls = Mock() + monkeypatch.setattr(logger, "log", calls) + monkeypatch.setattr(logger, "_appeared_keys", set()) + return calls + + +def _messages(log_calls: Mock, severity: str) -> list[str]: + return [ + " ".join(map(str, c.args[1:])) for c in log_calls.call_args_list if c.args[0] == severity + ] + + +@pytest.mark.parametrize( + ("raw", "pp_size", "disable_overlap_scheduler", "enabled", "logged"), + [ + (None, 1, True, False, None), + ("1", 1, True, True, "=1: dropping in-flight requests at shutdown is on"), + ("1", 2, True, False, "=1: dropping in-flight requests at shutdown is off (needs pp_size"), + ("1", 1, False, False, "=1: dropping in-flight requests at shutdown is off (needs pp_size"), + ("true", 1, True, False, "='true': dropping in-flight requests at shutdown is off (only"), + ], +) +def test_gate_needs_one_pipeline_stage_without_the_overlap_scheduler( + monkeypatch, log_calls, raw, pp_size, disable_overlap_scheduler, enabled, logged +) -> None: + """A pipeline-parallel or overlap loop keeps draining whatever the variable says. + + A PP loop can stop with microbatches still travelling between ranks, and the + overlap loop would skip processing its previous batch. A set variable logs the + value that takes effect, "1" bare and any other value quoted; unset is silent. + """ + if raw is None: + monkeypatch.delenv(_ENV, raising=False) + else: + monkeypatch.setenv(_ENV, raw) + + assert _shutdown_drop_inflight_enabled(pp_size, disable_overlap_scheduler) is enabled + + if raw is None: + log_calls.assert_not_called() + else: + (message,) = _messages(log_calls, logger.INFO) + assert message.startswith(_ENV + logged), message + + +@pytest.mark.parametrize("connector", [None, "idle", "load", "transfer"]) +@pytest.mark.parametrize(("active", "queued"), [(0, 0), (2, 0), (0, 3)]) +@pytest.mark.parametrize("is_shutdown", [False, True]) +@pytest.mark.parametrize("drop_inflight", [_UNSET, False, True]) +def test_a_requested_shutdown_drops_in_flight_work_only_when_set( + log_calls, drop_inflight, is_shutdown, active, queued, connector +) -> None: + """Set, work still active or queued no longer holds the loop; unset, it is drained. + + A connector load or save reads the KV pools that shutdown frees next, so it is + waited for either way. Dropping is warned once, with the counts. + """ + executor = _executor( + drop_inflight=drop_inflight, active=active, queued=queued, connector=connector + ) + executor.is_shutdown = is_shutdown + dropping = drop_inflight is True and active + queued > 0 + stops = is_shutdown and connector in (None, "idle") and (dropping or active + queued == 0) + + assert executor.should_stop_processing is stops + assert executor.should_stop_processing is stops + warnings = _messages(log_calls, logger.WARNING) + if stops and dropping: + (warning,) = warnings + assert warning.startswith( + f"shutdown requested: dropping {active} active and {queued} queued" + ) + assert f"({_ENV}=1)" in warning + else: + assert warnings == [] + + +@pytest.mark.parametrize("drop_inflight", [_UNSET, False, True]) +def test_dropped_requests_release_through_the_termination_path(drop_inflight) -> None: + """Each dropped request hands its KV cache and slots back before shutdown frees the pools. + + A request whose KV cache the transfer fabric may still read keeps it, as on + cancellation; the KV cache manager closes that cache at its own shutdown. + Unset, the release leaves everything in place, e.g. after the insufficient-KV + benchmark exit. + """ + executor = _executor(drop_inflight=drop_inflight, active=3, queued=2) + requests = list(executor.active_requests) + released, in_transfer_state, pipelined = requests + in_transfer_state.state = LlmRequestState.DISAGG_CONTEXT_TRANS_IN_PROGRESS + executor.kv_cache_transceiver = Mock() + executor.kv_cache_transceiver.has_inflight_transfer.side_effect = lambda req: req is pipelined + executor.result_wait_queues = {req.py_request_id: Mock() for req in requests} + + executor._release_requests_dropped_at_shutdown() + + if drop_inflight is not True: + assert executor.active_requests == requests and len(executor.waiting_queue) == 2 + executor.resource_manager.free_resources.assert_not_called() + return + executor.resource_manager.free_resources.assert_called_once_with(released) + executor._disagg_coordinator.forget_request.assert_called_once_with(released.py_request_id) + assert set(executor.result_wait_queues) == { + in_transfer_state.py_request_id, + pipelined.py_request_id, + } + assert executor.active_requests == [] and len(executor.waiting_queue) == 0 + + +@pytest.mark.parametrize(("active", "queued"), [(0, 0), (2, 3)]) +@pytest.mark.parametrize("drop_inflight", [_UNSET, False, True]) +def test_the_loop_releases_dropped_requests_after_its_last_collective( + monkeypatch, drop_inflight, active, queued +) -> None: + """The release follows the loop's final response flush and its completion mark. + + The flush is the loop's last collective. Releasing after it keeps the release + rank-local, and an error raised by it is reported as teardown rather than as a + crash that strands peer ranks. Unset, pending work keeps the loop going. + """ + for name in ("torch.cuda.set_device", "cudart.cudaSetDevice", "CUASSERT"): + monkeypatch.setattr(f"tensorrt_llm._torch.pyexecutor.py_executor.{name}", Mock()) + executor = _executor(drop_inflight=drop_inflight, active=active, queued=queued) + executor.is_shutdown = False + + def fetch_the_shutdown_request(): + executor.is_shutdown = True + return [] + + executor._fetch_and_activate_new_requests = Mock(side_effect=fetch_the_shutdown_request) + executor._handle_control_request = Mock(side_effect=_KeptDraining) + order = [] + executor._flush_pending_transfer_responses = Mock(side_effect=lambda: order.append("flush")) + executor.resource_manager.free_resources.side_effect = lambda req: order.append( + (req.py_request_id, executor._event_loop_completed) + ) + executor.device_id = 0 + executor._profiler = lambda: nullcontext(lambda: None) + executor.hang_detector = MagicMock() + executor.enable_iter_perf_stats = False + executor._resource_governor_enabled = False + executor.is_benchmark_disagg = False + executor._event_loop_completed = False + keeps_draining = active + queued > 0 and drop_inflight is not True + + with pytest.raises(_KeptDraining) if keeps_draining else nullcontext(): + executor._executor_loop() + + assert executor._event_loop_completed is not keeps_draining + assert order == ([] if keeps_draining else ["flush", *[(i, True) for i in range(active)]]) + left = len(executor.active_requests) + len(executor.waiting_queue) + assert left == (active + queued if keeps_draining else 0) diff --git a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py index 36766f504881..5c3899556c2f 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py +++ b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py @@ -791,6 +791,7 @@ def test_shutdown_drains_the_connector_before_the_pools_it_reads_are_freed() -> executor.virtual_memory_pools = None executor.sampler = Mock() executor.dwdp_manager = None + executor.dense_dwdp_manager = None executor.shutdown() diff --git a/tests/unittest/_torch/executor/test_pp_mpi_progress_thread.py b/tests/unittest/_torch/executor/test_pp_mpi_progress_thread.py index 17c6b3c4f0c5..bbbbb0691c03 100644 --- a/tests/unittest/_torch/executor/test_pp_mpi_progress_thread.py +++ b/tests/unittest/_torch/executor/test_pp_mpi_progress_thread.py @@ -204,6 +204,7 @@ def _executor(monkeypatch, pp_size: int) -> PyExecutor: executor.resource_manager = SimpleNamespace(resource_managers={}) executor.virtual_memory_pools = None executor.dwdp_manager = None + executor.dense_dwdp_manager = None return executor diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index b12ca615f74d..bbf872830d79 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -130,6 +130,7 @@ def test_shutdown_graph_lifecycle(shutdown_method: str, releases_model_owned_gra executor.virtual_memory_pools = None executor.sampler = object() executor.dwdp_manager = None + executor.dense_dwdp_manager = None executor._profile_manager = Mock() with patch("torch.cuda.is_available", return_value=False): diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_attached_gpu.py b/tests/unittest/_torch/modules/dwdp/test_dense_attached_gpu.py new file mode 100644 index 000000000000..0c7d8e362aae --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_attached_gpu.py @@ -0,0 +1,313 @@ +# 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. +"""Attached dense DWDP beside the expert prefetch on real streams, events and copies. + +A real ``DenseDWDP`` runtime, with staging slots from ``allocate_slot_ring`` and a launcher that +copies each layer's source bytes into the layer's slots, and a real ``DWDPWeightManager``, over a +stand-in for ``WeightBuffer`` whose prefetch slots hold the remote experts, share one copy stream. +The attached ``DenseDwdpManager`` holds both. A harness drives them the way the executor loop +(``DwdpManager.start_forward``), ``KimiLinearModel`` (H0-H3) and ``ConfigurableMoE`` do, with a +read kernel in every read window: the attention weights between H1 and H2, the MoE-dense weights +before the expert bind and after the expert finish, the experts between the bind and the finish. +Delaying the reads exposes a missing WAR wait (a copy overwrites a slot that is still being read), +delaying the copies a missing RAW wait. Every read is checked against the bytes of its layer. +""" + +import itertools +import time +from typing import Dict, List, Optional, Tuple + +import pytest +import torch + +from tensorrt_llm._torch.modules.dwdp.dense.manager import DenseDwdpManager +from tensorrt_llm._torch.modules.dwdp.dense.runtime import DenseDWDP +from tensorrt_llm._torch.modules.dwdp.dense.types import ( + RING_DEPTH, + DenseMode, + SegmentKind, + SlotSpec, + allocate_slot_ring, +) +from tensorrt_llm._torch.modules.dwdp.weight_manager import DWDPWeightManager +from tensorrt_llm._torch.pyexecutor.dwdp import DwdpManager + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +ATTN = SegmentKind.ATTN +MOE = SegmentKind.MOE +BOTH = (ATTN, MOE) +NUM_EXPERTS = 4 # experts per rank +DELAY_CYCLES = 400_000 +ABORT_POINTS = ("H1", "H2", "finish", "H3") + + +class _ExpertSlots: + """Stand-in for ``WeightBuffer`` of DWDP rank 0 of two: rows [0, NUM_EXPERTS) of a layer's + full tensor play the (unused) local experts, the other rows are its slot's remote experts.""" + + local_start = 0 + local_end = NUM_EXPERTS + + def __init__(self, layers: List[int], num_slots: int, dim: int) -> None: + self.device_id = torch.cuda.current_device() + self.num_slots = num_slots + self._layers = layers + self.slots = [ + torch.full((2 * NUM_EXPERTS, dim), -1.0, dtype=torch.float64, device="cuda") + for _ in range(num_slots) + ] + + def buffer_index_for_layer(self, layer: int) -> int: + return self._layers.index(layer) % self.num_slots + + def get_full_tensor(self, layer: int, name: str) -> torch.Tensor: + return self.slots[self.buffer_index_for_layer(layer)] + + def get_remote_slices(self, layer: int, name: str) -> List[Tuple[torch.Tensor, int, int]]: + full = self.get_full_tensor(layer, name) + return [(full[NUM_EXPERTS:], NUM_EXPERTS, full.shape[0])] + + +class _Backend(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.w = torch.nn.Parameter( + torch.empty(0, dtype=torch.float64, device="cuda"), requires_grad=False + ) + + +class _Abort(Exception): + pass + + +class _Harness: + """Attached dense DWDP and expert DWDP of ``layers``, driven in Kimi-K3's program order. + + Args: + layers: Covered layers, which are also the MoE layers. + segments: Enabled dense segments. + num_slots: Expert prefetch slots. + dense_bytes: Bytes of each layer's weights per dense segment. + dim: Columns of each expert row. + copy_delay: Cycles each copy waits on the copy stream before it starts. + """ + + def __init__( + self, + layers: List[int], + segments: Tuple[SegmentKind, ...], + num_slots: int, + dense_bytes: int, + dim: int, + copy_delay: int = 0, + ) -> None: + self.layers = layers + self.segments = segments + self.dense_bytes = dense_bytes + self.copy_delay = copy_delay + self.copies: List[Tuple[str, int]] = [] + self.dense_src: Dict[Tuple[SegmentKind, int], torch.Tensor] = { + (kind, layer): torch.randint(0, 256, (dense_bytes,), dtype=torch.uint8, device="cuda") + for kind in segments + for layer in layers + } + self.expert_src: Dict[int, torch.Tensor] = { + layer: torch.randint(0, 64, (NUM_EXPERTS, dim), device="cuda").to(torch.float64) + for layer in layers + } + self.buffer = _ExpertSlots(layers, num_slots, dim) + self.wm = DWDPWeightManager( + weight_buffer=self.buffer, + peer_views={(1, layer, "w"): self.expert_src[layer] for layer in layers}, + peer_ranges=[(0, NUM_EXPERTS), (NUM_EXPERTS, 2 * NUM_EXPERTS)], + moe_layer_indices=layers, + weight_names=["w"], + dwdp_rank=0, + dwdp_size=2, + ) + slots = { + kind: allocate_slot_ring(SlotSpec(kind, RING_DEPTH[kind], dense_bytes, 1), "cuda") + for kind in segments + } + self.dense = DenseDWDP( + layers, segments, slots, self.wm.copy_stream, prefetch_next_pass=True + ) + self.dense.set_launcher(self._launch_dense) + self.dense.commit(None, fill_byte=0) + self.dense_manager = DenseDwdpManager(DenseMode.ATTACHED) + self.dense_manager.attach(self.dense, expert=self.wm) + self.mgr = object.__new__(DwdpManager) + self.mgr._weight_manager = self.wm + self.mgr._dense_manager = self.dense_manager + self.backend = _Backend() + self.compute_stream = torch.cuda.Stream() + self.results: List[Tuple[str, int, torch.Tensor]] = [] + + def _launch_dense(self, layer: int) -> None: + """The dense launcher: copies ``L(layer)`` into its slots on the current (copy) stream.""" + if self.copy_delay: + torch.cuda._sleep(self.copy_delay) + self.copies.append(("L", layer)) + for kind in self.segments: + slot = self.dense.slots[kind][self.dense.slot_index(kind, layer)] + slot.blob_view(self.dense_bytes).copy_(self.dense_src[(kind, layer)]) + + def _read(self, what: str, layer: int, delay: int) -> None: + """A read kernel in the read window, after ``delay`` cycles of other work.""" + if delay: + torch.cuda._sleep(delay) + if what == "E": + value = self.backend.w[NUM_EXPERTS:].clone() + else: + kind = SegmentKind(what) + slot = self.dense.slots[kind][self.dense.slot_index(kind, layer)] + value = slot.blob_view(self.dense_bytes).clone() + self.results.append((what, layer, value)) + + def run_pass( + self, + read_delay: int = 0, + start: bool = True, + stop: Optional[Tuple[str, int]] = None, + ) -> None: + """One executor step; ``stop=(point, layer)`` raises ``_Abort`` right before that point.""" + if start: + self.mgr.start_forward() + self.compute_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.compute_stream): + self.dense_manager.ensure_primed() # H0 + for layer in self.layers: + point = stop[0] if stop is not None and stop[1] == layer else None + if point == "H1": + raise _Abort() + self.dense.wait_layer(layer) # H1 + self._read(ATTN.value, layer, read_delay) + if point == "H2": + raise _Abort() + self.dense.attention_done(layer) # H2 + if MOE in self.segments: + self._read(MOE.value, layer, read_delay) + self.mgr.wait_and_bind(self.backend, layer) + self._read("E", layer, read_delay) + if point == "finish": + raise _Abort() + self.mgr.finish_layer(layer) + if MOE in self.segments: + self._read(MOE.value, layer, read_delay) + if point == "H3": + raise _Abort() + self.dense.layer_done(layer) # H3 + torch.cuda.current_stream().wait_stream(self.compute_stream) + + def mismatches(self) -> List[Tuple[str, int]]: + """Reads that saw other bytes than their layer's; clears the reads.""" + torch.cuda.synchronize() + bad = [] + for what, layer, value in self.results: + if what == "E": + expected = self.expert_src[layer] + else: + expected = self.dense_src[(SegmentKind(what), layer)] + if not torch.equal(value, expected): + bad.append((what, layer)) + self.results.clear() + return bad + + +def _drive(h: _Harness, read_delay: int) -> None: + """Warmup, steps with idle gaps, every abort point with its recovery, an expert update.""" + layers = h.layers + # A forward the executor does not start: H0 primes, the end of the pass primes the next one. + h.run_pass(read_delay, start=False) + assert h.mismatches() == [] + assert h.dense.primed + for step, delay in enumerate([read_delay, 0, read_delay]): + del h.copies[:] + h.run_pass(delay) + assert h.mismatches() == [], f"step {step}" + # Every layer's dense weights are gathered exactly once per pass. + assert sorted(h.copies) == sorted(("L", layer) for layer in layers), f"step {step}" + time.sleep(0.01) + torch.cuda._sleep(DELAY_CYCLES) + # Aborts: the next pass recovers, with the executor's start or, as in warmup, from H0 alone, + # while the reads the aborted pass enqueued may still run. + stops = itertools.product(ABORT_POINTS, (layers[0], layers[len(layers) // 2], layers[-1])) + for index, stop in enumerate(stops): + with pytest.raises(_Abort): + h.run_pass(read_delay, stop=stop) + h.run_pass(read_delay, start=index % 2 == 0) + assert h.mismatches() == [], stop + # New expert weights: the copies issued ahead hold the old ones until invalidate_prefetch. + torch.cuda.synchronize() + for tensor in h.expert_src.values(): + tensor.add_(1000.0) + h.mgr.invalidate_prefetch() + h.run_pass(read_delay) + h.run_pass(read_delay) + assert h.mismatches() == [] + assert h.dense.demand_fetches == 0 and h.dense.war_fences == 0 + + +@pytest.mark.parametrize("segments", [BOTH, (ATTN,)], ids=["attn+moe", "attn"]) +@pytest.mark.parametrize("num_layers", [4, 5]) +@pytest.mark.parametrize("num_slots", [2, 3]) +@pytest.mark.parametrize("slow", ["read", "copy"]) +def test_dense_and_expert_reads_see_their_layers(segments, num_layers, num_slots, slow): + layers = [2 * i + 1 for i in range(num_layers)] # non-contiguous decoder indices + h = _Harness( + layers, + segments, + num_slots, + dense_bytes=1 << 16, + dim=1024, + copy_delay=DELAY_CYCLES if slow == "copy" else 0, + ) + _drive(h, DELAY_CYCLES if slow == "read" else 0) + + +@pytest.mark.parametrize("segments", [BOTH, (ATTN,)], ids=["attn+moe", "attn"]) +def test_kimi_k3_pass_boundary(segments): + """92 layers, two expert slots: the next pass's first layers are gathered at its end.""" + layers = list(range(1, 93)) + h = _Harness(layers, segments, 2, dense_bytes=1 << 12, dim=64, copy_delay=DELAY_CYCLES // 8) + _drive(h, DELAY_CYCLES // 8) + + +def test_release_right_after_a_pass(): + """The end of a pass leaves the next pass's first layers in flight; release waits for them.""" + h = _Harness([1, 2, 3, 4, 5], BOTH, 2, dense_bytes=1 << 16, dim=256, copy_delay=DELAY_CYCLES) + h.run_pass(DELAY_CYCLES) + assert h.dense.primed and h.copies[-2:] == [("L", 1), ("L", 2)] + h.dense.release() + assert h.mismatches() == [] + assert h.dense.slots == {} + + +def test_missing_dense_consume_record_is_detected(monkeypatch): + """Sensitivity check: without H2's consume record the fill overwrites attention slots in use.""" + original = DenseDWDP._close_window + + def without_attention_record(self, kind, slot, layer, hook): + if kind is ATTN: + self._consumed[kind][slot] = True + return + original(self, kind, slot, layer, hook) + + monkeypatch.setattr(DenseDWDP, "_close_window", without_attention_record) + h = _Harness([1, 2, 3, 4, 5], BOTH, 2, dense_bytes=1 << 16, dim=256) + for _ in range(3): + h.run_pass(DELAY_CYCLES) + assert any(what == ATTN.value for what, _ in h.mismatches()) diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_config.py b/tests/unittest/_torch/modules/dwdp/test_dense_config.py new file mode 100644 index 000000000000..ee5fbd720ce2 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_config.py @@ -0,0 +1,54 @@ +# 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. +"""CPU tests of the dense DWDP switches and how their values are parsed. + +* Booleans: ``1/true/yes/on`` and ``0/false/no/off`` in any case and with blanks around them; + unset or blank keeps the default; any other value logs a warning and keeps the default, so + ``TRTLLM_DWDP_DENSE_REQUIRED`` fails safe (stays on). +* ``TRTLLM_DWDP_DENSE_SEGMENTS`` is a comma-separated list that must include ``attn``. + +The tests pass an ``environ`` dict and do not touch the process environment. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.modules.dwdp.dense import config as dense_config +from tensorrt_llm._torch.modules.dwdp.dense.config import ( + DENSE_REQUIRED_ENV, + DENSE_SEGMENTS_ENV, + DenseConfig, +) + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize( + "value, required, warned", + [(None, True, False), ("", True, False), (" False ", False, False), ("onn", True, True)], +) +def test_required_defaults_on_and_fails_safe(monkeypatch, value, required, warned): + warnings = [] + monkeypatch.setattr(dense_config, "logger", SimpleNamespace(warning=warnings.append)) + environ = {} if value is None else {DENSE_REQUIRED_ENV: value} + assert DenseConfig.from_env(environ).required is required + assert len(warnings) == warned and all(message.endswith("using 1") for message in warnings) + + +@pytest.mark.parametrize("value", ["moe", "attn;moe", ","]) +def test_an_invalid_segment_list_raises(value): + with pytest.raises(ValueError, match=f"{DENSE_SEGMENTS_ENV}={value!r} must list attn"): + DenseConfig.from_env({DENSE_SEGMENTS_ENV: value}) diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_executor.py b/tests/unittest/_torch/modules/dwdp/test_dense_executor.py new file mode 100644 index 000000000000..500147dc2873 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_executor.py @@ -0,0 +1,234 @@ +# 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. +"""``PyExecutor`` and the DWDP managers: which one primes, which ones it accepts, the teardown. + +With expert DWDP the executor holds the ``DwdpManager``, which drives the attached dense DWDP +manager; with standalone dense DWDP it holds the ``DenseDwdpManager`` itself, whose dense group is +a tagged MPI sub-communicator (``manager.standalone_group``) and which reserves the SM gather's +SMs while it is registered. The executor is assembled with ``object.__new__`` and only the +attributes a method reads, as the executor tests do. +""" + +import os +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from torch import nn + +from tensorrt_llm._torch.cute_dsl_utils import sm_reserve +from tensorrt_llm._torch.modules.dwdp import fabric_gather +from tensorrt_llm._torch.modules.dwdp.dense import manager as dense_manager +from tensorrt_llm._torch.modules.dwdp.dense.config import DENSE_COPY_ENV, DENSE_ENV +from tensorrt_llm._torch.modules.dwdp.dense.manager import ( + STANDALONE_GROUP_TAG, + DenseDwdpManager, + get_global_dense_dwdp_manager, + standalone_group, +) +from tensorrt_llm._torch.modules.dwdp.dense.types import DenseMode +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + + +def test_prime_dwdp_primes_a_standalone_dense_manager(): + executor = object.__new__(PyExecutor) + executor.dwdp_manager = None + executor.dense_dwdp_manager = MagicMock() + + executor._prime_dwdp() + + executor.dense_dwdp_manager.prefetch_first_layers.assert_called_once_with() + + +def test_prime_dwdp_starts_expert_dwdp_which_drives_the_attached_dense_manager(): + executor = object.__new__(PyExecutor) + executor.dwdp_manager = MagicMock() + executor.dense_dwdp_manager = None + + executor._prime_dwdp() + + executor.dwdp_manager.start_forward.assert_called_once_with() + + +@pytest.mark.parametrize( + "expert, dense, disable_overlap_scheduler, error", + [ + (False, True, True, None), + (False, True, False, None), + (True, False, False, "DWDP requires disable_overlap_scheduler=True. Overlap scheduler"), + (True, True, True, "pass dense_dwdp_manager only without DWDP"), + ], + ids=["standalone-no_overlap", "standalone-overlap", "expert-overlap", "both"], +) +def test_the_executor_accepts_a_standalone_dense_manager_with_either_loop( + expert, dense, disable_overlap_scheduler, error +): + """Expert DWDP still refuses the overlap scheduler, and drives dense DWDP itself.""" + managers = (MagicMock() if expert else None, MagicMock() if dense else None) + if error is None: + PyExecutor._validate_dwdp_managers(*managers, disable_overlap_scheduler) + else: + with pytest.raises(ValueError, match=error): + PyExecutor._validate_dwdp_managers(*managers, disable_overlap_scheduler) + + +@pytest.mark.parametrize("mode", ["standalone", "attached"]) +def test_shutdown_exits_the_manager_it_holds_after_the_model_is_dropped(monkeypatch, mode): + # ``shutdown`` synchronizes the device when CUDA is available; this executor has none. + monkeypatch.setattr("torch.cuda.is_available", lambda: False) + log = [] + executor = object.__new__(PyExecutor) + manager = MagicMock() + manager.__exit__.side_effect = lambda *_args: log.append( + (f"{mode}.__exit__", hasattr(executor, "model_engine")) + ) + executor.executor_request_queue = SimpleNamespace(enqueue_shutdown_request=lambda: None) + executor.shutdown_event = SimpleNamespace(wait=lambda: None) + executor._profile_manager = SimpleNamespace(cleanup=lambda: None) + executor.hang_detector = SimpleNamespace(detected=lambda: False) + executor.worker_thread = SimpleNamespace(join=lambda: log.append("worker joined")) + executor.enable_attention_dp = False + executor.kv_connector_manager = None + executor._kv_connector_failed = False + executor.dist = SimpleNamespace(pp_size=1) + executor._shutdown_sleep_wakeup_listeners = lambda: None + executor.model_engine = SimpleNamespace(release_executor_resources=lambda **_kwargs: None) + executor.draft_model_engine = None + executor.resource_manager = SimpleNamespace( + resource_managers={"kv": SimpleNamespace(shutdown=lambda: log.append("kv.shutdown"))} + ) + executor.virtual_memory_pools = None + executor.sampler = SimpleNamespace() + executor.dwdp_manager = manager if mode == "attached" else None + executor.dense_dwdp_manager = manager if mode == "standalone" else None + + executor.shutdown() + + assert log == ["worker joined", "kv.shutdown", (f"{mode}.__exit__", False)] + manager.__exit__.assert_called_once_with(None, None, None) + assert executor.dwdp_manager is None and executor.dense_dwdp_manager is None + + +@pytest.mark.parametrize( + "mapping, tp_group", + [ + ( + Mapping(world_size=4, rank=2, tp_size=4, moe_ep_size=4, enable_attention_dp=True), + [0, 1, 2, 3], + ), + ( + Mapping(world_size=8, rank=5, tp_size=4, pp_size=2, enable_attention_dp=True), + [4, 5, 6, 7], + ), + ], + ids=["tp4", "second_pp_stage"], +) +def test_standalone_group_is_a_sub_communicator_over_the_tp_group(monkeypatch, mapping, tp_group): + """The group is created with a tag of its own, not the default tag 0 of ``MPIDist.tp_comm``. + + ``MPIDist.tp_comm`` creates a group over the same ranks of the same parent, and MPI pairs the + ``Create_group`` calls on one parent by their tag. + """ + world = MagicMock() + monkeypatch.setattr(dense_manager, "mpi_disabled", lambda: False) + monkeypatch.setattr(dense_manager, "mpi_comm", lambda: world) + + group = standalone_group(mapping) + + assert mapping.tp_group == tp_group + world.group.Incl.assert_called_once_with(tp_group) + world.Create_group.assert_called_once_with( + world.group.Incl.return_value, tag=STANDALONE_GROUP_TAG + ) + assert group is world.Create_group.return_value + # Not the default tag, and a valid tag under every MPI (MPI_TAG_UB >= 32767). + assert 0 < STANDALONE_GROUP_TAG <= 32767 + + +# --------------------------------------------------------------------------- +# The standalone SM reserve +# --------------------------------------------------------------------------- + + +def _standalone_manager() -> DenseDwdpManager: + mapping = Mapping(world_size=4, rank=1, tp_size=4, moe_ep_size=4, enable_attention_dp=True) + return DenseDwdpManager(DenseMode.STANDALONE, mapping=mapping) + + +@pytest.fixture +def gather_blocks(monkeypatch): + """A resident gather of 30 blocks of 1024 threads; no reserve pinned.""" + monkeypatch.delenv("TRTLLM_CUTEDSL_SM_RESERVE", raising=False) + monkeypatch.setenv("TRTLLM_DWDP_GATHER_BLOCKS", "30") + monkeypatch.setattr(fabric_gather, "THREADS", 1024) + for name in [name for name in os.environ if name.startswith(DENSE_ENV)]: + monkeypatch.delenv(name) + yield 30 + assert sm_reserve() == 0 + + +@pytest.mark.parametrize("copy, reserve", [("gather", 30), ("batched", 0), (None, 0)]) +def test_a_standalone_gather_reserves_its_sms_from_registration_to_exit( + monkeypatch, gather_blocks, copy, reserve +): + """The reserve is set before the model is built, so grids sized at construction see it.""" + if copy is not None: + monkeypatch.setenv(DENSE_COPY_ENV, copy) + manager = _standalone_manager() + assert sm_reserve() == 0 + manager.__enter__() + try: + assert get_global_dense_dwdp_manager() is manager + assert sm_reserve() == reserve + finally: + manager.__exit__(None, None, None) + assert sm_reserve() == 0 + assert get_global_dense_dwdp_manager() is None + + +@pytest.mark.parametrize("copy", ["batched", "per_slice", "gather"]) +def test_a_pinned_reserve_holds_with_every_standalone_transfer(monkeypatch, gather_blocks, copy): + """``TRTLLM_CUTEDSL_SM_RESERVE`` pins the reserve with the copy engines too, as it does for + ``DwdpManager``, from registration to exit; with the gather it replaces the gather's count.""" + monkeypatch.setenv(DENSE_COPY_ENV, copy) + monkeypatch.setenv("TRTLLM_CUTEDSL_SM_RESERVE", "12") + with _standalone_manager(): + assert sm_reserve() == 12 + assert sm_reserve() == 0 + + +def test_an_attached_dense_manager_leaves_the_reserve_to_dwdp_manager(monkeypatch, gather_blocks): + monkeypatch.setenv(DENSE_COPY_ENV, "gather") + with DenseDwdpManager(DenseMode.ATTACHED): + assert sm_reserve() == 0 + + +def test_a_setup_that_commits_nothing_drops_the_reserve(monkeypatch, gather_blocks): + """No dense roots and ``TRTLLM_DWDP_DENSE_REQUIRED=0``: no gather runs, so no SM is held.""" + monkeypatch.setenv(DENSE_COPY_ENV, "gather") + monkeypatch.setenv(f"{DENSE_ENV}_REQUIRED", "0") + warnings = [] + monkeypatch.setattr(dense_manager.logger, "warning", warnings.append) + model = nn.Module() + model.model = nn.Module() + model.model.layers = nn.ModuleList() + with _standalone_manager() as manager: + assert sm_reserve() == 30 + assert manager.setup(model) is None + assert sm_reserve() == 0 + assert any("keep leaving 30 SMs out" in message for message in warnings), warnings diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_gate.py b/tests/unittest/_torch/modules/dwdp/test_dense_gate.py new file mode 100644 index 000000000000..84ec7f7d18ea --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_gate.py @@ -0,0 +1,219 @@ +# 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. +"""CPU tests of the dense DWDP launch gate (``gate.validate_dense_gate``) in both modes. + +The gate returns the mode the launch runs dense DWDP in: attached, with expert DWDP +(``dwdp_config`` with ``dwdp_size > 1``, ``tp_size=1``, the overlap scheduler off), or standalone +on an attention-DP instance (no ``dwdp_config``, ``enable_attention_dp`` and ``tp_size > 1``). It +rejects a launch that fits neither, or that breaks a requirement, with a message naming it. Every +launch starts without ``TRTLLM_DWDP_DENSE*`` variables, with MPI on, FABRIC peer handles and a +registered gather op; a launch changes only what it checks. +""" + +import os +import re +from types import SimpleNamespace +from typing import NamedTuple + +import pytest + +from tensorrt_llm._torch.modules.dwdp import fabric_gather, vmm +from tensorrt_llm._torch.modules.dwdp.dense import gate +from tensorrt_llm._torch.modules.dwdp.dense.config import DENSE_COPY_ENV, DENSE_ENV +from tensorrt_llm._torch.modules.dwdp.dense.gate import validate_dense_gate +from tensorrt_llm._torch.modules.dwdp.dense.types import DenseMode +from tensorrt_llm._torch.peft.lora.config import LoraConfig +from tensorrt_llm.llmapi.llm_args import ( + CudaGraphConfig, + DwdpConfig, + MTPDecodingConfig, + PrefillCudaGraphBackend, + SleepConfig, + TorchCompileConfig, + TorchLlmArgs, +) +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + +ATTACHED = DenseMode.ATTACHED +STANDALONE = DenseMode.STANDALONE +_HANDLE_TYPE = vmm.cuda.CUmemAllocationHandleType + + +class _Launch(NamedTuple): + """The accepted context-worker launch of ``mode``, changed as given, and the gate's verdict. + + ``args`` are ``llm_args`` fields (applied unvalidated), ``mapping`` ``Mapping`` arguments, + ``env`` switches and ``lacks`` what the host lacks: ``fabric`` handles, ``mpi`` or the + ``gather_op``. With ``error`` the gate must reject the launch with it in its message, else + accept it with ``warning`` in its only warning, or with no warning. + """ + + mode: DenseMode + error: str = "" + warning: str = "" + args: dict = {} + mapping: dict = {} + env: dict = {} + lacks: str = "" + + def check(self, monkeypatch) -> DenseMode: + for name in [name for name in os.environ if name.startswith(DENSE_ENV)]: + monkeypatch.delenv(name) + for name, value in self.env.items(): + monkeypatch.setenv(name, value) + handle_type = _HANDLE_TYPE.CU_MEM_HANDLE_TYPE_FABRIC + if self.lacks == "fabric": + handle_type = _HANDLE_TYPE.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + monkeypatch.setattr(vmm, "peer_handle_type", lambda: handle_type) + monkeypatch.setattr(gate, "mpi_disabled", lambda: self.lacks == "mpi") + gather_op = None if self.lacks == "gather_op" else object() + monkeypatch.setattr(fabric_gather, "gather_op", lambda: gather_op) + dwdp_config = None + if self.mode is ATTACHED: + dwdp_config = DwdpConfig( + dwdp_size=16, num_groups=1, num_experts_per_worker=56, num_prefetch_experts=56 + ) + llm_args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + cuda_graph_config=None, + disable_overlap_scheduler=True, + dwdp_config=dwdp_config, + ).model_copy(update=self.args) + # Attached: tp 1. Standalone: a DEP16 context instance, experts parallel across it. + fields = dict(tp_size=1, pp_size=1, cp_size=1, enable_attention_dp=False) + if self.mode is STANDALONE: + fields.update(tp_size=16, enable_attention_dp=True) + fields.update(self.mapping) + if fields["enable_attention_dp"] and fields["cp_size"] == 1: + fields["moe_ep_size"] = fields["tp_size"] # CP has no MoE EP + world_size = fields["tp_size"] * fields["pp_size"] * fields["cp_size"] + return validate_dense_gate(llm_args, Mapping(world_size=world_size, rank=0, **fields)) + + +_COPY = "requires valid TRTLLM_DWDP_DENSE_* switches, but TRTLLM_DWDP_DENSE_COPY=" +_IGNORED = "is ignored with expert DWDP: the dense copies run on the expert DWDP copy stream" + +LAUNCHES = { + "attached": _Launch(ATTACHED), + "standalone": _Launch(STANDALONE), + # Without a DWDP group of two or more ranks, only an attention-DP group of them will do. + "dwdp_size_1": _Launch( + ATTACHED, + "but this launch has dwdp_config with dwdp_size=1.", + args={"dwdp_config": DwdpConfig(dwdp_size=1)}, + ), + "no_dwdp_config": _Launch( + ATTACHED, + "but this launch has no dwdp_config and enable_attention_dp=False.", + args={"dwdp_config": None}, + ), + "attention_dp_tp1": _Launch( + ATTACHED, + "but this launch has no dwdp_config and tp_size=1.", + args={"dwdp_config": None}, + mapping={"enable_attention_dp": True}, + ), + # Standalone, TRTLLM_DWDP_DENSE_COPY must be valid; with expert DWDP it is ignored. + "standalone-bogus_copy": _Launch( + STANDALONE, + f"{_COPY}'bogus' must be one of batched, per_slice, gather.", + env={DENSE_COPY_ENV: "bogus"}, + ), + "attached-gather_copy": _Launch( + ATTACHED, warning=f"{DENSE_COPY_ENV}='Gather' {_IGNORED}", env={DENSE_COPY_ENV: " Gather "} + ), + "attached-bogus_copy": _Launch( + ATTACHED, warning=f"{DENSE_COPY_ENV}='bogus' {_IGNORED}", env={DENSE_COPY_ENV: "bogus"} + ), + # One launch per requirement: attached also needs tp_size 1 and the overlap scheduler off; + # those of both modes alternate between the modes; then those of standalone alone. + "attached-tp_size": _Launch( + ATTACHED, "this launch has dwdp_config with tp_size=2", mapping={"tp_size": 2} + ), + "attached-overlap_scheduler": _Launch( + ATTACHED, + "this launch has dwdp_config with disable_overlap_scheduler=False", + args={"disable_overlap_scheduler": False}, + ), + **{ + f"{mode.value}-{name}": _Launch( + mode, f"requires {name}=None, but {name} is set", args={name: value} + ) + for mode, (name, value) in zip( + [ATTACHED, STANDALONE] * 3, + { + "cuda_graph_config": CudaGraphConfig(), + "torch_compile_config": TorchCompileConfig(), + "sleep_config": SleepConfig(), + "lora_config": LoraConfig(), + "speculative_config": MTPDecodingConfig(max_draft_len=1), + }.items(), + ) + }, + "standalone-prefill_cuda_graph_backend": _Launch( + STANDALONE, + "requires prefill_cuda_graph_backend='disabled', but prefill_cuda_graph_backend=", + args={"prefill_cuda_graph_backend": PrefillCudaGraphBackend.BREAKABLE}, + ), + "attached-enable_locality_domains": _Launch( + ATTACHED, + "requires enable_locality_domains=False, but enable_locality_domains=True", + args={"enable_locality_domains": True}, + ), + "standalone-pp_size": _Launch( + STANDALONE, "requires pp_size=1, but pp_size=2", mapping={"pp_size": 2} + ), + "attached-cp_size": _Launch( + ATTACHED, "requires cp_size=1, but cp_size=2", mapping={"cp_size": 2} + ), + "standalone-peer_handle_type": _Launch( + STANDALONE, + "requires FABRIC peer handles (CU_MEM_HANDLE_TYPE_FABRIC), but this host shares " + "CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR handles", + lacks="fabric", + ), + "standalone-mpi": _Launch(STANDALONE, "requires the MPI backend", lacks="mpi"), + "standalone-gather_op": _Launch( + STANDALONE, + f"requires trtllm::dwdp_fabric_gather for {DENSE_COPY_ENV}=gather, but this build", + env={DENSE_COPY_ENV: "gather"}, + lacks="gather_op", + ), + # A standalone gather launch passes with one warning naming the settings it was validated with. + "standalone-gather_copy": _Launch( + STANDALONE, + warning="TRTLLM_DWDP_GATHER_KERNEL=tma and TLLM_ROUTING_DISABLE_COOP_LAUNCH=1", + env={DENSE_COPY_ENV: "gather"}, + ), +} + + +@pytest.mark.parametrize("launch", LAUNCHES.values(), ids=LAUNCHES) +def test_dense_gate(monkeypatch, launch): + warnings = [] + monkeypatch.setattr(gate, "logger", SimpleNamespace(warning=warnings.append)) + if launch.error: + with pytest.raises(ValueError, match=re.escape(launch.error)) as error: + launch.check(monkeypatch) + assert str(error.value).startswith(f"Dense DWDP ({DENSE_ENV}) ") + assert str(error.value).endswith( + "Change the launch configuration or unset TRTLLM_DWDP_DENSE." + ) + else: + assert launch.check(monkeypatch) is launch.mode + assert len(warnings) == bool(launch.warning) and all(launch.warning in w for w in warnings) diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_guards.py b/tests/unittest/_torch/modules/dwdp/test_dense_guards.py new file mode 100644 index 000000000000..2e0174c6d511 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_guards.py @@ -0,0 +1,67 @@ +# 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. +"""CPU tests of the ``ModelLoader`` guards that refuse weight updates under dense DWDP. + +Dense DWDP setup marks the model and its decoder with ``_dense_dwdp_active`` once their dense +weights live in staging slots shared across layers. ``ModelLoader.reload`` and +``ModelLoader.begin_update_weights`` must then raise before they change anything. +""" + +from unittest.mock import MagicMock + +import pytest +from torch import nn + +from tensorrt_llm._torch.pyexecutor.model_loader import ModelLoader + +pytestmark = pytest.mark.cpu_only + + +class _CausalLM(nn.Module): + """Holds its decoder as ``model``, like ``DecoderModelForCausalLM``.""" + + def __init__(self) -> None: + super().__init__() + self.model = nn.Module() + + +class _Wrapper(nn.Module): + """Reaches the decoder through a property, like the Kimi-K3 VL wrapper.""" + + def __init__(self) -> None: + super().__init__() + self.llm = _CausalLM() + + @property + def model(self) -> nn.Module: + return self.llm.model + + +@pytest.mark.parametrize("action", ["reload", "update"]) +@pytest.mark.parametrize("owner", ["model", "decoder"]) +@pytest.mark.parametrize("model_class", [_CausalLM, _Wrapper]) +def test_weight_updates_refuse_a_dense_dwdp_model_before_any_change(model_class, owner, action): + model = model_class() + (model if owner == "model" else model.model)._dense_dwdp_active = True + loader = ModelLoader.__new__(ModelLoader) + loader.weight_mapper = MagicMock(name="weight_mapper", model=model) + + with pytest.raises(RuntimeError, match=f"Cannot {action} weights: dense DWDP"): + if action == "reload": + loader.reload(model, {"weight": None}) + else: + loader.begin_update_weights() + + assert loader.weight_mapper.mock_calls == [] diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_kimi_hooks.py b/tests/unittest/_torch/modules/dwdp/test_dense_kimi_hooks.py new file mode 100644 index 000000000000..e20c486b3bc0 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_kimi_hooks.py @@ -0,0 +1,395 @@ +# 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. +"""The Kimi-K3 side of dense DWDP: hook placement, the load guard and what the collector moves. + +Covers the hook calls H0-H3 of ``modeling_kimi_linear``, its load guard and the dense layout of its +real decoder layers, on a real model built from a tiny Kimi-Linear config. A real forward of the +attention and MoE submodules needs attention metadata, KV-cache pools and the routed-expert kernels, +so recorders replace exactly those submodules; everything else in the model forward runs for real, +on CPU in fp32. Stand-ins without tensors replace the routed experts and the MLA mixer +(``KimiK3MLAAttention``, whose attention backends allocate on CUDA); the MLA layer keeps its real +MLA branch and ``KimiMLARuntime`` wrapper. The model reaches dense DWDP through the dense DWDP +manager it captures when built (``get_global_dense_dwdp_manager``), in both modes: attached and +standalone. +""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +pytest.importorskip("fla") +from tensorrt_llm._torch.configs.kimi_linear import KimiLinearConfig +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models import modeling_kimi_linear as mkl +from tensorrt_llm._torch.modules import kimi_k3_mla +from tensorrt_llm._torch.modules.dwdp.dense.config import DenseConfig +from tensorrt_llm._torch.modules.dwdp.dense.layout import collect_dense_layout +from tensorrt_llm._torch.modules.dwdp.dense.manager import set_global_dense_dwdp_manager +from tensorrt_llm._torch.modules.dwdp.dense.types import SegmentKind +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.moe.fused_moe import ConfigurableMoE, CutlassFusedMoE +from tensorrt_llm._torch.pyexecutor.dwdp import set_global_dwdp_manager +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig + +pytestmark = pytest.mark.cpu_only + +_HIDDEN = 256 +_TOKENS = 3 +# Layers of the default ``_kimi_config``: 0 is a dense-MLP KDA layer, then KDA, MLA and KDA MoE. +_MOE_LAYERS = (1, 2, 3) + + +def _kimi_config(kda_layers=(1, 2, 4), mla_layers=(3,)) -> KimiLinearConfig: + """A tiny Kimi-K3; ``kda_layers`` and ``mla_layers`` are 1-indexed, as in the checkpoint.""" + return KimiLinearConfig( + vocab_size=128, + hidden_size=_HIDDEN, + intermediate_size=256, + num_hidden_layers=len(kda_layers) + len(mla_layers), + num_attention_heads=2, + q_lora_rank=256, + kv_lora_rank=128, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rms_norm_eps=1e-6, + attn_res_block_size=2, + mla_use_output_gate=True, + max_position_embeddings=64, + linear_attn_config={ + "kda_layers": list(kda_layers), + "full_attn_layers": list(mla_layers), + "num_heads": 2, + "head_dim": 128, + "short_conv_kernel_size": 4, + "use_full_rank_gate": True, + }, + num_experts=8, + num_experts_per_token=2, + moe_intermediate_size=128, + num_shared_experts=1, + routed_expert_hidden_size=128, + latent_moe_use_norm=True, + first_k_dense_replace=1, + activation_situ_beta=4.0, + activation_situ_linear_beta=25.0, + ) + + +def _model_config(cfg: KimiLinearConfig) -> ModelConfig: + """The deployed context-worker shape: tp1 with attention DP (fp32 router, replicated MLPs).""" + return ModelConfig( + pretrained_config=cfg, + mapping=Mapping(world_size=1, tp_size=1, rank=0, enable_attention_dp=True), + quant_config=QuantConfig(), + moe_backend="CUTLASS", + ) + + +class _CpuMlaMixer(nn.Module): + """Stands in for ``KimiK3MLAAttention``; ``KimiMLARuntime`` wraps it as the real mixer.""" + + def __init__(self, *, hidden_size: int, num_heads: int, v_head_dim: int, **kwargs) -> None: + super().__init__() + self.o_proj = nn.Linear(num_heads * v_head_dim, hidden_size, bias=False) + + +class _CpuRoutedBackend(CutlassFusedMoE): + """The CUTLASS backend the layers request, without its constructor and its tensors.""" + + def __init__(self, num_experts: int) -> None: + nn.Module.__init__(self) + self.initial_local_expert_ids = list(range(num_experts)) + + @classmethod + def can_implement(cls, p, d): + raise NotImplementedError("CPU stand-in") + + def _get_quant_method(self): + raise NotImplementedError("CPU stand-in") + + +class _CpuRoutedExperts(ConfigurableMoE): + """What ``KimiK3MoERuntime`` reads of the routed experts at construction, at tp1.""" + + def __init__(self, *, num_experts: int, **kwargs) -> None: + nn.Module.__init__(self) + self.backend = _CpuRoutedBackend(num_experts) + self.comm = None + self.layer_load_balancer = None + self.all_reduce = None + + +@pytest.fixture(autouse=True) +def _cpu_modules(monkeypatch: pytest.MonkeyPatch) -> None: + """CPU stand-ins for the routed experts, the MLA mixer and the model's auxiliary streams. + + Layers capture the global managers when built: every test starts and ends without one. + """ + monkeypatch.setattr(mkl, "create_moe", lambda **kwargs: _CpuRoutedExperts(**kwargs)) + monkeypatch.setattr(kimi_k3_mla, "KimiK3MLAAttention", _CpuMlaMixer) + monkeypatch.setattr(torch.cuda, "Stream", object) + set_global_dwdp_manager(None) + set_global_dense_dwdp_manager(None) + yield + set_global_dwdp_manager(None) + set_global_dense_dwdp_manager(None) + + +# --------------------------------------------------------------------------- +# Recorders +# --------------------------------------------------------------------------- + + +class _Recorder(nn.Module): + """Stands in for an attention, MoE, MLP or embedding: logs the call, returns fp32 states.""" + + def __init__(self, log: list, *entry) -> None: + super().__init__() + self._log = log + self._entry = entry + + def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor: + self._log.append(self._entry) + if hidden_states.dtype == torch.int32: # the token ids of the embedding + return torch.randn(hidden_states.shape[0], _HIDDEN) + return hidden_states * 0.5 + + +class _RecordingDense: + """The ``DenseDWDP`` calls the model makes, recorded; any other call raises.""" + + def __init__(self, log: list) -> None: + self._log = log + + def covers(self, layer: int) -> bool: + self._log.append(("covers", layer)) + return layer in _MOE_LAYERS + + def wait_layer(self, layer: int) -> None: + self._log.append(("wait_layer", layer)) + + def attention_done(self, layer: int) -> None: + self._log.append(("attention_done", layer)) + + def layer_done(self, layer: int) -> None: + self._log.append(("layer_done", layer)) + + +def _manager(log: list, dense: object | None) -> SimpleNamespace: + """``DenseDwdpManager.dense`` and ``DenseDwdpManager.ensure_primed``, recorded.""" + return SimpleNamespace(dense=dense, ensure_primed=lambda: log.append(("ensure_primed",))) + + +def _install_recorders(layer: mkl.KimiLinearDecoderLayer, log: list) -> None: + """Replace the layer's attention and MoE (or MLP) submodules; log its norms' calls too.""" + index = layer.layer_idx + setattr( + layer, "linear_attn" if layer.is_kda else "self_attn", _Recorder(log, "attention", index) + ) + ffn = "block_sparse_moe" if layer.is_moe else "mlp" + setattr(layer, ffn, _Recorder(log, "ffn", index)) + for norm in ("input_layernorm", "post_attention_layernorm"): + getattr(layer, norm).register_forward_pre_hook( + lambda module, args, name=norm: log.append((name, index)) + ) + + +def _layer_reads(index: int, dense: bool) -> list[tuple]: + """The calls of one decoder-layer forward; a covered layer brackets its reads with H1-H3.""" + hooked = dense and index in _MOE_LAYERS + return [ + *([("covers", index)] if dense else []), + *([("wait_layer", index)] if hooked else []), + ("input_layernorm", index), + ("attention", index), + *([("attention_done", index)] if hooked else []), + ("post_attention_layernorm", index), + ("ffn", index), + *([("layer_done", index)] if hooked else []), + ] + + +@pytest.mark.parametrize("registered", [True, False], ids=["registered", "unregistered"]) +def test_the_model_and_its_layers_capture_the_registered_dense_manager(registered): + """Built with or without a registered dense DWDP manager, the model and each layer keep it.""" + manager = _manager([], None) if registered else None + set_global_dense_dwdp_manager(manager) + + model = mkl.KimiLinearModel(_model_config(_kimi_config())) + + assert model._dense_dwdp is manager + assert model.layers and all(layer._dense_dwdp is manager for layer in model.layers) + + +@pytest.mark.parametrize("state", ["dense", "no_dense", "no_manager"]) +def test_the_model_primes_first_and_brackets_the_reads_of_covered_layers(state): + """H0 primes before the embedding; H1-H3 bracket the attention and MoE reads of KDA and MLA.""" + log = [] + model = mkl.KimiLinearModel(_model_config(_kimi_config())) + kinds = [(layer.is_moe, layer.is_kda) for layer in model.layers] + assert kinds == [(False, True), (True, True), (True, False), (True, True)] + model.embed_tokens = _Recorder(log, "embed") + manager = { + "dense": _manager(log, _RecordingDense(log)), + "no_dense": _manager(log, None), + "no_manager": None, + }[state] + model._dense_dwdp = manager + for layer in model.layers: + _install_recorders(layer, log) + layer._dense_dwdp = manager + # fp32 activations keep the residual mixing and the norms on their eager CPU paths. + metadata = SimpleNamespace(num_tokens=_TOKENS) + + output = model(metadata, input_ids=torch.zeros(_TOKENS, dtype=torch.int32)) + + assert output.shape == (_TOKENS, _HIDDEN) + dense = state == "dense" + reads = [("ensure_primed",)] * dense + [("embed",)] + assert log == reads + sum((_layer_reads(i, dense) for i in range(len(kinds))), []) + # The input check runs first: a rejected call does not prime. + log.clear() + with pytest.raises(ValueError, match="exactly one of input_ids or inputs_embeds"): + model(metadata) + assert log == [] + + +class _LoadReached(Exception): + """Raised by the stubbed loader once ``load_weights`` got past its guard.""" + + +def _causal_lm(active: bool | None) -> mkl.KimiLinearForCausalLM: + # The guard reads only ``self.model``; the first loader step is stubbed to report reaching it. + model = mkl.KimiLinearForCausalLM.__new__(mkl.KimiLinearForCausalLM) + nn.Module.__init__(model) + model.model = nn.Module() + if active is not None: + model.model._dense_dwdp_active = active + + def trunk_parameters(): + raise _LoadReached + + model._trunk_parameters = trunk_parameters + return model + + +def test_load_weights_refuses_while_dense_is_active(): + with pytest.raises(RuntimeError, match="dense DWDP is active"): + _causal_lm(active=True).load_weights({}) + for active in (False, None): + with pytest.raises(_LoadReached): + _causal_lm(active=active).load_weights({}) + + +# --------------------------------------------------------------------------- +# What the collector moves from the real layers +# --------------------------------------------------------------------------- + + +def _finish_loading(model: mkl.KimiLinearModel) -> None: + """What the loader does after the checkpoint: transform each Linear, build the KDA tensors. + + ``finalize_decode_weights`` skips CPU, so its fused projections and constants are built here. + """ + for module in model.modules(): + if isinstance(module, Linear): + module.post_load_weights() + for layer in model.layers: + if layer.is_kda: + kda = layer.linear_attn + kda._qkvg_proj_weight = kda._merge_projection_weights( + (kda.q_proj, kda.k_proj, kda.v_proj, kda.g_proj) + ) + kda._bfa_proj_weight = kda._merge_projection_weights( + (kda.f_a_proj, kda.b_proj), pad_rows_to=8 + ) + kda._build_decode_kernel_constants() + kda._build_mtp_conv_weights() + + +def _root_tensors(roots: list[tuple[str, nn.Module]]) -> set[str]: + """The names dense DWDP binds under ``roots``: the tensors that may move and their aliases. + + Parameters, buffers and the tensor attributes ``DWDP_DENSE_ATTR_ALLOW`` names may move unless + ``DWDP_DENSE_EXCLUDE`` names them; any other tensor reference under the roots moves only as a + view of a storage that moves. + """ + references = [] + for root_path, root in roots: + for path, module in root.named_modules(prefix=root_path): + for table in (module._parameters, module._buffers): + references += [ + (f"{path}.{name}", tensor, name not in mkl.DWDP_DENSE_EXCLUDE) + for name, tensor in table.items() + if tensor is not None + ] + for name, value in vars(module).items(): + allowed = name in mkl.DWDP_DENSE_ATTR_ALLOW - mkl.DWDP_DENSE_EXCLUDE + if isinstance(value, torch.Tensor): + references.append((f"{path}.{name}", value, allowed)) + elif isinstance(value, (tuple, list)) and value: + if all(isinstance(item, torch.Tensor) for item in value): + references += [ + (f"{path}.{name}[{index}]", item, allowed) + for index, item in enumerate(value) + ] + storage = lambda tensor: tensor.untyped_storage().data_ptr() # noqa: E731 + moving = {storage(tensor) for _, tensor, may_move in references if may_move} + return {name for name, tensor, _ in references if storage(tensor) in moving} + + +def test_the_collector_moves_every_root_tensor_of_the_moe_layers_and_nothing_else(): + """Each tensor under the attention and MoE-dense roots moves; layer 0 and the experts stay. + + Strict collection with a 1-byte threshold also fails on any plain tensor attribute under the + roots that ``DWDP_DENSE_ATTR_ALLOW`` and ``DWDP_DENSE_EXCLUDE`` do not account for. + """ + model = mkl.KimiLinearModel(_model_config(_kimi_config())) + _finish_loading(model) + + layout = collect_dense_layout( + list(enumerate(model.layers)), + DenseConfig(enabled=True, min_bytes=1, strict=True), + num_ranks=2, + device=torch.device("cpu"), + attr_allow=mkl.DWDP_DENSE_ATTR_ALLOW, + exclude=mkl.DWDP_DENSE_EXCLUDE, + ) + + assert layout.layers == _MOE_LAYERS + assert layout.warnings == [] + for index in _MOE_LAYERS: + layer = model.layers[index] + attention = "linear_attn" if layer.is_kda else "self_attn" + assert [path for path, _ in layer.dense_dwdp_roots("attn")] == [attention] + assert [path for path, _ in layer.dense_dwdp_roots("moe")] == [ + "block_sparse_moe.gate", + "block_sparse_moe.shared_experts", + "block_sparse_moe.routed_expert_down_proj", + "block_sparse_moe.routed_expert_up_proj", + "block_sparse_moe.routed_expert_norm", + ] + for kind in (SegmentKind.ATTN, SegmentKind.MOE): + moved = { + binding.qualified_name + for anchor in layout.segment(index, kind).anchors + for binding in anchor.bindings + } + assert moved == _root_tensors(layer.dense_dwdp_roots(kind.value)), (index, kind) + assert not any(name.startswith("block_sparse_moe.routed_experts") for name in moved) diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_layout.py b/tests/unittest/_torch/modules/dwdp/test_dense_layout.py new file mode 100644 index 000000000000..66fc43fa8d38 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_layout.py @@ -0,0 +1,464 @@ +# 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. +"""CPU tests of the dense DWDP collector. + +The toy decoder layers below imitate the Kimi-K3 tensors the collector has to handle: the fused +``qkvg_proj.weight`` buffer with ``q/k/v/g_proj.weight`` row-view buffers and the low-M GEMM's +scale copy on ``qkvg_proj``, the plain ``_bfa_proj_weight`` attribute with ``f_a_proj`` / +``b_proj`` Parameter row views and pad rows no view covers, the ``_mtp_conv_weights`` tuple, +tensors under the threshold, excluded scratch names (one of them aliasing an anchor), an MLA layer +and the MoE-dense roots. Aliases that are not plain row views ride along: ``_gate_b_t`` is a +transposed view of ``f_b_proj``'s storage, and ``_bfa_views`` holds a stride-0 view and an empty +view of the ``_bfa`` storage. The setup tests reuse the toy. The slot checksum watchdog of +``TRTLLM_DWDP_DENSE_DEBUG=2`` is checked against the toy's layout. +""" + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.modules.dwdp.dense.config import DenseConfig +from tensorrt_llm._torch.modules.dwdp.dense.layout import ( + collect_dense_layout, + compute_checksums, + write_blob, +) +from tensorrt_llm._torch.modules.dwdp.dense.types import ( + BindingKind, + DenseLayoutError, + SegmentKind, + allocate_slot_ring, +) +from tensorrt_llm._torch.modules.dwdp.dense.watchdog import SlotWatchdog + +pytestmark = pytest.mark.cpu_only + +ATTN = SegmentKind.ATTN +MOE = SegmentKind.MOE + +#: The model tables, as ``modeling_kimi_linear`` defines them. +ATTR_ALLOW = 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".split() +) +EXCLUDE = frozenset( + "_o_dense _fused_q_placeholder softmax_stats_tensor temp_softmax_stats_tensor " + "e_score_correction_bias".split() +) + +#: Toy threshold: the toy weights are 8-64 KiB, the small tensors at most 2 KiB. +TOY_MIN_BYTES = 8192 +NUM_RANKS = 4 +HIDDEN = 256 +ROWS = 64 +F_ROWS = 16 +B_ROWS = 2 +#: Decoder layer kinds: layer 0 stays replicated, layer 3 is an MLA layer. +LAYER_KINDS = ("dense", "kda", "kda", "mla", "kda", "kda") +COVERED = (1, 2, 3, 4, 5) + +FP8 = torch.float8_e4m3fn +BF16 = torch.bfloat16 + + +# --------------------------------------------------------------------------- +# Toy Kimi-K3 decoder layers +# --------------------------------------------------------------------------- + + +def _ints(generator, shape, dtype): + """Small integers in [-3, 3], exactly representable in every moved dtype.""" + return torch.randint(-3, 4, shape, generator=generator).to(torch.float32).to(dtype) + + +def _bytes(generator, shape): + return torch.randint(0, 256, shape, generator=generator, dtype=torch.uint8) + + +def _parameter(tensor): + return nn.Parameter(tensor, requires_grad=False) + + +class ToyFp8Linear(nn.Module): + """Stands in for ``_Fp8BlockScaleWeightReadLinear``: weight and scale buffers.""" + + def __init__(self, weight, scale): + super().__init__() + self._weights_transformed = True + self.register_buffer("weight", weight, persistent=False) + self.register_buffer("weight_scale", scale, persistent=False) + self._low_m_weight_scale = None + + +class ToyNorm(nn.Module): + def __init__(self, generator, size, dtype=BF16): + super().__init__() + self.weight = _parameter(_ints(generator, (size,), dtype)) + + +class ToyLinear(nn.Linear): + def __init__(self, generator, in_features, out_features): + super().__init__(in_features, out_features, bias=False, dtype=BF16) + self.weight.requires_grad_(False) + self.weight.data.copy_(_ints(generator, (out_features, in_features), BF16)) + + +class ToyKDA(nn.Module): + """KDA mixer after load: fused FP8 qkvg, fused BF16 [f_a | b], decode constants.""" + + def __init__(self, generator): + super().__init__() + g = generator + self.A_log = _parameter(_ints(g, (16,), torch.float32)) + self.dt_bias = _parameter(_ints(g, (64,), torch.float32)) + for name in ("q_proj", "k_proj", "v_proj", "g_proj"): + setattr(self, name, ToyFp8Linear(_ints(g, (ROWS, HIDDEN), FP8), _bytes(g, (ROWS, 128)))) + self.o_proj = ToyFp8Linear(_ints(g, (HIDDEN, 4 * ROWS), FP8), _bytes(g, (HIDDEN, 32))) + self.f_a_proj = ToyLinear(g, HIDDEN, F_ROWS) + self.b_proj = ToyLinear(g, HIDDEN, B_ROWS) + self.f_b_proj = ToyLinear(g, F_ROWS, HIDDEN) + self.q_conv1d = ToyNorm(g, 4 * HIDDEN) + self.o_norm = ToyNorm(g, HIDDEN) + # What the K3 loader's finalize step and FP8 fusion leave behind. + projections = [self.q_proj, self.k_proj, self.v_proj, self.g_proj] + fused = torch.cat([module.weight for module in projections]) + self.qkvg_proj = ToyFp8Linear(fused, _bytes(g, (4 * ROWS, 64))) + # The low-M GEMM's scale on SM107: a packed int32 copy, a transposed view of its storage. + self.qkvg_proj._low_m_weight_scale = ( + torch.arange(16 * 4 * ROWS, dtype=torch.int32).view(16, 4 * ROWS).t() + ) + for index, module in enumerate(projections): + module.weight = fused[index * ROWS : (index + 1) * ROWS] + bfa = torch.cat( + [self.f_a_proj.weight.data, self.b_proj.weight.data, torch.zeros(6, HIDDEN, dtype=BF16)] + ) + self.f_a_proj.weight.data = bfa[:F_ROWS] + self.b_proj.weight.data = bfa[F_ROWS : F_ROWS + B_ROWS] + self._bfa_proj_weight = bfa + # Names in neither model table that alias the _bfa storage: a stride-0 view and an empty + # view past the pad rows. + self._bfa_views = [bfa[:1].expand(4, -1), bfa[bfa.shape[0] :]] + self._qkvg_proj_weight = None + # The low-rank gate's [f_b; g_b] bmm operand: a transposed view of f_b_proj's storage. + self._gate_b_t = self.f_b_proj.weight.data.view(2, HIDDEN // 2, F_ROWS).transpose(1, 2) + self._mtp_conv_weights = tuple(_ints(g, (4, 1024), BF16) for _ in range(3)) + self._packed_conv_weight = _ints(g, (3, HIDDEN), BF16) + self._w_q_t = _ints(g, (4, 128), BF16) + self._A_log_f32 = self.A_log.detach().float().contiguous() + self._dt_bias_f32 = self.dt_bias.detach().float().contiguous() + self._o_dense = self.o_proj.weight[64:192] + + +class ToyAttentionBackend: + """Not an ``nn.Module``: the walk never enters it (like MLA's ``mha`` / ``mqa``).""" + + def __init__(self, generator): + self.workspace = _ints(generator, (4096,), BF16) + + +class ToyMLA(nn.Module): + """MLA attention after load, with the BF16 absorb copies and excluded scratch tensors.""" + + def __init__(self, generator): + super().__init__() + g = generator + self.kv_a_proj_with_mqa = ToyFp8Linear(_ints(g, (128, HIDDEN), FP8), _bytes(g, (128, 64))) + self.q_b_proj = ToyFp8Linear(_ints(g, (HIDDEN, 128), FP8), _bytes(g, (HIDDEN, 32))) + self.kv_b_proj = ToyFp8Linear(_ints(g, (HIDDEN, 64), FP8), _bytes(g, (HIDDEN, 32))) + self.o_proj = ToyFp8Linear(_ints(g, (HIDDEN, HIDDEN), FP8), _bytes(g, (HIDDEN, 32))) + self.k_b_proj_trans = _parameter(_ints(g, (8, 64, 64), FP8)) + self.k_b_proj_trans_dequant = _parameter(_ints(g, (8, 64, 64), BF16)) + self.v_b_proj_scale = _parameter(self.kv_b_proj.weight_scale[:128]) + self._fused_q_placeholder = _ints(g, (64, HIDDEN), BF16) + self.softmax_stats_tensor = torch.zeros(4096, dtype=torch.float32) + self.mha = ToyAttentionBackend(g) + + +class ToyGate(nn.Module): + def __init__(self, generator): + super().__init__() + self.weight = _parameter(_ints(generator, (32, HIDDEN), torch.float32)) + self.e_score_correction_bias = _parameter(_ints(generator, (32,), torch.float32)) + + +class ToySharedMLP(nn.Module): + def __init__(self, generator): + super().__init__() + g = generator + self.gate_up_proj = ToyFp8Linear(_ints(g, (HIDDEN, HIDDEN), FP8), _bytes(g, (HIDDEN, 32))) + self.down_proj = ToyFp8Linear(_ints(g, (HIDDEN, 128), FP8), _bytes(g, (HIDDEN, 32))) + + +class ToyMoE(nn.Module): + def __init__(self, generator): + super().__init__() + g = generator + self.gate = ToyGate(g) + self.shared_experts = ToySharedMLP(g) + self.routed_expert_down_proj = ToyFp8Linear( + _ints(g, (128, HIDDEN), FP8), _bytes(g, (128, 16)) + ) + self.routed_expert_up_proj = ToyFp8Linear( + _ints(g, (HIDDEN, 128), FP8), _bytes(g, (HIDDEN, 8)) + ) + self.routed_expert_norm = ToyNorm(g, 128) + # Managed by expert DWDP; never a dense root. + self.routed_experts = ToyLinear(g, HIDDEN, 4 * HIDDEN) + + +class ToyDecoderLayer(nn.Module): + """Decoder layer with the model seam dense DWDP reads: ``dense_dwdp_roots`` and the tables.""" + + dense_dwdp_attr_allow = ATTR_ALLOW + dense_dwdp_exclude = EXCLUDE + + def __init__(self, kind, generator): + super().__init__() + self.is_moe = kind != "dense" + self.is_kda = kind != "mla" + self.input_layernorm = ToyNorm(generator, HIDDEN) + if self.is_kda: + self.linear_attn = ToyKDA(generator) + else: + self.self_attn = ToyMLA(generator) + if self.is_moe: + self.block_sparse_moe = ToyMoE(generator) + else: + self.mlp = ToySharedMLP(generator) + + def dense_dwdp_roots(self, segment): + if not self.is_moe: + 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 ( + "gate", + "shared_experts", + "routed_expert_down_proj", + "routed_expert_up_proj", + "routed_expert_norm", + ) + ] + + +class ToyModel(nn.Module): + def __init__(self, seed=0, kinds=LAYER_KINDS): + super().__init__() + generator = torch.Generator().manual_seed(seed) + self.layers = nn.ModuleList(ToyDecoderLayer(kind, generator) for kind in kinds) + + +def make_config(**overrides): + return DenseConfig(enabled=True, **{"min_bytes": TOY_MIN_BYTES, **overrides}) + + +def collect(model, config=None, *, num_ranks=NUM_RANKS): + return collect_dense_layout( + list(enumerate(model.layers)), + config or make_config(), + num_ranks=num_ranks, + device=torch.device("cpu"), + attr_allow=ATTR_ALLOW, + exclude=EXCLUDE, + ) + + +def anchor_paths(layout, layer, segment): + return [ + [binding.qualified_name for binding in anchor.bindings] + for anchor in layout.segment(layer, segment).anchors + ] + + +def storage_ptr(tensor): + return tensor.untyped_storage().data_ptr() + + +#: Anchors of a toy KDA layer in walk order: the mixer's own attributes (sorted by name), then +#: the children in registration order (``qkvg_proj`` last, as the K3 loader adds it). Each entry +#: lists the anchor's bindings in walk order. +KDA_ANCHORS = [ + [ + "linear_attn._bfa_proj_weight", + "linear_attn._bfa_views[0]", + "linear_attn._bfa_views[1]", + "linear_attn.f_a_proj.weight", + "linear_attn.b_proj.weight", + ], + ["linear_attn._gate_b_t", "linear_attn.f_b_proj.weight"], + ["linear_attn._mtp_conv_weights[0]"], + ["linear_attn._mtp_conv_weights[1]"], + ["linear_attn._mtp_conv_weights[2]"], + [ + "linear_attn.q_proj.weight", + "linear_attn.k_proj.weight", + "linear_attn.v_proj.weight", + "linear_attn.g_proj.weight", + "linear_attn.qkvg_proj.weight", + ], + ["linear_attn.q_proj.weight_scale"], + ["linear_attn.k_proj.weight_scale"], + ["linear_attn.v_proj.weight_scale"], + ["linear_attn.g_proj.weight_scale"], + ["linear_attn._o_dense", "linear_attn.o_proj.weight"], + ["linear_attn.o_proj.weight_scale"], + ["linear_attn.qkvg_proj.weight_scale"], + ["linear_attn.qkvg_proj._low_m_weight_scale"], +] + +MLA_ANCHORS = [ + ["self_attn.k_b_proj_trans"], + ["self_attn.k_b_proj_trans_dequant"], + ["self_attn.v_b_proj_scale", "self_attn.kv_b_proj.weight_scale"], + ["self_attn.kv_a_proj_with_mqa.weight"], + ["self_attn.kv_a_proj_with_mqa.weight_scale"], + ["self_attn.q_b_proj.weight"], + ["self_attn.q_b_proj.weight_scale"], + ["self_attn.kv_b_proj.weight"], + ["self_attn.o_proj.weight"], + ["self_attn.o_proj.weight_scale"], +] + +MOE_ANCHORS = [ + ["block_sparse_moe.gate.weight"], + ["block_sparse_moe.shared_experts.gate_up_proj.weight"], + ["block_sparse_moe.shared_experts.gate_up_proj.weight_scale"], + ["block_sparse_moe.shared_experts.down_proj.weight"], + ["block_sparse_moe.shared_experts.down_proj.weight_scale"], + ["block_sparse_moe.routed_expert_down_proj.weight"], + ["block_sparse_moe.routed_expert_up_proj.weight"], +] + + +# --------------------------------------------------------------------------- +# Collection +# --------------------------------------------------------------------------- + + +def test_collects_k3_like_anchors_and_bindings(): + model = ToyModel() + layout = collect(model) + + assert layout.layers == COVERED + assert layout.warnings == [] + assert list(layout.segments) == [(layer, kind) for layer in COVERED for kind in (ATTN, MOE)] + for layer in COVERED: + expected = MLA_ANCHORS if LAYER_KINDS[layer] == "mla" else KDA_ANCHORS + assert anchor_paths(layout, layer, ATTN) == expected + assert anchor_paths(layout, layer, MOE) == MOE_ANCHORS + + kda = model.layers[1].linear_attn + anchors = { + anchor.bindings[0].qualified_name: anchor for anchor in layout.segment(1, ATTN).anchors + } + qkvg = anchors["linear_attn.q_proj.weight"] + assert qkvg.nbytes == 4 * ROWS * HIDDEN + assert [binding.byte_offset for binding in qkvg.bindings] == [0, 16384, 32768, 49152, 0] + assert [binding.kind for binding in qkvg.bindings] == [BindingKind.BUFFER] * 5 + # The whole _bfa storage moves, including the pad rows no Parameter view covers. + bfa = anchors["linear_attn._bfa_proj_weight"] + assert bfa.nbytes == 24 * HIDDEN * 2 + assert [(binding.kind, binding.byte_offset) for binding in bfa.bindings] == [ + (BindingKind.ATTRIBUTE, 0), + (BindingKind.SEQUENCE_ITEM, 0), + (BindingKind.SEQUENCE_ITEM, bfa.nbytes), + (BindingKind.PARAMETER, 0), + (BindingKind.PARAMETER, F_ROWS * HIDDEN * 2), + ] + assert bfa.bindings[3].original.data_ptr() == kda.f_a_proj.weight.data_ptr() + # Bindings that are not row views keep their geometry: stride 0, empty, transposed. + stride_0, empty = bfa.bindings[1:3] + assert (stride_0.size, stride_0.stride) == ((4, HIDDEN), (0, 1)) + assert stride_0.span_bytes == 2 * HIDDEN + assert (empty.size, empty.span_bytes) == ((0, HIDDEN), 0) + gate_b_t = anchors["linear_attn._gate_b_t"].bindings[0] + assert (gate_b_t.size, gate_b_t.stride) == ( + (2, F_ROWS, HIDDEN // 2), + (HIDDEN // 2 * F_ROWS, 1, F_ROWS), + ) + assert not gate_b_t.original.is_contiguous() + # The excluded scratch name that aliases o_proj is a binding of o_proj's anchor. + o_dense = anchors["linear_attn._o_dense"].bindings[0] + assert (o_dense.size, o_dense.stride, o_dense.byte_offset) == ((128, 256), (256, 1), 16384) + # The low-M scale copy has a storage of its own: an allowed attribute, so it moves too. + low_m = anchors["linear_attn.qkvg_proj._low_m_weight_scale"].bindings[0] + assert (low_m.kind, low_m.size, low_m.stride) == ( + BindingKind.ATTRIBUTE, + (4 * ROWS, 16), + (1, 4 * ROWS), + ) + mtp = anchors["linear_attn._mtp_conv_weights[1]"].bindings[0] + assert (mtp.kind, mtp.name, mtp.index) == (BindingKind.SEQUENCE_ITEM, "_mtp_conv_weights", 1) + # Sources are flat byte views of the whole original storage. + for anchor in layout.anchors(): + assert anchor.source.dtype == torch.uint8 and anchor.source.numel() == anchor.nbytes + assert storage_ptr(anchor.source) == storage_ptr(anchor.bindings[0].original) + for binding in anchor.bindings: + assert binding.owner_ref() is not None and binding.checksum is None + # Checksums are left to compute_checksums. + assert all(segment.checksum is None for segment in layout.segments.values()) + + +def test_storage_shared_across_layers_stays(): + model = ToyModel() + model.layers[4].block_sparse_moe.gate.weight = model.layers[1].block_sparse_moe.gate.weight + + layout = collect(model) + assert len(layout.warnings) == 1 + assert "layers.1/moe, layers.4/moe" in layout.warnings[0] + for layer in (1, 4): + assert anchor_paths(layout, layer, MOE) == MOE_ANCHORS[1:] + assert anchor_paths(layout, 2, MOE) == MOE_ANCHORS + + with pytest.raises(DenseLayoutError, match="layers.1.block_sparse_moe.gate.weight"): + collect(model, make_config(strict=True)) + + +@pytest.mark.parametrize("strict", [False, True]) +def test_untransformed_weights_raise(strict): + model = ToyModel() + model.layers[4].linear_attn.o_proj._weights_transformed = False + with pytest.raises(DenseLayoutError, match="_weights_transformed=False"): + collect(model, make_config(strict=strict)) + + +def test_slot_watchdog_accepts_exactly_the_gathered_blobs(): + """The P1 checksums are what a gather leaves in a slot: the blob, zero-padded.""" + layout = collect(ToyModel()) + with pytest.raises(DenseLayoutError, match="not computed"): + SlotWatchdog.from_layout(layout, "cpu") + compute_checksums(layout) + watchdog = SlotWatchdog.from_layout(layout, "cpu") + assert watchdog.layers == layout.layers and watchdog.segments == layout.segment_kinds + rings = {kind: allocate_slot_ring(spec, "cpu") for kind, spec in layout.slot_specs.items()} + for layer in layout.layers: + for segment in layout.layer_segments(layer): + slot = rings[segment.segment][layout.slot_index(segment.segment, layer)] + blob = slot.blob_view(layout.num_ranks * segment.row_bytes) + write_blob(segment, blob, num_ranks=layout.num_ranks) + watchdog.check(layer, segment.segment, slot) + if segment.blob_bytes: + slot.blob_view(segment.blob_bytes)[segment.blob_bytes - 1] ^= 0x80 + watchdog.check(layer, segment.segment, slot) + flipped = [key for key, segment in layout.segments.items() if segment.blob_bytes] + findings = watchdog.poll() + assert [(layer, kind) for layer, kind, _ in findings] == flipped + assert all(count == 1 for _, _, count in findings) + assert watchdog.checks == len(layout.segments) + len(flipped) diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_protocol.py b/tests/unittest/_torch/modules/dwdp/test_dense_protocol.py new file mode 100644 index 000000000000..0ade32852c42 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_protocol.py @@ -0,0 +1,954 @@ +# 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. +"""CPU tests of the dense DWDP runtime protocol (``runtime.DenseDWDP``), attached and standalone. + +A recording fake of the CUDA seam logs every event record, stream wait, copy launch and device +synchronization in host order; a wait binds to the latest record of its event before it, as +``cudaStreamWaitEvent`` does, and a device synchronization orders everything after it behind +everything before it. The runtime runs on the fake through its ``ops`` seam. Attached, the real +``DWDPWeightManager``, with its ``PrefetchSchedule``, runs on it through patched ``torch.cuda`` +calls and shares the copy stream, and the real ``DenseDwdpManager`` holds both. A driver calls +them the way the executor loop (``DwdpManager.start_forward``, or standalone +``DenseDwdpManager.prefetch_first_layers``, on the default stream), ``KimiLinearModel`` (H0-H3, on +the execution stream) and ``ConfigurableMoE`` (``wait_and_bind``, ``finish_layer``) do, and logs +the model's weight reads. The tests check where the runtime records its events, the copy order on +the copy stream, that the host synchronizes only to recover a forward abandoned inside a MoE layer +and, with vector clocks over the log, that no copy into a dense or expert slot can overlap a read +of that slot: a missing consume wait, fence, WAR barrier or synchronization fails them. +""" + +from __future__ import annotations + +import contextlib +import gc +import random +from collections import Counter +from collections.abc import Callable, Iterator +from dataclasses import dataclass +from unittest.mock import patch + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.modules.dwdp.dense.manager import DenseDwdpManager +from tensorrt_llm._torch.modules.dwdp.dense.runtime import DenseDWDP +from tensorrt_llm._torch.modules.dwdp.dense.types import ( + EMPTY_TAG, + RING_DEPTH, + ROW_ALIGN, + DenseMode, + SegmentKind, + SlotSpec, + allocate_slot_ring, +) +from tensorrt_llm._torch.modules.dwdp.weight_manager import DWDPWeightManager +from tensorrt_llm._torch.pyexecutor.dwdp import DwdpManager +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + +ATTN = SegmentKind.ATTN +MOE = SegmentKind.MOE +BOTH = (ATTN, MOE) +SEGMENT_SETS = pytest.mark.parametrize("segments", [BOTH, (ATTN,)], ids=["attn+moe", "attn"]) +K3_LAYERS = tuple(range(1, 93)) +#: What a launch or a read names: the dense segments, and the expert weights. +EXPERT = "E" +_WEIGHT_MANAGER = "tensorrt_llm._torch.modules.dwdp.weight_manager" + +# --------------------------------------------------------------------------- +# Recording fake of the CUDA seam +# --------------------------------------------------------------------------- + + +@dataclass(eq=False) +class Entry: + """One log entry, in host order. + + ``kind`` is ``record``, ``wait``, ``launch`` or ``read`` (stream operations), ``sync`` (a + device synchronization) or ``mark`` (host-side). ``bound`` is the log index of the record a + ``wait`` binds to (None: the event was never recorded). ``what`` is ``L``/``E`` for a + launch, the segment or ``E`` for a read and the hook for a mark. + """ + + kind: str + stream: FakeStream | None = None + event: FakeEvent | None = None + bound: int | None = None + what: str = "" + layer: int | None = None + + +class FakeStream: + def __init__(self, log: list[Entry], name: str) -> None: + self.log = log + self.name = name + + def wait_event(self, event: FakeEvent) -> None: + self.log.append(Entry("wait", stream=self, event=event, bound=event.last_record)) + + +class FakeEvent: + """An event; ``side`` tells the dense runtime's events from the weight manager's.""" + + def __init__(self, ops: FakeOps, side: str) -> None: + self.ops = ops + self.side = side + self.last_record: int | None = None + + def record(self, stream: FakeStream | None = None) -> None: + self.last_record = len(self.ops.log) + stream = self.ops.current if stream is None else stream + self.ops.log.append(Entry("record", stream=stream, event=self)) + + +class FakeOps: + """``DeviceOps`` that logs every stream operation in host order. + + Three streams: ``default`` (the executor loop and setup), ``compute`` (the execution stream + the model runs on) and ``copy``. + """ + + def __init__(self) -> None: + self.log: list[Entry] = [] + self.default = FakeStream(self.log, "default") + self.compute = FakeStream(self.log, "compute") + self.copy = FakeStream(self.log, "copy") + self.current = self.default + + def new_event(self) -> FakeEvent: + return FakeEvent(self, "dense") + + def current_stream(self) -> FakeStream: + return self.current + + @contextlib.contextmanager + def stream(self, stream: FakeStream) -> Iterator[FakeStream]: + previous, self.current = self.current, stream + try: + yield stream + finally: + self.current = previous + + def synchronize(self) -> None: + self.log.append(Entry("sync")) + + def is_current_stream_capturing(self) -> bool: + return False + + def do_multi_stream(self) -> bool: + return False + + def nvtx_range(self, message: str) -> contextlib.AbstractContextManager: + return contextlib.nullcontext() + + +class _SlotMap: + """The ``WeightBuffer`` calls ``DWDPWeightManager`` makes: MoE order ``o`` uses slot ``o % k``.""" + + device_id = 0 + local_start = 0 + local_end = 1 + + def __init__(self, layers: tuple[int, ...], num_slots: int) -> None: + self._layers = layers + self.num_slots = num_slots + + def buffer_index_for_layer(self, layer: int) -> int: + return self._layers.index(layer) % self.num_slots + + def get_full_tensor(self, layer: int, name: str) -> torch.Tensor: + return torch.empty(0) + + +# --------------------------------------------------------------------------- +# Happens-before over the log +# --------------------------------------------------------------------------- + + +def vector_clocks(log: list[Entry]) -> dict[int, dict[FakeStream, int]]: + """Vector clock of every stream operation: ``clock[s]`` = operations of ``s`` it follows. + + Streams run their operations in enqueue order; a wait makes the waiting stream follow every + operation its bound record follows, the record included; a device synchronization makes every + later operation follow every earlier one. + """ + stream_clocks: dict[FakeStream, dict[FakeStream, int]] = {} + clocks: dict[int, dict[FakeStream, int]] = {} + floor: dict[FakeStream, int] = {} + for index, entry in enumerate(log): + if entry.kind == "sync": + for clock in stream_clocks.values(): + for stream, count in clock.items(): + floor[stream] = max(floor.get(stream, 0), count) + for stream in stream_clocks: + stream_clocks[stream] = dict(floor) + continue + if entry.stream is None: + continue + clock = stream_clocks.setdefault(entry.stream, dict(floor)) + if entry.kind == "wait": + assert entry.bound is not None, f"wait #{index} binds to an event never recorded" + for stream, count in clocks[entry.bound].items(): + clock[stream] = max(clock.get(stream, 0), count) + clock[entry.stream] = clock.get(entry.stream, 0) + 1 + clocks[index] = dict(clock) + return clocks + + +def happens_before( + log: list[Entry], clocks: dict[int, dict[FakeStream, int]], first: int, second: int +) -> bool: + """Whether stream operation ``first`` completes before stream operation ``second`` starts.""" + stream = log[first].stream + return first != second and clocks[second].get(stream, 0) >= clocks[first][stream] + + +# --------------------------------------------------------------------------- +# Driver: the executor loop, the model, ConfigurableMoE and the managers around the runtime +# --------------------------------------------------------------------------- + + +class Abort(Exception): + """A forward that stops midway, as one whose kernel launch or collective raised.""" + + +class LaunchError(Exception): + """A dense copy launch that raised.""" + + +class Driver: + """Calls the managers and the runtime in program order; logs reads and copy launches. + + Every call is wrapped in a span keyed ``(forward, hook, layer)``; a ``mark`` entry opens it, so + each log entry can be attributed to the call that made it. + + Args: + segments: Enabled dense segments. + layers: Covered layers, which are also the expert DWDP MoE layers when attached. + num_slots: Expert prefetch slots (``DwdpConfig.num_prefetch_slots``). + standalone: Dense DWDP without expert DWDP: no expert copies or reads, and the runtime + does not prefetch the next forward. + """ + + def __init__( + self, + segments: tuple[SegmentKind, ...] = BOTH, + *, + layers: tuple[int, ...] = K3_LAYERS, + num_slots: int = 2, + standalone: bool = False, + ) -> None: + self.ops = FakeOps() + self.log = self.ops.log + self.segments = segments + self.layers = tuple(layers) + self.num_slots = num_slots + self.standalone = standalone + self.positions = {layer: position for position, layer in enumerate(self.layers)} + self.forward_index = 0 + self.spans: dict[tuple[int, str, int | None], tuple[int, int]] = {} + # Set to make the next dense launch raise, as a failing copy launch does. + self.fail_dense_launch = False + spec = {kind: SlotSpec(kind, RING_DEPTH[kind], ROW_ALIGN, 2) for kind in segments} + slots = {kind: allocate_slot_ring(spec[kind], "cpu") for kind in segments} + self.mark("construct", None) + self.runtime = DenseDWDP( + self.layers, + segments, + slots, + self.ops.copy, + prefetch_next_pass=not standalone, + ops=self.ops, + ) + self.runtime.set_launcher(self.launch_dense) + self.span("commit", None, lambda: self.runtime.commit(None, fill_byte=0)) + if standalone: + mapping = Mapping(world_size=2, rank=0, tp_size=2, enable_attention_dp=True) + self.dense_manager = DenseDwdpManager(DenseMode.STANDALONE, mapping=mapping) + # What DenseDwdpManager.setup stores once setup_dense_dwdp commits. + self.dense_manager._dense = self.runtime + self.weight_manager = self.manager = None + else: + with self.cuda(): + self.weight_manager = DWDPWeightManager( + weight_buffer=_SlotMap(self.layers, num_slots), + peer_views={}, + peer_ranges=[(0, 1), (1, 2)], + moe_layer_indices=list(self.layers), + weight_names=["w"], + dwdp_rank=0, + dwdp_size=2, + ) + # Nothing to release: no weight buffer, transport or device behind the fake. + self.weight_manager._released = True + self.dense_manager = DenseDwdpManager(DenseMode.ATTACHED) + self.dense_manager.attach(self.runtime, expert=self.weight_manager) + # DwdpManager after setup(): the weight manager and the attached dense manager. + self.manager = object.__new__(DwdpManager) + self.manager._weight_manager = self.weight_manager + self.manager._dense_manager = self.dense_manager + self.backend = nn.Module() + self.backend.w = nn.Parameter(torch.empty(0), requires_grad=False) + self.step_start = FakeEvent(self.ops, "executor") + self.step_end = FakeEvent(self.ops, "executor") + + # -- logging ------------------------------------------------------------- + + def mark(self, hook: str, layer: int | None) -> None: + self.log.append(Entry("mark", what=hook, layer=layer)) + + def span(self, hook: str, layer: int | None, action: Callable[[], object]) -> None: + start = len(self.log) + self.mark(hook, layer) + try: + action() + finally: + self.spans[(self.forward_index, hook, layer)] = (start, len(self.log)) + + def launch(self, what: str, layer: int) -> None: + self.log.append(Entry("launch", stream=self.ops.current, what=what, layer=layer)) + + def launch_dense(self, layer: int) -> None: + """The dense ``LayerLauncher``: one launch on the current stream.""" + if self.fail_dense_launch: + self.fail_dense_launch = False + raise LaunchError(f"L({layer})") + self.launch("L", layer) + + def read(self, what: str, layer: int) -> None: + """A kernel that reads ``what``'s weights of ``layer`` on the current stream.""" + self.log.append(Entry("read", stream=self.ops.current, what=what, layer=layer)) + + @contextlib.contextmanager + def cuda(self) -> Iterator[None]: + """``torch.cuda`` as the weight manager sees it: the fake seam, and its copies logged. + + The patches are process-wide, so the cyclic garbage collector is paused meanwhile: the + finalizer of another test's weight manager or runtime would call them and log a device + synchronization here. + """ + ops = self.ops + fakes = { + "Stream": lambda device=None: ops.copy, + "Event": lambda: FakeEvent(ops, "expert"), + "current_stream": lambda device=None: ops.current, + "stream": ops.stream, + "synchronize": lambda device=None: ops.synchronize(), + } + with contextlib.ExitStack() as stack: + if gc.isenabled(): + gc.disable() + stack.callback(gc.enable) + for name, fake in fakes.items(): + stack.enter_context(patch(f"{_WEIGHT_MANAGER}.torch.cuda.{name}", fake)) + stack.enter_context( + patch.object( + DWDPWeightManager, + "_prefetch_layer_per_slice", + lambda _manager, layer: self.launch(EXPERT, layer), + ) + ) + yield + + # -- executor loop, model and ConfigurableMoE --------------------------------- + + def start(self) -> None: + """The executor loop before a forward (``PyExecutor._prime_dwdp``), default stream.""" + if self.standalone: + start = self.dense_manager.prefetch_first_layers + else: + start = self.manager.start_forward + with self.cuda(), self.ops.stream(self.ops.default): + self.span("start", None, start) + + def h0(self) -> None: + """The model's forward entry: ``DenseDwdpManager.ensure_primed``.""" + self.dense_manager.ensure_primed() + + def invalidate(self) -> None: + """The executor after a control action: ``DwdpManager.invalidate_prefetch``.""" + with self.cuda(), self.ops.stream(self.ops.default): + self.span("invalidate", None, self.manager.invalidate_prefetch) + + def layer(self, layer: int, stop: str | None = None) -> None: + """One decoder layer in program order; raises :class:`Abort` right before ``stop``. + + ``stop`` is ``H1``, ``H2``, ``finish`` or ``H3``; ``fill`` makes the first dense copy + launch of the layer's hooks raise instead. + """ + runtime = self.runtime + if stop == "fill": + self.fail_dense_launch = True + try: + if stop == "H1": + raise Abort + self.span("H1", layer, lambda: runtime.wait_layer(layer)) + self.read(ATTN.value, layer) + if stop == "H2": + raise Abort + self.span("H2", layer, lambda: runtime.attention_done(layer)) + if MOE in self.segments: + self.read(MOE.value, layer) # the router and the latent down-projection + if not self.standalone: + self.span("bind", layer, lambda: self.manager.wait_and_bind(self.backend, layer)) + self.read(EXPERT, layer) + if stop == "finish": + raise Abort + if not self.standalone: + self.span("finish", layer, lambda: self.manager.finish_layer(layer)) + if MOE in self.segments: + self.read(MOE.value, layer) # the latent up-projection and the shared expert + if stop == "H3": + raise Abort + self.span("H3", layer, lambda: runtime.layer_done(layer)) + finally: + self.fail_dense_launch = False + + def forward( + self, + *, + start: bool = True, + h0: bool = True, + stop_at: tuple[str, int] | None = None, + fail_h0: bool = False, + ) -> bool: + """One executor step; False if the forward stopped midway. + + ``start``: the executor loop starts the forward (the non-overlap loop; warmup and the + KV-cache estimation do not). ``h0``: the model's forward entry runs H0. + ``stop_at=(stop, layer)``: the forward stops in ``layer`` (see :meth:`layer`). + ``fail_h0``: the first dense copy launch of H0 raises. The execution stream joins the + default stream after a forward that completes, as ``PyExecutor._forward_step`` does. + """ + if start: + self.start() + ops = self.ops + try: + with self.cuda(): + self.step_start.record(ops.default) + ops.compute.wait_event(self.step_start) + with ops.stream(ops.compute): + if h0: + self.fail_dense_launch = fail_h0 + try: + self.span("H0", None, self.h0) + finally: + self.fail_dense_launch = False + for layer in self.layers: + stop = stop_at[0] if stop_at is not None and stop_at[1] == layer else None + self.layer(layer, stop) + self.step_end.record(ops.compute) + ops.default.wait_event(self.step_end) + except (Abort, LaunchError): + return False + finally: + self.forward_index += 1 + return True + + # -- analysis -------------------------------------------------------------- + + def slots_of(self, what: str, layer: int) -> list[tuple[str, int]]: + """The slots a launch of ``what`` writes, or a read of ``what`` reads, for ``layer``.""" + position = self.positions[layer] + if what == EXPERT: + return [(EXPERT, position % self.num_slots)] + kinds = self.segments if what == "L" else (SegmentKind(what),) + return [(kind.value, position % RING_DEPTH[kind]) for kind in kinds] + + def assert_race_free(self) -> int: + """Check every logged read against every copy into its slot; returns the reads checked. + + RAW: the latest copy of the read layer into the slot completes before the read. WAR: the + next copy into the slot starts after the read completes; copies into one slot run on one + stream in enqueue order, so every later copy does too. + """ + log = self.log + clocks = vector_clocks(log) + writes: dict[tuple[str, int], list[tuple[int, int]]] = {} + for index, entry in enumerate(log): + if entry.kind == "launch": + for key in self.slots_of(entry.what, entry.layer): + writes.setdefault(key, []).append((index, entry.layer)) + for key, slot_writes in writes.items(): + assert len({log[index].stream for index, _ in slot_writes}) == 1, key + reads = 0 + for index, entry in enumerate(log): + if entry.kind != "read": + continue + (key,) = self.slots_of(entry.what, entry.layer) + slot_writes = writes.get(key, []) + earlier = [(w, layer) for w, layer in slot_writes if w < index] + sources = [w for w, layer in earlier if layer == entry.layer] + assert sources, f"read #{index} of {entry.what}({entry.layer}) precedes its copy" + source = sources[-1] + assert happens_before(log, clocks, source, index), ( + f"RAW: read #{index} of {entry.what}({entry.layer}) may start before copy #{source}" + ) + later = [(w, layer) for w, layer in slot_writes if w > source] + if later: + w, layer = later[0] + assert happens_before(log, clocks, index, w), ( + f"WAR: copy #{w} of {layer} into {key} may start before read #{index} of " + f"{entry.what}({entry.layer}) completes" + ) + reads += 1 + return reads + + def syncs(self, start: int = 0, end: int | None = None) -> list[int]: + """Log indices of the device synchronizations in ``[start, end)``.""" + stop = len(self.log) if end is None else end + return [i for i in range(start, stop) if self.log[i].kind == "sync"] + + def assert_no_host_sync(self) -> None: + """Check that the host never synchronized after commit, so it can run ahead of the GPU.""" + committed = self.spans[(0, "commit", None)][1] + assert self.syncs(committed) == [] + + def records(self, start: int, end: int, side: str, stream: FakeStream | None = None) -> list: + """Log indices of the records of ``side``'s events in ``[start, end)``, on ``stream``.""" + return [ + index + for index in range(start, end) + if self.log[index].kind == "record" + and self.log[index].event.side == side + and (stream is None or self.log[index].stream is stream) + ] + + def launches( + self, start: int = 0, end: int | None = None, what: str | None = None + ) -> list[tuple[str, int]]: + """``(what, layer)`` of the copy launches in ``[start, end)``.""" + return [ + (entry.what, entry.layer) + for entry in self.log[start:end] + if entry.kind == "launch" and (what is None or entry.what == what) + ] + + def launches_in(self, key: tuple[int, str, int | None]) -> list[tuple[str, int]]: + """The copy launches made inside the span ``key``.""" + start, end = self.spans[key] + return self.launches(start, end) + + def forward_bounds(self, forward: int) -> tuple[int, int]: + """Log range of forward ``forward``: from its first span to the next forward's first.""" + starts = {} + for (index, _hook, _layer), (start, _end) in self.spans.items(): + starts[index] = min(start, starts.get(index, start)) + end = starts.get(forward + 1, len(self.log)) + return starts[forward], end + + def stream_ops(self, start: int, end: int | None = None) -> list[Entry]: + return [e for e in self.log[start:end] if e.kind in ("record", "wait", "launch")] + + +def _check_steady(driver: Driver) -> None: + """No fence, demand fetch or host synchronization so far, and every read race-free.""" + assert driver.runtime.war_fences == 0 + assert driver.runtime.demand_fetches == 0 + driver.assert_no_host_sync() + driver.assert_race_free() + + +# --------------------------------------------------------------------------- +# Two complete Kimi-K3 forwards +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module", params=[BOTH, (ATTN,)], ids=["attn+moe", "attn"]) +def two_forwards(request: pytest.FixtureRequest) -> Driver: + """Two complete K3 forwards (92 layers each), started by the executor; tests only read it.""" + driver = Driver(segments=request.param) + assert driver.forward() and driver.forward() + return driver + + +def test_read_windows_end_at_h2_and_h3_whose_fill_runs_into_the_next_forward(two_forwards): + driver = two_forwards + log = driver.log + has_moe = MOE in driver.segments + window_ends = ("H2", "H3") if has_moe else ("H2",) + # Where the fill enqueues the next forward's first layers: H3 of the last two layers, whose + # MoE-dense windows hold those layers' slots until then; H2 with the attention segment alone. + boundary = "H3" if has_moe else "H2" + for forward in range(2): + for layer in driver.layers: + for hook in window_ends: + start, end = driver.spans[(forward, hook, layer)] + ops = driver.stream_ops(start, end) + # The hook ends the window first, so each copy it enqueues binds to that record. + assert ops[0].kind == "record" and ops[0].stream is driver.ops.compute + consume = driver.records(start, end, "dense", driver.ops.compute) + assert len(consume) == 1, (forward, hook, layer) + expected = {"H2": [], "H3": []} + if layer + 2 <= 92: + expected["H2"] = [("L", layer + 2)] + else: + expected[boundary] = [("L", layer - 90)] + for hook in window_ends: + assert driver.launches_in((forward, hook, layer)) == expected[hook], ( + forward, + hook, + layer, + ) + # Those are the only dense compute records after commit: no other hook ends a read window. + committed = driver.spans[(0, "commit", None)][1] + records = driver.records(committed, len(log), "dense", driver.ops.compute) + assert len(records) == 2 * len(driver.layers) * len(window_ends) + # The expert slots are released right after each MoE, on the execution stream. + finishes = driver.records(committed, len(log), "expert", driver.ops.compute) + assert len(finishes) == 2 * len(driver.layers) + # The first two layers: primed by the first forward's start, then by the end of each forward. + for what in ("L", EXPERT): + copies = Counter(layer for kind, layer in driver.launches() if kind == what) + assert copies == Counter({layer: 3 if layer <= 2 else 2 for layer in driver.layers}) + assert not driver.runtime.forward_open and driver.runtime.primed + reads_per_layer = (3 if has_moe else 1) + 1 + assert driver.assert_race_free() == 2 * len(driver.layers) * reads_per_layer + _check_steady(driver) + + +# --------------------------------------------------------------------------- +# Copy order on the shared copy stream +# --------------------------------------------------------------------------- + + +def copy_order_problems(driver: Driver) -> list[str]: + """How two complete K3 forwards with two expert slots break the expected copy order. + + * Cold start: the executor's start enqueues ``L(1), L(2), E(1), E(2)``. + * Within a forward: each ``L(l)`` of layers 3..92 right before ``E(l)``. + * At the end of a forward: the next one's ``E(1), L(1), E(2), L(2)`` (``L(1), E(1), L(2), + E(2)`` with the attention segment alone). + * The second forward's start and H0 enqueue nothing; it copies every layer once. + """ + problems = [] + cold = driver.launches_in((0, "start", None)) + if cold != [("L", 1), ("L", 2), (EXPERT, 1), (EXPERT, 2)]: + problems.append(f"cold start {cold}") + for forward in range(2): + first, end = driver.spans[(forward, "H1", 1)][0], driver.forward_bounds(forward)[1] + sequence = driver.launches(first, end) + for layer in driver.layers[2:]: + at = sequence.index((EXPERT, layer)) + if at == 0 or sequence[at - 1] != ("L", layer): + problems.append(f"forward {forward}: E({layer}) not right after L({layer})") + tail = sequence[-4:] + if MOE in driver.segments: + expected_tail = [(EXPERT, 1), ("L", 1), (EXPERT, 2), ("L", 2)] + else: + expected_tail = [("L", 1), (EXPERT, 1), ("L", 2), (EXPERT, 2)] + if tail != expected_tail: + problems.append(f"forward {forward} ends with {tail}") + for hook in ("start", "H0"): + if driver.launches_in((1, hook, None)): + problems.append( + f"the second forward's {hook} enqueued {driver.launches_in((1, hook, None))}" + ) + second = Counter(driver.launches(*driver.forward_bounds(1))) + if second != Counter((what, layer) for what in ("L", EXPERT) for layer in driver.layers): + problems.append("the second forward does not copy every layer once") + return problems + + +@SEGMENT_SETS +def test_the_attached_copy_order(segments): + driver = Driver(segments) + assert driver.forward() and driver.forward() + assert copy_order_problems(driver) == [] + _check_steady(driver) + + +@SEGMENT_SETS +def test_the_order_check_fails_if_h2_does_not_fill(monkeypatch, segments): + """Sensitivity: with the fill at H3 only, ``L(l)`` follows ``E(l)``, issued at the finish.""" + attention_done = DenseDWDP.attention_done + + def attention_done_without_fill(runtime: DenseDWDP, layer: int) -> None: + runtime._fill = lambda: None + try: + attention_done(runtime, layer) + finally: + del runtime._fill + + monkeypatch.setattr(DenseDWDP, "attention_done", attention_done_without_fill) + driver = Driver(segments) + assert driver.forward() and driver.forward() + problems = copy_order_problems(driver) + assert any("not right after" in problem for problem in problems), problems + + +def test_the_order_check_fails_if_the_expert_prefetch_starts_first(monkeypatch): + """Sensitivity: an expert start ahead of the dense priming breaks the cold start.""" + driver = Driver() + prefetch = DenseDWDP.prefetch + + def prefetch_after_the_expert_start(runtime: DenseDWDP, layer: int) -> None: + driver.weight_manager.start_forward() + prefetch(runtime, layer) + + monkeypatch.setattr(DenseDWDP, "prefetch", prefetch_after_the_expert_start) + assert driver.forward() and driver.forward() + problems = copy_order_problems(driver) + assert any(problem.startswith("cold start") for problem in problems), problems + + +# --------------------------------------------------------------------------- +# Expert prefetch slots and model sizes +# --------------------------------------------------------------------------- + + +@SEGMENT_SETS +@pytest.mark.parametrize("num_layers", [1, 2, 3, 4, 5, 92]) +@pytest.mark.parametrize("num_slots", [2, 3, 4]) +def test_attached_forwards_are_race_free(segments, num_layers, num_slots): + """A warmup forward the executor did not start, then three it did: every read is ordered.""" + layers = tuple(range(1, num_layers + 1)) + driver = Driver(segments, layers=layers, num_slots=num_slots) + assert driver.forward(start=False) + assert driver.runtime.primed + for _ in range(3): + assert driver.forward() + assert driver.runtime.primed and not driver.runtime.forward_open + # The last forward copies every layer once, dense and expert. + last = Counter(driver.launches(*driver.forward_bounds(3))) + assert last == Counter((what, layer) for what in ("L", EXPERT) for layer in layers) + _check_steady(driver) + + +# --------------------------------------------------------------------------- +# Fences, demand fetches, the WAR barrier and recovery +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("prime", [True, False], ids=["primed", "demand-fetched"]) +def test_fence_when_a_holder_has_an_open_read_window(prime): + """A forward dies inside attention(1); a prefetch that skips the WAR barrier gets the fence. + + H1(1) opened the read windows of ``A(1)`` and ``D(1)``, whether priming or, in a forward that + skipped priming, H1's demand fetch enqueued ``L(1)``, and no consume record closed them: only + the fence, recorded on compute after the reads, orders the copy after them. + """ + driver = Driver() + runtime = driver.runtime + assert not driver.forward(start=prime, h0=prime, stop_at=("H2", 1)) + assert runtime.forward_open and runtime.demand_fetches == int(not prime) + with driver.ops.stream(driver.ops.compute): + driver.span("again", None, lambda: runtime.prefetch(1)) + assert runtime.war_fences == 1 # one per prefetch, not one per slot + assert driver.launches_in((1, "again", None)) == [("L", 1)] + assert driver.forward() # H0 runs the barrier, then primes + assert runtime.war_fences == 1 and runtime.demand_fetches == int(not prime) + driver.assert_race_free() + driver.assert_no_host_sync() + + +def test_war_barrier_then_reprime_after_an_abort(): + driver = Driver() + runtime = driver.runtime + assert driver.forward() + assert not driver.forward(stop_at=("H3", 41)) # dies inside the MoE block of layer 41 + assert runtime.forward_open and not runtime.primed + start = len(driver.log) + with driver.ops.stream(driver.ops.compute): + driver.span("barrier", None, runtime.war_barrier) + records = driver.stream_ops(start) + assert [entry.kind for entry in records] == ["record"] * (RING_DEPTH[ATTN] + RING_DEPTH[MOE]) + assert all(entry.stream is driver.ops.compute for entry in records) + assert len({id(entry.event) for entry in records}) == len(records) + for kind in BOTH: + for index in range(RING_DEPTH[kind]): + assert runtime.slot_tag(kind, index) == EMPTY_TAG + assert runtime.slot_consumed(kind, index) + assert not runtime.primed and not runtime.forward_open + assert driver.forward() # the start primes; the expert prefetch restarts its sequence + assert driver.launches_in((2, "start", None)) == [ + ("L", 1), + ("L", 2), + (EXPERT, 1), + (EXPERT, 2), + ] + assert driver.forward() + _check_steady(driver) + + +ABORT_POINTS = ["H1", "H2", "finish", "H3"] + + +@SEGMENT_SETS +@pytest.mark.parametrize("recovery", ["executor", "model"]) +@pytest.mark.parametrize("position", [0, 1, 46, 90, 91]) +@pytest.mark.parametrize("stop", ABORT_POINTS) +def test_an_aborted_forward_recovers_both_sides(segments, recovery, position, stop): + """A K3 forward stops at ``stop`` of the layer at ``position``; the next forward recovers. + + The executor's start leaves an open forward to H0, which runs the dense WAR barrier on the + execution stream, primes ``L(1), L(2)`` and then starts the expert prefetch, which restarts + its sequence, synchronizing the device only if the stopped forward left a MoE layer bound. + """ + driver = Driver(segments) + runtime = driver.runtime + layer = driver.layers[position] + assert driver.forward() and driver.forward() + aborted_at = len(driver.log) + assert not driver.forward(stop_at=(stop, layer)) + was_open = runtime.forward_open + assert was_open == (stop != "H1" or position > 0) + bound = stop == "finish" + # The expert sequence restarts unless the stopped forward bound no MoE layer or finished all. + last = len(driver.layers) - 1 + restarts = bound or (position > 0 if stop in ("H1", "H2") else position != last) + assert driver.forward(start=recovery == "executor") + if recovery == "executor": + # An open forward is left to H0; otherwise the forward is primed already. + assert driver.launches_in((3, "start", None)) == [] + h0_start, h0_end = driver.spans[(3, "H0", None)] + barrier = driver.records(h0_start, h0_end, "dense", driver.ops.compute) + assert len(barrier) == (sum(RING_DEPTH[kind] for kind in segments) if was_open else 0) + assert driver.records(h0_start, h0_end, "expert", driver.ops.compute) == [] + expected = [("L", 1), ("L", 2)] * was_open + [(EXPERT, 1), (EXPERT, 2)] * restarts + assert driver.launches(h0_start, h0_end) == expected + assert driver.syncs(aborted_at) == driver.syncs(h0_start, h0_end) + assert len(driver.syncs(aborted_at)) == int(bound) + assert runtime.primed and not runtime.forward_open + assert driver.forward() + assert runtime.war_fences == 0 and runtime.demand_fetches == 0 + driver.assert_race_free() + + +def test_invalidate_prefetch_keeps_the_dense_copies(): + """New expert weights regather ``E(1), E(2)``; the dense weights cannot change.""" + driver = Driver() + assert driver.forward() and driver.forward() + driver.invalidate() + assert driver.launches_in((2, "invalidate", None)) == [] + assert driver.forward() + assert driver.launches_in((2, "start", None)) == [(EXPERT, 1), (EXPERT, 2)] + assert driver.launches_in((2, "H0", None)) == [] + assert driver.runtime.primed + _check_steady(driver) + + +# --------------------------------------------------------------------------- +# Standalone: no expert DWDP, the dense manager primes every forward +# --------------------------------------------------------------------------- + + +@SEGMENT_SETS +def test_standalone_keeps_the_per_pass_schedule(segments): + """The executor's start primes ``L(1), L(2)``; H2 enqueues ``L(l + 2)``; nothing more.""" + driver = Driver(segments, standalone=True) + runtime = driver.runtime + assert not runtime.prefetch_next_pass + for forward in range(2): + assert driver.forward() + assert driver.launches_in((forward, "start", None)) == [("L", 1), ("L", 2)] + assert driver.launches_in((forward, "H0", None)) == [] + for layer in driver.layers: + expected = [("L", layer + 2)] if layer + 2 <= 92 else [] + assert driver.launches_in((forward, "H2", layer)) == expected, (forward, layer) + if MOE in segments: + assert driver.launches_in((forward, "H3", layer)) == [], (forward, layer) + assert not runtime.primed and not runtime.forward_open + assert driver.launches(what=EXPERT) == [] + copies = Counter(layer for _what, layer in driver.launches()) + assert copies == Counter({layer: 2 for layer in driver.layers}) + _check_steady(driver) + + +@SEGMENT_SETS +@pytest.mark.parametrize("recovery", ["executor", "model"]) +@pytest.mark.parametrize("position", [0, 1, 46, 91]) +@pytest.mark.parametrize("stop", ["H1", "H2", "H3"]) +def test_a_standalone_forward_recovers_after_an_abort(segments, recovery, position, stop): + """H0 runs the WAR barrier and primes again; the executor's start leaves an open forward.""" + driver = Driver(segments, standalone=True) + runtime = driver.runtime + assert driver.forward() + assert not driver.forward(stop_at=(stop, driver.layers[position])) + was_open = runtime.forward_open + assert driver.forward(start=recovery == "executor") + start, end = driver.spans[(2, "H0", None)] + if was_open: + assert driver.launches(start, end) == [("L", 1), ("L", 2)] + assert len(driver.records(start, end, "dense", driver.ops.compute)) == sum( + RING_DEPTH[kind] for kind in segments + ) + assert not runtime.primed and not runtime.forward_open + assert driver.forward() + _check_steady(driver) + + +# --------------------------------------------------------------------------- +# Random sequences +# --------------------------------------------------------------------------- + + +def _random_sequence(seed: int) -> str | None: + """One random sequence of forwards; what went wrong, or None.""" + rng = random.Random(seed) + num_layers = rng.choice([1, 2, 3, 5, 8]) + segments = rng.choice([BOTH, (ATTN,)]) + standalone = rng.random() < 0.25 + driver = Driver( + segments, + layers=tuple(range(1, num_layers + 1)), + num_slots=rng.choice([2, 3, 4]), + standalone=standalone, + ) + runtime = driver.runtime + steps = [] + assert driver.forward(start=False) + for _ in range(12): + roll = rng.random() + start = rng.random() < 0.7 + if roll < 0.35: + stop = rng.choice(["H1", "H2", "finish", "H3", "fill"]) + layer = rng.choice(driver.layers) + steps.append(f"stop at {stop} of L{layer}, start={start}") + driver.forward(start=start, stop_at=(stop, layer)) + elif roll < 0.45: + steps.append("start, then the forward fails before the model") + driver.start() + elif roll < 0.55: + if not standalone: + steps.append("invalidate") + driver.invalidate() + elif roll < 0.62: + steps.append(f"a launch of H0 fails, start={start}") + driver.forward(start=start, fail_h0=True) + else: + steps.append(f"full forward, start={start}") + demand_fetches = runtime.demand_fetches + if not driver.forward(start=start): + return f"{steps}: the forward stopped" + if runtime.demand_fetches != demand_fetches or runtime.primed == standalone: + return f"{steps}: demand fetch, or primed state wrong after a complete forward" + if runtime.war_fences: + return f"{steps}: {runtime.war_fences} fences" + try: + driver.assert_race_free() + except AssertionError as exc: + return f"{steps}: {exc}" + return None + + +def test_random_sequences_are_race_free(): + """Aborts at every point (also inside a fill or H0's priming), starts whose forward fails + before the model and invalidations, in random order, attached and standalone: no race, fence + or window error.""" + problems = {seed: problem for seed in range(200) if (problem := _random_sequence(seed))} + assert problems == {} diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_setup.py b/tests/unittest/_torch/modules/dwdp/test_dense_setup.py new file mode 100644 index 000000000000..c0bbaf9bbfc1 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_setup.py @@ -0,0 +1,624 @@ +# 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. +"""CPU tests of the dense DWDP setup. + +Every rank of an ``N``-rank dense group is a thread that runs ``setup_dense_dwdp`` on its own copy +of the toy Kimi-K3 decoder of ``test_dense_layout.py``. The threads gather over +``StoreDwdpGroup`` on one in-process ``HashStore`` with every wait bounded, so a rank that skips or +adds a collective fails the test instead of hanging it. Each rank keeps its rows in the real +``LoopbackShardStore``; fakes stand in for the device: ``FakeOps`` of ``test_dense_protocol.py`` +for the stream and event seam and ``FakeLauncher``, which runs the real copy plans as host +memmoves, each checked to read a store chunk and write a live slot. Faults are injected in each +phase. +""" + +from __future__ import annotations + +import ctypes +import gc +import re +import sys +import threading +import traceback +import weakref +from collections import Counter +from dataclasses import dataclass, replace +from types import SimpleNamespace + +import pytest +import torch +from test_dense_layout import COVERED, LAYER_KINDS, TOY_MIN_BYTES, ToyModel, collect +from test_dense_protocol import FakeOps +from torch import nn + +from tensorrt_llm._torch.modules.dwdp.dense import repoint as dense_repoint +from tensorrt_llm._torch.modules.dwdp.dense import setup as dense_setup +from tensorrt_llm._torch.modules.dwdp.dense.config import DenseConfig +from tensorrt_llm._torch.modules.dwdp.dense.runtime import DenseDWDP +from tensorrt_llm._torch.modules.dwdp.dense.setup import PhaseReport, setup_dense_dwdp +from tensorrt_llm._torch.modules.dwdp.dense.store import LoopbackShardStore +from tensorrt_llm._torch.modules.dwdp.dense.types import ( + DenseRepointError, + DenseSetupError, + SegmentKind, +) +from tensorrt_llm._torch.modules.dwdp.rendezvous import StoreDwdpGroup + +pytestmark = pytest.mark.cpu_only + +#: Seconds one collective waits for the slowest rank; a rank that skips a collective times out. +COLLECTIVE_TIMEOUT_S = 20 +#: Seconds a test waits for its rank threads. +JOIN_TIMEOUT_S = 90 +#: Free device memory every rank reports. +PLENTY_FREE_BYTES = 64 << 30 +#: The last MoE-dense anchor of the last covered layer; the corruption faults flip its last byte. +LAST_MOE_ANCHOR = f"layers.{COVERED[-1]}.block_sparse_moe.routed_expert_up_proj.weight" + + +# --------------------------------------------------------------------------- +# Module tensor snapshots +# --------------------------------------------------------------------------- + + +def module_tensors(decoder): + """Every tensor the decoder layers hold, by ``layers..``.""" + found = {} + for index, layer in enumerate(decoder.layers): + for path, module in layer.named_modules(prefix=f"layers.{index}"): + for name, value in [*module._parameters.items(), *module._buffers.items()]: + if value is not None: + found[f"{path}.{name}"] = value + for name, value in vars(module).items(): + if isinstance(value, torch.Tensor): + found[f"{path}.{name}"] = value + elif isinstance(value, (tuple, list)) and value: + if all(isinstance(item, torch.Tensor) for item in value): + found.update({f"{path}.{name}[{i}]": item for i, item in enumerate(value)}) + return found + + +def placement(tensor): + storage = tensor.untyped_storage().data_ptr() + return ( + storage, + tensor.storage_offset(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + ) + + +def raw_bytes(tensor): + return tensor.detach().contiguous().reshape(-1).view(torch.uint8).clone() + + +class Snapshot: + """A decoder's tensors before setup: placements, bytes and Parameter objects. + + It holds no strong reference to an original storage, so the originals can die at commit. + """ + + def __init__(self, decoder): + tensors = module_tensors(decoder) + self.placements = {path: placement(tensor) for path, tensor in tensors.items()} + self.values = {path: raw_bytes(tensor) for path, tensor in tensors.items()} + self.parameters = {p: t for p, t in tensors.items() if isinstance(t, nn.Parameter)} + + def assert_restored(self, decoder): + """Every tensor is back: same storage, view, dtype and bytes; same Parameter objects.""" + now = module_tensors(decoder) + assert now.keys() == self.placements.keys() + changed = [ + path + for path, tensor in now.items() + if placement(tensor) != self.placements[path] + or not torch.equal(raw_bytes(tensor), self.values[path]) + ] + assert changed == [] + assert [path for path, param in self.parameters.items() if now[path] is not param] == [] + + +# --------------------------------------------------------------------------- +# Fakes +# --------------------------------------------------------------------------- + + +class HostMemory: + """The store chunks and slots of a test's ranks; every fake copy must stay inside them. + + Slots are tracked by weak reference. + """ + + def __init__(self): + self._lock = threading.Lock() + self.chunks = [] + self.slots = [] + + def add_chunk(self, start, end): + with self._lock: + self.chunks.append((start, end)) + + def add_slot(self, storage): + entry = (threading.current_thread(), storage.data_ptr(), storage.numel()) + with self._lock: + self.slots.append((*entry, weakref.ref(storage.untyped_storage()))) + + def live_slots(self, thread=None): + """Owners of slot storages that are still alive; only ``thread``'s if given.""" + with self._lock: + return [s[0].name for s in self.slots if s[3]() is not None and thread in (None, s[0])] + + def check_copy(self, dst, src, nbytes): + with self._lock: + chunks = list(self.chunks) + slots = [(start, start + size) for _, start, size, ref in self.slots if ref()] + if not any(start <= src and src + nbytes <= end for start, end in chunks): + raise AssertionError(f"copy source {src:#x}+{nbytes} is outside every store chunk") + if not any(start <= dst and dst + nbytes <= end for start, end in slots): + raise AssertionError(f"copy destination {dst:#x}+{nbytes} is outside every slot") + + +def _called_from_finalizer(): + """Whether a ``__del__`` is on the stack: garbage collection, not setup, releases.""" + frame = sys._getframe(1) + while frame is not None: + if frame.f_code.co_name == "__del__": + return True + frame = frame.f_back + return False + + +class LoopbackStore(LoopbackShardStore): + """The real loopback store, every rank's rows in this process, with its calls counted. + + ``spec.fail`` makes the step it names raise before it runs. + """ + + def __init__(self, layout, rank, device, *, owner): + super().__init__(layout, rank, device) + self.owner = owner + self.calls = Counter() + self.finalizer_releases = 0 + + def _step(self, name, step, *args): + self.owner.log.append(name) + self.calls[name] += 1 + if self.owner.spec.fail == name: + raise RuntimeError(f"injected: store {name} failed") + return step(*args) + + def allocate(self): + self._step("allocate", super().allocate) + + def write_local_rows(self): + self._step("write", super().write_local_rows) + + def export(self): + return self._step("export", super().export) + + def import_peers(self, payloads): + self._step("import", super().import_peers, payloads) + for bases in self.chunk_bases(): + for chunk, base in zip(self._layout.chunks, bases): + self.owner.memory.add_chunk(base, base + chunk.handle_bytes) + + def release(self): + self.finalizer_releases += _called_from_finalizer() + self._step("release", super().release) + + +class FakeLauncher: + """Runs a rank's copy plans as host memmoves; round 0 is Phase V (P4), round 1 the P5 check. + + With ``spec.corrupt`` (4 or 5) it flips the last byte of ``LAST_MOE_ANCHOR`` in its slot right + after it copies the last layer in that phase. + """ + + def __init__(self, plans, rank): + self.plans = dict(plans) + self.rank = rank + self.launches = 0 + self.releases = 0 + + def __call__(self, layer): + rank = self.rank + phase = 4 + self.launches // len(self.plans) + self.launches += 1 + rank.log.append(f"launch L{layer}") + for dst, src, nbytes in self.plans[layer].descriptors: + rank.memory.check_copy(dst, src, nbytes) + ctypes.memmove(dst, src, nbytes) + if (rank.spec.corrupt, layer) == (phase, COVERED[-1]): + ctypes.c_uint8.from_address(rank.corrupt_address).value ^= 0xFF + + def release(self): + self.releases += 1 + + +def corruption_address(layout, slot_bases): + """Address of the last byte of ``LAST_MOE_ANCHOR`` in the slot its layer is gathered into.""" + layer = COVERED[-1] + anchor = layout.segment(layer, SegmentKind.MOE).anchors[-1] + base = slot_bases[SegmentKind.MOE][layout.slot_index(SegmentKind.MOE, layer)] + return base + anchor.blob_offset + anchor.nbytes - 1 + + +class CountingComm: + """The rank's ``GroupComm``: a ``StoreDwdpGroup`` whose collectives are recorded.""" + + def __init__(self, group): + self.group = group + self.sent = [] + + def Get_rank(self): + return self.group.Get_rank() + + def Get_size(self): + return self.group.Get_size() + + def allgather(self, obj): + self.sent.append(obj) + return self.group.allgather(obj) + + +class CaptureLogger: + """Stands in for the setup and repoint modules' ``logger``; keeps each message's thread.""" + + def __init__(self): + self.records = [] + + def __getattr__(self, level): + return lambda message: self.records.append((threading.current_thread(), level, message)) + + +# --------------------------------------------------------------------------- +# Rank threads +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class RankSpec: + """One rank's switches, model, free device memory and injected fault. + + ``fail`` names a store step that raises; ``corrupt`` (4 or 5) flips a slot byte in that phase + (see ``FakeLauncher``); ``perturb`` adds 1 to the first element of a decoder parameter. + ``attached`` passes the expert DWDP MoE layers (``moe_layers``); False sets up standalone. + """ + + requested: bool = True + required: bool = True + strict: bool = False + debug: int = 0 + kinds: tuple[str, ...] = LAYER_KINDS + free_bytes: int = PLENTY_FREE_BYTES + fail: str | None = None + corrupt: int | None = None + perturb: str | None = None + attached: bool = True + + +class Rank: + """Everything one rank thread owns: model, fakes, its log and the outcome of setup.""" + + def __init__(self, index, spec, memory): + self.index = index + self.spec = spec + self.memory = memory + self.log = [] + # Like the deployed VL wrapper: ``model.model`` is the decoder that holds ``layers``. + self.model = nn.Module() + self.model.model = self.decoder = ToyModel(kinds=spec.kinds) + if spec.perturb is not None: + with torch.no_grad(): + self.decoder.get_parameter(spec.perturb).view(-1)[0] += 1 + self.snapshot = Snapshot(self.decoder) + self.ops = FakeOps() + self.transfer = SimpleNamespace( + copy_stream=self.ops.copy, use_gather=False, use_batched_prefetch=False + ) + self.stores = [] + self.launchers = [] + #: ``(stores released, live slots of this rank)`` at every ``torch.cuda.empty_cache``. + self.cache_states = [] + self.corrupt_address = None + self.thread = self.comm = self.result = self.error = None + self.trace = "" + + def store_factory(self, layout, rank, device): + store = LoopbackStore(layout, rank, device, owner=self) + self.stores.append(store) + return store + + def launcher_factory(self, plans, transfer, device): + assert transfer is self.transfer and device == torch.device("cpu") + assert self.ops.current is self.ops.copy, "the launcher is built on the copy stream" + self.launchers.append(FakeLauncher(plans, self)) + return self.launchers[-1] + + def on_empty_cache(self): + self.log.append("empty_cache") + released = all(store.calls["release"] for store in self.stores) + self.cache_states.append((released, self.memory.live_slots(self.thread))) + + def run(self, store, size): + group = StoreDwdpGroup( + global_rank=self.index, + dwdp_size=size, + master_addr="unused", + timeout_s=COLLECTIVE_TIMEOUT_S, + store=store, + ) + self.comm = CountingComm(group) + config = DenseConfig( + enabled=self.spec.requested, + required=self.spec.required, + min_bytes=TOY_MIN_BYTES, + strict=self.spec.strict, + debug=self.spec.debug, + ) + try: + self.result = setup_dense_dwdp( + self.model, + self.transfer, + self.comm, + dwdp_rank=self.index, + dwdp_size=size, + device_id=0, + moe_layers=list(COVERED) if self.spec.attached else None, + config=config, + store_factory=self.store_factory, + ops=self.ops, + device=torch.device("cpu"), + launcher_factory=self.launcher_factory, + mem_get_info=lambda: (self.spec.free_bytes, 4 * PLENTY_FREE_BYTES), + ) + except Exception as exc: # the test inspects every rank's outcome + self.error = (type(exc), str(exc)) + self.trace = traceback.format_exc() + + +class Group: + """``N`` rank threads that run dense setup together over one ``HashStore``.""" + + def __init__(self, specs, memory, capture, store): + self.ranks = [Rank(index, spec, memory) for index, spec in enumerate(specs)] + self.memory = memory + self.capture = capture + self.store = store + + def run(self): + for rank in self.ranks: + rank.thread = threading.Thread( + target=rank.run, + args=(self.store, len(self.ranks)), + name=f"dense-setup-rank{rank.index}", + daemon=True, + ) + for rank in self.ranks: + rank.thread.start() + for rank in self.ranks: + rank.thread.join(JOIN_TIMEOUT_S) + stuck = [rank.index for rank in self.ranks if rank.thread.is_alive()] + assert stuck == [], f"dense setup did not finish on ranks {stuck}" + return self + + def messages(self, rank, level): + records = self.capture.records + return [text for thread, kind, text in records if thread is rank.thread and kind == level] + + def assert_collectives(self, count): + """Every rank made ``count`` collectives of ``PhaseReport``s; returns the last gathered.""" + for rank in self.ranks: + assert len(rank.comm.sent) == count, rank.index + assert all(isinstance(report, PhaseReport) for report in rank.comm.sent) + return [rank.comm.sent[-1] for rank in self.ranks] + + def assert_abandoned(self, pattern, raised): + """Every rank raised ``DenseSetupError`` (else logged one error), and its model is intact.""" + for rank in self.ranks: + if raised: + assert rank.error is not None and rank.error[0] is DenseSetupError, rank.trace + message = rank.error[1] + else: + assert rank.error is None and rank.result is None, rank.trace + (message,) = self.messages(rank, "error") + assert re.search(pattern, message), message + rank.snapshot.assert_restored(rank.decoder) + assert not hasattr(rank.model, "_dense_dwdp_active") + assert not hasattr(rank.decoder, "_dense_dwdp_active") + + def assert_released(self): + """Each store and launcher was released once, by setup; the last step, ``empty_cache``, + came after the store release and the death of the slots.""" + gc.collect() + for rank in self.ranks: + released = [(store.calls["release"], store.finalizer_releases) for store in rank.stores] + assert released == [(1, 0)] * len(rank.stores), rank.index + assert [launcher.releases for launcher in rank.launchers] == [1] * len(rank.launchers) + assert rank.log[-1] == "empty_cache", rank.log + assert rank.cache_states[-1] == (True, []), rank.index + assert self.memory.live_slots() == [] + + +@pytest.fixture +def make_group(monkeypatch): + """Builds rank groups; routes setup's slots, plans, cache release and logs to the test.""" + torch_distributed = pytest.importorskip("torch.distributed") + memory = HostMemory() + capture = CaptureLogger() + groups = [] + allocate_slot_ring = dense_setup.allocate_slot_ring + build_copy_plans = dense_setup.build_copy_plans + + def current_rank(): + thread = threading.current_thread() + ranks = [rank for group in groups for rank in group.ranks if rank.thread is thread] + return ranks[0] if ranks else None + + def recording_allocate_slot_ring(spec, device): + ring = allocate_slot_ring(spec, device) + for slot in ring: + memory.add_slot(slot.storage) + return ring + + def recording_build_copy_plans(layout, *, rank, chunk_bases, slot_bases): + current_rank().corrupt_address = corruption_address(layout, slot_bases) + return build_copy_plans(layout, rank=rank, chunk_bases=chunk_bases, slot_bases=slot_bases) + + def recording_empty_cache(): + rank = current_rank() + if rank is not None: + rank.on_empty_cache() + + monkeypatch.setattr(dense_setup, "allocate_slot_ring", recording_allocate_slot_ring) + monkeypatch.setattr(dense_setup, "build_copy_plans", recording_build_copy_plans) + monkeypatch.setattr(dense_setup, "logger", capture) + monkeypatch.setattr(dense_repoint, "logger", capture) + monkeypatch.setattr(torch.cuda, "empty_cache", recording_empty_cache) + + def make(specs): + groups.append(Group(specs, memory, capture, torch_distributed.HashStore())) + return groups[-1] + + yield make + for group in groups: + for rank in group.ranks: + if isinstance(rank.result, DenseDWDP): + rank.result.release() + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("attached", [True, False], ids=["attached", "standalone"]) +def test_only_attached_dense_dwdp_prefetches_the_next_forward(make_group, attached): + """With expert DWDP (``moe_layers``) the runtime gathers the next forward's first layers at + the end of a forward, beside the expert prefetch; standalone the dense manager primes them.""" + group = make_group([RankSpec(attached=attached)] * 2).run() + group.assert_collectives(5) + for rank in group.ranks: + assert isinstance(rank.result, DenseDWDP), rank.trace + assert rank.result.layers == COVERED + assert rank.result.prefetch_next_pass is attached + + +@pytest.mark.parametrize("strict,debug", [(False, 1), (True, 0)], ids=["warn-debug", "strict"]) +def test_a_leaked_original_is_reported_after_the_commit_point(make_group, strict, debug): + group = make_group([RankSpec(strict=strict), RankSpec(strict=strict, debug=debug)]) + leaker = group.ranks[1] + # A holder the bindings do not know about keeps one original alive. + hidden = leaker.decoder.layers[1].linear_attn.o_proj.weight + group.run() + group.assert_collectives(5) + assert isinstance(group.ranks[0].result, DenseDWDP), group.ranks[0].trace + (warning,) = [m for m in group.messages(leaker, "warning") if "still alive after commit" in m] + assert "storage of layers.1.linear_attn._o_dense " in warning + # TRTLLM_DWDP_DENSE_DEBUG >= 1 also describes the holders (a scan of every live object). + assert ("; holders: " in warning) == (debug >= 1) + if not strict: + assert isinstance(leaker.result, DenseDWDP), leaker.trace + (line,) = [m for m in group.messages(leaker, "info") if "committed:" in m] + assert " leaks=1 " in line + return + # Only the leaking rank fails, after the commit point: its originals are gone, so its + # bindings are detached and its store released, and the model is flagged. + error_type, message = leaker.error + assert error_type is DenseRepointError and "still alive after commit" in message + assert leaker.result is None + moved = { + f"layers.{layer}.{binding.qualified_name}" + for (layer, _), segment in collect(ToyModel()).segments.items() + for anchor in segment.anchors + for binding in anchor.bindings + } + tensors = module_tensors(leaker.decoder) + assert all(tensors[path].numel() == 0 for path in moved) + (store,) = leaker.stores + assert (store.calls["release"], store.finalizer_releases) == (1, 0) + assert leaker.launchers[0].releases == 1 + assert leaker.model._dense_dwdp_active is True + expected = leaker.snapshot.values["layers.1.linear_attn.o_proj.weight"] + assert torch.equal(raw_bytes(hidden), expected) + + +#: ``(REQUIRED on every rank, change on one rank, what every rank raises or logs)``. +DISAGREEMENTS = { + "layouts": ( + False, + {"kinds": ("dense",) + ("kda",) * 5}, + r"abandoned at collective #1 .*the dense layouts differ across ranks \(digest .* on rank 2\)", + ), + "weights": ( + True, + {"perturb": "layers.4.block_sparse_moe.gate.weight"}, + r"#1 .*the dense weights differ across ranks: layers\.4 moe: ranks \[0, 1, 3\] vs rank 2$", + ), + "strict": ( + True, + {"strict": True}, + r"#1 .*TRTLLM_DWDP_DENSE_STRICT differs across ranks \(0 on ranks \[0, 1, 3\] vs 1 on", + ), + # Ranks that differ on TRTLLM_DWDP_DENSE itself raise even with REQUIRED off. + "switch": (False, {"requested": False}, r"TRTLLM_DWDP_DENSE differs across the DWDP group"), +} + + +@pytest.mark.parametrize("required,change,pattern", DISAGREEMENTS.values(), ids=DISAGREEMENTS) +def test_ranks_that_disagree_stop_at_the_first_collective(make_group, required, change, pattern): + specs = [RankSpec(required=required)] * 4 + specs[2] = replace(specs[2], **change) + group = make_group(specs).run() + group.assert_collectives(1) + group.assert_abandoned(pattern, raised=required or "requested" in change) + assert all(rank.stores == [] for rank in group.ranks) + + +@pytest.mark.parametrize( + "phase,num_ranks,failing,required,fault,report", + [ + (1, 4, 3, True, {"free_bytes": 1 << 20}, r"DenseSetupError: 0\.00 GB of device memory"), + (2, 2, 0, False, {"fail": "allocate"}, r"RuntimeError: injected: store allocate failed"), + (3, 4, 1, True, {"fail": "import"}, r"RuntimeError: injected: store import failed"), + (4, 2, 1, False, {"corrupt": 4}, r"Phase V: 1 gathered anchors differ from their"), + (5, 4, 2, True, {"corrupt": 5}, r"P5 self-check: 1 module tensors are not their slot"), + ], + ids=["p1-memory", "p2-allocate", "p3-import", "p4-phase-v", "p5-self-check"], +) +def test_a_failure_in_any_phase_abandons_everywhere( + make_group, phase, num_ranks, failing, required, fault, report +): + specs = [RankSpec(required=required)] * num_ranks + specs[failing] = replace(specs[failing], **fault) + group = make_group(specs).run() + + reports = group.assert_collectives(phase) + assert [report.ok for report in reports] == [rank != failing for rank in range(num_ranks)] + if phase >= 4: + # The failing rank names exactly the corrupted tensor; the others found nothing. + assert [report.data for report in reports] == [ + [LAST_MOE_ANCHOR] if rank == failing else [] for rank in range(num_ranks) + ] + group.assert_abandoned( + rf"abandoned at collective #{phase} .*: rank {failing}: {report}", required + ) + group.assert_released() + assert [len(rank.stores) for rank in group.ranks] == [int(phase >= 2)] * num_ranks + expected_launchers = [ + int(phase >= 4 or (phase == 3 and r != failing)) for r in range(num_ranks) + ] + assert [len(rank.launchers) for rank in group.ranks] == expected_launchers diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_store.py b/tests/unittest/_torch/modules/dwdp/test_dense_store.py new file mode 100644 index 000000000000..a7b6b8ec32eb --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_store.py @@ -0,0 +1,375 @@ +# 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. +"""CPU tests of the dense DWDP FABRIC row store over a recording fake of the CUDA VMM calls. + +The layouts come from toy decoder layers through ``layout.collect_dense_layout``; the big layers +have rows of about 1.25 MiB (``N = 4``), so ``STORE_CHUNK_MB`` splits them. The store maps its +chunks read-write and zero-filled, writes this rank's rows, exports one handle per chunk, imports +and maps every peer's handles read-only in rank order, and releases peer mappings, imported +handles, local mappings and local handles in that order, also after a failure at any step of +``allocate`` or ``import_peers``, which then does not run again. +""" + +import ctypes +import math +from collections import Counter +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.modules.dwdp import transport, vmm +from tensorrt_llm._torch.modules.dwdp.dense.config import DenseConfig +from tensorrt_llm._torch.modules.dwdp.dense.layout import collect_dense_layout, write_blob +from tensorrt_llm._torch.modules.dwdp.dense.store import FabricShardStore +from tensorrt_llm._torch.modules.dwdp.dense.types import ROW_ALIGN + +pytestmark = pytest.mark.cpu_only + +FABRIC = vmm.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC +PROT_READ = vmm.cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READ +PROT_READWRITE = vmm.cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE +GRANULARITY = 2 << 20 +GARBAGE = 0xEE # Fake device memory before the store zero-fills it. + +#: Four layers whose rows take about 1.25 MiB each with N = 4 (A 3 MiB + 5 * l bytes, D 2 MiB). +BIG_LAYERS = { + layer: ([(torch.uint8, ((3 << 20) + 5 * layer,))], [(torch.bfloat16, (1 << 20,))]) + for layer in range(1, 5) +} + + +class ToyLayer(nn.Module): + """A decoder layer whose ``attn`` and ``moe`` roots hold one random Parameter per spec.""" + + def __init__(self, spec, generator) -> None: + super().__init__() + for segment, tensors in zip(("attn", "moe"), spec): + root = nn.Module() + for index, (dtype, shape) in enumerate(tensors): + nbytes = math.prod(shape) * dtype.itemsize + data = torch.randint(0, 256, (nbytes,), generator=generator, dtype=torch.uint8) + weight = nn.Parameter(data.view(dtype).reshape(shape), requires_grad=False) + root.register_parameter(f"w{index}", weight) + setattr(self, segment, root) + + def dense_dwdp_roots(self, segment: str) -> list[tuple[str, nn.Module]]: + return [(segment, getattr(self, segment))] + + +def collect(num_ranks: int, chunk_mb: int): + """The layout of the ``BIG_LAYERS`` toy layers (seeded, so every call gives equal bytes).""" + generator = torch.Generator().manual_seed(1234) + layers = [(index, ToyLayer(BIG_LAYERS[index], generator)) for index in sorted(BIG_LAYERS)] + config = DenseConfig(enabled=True, min_bytes=64, store_chunk_mb=chunk_mb) + return collect_dense_layout( + layers, + config, + num_ranks=num_ranks, + device=torch.device("cpu"), + attr_allow=frozenset(), + exclude=frozenset(), + ) + + +def host_bytes(address: int, nbytes: int) -> torch.Tensor: + """A ``uint8`` tensor that aliases ``nbytes`` of host memory at ``address``.""" + return torch.frombuffer((ctypes.c_ubyte * nbytes).from_address(address), dtype=torch.uint8) + + +class FakeVmm: + """The VMM, export/import and synchronize calls of the FABRIC store, over host memory. + + Every call is recorded as ``(op, value)``; ``cuMemSetAccess`` is faked at the driver binding, + so the ``access_rw`` / ``access_ro`` records are the access flags the store asks for. + Reserved address ranges are host buffers filled with garbage, so the store's views alias them. + Like the driver, the fake knows the size of every handle (created, or imported from a payload + ``export`` or ``peer_payloads`` registered) and checks that handles and reservations are whole + granules, a mapping covers one whole handle and exactly its reservation, and access, unmap + and free name a whole reservation. ``fail=(op, n)`` makes the ``n``-th call of ``op`` raise. + """ + + def __init__(self, monkeypatch, *, fail=None) -> None: + self.calls: list[tuple[str, object]] = [] + self.fail = fail + self.counts: Counter = Counter() + self.memory: dict[int, torch.Tensor] = {} + self.owner: dict[int, int] = {} # va -> the handle mapped (or about to be) there + self.imported: set[int] = set() + self.handle_sizes: dict[int, int] = {} + self.payload_sizes: dict[bytes, int] = {} + self.va_sizes: dict[int, int] = {} + self._next_handle = 1000 + self._last_handle = None + patches = { + (vmm, "peer_handle_type"): lambda: FABRIC, + (vmm, "get_allocation_granularity"): lambda device_id: GRANULARITY, + (vmm, "create_fabric_handle"): self.create, + (vmm, "reserve_va"): self.reserve, + (vmm, "map_handle"): self.map, + (vmm.cuda, "cuMemSetAccess"): self.set_access, + (vmm, "tensor_from_ptr"): self.view, + (vmm, "unmap_va"): self.unmap, + (vmm, "free_va"): self.free, + (vmm, "release_handle"): lambda handle: self._record("release", handle), + (transport, "_export_handle"): self.export, + (transport, "_import_handle"): self.import_handle, + (torch.cuda, "synchronize"): lambda device=None: self._record("sync", device), + (torch.cuda, "current_stream"): self.current_stream, + } + for (owner, name), value in patches.items(): + monkeypatch.setattr(owner, name, value) + + def _record(self, op: str, value: object) -> None: + self.counts[op] += 1 + if self.fail == (op, self.counts[op]): + raise RuntimeError(f"injected failure of {op} #{self.counts[op]}") + self.calls.append((op, value)) + + def _new_handle(self, size: int) -> int: + self._next_handle += 1 + self.handle_sizes[self._next_handle] = size + self._last_handle = self._next_handle + return self._next_handle + + def _check_reservation(self, what: str, va: int, size: int) -> None: + assert self.va_sizes.get(va) == size, f"{what} of {size} bytes at {va:#x}" + + def create(self, size, device_id): + assert size > 0 and size % GRANULARITY == 0, f"cuMemCreate of {size} bytes" + self._record("create", (self._next_handle + 1, size)) + return self._new_handle(size) + + def import_handle(self, payload, handle_type): + assert handle_type == FABRIC and payload in self.payload_sizes, payload + self._record("import", (self._next_handle + 1, payload)) + self.imported.add(self._next_handle + 1) + return self._new_handle(self.payload_sizes[payload]) + + def export(self, handle, handle_type): + assert handle_type == FABRIC + self._record("export", handle) + payload = f"fabric-handle-{handle}".encode() + self.payload_sizes[payload] = self.handle_sizes[handle] + return payload + + def reserve(self, size, alignment): + assert alignment == GRANULARITY and size > 0 and size % GRANULARITY == 0, size + self._record("reserve", size) + buffer = torch.full((size + 2 * ROW_ALIGN,), GARBAGE, dtype=torch.uint8) + va = buffer.data_ptr() + (-buffer.data_ptr()) % ROW_ALIGN + self.memory[va] = buffer + self.va_sizes[va] = size + self.owner[va] = self._last_handle + self.calls[-1] = ("reserve", (va, size)) + return va + + def map(self, va, size, handle, offset=0): + assert self.owner[va] == handle and offset == 0 + self._check_reservation("cuMemMap", va, size) + assert size == self.handle_sizes[handle], f"cuMemMap of {size} bytes of handle {handle}" + self._record("map", (va, handle)) + + def set_access(self, va, size, descriptors, count): + """``cuMemSetAccess``: records ``access_rw`` or ``access_ro`` by the requested flags.""" + assert count == len(descriptors) == 1 and descriptors[0].location.id == 0 + self._check_reservation("cuMemSetAccess", va, size) + flags = int(descriptors[0].flags) + op = {int(PROT_READWRITE): "access_rw", int(PROT_READ): "access_ro"}.get(flags) + assert op is not None, f"unexpected access flags {flags}" + self._record(op, va) + return (vmm.cuda.CUresult.CUDA_SUCCESS,) + + def view(self, ptr, shape, dtype, device_id): + self._record("view", ptr) + return host_bytes(ptr, math.prod(shape) * dtype.itemsize).view(dtype).reshape(shape) + + def unmap(self, va, size): + self._check_reservation("cuMemUnmap", va, size) + self._record("unmap", va) + + def free(self, va, size): + self._check_reservation("cuMemAddressFree", va, size) + self._record("free", va) + del self.memory[va] + del self.va_sizes[va] + + def current_stream(self, device=None): + return SimpleNamespace(synchronize=lambda: self._record("stream_sync", device)) + + def ops(self, start: int = 0) -> list[str]: + return [op for op, _ in self.calls[start:]] + + def values(self, op: str, start: int = 0, stop: int | None = None) -> list: + return [value for name, value in self.calls[start:stop] if name == op] + + def is_peer(self, value) -> bool: + """Whether a recorded handle or address belongs to an imported (peer) handle.""" + return self.owner.get(value, value) in self.imported + + +def peer_payloads(fake: FakeVmm, store: FabricShardStore, layout, rank: int): + """Collective #2 as this rank sees it: its own export plus a distinct export per peer.""" + own = store.export() + payloads = [] + for peer in range(layout.num_ranks): + blobs = [f"peer-{peer}-chunk-{chunk.index}".encode() for chunk in layout.chunks] + for blob, own_blob in zip(blobs, own): + fake.payload_sizes[blob] = fake.payload_sizes[own_blob] + payloads.append(own if peer == rank else blobs) + return payloads + + +def check_release(fake: FakeVmm, start: int) -> None: + """The calls since ``start`` undo every handle and mapping created before it, once each. + + Order: a device synchronize first; then the peer mappings (unmap, free) and imported handles, + then the local mappings and handles; a mapping is unmapped before its VA is freed and both + happen before its handle is released. + """ + created = [handle for handle, _ in fake.values("create", 0, start)] + imported = [handle for handle, _ in fake.values("import", 0, start)] + reserved = [va for va, _ in fake.values("reserve", 0, start)] + mapped = [va for va, _ in fake.values("map", 0, start)] + tail = fake.calls[start:] + if not (created or imported): + assert tail == [] + return + assert tail[0][0] == "sync" + released = fake.values("release", start) + assert sorted(released) == sorted(created + imported) and len(set(released)) == len(released) + assert sorted(fake.values("free", start)) == sorted(reserved) + assert sorted(fake.values("unmap", start)) == sorted(mapped) + position = {(op, value): index for index, (op, value) in enumerate(tail)} + for va in reserved: + assert position[("free", va)] < position[("release", fake.owner[va])] + if va in mapped: + assert position[("unmap", va)] < position[("free", va)] + peer_steps = [i for i, (_, value) in enumerate(tail[1:]) if fake.is_peer(value)] + local_steps = [i for i, (_, value) in enumerate(tail[1:]) if not fake.is_peer(value)] + if peer_steps and local_steps: + assert max(peer_steps) < min(local_steps), "peer resources must go before local ones" + peer_frees = [position[("free", va)] for va in reserved if fake.is_peer(va)] + peer_releases = [position[("release", handle)] for handle in imported] + if peer_frees and peer_releases: + assert max(peer_frees) < min(peer_releases), "peer VA ranges go before imported handles" + + +@pytest.mark.parametrize("chunk_mb", [0, 1], ids=["one-chunk", "chunked"]) +def test_fabric_store_maps_writes_exports_and_imports(monkeypatch, chunk_mb): + """One handle per chunk mapped read-write and zeroed, this rank's rows, peers read-only.""" + num_ranks, rank = 4, 2 + layout = collect(num_ranks, chunk_mb) + num_chunks = len(layout.chunks) + fake = FakeVmm(monkeypatch) + store = FabricShardStore(layout, rank, torch.device("cuda", 0)) + + store.allocate() + per_chunk = ["create", "reserve", "map", "access_rw", "view"] + assert fake.ops() == per_chunk * num_chunks + ["stream_sync"] + assert [size for _, size in fake.values("create")] == [c.handle_bytes for c in layout.chunks] + local_vas = [va for va, _ in fake.values("reserve")] + for chunk, va in zip(layout.chunks, local_vas): + assert not host_bytes(va, chunk.handle_bytes).any(), f"chunk {chunk.index} not zeroed" + + store.write_local_rows() + for chunk, va in zip(layout.chunks, local_vas): + written = torch.zeros(chunk.handle_bytes, dtype=torch.uint8) + for layer in chunk.layers: + for segment in layout.layer_segments(layer): + width = segment.row_bytes + blob = torch.empty(num_ranks * width, dtype=torch.uint8) + write_blob(segment, blob, num_ranks=num_ranks) + row = blob[rank * width : (rank + 1) * width] + written[segment.chunk_offset : segment.chunk_offset + width] = row + assert torch.equal(host_bytes(va, chunk.handle_bytes), written), f"chunk {chunk.index}" + + mark = len(fake.calls) + payloads = peer_payloads(fake, store, layout, rank) + assert payloads[rank] == [f"fabric-handle-{h}".encode() for h, _ in fake.values("create")] + assert store.export() == payloads[rank] # Exported once, returned again. + assert fake.ops(mark) == ["export"] * num_chunks + + mark = len(fake.calls) + store.import_peers(payloads) + imports = [payload for _, payload in fake.values("import", mark)] + assert imports == [blob for peer in (0, 1, 3) for blob in payloads[peer]] + assert fake.ops(mark) == ["import", "reserve", "map", "access_ro"] * (3 * num_chunks) + peer_sizes = [size for _, size in fake.values("reserve", mark)] + assert peer_sizes == [chunk.handle_bytes for chunk in layout.chunks] * 3 + peer_vas = iter(va for va, _ in fake.values("reserve", mark)) + expected_bases = tuple( + tuple(local_vas) if peer == rank else tuple(next(peer_vas) for _ in layout.chunks) + for peer in range(num_ranks) + ) + assert store.chunk_bases() == expected_bases + + mark = len(fake.calls) + store.release() + check_release(fake, mark) + mark = len(fake.calls) + store.release() + assert fake.calls[mark:] == [] and fake.memory == {} + + +#: ``(op, n)`` failure points of ``allocate`` and ``import_peers`` on rank 1 of 3 with four chunks: +#: allocate makes four creates, reserves, maps, read-write accesses and views; import_peers makes +#: eight imports, reserves (calls 5-12 overall), maps and read-only accesses. +FAILURE_POINTS = [ + ("reserve", 2), + ("map", 3), + ("view", 4), + ("import", 6), + ("reserve", 11), + ("access_ro", 7), +] + + +@pytest.mark.parametrize("fail", FAILURE_POINTS, ids=[f"{op}-{n}" for op, n in FAILURE_POINTS]) +def test_fabric_release_after_a_partial_allocate_or_import(monkeypatch, fail): + """A step that fails part way does not run again and leaves what it created to ``release``.""" + num_ranks, rank = 3, 1 + layout = collect(num_ranks, chunk_mb=1) + assert len(layout.chunks) == 4 + fake = FakeVmm(monkeypatch, fail=fail) + store = FabricShardStore(layout, rank, torch.device("cuda", 0)) + failed = "allocate" + with pytest.raises(RuntimeError, match="injected failure"): + store.allocate() + failed = "write_local_rows" + store.write_local_rows() + failed = "export" + payloads = peer_payloads(fake, store, layout, rank) + failed = "import_peers" + store.import_peers(payloads) + with pytest.raises(RuntimeError): + store.chunk_bases() + if failed == "allocate": + with pytest.raises(RuntimeError, match="allocate already ran"): + store.allocate() + for step in (store.write_local_rows, store.export): + with pytest.raises(RuntimeError, match="before a successful allocate"): + step() + elif failed == "import_peers": + with pytest.raises(RuntimeError, match="import_peers already ran"): + store.import_peers(payloads) + mark = len(fake.calls) + store.release() + check_release(fake, mark) + assert fake.memory == {} + mark = len(fake.calls) + store.release() + assert fake.calls[mark:] == [] diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py index 33cbfc0e46d5..56b28dc6c013 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py @@ -15,12 +15,35 @@ """PyExecutor construction ownership tests for DWDP.""" from contextlib import contextmanager +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest +from tensorrt_llm._torch.cute_dsl_utils import sm_reserve +from tensorrt_llm._torch.modules.dwdp import fabric_gather +from tensorrt_llm._torch.modules.dwdp.dense.config import DENSE_COPY_ENV, DENSE_ENV +from tensorrt_llm._torch.modules.dwdp.dense.manager import ( + DenseDwdpManager, + get_global_dense_dwdp_manager, + set_global_dense_dwdp_manager, +) +from tensorrt_llm._torch.modules.dwdp.dense.types import DenseMode from tensorrt_llm._torch.pyexecutor import py_executor_creator from tensorrt_llm._torch.pyexecutor.dwdp import get_global_dwdp_manager, set_global_dwdp_manager +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor +from tensorrt_llm.llmapi.llm_args import DwdpConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.quantization import QuantAlgo + +__extra_import_path__ = ["../../executor"] +from test_py_executor_creator_mla_cache_reuse_sync import ( # noqa: E402 + _DummyCalibrator, + _DummyModelEngine, + _make_llm_args, +) + +pytestmark = pytest.mark.cpu_only @contextmanager @@ -74,3 +97,221 @@ def fail_without_registering(**_kwargs): existing.__exit__.assert_not_called() assert get_global_dwdp_manager() is existing + + +@contextmanager +def _restore_global_dense_dwdp_manager(): + """Keep the process-global dense manager from leaking into later tests.""" + previous = get_global_dense_dwdp_manager() + try: + yield + finally: + set_global_dense_dwdp_manager(previous) + + +def test_create_py_executor_rolls_back_a_new_standalone_dense_manager_on_error(): + """A standalone dense manager registered by the failed construction is exited.""" + with _restore_global_dwdp_manager(), _restore_global_dense_dwdp_manager(): + set_global_dwdp_manager(None) + set_global_dense_dwdp_manager(None) + manager = MagicMock() + manager.__exit__.side_effect = lambda *_args: set_global_dense_dwdp_manager(None) + + def fail_after_registration(**_kwargs): + set_global_dense_dwdp_manager(manager) + raise RuntimeError("construction failed") + + with patch.object( + py_executor_creator, + "_create_py_executor_impl", + side_effect=fail_after_registration, + ): + with pytest.raises(RuntimeError, match="construction failed"): + py_executor_creator.create_py_executor(llm_args=MagicMock()) + + manager.__exit__.assert_called_once_with(None, None, None) + assert get_global_dense_dwdp_manager() is None + + +def test_create_py_executor_leaves_a_pre_existing_dense_manager_alone(): + """The wrapper only exits a dense manager it created, not one it inherited.""" + with _restore_global_dense_dwdp_manager(): + existing = MagicMock() + set_global_dense_dwdp_manager(existing) + + def fail_without_registering(**_kwargs): + raise RuntimeError("construction failed") + + with patch.object( + py_executor_creator, + "_create_py_executor_impl", + side_effect=fail_without_registering, + ): + with pytest.raises(RuntimeError, match="construction failed"): + py_executor_creator.create_py_executor(llm_args=MagicMock()) + + existing.__exit__.assert_not_called() + assert get_global_dense_dwdp_manager() is existing + + +# ---------------------------------------------------------------------- +# The executor that serves gets the DWDP manager +# ---------------------------------------------------------------------- +# ``_create_py_executor_impl`` runs to its end over the stand-ins of +# ``test_py_executor_creator_mla_cache_reuse_sync.py``, with stand-in DWDP managers and a +# stand-in dense DWDP gate, which the creator runs only when ``TRTLLM_DWDP_DENSE`` asks for it. + +#: Creator dependencies replaced by no-ops. +_NOOPS = ("instantiate_sampler", "get_spec_resource_manager", "get_spec_drafter") +_NOOPS += ("_adjust_torch_mem_fraction", "log_memory_usage") + + +def _standalone_dense_manager(mode: DenseMode, *, mapping: Mapping) -> DenseDwdpManager: + """A real standalone manager, of an attention-DP instance whatever the creator's stand-in + mapping.""" + assert mode is DenseMode.STANDALONE + instance = Mapping(world_size=2, rank=0, tp_size=2, enable_attention_dp=True) + return DenseDwdpManager(mode, mapping=instance) + + +def _create( + monkeypatch, mode: str, kv_cache: str, *, real_dense_manager: bool = False +) -> tuple[list, object]: + """Run ``create_py_executor`` with DWDP in ``mode``, start a forward of the executor it returns. + + Args: + mode: ``"expert"`` (expert DWDP alone), ``"attached"`` (dense DWDP attached to it) or + ``"standalone"`` (dense DWDP alone). + kv_cache: ``"estimated"``, ``"not_estimated"`` or ``"none"`` (a model without one). + real_dense_manager: Standalone, the creator builds a real ``DenseDwdpManager``, whose + setup only logs; the model engine logs the SM reserve it is built with. + + Returns: + The call log and the one DWDP manager created. + """ + log, managers = [], [] + + def manager(*_args, **_kwargs): + """Stands in for ``DwdpManager`` and ``DenseDwdpManager``.""" + managers.append( + SimpleNamespace( + __enter__=lambda: None, + setup=lambda _model: log.append("setup"), + start_forward=lambda: log.append("start_forward"), + prefetch_first_layers=lambda: log.append("prefetch_first_layers"), + ) + ) + return managers[-1] + + def kv_cache_creator(**_kwargs): + return SimpleNamespace( + try_prepare_estimation=lambda: kv_cache == "estimated", + build_managers=lambda _resources, estimating: log.append( + ("build_managers", estimating) + ), + configure_kv_cache_capacity=lambda _executor: None, + teardown_managers=lambda _resources: None, + ) + + def create_py_executor_instance(**kwargs): + given = kwargs["dwdp_manager"], kwargs["dense_dwdp_manager"] + log.append(("executor", *given)) + executor = object.__new__(PyExecutor) # with what the creator and _prime_dwdp read + executor._metrics, executor.start_worker = {}, lambda: None + executor.peft_cache_config = kwargs["peft_cache_config"] + executor.dwdp_manager, executor.dense_dwdp_manager = given + return executor + + def model_engine(attn_runtime_features, **_kwargs): + if real_dense_manager: + log.append(("model built with an SM reserve of", sm_reserve())) + engine = _DummyModelEngine( + attn_runtime_features=attn_runtime_features, kv_cache_quant_algo=QuantAlgo.NO_QUANT + ) + engine.model.model_config.is_generation = kv_cache != "none" + return engine + + dense_mode = {"attached": DenseMode.ATTACHED, "standalone": DenseMode.STANDALONE}.get(mode) + monkeypatch.setenv(DENSE_ENV, "0" if dense_mode is None else "1") + fakes = dict.fromkeys(_NOOPS, lambda *_args, **_kwargs: None) + fakes.update( + validate_dense_gate=lambda *_args: log.append("gate") or dense_mode, + DwdpManager=manager, + DenseDwdpManager=_standalone_dense_manager if real_dense_manager else manager, + KvCacheCreator=kv_cache_creator, + create_py_executor_instance=create_py_executor_instance, + ModelEngine=model_engine, + _load_config_and_create_checkpoint_loader=lambda llm_args, _dir: (llm_args, None), + _get_mapping=lambda _: Mapping(), + dwdp_global_rank=lambda: 0, + get_calibrator=_DummyCalibrator, + is_mla=lambda _: True, + is_hybrid_linear=lambda _: False, + get_sm_version=lambda: 100, + ) + for name, fake in fakes.items(): + monkeypatch.setattr(py_executor_creator, name, fake) + monkeypatch.setattr(py_executor_creator.Distributed, "get", staticmethod(lambda _: None)) + cuda = dict( + mem_get_info=lambda: (2 << 30, 4 << 30), + empty_cache=lambda: None, + reset_peak_memory_stats=lambda: None, + memory_stats=lambda: {"allocated_bytes.all.current": 0}, + Stream=lambda: SimpleNamespace(cuda_stream=123), + ) + for name, fake in cuda.items(): + monkeypatch.setattr(py_executor_creator.torch.cuda, name, fake) + + llm_args = _make_llm_args() + if mode != "standalone": + llm_args.dwdp_config = DwdpConfig( + dwdp_size=2, num_experts_per_worker=4, num_prefetch_experts=4 + ) + if real_dense_manager: + monkeypatch.setattr(DenseDwdpManager, "setup", lambda _self, _model: log.append("setup")) + executor = py_executor_creator.create_py_executor(llm_args=llm_args, checkpoint_dir=None) + executor._prime_dwdp() + (manager,) = [executor.dense_dwdp_manager] if real_dense_manager else managers + return log, manager + + +@pytest.mark.parametrize("kv_cache", ["estimated", "not_estimated", "none"]) +@pytest.mark.parametrize("mode", ["expert", "attached", "standalone"]) +def test_only_the_serving_executor_gets_the_manager(monkeypatch, mode, kv_cache): + """The manager is set up once; only the executor that serves gets it and starts its forwards. + + With the KV-cache estimation, the estimation executor gets none, so its shutdown exits none. + Without it the first executor is the only one, and dense DWDP is set up before + ``build_managers`` sizes the final KV cache, which then gets the memory dense DWDP frees; + expert DWDP alone is set up after it. + """ + log, manager = _create(monkeypatch, mode, kv_cache) + + final = ("executor", None, manager) if mode == "standalone" else ("executor", manager, None) + final_kv_cache = ("build_managers", False) + setup_first = mode != "expert" + expected = { + "estimated": [("build_managers", True), "setup", ("executor", None, None), final_kv_cache], + "not_estimated": ["setup", final_kv_cache] if setup_first else [final_kv_cache, "setup"], + "none": ["setup"], + }[kv_cache] + start = "prefetch_first_layers" if mode == "standalone" else "start_forward" + assert log == ["gate"] * (mode != "expert") + expected + [final, start] + + +def test_a_standalone_gather_reserves_its_sms_before_the_model_is_built(monkeypatch): + """Kernels that size their persistent grids when the model is built (the KDA prefill) leave + the SM gather's resident SMs out; exiting the manager drops the reserve.""" + monkeypatch.setenv(DENSE_COPY_ENV, "gather") + monkeypatch.delenv("TRTLLM_CUTEDSL_SM_RESERVE", raising=False) + monkeypatch.setenv("TRTLLM_DWDP_GATHER_BLOCKS", "30") + monkeypatch.setattr(fabric_gather, "THREADS", 1024) + with _restore_global_dense_dwdp_manager(): + log, manager = _create(monkeypatch, "standalone", "not_estimated", real_dense_manager=True) + try: + assert ("model built with an SM reserve of", 30) in log + assert log.index("setup") > log.index(("model built with an SM reserve of", 30)) + assert get_global_dense_dwdp_manager() is manager and sm_reserve() == 30 + finally: + manager.__exit__(None, None, None) + assert sm_reserve() == 0 diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py index 31217886e537..4d28bd596e11 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py @@ -16,22 +16,39 @@ DwdpManager is a thin lifecycle layer that: - creates a DWDP MPI sub-communicator +- resolves the expert partition each MoE backend is sized from - owns the Single Source of Truth for registered MoE layer indices - forwards runtime calls to the DWDPWeightManager produced by setup_dwdp +- owns the attached DenseDwdpManager: sets dense DWDP up before setup_dwdp These tests mock out COMM_WORLD / dwdp_global_rank / setup_dwdp so the full lifecycle can be exercised on a single CPU without MPI / CUDA / MNNVL. """ +import os import unittest +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest +import torch +import torch.nn as nn from tensorrt_llm._torch.cute_dsl_utils import sm_reserve from tensorrt_llm._torch.distributed import MPIDist +from tensorrt_llm._torch.modules.dwdp.dense.config import DENSE_ENV +from tensorrt_llm._torch.modules.dwdp.dense.manager import ( + H0_HOOK, + DenseDwdpManager, + get_global_dense_dwdp_manager, + set_global_dense_dwdp_manager, +) +from tensorrt_llm._torch.modules.dwdp.dense.types import DenseMode, DenseSetupError +from tensorrt_llm._torch.modules.dwdp.weight_buffer import PrefetchRingMemoryError +from tensorrt_llm._torch.moe.fused_moe.interface import MoE from tensorrt_llm._torch.pyexecutor.dwdp import ( DwdpManager, + _AttachedTransfer, get_global_dwdp_manager, set_global_dwdp_manager, ) @@ -40,11 +57,13 @@ pytestmark = pytest.mark.cpu_only +_FAC = "tensorrt_llm._torch.pyexecutor.dwdp" + -def _make_config(dwdp_size: int = 2) -> DwdpConfig: +def _make_config(dwdp_size: int = 2, num_groups: int = 1) -> DwdpConfig: return DwdpConfig( dwdp_size=dwdp_size, - num_groups=1, + num_groups=num_groups, num_experts_per_worker=4, num_prefetch_experts=4, ) @@ -74,16 +93,42 @@ def _configure_mock_comm_world(mock_comm_world, world_size: int = 1024): return sub_comm +def _start(test: unittest.TestCase, patcher): + """Start ``patcher`` for the rest of ``test`` and return its mock.""" + mock = patcher.start() + test.addCleanup(patcher.stop) + return mock + + +def _set_dense_env(test: unittest.TestCase, value: str | None) -> None: + """Set ``TRTLLM_DWDP_DENSE`` (None: unset) for the rest of ``test``.""" + _start(test, patch.dict(os.environ)) + os.environ.pop(DENSE_ENV, None) + if value is not None: + os.environ[DENSE_ENV] = value + + +def _clear_global_managers() -> None: + """Unregister both managers: a leftover one fails the next ``__init__``.""" + set_global_dwdp_manager(None) + set_global_dense_dwdp_manager(None) + + @patch("tensorrt_llm._torch.pyexecutor.dwdp.setup_dwdp") @patch("tensorrt_llm._torch.pyexecutor.dwdp.dwdp_global_rank", return_value=0) @patch("tensorrt_llm._torch.pyexecutor.dwdp.COMM_WORLD") class TestDwdpManagerLifecycle(unittest.TestCase): def setUp(self): # Ensure no leftover global singleton from a prior test. - set_global_dwdp_manager(None) + _clear_global_managers() + # Dense DWDP off: a model without dense roots, TRTLLM_DWDP_DENSE unset. + _set_dense_env(self, None) + _start(self, patch(f"{_FAC}.dense_roots_available", return_value=False)) + self.setup_dense = _start(self, patch(f"{_FAC}.setup_dense_dwdp")) def tearDown(self): - set_global_dwdp_manager(None) + _clear_global_managers() + self.setup_dense.assert_not_called() # ------------------------------------------------------------------ # Construction @@ -119,11 +164,8 @@ def test_init_stores_fields(self, mock_comm_world, _rank, _setup): self.assertEqual(mgr._registered_layers, []) self.assertIs(mgr.dwdp_group, sub_comm) self.assertIsNone(mgr._weight_manager) - # start_expert_id / end_expert_id derive from - # num_prefetch_experts (stride) and num_experts_per_worker (size). - # rank=0: start=0, end=4. - self.assertEqual(mgr.start_expert_id, 0) - self.assertEqual(mgr.end_expert_id, 4) + # rank=0 of a uniform 4/4 partition of 8 experts: [0, 4). + self.assertEqual(mgr.local_expert_range(8), (0, 4)) def test_init_expert_range_uniform(self, mock_comm_world, _rank, _setup): """Uniform partition (size == stride): each rank's range is a clean slice.""" @@ -134,8 +176,7 @@ def test_init_expert_range_uniform(self, mock_comm_world, _rank, _setup): dist=MagicMock(spec=MPIDist), mapping=_make_mapping(dwdp_size=2, dwdp_rank=1), ) - self.assertEqual(mgr.start_expert_id, 4) - self.assertEqual(mgr.end_expert_id, 8) + self.assertEqual(mgr.local_expert_range(8), (4, 8)) def test_init_expert_range_redundancy(self, mock_comm_world, _rank, _setup): """Redundancy partition (size > stride): adjacent ranges overlap.""" @@ -146,18 +187,60 @@ def test_init_expert_range_redundancy(self, mock_comm_world, _rank, _setup): num_experts_per_worker=6, # storage size num_prefetch_experts=4, # stride < size -> 2-expert overlap ) - # rank=0: [0, 6). rank=1 would be [4, 10) -> 2-expert overlap [4, 6). + # rank=0: [0, 6). rank=1: [4, 10) -> 2-expert overlap [4, 6). mgr = DwdpManager( config=redundant_config, dist=MagicMock(spec=MPIDist), mapping=_make_mapping(dwdp_size=2, dwdp_rank=0), ) - self.assertEqual(mgr.start_expert_id, 0) - self.assertEqual(mgr.end_expert_id, 6) + self.assertEqual(mgr.expert_partition(10).peer_ranges, [(0, 6), (4, 10)]) # Stored config knobs match what's plumbed into setup_dwdp. self.assertEqual(mgr.num_experts_per_worker, 6) self.assertEqual(mgr.num_prefetch_experts, 4) + # ------------------------------------------------------------------ + # Expert partition: DwdpConfig -> the range each MoE backend is sized from + # ------------------------------------------------------------------ + + def _k3_manager(self, mock_comm_world, size: int, stride: int, dwdp_rank: int) -> DwdpManager: + """DWDP rank ``dwdp_rank`` of 12 with the given partition fields.""" + _configure_mock_comm_world(mock_comm_world) + config = DwdpConfig(dwdp_size=12, num_experts_per_worker=size, num_prefetch_experts=stride) + mapping = _make_mapping(dwdp_size=12, dwdp_rank=dwdp_rank) + return DwdpManager(config=config, dist=MagicMock(spec=MPIDist), mapping=mapping) + + def test_config_refuses_a_half_set_or_negative_partition(self, *_mocks): + for size, stride in ((0, 74), (82, 0), (-1, -1), (-1, 4)): + with self.subTest(size=size, stride=stride), self.assertRaises(ValueError): + DwdpConfig(dwdp_size=12, num_experts_per_worker=size, num_prefetch_experts=stride) + + def test_balanced_partition_is_resolved_per_expert_count(self, mock_comm_world, _rank, _setup): + mgr = self._k3_manager(mock_comm_world, 0, 0, dwdp_rank=8) + self.assertEqual(mgr.local_expert_range(896), (600, 674)) # 8 x 75 + 4 x 74 + self.assertIs(mgr.expert_partition(896), mgr.expert_partition(896)) + self.assertEqual(mgr.local_expert_range(72), (48, 54)) + + def test_moe_backend_is_sized_from_the_rank_range(self, mock_comm_world, _rank, _setup): + """``MoE._init_dwdp_expert_layout`` sizes a backend with exactly its rank's rows.""" + # balanced, capped (the last rank stores 71), redundant + cases = (((0, 0), 8, (600, 674)), ((75, 75), 11, (825, 896)), ((82, 74), 11, (814, 896))) + for (size, stride), rank, (start, end) in cases: + with self.subTest(size=size, stride=stride): + set_global_dwdp_manager(self._k3_manager(mock_comm_world, size, stride, rank)) + moe = SimpleNamespace(num_experts=896, layer_load_balancer=None) + MoE._init_dwdp_expert_layout(moe) + set_global_dwdp_manager(None) + + self.assertEqual( + (moe.num_slots, moe.expert_size_per_partition, moe.slot_start, moe.slot_end), + (896, end - start, start, end), + ) + self.assertEqual(moe.initial_local_expert_ids, list(range(start, end))) + self.assertEqual(moe.initial_global_assignments, list(range(896))) + set_global_dwdp_manager(self._k3_manager(mock_comm_world, 82, 74, dwdp_rank=0)) + with self.assertRaises(ValueError): # 82/74 does not fit 900 experts + MoE._init_dwdp_expert_layout(SimpleNamespace(num_experts=900, layer_load_balancer=None)) + # ------------------------------------------------------------------ # num_groups topology validation (was a dead schema field through 1fbc0d49) # ------------------------------------------------------------------ @@ -191,10 +274,39 @@ def test_enter_registers_global(self, mock_comm_world, _rank, _setup): mapping=_make_mapping(), ) self.assertIsNone(get_global_dwdp_manager()) + self.assertIsNone(get_global_dense_dwdp_manager()) mgr.__enter__() self.assertIs(get_global_dwdp_manager(), mgr) + self.assertIs(get_global_dense_dwdp_manager(), mgr.dense_manager) + self.assertIs(mgr.dense_manager.mode, DenseMode.ATTACHED) mgr.__exit__(None, None, None) self.assertIsNone(get_global_dwdp_manager()) + self.assertIsNone(get_global_dense_dwdp_manager()) + + def test_enter_frees_the_group_when_the_dense_manager_cannot_register( + self, mock_comm_world, _rank, _setup + ): + """Neither stays registered; a failing Free is logged and the registration error raised.""" + sub_comm = _configure_mock_comm_world(mock_comm_world) + sub_comm.Free.side_effect = RuntimeError("free failed") + mgr = DwdpManager( + config=_make_config(), dist=MagicMock(spec=MPIDist), mapping=_make_mapping() + ) + stray = DenseDwdpManager(DenseMode.ATTACHED).__enter__() + + with patch(f"{_FAC}.logger.error") as log_error: + with self.assertRaisesRegex( + RuntimeError, "^DenseDwdpManager already registered globally$" + ): + mgr.__enter__() + + sub_comm.Free.assert_called_once_with() + self.assertIn("free failed", log_error.call_args.args[0]) + self.assertIsNone(mgr.dwdp_group) + self.assertIsNone(get_global_dwdp_manager()) + self.assertIs(get_global_dense_dwdp_manager(), stray) + with self.assertRaisesRegex(RuntimeError, r"setup\(\) called after cleanup\(\)"): + mgr.setup(MagicMock()) def test_enter_reserves_the_gather_sms(self, mock_comm_world, _rank, _setup): """The gather's resident SMs reach the persistent grids through set_sm_reserve, from @@ -246,6 +358,7 @@ def test_exit_clears_global_even_when_cleanup_fails(self, mock_comm_world, _rank mgr.__exit__(None, None, None) self.assertIsNone(get_global_dwdp_manager()) + self.assertIsNone(get_global_dense_dwdp_manager()) def test_cleanup_frees_comm_when_weight_release_fails(self, mock_comm_world, _rank, _setup): sub_comm = _configure_mock_comm_world(mock_comm_world) @@ -354,6 +467,9 @@ def test_setup_forwards_layer_indices(self, mock_comm_world, _rank, mock_setup): self.assertIs(result, fake_weight_manager) self.assertIs(mgr._weight_manager, fake_weight_manager) self.assertEqual(kwargs["repeated_layer_indices"], []) + # Dense DWDP is off: the weight manager makes its own copy stream, nothing is reserved. + self.assertIsNone(kwargs["copy_stream"]) + self.assertEqual(kwargs["reserved_storages"], frozenset()) def test_setup_forwards_repeated_dwdp_layers(self, mock_comm_world, _rank, _setup): """Repeated layers reach setup_dwdp; layers without a DWDP MoE are dropped.""" @@ -443,5 +559,257 @@ def test_wait_and_bind_forwards(self, mock_comm_world, _rank, mock_setup): fake_wm.wait_and_bind.assert_called_once_with(backend, 5) +class _FakeWeightManager: + """The DWDPWeightManager calls DwdpManager and its dense manager make, logged.""" + + def __init__(self, log: list[tuple]) -> None: + self.log = log + # What DwdpManager.setup compares with the transfer dense DWDP ran on. + self.copy_stream, self.use_gather, self.use_batched_prefetch = None, False, False + + def start_forward(self) -> None: + self.log.append(("start_forward",)) + + +class _FakeDense: + """The DenseDWDP calls the attached dense manager makes, logged; ``slots`` are its rings.""" + + def __init__(self, log: list[tuple], layers: list[int], slots: dict | None = None) -> None: + self.log, self.layers, self.slots = log, tuple(layers), slots or {} + self.primed = self.forward_open = False + self.eager_checks: list[str] = [] + + def require_eager(self, hook: str) -> None: + self.eager_checks.append(hook) + + def prefetch(self, layer: int) -> None: + self.log.append(("L", layer)) + + def mark_primed(self) -> None: + self.log.append(("mark_primed",)) + self.primed = True + + def war_barrier(self) -> None: + self.log.append(("war_barrier",)) + self.primed = self.forward_open = False + + def release(self) -> None: + self.log.append(("dense.release",)) + + +class _DenseModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.model = nn.Module() + self.model.layers = nn.ModuleList() + + +class TestDwdpManagerDenseSetup(unittest.TestCase): + """A DwdpManager over fakes, not entered unless a test enters it; its dense manager is real.""" + + LAYERS = [3, 5, 7] + #: Global rank 3 is DWDP rank 1 of group 1 on device 2: passing one for another shows. + GLOBAL_RANK, DWDP_RANK, DEVICE_ID = 3, 1, 2 + + def setUp(self): + _clear_global_managers() + self.addCleanup(_clear_global_managers) + self.log: list[tuple] = [] + self.sub_comm = _configure_mock_comm_world(_start(self, patch(f"{_FAC}.COMM_WORLD"))) + self.sub_comm.Free.side_effect = lambda: self.log.append(("group.Free",)) + _start(self, patch(f"{_FAC}.dwdp_global_rank", return_value=self.GLOBAL_RANK)) + _start(self, patch(f"{_FAC}.torch.cuda.current_device", return_value=self.DEVICE_ID)) + self.weight_manager = _FakeWeightManager(self.log) + self.setup_dwdp = _start( + self, patch(f"{_FAC}.setup_dwdp", return_value=self.weight_manager) + ) + teardown = patch( + f"{_FAC}.teardown_dwdp", + side_effect=lambda *_a, **_k: self.log.append(("teardown_dwdp",)), + ) + _start(self, teardown) + self.roots_available = _start( + self, patch(f"{_FAC}.dense_roots_available", return_value=True) + ) + self.setup_dense = _start(self, patch(f"{_FAC}.setup_dense_dwdp", return_value=None)) + _set_dense_env(self, None) + self.model = _DenseModel() + self.manager = self.new_manager() + + def new_manager(self) -> DwdpManager: + """A DwdpManager of DWDP rank ``DWDP_RANK`` with ``LAYERS`` registered in reverse order.""" + mapping = _make_mapping(dwdp_rank=self.DWDP_RANK) + manager = DwdpManager( + config=_make_config(num_groups=2), dist=MagicMock(spec=MPIDist), mapping=mapping + ) + for layer in reversed(self.LAYERS): + manager.add_layer(layer) + return manager + + def test_dense_setup_runs_first_and_hands_its_stream_and_slots_to_the_expert_setup(self): + """On every rank of a model with roots, requested or not, before the expert setup. + + The expert setup gets the copy stream dense setup made (none if it made none) and the + committed slots. If it then fails (the ring does not fit), the runtime no manager took + over is released and the error propagates as raised. + """ + os.environ.pop("TRTLLM_DWDP_PREFETCH_GATHER", None) # restored by _set_dense_env's patch + stream = object() + _start(self, patch(f"{_FAC}.new_expert_copy_stream", return_value=stream)) + slots = {"ring": [SimpleNamespace(storage=torch.zeros(1)) for _ in range(2)]} + dense, seen = _FakeDense(self.log, self.LAYERS, slots), [] + + def dense_setup(model, transfer, comm, **kwargs): + seen.append((type(transfer), self.setup_dwdp.call_count, self.manager._model_ref)) + self.assertIs(transfer.copy_stream, stream) # the first read, in P2, creates it + return dense + + self.setup_dense.side_effect = dense_setup + self.weight_manager.copy_stream = stream + self.assertIs(self.manager.setup(self.model), self.weight_manager) + + self.setup_dense.assert_called_once_with( + self.model, + unittest.mock.ANY, + self.sub_comm, + dwdp_rank=self.DWDP_RANK, + dwdp_size=2, + device_id=self.DEVICE_ID, + moe_layers=[3, 5, 7], + ) + self.assertEqual(seen, [(_AttachedTransfer, 0, None)]) + kwargs = self.setup_dwdp.call_args.kwargs + self.assertEqual((kwargs["device_id"], kwargs["copy_stream"]), (self.DEVICE_ID, stream)) + self.assertEqual(kwargs["num_prefetch_slots"], 2) + self.assertEqual(kwargs["repeated_layer_indices"], []) + storages = {slot.storage.untyped_storage().data_ptr() for slot in slots["ring"]} + self.assertEqual(kwargs["reserved_storages"], storages) + self.assertIs(self.manager.dense, dense) + + self.setup_dense.side_effect = None # abandoned: dense DWDP stays off, no stream made + self.new_manager().setup(_DenseModel()) + self.assertIsNone(self.setup_dwdp.call_args.kwargs["copy_stream"]) + + self.setup_dense.side_effect = lambda *_args, **_kwargs: _FakeDense(self.log, self.LAYERS) + error = self.setup_dwdp.side_effect = PrefetchRingMemoryError("the ring does not fit") + manager = self.new_manager() + with self.assertRaises(PrefetchRingMemoryError) as raised: + manager.setup(_DenseModel()) + + self.assertIs(raised.exception, error) + self.assertEqual(self.log, [("dense.release",)]) + self.assertIsNone(manager.dense) + + def test_an_expert_setup_on_another_copy_stream_or_transfer_is_refused(self): + """The dense and the expert copies must share one copy stream and one transfer. + + A weight manager on another copy stream than dense setup made, or with another + ``use_batched_prefetch``, is refused after the expert setup; the runtime no manager took + over is released. + """ + os.environ.pop("TRTLLM_DWDP_PREFETCH_GATHER", None) # restored by _set_dense_env's patch + stream = object() + _start(self, patch(f"{_FAC}.new_expert_copy_stream", return_value=stream)) + + def dense_setup(model, transfer, comm, **kwargs): + self.assertIs(transfer.copy_stream, stream) # the first read, in P2, creates it + return _FakeDense(self.log, self.LAYERS) + + self.setup_dense.side_effect = dense_setup + for copy_stream, batched in ((object(), False), (stream, True)): + with self.subTest(other_stream=copy_stream is not stream, batched=batched): + self.log.clear() + self.weight_manager.copy_stream = copy_stream + self.weight_manager.use_batched_prefetch = batched + manager = self.new_manager() + with self.assertRaisesRegex(RuntimeError, "must share one copy stream"): + manager.setup(_DenseModel()) + + self.assertEqual(self.log, [("dense.release",)]) + self.assertIsNone(manager.dense) + + def test_a_model_without_roots_makes_no_dense_collective(self): + """A rank that asks for dense DWDP decides alone, after the expert setup. + + It raises unless its own ``TRTLLM_DWDP_DENSE_REQUIRED`` is off; switches it cannot read + count as required. A rank that does not ask serves expert DWDP alone. + """ + self.roots_available.return_value = False + log_error = _start(self, patch("tensorrt_llm._torch.modules.dwdp.dense.setup.logger.error")) + cases = ( + (None, {}, False), + ("0", {}, False), + ("1", {}, True), + ("1", {"REQUIRED": "0"}, False), + ("1", {"REQUIRED": "0", "SEGMENTS": "atn"}, True), + ) + for request, switches, raises in cases: + with self.subTest(request=request, **switches): + _set_dense_env(self, request) + for name in [name for name in os.environ if name.startswith(f"{DENSE_ENV}_")]: + del os.environ[name] + os.environ.update({f"{DENSE_ENV}_{key}": value for key, value in switches.items()}) + if raises: + with self.assertRaisesRegex( + DenseSetupError, "_DenseModel do not provide dense_dwdp_roots" + ): + self.manager.setup(self.model) + else: + self.assertIs(self.manager.setup(self.model), self.weight_manager) + self.assertIsNone(self.manager.dense) + self.assertEqual(self.setup_dwdp.call_count, len(cases)) + self.setup_dense.assert_not_called() + self.sub_comm.allgather.assert_not_called() + self.assertIn("_REQUIRED=0, so dense DWDP is off on this rank", log_error.call_args.args[0]) + + def test_setup_priming_abort_recovery_and_teardown_keep_their_calls(self): + """The attached lifecycle against one call log. + + The dense priming of the first two layers before every expert start; after an aborted + forward the executor's start leaves the recovery to H0: the WAR barrier, the dense + priming, then the expert start; dense released first. + """ + dense = _FakeDense(self.log, self.LAYERS) + self.setup_dense.return_value = dense + + with self.manager: + self.manager.setup(self.model) + self.assertIs(self.manager.dense_manager.dense, dense) + self.manager.start_forward() # the executor loop, before the forward + self.manager.ensure_primed() # H0 of that forward: already primed + dense.primed, dense.forward_open = False, True # the forward aborts midway + self.manager.start_forward() # skipped: H0 recovers on the execution stream + self.manager.dense_manager.ensure_primed() # H0 as the model calls it + self.manager.start_forward() # dense primed again: the expert start only + + priming = [("L", 3), ("L", 5), ("mark_primed",)] + start = [("start_forward",)] + recovery = [("war_barrier",)] + priming + start + teardown = [("dense.release",), ("teardown_dwdp",), ("group.Free",)] + self.assertEqual(self.log, priming + start + start + recovery + start + teardown) + self.assertEqual(dense.eager_checks, [H0_HOOK, H0_HOOK]) + self.assertIsNone(get_global_dense_dwdp_manager()) + + def test_more_than_two_expert_slots_are_logged_when_dense_dwdp_commits(self): + """With more than two expert prefetch slots each expert copy runs ahead of its layer's + dense copy on the shared copy stream; setup says so at INFO.""" + info = _start(self, patch(f"{_FAC}.logger.info")) + for slots, committed in ((2, True), (3, False), (3, True)): + with self.subTest(slots=slots, committed=committed): + self.setup_dense.return_value = MagicMock() if committed else None + manager = self.new_manager() + manager.config = manager.config.model_copy(update={"num_prefetch_slots": slots}) + info.reset_mock() + manager.setup(_DenseModel()) + messages = [call.args[0] for call in info.call_args_list] + logged = [ + message + for message in messages + if f"{slots} expert prefetch slots: each expert copy runs ahead of its " + "layer's dense copy" in message + ] + self.assertEqual(len(logged), int(slots > 2 and committed), messages) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_mapping.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_mapping.py index f7ba7942f80b..380bbf03e08d 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_mapping.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_mapping.py @@ -132,7 +132,7 @@ def test_moe_ep_rank_zero_when_dwdp_enabled(self): """DWDP on: moe_ep_rank is always 0 (matches moe_ep_size = 1). The actual rank-specific expert range comes from - ``DwdpManager.start_expert_id`` via ``_init_dwdp_expert_layout``, + ``DwdpManager.expert_partition`` via ``_init_dwdp_expert_layout``, not from ``moe_ep_rank``. """ for dwdp_rank in range(4): diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_peer_ranges.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_peer_ranges.py index 411b68c471d0..f5940218b16e 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_peer_ranges.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_peer_ranges.py @@ -12,103 +12,133 @@ # 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. -"""Unit tests for ``compute_peer_ranges`` and ``lookup_owner``. - -These two helpers are the core of Phase 2 (non-uniform / redundant -partition support). They replace the legacy -``cursor // experts_per_rank`` formula used everywhere a remote expert -needed to be mapped to its owning peer rank. - -Cases covered: - - * Uniform partition (size == stride == num_experts // dwdp_size): - behavior must match the legacy ``//`` formula exactly. - * Non-uniform tail padding (size * dwdp_size > num_experts): the tail - rank's valid range is capped at ``num_experts``. - * Redundancy (size > stride): adjacent ranges overlap; - ``lookup_owner`` picks the lowest-rank owner (deterministic). - * Boundary expert ids (0, num_experts - 1). - * Out-of-range expert ids raise. +"""Unit tests for ``ExpertPartition``, ``handle_geometry`` and ``lookup_owner``. + +0/0 is the EP partition of the MoE modules (balanced); otherwise a strided table whose last range +ends at N (uniform, redundant) or is cut at N (capped). ``lookup_owner`` gives the lowest owner. """ +import itertools import unittest import pytest -from tensorrt_llm._torch.modules.dwdp.specs import compute_peer_ranges, lookup_owner +from tensorrt_llm._torch.modules.dwdp.specs import ( + ExpertPartition, + PageAlignedLayout, + handle_geometry, + lookup_owner, +) +from tensorrt_llm._torch.moe.fused_moe.interface import _compute_ep_partition pytestmark = pytest.mark.cpu_only +K3_EXPERTS = 896 -class TestComputePeerRanges(unittest.TestCase): - def test_uniform_dwdp4_matches_floor_div(self): - # 256 experts / 4 ranks = 64 each. Every rank's range starts at - # rank * 64 and ends at the next rank's start. - peer_ranges = compute_peer_ranges( - dwdp_size=4, - num_experts_per_worker=64, - num_prefetch_experts=64, - num_experts_total=256, - ) - self.assertEqual(peer_ranges, [(0, 64), (64, 128), (128, 192), (192, 256)]) - - def test_non_uniform_tail_padding_capped(self): - # Defensive: the function caps the last rank's end at - # ``num_experts_total`` even when ``(dwdp-1)*stride + size > - # num_experts``. In production this configuration is rejected - # upstream by ``_validate_partition_config`` (tail padding is - # incompatible with GB200 fabric-handle partial mapping), but - # ``compute_peer_ranges`` itself stays robust. - peer_ranges = compute_peer_ranges( - dwdp_size=3, - num_experts_per_worker=86, - num_prefetch_experts=86, - num_experts_total=256, - ) - self.assertEqual(peer_ranges, [(0, 86), (86, 172), (172, 256)]) - def test_mode_b_dwdp3_overlap(self): - # dwdp=3, 256 experts, Mode B: size=86, stride=85 — the - # production recipe for ``num_experts not divisible by dwdp_size``. - # ``(dwdp-1)*stride + size == num_experts`` exactly, with 1 - # expert overlapping between adjacent ranks. - peer_ranges = compute_peer_ranges( - dwdp_size=3, - num_experts_per_worker=86, - num_prefetch_experts=85, - num_experts_total=256, - ) - self.assertEqual(peer_ranges, [(0, 86), (85, 171), (170, 256)]) - - def test_redundancy_overlap(self): - # dwdp=4, size=70, stride=62: 8-expert overlap between adjacent - # ranks. Last rank's valid end is min(186 + 70, 256) = 256. - peer_ranges = compute_peer_ranges( - dwdp_size=4, - num_experts_per_worker=70, - num_prefetch_experts=62, - num_experts_total=256, +def _from_config(num_experts, dwdp_size, size, stride): + return ExpertPartition.from_config( + num_experts=num_experts, + dwdp_size=dwdp_size, + num_experts_per_worker=size, + num_prefetch_experts=stride, + ) + + +@pytest.mark.parametrize("num_experts", [7, 72, 73, 128, 160, 256, 384, K3_EXPERTS]) +def test_balanced_partition_is_the_ep_partition(num_experts): + for dwdp_size in range(2, 33): + if num_experts < dwdp_size: + with pytest.raises(ValueError): + ExpertPartition.balanced(num_experts=num_experts, dwdp_size=dwdp_size) + continue + partition = ExpertPartition.balanced(num_experts=num_experts, dwdp_size=dwdp_size) + assert partition.ranges == tuple( + _compute_ep_partition(num_experts, dwdp_size, rank)[1:] for rank in range(dwdp_size) ) - self.assertEqual(peer_ranges, [(0, 70), (62, 132), (124, 194), (186, 256)]) - # Verify overlap exists between every adjacent pair. - for r in range(len(peer_ranges) - 1): - cur_end = peer_ranges[r][1] - next_start = peer_ranges[r + 1][0] - self.assertGreater( - cur_end, next_start, f"Expected overlap between rank {r} and {r + 1}" - ) + + +def test_strided_partition_accepts_exact_and_capped_tables(): + """Every (size, stride) on a grid: accepted iff exact or capped, with the strided ranges.""" + mismatches = [] + for num_experts, dwdp_size in itertools.product((8, 9, 10, 12, 72, 96), (2, 3, 4, 5)): + for size, stride in itertools.product(range(1, num_experts + 2), repeat=2): + try: + partition = ExpertPartition.strided( + num_experts=num_experts, dwdp_size=dwdp_size, size=size, stride=stride + ) + got = (partition.kind, partition.peer_ranges) + except ValueError: + got = None + # Accepted: ranges that end exactly at N (stride <= size), or a capped table. + exact = stride <= size and (dwdp_size - 1) * stride + size == num_experts + capped = stride == size and (dwdp_size - 1) * size < num_experts < dwdp_size * size + want = None + if exact or capped: + kind = "capped" if capped else ("uniform" if size == stride else "redundant") + ranges = [ + (r * stride, min(r * stride + size, num_experts)) for r in range(dwdp_size) + ] + want = (kind, ranges) + if got != want: + mismatches.append((num_experts, dwdp_size, size, stride)) + assert mismatches == [] + + +@pytest.mark.parametrize( + "num_experts, dwdp_size, size, stride, counts", + [ + (K3_EXPERTS, 12, 0, 0, [75] * 8 + [74] * 4), + (K3_EXPERTS, 12, 82, 74, [82] * 12), + (72, 2, 40, 40, [40, 32]), + (K3_EXPERTS, 12, 75, 75, [75] * 11 + [71]), + # Rejected: the last rank would store nothing, a gap, overlap past the experts, one field + # set, negative fields, fewer experts than ranks. + (K3_EXPERTS, 33, 28, 28, None), + (72, 10, 8, 8, None), + (K3_EXPERTS, 12, 74, 75, None), + (K3_EXPERTS, 12, 83, 74, None), + (K3_EXPERTS, 12, 0, 74, None), + (K3_EXPERTS, 12, 82, 0, None), + (K3_EXPERTS, 12, -1, -1, None), + (7, 8, 0, 0, None), + ], +) +def test_partition_from_config(num_experts, dwdp_size, size, stride, counts): + if counts is None: + with pytest.raises(ValueError): + _from_config(num_experts, dwdp_size, size, stride) + return + partition = _from_config(num_experts, dwdp_size, size, stride) + assert [partition.count_of(rank) for rank in range(dwdp_size)] == counts + + +@pytest.mark.parametrize( + "ranges", + [((0, 4), (5, 8)), ((1, 4), (4, 8)), ((0, 4), (4, 7)), ((0, 4), (4, 4), (4, 8))], + ids=["gap", "late_start", "short_end", "empty_range"], +) +def test_partition_refuses_an_expert_without_owner_or_a_rank_without_expert(ranges): + with pytest.raises(ValueError): + ExpertPartition(num_experts=8, ranges=ranges, kind="custom") + + +@pytest.mark.parametrize("expert_bytes", [11010048, 5505024, 1376256, 688128]) +def test_handle_geometry_is_the_page_aligned_layout_of_the_range(expert_bytes): + """The span ``WeightBuffer`` maps whole, at its data offset; Kimi-K3's four expert weights.""" + page = 2 << 20 + for partition in (_from_config(K3_EXPERTS, 12, 0, 0), _from_config(K3_EXPERTS, 12, 82, 74)): + for start, end in partition.ranges: + layout = PageAlignedLayout.compute(expert_bytes, K3_EXPERTS, start, end, page, 1 << 62) + geometry = (layout.mnnvl_size, layout.data_offset) + assert handle_geometry(expert_bytes, start, end, page) == geometry class TestLookupOwner(unittest.TestCase): def test_uniform_matches_floor_div(self): # The whole point: lookup_owner must agree with the # legacy formula on every uniform expert id. - peer_ranges = compute_peer_ranges( - dwdp_size=4, - num_experts_per_worker=64, - num_prefetch_experts=64, - num_experts_total=256, - ) + peer_ranges = _from_config(256, 4, 64, 64).peer_ranges for expert_id in range(256): self.assertEqual( lookup_owner(expert_id, peer_ranges), @@ -121,12 +151,7 @@ def test_mode_b_dwdp3_overlap(self): # recipe for non-divisible dwdp_size. Adjacent ranks overlap by # 1 expert; lookup_owner picks the lowest-rank owner for shared # experts (deterministic). - peer_ranges = compute_peer_ranges( - dwdp_size=3, - num_experts_per_worker=86, - num_prefetch_experts=85, - num_experts_total=256, - ) + peer_ranges = _from_config(256, 3, 86, 85).peer_ranges # Rank 0's exclusive range self.assertEqual(lookup_owner(0, peer_ranges), 0) self.assertEqual(lookup_owner(84, peer_ranges), 0) @@ -146,12 +171,7 @@ def test_redundancy_picks_lowest_rank(self): # dwdp=4, size=70, stride=62: experts in the overlap zone # [62, 70) are owned by both rank 0 and rank 1; lookup_owner # must consistently return the lowest rank (0) for these. - peer_ranges = compute_peer_ranges( - dwdp_size=4, - num_experts_per_worker=70, - num_prefetch_experts=62, - num_experts_total=256, - ) + peer_ranges = _from_config(256, 4, 70, 62).peer_ranges # Experts in rank 0's exclusive range self.assertEqual(lookup_owner(0, peer_ranges), 0) self.assertEqual(lookup_owner(61, peer_ranges), 0) diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py index 96f21c46f79b..8d1b8f2dbe88 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py @@ -30,6 +30,8 @@ import torch from tensorrt_llm._torch.modules.dwdp import fabric_gather +from tensorrt_llm._torch.modules.dwdp.dense.manager import DenseDwdpManager +from tensorrt_llm._torch.modules.dwdp.dense.types import DenseMode from tensorrt_llm._torch.modules.dwdp.weight_manager import DWDPWeightManager from tensorrt_llm._torch.pyexecutor.dwdp import DwdpManager @@ -115,6 +117,8 @@ def __init__( ) self.mgr = object.__new__(DwdpManager) self.mgr._weight_manager = self.wm + # Expert DWDP alone: the attached dense manager has no runtime. + self.mgr._dense_manager = DenseDwdpManager(DenseMode.ATTACHED) self.backend = _Backend() self.compute_stream = torch.cuda.Stream() self.copies: List[int] = [] diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_pipeline.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_pipeline.py index c8f121c26ada..19c71bb3d04b 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_pipeline.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_pipeline.py @@ -207,6 +207,48 @@ def test_unknown_layer_raises(self): with pytest.raises(KeyError): self.manager.finish_layer(4) + def test_a_given_copy_stream_on_the_buffer_device_is_used_and_none_is_created(self): + kwargs = dict( + weight_buffer=self.weight_buffer, + peer_views={}, + peer_ranges=[(0, 2), (2, 4)], + moe_layer_indices=LAYERS, + weight_names=["gate_up_proj", "down_proj"], + dwdp_rank=0, + dwdp_size=2, + ) + other = _Stream("other", self.log) + other.device = torch.device("cuda", 1) + with pytest.raises( + ValueError, match="copy_stream is on cuda:1, the weight buffer on cuda:0" + ): + DWDPWeightManager(**kwargs, copy_stream=other) + + given = _Stream("given", self.log) + given.device = torch.device("cuda", 0) + names = [f"prefetch{slot}" for slot in range(self.num_slots)] + names += [f"consume{slot}" for slot in range(self.num_slots)] + events = iter([_Event(name, self.log) for name in names]) + module = "tensorrt_llm._torch.modules.dwdp.weight_manager.torch.cuda" + with patch(f"{module}.Stream") as new_stream: + with patch(f"{module}.Event", side_effect=lambda: next(events)): + manager = DWDPWeightManager(**kwargs, copy_stream=given) + manager._released = True + new_stream.assert_not_called() + assert manager.copy_stream is given + self.log.clear() + manager.start_forward() + # Every copy waits for its slot's consume event and records its prefetch event on it. + assert self.log == [ + entry + for slot in range(self.num_slots) + for entry in ( + ("wait", "given", f"consume{slot}"), + ("copy", LAYERS[slot]), + ("record", f"prefetch{slot}", "given"), + ) + ] + class TestDWDPPrefetchPipelineRepeatedLayer(TestDWDPPrefetchPipeline): """The last MoE layer runs three times per pass, like the one-model MTP-Eagle layer.""" diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py index 1ae2e167e0da..e2d9fb08b501 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py @@ -14,14 +14,23 @@ # limitations under the License. """Failure rollback tests for the DWDP setup ownership chain.""" +import gc +import pickle from contextlib import ExitStack from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest +import torch import torch.nn as nn +from tensorrt_llm._torch.modules.dwdp import page_pool, weight_buffer from tensorrt_llm._torch.modules.dwdp import setup as dwdp_setup_module +from tensorrt_llm._torch.modules.dwdp.page_pool import PagePool +from tensorrt_llm._torch.modules.dwdp.specs import MnnvlHandleSet, WeightSpec +from tensorrt_llm._torch.modules.dwdp.weight_buffer import PrefetchRingMemoryError, WeightBuffer + +pytestmark = pytest.mark.cpu_only def _model(): @@ -38,7 +47,6 @@ def _mapping(): def _setup_mocks(stack: ExitStack): transport = MagicMock() transport.get_handle_set.return_value = MagicMock() - transport.get_peer_ranges.return_value = [(0, 2), (2, 4)] transport.get_peer_views.return_value = {} weight_buffer = MagicMock() weight_manager = MagicMock() @@ -57,14 +65,10 @@ def _setup_mocks(stack: ExitStack): stack.enter_context( patch.object(dwdp_setup_module, "build_weight_specs", return_value=MagicMock()) ) + stack.enter_context(patch.object(dwdp_setup_module, "_agree_on_partition")) stack.enter_context( - patch.object( - dwdp_setup_module, - "_get_first_spec", - return_value=SimpleNamespace(local_experts=2), - ) + patch.object(dwdp_setup_module, "get_allocation_granularity", return_value=2 << 20) ) - stack.enter_context(patch.object(dwdp_setup_module, "_validate_partition_config")) stack.enter_context( patch.object(dwdp_setup_module.DWDPTransport, "create", return_value=transport) ) @@ -75,11 +79,21 @@ def _setup_mocks(stack: ExitStack): stack.enter_context( patch.object(dwdp_setup_module, "DWDPWeightManager", return_value=weight_manager) ) + stack.enter_context( + patch.object(dwdp_setup_module, "_exchange_expert_scalars", side_effect=_vote) + ) stack.enter_context(patch.object(dwdp_setup_module, "fixup_moe_backends")) return transport, weight_buffer, weight_manager -def _run_setup(): +def _vote(*_args, error=None, **_kwargs): + """Stands in for ``_exchange_expert_scalars``: a rank that failed raises its own error there.""" + if error is not None: + raise error + return {} + + +def _run_setup(**kwargs): return dwdp_setup_module.setup_dwdp( model=_model(), mapping=_mapping(), @@ -89,6 +103,7 @@ def _run_setup(): num_experts_per_worker=2, num_prefetch_experts=2, num_prefetch_slots=2, + **kwargs, ) @@ -153,3 +168,123 @@ def test_backend_fixup_failure_releases_owned_manager(): assert weight_manager._transport is transport weight_manager.release.assert_called_once_with() + + +def test_an_expert_weight_in_a_dense_slot_is_refused_before_the_transport(): + """The transport would free a reserved storage (a committed dense slot) under its owner.""" + slot = torch.zeros(64, dtype=torch.uint8) + params = { + (3, "w3_w1_weight"): slot[16:32], + (3, "w2_weight"): torch.zeros(16, dtype=torch.uint8), + } + for reserved_storages, refused in (({slot.untyped_storage().data_ptr()}, True), ({1}, False)): + with ExitStack() as stack: + _setup_mocks(stack) + dwdp_setup_module.collect_moe_params.return_value = (params, ["w3_w1_weight"]) + if refused: + with pytest.raises( + RuntimeError, + match=r"Expert weights in reserved storage, which the transport would free: " + r"layers\.3\.w3_w1_weight\. They are views of the committed dense DWDP slots", + ): + _run_setup(reserved_storages=frozenset(reserved_storages)) + dwdp_setup_module.DWDPTransport.create.assert_not_called() + else: + _run_setup(reserved_storages=frozenset(reserved_storages)) + dwdp_setup_module.DWDPTransport.create.assert_called_once() + + +GRANULARITY, PAGE, TOTAL = 2 << 20, 16 << 20, 278 << 30 + + +@pytest.mark.parametrize("num_slots", [2, 3]) +def test_a_ring_larger_than_free_memory_is_refused_before_any_allocation(monkeypatch, num_slots): + """One MoE layer per slot, of two 1 MiB-per-expert weights, 4 of 16 experts local: 2 pages + per slot.""" + shapes = {"w1": (1024, 1024), "w2": (512, 2048)} + specs = { + layer: { + name: WeightSpec( + 16, chunk_shape=(4, *shape), full_shape=(16, *shape), dtype=torch.uint8 + ) + for name, shape in shapes.items() + } + for layer in range(3, 3 + num_slots) + } + keys = [(layer, name) for layer in specs for name in shapes] + handles = MnnvlHandleSet( + handles={key: index for index, key in enumerate(keys, 1)}, + sizes=dict.fromkeys(keys, 2 * GRANULARITY), # the 4 local experts' pages + ) + needed, calls = 2 * num_slots * PAGE, [] + monkeypatch.setattr(weight_buffer, "get_allocation_granularity", lambda *_a, **_k: GRANULARITY) + monkeypatch.setattr(weight_buffer.torch.cuda, "empty_cache", lambda: calls.append("empty")) + monkeypatch.setattr( + weight_buffer.torch.cuda, + "mem_get_info", + lambda device_id: calls.append(("mem_get_info", device_id)) or (needed - 1, TOTAL), + ) + create = MagicMock() + monkeypatch.setattr(weight_buffer.PagePool, "create", create) + monkeypatch.setattr(weight_buffer.WeightBuffer, "_setup_layer", MagicMock()) + + with pytest.raises(PrefetchRingMemoryError) as raised: + WeightBuffer.create( + layer_weight_specs=specs, + handles=handles, + local_start=0, + local_end=4, + dwdp_size=4, + device_id=0, + num_slots=num_slots, + ) + + create.assert_not_called() + assert calls == ["empty", ("mem_get_info", 0)] + error = raised.value + assert isinstance(error, RuntimeError) + numbers = (error.needed_bytes, error.free_bytes, error.total_bytes, error.device_id) + assert numbers == (needed, needed - 1, TOTAL, 0) + assert (error.slot_pages, error.page_size) == ((2,) * num_slots, PAGE) + message = str(error) + pages = " + ".join(["2"] * num_slots) + gib = needed / (1 << 30) + assert message.startswith( + f"[DWDP Setup] The expert prefetch ring does not fit on GPU 0: its {num_slots} slots need " + f"{gib:.2f} GiB ({pages} pages of 16 MiB, the remote experts of {num_slots} MoE layers), " + f"but {gib:.2f} GiB of 278.00 GiB is free." + ) + assert "use a larger dwdp_size" in message and "dense DWDP (TRTLLM_DWDP_DENSE=1)" in message + # Fewer slots is a remedy only above the two the prefetch needs. + assert ("fewer prefetch slots (DwdpConfig.num_prefetch_slots)" in message) == (num_slots > 2) + # The executor sends a setup error to the proxy pickled: the error survives it. + copy = pickle.loads(pickle.dumps(error)) + assert type(copy) is PrefetchRingMemoryError and str(copy) == message + assert (copy.needed_bytes, copy.slot_pages, copy.page_size) == (needed, (2,) * num_slots, PAGE) + + +def test_a_page_pool_allocation_failure_releases_every_handle_once(monkeypatch): + """Slot 1 page 2 fails: the 5 handles made so far are released once, ``__del__`` adds none.""" + created, released = [], [] + + def create_local_handle(size: int, device_id: int) -> int: + assert (size, device_id) == (PAGE, 0) + if len(created) == 5: + raise RuntimeError("CUDA_ERROR_OUT_OF_MEMORY") + created.append(1000 + len(created)) + return created[-1] + + monkeypatch.setattr(page_pool, "create_local_handle", create_local_handle) + monkeypatch.setattr(page_pool, "release_handle", released.append) + error = None + try: + PagePool([3 * PAGE, 3 * PAGE], 0, granularity=GRANULARITY, page_size=PAGE) + except RuntimeError as exc: + error = exc + assert str(error) == "Failed to allocate page 2 in slot 1: CUDA_ERROR_OUT_OF_MEMORY" + assert str(error.__cause__) == "CUDA_ERROR_OUT_OF_MEMORY" + # Drop the traceback, the last reference to the half-built pool: its __del__ runs now. + del error + gc.collect() + assert created == [1000, 1001, 1002, 1003, 1004] + assert sorted(released) == created diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_protocol.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_protocol.py new file mode 100644 index 000000000000..f63df98dc035 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_protocol.py @@ -0,0 +1,345 @@ +# 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. +"""``setup_dwdp`` on one thread per DWDP rank, over ``StoreDwdpGroup`` on one ``HashStore``. + +``_FakeDevice`` runs the transport's VMM calls on host memory; the weight buffer, edge-byte fill +and weight manager are stand-ins. A rank that skips a collective times out instead of hanging. +""" + +import math +import re +import threading +from collections import Counter +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from torch.distributed import HashStore + +from tensorrt_llm._torch.modules.dwdp import setup as dwdp_setup +from tensorrt_llm._torch.modules.dwdp import transport +from tensorrt_llm._torch.modules.dwdp.rendezvous import StoreDwdpGroup +from tensorrt_llm._torch.modules.dwdp.specs import ExpertPartition +from tensorrt_llm._torch.modules.dwdp.weight_buffer import PrefetchRingMemoryError + +pytestmark = pytest.mark.cpu_only + +GRANULARITY = 4096 +LAYER = 1 +#: int32 columns per expert: 160 and 384 bytes, so most ranges start and end inside a page. +WEIGHTS = {"w2_weight": 40, "w3_w1_weight": 96} +ALL_COLLECTIVES = ["handshake", "export", "export", "import status", "vote"] + +_RANK = threading.local() + + +def _rows(start, end, width): + """Experts ``[start, end)`` of a toy weight: each row holds its expert's id.""" + return torch.arange(start, end, dtype=torch.int32)[:, None].repeat(1, width) + + +def _pages(start, end, expert_bytes): + """Bytes of the pages experts ``[start, end)`` of one weight touch: that rank's handle size.""" + page_end = -(-end * expert_bytes // GRANULARITY) * GRANULARITY + return page_end - start * expert_bytes // GRANULARITY * GRANULARITY + + +class _Experts(nn.Module): + """A rank's MoE backend after loading: the rows of experts ``[start, end)``.""" + + def __init__(self, num_experts, start, end, weights): + super().__init__() + self.num_experts = num_experts + for name, width in WEIGHTS.items(): + rows = weights.get(name, _rows(start, end, width)) + setattr(self, name, nn.Parameter(rows, requires_grad=False)) + alpha = torch.arange(start, end, dtype=torch.float32) + self.fc31_alpha = nn.Parameter(alpha, requires_grad=False) + self.fc31_input_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False) + + +class _Group(StoreDwdpGroup): + """A rank's DWDP group on the test's store; it logs which collective each allgather is.""" + + def __init__(self, store, rank, size): + super().__init__(rank, size, "unused", timeout_s=20, store=store) + self.log = [] + + def allgather(self, obj): + if not isinstance(obj, tuple): + self.log.append("import status") + elif len(obj) == 6: + self.log.append("handshake") + else: # (status, payload) of an export, (status, small parameters) of the vote + self.log.append("vote" if isinstance(obj[1], dict) else "export") + return super().allgather(obj) + + +class _WeightManager: + """Stands in for ``DWDPWeightManager``: records its arguments, releases the transport.""" + + def __init__(self, **kwargs): + self.kwargs = kwargs + self._transport = None + + def release(self): + self._transport.release() + + +class _FakeDevice: + """The transport's VMM calls over host memory: a handle is a ``uint8`` tensor, a VA a number. + + As on the device, a handle maps only whole. ``fault = (rank, step)``: that rank raises + ``injected`` at its second handle (``create``), at the access grant of its first imported + handle (``import``) or at its prefetch ring (``ring``). + """ + + def __init__(self, monkeypatch): + self.fault = self.injected = None + self.handles, self.va_ranges, self.mappings = {}, {}, {} + self.imported = set() + self.calls = Counter() # by step and by (rank, step) + self._lock = threading.Lock() + self._next_id = GRANULARITY + fabric = transport.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + fakes = { + "get_allocation_granularity": lambda device_id, use_cache=True: GRANULARITY, + "peer_handle_type": lambda: fabric, + "create_fabric_handle": self.create, + "_export_handle": lambda handle, htype: handle, + "_import_handle": self.import_handle, + "reserve_va": lambda size, alignment: self._add(self.va_ranges, size, size + alignment), + "map_handle": self.map, + "set_access": self.set_access, + "tensor_from_ptr": self.view, + "unmap_va": lambda va, size: self._drop(self.mappings, va), + "free_va": lambda va, size: self._drop(self.va_ranges, va), + "release_handle": lambda handle: self._drop(self.handles, handle), + } + for name, fake in fakes.items(): + monkeypatch.setattr(transport, name, fake) + monkeypatch.setattr( + dwdp_setup, "get_allocation_granularity", fakes["get_allocation_granularity"] + ) + monkeypatch.setattr(dwdp_setup.WeightBuffer, "create", self.create_weight_buffer) + monkeypatch.setattr(dwdp_setup, "fill_edge_bytes", lambda **_kwargs: None) + monkeypatch.setattr(dwdp_setup, "DWDPWeightManager", _WeightManager) + monkeypatch.setattr(torch.cuda, "synchronize", lambda device=None: None) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (1 << 40, 1 << 41)) + + def _step(self, step, failing_call=1): + with self._lock: + self.calls[step] += 1 + self.calls[_RANK.value, step] += 1 + count = self.calls[_RANK.value, step] + if self.fault == (_RANK.value, step) and count == failing_call: + raise self.injected + + def _add(self, table, value, span=1): + with self._lock: + key = self._next_id + self._next_id += span + table[key] = value + return key + + def _drop(self, table, key): + with self._lock: + del table[key] + + def create(self, size, device_id): + self._step("create", failing_call=2) + return self._add(self.handles, torch.zeros(size, dtype=torch.uint8)) + + def import_handle(self, payload, htype): + handle = self._add(self.handles, self.handles[payload]) + self.imported.add(handle) + return handle + + def map(self, va, size, handle, offset=0): + if offset or size != self.handles[handle].numel(): + raise RuntimeError("CUDA_ERROR_NOT_SUPPORTED: a handle maps only whole") + with self._lock: + self.mappings[va] = handle + + def set_access(self, va, size, device_id): + if self.mappings[va] in self.imported: + self._step("import") + + def view(self, ptr, shape, dtype, device_id): + with self._lock: + va = max(mapped for mapped in self.mappings if mapped <= ptr) + memory = self.handles[self.mappings[va]] + start, nbytes = ptr - va, math.prod(shape) * dtype.itemsize + assert start + nbytes <= memory.numel(), "a view past the end of its mapping" + return memory[start : start + nbytes].view(dtype).reshape(shape) + + def create_weight_buffer(self, **kwargs): + self._step("ring") + return SimpleNamespace(release=lambda: None, kwargs=kwargs) + + +@pytest.fixture +def device(monkeypatch): + return _FakeDevice(monkeypatch) + + +def _run_setup(num_experts, configs, loaded, weights=None): + """Rank ``r`` runs ``setup_dwdp`` with ``configs[r]`` (size, stride) on experts ``loaded[r]``. + + ``weights[r]`` replaces some of rank ``r``'s toy weights. Returns each rank's outcome. + """ + store, size = HashStore(), len(configs) + outcomes = [SimpleNamespace(result=None, error=None, log=None, model=None) for _ in configs] + + def run(rank): + _RANK.value = rank + outcome = outcomes[rank] + group = _Group(store, rank, size) + outcome.log = group.log + outcome.model = nn.Module() + outcome.model.layers = nn.ModuleList([nn.Module(), nn.Module()]) + experts = _Experts(num_experts, *loaded[rank], (weights or {}).get(rank, {})) + outcome.model.layers[LAYER].mlp = experts + try: + outcome.result = dwdp_setup.setup_dwdp( + model=outcome.model, + mapping=SimpleNamespace(dwdp_enabled=True, dwdp_rank=rank, dwdp_size=size), + device_id=0, + comm=group, + layer_indices=[LAYER], + num_experts_per_worker=configs[rank][0], + num_prefetch_experts=configs[rank][1], + num_prefetch_slots=2, + ) + except Exception as exc: # the test inspects every rank's outcome + outcome.error = exc + + threads = [threading.Thread(target=run, args=(rank,), daemon=True) for rank in range(size)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(60) + assert not any(thread.is_alive() for thread in threads), "setup_dwdp did not finish" + return outcomes + + +@pytest.mark.parametrize( + "num_experts, dwdp_size, size, stride", + [(896, 12, 0, 0), (72, 2, 40, 40), (64, 4, 16, 16), (256, 3, 86, 85)], + ids=["balanced", "capped", "uniform", "redundant"], +) +def test_every_peer_is_mapped_whole_with_its_own_size_and_rows( + device, num_experts, dwdp_size, size, stride +): + """Every peer maps a rank's handle, the pages of its range, whole and views its rows. + + Every rank makes the same collectives, gets the full ``fc31_alpha`` and hands the + partition's ranges to the weight buffer and manager. + """ + partition = ExpertPartition.from_config( + num_experts=num_experts, + dwdp_size=dwdp_size, + num_experts_per_worker=size, + num_prefetch_experts=stride, + ) + + outcomes = _run_setup(num_experts, [(size, stride)] * dwdp_size, partition.ranges) + + for rank, outcome in enumerate(outcomes): + assert outcome.error is None, f"rank {rank}: {outcome.error!r}" + assert outcome.log == ALL_COLLECTIVES + handles = outcome.result._transport.get_handle_set() + views = outcome.result._transport.get_peer_views() + for name, width in WEIGHTS.items(): + assert handles.get_size(LAYER, name) == _pages(*partition.range_of(rank), 4 * width) + for peer, (start, end) in enumerate(partition.ranges): + if peer != rank: + assert torch.equal(views[(peer, LAYER, name)], _rows(start, end, width)) + manager = outcome.result.kwargs + assert manager["peer_ranges"] == partition.peer_ranges + assert manager["repeated_layer_indices"] == () + ring = manager["weight_buffer"].kwargs + assert (ring["local_start"], ring["local_end"]) == partition.range_of(rank) + assert ring["num_slots"] == 2 + alpha = outcome.model.layers[LAYER].mlp.fc31_alpha + assert torch.equal(alpha, torch.arange(num_experts, dtype=torch.float32)) + for outcome in outcomes: + outcome.result.release() + assert device.handles == device.va_ranges == device.mappings == {} + + +FAILING = 2 + + +@pytest.mark.parametrize( + "fault, collectives", + [ + ("rows", ["handshake"]), + ("partition", ["handshake"]), + ("dtype", ["handshake"]), + ("create", ["handshake", "export", "export"]), + ("import", ALL_COLLECTIVES[:4]), + ("ring", ALL_COLLECTIVES), + ], +) +def test_a_failure_on_one_rank_raises_on_every_rank_after_the_same_collective( + device, fault, collectives +): + """Rank 2 fails; every rank raises after the collective that reports it and releases all. + + A handshake mismatch (a row more, another partition, another dtype of the same bytes) raises + the same error everywhere before any handle exists. Otherwise the failing rank raises its + own error, the others one naming it. + """ + partition = ExpertPartition.balanced(num_experts=13, dwdp_size=4) # 4 + 3 + 3 + 3 + configs, loaded, weights = [(0, 0)] * 4, list(partition.ranges), {} + if fault == "rows": + loaded[FAILING] = (loaded[FAILING][0], loaded[FAILING][1] + 1) + elif fault == "partition": + configs[FAILING], loaded[FAILING] = (4, 3), (6, 10) # redundant: rank 2 stores [6, 10) + elif fault == "dtype": + weights = {"w2_weight": _rows(*loaded[FAILING], WEIGHTS["w2_weight"]).view(torch.float32)} + else: + device.fault = (FAILING, fault) + device.injected = (PrefetchRingMemoryError if fault == "ring" else RuntimeError)("injected") + + outcomes = _run_setup(13, configs, loaded, {FAILING: weights}) + + assert [outcome.log for outcome in outcomes] == [collectives] * 4 + assert device.handles == device.va_ranges == device.mappings == {} + if collectives == ["handshake"]: + assert device.calls["create"] == 0 + assert all(type(outcome.error) is RuntimeError for outcome in outcomes) + assert len({str(outcome.error) for outcome in outcomes}) == 1 + return + assert fault != "create" or not device.imported + for rank, outcome in enumerate(outcomes): + if rank == FAILING: + assert outcome.error is device.injected + else: # a plain RuntimeError: a store timeout would be a DistStoreError + assert type(outcome.error) is RuntimeError, repr(outcome.error) + assert re.search(rf"\brank {FAILING}\b", str(outcome.error)), str(outcome.error) + + +def test_the_handshake_names_rank_0_when_it_alone_computed_another_partition(device): + """Each rank is compared with the most common report, so the error names rank 0 here.""" + loaded = ExpertPartition.balanced(num_experts=13, dwdp_size=4).ranges # rank 0: [0, 4) too + outcomes = _run_setup(13, [(4, 3)] + [(0, 0)] * 3, loaded) # rank 0 redundant, 1-3 balanced + + assert len({str(outcome.error) for outcome in outcomes}) == 1 + named = re.findall(r"\brank (\d+) computed the (\w+) partition", str(outcomes[0].error)) + assert named == [("0", "redundant")] diff --git a/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py b/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py index f7cf0ff253fc..72e987d4a129 100644 --- a/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py +++ b/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py @@ -108,6 +108,22 @@ def test_resident_sms_is_the_block_count_of_a_resident_gather(monkeypatch): assert fabric_gather.resident_sms() == 24 +def test_resident_sms_of_a_gather_chosen_by_the_caller(monkeypatch): + """Standalone dense DWDP passes its own transport choice instead of the expert one.""" + _unset(monkeypatch, "TRTLLM_CUTEDSL_SM_RESERVE") + monkeypatch.setenv("TRTLLM_DWDP_GATHER_BLOCKS", "30") + monkeypatch.setattr(fabric_gather, "available", lambda: False) + monkeypatch.setattr(fabric_gather, "THREADS", 1024) + assert fabric_gather.resident_sms(gather=True) == 30 + assert fabric_gather.resident_sms(gather=False) == 0 + monkeypatch.setattr(fabric_gather, "available", lambda: True) + assert fabric_gather.resident_sms(gather=False) == 0 + monkeypatch.setattr(fabric_gather, "THREADS", 128) + assert fabric_gather.resident_sms(gather=True) == 0 + monkeypatch.setenv("TRTLLM_CUTEDSL_SM_RESERVE", "24") + assert fabric_gather.resident_sms(gather=False) == 24 + + def test_sm_reserve_is_set_explicitly(): from tensorrt_llm._torch import cute_dsl_utils diff --git a/tests/unittest/_torch/modules/dwdp/test_prefetch_copy_plan.py b/tests/unittest/_torch/modules/dwdp/test_prefetch_copy_plan.py index 8f560730581e..52ee8129a31a 100644 --- a/tests/unittest/_torch/modules/dwdp/test_prefetch_copy_plan.py +++ b/tests/unittest/_torch/modules/dwdp/test_prefetch_copy_plan.py @@ -17,15 +17,19 @@ ``_build_peer_descriptor_groups`` (one coalesced group per peer, consumed by the SM gather) and ``_build_batched_prefetch_copy_plan`` (the same bytes re-sliced and interleaved for ``cudaMemcpyBatchAsync``) must move identical bytes, so each plan is reduced to a canonical -interval set and compared across slice geometries, ranks and slice sizes. +interval set and compared across slice geometries, ranks and slice sizes, also with the unequal +peer views of the balanced expert partition. """ +import itertools from typing import Dict, List, Tuple import pytest import torch +from tensorrt_llm._torch.modules.dwdp import fabric_gather from tensorrt_llm._torch.modules.dwdp import weight_manager as wm +from tensorrt_llm._torch.modules.dwdp.specs import ExpertPartition, PeerRanges pytestmark = pytest.mark.cpu_only @@ -38,7 +42,8 @@ DTYPE = torch.bfloat16 ELEM_BYTES = DTYPE.itemsize LAYER_BYTES = len(WEIGHT_NAMES) * TOTAL_EXPERTS * EXPERT_ELEMS * ELEM_BYTES -PEER_WEIGHT_BYTES = EXPERTS_PER_PEER * EXPERT_ELEMS * ELEM_BYTES +EXPERT_BYTES = EXPERT_ELEMS * ELEM_BYTES +PEER_WEIGHT_BYTES = EXPERTS_PER_PEER * EXPERT_BYTES PRODUCTION_BOUNDS = [0, TOTAL_EXPERTS] FRAGMENTED_BOUNDS = [0, 10, 100, 101, 200, 373, TOTAL_EXPERTS] # straddles peer boundaries @@ -61,34 +66,39 @@ def get_remote_slices(self, layer_idx, name): def _make_manager(slice_bounds: List[int], dwdp_rank: int = 0): - """Return (manager, backings, src_ptr_to_peer); keep backings alive while a plan is in use.""" + """``_make_manager_over`` with equal peers and consecutive remote slices at ``slice_bounds``.""" assert slice_bounds[0] == 0 and slice_bounds[-1] == TOTAL_EXPERTS + peer_ranges = [(p * EXPERTS_PER_PEER, (p + 1) * EXPERTS_PER_PEER) for p in range(DWDP_SIZE)] + return _make_manager_over(peer_ranges, list(zip(slice_bounds, slice_bounds[1:])), dwdp_rank) + +def _make_manager_over(peer_ranges: PeerRanges, remote: List[Tuple[int, int]], dwdp_rank: int): + """Return (manager, backings, src_ptr_to_peer); keep backings alive while a plan is in use. + + Peer ``p``'s views have the rows of ``peer_ranges[p]``; ``remote`` holds the ``(start, end)`` + of each remote slice of the layer. + """ peer_views: Dict[Tuple[int, int, str], torch.Tensor] = {} src_ptr_to_peer: Dict[int, int] = {} - for peer in range(DWDP_SIZE): + for peer, (start, end) in enumerate(peer_ranges): for name in WEIGHT_NAMES: - t = torch.empty((EXPERTS_PER_PEER, EXPERT_ELEMS), dtype=DTYPE) + t = torch.empty((end - start, EXPERT_ELEMS), dtype=DTYPE) peer_views[(peer, LAYER, name)] = t src_ptr_to_peer[t.data_ptr()] = peer slices, backings = {}, [] for name in WEIGHT_NAMES: - backing = torch.empty((TOTAL_EXPERTS, EXPERT_ELEMS), dtype=DTYPE) + backing = torch.empty((peer_ranges[-1][1], EXPERT_ELEMS), dtype=DTYPE) backings.append(backing) - slices[(LAYER, name)] = [ - (backing[a:b], a, b) for a, b in zip(slice_bounds, slice_bounds[1:]) - ] + slices[(LAYER, name)] = [(backing[a:b], a, b) for a, b in remote] mgr = object.__new__(wm.DWDPWeightManager) mgr._weight_names = WEIGHT_NAMES mgr._weight_buffer = _FakeWeightBuffer(slices) mgr._peer_views = peer_views - mgr._peer_ranges = [ - (p * EXPERTS_PER_PEER, (p + 1) * EXPERTS_PER_PEER) for p in range(DWDP_SIZE) - ] + mgr._peer_ranges = list(peer_ranges) mgr._dwdp_rank = dwdp_rank - mgr._dwdp_size = DWDP_SIZE + mgr._dwdp_size = len(peer_ranges) return mgr, backings, src_ptr_to_peer @@ -168,3 +178,31 @@ def test_batch_interleaves_peers_round_robin(monkeypatch): assert len(src) > DWDP_SIZE first_round = [_peer_of(s, ptrs) for s in src[:DWDP_SIZE]] assert sorted(first_round) == list(range(DWDP_SIZE)), first_round + + +@pytest.mark.parametrize("dwdp_size", [12, 20, 24], ids=["75-74", "45-44", "38-37"]) +def test_unequal_peer_views_read_each_remote_expert_from_its_owner(dwdp_size, monkeypatch): + """Every rank: each descriptor copies its experts' rows from their owner, each remote once.""" + partition = ExpertPartition.balanced(num_experts=896, dwdp_size=dwdp_size) + for rank, (start, end) in enumerate(partition.ranges): + remote = [(a, b) for a, b in ((0, start), (end, 896)) if a < b] + mgr, backings, _ptrs = _make_manager_over(partition.peer_ranges, remote, rank) + groups = _groups(mgr) + flat = _flat_groups(groups) + for dst, src, size in zip(*flat): + (peer, _layer, name), view = next( + (key, view) + for key, view in mgr._peer_views.items() + if view.data_ptr() <= src and src + size <= view.data_ptr() + view.nbytes + ) + # Row i of peer p's view is expert start_p + i of the same weight. + dst_offset = dst - backings[WEIGHT_NAMES.index(name)].data_ptr() + src_offset = src - view.data_ptr() + assert peer != rank + assert dst_offset == src_offset + partition.range_of(peer)[0] * EXPERT_BYTES + assert sum(flat[2]) == len(WEIGHT_NAMES) * (896 - (end - start)) * EXPERT_BYTES + assert _overlap_count(flat) == 0 + assert _canonical(_plan(mgr, wm.MEMCPY_BATCH_SLICE_BYTES, monkeypatch)) == _canonical(flat) + _dst, _src, lines, row_starts, _ = fabric_gather.build_plan(groups, torch.device("cpu")) + assert row_starts.tolist() == [0, *itertools.accumulate(len(group) for group in groups)] + assert lines.sum().item() * fabric_gather.LINE_BYTES == sum(flat[2]) diff --git a/tests/unittest/_torch/moe/test_dwdp_fixup_moe_backends.py b/tests/unittest/_torch/moe/test_dwdp_fixup_moe_backends.py index a845ab192455..b14e997778cc 100644 --- a/tests/unittest/_torch/moe/test_dwdp_fixup_moe_backends.py +++ b/tests/unittest/_torch/moe/test_dwdp_fixup_moe_backends.py @@ -12,20 +12,11 @@ # 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. -"""Unit tests for the MPI-allgather refactor of DWDP's ``fixup_moe_backends``. +"""Unit tests for the small-parameter exchange of DWDP's setup and ``fixup_moe_backends``. -These tests verify that the MPI-based allgather (replacing tekit's TCPDWDPStore) -produces the same semantics as the tekit original: - -- Each rank's local shard is contributed via comm.allgather. -- Ranks receive a list of shards (one per peer, indexed by rank). -- The function concatenates along dim 0 and updates the parameter in place. - -We mock the MPI communicator since the tests run single-process on one GPU. -The mock's ``allgather`` returns a pre-built shard list that the test provides -(simulating "what this rank would receive if 4 ranks each contributed their -shard"). This lets us cover the multi-rank semantics end-to-end without a -real MPI environment. +``_exchange_expert_scalars`` gathers every rank's small per-expert parameters and gate bias in one +``comm.allgather`` and puts them together along the expert partition; ``fixup_moe_backends`` +installs them. A mock comm returns every simulated rank's report. """ import unittest @@ -36,79 +27,105 @@ import torch.nn as nn from tensorrt_llm._torch.modules.dwdp.setup import ( - _allgather_e_score_correction_bias, - _allgather_expert_scales, + _collect_expert_scalars, + _exchange_expert_scalars, _scatter_shards_to_full, + fixup_moe_backends, ) +from tensorrt_llm._torch.modules.dwdp.specs import ExpertPartition +from tensorrt_llm._torch.moe.fused_moe.quantization import NVFP4CuteDslFusedMoEMethod pytestmark = pytest.mark.cpu_only +LAYER = 3 + + +class _Experts(nn.Module): + """An MoE backend: an expert weight plus the given small parameters.""" + + def __init__(self, num_experts, rows, params): + super().__init__() + self.num_experts = num_experts + self.w3_w1_weight = nn.Parameter(torch.zeros(rows, 4), requires_grad=False) + for name, tensor in params.items(): + setattr(self, name, nn.Parameter(tensor, requires_grad=False)) + + +def _make_model(experts, bias=None, bias_as_parameter=True): + """A model whose layer ``LAYER`` is an MoE layer; its gate holds ``bias`` when given one.""" + model = nn.Module() + model.layers = nn.ModuleList([nn.Module() for _ in range(LAYER + 1)]) + mlp = model.layers[LAYER].mlp = nn.Module() + mlp.experts, mlp.gate = experts, nn.Module() + if bias is not None and bias_as_parameter: + bias = nn.Parameter(bias, requires_grad=False) + if bias is not None: + mlp.gate.e_score_correction_bias = bias + return model + + +def _make_rank_models(partition, params_per_rank, biases=None, bias_as_parameter=True): + """One model per DWDP rank, holding that rank's shards.""" + models = [] + for rank in range(partition.dwdp_size): + experts = _Experts(partition.num_experts, partition.count_of(rank), params_per_rank[rank]) + bias = None if biases is None else biases[rank] + models.append(_make_model(experts, bias, bias_as_parameter)) + return models + -def _make_mock_comm(all_shards): - """Create a mock MPI comm whose allgather() returns ``all_shards``. +def _exchange_and_fixup(models, rank, partition): + """Run the exchange and the fixup on ``rank``; return its comm and model. - ``all_shards`` is a list of torch.Tensor objects, one per simulated rank. - The mock ignores the input argument and returns the list, which is the - semantic shape of a real ``Intracomm.allgather(local_shard)`` call. + The comm's allgather returns every rank's report, this rank's as passed in. """ + reports = [(None, _collect_expert_scalars(m, [LAYER], partition.num_experts)) for m in models] comm = MagicMock() - comm.allgather = MagicMock(side_effect=lambda _local: all_shards) - comm.Barrier = MagicMock() - comm.Get_rank = MagicMock(return_value=0) - comm.Get_size = MagicMock(return_value=len(all_shards)) - return comm + comm.allgather.side_effect = lambda local: [ + local if r == rank else report for r, report in enumerate(reports) + ] + model = models[rank] + scalars = _exchange_expert_scalars(comm, model, [LAYER], partition) + fixup_moe_backends(model, [LAYER], partition.num_experts, scalars) + return comm, model class TestAllgatherEScoreCorrectionBias(unittest.TestCase): - def _make_gate(self, bias_tensor, as_parameter=True): - """Minimal nn.Module with an e_score_correction_bias attribute.""" - gate = nn.Module() - if as_parameter: - gate.e_score_correction_bias = nn.Parameter(bias_tensor, requires_grad=False) - else: - gate.e_score_correction_bias = bias_tensor - return gate - def test_missing_bias_is_noop(self): - gate = nn.Module() - comm = _make_mock_comm([torch.zeros(2)]) - _allgather_e_score_correction_bias( - gate, layer_idx=3, dwdp_rank=0, dwdp_size=4, num_experts_total=16, comm=comm - ) - comm.allgather.assert_not_called() - self.assertFalse( - hasattr(gate, "e_score_correction_bias") - and getattr(gate, "e_score_correction_bias") is not None - and isinstance(getattr(gate, "e_score_correction_bias"), (torch.Tensor, nn.Parameter)) - ) + partition = ExpertPartition.from_counts([2, 2, 2, 2]) + params = [{"fc31_alpha": torch.ones(2)} for _ in range(4)] + models = _make_rank_models(partition, params) + + comm, model = _exchange_and_fixup(models, 0, partition) + + # One exchange for every layer, also without a bias. + comm.allgather.assert_called_once() + self.assertFalse(hasattr(model.layers[LAYER].mlp.gate, "e_score_correction_bias")) def test_allgather_sharded_parameter(self): """Sharded bias: allgather 4 shards, concat, replace nn.Parameter.""" dwdp_size = 4 shard_size = 2 num_experts_total = shard_size * dwdp_size + partition = ExpertPartition.strided( + num_experts=num_experts_total, dwdp_size=dwdp_size, size=shard_size, stride=shard_size + ) # Each rank's shard has unique values shards = [ torch.arange(r * shard_size, (r + 1) * shard_size, dtype=torch.float32) for r in range(dwdp_size) ] expected = torch.cat(shards, dim=0) + params = [{"fc31_alpha": torch.ones(shard_size)} for _ in range(dwdp_size)] for dwdp_rank in range(dwdp_size): with self.subTest(dwdp_rank=dwdp_rank): - gate = self._make_gate(shards[dwdp_rank].clone()) - comm = _make_mock_comm(shards) - - _allgather_e_score_correction_bias( - gate, - layer_idx=3, - dwdp_rank=dwdp_rank, - dwdp_size=dwdp_size, - num_experts_total=num_experts_total, - comm=comm, - ) + models = _make_rank_models(partition, params, [s.clone() for s in shards]) + + comm, model = _exchange_and_fixup(models, dwdp_rank, partition) comm.allgather.assert_called_once() + gate = model.layers[LAYER].mlp.gate self.assertIsInstance(gate.e_score_correction_bias, nn.Parameter) self.assertEqual(gate.e_score_correction_bias.shape, (num_experts_total,)) torch.testing.assert_close(gate.e_score_correction_bias.data, expected) @@ -118,108 +135,70 @@ def test_allgather_sharded_plain_tensor(self): dwdp_size = 2 shard_size = 3 num_experts_total = shard_size * dwdp_size + partition = ExpertPartition.strided( + num_experts=num_experts_total, dwdp_size=dwdp_size, size=shard_size, stride=shard_size + ) shards = [ torch.arange(r * shard_size, (r + 1) * shard_size, dtype=torch.float32) for r in range(dwdp_size) ] - gate = self._make_gate(shards[0].clone(), as_parameter=False) - comm = _make_mock_comm(shards) - - _allgather_e_score_correction_bias( - gate, - layer_idx=3, - dwdp_rank=0, - dwdp_size=dwdp_size, - num_experts_total=num_experts_total, - comm=comm, + params = [{"fc31_alpha": torch.ones(shard_size)} for _ in range(dwdp_size)] + models = _make_rank_models( + partition, params, [s.clone() for s in shards], bias_as_parameter=False ) + comm, model = _exchange_and_fixup(models, 0, partition) + comm.allgather.assert_called_once() - self.assertEqual(gate.e_score_correction_bias.shape, (num_experts_total,)) - torch.testing.assert_close(gate.e_score_correction_bias, torch.cat(shards, dim=0)) + bias = model.layers[LAYER].mlp.gate.e_score_correction_bias + self.assertNotIsInstance(bias, nn.Parameter) + self.assertEqual(bias.shape, (num_experts_total,)) + torch.testing.assert_close(bias, torch.cat(shards, dim=0)) class TestAllgatherExpertScales(unittest.TestCase): - def _make_experts_module(self, params): - """Return an nn.Module with the given {name: tensor} parameters.""" - module = nn.Module() - for name, tensor in params.items(): - setattr(module, name, nn.Parameter(tensor, requires_grad=False)) - return module - def test_skip_non_scale_params(self): """Parameters without scale keywords are not touched.""" experts_per_rank = 2 - module = self._make_experts_module( - { - "random_weight": torch.zeros(experts_per_rank, 4), - "fc31_alpha": torch.tensor([1.0, 2.0]), - } - ) + partition = ExpertPartition.from_counts([experts_per_rank, experts_per_rank]) all_alpha_shards = [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])] - comm = _make_mock_comm(all_alpha_shards) - - _allgather_expert_scales( - module, - layer_idx=3, - dwdp_rank=0, - dwdp_size=2, - comm=comm, - experts_per_rank=experts_per_rank, - ) + params = [ + {"random_weight": torch.zeros(experts_per_rank, 4), "fc31_alpha": shard.clone()} + for shard in all_alpha_shards + ] + models = _make_rank_models(partition, params) + + comm, model = _exchange_and_fixup(models, 0, partition) # Called exactly once — only fc31_alpha matched. - self.assertEqual(comm.allgather.call_count, 1) + comm.allgather.assert_called_once() + _status, reported = comm.allgather.call_args.args[0] + self.assertEqual(sorted(reported[LAYER][0]), ["fc31_alpha"]) + experts = model.layers[LAYER].mlp.experts # random_weight unchanged (still size (2,4)) - self.assertEqual(module.random_weight.shape, (experts_per_rank, 4)) + self.assertEqual(experts.random_weight.shape, (experts_per_rank, 4)) # fc31_alpha concatenated to full size (4,) - self.assertEqual(module.fc31_alpha.shape, (4,)) - torch.testing.assert_close(module.fc31_alpha.data, torch.cat(all_alpha_shards, dim=0)) + self.assertEqual(experts.fc31_alpha.shape, (4,)) + torch.testing.assert_close(experts.fc31_alpha.data, torch.cat(all_alpha_shards, dim=0)) def test_allgather_multiple_scales(self): - """Two matching scale params: each gets its own allgather call.""" - experts_per_rank = 2 - dwdp_size = 2 + """Two matching scale params, one with a trailing dim: both put together.""" + partition = ExpertPartition.from_counts([2, 2]) alpha_shards = [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])] - beta_shards = [torch.tensor([[5.0]]), torch.tensor([[6.0]])] - # fc31_alpha has shape (2,); fc2_alpha has shape (2,1) — both match - # experts_per_rank=2 on dim 0. - module = self._make_experts_module( - { - "fc31_alpha": alpha_shards[0].clone(), - "fc2_alpha": torch.tensor([[5.0]]), # shape (1,1) — shape[0]=1 != 2 - } - ) - # Rewrite fc2_alpha with shape[0]=2 so it matches - module.fc2_alpha = nn.Parameter(torch.tensor([[1.0], [2.0]]), requires_grad=False) - # Two shards, each shape (2,1) + # fc2_alpha has shape (2, 1) — dim 0 is the rank's expert count too. beta_shards = [torch.tensor([[1.0], [2.0]]), torch.tensor([[3.0], [4.0]])] + params = [ + {"fc31_alpha": alpha.clone(), "fc2_alpha": beta.clone()} + for alpha, beta in zip(alpha_shards, beta_shards) + ] + models = _make_rank_models(partition, params) - # The mock returns the pre-set shard list for whichever param is being - # allgathered. named_parameters() order is insertion order — the - # ordering here is fc31_alpha first, fc2_alpha second. - call_idx = {"n": 0} - - def _side_effect(_local): - idx = call_idx["n"] - call_idx["n"] += 1 - return alpha_shards if idx == 0 else beta_shards - - comm = MagicMock() - comm.allgather = MagicMock(side_effect=_side_effect) - - _allgather_expert_scales( - module, - layer_idx=3, - dwdp_rank=0, - dwdp_size=dwdp_size, - comm=comm, - experts_per_rank=experts_per_rank, - ) + comm, model = _exchange_and_fixup(models, 0, partition) - self.assertEqual(comm.allgather.call_count, 2) - torch.testing.assert_close(module.fc31_alpha.data, torch.cat(alpha_shards, dim=0)) - torch.testing.assert_close(module.fc2_alpha.data, torch.cat(beta_shards, dim=0)) + self.assertEqual(comm.allgather.call_count, 1) + experts = model.layers[LAYER].mlp.experts + torch.testing.assert_close(experts.fc31_alpha.data, torch.cat(alpha_shards, dim=0)) + torch.testing.assert_close(experts.fc2_alpha.data, torch.cat(beta_shards, dim=0)) class TestScatterShardsToFull(unittest.TestCase): @@ -292,16 +271,9 @@ def test_higher_dim_trailing_shape(self): class TestAllgatherExpertScalesNonUniform(unittest.TestCase): - """Validate ``_allgather_expert_scales`` end-to-end with the new - Phase 2 ``peer_ranges`` argument: shards loaded under Mode B - overlap (size > stride) must reconstruct the full bias correctly, - with overlapping experts getting the same value from any peer.""" - - def _make_experts_module(self, params): - module = torch.nn.Module() - for name, tensor in params.items(): - setattr(module, name, torch.nn.Parameter(tensor, requires_grad=False)) - return module + """Validate ``_exchange_expert_scalars`` end-to-end on a redundant partition: shards loaded + under Mode B overlap (size > stride) must reconstruct the full parameter correctly, with + overlapping experts getting the same value from any peer.""" def test_mode_b_overlap_dwdp3_via_peer_ranges(self): # dwdp=3, 256 experts, Mode B with size=86, stride=85 @@ -309,37 +281,97 @@ def test_mode_b_overlap_dwdp3_via_peer_ranges(self): # adjacent ranks). Shards loaded by each rank cover overlapping # ranges; ``_scatter_shards_to_full`` writes the shared experts # multiple times but the values agree (same checkpoint). - peer_ranges = [(0, 86), (85, 171), (170, 256)] - all_shards = [ - torch.arange(0, 86, dtype=torch.float32), # rank 0: [0..85] - torch.arange(85, 171, dtype=torch.float32), # rank 1: [85..170] - torch.arange(170, 256, dtype=torch.float32), # rank 2: [170..255] + partition = ExpertPartition.strided(num_experts=256, dwdp_size=3, size=86, stride=85) + self.assertEqual(partition.peer_ranges, [(0, 86), (85, 171), (170, 256)]) + params = [ + {"fc31_alpha": torch.arange(start, end, dtype=torch.float32)} + for start, end in partition.ranges ] - module = self._make_experts_module( - { - "fc31_alpha": torch.arange(0, 86, dtype=torch.float32), - } - ) - comm = _make_mock_comm(all_shards) - - _allgather_expert_scales( - module, - layer_idx=3, - dwdp_rank=0, - dwdp_size=3, - comm=comm, - experts_per_rank=86, - num_experts_total=256, - peer_ranges=peer_ranges, - ) + models = _make_rank_models(partition, params) + + comm, model = _exchange_and_fixup(models, 0, partition) comm.allgather.assert_called_once() # Reconstructed alpha covers ALL 256 experts with their canonical values. torch.testing.assert_close( - module.fc31_alpha.data, + model.layers[LAYER].mlp.experts.fc31_alpha.data, torch.arange(0, 256, dtype=torch.float32), ) +def _alphas(partition): + """Each rank's ``fc31_alpha`` shard: the ids of the experts it stores.""" + return [{"fc31_alpha": torch.arange(s, e, dtype=torch.float32)} for s, e in partition.ranges] + + +class TestExchangeRefusesReportsThatDoNotFit(unittest.TestCase): + def test_every_rank_raises_the_same_error(self): + """The verdict is made from the gathered reports, the same on every rank.""" + uneven = ExpertPartition.from_counts([3, 2]) + # Experts 2-3 are stored on ranks 0 and 1, experts 4-5 on ranks 1 and 2. + overlap = ExpertPartition.strided(num_experts=8, dwdp_size=3, size=4, stride=2) + two_copies = _alphas(overlap) + two_copies[1]["fc31_alpha"][0] = -1.0 + alphas = _alphas(uneven) + scales = [dict(a, fc31_input_scale=torch.tensor(s)) for a, s in zip(alphas, (1.0, 0.5))] + other_names = [alphas[0], dict(alphas[1], fc2_alpha=torch.ones(2))] + bias = "e_score_correction_bias" + cases = { + # case: (partition, small parameters per rank, bias per rank, offending name) + "two values of an overlapping expert": (overlap, two_copies, None, "fc31_alpha"), + "input scales that differ": (uneven, scales, None, "fc31_input_scale"), + "other parameter names": (uneven, other_names, None, "fc2_alpha"), + "a bias neither full nor per expert": (uneven, alphas, [torch.ones(3)] * 2, bias), + "a bias on one rank only": (uneven, alphas, [torch.ones(3), None], bias), + } + for case, (partition, params, biases, name) in cases.items(): + with self.subTest(case): + messages = set() + for rank in range(partition.dwdp_size): + models = _make_rank_models(partition, params, biases) + with self.assertRaisesRegex(RuntimeError, name) as caught: + _exchange_and_fixup(models, rank, partition) + messages.add(str(caught.exception)) + self.assertEqual(len(messages), 1, messages) + + +class _ConfigurableMoE(nn.Module): + """ConfigurableMoE's shape: the backend holds the weights, the wrapper delegates to it.""" + + def __init__(self, backend): + super().__init__() + self.backend = backend + + @property + def w3_w1_weight(self): + return self.backend.w3_w1_weight + + +class TestFixupConfigurableMoE(unittest.TestCase): + def test_both_modules_see_every_expert_and_the_quant_scales_the_full_alphas(self): + """K3's routed experts: a CuteDSL backend in a wrapper; the gate bias is full-sized.""" + partition = ExpertPartition.balanced(num_experts=7, dwdp_size=2) # 4 + 3 + bias, models = torch.linspace(-1.0, 1.0, 7), [] + for start, end in partition.ranges: + alpha = torch.arange(start, end, dtype=torch.float32) + block = torch.zeros(end - start, 1) + params = dict(fc31_alpha=alpha, fc2_alpha=alpha + 0.5, w3_w1_weight_scale=block) + params.update(w2_weight_scale=block.clone(), fc31_input_scale=torch.tensor(1.0)) + params.update(fc2_input_scale=torch.tensor(1.0)) + backend = _Experts(partition.num_experts, end - start, params) + NVFP4CuteDslFusedMoEMethod().setup_quant_scales(backend) + models.append(_make_model(_ConfigurableMoE(backend), bias.clone())) + + _, model = _exchange_and_fixup(models, 1, partition) + + wrapper = model.layers[LAYER].mlp.experts + self.assertEqual(wrapper.expert_size_per_partition, 7) + self.assertEqual(wrapper.backend.expert_size_per_partition, 7) + scales = wrapper.backend.quant_scales + torch.testing.assert_close(scales.fc1_global, torch.arange(7.0)) + torch.testing.assert_close(scales.fc2_global, torch.arange(7.0) + 0.5) + torch.testing.assert_close(model.layers[LAYER].mlp.gate.e_score_correction_bias.data, bias) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index ad1e6829630e..756c74feddfe 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -646,6 +646,7 @@ def test_shutdown_drains_save_before_loop_exit( executor.virtual_memory_pools = {} executor.sampler = None executor.dwdp_manager = None + executor.dense_dwdp_manager = None executor.resource_manager = MagicMock() resource = MagicMock() executor.resource_manager.resource_managers = {"kv": resource} diff --git a/tests/unittest/executor/test_forward_error_survival.py b/tests/unittest/executor/test_forward_error_survival.py index b37aa5edc235..e15d934f9714 100644 --- a/tests/unittest/executor/test_forward_error_survival.py +++ b/tests/unittest/executor/test_forward_error_survival.py @@ -110,6 +110,7 @@ def _loop_executor(monkeypatch: pytest.MonkeyPatch) -> PyExecutor: executor.guided_decoder = None executor.drafter = None executor.dwdp_manager = None + executor.dense_dwdp_manager = None executor.speculation_gate = None executor.model_engine = SimpleNamespace( attn_metadata=None, kv_cache_dtype_byte_size=None, spec_metadata=None diff --git a/tests/unittest/llmapi/apps/test_openai_server_graceful_shutdown.py b/tests/unittest/llmapi/apps/test_openai_server_graceful_shutdown.py new file mode 100644 index 000000000000..b209d453c220 --- /dev/null +++ b/tests/unittest/llmapi/apps/test_openai_server_graceful_shutdown.py @@ -0,0 +1,217 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The graceful-shutdown bound that ``OpenAIServer`` takes from its environment. + +``TRTLLM_SERVE_GRACEFUL_SHUTDOWN_TIMEOUT_S`` bounds how long uvicorn waits for running request +handlers after a shutdown signal. The ``uvicorn.Config`` call is checked in this process. The +shutdown itself runs the real ``OpenAIServer.__call__`` under uvloop, as ``trtllm-serve`` does, in +a fresh interpreter that receives SIGINT. Its ASGI app has a request handler that never finishes, +like a handler waiting for the final response of a request that the executor dropped. +""" + +import asyncio +import inspect +import os +import signal +import socket +import subprocess +import sys +import time +from pathlib import Path +from tempfile import TemporaryDirectory + +import pytest +import uvicorn + +from tensorrt_llm import visual_gen +from tensorrt_llm.commands import serve +from tensorrt_llm.serve.openai_server import OpenAIServer + +# The CPU stage collects with `-m cpu_only`; unittest/conftest.py also skips collecting any file +# that does not contain the cpu_only marker. +pytestmark = pytest.mark.cpu_only + +_TIMEOUT_ENV = "TRTLLM_SERVE_GRACEFUL_SHUTDOWN_TIMEOUT_S" +# The parameters of the installed uvicorn.Config, read before any test replaces it. +_CONFIG_SIGNATURE = inspect.signature(uvicorn.Config) +# Graceful-shutdown bound of the child server, in seconds. +_BOUND_S = 1 +# Allowance over the bound for the lifespan shutdown to start after SIGINT, in seconds: generous +# for a loaded CI host, yet it fails a shutdown that cancels the handler well past the bound. +_SLACK_S = 10 + + +class _ConfigCalled(Exception): + """Stops ``OpenAIServer.__call__`` at its ``uvicorn.Config`` call.""" + + +def _config_keywords(monkeypatch: pytest.MonkeyPatch, value: str | None) -> dict: + """Return the keywords that ``OpenAIServer.__call__`` passes to ``uvicorn.Config``.""" + if value is None: + monkeypatch.delenv(_TIMEOUT_ENV, raising=False) + else: + monkeypatch.setenv(_TIMEOUT_ENV, value) + + def config(app, **kwargs): + # Binding raises TypeError for a keyword that the installed uvicorn does not accept. + _CONFIG_SIGNATURE.bind(app, **kwargs) + raise _ConfigCalled(kwargs) + + monkeypatch.setattr(uvicorn, "Config", config) + server = object.__new__(OpenAIServer) + server.app = None + with pytest.raises(_ConfigCalled) as called: + asyncio.run(server("localhost", 0)) + return called.value.args[0] + + +@pytest.mark.parametrize( + ("value", "seconds"), [(None, None), ("", None), ("45", 45)], ids=["unset", "empty", "45"] +) +def test_timeout_is_passed_to_uvicorn_only_when_set( + monkeypatch: pytest.MonkeyPatch, value: str | None, seconds: int | None +) -> None: + """Unset or empty adds no keyword; a positive integer adds ``timeout_graceful_shutdown`` only.""" + unset = _config_keywords(monkeypatch, None) + assert "timeout_graceful_shutdown" not in unset + + expected = unset if seconds is None else {**unset, "timeout_graceful_shutdown": seconds} + assert _config_keywords(monkeypatch, value) == expected + + +@pytest.mark.parametrize("value", ["1.5", "0", " "], ids=["fraction", "zero", "blank"]) +def test_malformed_timeout_fails_before_uvicorn_config( + monkeypatch: pytest.MonkeyPatch, value: str +) -> None: + """A malformed value raises a ValueError that names the variable and the value.""" + with pytest.raises(ValueError) as error: + _config_keywords(monkeypatch, value) + + assert _TIMEOUT_ENV in str(error.value) and repr(value) in str(error.value) + + +# _detect_external_launch() without an external launcher, and on worker rank 1 of 2. +@pytest.mark.parametrize( + "external_launch", [None, (1, 1, 2, "127.0.0.1", 29500)], ids=["rank0", "worker_rank"] +) +def test_malformed_timeout_fails_before_visual_gen_is_constructed( + monkeypatch: pytest.MonkeyPatch, external_launch: tuple | None +) -> None: + """The VisualGen launcher rejects a malformed value before any rank constructs ``VisualGen``.""" + monkeypatch.setenv(_TIMEOUT_ENV, "1.5") + monkeypatch.setattr( + "tensorrt_llm._torch.visual_gen.executor._detect_external_launch", lambda: external_launch + ) + + def construct(*args, **kwargs) -> None: + pytest.fail("VisualGen was constructed before the malformed value was rejected") + + monkeypatch.setattr(visual_gen, "VisualGen", construct) + with pytest.raises(ValueError, match=_TIMEOUT_ENV): + serve.launch_visual_gen_server("127.0.0.1", 0, "test-model") + + +def _wait_for(path: Path, child: subprocess.Popen, timeout: float) -> None: + """Wait for a marker that the child creates, failing early if the child exits first.""" + deadline = time.monotonic() + timeout + while not path.exists(): + if (child.poll() is not None or time.monotonic() > deadline) and not path.exists(): + pytest.fail(f"No {path.name}; child exit code {child.poll()}") + time.sleep(0.01) + + +def _serve_never_finishing_handler(directory: str) -> None: + """Serve a request handler that never finishes through ``OpenAIServer.__call__``.""" + import uvloop + + workdir = Path(directory) + + async def app(scope, receive, send) -> None: + if scope["type"] == "lifespan": + while True: + message = await receive() + if message["type"] == "lifespan.startup": + (workdir / "ready").touch() + await send({"type": "lifespan.startup.complete"}) + elif message["type"] == "lifespan.shutdown": + (workdir / "lifespan_shutdown").touch() + await send({"type": "lifespan.shutdown.complete"}) + return + (workdir / "handler_started").touch() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + (workdir / "handler_cancelled").touch() + raise + + server = object.__new__(OpenAIServer) + server.app = app + server.disagg_cluster_worker = None + listener = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + listener.bind(str(workdir / "http.sock")) + listener.listen() + # The same event loop as trtllm-serve. + uvloop.run(server("localhost", 0, sockets=[listener])) + + +@pytest.mark.skipif(os.name != "posix", reason="The test verifies POSIX SIGINT semantics.") +def test_bounded_shutdown_cancels_handler_then_runs_lifespan_shutdown() -> None: + """With the bound, a stuck handler is cancelled and the lifespan shutdown still runs.""" + env = {k: v for k, v in os.environ.items() if k != "TRTLLM_USAGE_FORCE_ENABLED"} + env[_TIMEOUT_ENV] = str(_BOUND_S) + # Info logs, unbuffered so the log is complete while the child runs, and no telemetry. + env.update(TLLM_LOG_LEVEL="INFO", PYTHONUNBUFFERED="1", TRTLLM_NO_USAGE_STATS="1") + # Runs this file as __main__. Unlike `python `, it keeps this directory, whose + # helper modules (utils.py, openai_server.py) would shadow installed ones, off sys.path. + run_as_main = ( + f"import runpy; runpy.run_path({str(Path(__file__).resolve())!r}, run_name='__main__')" + ) + # A short directory keeps the socket path under the AF_UNIX pathname limit. + with TemporaryDirectory(prefix="trtllm-gs-") as directory: + workdir = Path(directory) + with (workdir / "child.log").open("wb") as log_file: + child = subprocess.Popen( + [sys.executable, "-c", run_as_main, directory], + cwd=directory, + env=env, + stdout=log_file, + stderr=subprocess.STDOUT, + ) + try: + # The first wait covers the full tensorrt_llm import in the child. + _wait_for(workdir / "ready", child, 300) + # The client leaves once its handler runs, as the disaggregated router does at + # shutdown, so uvicorn then waits on the running handler, not on a connection. + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as client: + client.connect(str(workdir / "http.sock")) + client.sendall(b"GET /never-finishes HTTP/1.1\r\nHost: localhost\r\n\r\n") + _wait_for(workdir / "handler_started", child, 60) + signalled = time.monotonic() + child.send_signal(signal.SIGINT) + _wait_for(workdir / "lifespan_shutdown", child, 60) + lifespan_shutdown_after = time.monotonic() - signalled + # The exit code comes from interpreter teardown after uvicorn re-raises SIGINT, so + # only the exit itself is checked. + child.wait(timeout=60) + finally: + if child.poll() is None: + child.kill() + child.wait() + log = (workdir / "child.log").read_text(errors="replace") + print(log) # Reported with any failure. + assert _BOUND_S <= lifespan_shutdown_after < _BOUND_S + _SLACK_S + assert (workdir / "handler_cancelled").exists() + assert f"{_TIMEOUT_ENV}={_BOUND_S}:" in log + # uvicorn waits on the handler, cancels it at the bound, then runs the lifespan shutdown. + stages = [ + "Waiting for background tasks to complete", + "Cancel 1 running task(s), timeout graceful shutdown exceeded", + "Waiting for application shutdown", + "Finished server process", + ] + positions = [log.find(stage) for stage in stages] + assert positions[0] >= 0 and positions == sorted(positions) + + +if __name__ == "__main__": + _serve_never_finishing_handler(sys.argv[1]) diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 38cf3edc82ba..9da7c021e197 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -145,6 +145,13 @@ def schemas(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]: (generator.SERVE_SCHEMA, {"kv_cache_config": {"enable_mix_precision": "nvfp4"}}, False), (generator.SERVE_SCHEMA, {"checkpoint_loader": {}}, False), (generator.SERVE_SCHEMA, {"tokenizer": 42}, False), + ( + generator.SERVE_SCHEMA, + {"dwdp_config": {"num_experts_per_worker": 0, "num_prefetch_experts": 0}}, + True, + ), + (generator.SERVE_SCHEMA, {"dwdp_config": {"num_experts_per_worker": -1}}, False), + (generator.SERVE_SCHEMA, {"dwdp_config": {"num_prefetch_experts": -1}}, False), (generator.DISAGG_SCHEMA, {"backend": "pytorch"}, True), (generator.DISAGG_SCHEMA, {"backend": "_autodeploy"}, False), (generator.DISAGG_SCHEMA, {"context_servers": {"urls": ["ctx:8001"]}}, True), @@ -156,6 +163,11 @@ def schemas(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]: True, ), (generator.DISAGG_SCHEMA, {"context_servers": {"env_overrides": {"FLAG": 1}}}, True), + ( + generator.DISAGG_SCHEMA, + {"context_servers": {"dwdp_config": {"dwdp_size": 12, "num_experts_per_worker": -1}}}, + False, + ), ( generator.DISAGG_SCHEMA, {"context_servers": {"moe_config": {"backend_typo": "CUTLASS"}}}, @@ -484,6 +496,39 @@ def test_disagg_node_id_matches_runtime(schemas: dict[str, dict]) -> None: assert not validator.is_valid({"node_id": value}) +@pytest.mark.parametrize( + ("partition", "valid"), + [ + ({}, True), + ({"num_experts_per_worker": 0, "num_prefetch_experts": 0}, True), + ({"num_experts_per_worker": 1, "num_prefetch_experts": 1}, True), + ({"num_experts_per_worker": 2, "num_prefetch_experts": 1}, True), + ({"num_experts_per_worker": 1}, False), + ({"num_prefetch_experts": 1}, False), + ({"num_experts_per_worker": 0, "num_prefetch_experts": 1}, False), + ({"num_experts_per_worker": 1, "num_prefetch_experts": 0}, False), + ], +) +def test_dwdp_expert_partition_matches_runtime( + schemas: dict[str, dict], partition: dict[str, int], valid: bool +) -> None: + from tensorrt_llm.llmapi.llm_args import DwdpConfig + + if valid: + DwdpConfig(**partition) + else: + with pytest.raises(ValueError, match="set both"): + DwdpConfig(**partition) + + config = {"dwdp_config": partition} + for filename in (generator.SERVE_SCHEMA, generator.DISAGG_SCHEMA): + validator = Draft202012Validator(schemas[filename]) + assert validator.is_valid(config) == valid, filename + if filename == generator.DISAGG_SCHEMA: + for section in ("context_servers", "generation_servers"): + assert validator.is_valid({section: config}) == valid, section + + def test_runtime_only_union_does_not_disable_validation() -> None: class RuntimeObject: pass