From d31422ba256443240941017b054602b806e4db57 Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:27:28 -0700 Subject: [PATCH 1/9] [None][fix] Bound the context-only server shutdown: drop in-flight requests and time out the graceful stop (opt-in) Context-only benchmark jobs can finish the benchmark and then hang at teardown: requests orphaned on the context ranks meet three unbounded waits on the stop path (uvicorn waiting for handlers, a handler's aresult() with no timeout, PyExecutor draining every request). Both changes are off unless their variable is set, and behave as before when it is not. - TRTLLM_EXECUTOR_SHUTDOWN_DROP_INFLIGHT=1 (py_executor.py): once shutdown is requested, the executor stops with requests still active or queued instead of draining them. A pending KV connector load or transfer (_has_pending_connector_transfers) is still waited for, since it reads the KV pools that shutdown frees next. After the loop ends normally, after its last collective, the dropped requests are released through the existing termination path; a request whose KV cache the transfer fabric may still read keeps it, as on cancellation, and the KV cache manager closes that cache at its own shutdown. The release is gated on the loop's completion mark, so the loop-ending break still directly follows the mark and the insufficient-KV benchmark exit releases nothing. Effective only with pp_size == 1 and the overlap scheduler disabled; the effective value is logged once. Meant for context-only servers: do not set it on generation servers. - TRTLLM_SERVE_GRACEFUL_SHUTDOWN_TIMEOUT_S= (openai_server.py): passes uvicorn's timeout_graceful_shutdown, so handlers that never finish are cancelled and the lifespan shutdown (which sends the engine its quit sentinel) still runs. launch_server and launch_visual_gen_server read it first, the latter before any rank builds VisualGen, so a malformed value fails before the model loads. Tests (CPU): - test_executor_shutdown_drop_inflight.py: the gate (one pipeline stage, no overlap scheduler, logged once), should_stop_processing set and unset, with and without pending connector transfers, the release through the termination path, and the executor loop releasing after its last collective. - test_openai_server_graceful_shutdown.py: the uvicorn.Config keywords for set, unset and malformed values; a malformed value rejected before the VisualGen launcher constructs VisualGen, on rank 0 and on a worker rank; and a SIGINT child server whose hung handler is cancelled at the bound. Registered in l0_cpu. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 104 +++++++- tensorrt_llm/commands/serve.py | 10 + tensorrt_llm/serve/openai_server.py | 44 +++- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_executor_shutdown_drop_inflight.py | 230 ++++++++++++++++++ .../test_openai_server_graceful_shutdown.py | 217 +++++++++++++++++ 6 files changed, 603 insertions(+), 3 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_executor_shutdown_drop_inflight.py create mode 100644 tests/unittest/llmapi/apps/test_openai_server_graceful_shutdown.py diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a0e323196a53..1375706cdfc3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -206,6 +206,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). @@ -998,6 +1045,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 +1058,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 @@ -2320,8 +2374,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) @@ -5527,6 +5621,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/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/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/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/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/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]) From fab4f23fd8d51f892a6014b08ac291e09f2ac451 Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:27:40 -0700 Subject: [PATCH 2/9] [None][feat] Dense DWDP for Kimi-K3 context servers (env form) Dense DWDP row-shards the replicated dense weights of Kimi-K3's MoE decoder layers (attention, shared expert, latent projections and norm, router; layer 0 stays replicated) over the DWDP group, and gathers each layer's shard set two layers ahead on the expert DWDP copy stream into small staging slots. On a context-only rank this frees about 56 GB for the KV cache. It is opt-in through environment switches (TRTLLM_DWDP_DENSE, _SEGMENTS, _REQUIRED, _MIN_BYTES, _STORE_CHUNK_MB, _DEBUG, _STRICT), and its copies use the expert transfer: trtllm::dwdp_fabric_gather when TRTLLM_DWDP_PREFETCH_GATHER is on (kernel chosen by TRTLLM_DWDP_GATHER_KERNEL), otherwise the copy engines. - modules/dwdp/dense (new): the switches, the collector and blob layout, the row store (FABRIC handles, plus an in-process loopback for tests), the per-layer copy plans and launcher, repointing of the module tensors into the slots, the runtime (slot rings, RAW/WAR events, demand fetch and fences, the WAR barrier after an aborted forward), the DEBUG=2 slot watchdog, and setup: phases P1-P6 over five fixed collectives, so every rank reaches the same verdict, with REQUIRED deciding between raising everywhere and keeping the weights replicated. - The runtime's copy schedule. Once a layer's attention (H2) or MoE block (H3) has closed its read window, the runtime enqueues, in layer order, each next layer whose staging slots are free (empty, or their holder's window closed), so these copies never need the fence. Within a forward that is L(l+2) at H2 of layer l. At the end of a forward it goes on into the first two layers of the next one, as the expert prefetch schedule does with its last MoE layers, so a completed forward leaves the next one primed. With two expert prefetch slots each L(l) is enqueued right before E(l): H2 of layer l-2 enqueues L(l) before that layer's MoE finishes and issues E(l). For Kimi-K3 the MoE-dense slots of layers 1 and 2 are read by layers 91 and 92 until H3, so the next forward's copies are E(1), L(1), E(2), L(2) (L before E with the attention segment alone); it 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; DwdpManager.setup says so at INFO. - dense/gate.py: the launch gate py_executor_creator runs before the DWDP manager and the model are built. It needs dwdp_config with dwdp_size > 1, eager single-stream execution (no CUDA graphs, torch.compile or prefill graphs), no sleep, LoRA or speculative config, locality domains off (Linear and MLA keep _locality_domain_* shard copies dense DWDP does not move), the overlap scheduler off, tp/pp/cp size 1 and FABRIC peer handles. - pyexecutor/dwdp.py: DwdpManager sets dense DWDP up after the expert setup on a model with dense roots (on every rank, so ranks that disagree on TRTLLM_DWDP_DENSE stop at the first collective), exposes the runtime as `dense`, and releases dense DWDP before the expert teardown. - start_forward (the executor loop) enqueues L(first), L(second) unless they are enqueued already, then starts the expert prefetch (DWDPWeightManager.start_forward). While a forward is still open, i.e. the previous one aborted, it enqueues nothing and leaves both to the model. - ensure_primed (H0, at the model's forward entry, on the execution stream) runs the dense WAR barrier after an aborted forward, enqueues L(first), L(second) unless already primed, then starts the expert prefetch, which enqueues nothing when the first MoE layers are in flight and otherwise restarts its sequence, synchronizing the device first if the aborted forward left a MoE layer bound. So L(first), L(second), E(first), E(second) is also the order after an abort that restarts the expert sequence. With a store rendezvous, mark_setup_done stays the last step of setup, after the dense collectives that also run over the store. Cleanup errors are logged with logger.error and the traceback, since the logger has no exception() method. - modeling_kimi_linear.py: the dense roots and the allow/exclude tables of each decoder layer, the read-window hooks H1-H3 in the decoder layer, H0 priming at the model's forward entry, and a load_weights guard. modeling_kimi_k25.py: the same load guard. - model_loader.py: reload and begin_update_weights refuse a model dense DWDP sharded. - weight_manager.py: copy_stream, use_gather and use_batched_prefetch properties for dense setup. - MOE_DEVELOPER_GUIDE.md: how dense DWDP relates to the expert DWDP protocol. Tests (cpu_only): the switches, the collector and the DEBUG=2 watchdog on a K3-like toy decoder, setup on N rank threads over the loopback store and fake devices (disagreeing ranks and failures in every phase abandon everywhere; an original still alive after the commit point is reported), the FABRIC store over a fake VMM, the Kimi-K3 hook placement and load guard on a real model of a tiny config, the ModelLoader guards, dense setup inside DwdpManager (with its INFO line for more than two expert prefetch slots), and the gate. test_dense_protocol.py runs the runtime over a recording fake of the CUDA seam beside the real DWDPWeightManager and its prefetch schedule, driven like the executor loop, the model and ConfigurableMoE, and checks every dense and expert read with vector clocks over the log: - where H2 and H3 record and what their fill enqueues, into the next forward included, for both segment sets; the copy order (cold start L(1), L(2), E(1), E(2); L(l) right before E(l) within a forward; E(1), L(1), E(2), L(2) at the K3 boundary; nothing at the next start), with two sensitivity checks (no fill at H2; the expert prefetch started before the dense priming); - 1 to 5 and 92 layers with 2, 3 and 4 expert slots: race-free, no fence, demand fetch or host synchronization, primed after every forward; - a forward aborted before H1, before H2, between the expert bind and finish, and before H3, at the first, second, middle and last two layers, recovered with and without the executor's start: the barrier on the execution stream, L(1), L(2) before the expert restart, a device synchronization only after an abandoned bind; - invalidate_prefetch regathers E(1), E(2) only; the fence and the WAR barrier after an abort; 200 random sequences of aborts (also inside a fill and inside H0's priming), starts whose forward fails before the model and invalidations. test_dense_attached_gpu.py (GPU, registered in l0_b200 and l0_gb300_multi_gpus) runs the real runtime and weight manager on one copy stream with a read kernel in every read window and checks every read's bytes: 4 and 5 layers with 2 and 3 expert slots, slow reads or slow copies, both segment sets, a warmup forward, steps, every abort point with its recovery and new expert weights; 92 layers at the K3 boundary; a release right after a forward, with the next forward's copies in flight; skipping H2's consume record is detected. test_dwdp_creator.py gets the cpu_only marker, without which the CPU stage skips the file. test_dwdp_prefetch_gpu.py's stand-in DwdpManager has no dense runtime. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- .../_torch/models/modeling_kimi_k25.py | 13 +- .../_torch/models/modeling_kimi_linear.py | 124 ++ tensorrt_llm/_torch/modules/dwdp/__init__.py | 35 + .../_torch/modules/dwdp/dense/__init__.py | 33 + .../_torch/modules/dwdp/dense/config.py | 228 ++++ .../_torch/modules/dwdp/dense/gate.py | 113 ++ .../_torch/modules/dwdp/dense/layout.py | 772 ++++++++++++ .../_torch/modules/dwdp/dense/plan.py | 548 ++++++++ .../_torch/modules/dwdp/dense/repoint.py | 577 +++++++++ .../_torch/modules/dwdp/dense/runtime.py | 965 ++++++++++++++ .../_torch/modules/dwdp/dense/setup.py | 865 +++++++++++++ .../_torch/modules/dwdp/dense/store.py | 611 +++++++++ .../_torch/modules/dwdp/dense/types.py | 1118 +++++++++++++++++ .../_torch/modules/dwdp/dense/watchdog.py | 200 +++ .../_torch/modules/dwdp/weight_manager.py | 15 + .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 14 + tensorrt_llm/_torch/pyexecutor/dwdp.py | 202 ++- .../_torch/pyexecutor/model_loader.py | 39 +- .../_torch/pyexecutor/py_executor_creator.py | 22 +- .../test_lists/test-db/l0_b200.yml | 2 + .../test-db/l0_gb300_multi_gpus.yml | 2 + .../modules/dwdp/test_dense_attached_gpu.py | 307 +++++ .../_torch/modules/dwdp/test_dense_config.py | 54 + .../_torch/modules/dwdp/test_dense_guards.py | 67 + .../modules/dwdp/test_dense_kimi_hooks.py | 277 ++++ .../_torch/modules/dwdp/test_dense_layout.py | 464 +++++++ .../modules/dwdp/test_dense_protocol.py | 867 +++++++++++++ .../_torch/modules/dwdp/test_dense_setup.py | 602 +++++++++ .../_torch/modules/dwdp/test_dense_store.py | 375 ++++++ .../_torch/modules/dwdp/test_dwdp_creator.py | 68 +- .../_torch/modules/dwdp/test_dwdp_manager.py | 172 ++- .../modules/dwdp/test_dwdp_prefetch_gpu.py | 1 + 32 files changed, 9734 insertions(+), 18 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/__init__.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/config.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/gate.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/layout.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/plan.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/repoint.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/runtime.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/setup.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/store.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/types.py create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/watchdog.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_attached_gpu.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_config.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_guards.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_kimi_hooks.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_layout.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_protocol.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_setup.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_store.py 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..4d3c3799e0d7 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -138,6 +138,7 @@ from ..moe.fused_moe.interface import MoESchedulerKind from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..nccl_window_tensor_scope import nccl_window_tensor_scoped_module +from ..pyexecutor.dwdp import get_global_dwdp_manager from ..utils import AuxStreamType from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently @@ -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,43 @@ 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 DWDP manager exists before the model is built; forward() reaches the dense DWDP + # runtime through it. + self._dwdp = get_global_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 +2320,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. + dwdp = getattr(self, "_dwdp", None) + dense = dwdp.dense if dwdp 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 +2377,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 +2407,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 +2500,10 @@ def __init__(self, model_config: ModelConfig): key="kimi_k3_aux_capture_mode", ) + # The DWDP manager exists before the model is built; forward() primes dense DWDP + # through it. + self._dwdp = get_global_dwdp_manager() + def forward( self, attn_metadata: AttentionMetadata, @@ -2406,6 +2516,13 @@ 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 and expert weights are enqueued + # before any layer runs (the previous forward's end or the executor loop usually already + # did). After a forward that stopped midway, a WAR barrier on this stream comes first. + dwdp = getattr(self, "_dwdp", None) + if dwdp is not None and dwdp.dense is not None: + dwdp.ensure_primed() + if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) hidden_states = inputs_embeds @@ -3113,6 +3230,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..dc9f36232542 100644 --- a/tensorrt_llm/_torch/modules/dwdp/__init__.py +++ b/tensorrt_llm/_torch/modules/dwdp/__init__.py @@ -28,8 +28,18 @@ 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 @@ -57,6 +67,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", @@ -90,4 +111,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..8ec8346d6b90 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/__init__.py @@ -0,0 +1,33 @@ +# 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 DWDP 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 the expert DWDP copy 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. +* ``gate`` -- the launch gate ``py_executor_creator`` runs before the model is built. + +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..5ba6cc8a6e0e --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/config.py @@ -0,0 +1,228 @@ +# 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. +* ``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. +* ``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. + +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. +""" + +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, + 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" + +_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 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``. + """ + + 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 + + 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. + + 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) + 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), + ) + + @property + def has_moe(self) -> bool: + """Whether the MoE-dense segment is enabled.""" + return SegmentKind.MOE in self.segments + + @property + def store_chunk_bytes(self) -> int: + """Store chunk limit in bytes; 0 means a single chunk.""" + return self.store_chunk_mb << 20 + + @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..8b265cffe3d4 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/gate.py @@ -0,0 +1,113 @@ +# 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 the ``DwdpManager`` and the model, so an unsupported launch fails before the weights +load rather than in setup or in the first forward. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, NoReturn + +from tensorrt_llm.llmapi.llm_args import PrefillCudaGraphBackend + +from .. import vmm +from .config import DENSE_ENV, DenseConfig + +if TYPE_CHECKING: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + from tensorrt_llm.mapping import Mapping + + +def validate_dense_gate(llm_args: TorchLlmArgs, mapping: Mapping) -> None: + """Reject launches dense DWDP cannot serve, before the model is built. + + Called only when ``dense_dwdp_requested()``. The requirements, checked in this order: + + * 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. + * ``dwdp_config`` with ``dwdp_size > 1``: the dense rows are sharded over the DWDP group. + * ``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. + * ``disable_overlap_scheduler``: DWDP runs only with the overlap scheduler off. + * ``mapping.tp_size``, ``pp_size`` and ``cp_size`` all 1: every rank of the group holds the + same unsharded dense weights of every layer. + * FABRIC peer handles (``vmm.peer_handle_type()``): the row store is shared only through + FABRIC handles. + + Args: + llm_args: The launch configuration. + mapping: The rank mapping, with the DWDP fields filled in. + + Raises: + ValueError: Naming the first violated requirement. + """ + try: + DenseConfig.from_env() + except ValueError as exc: + _reject(f"valid {DENSE_ENV}_* switches", str(exc)) + dwdp_config = llm_args.dwdp_config + if dwdp_config is None: + _reject("dwdp_config with dwdp_size > 1", "dwdp_config is not set") + if dwdp_config.dwdp_size <= 1: + _reject("dwdp_config with dwdp_size > 1", f"dwdp_size={dwdp_config.dwdp_size}") + _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") + if not llm_args.disable_overlap_scheduler: + _reject("disable_overlap_scheduler=True", "disable_overlap_scheduler=False") + for name in ("tp_size", "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", + ) + + +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}." + ) 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/plan.py b/tensorrt_llm/_torch/modules/dwdp/dense/plan.py new file mode 100644 index 000000000000..04edf9e50826 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/plan.py @@ -0,0 +1,548 @@ +# 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 that follows the expert one: gather, else batched, else per-slice.""" + 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..b9d7118f2d2e --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py @@ -0,0 +1,965 @@ +# 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 ``DwdpManager.ensure_primed()``: barrier + after an abort, then ``L(1), L(2)`` unless + already primed, 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 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. 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 fill never targets a slot +that is unread or being read, so it never needs the fence. + +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. 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 ``DwdpManager``. + + 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 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``; 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; 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 expert DWDP copy stream (``DWDPWeightManager.copy_stream``). + 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, + *, + 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 the expert DWDP 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._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. + 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"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 primed(self) -> bool: + """Whether ``L(first)`` and ``L(second)`` are enqueued for a forward that has not started. + + Set by :meth:`mark_primed`, and 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 (``DwdpManager.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; 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 ``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)) + self._finish_forward() + if self._watchdog is not None: + self._report_watchdog(self._watchdog.poll(), strict=self._strict) + + # ------------------------------------------------------------------ + # Priming and recovery (driven by DwdpManager) + # ------------------------------------------------------------------ + + def prefetch(self, layer: int) -> None: + """Enqueue ``L(layer)`` (every enabled segment, one launch) on the copy stream. + + Used by ``DwdpManager`` 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 ``DwdpManager`` 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. 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. ``DwdpManager.cleanup`` calls it 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; 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) + 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..02e272e4ece2 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/setup.py @@ -0,0 +1,865 @@ +# 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 (design 3.2.6, 3.3.3, 3.7). + +Setup runs once per context worker, from ``DwdpManager.setup`` right after the expert DWDP setup +returned, so the model is fully loaded and post-processed and the expert transport, buffers and +backend fixups are done. 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 (design 3.3.3). +""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, 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, + TorchCudaOps, + allocate_slot_ring, +) +from .watchdog import SlotWatchdog + +if TYPE_CHECKING: + from ..weight_manager import DWDPWeightManager + +#: Free device memory P1 requires on top of the store and the slots (design 3.1.6). +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 DWDP group calls setup makes: an mpi4py sub-communicator or ``StoreDwdpGroup``.""" + + 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 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 (design 3.3.3).""" + + 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.""" + + +#: 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, weight_manager, device)``; setup calls it with the +#: expert DWDP copy stream current. +LauncherFactory = Callable[ + [Mapping[int, LayerCopyPlan], "DWDPWeightManager", 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`` (design 3.2.1). + + ``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. 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, 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`` (design 3.2.1). + + ``DwdpManager.setup`` calls it instead of :func:`setup_dense_dwdp` on a rank whose + ``TRTLLM_DWDP_DENSE`` is on when :func:`dense_roots_available` is False. 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`` (design 3.2.8). + + 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, + weight_manager: DWDPWeightManager, + comm: GroupComm, + *, + dwdp_rank: int, + dwdp_size: int, + device_id: int, + moe_layers: list[int], + 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`` (design 3.2.6, 3.3.3). + + 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`` 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)``. 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 and checksums: continue; else abandon. + 2. P2: allocate the store and the slots, write this rank's rows, ``synchronize``, construct + the runtime (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. + weight_manager: The expert DWDP weight manager: copy stream and transfer mode. + comm: The DWDP group. + dwdp_rank: This rank in the group. + dwdp_size: Group size ``N``. + device_id: CUDA device of this rank. + moe_layers: Expert DWDP MoE layers; the covered dense layers must equal them. + 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=weight_manager.use_gather, + use_batched=weight_manager.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). + """ + if config is not None: + requested = config.enabled + else: + requested = dense_dwdp_requested() + run = _DenseSetup( + model, + weight_manager, + 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, + weight_manager: DWDPWeightManager, + comm: GroupComm, + *, + rank: int, + size: int, + device_id: int, + device: torch.device | None, + moe_layers: Sequence[int], + config: DenseConfig | None, + requested: bool, + store_factory: StoreFactory, + ops: DeviceOps, + launcher_factory: LauncherFactory, + mem_get_info: MemGetInfo, + ) -> None: + self._model = model + self._weight_manager = weight_manager + 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, + ) + _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, + ) + 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) + self._runtime = DenseDWDP( + layout.layers, + layout.segment_kinds, + self._slots, + self._weight_manager.copy_stream, + 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._weight_manager.copy_stream): + self._launcher = self._launcher_factory(plans, self._weight_manager, 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 (design 3.2.8) 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], weight_manager: DWDPWeightManager, device: torch.device +) -> LayerLauncher: + """The production launcher: the transfer the expert DWDP prefetch uses (design 3.5.3).""" + mode = select_copy_mode( + use_gather=weight_manager.use_gather, use_batched=weight_manager.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 (design 3.2.1). + + 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`` (design 3.1.6). + + 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.""" + 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" + ) + 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..e9c313959de7 --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/types.py @@ -0,0 +1,1118 @@ +# 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), +* 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; follows the expert DWDP transfer.""" + + 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.""" + + +#: 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/weight_manager.py b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py index a918e2e2c7af..8596f57ee378 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py @@ -275,6 +275,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..bf9bb6add4d5 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,20 @@ 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 `DwdpManager` and the +model's decoder-layer hooks, which 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 (`DwdpManager.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. `ConfigurableMoE` is unaware of +dense DWDP and keeps the expert wait/bind/finish protocol above unchanged. + 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/pyexecutor/dwdp.py b/tensorrt_llm/_torch/pyexecutor/dwdp.py index 0497b1299119..a71571ca4013 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,14 +23,20 @@ plus runtime entry points that forward to ``DWDPWeightManager`` (``start_forward``, ``wait_and_bind``, ``finish_layer``, ``invalidate_prefetch``) + * Dense DWDP ownership: the ``DenseDWDP`` runtime that ``setup(model)`` + commits, exposed as ``dense`` to the model's layer hooks, started ahead + of the expert prefetch (``start_forward``, ``ensure_primed``) and + released 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 @@ -41,14 +47,22 @@ 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, 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.setup import ( + dense_dwdp_active, + dense_roots_available, + refuse_dense_dwdp_without_roots, +) from tensorrt_llm.llmapi.llm_args import DwdpConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -79,6 +93,15 @@ class DwdpManager: 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, ``dense`` returns its runtime; the model's + decoder-layer hooks drive it and the model's forward calls + ``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)``. """ def __init__( @@ -172,6 +195,9 @@ def __init__( # Set by setup(model); None until then self._weight_manager: Optional[DWDPWeightManager] = None self._model_ref: Optional[weakref.ReferenceType[nn.Module]] = None + # The committed dense DWDP runtime; None when dense DWDP is off, before + # setup(model) and after cleanup() + self._dense: DenseDWDP | None = None # ------------------------------------------------------------------ # DWDP group @@ -221,20 +247,53 @@ def __exit__(self, exc_type, exc_val, exc_tb) -> bool: def is_enabled(self) -> bool: return self.dwdp_size > 1 + @property + def dense(self) -> DenseDWDP | None: + """The committed dense DWDP runtime, or None when dense DWDP is not active. + + 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 + 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: its 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 + dense, self._dense = self._dense, None + if dense is not None: + try: + dense.release() + 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,7 +303,10 @@ 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 @@ -279,8 +341,47 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: manager + MoE backend fixup). Stores the returned ``DWDPWeightManager`` for subsequent runtime calls. - Called once from ``py_executor_creator`` after weight loading. + Then, 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``. + The committed runtime, if any, becomes ``dense``; dense setup errors + propagate. 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 + (``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. + + 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. + DenseSetupError: Dense DWDP setup failed, or was requested on a + model without dense roots, with ``TRTLLM_DWDP_DENSE_REQUIRED`` + on. """ + if self.dwdp_group is None: + raise RuntimeError("DwdpManager.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( + "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( @@ -298,6 +399,24 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: ], ) self._model_ref = weakref.ref(model) + if dense_roots_available(model): + self._dense = setup_dense_dwdp( + model, + self._weight_manager, + self.dwdp_group, + dwdp_rank=self.dwdp_rank, + dwdp_size=self.dwdp_size, + device_id=device_id, + moe_layers=layer_indices, + ) + elif 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 +427,74 @@ 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 weights of the first two + layers go first, unless they are already enqueued (``dense.primed``), + 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 (``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 + if dense is not None: + if dense.forward_open: + return + if not dense.primed: + self._prime(dense) self._weight_manager.start_forward() + def ensure_primed(self) -> None: + """Model-side priming (H0): called at the model's forward entry when ``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). + 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, and starts the expert prefetch + (``DWDPWeightManager.start_forward``), which enqueues nothing when the + first MoE layers are already in flight and otherwise prefetches them; + after an aborted forward it restarts the expert prefetch sequence, + synchronizing the device first 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 + weight_manager = self._weight_manager + if dense is None or weight_manager is None: + raise RuntimeError("DwdpManager.ensure_primed() needs an active dense DWDP runtime") + dense.require_eager("DwdpManager.ensure_primed (H0)") + if dense.forward_open: + dense.war_barrier() + if not dense.primed: + self._prime(dense) + weight_manager.start_forward() + + @staticmethod + def _prime(dense: DenseDWDP) -> None: + """Enqueue the dense weights of the first two covered layers, then mark them primed. + + The covered layers are the expert DWDP MoE layers, so the callers then + start the expert prefetch, whose first copies queue behind these on + the shared copy stream. + """ + for layer in dense.layers[:2]: + dense.prefetch(layer) + dense.mark_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_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f0e465944c5a..08f8d3869ae8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -18,7 +18,9 @@ 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.pyexecutor.resource_manager import ( BaseResourceManager, ResourceManagerType) from tensorrt_llm._utils import get_sm_version @@ -207,6 +209,20 @@ def observe_creation_stage(self, current_stage: ExecutorMemoryType): )) +def _check_dense_dwdp_launch(llm_args: TorchLlmArgs, mapping: Mapping) -> None: + """Run the dense DWDP launch gate if ``TRTLLM_DWDP_DENSE`` asks for it. + + It runs before the DWDP manager and the model are built, so a launch that + dense DWDP cannot serve fails before the weights load. + + Raises: + ValueError: Dense DWDP is requested and the launch violates one of its + requirements (see ``validate_dense_gate``). + """ + if dense_dwdp_requested(): + validate_dense_gate(llm_args, mapping) + + def _flashinfer_one_engine_spec_supported(attn_backend: str, spec_config) -> bool: """Whether this speculation is qualified on the FlashInfer target backend. @@ -733,6 +749,10 @@ def _create_py_executor( ) logger.info("ATTENTION RUNTIME FEATURES: ", attn_runtime_features) + # The dense DWDP gate runs outside the DWDP block below, so that it also + # rejects a dense DWDP launch without dwdp_config. + _check_dense_dwdp_launch(llm_args, mapping) + # Initialize DWDP Manager. # # DWDP needs every global DWDP rank to be a complete, unsharded model replica 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_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/modules/dwdp/test_dense_attached_gpu.py b/tests/unittest/_torch/modules/dwdp/test_dense_attached_gpu.py new file mode 100644 index 000000000000..9f52603485ea --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_attached_gpu.py @@ -0,0 +1,307 @@ +# 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. +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.runtime import DenseDWDP +from tensorrt_llm._torch.modules.dwdp.dense.types import ( + RING_DEPTH, + 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) + self.dense.set_launcher(self._launch_dense) + self.dense.commit(None, fill_byte=0) + self.mgr = object.__new__(DwdpManager) + self.mgr._weight_manager = self.wm + self.mgr._dense = self.dense + 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.mgr.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_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..ea4e89986e80 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_kimi_hooks.py @@ -0,0 +1,277 @@ +# 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 and the load guard. + +Covers the hook calls H0-H3 of ``modeling_kimi_linear`` and its load guard 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 ``DwdpManager`` it captures +when built (``get_global_dwdp_manager``). +""" + +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.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 DWDP manager 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) + yield + set_global_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: + """``DwdpManager.dense`` and ``DwdpManager.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("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._dwdp = manager + for layer in model.layers: + _install_recorders(layer, log) + layer._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({}) 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..b753969ff086 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_protocol.py @@ -0,0 +1,867 @@ +# 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``) beside expert DWDP. + +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; the real +``DWDPWeightManager``, with its ``PrefetchSchedule``, runs on it through patched ``torch.cuda`` +calls and shares the copy stream. A driver calls them the way the executor loop +(``DwdpManager.start_forward`` 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 +shared 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.runtime import DenseDWDP +from tensorrt_llm._torch.modules.dwdp.dense.types import ( + EMPTY_TAG, + RING_DEPTH, + ROW_ALIGN, + 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.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 DwdpManager 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 runtime and the weight manager 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. + num_slots: Expert prefetch slots (``DwdpConfig.num_prefetch_slots``). + """ + + def __init__( + self, + segments: tuple[SegmentKind, ...] = BOTH, + *, + layers: tuple[int, ...] = K3_LAYERS, + num_slots: int = 2, + ) -> None: + self.ops = FakeOps() + self.log = self.ops.log + self.segments = segments + self.layers = tuple(layers) + self.num_slots = num_slots + 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, ops=self.ops) + self.runtime.set_launcher(self.launch_dense) + self.span("commit", None, lambda: self.runtime.commit(None, fill_byte=0)) + 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.manager = object.__new__(DwdpManager) + self.manager._weight_manager = self.weight_manager + self.manager._dense = self.runtime + 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: ``DwdpManager.start_forward``, default stream.""" + with self.cuda(), self.ops.stream(self.ops.default): + self.span("start", None, self.manager.start_forward) + + def h0(self) -> None: + """The model's forward entry: ``DwdpManager.ensure_primed``.""" + self.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 + self.span("bind", layer, lambda: self.manager.wait_and_bind(self.backend, layer)) + self.read(EXPERT, layer) + if stop == "finish": + raise Abort + 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) + + +# --------------------------------------------------------------------------- +# 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,)]) + driver = Driver( + segments, layers=tuple(range(1, num_layers + 1)), num_slots=rng.choice([2, 3, 4]) + ) + 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: + 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 not runtime.primed: + return f"{steps}: demand fetch, or not primed 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: 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..5ba8231f1b78 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_setup.py @@ -0,0 +1,602 @@ +# 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 DWDP 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. + """ + + 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 + + +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=True, + 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), + 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("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$", + ), +} + + +@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) + 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..99d645afc380 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py @@ -12,15 +12,26 @@ # 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 construction ownership tests for DWDP.""" +"""PyExecutor construction tests for DWDP: manager ownership and the dense DWDP launch gate.""" +import os +import re from contextlib import contextmanager from unittest.mock import MagicMock, patch import pytest +from tensorrt_llm._torch.modules.dwdp import vmm +from tensorrt_llm._torch.modules.dwdp.dense.config import DENSE_ENV +from tensorrt_llm._torch.modules.dwdp.dense.gate import validate_dense_gate 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.llmapi.llm_args import DwdpConfig, TorchLlmArgs +from tensorrt_llm.mapping import Mapping + +pytestmark = pytest.mark.cpu_only + +_HANDLE_TYPE = vmm.cuda.CUmemAllocationHandleType @contextmanager @@ -74,3 +85,58 @@ def fail_without_registering(**_kwargs): existing.__exit__.assert_not_called() assert get_global_dwdp_manager() is existing + + +# ---------------------------------------------------------------------- +# Dense DWDP launch gate +# ---------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _no_dense_switches(monkeypatch): + """Start every test without ``TRTLLM_DWDP_DENSE*`` variables; tests set the ones they need.""" + for name in [name for name in os.environ if name.startswith(DENSE_ENV)]: + monkeypatch.delenv(name) + + +@pytest.fixture +def fabric_host(monkeypatch): + """Make this host share FABRIC handles, whatever its architecture.""" + monkeypatch.setattr(vmm, "peer_handle_type", lambda: _HANDLE_TYPE.CU_MEM_HANDLE_TYPE_FABRIC) + + +def _dense_launch(**overrides) -> TorchLlmArgs: + """A dense DWDP context-worker launch, with ``overrides`` applied without validation.""" + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + cuda_graph_config=None, + disable_overlap_scheduler=True, + dwdp_config=DwdpConfig( + dwdp_size=16, num_groups=1, num_experts_per_worker=56, num_prefetch_experts=56 + ), + ) + return args.model_copy(update=overrides) + + +def _mapping(tp_size: int = 1, pp_size: int = 1, cp_size: int = 1) -> Mapping: + return Mapping( + world_size=tp_size * pp_size * cp_size, + rank=0, + tp_size=tp_size, + pp_size=pp_size, + cp_size=cp_size, + ) + + +def test_dense_gate_accepts_a_dense_dwdp_context_worker(fabric_host): + validate_dense_gate(_dense_launch(), _mapping()) + + +@pytest.mark.parametrize( + "dwdp_config, actual", + [(None, "dwdp_config is not set"), (DwdpConfig(dwdp_size=1), "dwdp_size=1")], +) +def test_dense_gate_requires_a_dwdp_group(fabric_host, dwdp_config, actual): + with pytest.raises(ValueError, match=re.escape(f"dwdp_size > 1, but {actual}")): + validate_dense_gate(_dense_launch(dwdp_config=dwdp_config), _mapping()) diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py index 31217886e537..76275df692e2 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py @@ -18,18 +18,23 @@ - creates a DWDP MPI sub-communicator - owns the Single Source of Truth for registered MoE layer indices - forwards runtime calls to the DWDPWeightManager produced by setup_dwdp +- owns the dense DWDP runtime: sets it up after 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 unittest.mock import MagicMock, patch import pytest +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.types import DenseSetupError from tensorrt_llm._torch.pyexecutor.dwdp import ( DwdpManager, get_global_dwdp_manager, @@ -40,11 +45,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,6 +81,21 @@ 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 + + @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") @@ -81,9 +103,14 @@ class TestDwdpManagerLifecycle(unittest.TestCase): def setUp(self): # Ensure no leftover global singleton from a prior test. set_global_dwdp_manager(None) + # 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) + self.setup_dense.assert_not_called() # ------------------------------------------------------------------ # Construction @@ -443,5 +470,146 @@ 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 _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.""" + + 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): + set_global_dwdp_manager(None) + self.addCleanup(set_global_dwdp_manager, None) + 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 = object() # what setup_dwdp returns + 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_setup_runs_dense_setup_when_roots_are_available_but_not_requested(self): + """Collective #1 runs on every rank of a model with roots, requested or not. + + It runs after the expert setup, once the model reference is kept. + """ + seen = [] + + def dense_setup(model, weight_manager, comm, **kwargs): + seen.append((self.manager._model_ref(), self.manager._weight_manager)) + + self.setup_dense.side_effect = dense_setup + self.assertIs(self.manager.setup(self.model), self.weight_manager) + + self.setup_dense.assert_called_once_with( + self.model, + self.weight_manager, + self.sub_comm, + dwdp_rank=self.DWDP_RANK, + dwdp_size=2, + device_id=self.DEVICE_ID, + moe_layers=[3, 5, 7], + ) + self.assertEqual(seen, [(self.model, self.weight_manager)]) + self.assertEqual(self.setup_dwdp.call_args.kwargs["device_id"], self.DEVICE_ID) + self.assertIsNone(self.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_dense_setup_error_propagates_and_cleanup_still_tears_down(self): + self.setup_dense.side_effect = DenseSetupError("dense setup abandoned") + + with self.assertRaisesRegex(DenseSetupError, "dense setup abandoned"): + self.manager.setup(self.model) + self.assertIsNone(self.manager.dense) + + self.manager.cleanup() + + self.assertEqual(self.log, [("teardown_dwdp",), ("group.Free",)]) + + 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_prefetch_gpu.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py index 96f21c46f79b..6c066f2c3ef2 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py @@ -115,6 +115,7 @@ def __init__( ) self.mgr = object.__new__(DwdpManager) self.mgr._weight_manager = self.wm + self.mgr._dense = None # expert DWDP alone self.backend = _Backend() self.compute_stream = torch.cuda.Stream() self.copies: List[int] = [] From a0ad3f2a0720f85fa25401af2544226dfb3ae726 Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:27:52 -0700 Subject: [PATCH 3/9] [None][feat] Dense DWDP without expert DWDP: standalone on attention-DP instances (env form) Dense DWDP no longer needs expert DWDP. With TRTLLM_DWDP_DENSE=1 it runs in one of two modes: - Attached (dwdp_config with dwdp_size > 1, tp_size 1, overlap scheduler off): the behaviour of the previous commit. DwdpManager attaches the dense runtime to an attached DenseDwdpManager, which keeps the copy schedule on the expert copy stream (the next forward's first layers gathered at the end of a forward, the dense priming before the expert prefetch starts) and the teardown order (dense first). - Standalone (no dwdp_config, enable_attention_dp, tp_size > 1, e.g. a DEP16 context server): the dense group is the instance's tp group, created as an MPI sub-communicator with its own tag (0x4457), so a rank that disagrees on TRTLLM_DWDP_DENSE waits in Create_group instead of pairing with MPIDist.tp_comm's tag-0 group. The new DenseDwdpManager owns the runtime and a lazily created copy stream of its own. TRTLLM_DWDP_DENSE_COPY picks the transfer: batched (default; one cudaMemcpyBatchAsync per layer on the copy engines), per_slice, or gather (the SM fabric gather; it holds SMs next to the MoE all-to-all). The overlap scheduler is allowed: the non-overlap executor loop primes the first layers, otherwise the forward entry (H0) does. The runtime does not gather the next forward's first layers at the end of a forward: no expert copies share its stream, and every forward is primed. With the gather, the manager keeps the gather's resident SMs out of the persistent compute grids from its registration, before the model is built, to its exit, as DwdpManager does for the expert gather. As with DwdpManager, a count pinned with TRTLLM_CUTEDSL_SM_RESERVE is reserved with every transfer, the copy engines included. A setup that commits nothing drops the reserve. Standalone is meant for context-only servers; the gate does not check the server role. Expert DWDP over an attention-DP instance (dwdp_config with tp_size > 1) is not a dense DWDP mode: the gate rejects that launch with a message naming both modes. Changes: - dense/manager.py (new): DenseDwdpManager for both modes, standalone_group, StandaloneTransfer and get_global_dense_dwdp_manager. prefetch_first_layers enqueues L(first), L(second) unless primed or a forward is open; ensure_primed (H0) runs the dense WAR barrier after an aborted forward, primes, and attached then starts the expert prefetch (the weight manager's start_forward). - dense/gate.py: the two modes and their checks. Common to both: locality domains off, pp_size 1, cp_size 1. Standalone also needs the MPI backend and refuses an invalid TRTLLM_DWDP_DENSE_COPY; with expert DWDP that variable is ignored with a warning. - dense/config.py, setup.py, runtime.py, plan.py, types.py, __init__.py: mode-aware setup and copy launcher; standalone does not depend on the expert weight manager. DenseDWDP takes a required prefetch_next_pass, which setup turns on when it gets the expert MoE layers. In both modes setup also compares TRTLLM_DWDP_DENSE_STRICT at collective #1: STRICT decides a raise after the commit point, so ranks that differ on it abandon there instead of raising on some ranks only. - fabric_gather.resident_sms takes the caller's transport choice (gather=None keeps the expert one). - pyexecutor/dwdp.py: DwdpManager creates and registers an attached DenseDwdpManager and hands it the runtime instead of owning it. start_forward has the dense manager prime, then starts the expert prefetch, and still leaves an open forward to H0. __enter__ registers both managers, then sets the SM reserve; a failure unregisters both and frees the group. - py_executor_creator.py, py_executor.py, _util.py: create and register the standalone manager before the model is built, hand it to the final executor only, and prime from the non-overlap loop (PyExecutor._prime_dwdp: DwdpManager.start_forward with expert DWDP, else the standalone manager's prefetch_first_layers); PyExecutor validates the two managers together. With KV-cache estimation off, standalone setup runs before the final KV cache is sized, so the freed memory reaches it. - modeling_kimi_linear.py: the K3 hooks read the dense manager (_dense_dwdp) in both modes. - MOE_DEVELOPER_GUIDE.md: ConfigurableMoE is unaware of dense DWDP. Tests: - New: test_dense_gate (the two modes are accepted; a launch that fits neither is refused with what it found; the copy switch is validated standalone and ignored with one warning attached) and test_dense_executor (who primes, either executor loop, expert DWDP still refusing the overlap scheduler, both managers refused, shutdown after the model is dropped, the tagged standalone group; the standalone gather's SM reserve from registration to exit, none for the copy engines or an attached manager, a pinned TRTLLM_CUTEDSL_SM_RESERVE with every transfer, and dropped by a setup that commits nothing). - Updated: test_dense_protocol (the driver goes through the real attached DenseDwdpManager; two standalone K3 forwards keep the per-forward schedule, L(1), L(2) at the start, L(l+2) at H2 and nothing at H3, and recover from an abort at H1, H2 or H3; the random sequences include standalone), test_dense_kimi_hooks (both modes), test_dense_setup (ranks that differ on TRTLLM_DWDP_DENSE or on TRTLLM_DWDP_DENSE_STRICT stop at the first collective; only attached setup turns prefetch_next_pass on), test_fabric_gather (resident_sms with the caller's choice), test_dwdp_creator, test_dwdp_manager, and the stand-in DwdpManager of test_dwdp_prefetch_gpu and test_dense_attached_gpu (an attached DenseDwdpManager). - The executor tests that assemble a PyExecutor without __init__ and reach the DWDP managers (shutdown, the executor loop) set dense_dwdp_manager as they set dwdp_manager. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 33 +- .../_torch/modules/dwdp/dense/__init__.py | 15 +- .../_torch/modules/dwdp/dense/config.py | 73 ++- .../_torch/modules/dwdp/dense/gate.py | 161 +++-- .../_torch/modules/dwdp/dense/manager.py | 595 ++++++++++++++++++ .../_torch/modules/dwdp/dense/plan.py | 7 +- .../_torch/modules/dwdp/dense/runtime.py | 130 ++-- .../_torch/modules/dwdp/dense/setup.py | 201 ++++-- .../_torch/modules/dwdp/dense/types.py | 20 +- .../_torch/modules/dwdp/fabric_gather.py | 11 +- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 34 +- tensorrt_llm/_torch/pyexecutor/_util.py | 3 + tensorrt_llm/_torch/pyexecutor/dwdp.py | 215 ++++--- tensorrt_llm/_torch/pyexecutor/py_executor.py | 59 +- .../_torch/pyexecutor/py_executor_creator.py | 79 ++- .../test_kv_connector_executor_lifetime.py | 1 + .../executor/test_pp_mpi_progress_thread.py | 1 + .../_torch/executor/test_py_executor.py | 1 + .../modules/dwdp/test_dense_attached_gpu.py | 24 +- .../modules/dwdp/test_dense_executor.py | 234 +++++++ .../_torch/modules/dwdp/test_dense_gate.py | 150 +++++ .../modules/dwdp/test_dense_kimi_hooks.py | 146 ++++- .../modules/dwdp/test_dense_protocol.py | 169 +++-- .../_torch/modules/dwdp/test_dense_setup.py | 30 +- .../_torch/modules/dwdp/test_dwdp_creator.py | 101 ++- .../_torch/modules/dwdp/test_dwdp_manager.py | 65 +- .../modules/dwdp/test_dwdp_prefetch_gpu.py | 5 +- .../_torch/modules/dwdp/test_fabric_gather.py | 16 + tests/unittest/_torch/test_connector_adp.py | 1 + .../executor/test_forward_error_survival.py | 1 + 30 files changed, 2132 insertions(+), 449 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/dwdp/dense/manager.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_executor.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dense_gate.py diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 4d3c3799e0d7..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 @@ -138,7 +139,6 @@ from ..moe.fused_moe.interface import MoESchedulerKind from ..moe.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..nccl_window_tensor_scope import nccl_window_tensor_scoped_module -from ..pyexecutor.dwdp import get_global_dwdp_manager from ..utils import AuxStreamType from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently @@ -2248,9 +2248,10 @@ 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 DWDP manager exists before the model is built; forward() reaches the dense DWDP - # runtime through it. - self._dwdp = get_global_dwdp_manager() + # 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``. @@ -2325,8 +2326,8 @@ def _full_rows(t: torch.Tensor) -> torch.Tensor: # 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. - dwdp = getattr(self, "_dwdp", None) - dense = dwdp.dense if dwdp is not None else None + 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: @@ -2500,9 +2501,10 @@ def __init__(self, model_config: ModelConfig): key="kimi_k3_aux_capture_mode", ) - # The DWDP manager exists before the model is built; forward() primes dense DWDP - # through it. - self._dwdp = get_global_dwdp_manager() + # 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, @@ -2516,12 +2518,13 @@ 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 and expert weights are enqueued - # before any layer runs (the previous forward's end or the executor loop usually already - # did). After a forward that stopped midway, a WAR barrier on this stream comes first. - dwdp = getattr(self, "_dwdp", None) - if dwdp is not None and dwdp.dense is not None: - dwdp.ensure_primed() + # 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) diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/__init__.py b/tensorrt_llm/_torch/modules/dwdp/dense/__init__.py index 8ec8346d6b90..9ccf832f4be5 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/__init__.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/__init__.py @@ -14,9 +14,12 @@ # limitations under the License. """Dense DWDP: row-sharded replicated dense weights, gathered two layers ahead of their use. -Every rank of the DWDP 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 the expert DWDP copy stream. Submodules: +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. @@ -27,7 +30,11 @@ * ``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. -* ``gate`` -- the launch gate ``py_executor_creator`` runs before the model is built. +* ``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 index 5ba6cc8a6e0e..0d25a871d528 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/config.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/config.py @@ -14,7 +14,9 @@ # 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. +* ``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. @@ -22,17 +24,31 @@ * ``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. + 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; experimental standalone: it holds SMs next to + the MoE all-to-all). 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. +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 @@ -48,6 +64,7 @@ DEFAULT_MIN_BYTES, LOG_PREFIX, SEGMENT_ORDER, + CopyMode, DigestSwitches, SegmentKind, ) @@ -59,6 +76,7 @@ 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") @@ -121,6 +139,26 @@ def parse_segments(value: str) -> tuple[SegmentKind, ...]: 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. @@ -143,6 +181,9 @@ class DenseConfig: 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 @@ -152,6 +193,7 @@ class DenseConfig: 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) @@ -166,6 +208,9 @@ def __post_init__(self) -> None: 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. @@ -178,6 +223,7 @@ def from_env(cls, environ: Mapping[str, str] | None = None) -> DenseConfig: 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, @@ -186,17 +232,22 @@ def from_env(cls, environ: Mapping[str, str] | None = None) -> DenseConfig: 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, ) - @property - def has_moe(self) -> bool: - """Whether the MoE-dense segment is enabled.""" - return SegmentKind.MOE in self.segments + def standalone_copy_mode(self) -> CopyMode: + """The transfer of standalone dense DWDP: ``copy`` parsed, ``BATCHED`` when it is None. - @property - def store_chunk_bytes(self) -> int: - """Store chunk limit in bytes; 0 means a single chunk.""" - return self.store_chunk_mb << 20 + 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: diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/gate.py b/tensorrt_llm/_torch/modules/dwdp/dense/gate.py index 8b265cffe3d4..bf11ef0cac3c 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/gate.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/gate.py @@ -15,65 +15,102 @@ """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 the ``DwdpManager`` and the model, so an unsupported launch fails before the weights -load rather than in setup or in the first forward. +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 vmm -from .config import DENSE_ENV, DenseConfig +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 - -def validate_dense_gate(llm_args: TorchLlmArgs, mapping: Mapping) -> None: - """Reject launches dense DWDP cannot serve, before the model is built. - - Called only when ``dense_dwdp_requested()``. The requirements, checked in this order: - - * 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. - * ``dwdp_config`` with ``dwdp_size > 1``: the dense rows are sharded over the DWDP group. - * ``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. - * ``disable_overlap_scheduler``: DWDP runs only with the overlap scheduler off. - * ``mapping.tp_size``, ``pp_size`` and ``cp_size`` all 1: every rank of the group holds the - same unsharded dense weights of every layer. - * FABRIC peer handles (``vmm.peer_handle_type()``): the row store is shared only through - FABRIC handles. +#: 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: the gather is experimental without expert DWDP. 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: - DenseConfig.from_env() + config = DenseConfig.from_env() except ValueError as exc: _reject(f"valid {DENSE_ENV}_* switches", str(exc)) - dwdp_config = llm_args.dwdp_config - if dwdp_config is None: - _reject("dwdp_config with dwdp_size > 1", "dwdp_config is not set") - if dwdp_config.dwdp_size <= 1: - _reject("dwdp_config with dwdp_size > 1", f"dwdp_size={dwdp_config.dwdp_size}") + 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 @@ -87,9 +124,7 @@ def validate_dense_gate(llm_args: TorchLlmArgs, mapping: Mapping) -> None: _require_unset("speculative_config", llm_args.speculative_config) if llm_args.enable_locality_domains: _reject("enable_locality_domains=False", "enable_locality_domains=True") - if not llm_args.disable_overlap_scheduler: - _reject("disable_overlap_scheduler=True", "disable_overlap_scheduler=False") - for name in ("tp_size", "pp_size", "cp_size"): + for name in ("pp_size", "cp_size"): size = getattr(mapping, name) if size != 1: _reject(f"{name}=1", f"{name}={size}") @@ -99,6 +134,49 @@ def validate_dense_gate(llm_args: TorchLlmArgs, mapping: Mapping) -> None: "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 is experimental without expert DWDP: the gather " + f"kernel holds SMs next to the MoE all-to-all, a combination not validated yet; batched, " + f"the default, runs on the copy engines" + ) def _require_unset(name: str, value: object) -> None: @@ -111,3 +189,10 @@ def _reject(requirement: str, actual: str) -> NoReturn: 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/manager.py b/tensorrt_llm/_torch/modules/dwdp/dense/manager.py new file mode 100644 index 000000000000..f14f9fa14d8e --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/dense/manager.py @@ -0,0 +1,595 @@ +# 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 with the expert weight manager as the + transfer, 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 index 04edf9e50826..6224fd5b5b45 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/plan.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/plan.py @@ -364,7 +364,12 @@ def batched_copies(plan: LayerCopyPlan, *, slice_bytes: int = BATCH_SLICE_BYTES) def select_copy_mode(*, use_gather: bool, use_batched: bool) -> CopyMode: - """The dense transfer that follows the expert one: gather, else batched, else per-slice.""" + """The dense transfer from a transfer's two flags: gather, else batched, else per-slice. + + With expert DWDP the flags are the expert DWDP weight manager's, 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: diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py b/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py index b9d7118f2d2e..6f4539e43a05 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py @@ -19,10 +19,10 @@ ===== ========================================== ========================================== Hook Call site Action ===== ========================================== ========================================== -H0 ``KimiLinearModel.forward`` entry ``DwdpManager.ensure_primed()``: barrier - after an abort, then ``L(1), L(2)`` unless - already primed, then the expert - ``start_forward`` +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 @@ -32,22 +32,26 @@ 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 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. 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 fill never targets a slot -that is unread or being read, so it never needs the fence. - -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. 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. +``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 @@ -99,7 +103,10 @@ class _Phase(enum.Enum): class DenseDWDP: - """Slot rings, events and per-layer hooks of dense DWDP; owned by ``DwdpManager``. + """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 @@ -127,14 +134,14 @@ class DenseDWDP: 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 the first two layers of the next forward - once their slots free up at the end of this one. + 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``; H3 of layers 91 and 92 enqueue layers 1 and 2 of the next forward (H2 of them with - the attention segment alone). + ``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 @@ -152,11 +159,18 @@ class DenseDWDP: ``(layer, segment)`` whose slot held other bytes (``strict``: raises). Args: - layers: Covered layers, ascending; the same list as the expert DWDP MoE layers. + 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 expert DWDP copy stream (``DWDPWeightManager.copy_stream``). + copy_stream: The copy stream: the expert DWDP copy stream with expert DWDP + (``DWDPWeightManager.copy_stream``), 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. @@ -178,6 +192,7 @@ def __init__( slots: Mapping[SegmentKind, Sequence[SlotBuffer]], copy_stream: StreamLike, *, + prefetch_next_pass: bool, ops: DeviceOps | None = None, strict: bool = False, store: DenseShardStore | None = None, @@ -187,7 +202,7 @@ def __init__( kinds = _segment_kinds(segments) rings = _slot_rings(kinds, slots) if copy_stream is None: - raise ValueError("dense DWDP needs the expert DWDP copy stream") + 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 " @@ -205,6 +220,7 @@ def __init__( 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, ...]] = { @@ -226,7 +242,8 @@ def __init__( 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. + # 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 @@ -242,6 +259,7 @@ def __init__( 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'}" ) @@ -264,12 +282,17 @@ 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 by H3 of the last covered layer once the fill has enqueued - them; H1 of the first covered layer clears it. + 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 @@ -426,7 +449,7 @@ def require_eager(self, hook: str) -> None: 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 (``DwdpManager.ensure_primed``) calls this before it records or + 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. @@ -519,13 +542,14 @@ def layer_done(self, layer: int) -> None: 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; 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 ``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. + 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 @@ -543,18 +567,19 @@ def layer_done(self, layer: int) -> None: if position == len(self._layers) - 1: self._forward_open = False self._next_fill = max(0, self._next_fill - len(self._layers)) - self._finish_forward() + 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 DwdpManager) + # 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 ``DwdpManager`` priming (``L(first), L(second)``, then :meth:`mark_primed`). + 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 @@ -570,7 +595,7 @@ def prefetch(self, layer: int) -> None: self._prefetch(layer, fence=False) def mark_primed(self) -> None: - """Record that ``DwdpManager`` enqueued the first two layers; the fill goes on after them. + """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 @@ -598,9 +623,9 @@ def war_barrier(self) -> None: 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. 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. + 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`. @@ -669,8 +694,8 @@ def release(self) -> None: 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. ``DwdpManager.cleanup`` calls it before the expert - teardown. + 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 @@ -811,12 +836,13 @@ def _slot_free(self, kind: SegmentKind, slot: int) -> bool: 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; 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. + 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) + 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]): diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/setup.py b/tensorrt_llm/_torch/modules/dwdp/dense/setup.py index 02e272e4ece2..85db39641b61 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/setup.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/setup.py @@ -12,21 +12,22 @@ # 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 (design 3.2.6, 3.3.3, 3.7). - -Setup runs once per context worker, from ``DwdpManager.setup`` right after the expert DWDP setup -returned, so the model is fully loaded and post-processed and the expert transport, buffers and -backend fixups are done. 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 (design 3.3.3). +"""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`` right after the expert DWDP setup returned (with expert DWDP, so the expert +transport, buffers and backend fixups are done), 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 (design 3.3.3). """ from __future__ import annotations import time from collections.abc import Callable, Iterable, Mapping, Sequence -from typing import TYPE_CHECKING, NamedTuple, Protocol +from typing import NamedTuple, Protocol import torch from torch import nn @@ -51,15 +52,13 @@ LayerLauncher, SegmentKind, SlotBuffer, + StreamLike, TorchCudaOps, allocate_slot_ring, ) from .watchdog import SlotWatchdog -if TYPE_CHECKING: - from ..weight_manager import DWDPWeightManager - -#: Free device memory P1 requires on top of the store and the slots (design 3.1.6). +#: 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. @@ -76,7 +75,12 @@ class GroupComm(Protocol): - """The DWDP group calls setup makes: an mpi4py sub-communicator or ``StoreDwdpGroup``.""" + """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.""" @@ -91,6 +95,32 @@ def Get_size(self) -> int: ... +class DenseTransfer(Protocol): + """Where and how the dense copies run. + + ``DWDPWeightManager`` (with expert DWDP: the expert DWDP copy stream and transfer) 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. @@ -106,7 +136,7 @@ class PhaseReport(NamedTuple): class _CollectData(NamedTuple): - """The data of one rank in collective #1 (design 3.3.3).""" + """The data of one rank in collective #1.""" requested: bool """``TRTLLM_DWDP_DENSE`` on this rank.""" @@ -129,6 +159,9 @@ class _CollectData(NamedTuple): 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) @@ -136,10 +169,10 @@ class _CollectData(NamedTuple): #: 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, weight_manager, device)``; setup calls it with the -#: expert DWDP copy stream current. +#: Builds the copy launcher for ``(plans, transfer, device)``; setup calls it with the transfer's +#: copy stream current. LauncherFactory = Callable[ - [Mapping[int, LayerCopyPlan], "DWDPWeightManager", torch.device], LayerLauncher + [Mapping[int, LayerCopyPlan], DenseTransfer, torch.device], LayerLauncher ] #: ``(free, total)`` bytes of the dense device, as ``torch.cuda.mem_get_info`` returns them. @@ -147,14 +180,16 @@ class _CollectData(NamedTuple): def dense_roots_available(model: nn.Module) -> bool: - """Whether the model's decoder layers provide ``dense_dwdp_roots`` (design 3.2.1). - - ``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. 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, 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. + """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) @@ -164,15 +199,17 @@ def dense_roots_available(model: nn.Module) -> bool: def refuse_dense_dwdp_without_roots(model: nn.Module) -> None: - """Turn down a request for dense DWDP on a model without ``dense_dwdp_roots`` (design 3.2.1). - - ``DwdpManager.setup`` calls it instead of :func:`setup_dense_dwdp` on a rank whose - ``TRTLLM_DWDP_DENSE`` is on when :func:`dense_roots_available` is False. 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. + """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. @@ -197,7 +234,7 @@ def refuse_dense_dwdp_without_roots(model: nn.Module) -> None: def dense_dwdp_active(model: nn.Module) -> bool: - """Whether dense DWDP moved the dense weights of ``model`` (design 3.2.8). + """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 @@ -215,13 +252,13 @@ def dense_dwdp_active(model: nn.Module) -> bool: def setup_dense_dwdp( model: nn.Module, - weight_manager: DWDPWeightManager, + transfer: DenseTransfer, comm: GroupComm, *, dwdp_rank: int, dwdp_size: int, device_id: int, - moe_layers: list[int], + moe_layers: Sequence[int] | None = None, config: DenseConfig | None = None, store_factory: StoreFactory | None = None, ops: DeviceOps | None = None, @@ -229,22 +266,25 @@ def setup_dense_dwdp( launcher_factory: LauncherFactory | None = None, mem_get_info: MemGetInfo | None = None, ) -> DenseDWDP | None: - """Shard, verify and repoint the dense weights of ``model`` (design 3.2.6, 3.3.3). + """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`` and that ``free >= store + slots + + 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)``. 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 and checksums: continue; else abandon. + 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 (it records the consume events; with ``config.slot_watchdog`` it gets the - slot checksum watchdog over the P1 checksums), export. Data: the exported handles. + 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. @@ -274,20 +314,25 @@ def setup_dense_dwdp( Args: model: The loaded model (the VL wrapper or the causal LM); the decoder is found like the expert setup finds it. - weight_manager: The expert DWDP weight manager: copy stream and transfer mode. - comm: The DWDP group. - dwdp_rank: This rank in the group. - dwdp_size: Group size ``N``. + transfer: Copy stream and transfer mode: the expert DWDP weight manager with expert DWDP, + 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. + 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=weight_manager.use_gather, - use_batched=weight_manager.use_batched_prefetch), device=device)``. + ``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. @@ -299,7 +344,8 @@ def setup_dense_dwdp( 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). + 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 @@ -307,7 +353,7 @@ def setup_dense_dwdp( requested = dense_dwdp_requested() run = _DenseSetup( model, - weight_manager, + transfer, comm, rank=dwdp_rank, size=dwdp_size, @@ -336,14 +382,14 @@ class _DenseSetup: def __init__( self, model: nn.Module, - weight_manager: DWDPWeightManager, + transfer: DenseTransfer, comm: GroupComm, *, rank: int, size: int, device_id: int, device: torch.device | None, - moe_layers: Sequence[int], + moe_layers: Sequence[int] | None, config: DenseConfig | None, requested: bool, store_factory: StoreFactory, @@ -352,7 +398,7 @@ def __init__( mem_get_info: MemGetInfo, ) -> None: self._model = model - self._weight_manager = weight_manager + self._transfer = transfer self._comm = comm self._rank = rank self._size = size @@ -499,7 +545,8 @@ def _collect(self) -> PhaseReport: attr_allow=attr_allow, exclude=exclude, ) - _check_coverage(layout, self._moe_layers) + if self._moe_layers is not None: + _check_coverage(layout, self._moe_layers) _check_free_memory(layout, free_bytes) compute_checksums(layout) data = _CollectData( @@ -510,6 +557,7 @@ def _collect(self) -> PhaseReport: 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( @@ -535,11 +583,16 @@ def _allocate(self) -> PhaseReport: 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._weight_manager.copy_stream, + 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, @@ -562,8 +615,8 @@ def _map(self, payloads: Sequence[StorePayload]) -> PhaseReport: chunk_bases=store.chunk_bases(), slot_bases=slot_bases, ) - with self._ops.stream(self._weight_manager.copy_stream): - self._launcher = self._launcher_factory(plans, self._weight_manager, self._device) + 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 @@ -654,7 +707,7 @@ def _check_group(self) -> None: 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 (design 3.2.8) read this flag. + The weight reload, update and load guards read this flag. """ self._model._dense_dwdp_active = True if self._decoder is not None: @@ -708,11 +761,15 @@ def _release(self) -> bool: def _copy_launcher( - plans: Mapping[int, LayerCopyPlan], weight_manager: DWDPWeightManager, device: torch.device + plans: Mapping[int, LayerCopyPlan], transfer: DenseTransfer, device: torch.device ) -> LayerLauncher: - """The production launcher: the transfer the expert DWDP prefetch uses (design 3.5.3).""" + """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=weight_manager.use_gather, use_batched=weight_manager.use_batched_prefetch + use_gather=transfer.use_gather, use_batched=transfer.use_batched_prefetch ) return DenseCopyLauncher(plans, mode, device=device) @@ -730,7 +787,7 @@ def _allocated_bytes(device: torch.device) -> int: 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 (design 3.2.1). + """The model's allow and exclude tables, read from the decoder layers. Raises: DenseLayoutError: Decoder layers carry different tables; the collector applies one pair @@ -778,7 +835,7 @@ def _check_coverage(layout: DenseLayout, moe_layers: Sequence[int]) -> None: def _check_free_memory(layout: DenseLayout, free_bytes: int) -> None: - """P1 precheck: ``free >= store + slots + SETUP_HEADROOM_BYTES`` (design 3.1.6). + """P1 precheck: ``free >= store + slots + SETUP_HEADROOM_BYTES``. Raises: DenseSetupError: Too little free device memory. @@ -815,7 +872,11 @@ def _failures(reports: Sequence[PhaseReport]) -> str: def _disagreement(layout: DenseLayout | None, data: Sequence[_CollectData]) -> str: - """Why ok ranks cannot share rows: different layouts, or the same layout with other bytes.""" + """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()) @@ -823,6 +884,10 @@ def _disagreement(layout: DenseLayout | None, data: Sequence[_CollectData]) -> s 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) diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/types.py b/tensorrt_llm/_torch/modules/dwdp/dense/types.py index e9c313959de7..a3be8e3db14d 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/types.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/types.py @@ -22,7 +22,7 @@ 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), +* 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, @@ -167,7 +167,11 @@ class BindingKind(str, enum.Enum): class CopyMode(str, enum.Enum): - """Transfer used to gather a layer into its slots; follows the expert DWDP transfer.""" + """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``).""" @@ -180,6 +184,18 @@ class CopyMode(str, enum.Enum): """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] 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/moe/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md index bf9bb6add4d5..647ef92562b8 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -104,18 +104,28 @@ 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 `DwdpManager` and the -model's decoder-layer hooks, which 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 (`DwdpManager.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. `ConfigurableMoE` is unaware of -dense DWDP and keeps the expert wait/bind/finish protocol above unchanged. +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 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 a71571ca4013..946ef89f57d4 100644 --- a/tensorrt_llm/_torch/pyexecutor/dwdp.py +++ b/tensorrt_llm/_torch/pyexecutor/dwdp.py @@ -23,10 +23,13 @@ plus runtime entry points that forward to ``DWDPWeightManager`` (``start_forward``, ``wait_and_bind``, ``finish_layer``, ``invalidate_prefetch``) - * Dense DWDP ownership: the ``DenseDWDP`` runtime that ``setup(model)`` - commits, exposed as ``dense`` to the model's layer hooks, started ahead - of the expert prefetch (``start_forward``, ``ensure_primed``) and - released before the expert teardown + * 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 and hands it to the dense manager together with the weight + manager; 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`` @@ -58,11 +61,16 @@ 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.llmapi.llm_args import DwdpConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -84,24 +92,28 @@ class DwdpManager: 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, ``dense`` returns its runtime; the model's - decoder-layer hooks drive it and the model's forward calls - ``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)``. + 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. """ def __init__( @@ -112,6 +124,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() @@ -195,9 +211,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 - # The committed dense DWDP runtime; None when dense DWDP is off, before - # setup(model) and after cleanup() - self._dense: DenseDWDP | None = 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 @@ -230,10 +247,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: @@ -242,46 +277,52 @@ 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. - 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. + 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 + return self._dense_manager.dense def cleanup(self) -> None: """Release dense DWDP, tear down the weight manager, free the DWDP group. - Dense DWDP goes first: its 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. + 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 - dense, self._dense = self._dense, None - if dense is not None: - try: - dense.release() - 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()}" - ) + 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: @@ -345,11 +386,12 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: ``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``. - The committed runtime, if any, becomes ``dense``; dense setup errors - propagate. 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 - (``refuse_dense_dwdp_without_roots``: it raises unless its + The committed runtime, if any, goes to the dense manager together with + the weight manager (``DenseDwdpManager.attach``) and becomes + ``dense``; dense setup errors propagate. 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 (``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 @@ -369,7 +411,7 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: """ if self.dwdp_group is None: raise RuntimeError("DwdpManager.setup() called after cleanup()") - if self._dense is not None: + 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 " @@ -400,7 +442,7 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: ) self._model_ref = weakref.ref(model) if dense_roots_available(model): - self._dense = setup_dense_dwdp( + dense = setup_dense_dwdp( model, self._weight_manager, self.dwdp_group, @@ -409,6 +451,8 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: device_id=device_id, moe_layers=layer_indices, ) + if dense is not None: + self._dense_manager.attach(dense, expert=self._weight_manager) elif dense_dwdp_requested(): refuse_dense_dwdp_without_roots(model) if self.dense is not None and self.config.num_prefetch_slots > 2: @@ -431,69 +475,42 @@ def start_forward(self) -> None: 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 weights of the first two - layers go first, unless they are already enqueued (``dense.primed``), - 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 (``ensure_primed``) - recovers on the execution stream, dense first. + 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 - if dense is not None: - if dense.forward_open: - return - if not dense.primed: - self._prime(dense) + 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): called at the model's forward entry when ``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). - 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, and starts the expert prefetch - (``DWDPWeightManager.start_forward``), which enqueues nothing when the - first MoE layers are already in flight and otherwise prefetches them; - after an aborted forward it restarts the expert prefetch sequence, - synchronizing the device first if that forward left a MoE layer bound. + """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. """ - dense = self._dense - weight_manager = self._weight_manager - if dense is None or weight_manager is None: + if self._dense_manager.dense is None or self._weight_manager is None: raise RuntimeError("DwdpManager.ensure_primed() needs an active dense DWDP runtime") - dense.require_eager("DwdpManager.ensure_primed (H0)") - if dense.forward_open: - dense.war_barrier() - if not dense.primed: - self._prime(dense) - weight_manager.start_forward() - - @staticmethod - def _prime(dense: DenseDWDP) -> None: - """Enqueue the dense weights of the first two covered layers, then mark them primed. - - The covered layers are the expert DWDP MoE layers, so the callers then - start the expert prefetch, whose first copies queue behind these on - the shared copy stream. - """ - for layer in dense.layers[:2]: - dense.prefetch(layer) - dense.mark_primed() + 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.""" diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 1375706cdfc3..6b581b5a4d14 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 @@ -535,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, @@ -1244,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__: @@ -1287,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() @@ -1311,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 @@ -2221,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: """ @@ -5323,6 +5353,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. @@ -5485,8 +5533,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 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 08f8d3869ae8..7ef528916ba0 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -21,6 +21,9 @@ 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 @@ -209,18 +212,26 @@ def observe_creation_stage(self, current_stage: ExecutorMemoryType): )) -def _check_dense_dwdp_launch(llm_args: TorchLlmArgs, mapping: Mapping) -> None: +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 and the model are built, so a launch that - dense DWDP cannot serve fails before the weights load. + 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(): - validate_dense_gate(llm_args, mapping) + return validate_dense_gate(llm_args, mapping) + return None def _flashinfer_one_engine_spec_supported(attn_backend: str, @@ -749,9 +760,10 @@ def _create_py_executor( ) logger.info("ATTENTION RUNTIME FEATURES: ", attn_runtime_features) - # The dense DWDP gate runs outside the DWDP block below, so that it also - # rejects a dense DWDP launch without dwdp_config. - _check_dense_dwdp_launch(llm_args, mapping) + # 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. # @@ -794,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() @@ -1149,6 +1174,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): resources = {} estimating_kv_cache = False kv_cache_creator = None + # Whether standalone dense DWDP was set up before the KV caches (below). + dense_dwdp_set_up = False # Create the execution stream for model forward operations # for proper synchronization with KVCacheTransferManager's onboard/offload operations. @@ -1205,15 +1232,28 @@ 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 standalone dense DWDP is set up first: + # 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. + if dense_dwdp_manager is not None and not estimating_kv_cache: + dense_dwdp_manager.setup(model_engine.model) + dense_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). + # weight manager + MoE backend fixup (single entry point). Dense DWDP is set + # up there with expert DWDP, else here when it runs standalone (unless it + # was set up before the KV caches above). if dwdp_manager is not None: dwdp_manager.setup(model_engine.model) + elif dense_dwdp_manager is not None and not dense_dwdp_set_up: + dense_dwdp_manager.setup(model_engine.model) # Resource managers for speculative decoding # For user-specified drafters, use extra_resource_managers in PyTorchBackend config @@ -1266,6 +1306,10 @@ 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 standalone dense DWDP manager and exits it at shutdown. + dense_dwdp_manager=dense_dwdp_manager + if not estimating_kv_cache else None, max_num_sequences=max_num_seq_slots, ) @@ -1341,6 +1385,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, ) @@ -1380,8 +1425,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: @@ -1403,4 +1455,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/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 index 9f52603485ea..0c7d8e362aae 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dense_attached_gpu.py +++ b/tests/unittest/_torch/modules/dwdp/test_dense_attached_gpu.py @@ -17,12 +17,12 @@ 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. -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. +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 @@ -32,9 +32,11 @@ 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, @@ -140,12 +142,16 @@ def __init__( 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) + 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 = self.dense + self.mgr._dense_manager = self.dense_manager self.backend = _Backend() self.compute_stream = torch.cuda.Stream() self.results: List[Tuple[str, int, torch.Tensor]] = [] @@ -182,7 +188,7 @@ def run_pass( self.mgr.start_forward() self.compute_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.compute_stream): - self.mgr.ensure_primed() # H0 + 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": 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..055168e9fbb2 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dense_gate.py @@ -0,0 +1,150 @@ +# 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.llmapi.llm_args import DwdpConfig, 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"} + ), +} + + +@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_kimi_hooks.py b/tests/unittest/_torch/modules/dwdp/test_dense_kimi_hooks.py index ea4e89986e80..e20c486b3bc0 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dense_kimi_hooks.py +++ b/tests/unittest/_torch/modules/dwdp/test_dense_kimi_hooks.py @@ -12,16 +12,17 @@ # 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 and the load guard. - -Covers the hook calls H0-H3 of ``modeling_kimi_linear`` and its load guard 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 ``DwdpManager`` it captures -when built (``get_global_dwdp_manager``). +"""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 @@ -35,6 +36,11 @@ 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 @@ -133,14 +139,16 @@ def __init__(self, *, num_experts: int, **kwargs) -> None: 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 DWDP manager when built: every test starts and ends without one. + 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) # --------------------------------------------------------------------------- @@ -184,7 +192,7 @@ def layer_done(self, layer: int) -> None: def _manager(log: list, dense: object | None) -> SimpleNamespace: - """``DwdpManager.dense`` and ``DwdpManager.ensure_primed``, recorded.""" + """``DenseDwdpManager.dense`` and ``DenseDwdpManager.ensure_primed``, recorded.""" return SimpleNamespace(dense=dense, ensure_primed=lambda: log.append(("ensure_primed",))) @@ -217,6 +225,18 @@ def _layer_reads(index: int, dense: bool) -> list[tuple]: ] +@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.""" @@ -230,10 +250,10 @@ def test_the_model_primes_first_and_brackets_the_reads_of_covered_layers(state): "no_dense": _manager(log, None), "no_manager": None, }[state] - model._dwdp = manager + model._dense_dwdp = manager for layer in model.layers: _install_recorders(layer, log) - layer._dwdp = manager + layer._dense_dwdp = manager # fp32 activations keep the residual mixing and the norms on their eager CPU paths. metadata = SimpleNamespace(num_tokens=_TOKENS) @@ -275,3 +295,101 @@ def test_load_weights_refuses_while_dense_is_active(): 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_protocol.py b/tests/unittest/_torch/modules/dwdp/test_dense_protocol.py index b753969ff086..0ade32852c42 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dense_protocol.py +++ b/tests/unittest/_torch/modules/dwdp/test_dense_protocol.py @@ -12,20 +12,21 @@ # 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``) beside expert DWDP. +"""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; the real +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. A driver calls them the way the executor loop -(``DwdpManager.start_forward`` 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 -shared 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. +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 @@ -42,17 +43,20 @@ 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 @@ -214,7 +218,7 @@ def happens_before( # --------------------------------------------------------------------------- -# Driver: the executor loop, the model, ConfigurableMoE and DwdpManager around the runtime +# Driver: the executor loop, the model, ConfigurableMoE and the managers around the runtime # --------------------------------------------------------------------------- @@ -227,15 +231,17 @@ class LaunchError(Exception): class Driver: - """Calls the runtime and the weight manager in program order; logs reads and copy launches. + """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. + 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__( @@ -244,12 +250,14 @@ def __init__( *, 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]] = {} @@ -258,24 +266,41 @@ def __init__( 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, ops=self.ops) + 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)) - 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.manager = object.__new__(DwdpManager) - self.manager._weight_manager = self.weight_manager - self.manager._dense = self.runtime + 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") @@ -342,13 +367,17 @@ def cuda(self) -> Iterator[None]: # -- executor loop, model and ConfigurableMoE --------------------------------- def start(self) -> None: - """The executor loop before a forward: ``DwdpManager.start_forward``, default stream.""" + """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, self.manager.start_forward) + self.span("start", None, start) def h0(self) -> None: - """The model's forward entry: ``DwdpManager.ensure_primed``.""" - self.manager.ensure_primed() + """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``.""" @@ -374,11 +403,13 @@ def layer(self, layer: int, stop: str | None = None) -> None: 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 - self.span("bind", layer, lambda: self.manager.wait_and_bind(self.backend, layer)) - self.read(EXPERT, layer) + 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 - self.span("finish", layer, lambda: self.manager.finish_layer(layer)) + 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": @@ -811,6 +842,56 @@ def test_invalidate_prefetch_keeps_the_dense_copies(): _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 # --------------------------------------------------------------------------- @@ -821,8 +902,12 @@ def _random_sequence(seed: int) -> str | 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]) + segments, + layers=tuple(range(1, num_layers + 1)), + num_slots=rng.choice([2, 3, 4]), + standalone=standalone, ) runtime = driver.runtime steps = [] @@ -839,8 +924,9 @@ def _random_sequence(seed: int) -> str | None: steps.append("start, then the forward fails before the model") driver.start() elif roll < 0.55: - steps.append("invalidate") - driver.invalidate() + 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) @@ -849,8 +935,8 @@ def _random_sequence(seed: int) -> str | None: demand_fetches = runtime.demand_fetches if not driver.forward(start=start): return f"{steps}: the forward stopped" - if runtime.demand_fetches != demand_fetches or not runtime.primed: - return f"{steps}: demand fetch, or not primed after a complete forward" + 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: @@ -862,6 +948,7 @@ def _random_sequence(seed: int) -> str | 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: no race, fence or window error.""" + 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 index 5ba8231f1b78..c0bbaf9bbfc1 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dense_setup.py +++ b/tests/unittest/_torch/modules/dwdp/test_dense_setup.py @@ -14,7 +14,7 @@ # limitations under the License. """CPU tests of the dense DWDP setup. -Every rank of an ``N``-rank DWDP group is a thread that runs ``setup_dense_dwdp`` on its own copy +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 @@ -295,8 +295,10 @@ class RankSpec: ``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 @@ -305,6 +307,7 @@ class RankSpec: fail: str | None = None corrupt: int | None = None perturb: str | None = None + attached: bool = True class Rank: @@ -360,7 +363,7 @@ def run(self, store, size): ) self.comm = CountingComm(group) config = DenseConfig( - enabled=True, + enabled=self.spec.requested, required=self.spec.required, min_bytes=TOY_MIN_BYTES, strict=self.spec.strict, @@ -374,7 +377,7 @@ def run(self, store, size): dwdp_rank=self.index, dwdp_size=size, device_id=0, - moe_layers=list(COVERED), + moe_layers=list(COVERED) if self.spec.attached else None, config=config, store_factory=self.store_factory, ops=self.ops, @@ -502,6 +505,18 @@ def make(specs): # --------------------------------------------------------------------------- +@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)]) @@ -553,6 +568,13 @@ def test_a_leaked_original_is_reported_after_the_commit_point(make_group, strict {"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"), } @@ -562,7 +584,7 @@ def test_ranks_that_disagree_stop_at_the_first_collective(make_group, required, specs[2] = replace(specs[2], **change) group = make_group(specs).run() group.assert_collectives(1) - group.assert_abandoned(pattern, raised=required) + group.assert_abandoned(pattern, raised=required or "requested" in change) assert all(rank.stores == [] for rank in group.ranks) diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py index 99d645afc380..d4b5ef0d604e 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py @@ -12,27 +12,22 @@ # 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 construction tests for DWDP: manager ownership and the dense DWDP launch gate.""" +"""PyExecutor construction ownership tests for DWDP.""" -import os -import re from contextlib import contextmanager from unittest.mock import MagicMock, patch import pytest -from tensorrt_llm._torch.modules.dwdp import vmm -from tensorrt_llm._torch.modules.dwdp.dense.config import DENSE_ENV -from tensorrt_llm._torch.modules.dwdp.dense.gate import validate_dense_gate +from tensorrt_llm._torch.modules.dwdp.dense.manager import ( + get_global_dense_dwdp_manager, + set_global_dense_dwdp_manager, +) 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.llmapi.llm_args import DwdpConfig, TorchLlmArgs -from tensorrt_llm.mapping import Mapping pytestmark = pytest.mark.cpu_only -_HANDLE_TYPE = vmm.cuda.CUmemAllocationHandleType - @contextmanager def _restore_global_dwdp_manager(): @@ -87,56 +82,56 @@ def fail_without_registering(**_kwargs): assert get_global_dwdp_manager() is existing -# ---------------------------------------------------------------------- -# Dense DWDP launch gate -# ---------------------------------------------------------------------- - - -@pytest.fixture(autouse=True) -def _no_dense_switches(monkeypatch): - """Start every test without ``TRTLLM_DWDP_DENSE*`` variables; tests set the ones they need.""" - for name in [name for name in os.environ if name.startswith(DENSE_ENV)]: - monkeypatch.delenv(name) +@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) -@pytest.fixture -def fabric_host(monkeypatch): - """Make this host share FABRIC handles, whatever its architecture.""" - monkeypatch.setattr(vmm, "peer_handle_type", lambda: _HANDLE_TYPE.CU_MEM_HANDLE_TYPE_FABRIC) +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") -def _dense_launch(**overrides) -> TorchLlmArgs: - """A dense DWDP context-worker launch, with ``overrides`` applied without validation.""" - args = TorchLlmArgs( - model="/tmp/dummy_model", - skip_tokenizer_init=True, - cuda_graph_config=None, - disable_overlap_scheduler=True, - dwdp_config=DwdpConfig( - dwdp_size=16, num_groups=1, num_experts_per_worker=56, num_prefetch_experts=56 - ), - ) - return args.model_copy(update=overrides) + 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 _mapping(tp_size: int = 1, pp_size: int = 1, cp_size: int = 1) -> Mapping: - return Mapping( - world_size=tp_size * pp_size * cp_size, - rank=0, - tp_size=tp_size, - pp_size=pp_size, - cp_size=cp_size, - ) +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 test_dense_gate_accepts_a_dense_dwdp_context_worker(fabric_host): - validate_dense_gate(_dense_launch(), _mapping()) + 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()) -@pytest.mark.parametrize( - "dwdp_config, actual", - [(None, "dwdp_config is not set"), (DwdpConfig(dwdp_size=1), "dwdp_size=1")], -) -def test_dense_gate_requires_a_dwdp_group(fabric_host, dwdp_config, actual): - with pytest.raises(ValueError, match=re.escape(f"dwdp_size > 1, but {actual}")): - validate_dense_gate(_dense_launch(dwdp_config=dwdp_config), _mapping()) + existing.__exit__.assert_not_called() + assert get_global_dense_dwdp_manager() is existing diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py index 76275df692e2..94e8d57b79fc 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py @@ -18,7 +18,7 @@ - creates a DWDP MPI sub-communicator - owns the Single Source of Truth for registered MoE layer indices - forwards runtime calls to the DWDPWeightManager produced by setup_dwdp -- owns the dense DWDP runtime: sets it up after setup_dwdp +- owns the attached DenseDwdpManager: sets dense DWDP up after 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. @@ -34,7 +34,12 @@ 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.types import DenseSetupError +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, DenseSetupError from tensorrt_llm._torch.pyexecutor.dwdp import ( DwdpManager, get_global_dwdp_manager, @@ -96,20 +101,26 @@ def _set_dense_env(test: unittest.TestCase, value: str | None) -> 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() # ------------------------------------------------------------------ @@ -218,10 +229,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 @@ -273,6 +313,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) @@ -478,15 +519,15 @@ def __init__(self) -> None: class TestDwdpManagerDenseSetup(unittest.TestCase): - """A DwdpManager over fakes, not entered unless a test enters it.""" + """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): - set_global_dwdp_manager(None) - self.addCleanup(set_global_dwdp_manager, None) + _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",)) @@ -545,6 +586,16 @@ def dense_setup(model, weight_manager, comm, **kwargs): self.assertEqual(self.setup_dwdp.call_args.kwargs["device_id"], self.DEVICE_ID) self.assertIsNone(self.manager.dense) + def test_setup_hands_the_committed_runtime_and_the_weight_manager_to_the_dense_manager(self): + dense_manager, dense = self.manager.dense_manager, object() + attach = _start(self, patch.object(dense_manager, "attach", wraps=dense_manager.attach)) + self.setup_dense.return_value = dense + + self.manager.setup(self.model) + + attach.assert_called_once_with(dense, expert=self.weight_manager) + self.assertIs(self.manager.dense, 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. 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 6c066f2c3ef2..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,7 +117,8 @@ def __init__( ) self.mgr = object.__new__(DwdpManager) self.mgr._weight_manager = self.wm - self.mgr._dense = None # expert DWDP alone + # 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_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/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 From c7b363d18978004c6591f72efa0c8d626e278f5c Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:27:54 -0700 Subject: [PATCH 4/9] [None][fix] DWDP: the only executor gets the DwdpManager when KV-cache estimation is off With KV-cache estimation off, the first PyExecutor the creator builds is the final one. That happens with TRTLLM_SKIP_KV_CACHE_ESTIMATION=1; with context parallelism, VANILLA attention or an encoder-decoder model; and with a model without a KV cache. Only the second, post-estimation executor got dwdp_manager. So with expert DWDP, whether or not dense DWDP is attached, the final executor: - did not start its forwards through the manager: _prime_dwdp found none, so the non-overlap loop never called DwdpManager.start_forward. The first forward, and the first after a forward that stopped inside the expert sequence, issued their first expert copies only at their first bind instead of ahead of it; every other forward bound the copies the previous forward's last MoE layers had issued. With dense DWDP committed, the model's forward entry (H0) starts the dense and the expert prefetch itself. - did not drop the prefetch after a control action: the executor calls invalidate_prefetch on the manager it holds, so the copies the last forward had issued for the next one kept the expert weights from before the action. - never exited the manager at shutdown. Dense DWDP was not released, expert DWDP was not torn down and the DWDP group was not freed. Both managers stayed registered, so the process refused a later DwdpManager. The first executor now gets dwdp_manager=dwdp_manager if not estimating_kv_cache else None, as it already gets the standalone dense manager. - With the estimation on, nothing changes: the estimation executor still gets None, so its shutdown in configure_kv_cache_capacity exits nothing. - With it off, the only executor starts every forward through the manager (_prime_dwdp -> start_forward), drops its prefetch after a control action and exits it once at shutdown. - Behaviour change with it off: like the final executor with the estimation on, the only executor now refuses expert DWDP with the overlap scheduler on (a ValueError at construction, rolled back). Before, it was built without the manager. - A construction error after the executor got the manager is still rolled back once by create_py_executor. Unchanged: with the estimation off, expert DWDP setup still runs after the final KV cache is sized. Tests (test_dwdp_creator.py): _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 PyExecutors built without __init__, for expert-only and attached DWDP, with the estimation on and off and for a model without a KV cache. The manager is set up once. - Estimation off, or no KV cache: the only executor gets the DwdpManager and starts its forwards through it (_prime_dwdp -> start_forward). - Estimation on: unchanged, only the final executor gets the manager and starts its forwards through it; the estimation executor gets none, so its shutdown exits none. That an executor's shutdown exits the manager it holds is tested in test_dense_executor.py. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- .../_torch/pyexecutor/py_executor_creator.py | 4 +- .../_torch/modules/dwdp/test_dwdp_creator.py | 139 ++++++++++++++++++ 2 files changed, 142 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7ef528916ba0..2cd9a9a39d86 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1307,7 +1307,9 @@ def allocation_scope(current_stage: ExecutorMemoryType): 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 standalone dense DWDP manager and exits it at shutdown. + # 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, diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py index d4b5ef0d604e..964d20515437 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py @@ -15,16 +15,30 @@ """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.modules.dwdp.dense.config import DENSE_ENV from tensorrt_llm._torch.modules.dwdp.dense.manager import ( 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 @@ -135,3 +149,128 @@ def fail_without_registering(**_kwargs): 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 _create(monkeypatch, mode: str, kv_cache: str) -> 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). + + 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): + 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=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 + ) + executor = py_executor_creator.create_py_executor(llm_args=llm_args, checkpoint_dir=None) + executor._prime_dwdp() + (manager,) = managers + return log, manager + + +@pytest.mark.parametrize("kv_cache", ["estimated", "not_estimated", "none"]) +@pytest.mark.parametrize("mode", ["expert", "attached"]) +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. + """ + log, manager = _create(monkeypatch, mode, kv_cache) + + final = ("executor", manager, None) + final_kv_cache = ("build_managers", False) + expected = { + "estimated": [("build_managers", True), "setup", ("executor", None, None), final_kv_cache], + "not_estimated": [final_kv_cache, "setup"], + "none": ["setup"], + }[kv_cache] + assert log == ["gate"] * (mode != "expert") + expected + [final, "start_forward"] From 9bc1d573ad5fa6b4df39646d5b597ddb6a82aa87 Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:27:59 -0700 Subject: [PATCH 5/9] [None][fix] DWDP: commit attached dense DWDP before the expert prefetch ring is allocated A DWDP8 context-only run (112 resident experts per rank, attached dense DWDP) died at startup on every rank: "Failed to allocate page 158 in slot 1: CUDA_ERROR_OUT_OF_MEMORY" in PagePool.create. DwdpManager.setup ran the whole expert setup first, so the expert prefetch ring (two slots of 871-872 pages of 16 MiB, 27.2 GiB) was allocated while the replicated dense weights were still resident: 16.09 GiB was free after the transport. Dense DWDP, which frees them, ran only after. DwdpManager.setup now sets dense DWDP up first: - setup_dense_dwdp runs before setup_dwdp, on every rank of a model with dense roots, over the same group, ranks, device and MoE layers. One INFO line precedes it: its collective #1 is now the group's first collective of the setup and waits for every rank. Its transfer is an _AttachedTransfer: a copy stream made on the first read (new_expert_copy_stream, in P2), and use_gather/use_batched_prefetch computed as DWDPWeightManager computes them. Dense setup needs nothing else from the expert setup: the dense roots never reach the routed experts, and the transport, the ring and the backend fixup do not change what dense DWDP collects. - setup_dwdp and DWDPWeightManager take an optional copy_stream: the weight manager adopts the stream dense DWDP already runs on, so the dense and the expert copies still share one copy stream. A check refuses a weight manager on another stream or with other transfer flags, and, on a repeated setup where dense DWDP ran on the weight manager itself, any other weight manager. - setup_dwdp also gets the committed runtime's slot storages (reserved_storages) and refuses an expert weight in one of them before the transport. In the new order a dense root that reached a routed-expert weight would otherwise make the transport copy a slot as the experts and free it under the dense runtime. - attach(runtime, expert=weight_manager) comes last, as before. If the expert setup or a check fails after the commit, the unattached runtime is released and the error propagates. - On a model without dense roots, the refusal of a TRTLLM_DWDP_DENSE request stays after the expert setup. A repeated setup for the same model after dense DWDP stayed off runs dense setup on the weight manager, as before. - With a store rendezvous, mark_setup_done stays the last step of setup, after the dense and the expert collectives. The setup peak becomes the larger of (replicated dense + dense store + slots) and (dense shard + ring) instead of their sum. The state after setup, and so the KV-cache estimate, is unchanged. Also: - WeightBuffer.create compares the ring, the page pool of the expert prefetch slots (DwdpConfig.num_prefetch_slots, two by default), with the free memory (mem_get_info after empty_cache) before PagePool allocates a page, and raises PrefetchRingMemoryError, a picklable RuntimeError naming the sizes of every slot and the remedies (dense DWDP for models with dense roots; fewer prefetch slots when there are more than two), instead of failing in cuMemCreate partway through. A ring that fits is allocated as before. This applies with expert DWDP alone too. "[WeightBuffer] Created" now names the ring, and a new "[DWDP Setup] ... memory after weight buffer" line gives the free memory after it. - PagePool no longer releases its handles twice when a page allocation fails: its failure path released them and left them in place for __del__ to release again. It now releases them once, through release(), which logs a handle that fails to release and goes on, so no other handle leaks and the allocation error stays the one raised. Unchanged: the forward-time schedule and events, the teardown order (dense first), the dense verdicts, both KV-estimation flows, expert-only (a dense setup that is off creates no stream, so the weight manager makes its own) and standalone. Changed, attached on a model with dense roots: the dense collectives run before the expert ones; a REQUIRED=1 dense failure raises before any expert setup; the [DWDP Dense] lines come before "[DWDP Setup] ... memory before setup". Tests: - test_dwdp_manager.py: dense setup runs on the attached transfer before the expert setup and before the model reference is kept, and setup_dwdp gets no copy stream when dense DWDP stays off. The expert setup gets the copy stream dense setup created and the committed slots' storages as reserved_storages; when it then fails (PrefetchRingMemoryError), the unattached runtime is released and the same error propagates; it also gets num_prefetch_slots and the repeated layers. A weight manager on another copy stream, or with another use_batched_prefetch, is refused, and the unattached runtime is released. A golden trace covers the attached lifecycle from setup through the executor's start and H0, and the recovery after an aborted forward, to teardown. - test_dwdp_prefetch_pipeline.py: a given copy stream is used and none is created; a copy stream on another device is refused. - test_dwdp_setup_lifecycle.py: an expert weight in a reserved storage is refused before the transport; a ring larger than the free memory is refused before any allocation, with two and with three slots; a page pool whose allocation fails releases every handle once. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- .../_torch/modules/dwdp/dense/manager.py | 17 +- .../_torch/modules/dwdp/dense/plan.py | 7 +- .../_torch/modules/dwdp/dense/runtime.py | 4 +- .../_torch/modules/dwdp/dense/setup.py | 22 +- tensorrt_llm/_torch/modules/dwdp/page_pool.py | 15 +- tensorrt_llm/_torch/modules/dwdp/setup.py | 49 ++++ .../_torch/modules/dwdp/weight_buffer.py | 103 +++++++- .../_torch/modules/dwdp/weight_manager.py | 20 +- tensorrt_llm/_torch/pyexecutor/dwdp.py | 226 +++++++++++++++--- .../_torch/modules/dwdp/test_dwdp_manager.py | 163 ++++++++++--- .../dwdp/test_dwdp_prefetch_pipeline.py | 42 ++++ .../modules/dwdp/test_dwdp_setup_lifecycle.py | 132 +++++++++- 12 files changed, 702 insertions(+), 98 deletions(-) diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/manager.py b/tensorrt_llm/_torch/modules/dwdp/dense/manager.py index f14f9fa14d8e..06aa3752981f 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/manager.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/manager.py @@ -20,14 +20,15 @@ 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 with the expert weight manager as the - transfer, 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. + 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 diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/plan.py b/tensorrt_llm/_torch/modules/dwdp/dense/plan.py index 6224fd5b5b45..49249f03b311 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/plan.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/plan.py @@ -366,9 +366,10 @@ def batched_copies(plan: LayerCopyPlan, *, slice_bytes: int = BATCH_SLICE_BYTES) 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 the expert DWDP weight manager's, so the dense transfer follows - the expert one; standalone they are ``manager.StandaloneTransfer``'s, which - ``TRTLLM_DWDP_DENSE_COPY`` picks. + 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 diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py b/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py index 6f4539e43a05..a971d1aba243 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/runtime.py @@ -165,8 +165,8 @@ class DenseDWDP: 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 - (``DWDPWeightManager.copy_stream``), the dense manager's own stream standalone - (``manager.StandaloneTransfer.copy_stream``). + (``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 diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/setup.py b/tensorrt_llm/_torch/modules/dwdp/dense/setup.py index 85db39641b61..d76c99f3c1cd 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/setup.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/setup.py @@ -15,12 +15,12 @@ """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`` right after the expert DWDP setup returned (with expert DWDP, so the expert -transport, buffers and backend fixups are done), 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 (design 3.3.3). +``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 @@ -98,8 +98,9 @@ def Get_size(self) -> int: class DenseTransfer(Protocol): """Where and how the dense copies run. - ``DWDPWeightManager`` (with expert DWDP: the expert DWDP copy stream and transfer) and - ``manager.StandaloneTransfer`` (standalone: the dense manager's own stream and + ``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. @@ -314,8 +315,9 @@ def setup_dense_dwdp( 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: the expert DWDP weight manager with expert DWDP, - a ``manager.StandaloneTransfer`` standalone. + 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``. diff --git a/tensorrt_llm/_torch/modules/dwdp/page_pool.py b/tensorrt_llm/_torch/modules/dwdp/page_pool.py index 84c09788a68c..a3d32d5e113e 100644 --- a/tensorrt_llm/_torch/modules/dwdp/page_pool.py +++ b/tensorrt_llm/_torch/modules/dwdp/page_pool.py @@ -156,15 +156,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..df769d8b2982 100644 --- a/tensorrt_llm/_torch/modules/dwdp/setup.py +++ b/tensorrt_llm/_torch/modules/dwdp/setup.py @@ -85,6 +85,37 @@ def _release_during_rollback(resource, what: str) -> None: ) +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 +127,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. @@ -132,6 +165,14 @@ 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. @@ -171,6 +212,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}" @@ -251,6 +293,12 @@ def setup_dwdp( logger.info( f"[DWDP Setup] Weight buffer created successfully ({num_prefetch_slots} prefetch 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...") @@ -282,6 +330,7 @@ def setup_dwdp( dwdp_size=mapping.dwdp_size, contention_opt=contention_opt, repeated_layer_indices=repeated_layer_indices, + copy_stream=copy_stream, ) except BaseException: _release_during_rollback(weight_buffer, "weight buffer") diff --git a/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py b/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py index 42ac02afb167..64e3884bbfd6 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 diff --git a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py index 8596f57ee378..7e5aa8c3afa9 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. @@ -162,10 +163,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}") @@ -190,7 +196,15 @@ def __init__( # 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)} diff --git a/tensorrt_llm/_torch/pyexecutor/dwdp.py b/tensorrt_llm/_torch/pyexecutor/dwdp.py index 946ef89f57d4..cac3b11802bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/dwdp.py +++ b/tensorrt_llm/_torch/pyexecutor/dwdp.py @@ -26,8 +26,10 @@ * 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 and hands it to the dense manager together with the weight - manager; each forward starts the dense prefetch ahead of the expert + 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 @@ -71,6 +73,7 @@ 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 @@ -87,6 +90,106 @@ 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. @@ -113,7 +216,9 @@ class DwdpManager: 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. + 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__( @@ -376,38 +481,59 @@ 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. + """Set dense DWDP, then the VA pipeline of expert DWDP, up on ``model``. - Then, if the model's decoder layers provide dense roots, runs + 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``. - The committed runtime, if any, goes to the dense manager together with - the weight manager (``DenseDwdpManager.attach``) and becomes - ``dense``; dense setup errors propagate. 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 (``refuse_dense_dwdp_without_roots``: it raises unless its - ``TRTLLM_DWDP_DENSE_REQUIRED`` is off). + 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. + 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. + 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. DenseSetupError: Dense DWDP setup failed, or was requested on a model without dense roots, with ``TRTLLM_DWDP_DENSE_REQUIRED`` on. + PrefetchRingMemoryError: The expert prefetch ring does not fit in + the free device memory. """ if self.dwdp_group is None: raise RuntimeError("DwdpManager.setup() called after cleanup()") @@ -426,34 +552,64 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: ) 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) - if dense_roots_available(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, - self._weight_manager, + 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) - elif dense_dwdp_requested(): + 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( diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py index 94e8d57b79fc..95d189c7d023 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py @@ -18,7 +18,7 @@ - creates a DWDP MPI sub-communicator - 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 after 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. @@ -26,22 +26,27 @@ 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.pyexecutor.dwdp import ( DwdpManager, + _AttachedTransfer, get_global_dwdp_manager, set_global_dwdp_manager, ) @@ -422,6 +427,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.""" @@ -511,6 +519,44 @@ 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__() @@ -533,7 +579,7 @@ def setUp(self): 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 = object() # what setup_dwdp returns + self.weight_manager = _FakeWeightManager(self.log) self.setup_dwdp = _start( self, patch(f"{_FAC}.setup_dwdp", return_value=self.weight_manager) ) @@ -560,41 +606,87 @@ def new_manager(self) -> DwdpManager: manager.add_layer(layer) return manager - def test_setup_runs_dense_setup_when_roots_are_available_but_not_requested(self): - """Collective #1 runs on every rank of a model with roots, requested or not. + 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. - It runs after the expert setup, once the model reference is kept. + 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. """ - seen = [] + 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, weight_manager, comm, **kwargs): - seen.append((self.manager._model_ref(), self.manager._weight_manager)) + 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, - self.weight_manager, + 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, [(self.model, self.weight_manager)]) - self.assertEqual(self.setup_dwdp.call_args.kwargs["device_id"], self.DEVICE_ID) - self.assertIsNone(self.manager.dense) + 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) - def test_setup_hands_the_committed_runtime_and_the_weight_manager_to_the_dense_manager(self): - dense_manager, dense = self.manager.dense_manager, object() - attach = _start(self, patch.object(dense_manager, "attach", wraps=dense_manager.attach)) - self.setup_dense.return_value = 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.manager.setup(self.model) + 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()) - attach.assert_called_once_with(dense, expert=self.weight_manager) - self.assertIs(self.manager.dense, dense) + 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. @@ -630,16 +722,33 @@ def test_a_model_without_roots_makes_no_dense_collective(self): 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_dense_setup_error_propagates_and_cleanup_still_tears_down(self): - self.setup_dense.side_effect = DenseSetupError("dense setup abandoned") - - with self.assertRaisesRegex(DenseSetupError, "dense setup abandoned"): - self.manager.setup(self.model) - self.assertIsNone(self.manager.dense) + def test_setup_priming_abort_recovery_and_teardown_keep_their_calls(self): + """The attached lifecycle against one call log. - self.manager.cleanup() + 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 - self.assertEqual(self.log, [("teardown_dwdp",), ("group.Free",)]) + 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 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..8cc2425e2507 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(): @@ -79,7 +88,7 @@ def _setup_mocks(stack: ExitStack): return transport, weight_buffer, weight_manager -def _run_setup(): +def _run_setup(**kwargs): return dwdp_setup_module.setup_dwdp( model=_model(), mapping=_mapping(), @@ -89,6 +98,7 @@ def _run_setup(): num_experts_per_worker=2, num_prefetch_experts=2, num_prefetch_slots=2, + **kwargs, ) @@ -153,3 +163,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 From b2c698e2ade82e1f6ed59050b90ae9d909f13d3c Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:28:01 -0700 Subject: [PATCH 6/9] [None][fix] DWDP: set attached dense DWDP up before the final KV cache when KV-cache estimation is off With KV-cache estimation off (TRTLLM_SKIP_KV_CACHE_ESTIMATION=1, among other cases), build_managers sizes the final KV cache from the device memory free when it runs. Standalone dense DWDP was already set up before it. Attached dense DWDP was not: it is set up inside DwdpManager.setup, which ran after the final KV cache was sized, so the replicated dense weights it frees never reached the KV cache. On DWDP16 at free_gpu_memory_fraction 0.7 that is about 14 GiB less KV per rank; on DWDP8 at the default fraction 0.9 the launch fails, because too little memory is left for dense setup. When the gate chose a dense mode and the estimation is off, the creator now sets the dense owner up right after try_prepare_estimation() and before build_managers: DwdpManager.setup attached (dense DWDP, then the expert setup), DenseDwdpManager.setup standalone. dwdp_set_up replaces dense_dwdp_set_up and skips the later setup. Unchanged: with the estimation on, the manager is still set up before the estimation run measures the memory; expert DWDP without a dense request keeps its place after the final KV cache, so its prefetch ring still does not count against the KV cache. Tests (test_dwdp_creator.py): with the estimation off, the only executor's DwdpManager is set up before build_managers with dense DWDP attached, and after it with expert DWDP alone. The standalone manager is set up once: before build_managers without the estimation, before the estimation run with it; only the final executor gets it, and starts its forwards with prefetch_first_layers. With TRTLLM_DWDP_DENSE_COPY=gather, the real standalone manager's SM reserve is in place when the model is built, before its setup, and is dropped when it exits. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- .../_torch/pyexecutor/py_executor_creator.py | 34 ++++++----- .../_torch/modules/dwdp/test_dwdp_creator.py | 59 ++++++++++++++++--- 2 files changed, 69 insertions(+), 24 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 2cd9a9a39d86..7dbec12299d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1174,8 +1174,8 @@ def allocation_scope(current_stage: ExecutorMemoryType): resources = {} estimating_kv_cache = False kv_cache_creator = None - # Whether standalone dense DWDP was set up before the KV caches (below). - dense_dwdp_set_up = False + # 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. @@ -1233,13 +1233,16 @@ 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 standalone dense DWDP is set up first: - # 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. - if dense_dwdp_manager is not None and not estimating_kv_cache: - dense_dwdp_manager.setup(model_engine.model) - dense_dwdp_set_up = True + # 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 @@ -1248,12 +1251,13 @@ def allocation_scope(current_stage: ExecutorMemoryType): # DWDP setup: MNNVL handle exchange + composite VA weight buffer + # weight manager + MoE backend fixup (single entry point). Dense DWDP is set - # up there with expert DWDP, else here when it runs standalone (unless it - # was set up before the KV caches above). - if dwdp_manager is not None: - dwdp_manager.setup(model_engine.model) - elif dense_dwdp_manager is not None and not dense_dwdp_set_up: - dense_dwdp_manager.setup(model_engine.model) + # 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 diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py index 964d20515437..56b28dc6c013 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_creator.py @@ -20,8 +20,11 @@ import pytest -from tensorrt_llm._torch.modules.dwdp.dense.config import DENSE_ENV +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, ) @@ -163,13 +166,25 @@ def fail_without_registering(**_kwargs): _NOOPS += ("_adjust_torch_mem_fraction", "log_memory_usage") -def _create(monkeypatch, mode: str, kv_cache: str) -> tuple[list, object]: +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. @@ -208,6 +223,8 @@ def create_py_executor_instance(**kwargs): 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 ) @@ -220,7 +237,7 @@ def model_engine(attn_runtime_features, **_kwargs): fakes.update( validate_dense_gate=lambda *_args: log.append("gate") or dense_mode, DwdpManager=manager, - DenseDwdpManager=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, @@ -250,27 +267,51 @@ def model_engine(attn_runtime_features, **_kwargs): 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,) = managers + (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"]) +@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. + 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", manager, None) + 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": [final_kv_cache, "setup"], + "not_estimated": ["setup", final_kv_cache] if setup_first else [final_kv_cache, "setup"], "none": ["setup"], }[kv_cache] - assert log == ["gate"] * (mode != "expert") + expected + [final, "start_forward"] + 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 From 4631f23105028c397b355d3c32fddab535c06796 Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:28:03 -0700 Subject: [PATCH 7/9] [None][chore] Dense DWDP standalone: copy-mode guidance from the DEP16 validation A DEP16 context-only validation of standalone dense DWDP, with one 16-rank attention-DP instance, exercised both copy modes. The SM gather (TRTLLM_DWDP_DENSE_COPY=gather) ran without stalls or protocol trouble with TRTLLM_DWDP_GATHER_KERNEL=tma and TLLM_ROUTING_DISABLE_COOP_LAUNCH=1. The gather's texts now say so. The default stays batched and gets no warning. Every accept or reject decision stays the same. - A standalone gather launch's warning no longer calls the gather experimental. It names the settings the gather was validated with, says that setup applies TLLM_ROUTING_DISABLE_COOP_LAUNCH=1 unless it is set (fabric_gather.apply_compute_side_defaults), and why cooperative routing matters. - The TRTLLM_DWDP_DENSE_COPY entry in dense/config.py names the validated settings too, and the gate's docstring says that the warning names them. Tests (test_dense_gate.py): a launch that breaks one of the gate's requirements is rejected with that requirement named, one launch per requirement. Those of both modes, alternating between the modes: no CUDA graphs, torch_compile_config, prefill CUDA graphs, sleep_config, lora_config or speculative_config; locality domains off; pp_size 1 and cp_size 1; FABRIC peer handles. Attached also: tp_size 1 and the overlap scheduler off. Standalone also: the MPI backend and, for TRTLLM_DWDP_DENSE_COPY=gather, the gather op. A standalone gather launch passes with one warning, which names the settings the gather was validated with. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- .../_torch/modules/dwdp/dense/config.py | 8 ++- .../_torch/modules/dwdp/dense/gate.py | 10 +-- .../_torch/modules/dwdp/test_dense_gate.py | 71 ++++++++++++++++++- 3 files changed, 81 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/config.py b/tensorrt_llm/_torch/modules/dwdp/dense/config.py index 0d25a871d528..bed4e8e240ea 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/config.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/config.py @@ -34,9 +34,11 @@ 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; experimental standalone: it holds SMs next to - the MoE all-to-all). 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. + ``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 diff --git a/tensorrt_llm/_torch/modules/dwdp/dense/gate.py b/tensorrt_llm/_torch/modules/dwdp/dense/gate.py index bf11ef0cac3c..0c19128bbc7e 100644 --- a/tensorrt_llm/_torch/modules/dwdp/dense/gate.py +++ b/tensorrt_llm/_torch/modules/dwdp/dense/gate.py @@ -83,7 +83,7 @@ def validate_dense_gate(llm_args: TorchLlmArgs, mapping: Mapping) -> DenseMode: 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: the gather is experimental without expert DWDP. + one warning, which names the settings the gather was validated with. Args: llm_args: The launch configuration. @@ -173,9 +173,11 @@ def _check_standalone(copy_mode: CopyMode) -> None: "this build does not register it", ) logger.warning( - f"{LOG_PREFIX} {DENSE_COPY_ENV}=gather is experimental without expert DWDP: the gather " - f"kernel holds SMs next to the MoE all-to-all, a combination not validated yet; batched, " - f"the default, runs on the copy engines" + 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" ) diff --git a/tests/unittest/_torch/modules/dwdp/test_dense_gate.py b/tests/unittest/_torch/modules/dwdp/test_dense_gate.py index 055168e9fbb2..84ec7f7d18ea 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dense_gate.py +++ b/tests/unittest/_torch/modules/dwdp/test_dense_gate.py @@ -34,7 +34,16 @@ 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.llmapi.llm_args import DwdpConfig, TorchLlmArgs +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 @@ -131,6 +140,66 @@ def check(self, monkeypatch) -> DenseMode: "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"}, + ), } From a4a3c94c7eeac5608cf68329d4a21636a51ad9a7 Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Mon, 26 Oct 2026 06:28:14 -0700 Subject: [PATCH 8/9] [None][feat] Expert DWDP: a balanced expert partition when dwdp_size does not divide the experts Expert DWDP gives every rank one contiguous range of routed experts, set by two DwdpConfig scalars, num_experts_per_worker (size) and num_prefetch_experts (stride). Setup required (dwdp_size - 1) * stride + size == num_experts with stride <= size, so when dwdp_size does not divide the expert count only overlapping configs passed. Kimi-K3 DWDP12 runs 82/74 and stores 88 duplicate experts per layer. On a K3 DWDP12 context server (attached dense DWDP, max_num_tokens 32768, KV cache fraction 0.70), the balanced partition below lowers each rank's profiling peak by 10.64-12.42 GiB and raises its KV cache by 7.45-8.70 GiB against 82/74. The VMM does not need equal ranges: a handle maps only whole, whatever its size. The software assumed equal ranges in these places: - the range formula, computed four times; - each MoE backend, which allocated num_experts_per_worker rows; - every peer's handle size and view rows, taken from this rank's; - the fixup, which allgathered a small parameter only when its first dimension was num_experts_per_worker, so with unequal counts some ranks would skip a collective. ExpertPartition (dwdp/specs.py) is one table of every rank's [start, end). Every rank computes it from DwdpConfig, dwdp_size and the model's expert count N. It is the only source of any rank's range for DwdpManager, MoE._init_dwdp_expert_layout, setup_dwdp and the transport. - Both fields 0: balanced. With divmod(N, dwdp_size) == (base, rem), the ranks below rem store base + 1 experts and the others base, as _compute_ep_partition splits EP ranks (K3 DWDP12: 8 x 75 + 4 x 74). Needs N >= dwdp_size. - Both positive: strided, where rank r stores [r * stride, min(r * stride + size, N)). Accepted: - uniform: stride == size and dwdp_size * size == N; - redundant: stride < size and (dwdp_size - 1) * stride + size == N; the lowest rank storing an expert serves it; - capped: stride == size and (dwdp_size - 1) * size < N < dwdp_size * size; the last rank allocates only the experts left, with no padding. - Anything else raises ValueError: one field set or a negative value (DwdpConfig also rejects these when the config is parsed), a stride larger than the size, ranges that end before N, overlap past N, or an empty last rank. No config field is added. Both fields become NonNegativeInt, which the golden manifest records as int, so the manifest and the API references are unchanged. The two field descriptions change, and the serve JSON schemas (tensorrt_llm/schemas/trtllm-serve-config.schema.json and trtllm-serve-disagg-config.schema.json, regenerated with scripts/generate_trtllm_serve_schemas.py) gain minimum 0 for both. Annotations cannot express the pairing, so generate_serve_schema() mirrors DwdpConfig._validate_expert_partition in the DwdpConfig definition: both fields 0 (or absent), or both positive. Both schemas thus reject one field set alone or with the other at 0, as DwdpConfig does when the config is parsed. What reads the table: - DwdpManager.expert_partition(N), cached per N with no collective, and local_expert_range(N). MoE._init_dwdp_expert_layout sizes each backend with exactly its rank's rows, so an invalid config raises at model construction, before create_weights. initial_global_assignments is range(N) for every partition: only a load balancer reads it, and DWDP excludes one. - DWDPTransport.create takes the ExpertPartition. The transport allocates handle_geometry(expert_bytes, start, end, granularity) for the rank's range. It maps every peer's handle whole, with that peer's size and offset, and gives each peer view that peer's rows. - From the handshake to the vote, setup_dwdp collects first, then decides the same way on every rank. 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. - _agree_on_partition makes one allgather: N, every rank's range, the kind, the granularity, and for each (layer, weight) expert chunk its rows, per-expert shape and dtype. Each rank is compared with the most common report, so the error names the ranks that differ, rank 0 included (it lists at most 8 problems); a rank that sends no handshake is named too. - Each Phase-1 handle allgather carries (status, payload), and every rank checks the statuses before any import. The handles need no geometry check: every rank computes each handle's size and data offset with handle_geometry from what the handshake agreed on. Phase 3's Barrier becomes allgather(status). On POSIX_FD, Phase 2 opens each peer's pidfd inside the import, so that failure is reported in Phase 3 as well; a failed map or access grant releases its VA before the error is reported. - Steps 5-7 (the weight buffer and ring, the edge bytes, the weight manager) only record a local error. The fixup's one collective, _exchange_expert_scalars, is the group's vote. If any rank failed, that rank raises its own error and the others raise a RuntimeError naming it. Otherwise it assembles the per-expert alphas (and a per-expert gate bias) along the partition. It also requires overlapping copies, and per-layer parameters such as fc31_input_scale, to be bitwise equal on every rank. fixup_moe_backends makes no collective, so K3's 184 fixup allgathers become one. - If a verdict collective itself fails (a peer died, the store timed out), a rank that had failed raises its own error, with the collective's error as its context. The backend fixup after the vote is local: it re-materializes the activation constants (install_activation_params) for the full expert count, installs the exchanged tensors, and a failure there raises on that rank only. A rank that dies, or fails before the handshake, leaves its peers waiting until the rendezvous times out. - Removed: _validate_partition_config, _allgather_e_score_correction_bias, _allgather_expert_scales, _get_first_spec, WeightBuffer.compute_peer_data_offset, compute_peer_ranges, DWDPTransport.get_peer_ranges, and DwdpManager's start_expert_id and end_expert_id. Stale comments are fixed: ring pages are local NONE handles, the whole-handle rule applies to every VMM handle, the mapping limit, the composite-VA padding, and the edge bytes. Compatibility: explicit size/stride configs keep the same ranges, rows, handle sizes and offsets, peer views, ring pages, copy plans and fixup results, byte for byte. ExpertPartition.strided accepts exactly the old validator's configs plus capped tables, so the uniform and redundant configs of test_dwdp_aggregated.py, test_dwdp_disaggregated_serving.py and the Kimi-K3 perf-sanity configs (56/56) stay valid. For these configs only the following change: the collectives (19 -> 15 per rank on a 3-layer model), the log wording, the point where a config error is raised, and initial_global_assignments before the fixup (range(N), where a redundant config had a round-robin list); ranks that disagree now raise together. 0/0, which was rejected before (Kimi-K3 while the model is built, with an IndexError; DeepSeek-V3 in setup, after loading), now means balanced. Tests: - test_dwdp_peer_ranges.py: balanced equals _compute_ep_partition. On a 60,836-config grid, strided accepts exactly the configs whose ranges end at N (the old validator's set) plus capped tables, with the strided ranges. Also covers from_config accept and reject cases, invalid tables, and handle_geometry equal to PageAlignedLayout. - test_dwdp_manager.py: DwdpConfig validation; 0/0 resolved and cached per expert count; backend rows from _init_dwdp_expert_layout (74 on a balanced rank, 71 on the capped last rank, 82 on a redundant one); a config that does not fit raises at layout time. The existing range checks read local_expert_range and expert_partition instead of the removed start_expert_id and end_expert_id. - test_prefetch_copy_plan.py: with peers of 75/74, 45/44 and 38/37 experts, every descriptor reads its owner's rows, the remote experts are covered once, the batched plan moves the same bytes, and the gather plan has one row per peer group with all of the group's bytes. - tests/unittest/_torch/moe/test_dwdp_fixup_moe_backends.py: ported to the exchange. Every rank makes one allgather, and every inconsistency raises on every rank. A ConfigurableMoE-shaped wrapper around a backend set up by the real NVFP4CuteDslFusedMoEMethod.setup_quant_scales, on a 4 + 3 split: after the fixup both see every expert, quant_scales holds the full alphas, and a full-sized gate bias stays as it is. - test_dwdp_setup_lifecycle.py: patches the handshake, the granularity query and the vote. - test_trtllm_serve_config_schemas.py: the serve schema accepts 0/0 and rejects a negative num_experts_per_worker or num_prefetch_experts; the disaggregated schema rejects a negative one under context_servers. test_dwdp_expert_partition_matches_runtime checks eight cases against DwdpConfig and both schemas, the disaggregated one at the top level and under context_servers and generation_servers: both absent, 0/0, 1/1 and 2/1 are accepted; one field at 1 with the other absent or 0 is rejected. - New test_dwdp_setup_protocol.py (cpu_only): one thread per rank runs the real setup_dwdp over StoreDwdpGroup on a HashStore, with a fake VMM that maps handles only whole. The weight buffer, edge-byte fill and weight manager are stand-ins; those of the buffer and the manager record their arguments. - Every peer is mapped with its own size and rows (balanced, capped, uniform, redundant), steps 5-7 get the partition's local range and peer ranges, and the exchange gives every rank the full fc31_alpha (from 75/74 shards on the balanced DWDP12). - A failure on one rank raises on every rank after the same collective, and nothing leaks. Covered failures: the handshake (rows, partition, dtype), an export, the access grant of an imported handle, and the ring. - When rank 0 alone computed another partition, the handshake error names rank 0. Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- scripts/generate_trtllm_serve_schemas.py | 6 + tensorrt_llm/_torch/modules/dwdp/__init__.py | 6 +- tensorrt_llm/_torch/modules/dwdp/page_pool.py | 25 +- tensorrt_llm/_torch/modules/dwdp/setup.py | 1116 +++++++++-------- tensorrt_llm/_torch/modules/dwdp/specs.py | 298 ++++- tensorrt_llm/_torch/modules/dwdp/transport.py | 545 ++++---- .../_torch/modules/dwdp/weight_buffer.py | 45 +- .../_torch/modules/dwdp/weight_manager.py | 21 +- .../_torch/moe/fused_moe/interface.py | 94 +- tensorrt_llm/_torch/pyexecutor/dwdp.py | 73 +- tensorrt_llm/llmapi/llm_args.py | 36 +- tensorrt_llm/mapping.py | 11 +- .../schemas/trtllm-serve-config.schema.json | 32 +- .../trtllm-serve-disagg-config.schema.json | 32 +- .../test_dwdp_disaggregated_serving.py | 14 +- .../_torch/modules/dwdp/test_dwdp_manager.py | 60 +- .../_torch/modules/dwdp/test_dwdp_mapping.py | 2 +- .../modules/dwdp/test_dwdp_peer_ranges.py | 202 +-- .../modules/dwdp/test_dwdp_setup_lifecycle.py | 19 +- .../modules/dwdp/test_dwdp_setup_protocol.py | 345 +++++ .../modules/dwdp/test_prefetch_copy_plan.py | 64 +- .../moe/test_dwdp_fixup_moe_backends.py | 366 +++--- .../test_trtllm_serve_config_schemas.py | 45 + 23 files changed, 2234 insertions(+), 1223 deletions(-) create mode 100644 tests/unittest/_torch/modules/dwdp/test_dwdp_setup_protocol.py 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/modules/dwdp/__init__.py b/tensorrt_llm/_torch/modules/dwdp/__init__.py index dc9f36232542..882a15bc2693 100644 --- a/tensorrt_llm/_torch/modules/dwdp/__init__.py +++ b/tensorrt_llm/_torch/modules/dwdp/__init__.py @@ -45,13 +45,14 @@ 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 @@ -87,7 +88,8 @@ "PageAlignedLayout", "PeerRange", "PeerRanges", - "compute_peer_ranges", + "ExpertPartition", + "handle_geometry", "lookup_owner", # VMM utilities "align_up", diff --git a/tensorrt_llm/_torch/modules/dwdp/page_pool.py b/tensorrt_llm/_torch/modules/dwdp/page_pool.py index a3d32d5e113e..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. diff --git a/tensorrt_llm/_torch/modules/dwdp/setup.py b/tensorrt_llm/_torch/modules/dwdp/setup.py index df769d8b2982..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,23 @@ 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: @@ -144,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 @@ -178,6 +216,17 @@ def setup_dwdp( 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: @@ -220,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) @@ -236,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 @@ -260,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(), @@ -287,23 +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)." - ) - 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)" - ) + 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(), @@ -311,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(), @@ -332,30 +382,36 @@ def setup_dwdp( 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.") @@ -496,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``. @@ -585,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 # --------------------------------------------------------------------------- @@ -600,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. @@ -612,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 @@ -634,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] @@ -648,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( @@ -659,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) @@ -679,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. @@ -709,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 @@ -795,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], @@ -1052,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. """ @@ -1094,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. @@ -1116,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, @@ -1131,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( @@ -1247,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 64e3884bbfd6..74b8c3039b25 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py @@ -419,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) + - post_region : page-pool 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. - - 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) @@ -490,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 --- @@ -504,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: @@ -660,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 7e5aa8c3afa9..3ce78a52474a 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py @@ -142,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). @@ -187,10 +188,9 @@ 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. @@ -269,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}" ) 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/dwdp.py b/tensorrt_llm/_torch/pyexecutor/dwdp.py index cac3b11802bb..c9ff0ad50995 100644 --- a/tensorrt_llm/_torch/pyexecutor/dwdp.py +++ b/tensorrt_llm/_torch/pyexecutor/dwdp.py @@ -43,7 +43,7 @@ import traceback import weakref -from typing import List, Optional +from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn @@ -54,6 +54,7 @@ from tensorrt_llm._torch.modules.dwdp import ( DenseDWDP, DWDPWeightManager, + ExpertPartition, StoreDwdpGroup, dense_dwdp_requested, dwdp_global_rank, @@ -289,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). @@ -459,6 +452,47 @@ def cleanup(self) -> 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 # ------------------------------------------------------------------ @@ -528,12 +562,19 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: 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. + 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: The expert prefetch ring does not fit in - the free device memory. + 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()") 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/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/unittest/_torch/modules/dwdp/test_dwdp_manager.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py index 95d189c7d023..4d28bd596e11 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py @@ -16,6 +16,7 @@ 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 @@ -44,6 +45,7 @@ ) 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, @@ -162,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.""" @@ -177,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.""" @@ -189,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) # ------------------------------------------------------------------ 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_setup_lifecycle.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py index 8cc2425e2507..e2d9fb08b501 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py @@ -47,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() @@ -66,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) ) @@ -84,10 +79,20 @@ 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 _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(), 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_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/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 From 0d53826a6016fc0bf3012760b21a2147c65c800c Mon Sep 17 00:00:00 2001 From: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> Date: Wed, 28 Oct 2026 18:27:35 -0700 Subject: [PATCH 9/9] [None][fix] Release the attention-DP dummy's draft KV mirror when a skipped iteration rolls it back With KV cache manager V2, a separate draft KV cache and attention DP, KVCacheV2Scheduler admits the draft pool's mirror while it schedules the attention-DP pad dummy (_admit_unpaired_draft -> admit_mirror): 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 _finalize_adp_dummy_allocation rolls the dummy back before prepare_resources. It freed the dummy in the KV cache manager and the speculative resource manager, but not in the draft KV cache manager, so the empty mirror outlived the dummy. On the next queued iteration, the deferred previous batch, which still holds the earlier, finished dummy with the same ID, reaches the draft pool's update_resources, which resizes that empty cache to the old dummy's history and raises "History length cannot exceed capacity". A disaggregated DSpark generation server (overlap scheduler, attention DP) crashed this way one iteration after a skipped one. The rollback now also frees the dummy in the draft KV cache manager when KV cache manager V2 is used. Tests: - test_kv_cache_v2_adp_dummy_draft_rollback.py: with stand-in managers, a rolled-back dummy releases its draft mirror and a queued dummy keeps it (CPU); with real V2 target and draft pools and the V2 scheduler, a skipped iteration followed by the deferred previous batch no longer raises (GPU, registered in l0_a10). Signed-off-by: Zongfei Jing <20381269+zongfeijing@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 11 + .../integration/test_lists/test-db/l0_a10.yml | 1 + ...st_kv_cache_v2_adp_dummy_draft_rollback.py | 217 ++++++++++++++++++ 3 files changed, 229 insertions(+) create mode 100644 tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_adp_dummy_draft_rollback.py diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6b581b5a4d14..1dce7afd26cf 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4404,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): 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/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 == []