From 1c2f659097d4d50b63690af31e365673dc29f730 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:02:25 -0700 Subject: [PATCH 1/8] [None][perf] Kimi K3 KDA prefill: size the persistent grids to the usable SMs The FlashInfer Small-BH KDA prefill launched its persistent grid on every SM of the device. It now leaves out an explicit SM reserve: trtllm::kda_prefill (and KDAKernelDispatch.prefill_chunk_kda) take sm_reserve, the SMs a concurrent kernel keeps resident, and the KDA mixer passes the reserve it reads once at construction (cute_dsl_utils.sm_reserve()). The Small-BH/BT16 routing stays on the device SM count so it matches BT16's own router; the BT16 engine/decomposition kernels size their grids inside the vendored kernel and are unchanged. The reserve is an explicit value set once: DwdpManager computes it at construction (fabric_gather.resident_sms(): the gather's resident block count, or TRTLLM_CUTEDSL_SM_RESERVE where the deployment pins it) and sets it with cute_dsl_utils.set_sm_reserve() before the model is built. The CuTe DSL occupancy hook, the dense GEMM op and the KDA prefill read that value, and no module passes it through the environment. Blocks beyond the free SMs would otherwise wait behind the whole first wave. A reserve that leaves no SM is rejected. Rubin evidence (measured on the earlier CuTe DSL K123/K4 prefill, which the FlashInfer prefill replaced; hecate R200, 212 SMs; KDA layer prefill, H=96, bf16 state pool), outputs bitwise equal: T=16k -1.4%, T=7680 -6.3%, 2x16k -4.2% (L0 jobs 768150, 768235). Part of the context switch set measured on Rubin ctx-only DWDP16, mnt 16384, conc 128: uncached prefill tok/s/GPU 6,595 / 6,586 (jobs 772280, 781995) -> 6,957 (job 773435), +5.5%. Platform: S (shared kernels). Tests: test_kda_prefill_op.py::test_kda_prefill_op_sm_reserve_keeps_results (varlen and single sequence, BF16 Small-BH and FP32 auto; output and state pool bitwise equal across grid sizes; an all-SM reserve is rejected); test_fabric_gather.py (resident_sms, explicit pin, set_sm_reserve); test_dwdp_manager.py (test_enter_reserves_the_gather_sms); test_kda_prefill_op.py::test_kda_mixer_prefill_passes_the_reserve_set_before_construction. GB300 (OCI AGA) jobs 1369171, 1387340 (CuTe DSL K123/K4 prefill). Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../_torch/custom_ops/cute_dsl_custom_ops.py | 2 +- .../custom_ops/cute_dsl_kimi_k3_custom_ops.py | 11 ++- tensorrt_llm/_torch/cute_dsl_utils.py | 52 ++++++----- .../_torch/modules/dwdp/fabric_gather.py | 43 ++++++--- .../_torch/modules/dwdp/weight_manager.py | 8 +- .../_torch/modules/kimi_kda/_kda_kernels.py | 5 + .../_torch/modules/kimi_kda/kimi_kda_mixer.py | 5 + tensorrt_llm/_torch/pyexecutor/dwdp.py | 8 ++ .../_torch/modules/dwdp/test_dwdp_manager.py | 20 ++++ .../_torch/modules/dwdp/test_fabric_gather.py | 43 ++++++--- .../modules/kimi_kda/test_kda_prefill_op.py | 93 +++++++++++++++++++ 11 files changed, 236 insertions(+), 54 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py index d5a55e7e132c..54792dab7866 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py @@ -737,7 +737,7 @@ def _get_full_device_max_active_clusters(device_id: int, cluster_size: int, reserved_sms: int) -> int: """Return the cached full-device occupancy for a cluster shape. - ``reserved_sms`` (``TRTLLM_CUTEDSL_SM_RESERVE``, applied inside + ``reserved_sms`` (``cute_dsl_utils.sm_reserve()``, applied inside ``HardwareInfo.get_max_active_clusters``) is part of the key so a reserve set after the first query is not hidden by the cache. """ diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py index 5210664f89c3..a331f2747d42 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_custom_ops.py @@ -53,6 +53,7 @@ def kda_prefill( max_seqlen: int = 0, seq_order: torch.Tensor | None = None, cu_chunks: torch.Tensor | None = None, + sm_reserve: int = 0, ) -> torch.Tensor: """Update unique selected V-first state rows using packed raw Q/K/V. @@ -70,6 +71,11 @@ def kda_prefill( upper bound for all replays, or zero to use the total token capacity. Small-BH and BT16 decomposition consume aligned, head-contiguous beta views directly; the engine requires compact beta storage. + + ``sm_reserve`` SMs, held by a kernel resident for the whole step, are left + out of the Small-BH persistent grid; blocks beyond the free SMs would wait + behind the whole first wave. The routing above stays on the device SM count + so it matches BT16's own router. """ device = q.device if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3), (10, 7)): @@ -147,6 +153,8 @@ def kda_prefill( if max_seqlen < 0 or max_seqlen > q.shape[1]: raise ValueError("max_seqlen must be zero or a sequence-length bound within total tokens") sm_count = torch.cuda.get_device_properties(device).multi_processor_count + if not 0 <= sm_reserve < sm_count: + raise ValueError(f"sm_reserve={sm_reserve} must leave at least one of {sm_count} SMs") small_bh_eligible = state_pool.dtype == torch.bfloat16 and q.shape[1] > 1 if mode == "small-bh" and not small_bh_eligible: raise ValueError("Small-BH requires BF16 state and more than one packed token") @@ -193,7 +201,7 @@ def kda_prefill( state_input=state_pool, state_output=state_pool, state_indices=state_indices, - max_active_clusters=sm_count, + max_active_clusters=sm_count - sm_reserve, ) return output # Torch allocations belong to the CUDA graph's private memory pool during @@ -252,6 +260,7 @@ def _kda_prefill_fake( max_seqlen=0, seq_order=None, cu_chunks=None, + sm_reserve=0, ): return torch.empty_like(q) diff --git a/tensorrt_llm/_torch/cute_dsl_utils.py b/tensorrt_llm/_torch/cute_dsl_utils.py index 0cb30e8bb124..e6baaf9f0777 100644 --- a/tensorrt_llm/_torch/cute_dsl_utils.py +++ b/tensorrt_llm/_torch/cute_dsl_utils.py @@ -14,7 +14,6 @@ # limitations under the License. import math -import os import platform import sys import types @@ -77,24 +76,35 @@ def install_cutlass_dsl_compatibility() -> None: cute.make_fragment = cute.make_rmem_tensor -# ``TRTLLM_CUTEDSL_SM_RESERVE=`` (default 0) keeps ``n`` SMs out of every -# CuTeDSL persistent grid, so CTAs that would land on SMs another kernel holds -# resident do not run their static tile share as a second wave. The hook wraps -# ``HardwareInfo.get_max_active_clusters`` on the class, is installed when this -# module is imported (before any kernel queries occupancy), and reads the -# variable at each grid sizing; callers that memoize the occupancy must key on -# ``sm_reserve()`` (the dense GEMM custom op does) or a reserve set after their -# first query is lost. -_SM_RESERVE_ENV = "TRTLLM_CUTEDSL_SM_RESERVE" -_sm_reserve_announced = set() +# SMs that a kernel resident for the whole step holds (the DWDP fabric gather; +# ``DwdpManager`` sets it once, before the model is built). Persistent grids +# leave them out, so CTAs that would land on those SMs do not run their static +# tile share as a second wave. The hook below wraps +# ``HardwareInfo.get_max_active_clusters`` on the class (installed when this +# module is imported, before any kernel queries occupancy) so every CuTeDSL +# persistent grid honors the reserve. Kernels that size their own persistent +# grids take the reserve as an argument from their module, which reads +# ``sm_reserve()`` (the KDA mixer reads it at construction). Callers that +# memoize the occupancy key it on ``sm_reserve()`` (the dense GEMM custom op +# does). +_sm_reserve = 0 + + +def set_sm_reserve(reserve: int) -> None: + """Keep ``reserve`` SMs out of persistent grids from now on (0 = none).""" + global _sm_reserve + if reserve < 0: + raise ValueError(f"SM reserve must be non-negative, got {reserve}") + if reserve and reserve != _sm_reserve: + logger.info( + f"[cute_dsl_utils] persistent grids launch with {reserve} SMs reserved" + ) + _sm_reserve = reserve def sm_reserve() -> int: - """SMs to keep out of persistent grids, from ``TRTLLM_CUTEDSL_SM_RESERVE`` (0 = none).""" - try: - return max(0, int(os.environ.get(_SM_RESERVE_ENV, "0"))) - except ValueError: - return 0 + """SMs kept out of persistent grids (see :func:`set_sm_reserve`).""" + return _sm_reserve def reserved_max_active_clusters(full: int, cluster_size: int, @@ -125,15 +135,7 @@ def _install_sm_reserve() -> None: def get_max_active_clusters(self, cluster_size: int, *args, **kwargs) -> int: full = original(self, cluster_size, *args, **kwargs) - reserve = sm_reserve() - if reserve <= 0: - return full - if reserve not in _sm_reserve_announced: - _sm_reserve_announced.add(reserve) - logger.info( - f"[cute_dsl_utils] {_SM_RESERVE_ENV}={reserve}: persistent grids " - f"launch with {reserve} SMs reserved") - return reserved_max_active_clusters(full, cluster_size, reserve) + return reserved_max_active_clusters(full, cluster_size, sm_reserve()) get_max_active_clusters._trtllm_sm_reserve_hook = True HardwareInfo.get_max_active_clusters = get_max_active_clusters diff --git a/tensorrt_llm/_torch/modules/dwdp/fabric_gather.py b/tensorrt_llm/_torch/modules/dwdp/fabric_gather.py index 0c5a7738cbc8..43311d9189d8 100644 --- a/tensorrt_llm/_torch/modules/dwdp/fabric_gather.py +++ b/tensorrt_llm/_torch/modules/dwdp/fabric_gather.py @@ -97,26 +97,45 @@ def blocks_per_row(rows: int) -> int: return max(1, total_blocks() // max(rows, 1)) -#: Compute-side knobs the SM gather sets for itself unless the user pinned them. +#: SMs the persistent grids leave to the gather (``resident_sms``), if the user pins the count. _SM_RESERVE_ENV = "TRTLLM_CUTEDSL_SM_RESERVE" +#: Compute-side knobs the SM gather sets for itself unless the user pinned them. _DENSE_GEMM_SCHEDULER_ENV = "TRTLLM_CUTEDSL_DENSE_GEMM_SCHEDULER" _ROUTING_NO_COOP_ENV = "TLLM_ROUTING_DISABLE_COOP_LAUNCH" +def available() -> bool: + """True if the gather transport is requested and its op is registered.""" + 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``. + + ``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. + """ + pinned = _int_env(_SM_RESERVE_ENV, -1) + if pinned >= 0: + return pinned + if not available() or THREADS < 512: + return 0 + return total_blocks() + + def apply_compute_side_defaults() -> None: - """Tune the CuTeDSL persistent kernels for a resident gather; call before any kernel compiles. - - Reserves the gather's block count of SMs (``TRTLLM_CUTEDSL_SM_RESERVE``) and pins the SM107 - dense GEMM to the CLC dynamic scheduler (``TRTLLM_CUTEDSL_DENSE_GEMM_SCHEDULER``) so static - persistent grids do not run a second wave on the occupied SMs. Also turns off the - grid-cooperative MoE routing kernel (``TLLM_ROUTING_DISABLE_COOP_LAUNCH=1``), which needs every - SM co-resident and so would wait behind the gather. The routing code reads that knob once, at - its first launch, so this must run before the first forward. An explicit user setting wins. + """Tune the compute kernels for a resident gather; call before the first forward. + + Pins the SM107 dense GEMM to the CLC dynamic scheduler + (``TRTLLM_CUTEDSL_DENSE_GEMM_SCHEDULER``) so static persistent grids do not run a second wave + on the SMs the gather holds (the persistent grids themselves leave :func:`resident_sms` SMs + out, see ``DwdpManager``). Also turns off the grid-cooperative MoE routing kernel + (``TLLM_ROUTING_DISABLE_COOP_LAUNCH=1``), which needs every SM co-resident and so would wait + behind the gather. The routing code reads that knob once, at its first launch, so this must + run before the first forward. An explicit user setting wins. """ - reserve = str(total_blocks() if THREADS >= 512 else 0) applied = [] - if os.environ.setdefault(_SM_RESERVE_ENV, reserve) == reserve: - applied.append(f"{_SM_RESERVE_ENV}={reserve}") if os.environ.setdefault(_DENSE_GEMM_SCHEDULER_ENV, "clc_dynamic") == "clc_dynamic": applied.append(f"{_DENSE_GEMM_SCHEDULER_ENV}=clc_dynamic") if os.environ.setdefault(_ROUTING_NO_COOP_ENV, "1") == "1": diff --git a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py index c146cc0219e7..2ec645f18998 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py @@ -202,11 +202,11 @@ def __init__( # engine (see ``fabric_gather``). Opt-in; off if the op is not registered. # Layers whose descriptors are not 16-byte aligned cache None and fall back # to the copy engine. - self._use_gather = fabric_gather.is_enabled() and fabric_gather.gather_op() is not None + self._use_gather = fabric_gather.available() if self._use_gather: - # Reserve the gather's SMs in the persistent CuTeDSL grids, pin the dense - # GEMM to the CLC scheduler and turn off cooperative MoE routing unless the - # user set them explicitly. Runs during setup, before the first forward. + # Pin the dense GEMM to the CLC scheduler and turn off cooperative MoE + # routing unless the user set them explicitly. Runs during setup, before + # the first forward. The persistent grids' SM reserve is set by DwdpManager. fabric_gather.apply_compute_side_defaults() if fabric_gather.is_enabled() and not self._use_gather: logger.warning( diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index 555bf4b28cae..3d8119d3e35d 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -708,11 +708,15 @@ def prefill_chunk_kda( max_seqlen: int = 0, seq_order: Optional[torch.Tensor] = None, cu_chunks: Optional[torch.Tensor] = None, + sm_reserve: int = 0, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Run prefill on raw Q/K/V and gate/beta logits. FlashInfer updates indexed BF16/FP32 V-first state directly; FLA uses dense FP32 initial/final states. Both normalize Q/K internally. + + ``sm_reserve`` SMs are left out of the optimized kernels' persistent + grids for a kernel that runs concurrently with prefill. """ if state_pool is not None: if self.prefill_kernel_path != "optimized" or initial_state is not None: @@ -737,6 +741,7 @@ def prefill_chunk_kda( max_seqlen, seq_order, cu_chunks, + sm_reserve, ) return out, None diff --git a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py index 33f388bad319..fe06621c4cc2 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py @@ -30,6 +30,7 @@ from ....models.modeling_utils import QuantConfig from ....quantization import QuantAlgo from ...attention.backends import AttentionMetadata +from ...cute_dsl_utils import sm_reserve from ...distributed import AllReduce, AllReduceStrategy from ...model_config import ModelConfig from ...pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph @@ -310,6 +311,9 @@ def projection(name: str, in_features: int, out_features: int) -> nn.Module: use_optimized_decode=kernel_shape_ok, use_optimized_verify=kernel_shape_ok, ) + # SMs kept out of the prefill's persistent grids for a kernel resident + # for the whole step; fixed before the model is built. + self._prefill_sm_reserve = sm_reserve() self._fuse_prefill_conv = True @@ -885,6 +889,7 @@ def forward_prefill_core( max_seqlen=getattr(mamba_metadata, "kda_max_seqlen", 0), seq_order=getattr(mamba_metadata, "kda_seq_order", None), cu_chunks=getattr(mamba_metadata, "kda_cu_chunks", None), + sm_reserve=self._prefill_sm_reserve, ) # The packed convolution persisted the live convolution pool in place. diff --git a/tensorrt_llm/_torch/pyexecutor/dwdp.py b/tensorrt_llm/_torch/pyexecutor/dwdp.py index 78bd8b08c90b..2bcbebd92e3f 100644 --- a/tensorrt_llm/_torch/pyexecutor/dwdp.py +++ b/tensorrt_llm/_torch/pyexecutor/dwdp.py @@ -38,11 +38,13 @@ import torch.nn as nn from mpi4py.MPI import COMM_WORLD +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 ( DWDPWeightManager, StoreDwdpGroup, dwdp_global_rank, + fabric_gather, setup_dwdp, store_rendezvous_enabled, teardown_dwdp, @@ -156,6 +158,10 @@ def __init__( self.start_expert_id = self.dwdp_rank * self.num_prefetch_experts self.end_expert_id = self.start_expert_id + self.num_experts_per_worker + # SMs the expert gather holds for the whole step; persistent grids leave + # them out (``set_sm_reserve`` on ``__enter__``, before the model is built). + self.resident_sms = fabric_gather.resident_sms() + self.dwdp_group = self._create_dwdp_comm() # SSOT for MoE layer indices; populated by configurable_moe.add_layer() @@ -199,12 +205,14 @@ def __enter__(self) -> "DwdpManager": 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) return self def __exit__(self, exc_type, exc_val, exc_tb) -> bool: try: self.cleanup() finally: + set_sm_reserve(0) set_global_dwdp_manager(None) return False diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py index 0ed34c004890..9aaf54ac6463 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py @@ -28,6 +28,7 @@ import pytest +from tensorrt_llm._torch.cute_dsl_utils import sm_reserve from tensorrt_llm._torch.distributed import MPIDist from tensorrt_llm._torch.pyexecutor.dwdp import ( DwdpManager, @@ -195,6 +196,25 @@ def test_enter_registers_global(self, mock_comm_world, _rank, _setup): mgr.__exit__(None, None, None) self.assertIsNone(get_global_dwdp_manager()) + 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 + __enter__ (before the model is built) to __exit__.""" + _configure_mock_comm_world(mock_comm_world) + with patch( + "tensorrt_llm._torch.pyexecutor.dwdp.fabric_gather.resident_sms", return_value=30 + ): + mgr = DwdpManager( + config=_make_config(), + dist=MagicMock(spec=MPIDist), + mapping=_make_mapping(), + ) + self.assertEqual(mgr.resident_sms, 30) + self.assertEqual(sm_reserve(), 0) + mgr.__enter__() + self.assertEqual(sm_reserve(), 30) + mgr.__exit__(None, None, None) + self.assertEqual(sm_reserve(), 0) + def test_cleanup_frees_comm_and_idempotent(self, mock_comm_world, _rank, _setup): sub_comm = MagicMock() mock_comm_world.Create_group.return_value = sub_comm diff --git a/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py b/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py index 44b91a060f1c..5b7f9827c9d6 100644 --- a/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py +++ b/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py @@ -12,7 +12,7 @@ # 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 for the host side of the DWDP fabric gather (plan building and env knobs).""" +"""CPU tests for the host side of the DWDP fabric gather (plan building, knobs, SM reserve).""" import os @@ -78,28 +78,49 @@ def test_blocks_per_row_never_drops_below_one(monkeypatch): def test_compute_side_defaults_respect_explicit_values(monkeypatch): - monkeypatch.setattr(fabric_gather, "THREADS", 1024) - monkeypatch.setenv("TRTLLM_DWDP_GATHER_BLOCKS", "30") - _unset(monkeypatch, "TRTLLM_CUTEDSL_SM_RESERVE") monkeypatch.setenv("TRTLLM_CUTEDSL_DENSE_GEMM_SCHEDULER", "static") monkeypatch.setenv("TLLM_ROUTING_DISABLE_COOP_LAUNCH", "0") fabric_gather.apply_compute_side_defaults() - assert os.environ["TRTLLM_CUTEDSL_SM_RESERVE"] == "30" assert os.environ["TRTLLM_CUTEDSL_DENSE_GEMM_SCHEDULER"] == "static" assert os.environ["TLLM_ROUTING_DISABLE_COOP_LAUNCH"] == "0" - - -def test_small_blocks_do_not_reserve_sms(monkeypatch): - monkeypatch.setattr(fabric_gather, "THREADS", 128) - _unset(monkeypatch, "TRTLLM_CUTEDSL_SM_RESERVE") _unset(monkeypatch, "TRTLLM_CUTEDSL_DENSE_GEMM_SCHEDULER") _unset(monkeypatch, "TLLM_ROUTING_DISABLE_COOP_LAUNCH") fabric_gather.apply_compute_side_defaults() - assert os.environ["TRTLLM_CUTEDSL_SM_RESERVE"] == "0" assert os.environ["TRTLLM_CUTEDSL_DENSE_GEMM_SCHEDULER"] == "clc_dynamic" assert os.environ["TLLM_ROUTING_DISABLE_COOP_LAUNCH"] == "1" +def test_resident_sms_is_the_block_count_of_a_resident_gather(monkeypatch): + _unset(monkeypatch, "TRTLLM_CUTEDSL_SM_RESERVE") + monkeypatch.setenv("TRTLLM_DWDP_GATHER_BLOCKS", "30") + monkeypatch.setattr(fabric_gather, "available", lambda: True) + monkeypatch.setattr(fabric_gather, "THREADS", 1024) + assert fabric_gather.resident_sms() == 30 + # Small blocks share SMs with the compute kernels. + monkeypatch.setattr(fabric_gather, "THREADS", 128) + assert fabric_gather.resident_sms() == 0 + # No gather, no reserve. + monkeypatch.setattr(fabric_gather, "THREADS", 1024) + monkeypatch.setattr(fabric_gather, "available", lambda: False) + assert fabric_gather.resident_sms() == 0 + # An explicit count wins, with any gather configuration. + monkeypatch.setenv("TRTLLM_CUTEDSL_SM_RESERVE", "24") + assert fabric_gather.resident_sms() == 24 + + +def test_sm_reserve_is_set_explicitly(): + from tensorrt_llm._torch import cute_dsl_utils + + assert cute_dsl_utils.sm_reserve() == 0 + cute_dsl_utils.set_sm_reserve(30) + try: + assert cute_dsl_utils.sm_reserve() == 30 + finally: + cute_dsl_utils.set_sm_reserve(0) + with pytest.raises(ValueError): + cute_dsl_utils.set_sm_reserve(-1) + + @pytest.mark.parametrize( "full, cluster_size, reserve, expected", [(148, 1, 30, 118), (74, 2, 30, 59), (2, 2, 30, 1), (0, 2, 30, 0), (148, 1, 0, 148)], diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py index 489dca8ef7d5..7d60305b8479 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py @@ -487,6 +487,99 @@ def test_kda_mixer_empty_prefill_preserves_initial_conv_state(): assert torch.equal(conv_pool, conv_seed) +@torch.no_grad() +@pytest.mark.parametrize( + ("state_dtype", "mode"), + [(torch.bfloat16, "small-bh"), (torch.float32, "auto")], + ids=["bf16-small-bh", "fp32-auto"], +) +@pytest.mark.parametrize( + "sequence_lengths", + [[640, 100, 1000, 37], [1024]], + ids=["varlen", "single"], +) +def test_kda_prefill_op_sm_reserve_keeps_results(sequence_lengths, state_dtype, mode) -> None: + """Persistent grids of every size compute bitwise the same output and states. + + ``sm_reserve`` only shrinks the Small-BH persistent grid, so each work + unit is computed the same way by another block. + """ + from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops # noqa: F401 + + heads = 8 + tokens = sum(sequence_lengths) + rows = len(sequence_lengths) + cu_seqlens = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.int64, + device="cuda", + ) + gen = torch.Generator(device="cuda").manual_seed(0) + shape = (1, tokens, heads, HEAD_DIM) + + def bf16(scale: float = 1.0) -> torch.Tensor: + return (torch.randn(*shape, generator=gen, device="cuda") * scale).to(torch.bfloat16) + + inputs = dict( + q=bf16(), + k=bf16(), + v=bf16(0.5), + g=bf16(), + beta=torch.randn(1, tokens, heads, generator=gen, device="cuda").to(torch.bfloat16), + ) + state_seed = ( + torch.randn(rows + 1, heads, HEAD_DIM, HEAD_DIM, generator=gen, device="cuda") * 0.1 + ).to(state_dtype) + state_indices = torch.arange(1, rows + 1, dtype=torch.int32, device="cuda") + + def call(sm_reserve: int) -> tuple[torch.Tensor, torch.Tensor]: + state_pool = state_seed.clone() + output = torch.ops.trtllm.kda_prefill( + **inputs, + state_pool=state_pool, + state_indices=state_indices, + cu_seqlens=cu_seqlens, + A_log=torch.zeros(heads, device="cuda"), + dt_bias=torch.zeros(heads * HEAD_DIM, device="cuda"), + scale=HEAD_DIM**-0.5, + lower_bound=-5.0, + mode=mode, + sm_reserve=sm_reserve, + ) + return output, state_pool + + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + expected = call(0) + for sm_reserve in (1, num_sms // 2, num_sms - 3): + for got, want in zip(call(sm_reserve), expected, strict=True): + assert torch.equal(got, want), f"sm_reserve={sm_reserve}" + with pytest.raises(ValueError, match="sm_reserve"): + call(num_sms) + + +@torch.no_grad() +def test_kda_mixer_prefill_passes_the_reserve_set_before_construction() -> None: + """The mixer passes the SM reserve that was set before it was built to the prefill.""" + from tensorrt_llm._torch.cute_dsl_utils import set_sm_reserve + + set_sm_reserve(16) + try: + kda = _make_kda(num_heads=2) + finally: + set_sm_reserve(0) + reserves = [] + prefill_chunk_kda = kda._dispatch.prefill_chunk_kda + + def spy(**kwargs): + reserves.append(kwargs["sm_reserve"]) + return prefill_chunk_kda(**kwargs) + + kda._dispatch.prefill_chunk_kda = spy + hidden_states = torch.randn(1, 64, HIDDEN_SIZE, dtype=torch.bfloat16, device="cuda") * 0.05 + _run_production_prefill(kda, hidden_states) + assert reserves == [16] + + @torch.no_grad() def test_kda_prefill_op_partial_final_chunk_large_batch(): """Large packed batches with short final sequences match the FLA reference. From b3e52f166f64e50ea3ad8e8d5b860f5a222c754b Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:07:47 -0700 Subject: [PATCH 2/8] [None][perf] Kimi K3 KDA MTP verify: lower-latency v2 schedule of the verify kernel Add kda_decode_mtp_v2_kernel, a reschedule of the CuTe DSL KDA MTP-verify kernel with identical math: slot-independent loads are issued under the slot load and slot-dependent loads as soon as the slot arrives, the q/k precompute runs from shared memory with vectorized replay-cache writes moved under GEMM1/GEMM3, and GEMM3 overlaps the S0 copy into TMEM. Selection: _select_verify_kernel in cute_dsl_kimi_k3_kda_mtp_ops declares the condition in one place. Calls on the runtime verify contract (packed token rows, fused gated RMSNorm with a bf16 output, 3xTF32 recurrence, at most MAX_NUM_SPEC drafts, vector-aligned operands) run the v2 kernel; every other variant (MXFP8 output, BF16 math, legacy layouts, replay-only commits, misaligned operands) stays on kda_decode_mtp_kernel. No switch: the v2 kernel is bitwise identical where it applies. The v2 kernel takes kda_mtp_decode.py's constants and helpers (NUM_THREADS, TILE_K, T_PAD, TMEM_COLS, SCRATCH_FLOATS, TP_STRIDE, swz64_k16, cvt_bf16x2, round_tf32, tcgen05 fences). Both host functions build their MMA atoms, operand layouts and state TMA atoms with shared helpers (_verify_mma_atoms, _state_tma_atoms, _launch_verify). The v1 kernel's code is unchanged. Outputs, committed state, conv pools, replay caches and beta cache are bitwise identical to kda_decode_mtp_kernel. For the cluster-pair form (SPLIT_V == 2) the kernel orders CTA 0's replay-cache and conv-window writes after the peer CTA's load phase (CTA 0 waits on the pair's cluster barrier before those writes, CTA 1 before its state store), the ordering kda_decode_mtp_kernel already has. Evidence: GB300, per call (H=6, cold state, in a CUDA graph after the QKVG GEMM): 8-14% faster for B in {1..128}, T in {7, 8} (54108b3447). Rubin: the SM107 target configurations run the MXFP8-output verify, which this kernel does not serve; no direct Rubin measurement. GB300 single-wave grids vs the two-CTAs-per-SM kernel of the next commit (job 1368796, cold L2, H in {6, 12, 96}): with the bf16 state of the deployed configurations v2 is 7.6% / 6.4% faster at 3 / 5 drafts and 3.2% slower at 7 drafts; with the fp32 state (the Kimi K3 default) 16-20% faster at every draft count. So v2 stays for BF16-output grids that fit one wave. Platform: S (SM100/103 BF16-output verify path; bitwise vs the existing kernel). Tests: test_kda_mtp_decode_cute_parity.py::test_v2_kernel_matches_v1_bitwise (cluster-pair and single-CTA grids, fp32/bf16 state, strided state; out, state, conv pools and replay caches bitwise) and ::test_verify_kernel_selection. GB300 (OCI AGA) job 1387339. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../kimi_k3_kda/kda_mtp_decode_v2.py | 1320 +++++++++++++++++ .../cute_dsl_kimi_k3_kda_mtp_ops.py | 594 +++++--- .../test_kda_mtp_decode_cute_parity.py | 124 ++ 3 files changed, 1858 insertions(+), 180 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_v2.py diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_v2.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_v2.py new file mode 100644 index 000000000000..dcb520c8c67b --- /dev/null +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_v2.py @@ -0,0 +1,1320 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. + +"""Latency-scheduled CuTe DSL kernel for KDA multi-token speculative verify (packed rows). + +Computes exactly what ``kda_mtp_decode.kda_decode_mtp_kernel`` computes on its runtime +contract -- packed token rows, fused gated RMSNorm with a bf16 output, 3xTF32 tcgen05 chunk +recurrence, fp32 or bf16 state, K = V = 128, conv width 4, NUM_SPEC <= 7 -- with the same +arithmetic in the same order, so outputs, committed state, conv pools and replay caches are +bitwise identical. See that module for the math (replay of accepted drafts, chunk WY form, +GEMM1/GEMM3/forward substitution/GEMM2, in-place state commit after the golden token, the +cluster-pair ``SPLIT_V == 2`` form). The difference is the schedule of the latency-bound +CTA: + +* **Load phase.** The state slot is the first load. The slot-independent operands (new-token + projections, conv weights, per-head constants, output-norm operands) are issued under it; the + slot-dependent ones (the S0 TMA, the committed conv windows and every pending replay entry, + accepted or not) as soon as it arrives. Each lane moves four consecutive channels of a row + (8 B / 16 B vector loads) and keeps the raw words in registers until the accepted count fixes + the row's smem position, so the CTA waits for one round trip per dependency level (slot, then + operands) instead of one per staging step. +* **Precompute** works in smem only, with the q and k conv / SiLU / L2-norm chains of a token + interleaved in one warp. +* **Tensor-core ordering.** The GEMM3 operand rows lie clear of the S0 boxes, so GEMM3 is issued + before S0 moves into TMEM and runs under that copy; GEMM1 follows once its rows are written over + the read boxes. While the tensor cores run, warps 1..15 stage Kbar, seed the GEMM2 + accumulator, build the B2 tiles and write the new tokens' replay-cache rows and conv-window + commits (vectorized, from smem). +""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import Int64 +from cutlass.cutlass_dsl import T, dsl_user_op + +from .kda_mtp_decode import ( + NUM_THREADS, + SCRATCH_FLOATS, + T_PAD, + TILE_K, + TMEM_COLS, + TP_STRIDE, + cvt_bf16x2, + round_tf32, + swz64_k16, + tcgen05_fence_after_thread_sync, + tcgen05_fence_before_thread_sync, +) + +# Offsets (in floats) of the operand tiles in kda_mtp_decode's SCRATCH_FLOATS region; all +# multiples of 256 (1024 B) for the 128B-swizzle atoms. +OFF_B4HI = 10240 # 16 x 16 Bcoef hi, immediately followed by Kbar hi -> N=80 GEMM2 half-0 tile +OFF_B4LO = 12544 # 16 x 16 Bcoef lo, immediately followed by Kbar lo +OFF_B1K = 6144 # 16 x 128 K~hi (duplicate, N=16 tile for the S0_lo sweep; A2 reuses it) +OFF_B1 = 0 # 48 x 128 [K~hi; Q~; K~lo] +OFF_A2HI = 6144 # 128 x 16 V'^T hi +OFF_A2LO = 8192 # 128 x 16 V'^T lo +OFF_B2HI = 10496 # 128 x 16 Kbar hi +OFF_B2LO = 12800 # 128 x 16 Kbar lo +OFF_A3 = 16384 # 64 x 128 [Kf; Qf hi; Qf lo; -] (clear of the S0 boxes) +OFF_B3HI = 24576 # 16 x 128 Kb hi +OFF_B3LO = 26624 # 16 x 128 Kb lo + + +def _asm_words(n, asm, constraints, operands, *, loc, ip): + res = llvm.inline_asm( + llvm.StructType.get_literal([T.i32()] * n), + operands, + asm, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + return [cutlass.Int32(llvm.extractvalue(T.i32(), res, [i], loc=loc, ip=ip)) for i in range(n)] + + +@dsl_user_op +def ld_global_b32x2(addr: Int64, *, loc=None, ip=None): + """Two 32-bit words (8 B, 8-byte aligned) from global memory, unconverted: the load has no + consumer until its value is placed, so a warp keeps all of its loads in flight.""" + return _asm_words( + 2, + "ld.global.v2.b32 {$0, $1}, [$2];", + "=r,=r,l", + [Int64(addr).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_global_b32x4(addr: Int64, *, loc=None, ip=None): + """Four 32-bit words (16 B, 16-byte aligned) from global memory, unconverted.""" + return _asm_words( + 4, + "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,l", + [Int64(addr).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def ld_shared_f32x4(addr: cutlass.Int32, *, loc=None, ip=None): + """Four consecutive fp32 (16-byte aligned) from shared memory.""" + words = _asm_words( + 4, + "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", + "=r,=r,=r,=r,r", + [cutlass.Int32(addr).ir_value(loc=loc, ip=ip)], + loc=loc, + ip=ip, + ) + return [w.bitcast(cutlass.Float32) for w in words] + + +@dsl_user_op +def _st_asm(asm, constraints, operands, *, loc=None, ip=None): + llvm.inline_asm( + None, + operands, + asm, + constraints, + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + + +def st_global_b32x2(addr, a, b): + _st_asm( + "st.global.v2.b32 [$0], {$1, $2};", + "l,r,r", + [Int64(addr).ir_value(), cutlass.Int32(a).ir_value(), cutlass.Int32(b).ir_value()], + ) + + +def st_global_b32x4(addr, a, b, c, d): + _st_asm( + "st.global.v4.b32 [$0], {$1, $2, $3, $4};", + "l,r,r,r,r", + [Int64(addr).ir_value()] + [cutlass.Int32(x).ir_value() for x in (a, b, c, d)], + ) + + +def st_shared_f32x4(addr, a, b, c, d): + _st_asm( + "st.shared.v4.f32 [$0], {$1, $2, $3, $4};", + "r,f,f,f,f", + [cutlass.Int32(addr).ir_value()] + [cutlass.Float32(x).ir_value() for x in (a, b, c, d)], + ) + + +def gaddr(tensor, coord): + """Global address of `tensor[coord]`.""" + return (tensor.iterator + cute.crd2idx(coord, tensor.layout)).toint() + + +def saddr(tensor, offset): + """Shared-memory address of element `offset` of an smem tensor (row-major offset).""" + return (tensor.iterator + offset).toint() + + +def ld_row4(tensor, coord, words, base): + """Issue the load of four consecutive elements of `tensor` (16- or 32-bit) at `coord` into the + raw-word buffer `words[base : base + 4]` (two packed words for a 16-bit type).""" + if cutlass.const_expr(tensor.element_type.width == 16): + w = ld_global_b32x2(gaddr(tensor, coord)) + words[base] = w[0] + words[base + 1] = w[1] + else: + w = ld_global_b32x4(gaddr(tensor, coord)) + for i in range(4): + words[base + i] = w[i] + + +def unpack4(dtype, words, base): + """Four fp32 values from the raw words of `ld_row4` (bf16 -> fp32 is exact).""" + if cutlass.const_expr(dtype.width == 16): + assert dtype == cutlass.BFloat16, "16-bit operands are bf16" + w0 = words[base] + w1 = words[base + 1] + return [ + (w0 << 16).bitcast(cutlass.Float32), + (w0 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + (w1 << 16).bitcast(cutlass.Float32), + (w1 & cutlass.Int32(-65536)).bitcast(cutlass.Float32), + ] + return [words[base + i].bitcast(cutlass.Float32) for i in range(4)] + + +def st_row4(tensor, coord, v): + """Store four fp32 values `v` as consecutive elements of `tensor` (rounded to nearest even + for bf16, like a scalar `.to(BFloat16)`).""" + if cutlass.const_expr(tensor.element_type.width == 16): + assert tensor.element_type == cutlass.BFloat16, "16-bit caches are bf16" + st_global_b32x2( + gaddr(tensor, coord), + cvt_bf16x2(v[1], v[0]).bitcast(cutlass.Int32), + cvt_bf16x2(v[3], v[2]).bitcast(cutlass.Int32), + ) + else: + st_global_b32x4(gaddr(tensor, coord), *[x.bitcast(cutlass.Int32) for x in v]) + + +@cute.kernel +def kda_decode_mtp_v2_kernel( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + o: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + ssm_state_indices: cute.Tensor, + num_accepted_tokens: cute.Tensor, + smem_qk_layout: cute.Layout, + tiled_mma1: cute.TiledMma, + tiled_mma3: cute.TiledMma, + tiled_mma1k: cute.TiledMma, + smem_layout_b1: cute.ComposedLayout, + smem_layout_a2: cute.ComposedLayout, + smem_layout_a3: cute.ComposedLayout, + smem_layout_b3: cute.ComposedLayout, + smem_layout_b1k: cute.ComposedLayout, + tiled_mma4: cute.TiledMma, + tma_atom_t: cute.CopyAtom, + mT_tma: cute.Tensor, + box_layout: cute.ComposedLayout, + tma_atom_l: cute.CopyAtom, + mS_tma: cute.Tensor, + tiled_mma2h: cute.TiledMma, + smem_layout_b2h: cute.ComposedLayout, + box_layout_t: cute.ComposedLayout, + tiled_mma2q: cute.TiledMma, + smem_layout_b2q: cute.ComposedLayout, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + SPLIT_V: cutlass.Constexpr[int], +): + """KDA MTP verify, packed-token / fused-norm contract (see the module docstring). + + Grid ``(HV, N, SPLIT_V)``: one (head, request) tile per CTA, or per cluster pair of CTAs + with ``SPLIT_V == 2`` (each CTA seeds and stores half of the state rows).""" + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + i_hv, i_n, vh = cute.arch.block_idx() + i_h = i_hv + # The request's state slot is the first load of the kernel; its round trip overlaps the + # slot-independent loads issued below. + slot32 = ssm_state_indices[i_n] + hk_off = i_h * K + hv_off = i_hv * V + # Recurrent-state element type (fp32 or bf16). The TMA boxes have 128 B swizzled rows, + # i.e. BOX_COLS = 32 fp32 or 64 bf16 columns; the state tile is NUM_SLABS such boxes. + STATE_DTYPE = h0.element_type + STATE_BF16 = STATE_DTYPE == cutlass.BFloat16 + # 3xTF32 recurrence (hi/lo splits on every product that feeds the state); S0_LO: a fp32 + # state needs the lo part of S0 (a bf16 state is exact in tf32). + S0_LO = not STATE_BF16 + N1 = 3 * T_PAD # GEMM1 N: [K~hi; Q~; K~lo] + CONV_DTYPE = cs_q.element_type + K_CACHE_DTYPE = k_cache.element_type + V_CACHE_DTYPE = v_cache.element_type + BETA_CACHE_DTYPE = beta_cache.element_type + BOX_COLS = 128 // (STATE_DTYPE.width // 8) + NUM_SLABS = K // BOX_COLS + t_max = 2 * NUM_SPEC + 1 + NEW_TOK = 1 + NUM_SPEC + vec_size = TILE_K // 32 + # Load-phase warp roles: warp j < NEW_TOK stages new token j, warp REPLAY_WARP0 + r the + # pending draft r (all NUM_SPEC entries are loaded, the accepted ones placed), the last warp + # the committed conv windows and the S0 TMA. + REPLAY_WARP0 = NUM_THREADS // 64 + TMA_WARP = NUM_THREADS // 32 - 1 + assert NUM_SPEC <= 7 and t_max <= T_PAD, "chunk form pads the token axis to 16" + assert NEW_TOK <= REPLAY_WARP0 and REPLAY_WARP0 + NUM_SPEC <= TMA_WARP + assert K == 128 and V == 128 and KERNEL_WIDTH == 4 and NUM_THREADS == 512, ( + "the thread mappings assume K = V = 128, conv width 4 and 512 threads" + ) + k_weight_base = KERNEL_WIDTH * K + v_weight_base = 2 * KERNEL_WIDTH * K + conv_weight_elems = 2 * KERNEL_WIDTH * K + KERNEL_WIDTH * V + smem = cutlass.utils.SmemAllocator() + sQ = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sK = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sG = smem.allocate_tensor(cutlass.Float32, smem_qk_layout, 16) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) + sConvW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((conv_weight_elems,)), 16) + sOnG = smem.allocate_tensor(cutlass.Float32, cute.make_layout((NEW_TOK * V,)), 16) + sOnW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V,)), 16) + # Staged per-token inputs (raw conv sequence with the committed window + # prepended in rows 0..W-2, raw gates, raw betas). + sXQ = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * K,)), 16) + sXK = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * K,)), 16) + sXV = smem.allocate_tensor(cutlass.Float32, cute.make_layout(((t_max + 3) * V,)), 16) + sGin = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * K,)), 16) + sBetaRaw = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + sDtBias = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + sAlog = smem.allocate_tensor(cutlass.Float32, cute.make_layout((32,)), 16) + tmem_holder = smem.allocate_array(cutlass.Int32, 1) + mbar1 = smem.allocate_array(cutlass.Int64, 1) + mbar2 = smem.allocate_array(cutlass.Int64, 1) + mbar_ld = smem.allocate_array(cutlass.Int64, 1) + mbar2b = smem.allocate_array(cutlass.Int64, 1) + sScratch = smem.allocate_tensor(cutlass.Float32, cute.make_layout((SCRATCH_FLOATS,)), 1024) + sB1 = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B1, smem_layout_b1.inner, dtype=cutlass.Float32), + smem_layout_b1.outer, + ) + sA2hi = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_A2HI, smem_layout_a2.inner, dtype=cutlass.Float32), + smem_layout_a2.outer, + ) + sA2lo = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_A2LO, smem_layout_a2.inner, dtype=cutlass.Float32), + smem_layout_a2.outer, + ) + sA3 = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_A3, smem_layout_a3.inner, dtype=cutlass.Float32), + smem_layout_a3.outer, + ) + sB3hi = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B3HI, smem_layout_b3.inner, dtype=cutlass.Float32), + smem_layout_b3.outer, + ) + sB3lo = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B3LO, smem_layout_b3.inner, dtype=cutlass.Float32), + smem_layout_b3.outer, + ) + sB1k = cute.make_tensor( + cute.recast_ptr(sScratch.iterator + OFF_B1K, smem_layout_b1k.inner, dtype=cutlass.Float32), + smem_layout_b1k.outer, + ) + sB1kf = cute.make_tensor(sScratch.iterator + OFF_B1K, cute.make_layout((T_PAD * 128,))) + # Flat (unswizzled-pointer) views for element writes with closed-form swizzle offsets. + sB1f = cute.make_tensor(sScratch.iterator + OFF_B1, cute.make_layout((48 * 128,))) + sA2hif = cute.make_tensor(sScratch.iterator + OFF_A2HI, cute.make_layout((128 * 16,))) + sA2lof = cute.make_tensor(sScratch.iterator + OFF_A2LO, cute.make_layout((128 * 16,))) + sB2hif = cute.make_tensor(sScratch.iterator + OFF_B2HI, cute.make_layout((128 * 16,))) + sB2lof = cute.make_tensor(sScratch.iterator + OFF_B2LO, cute.make_layout((128 * 16,))) + sB4hif = cute.make_tensor(sScratch.iterator + OFF_B4HI, cute.make_layout((T_PAD * T_PAD,))) + sB4lof = cute.make_tensor(sScratch.iterator + OFF_B4LO, cute.make_layout((T_PAD * T_PAD,))) + sA3f = cute.make_tensor(sScratch.iterator + OFF_A3, cute.make_layout((64 * 128,))) + sB3hif = cute.make_tensor(sScratch.iterator + OFF_B3HI, cute.make_layout((16 * 128,))) + sB3lof = cute.make_tensor(sScratch.iterator + OFF_B3LO, cute.make_layout((16 * 128,))) + # S0 / committed-state TMA boxes at the start of the scratch region (fp32 and bf16 views) + sBoxf = cute.make_tensor(sScratch.iterator, cute.make_layout((V * K,))) + sBoxh = cute.make_tensor( + cute.recast_ptr(sScratch.iterator, dtype=cutlass.BFloat16), cute.make_layout((V * K,)) + ) + # GEMM3 result [L; B hi; B lo; -] as a plain (64 x 16) row-major matrix. + sLB = smem.allocate_tensor( + cutlass.Float32, cute.make_layout((4 * T_PAD, T_PAD), stride=(T_PAD, 1)), 16 + ) + sEc = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K,)), 16) + # V'^T (v, s) and Kbar^T (k, s) staged with a padded row stride (conflict-free for both + # the row-per-thread producers and the 16-lanes-per-row tile writers). + sVpT = smem.allocate_tensor(cutlass.Float32, cute.make_layout((V * TP_STRIDE,)), 16) + sKbT = smem.allocate_tensor(cutlass.Float32, cute.make_layout((K * TP_STRIDE,)), 16) + # Cumulative gates G_t (per k) reuse the raw-gate staging buffer (no longer needed after the + # precompute). + sE = cute.make_tensor(sGin.iterator, cute.make_layout((t_max * K,))) + + # ---- load phase. Every load is issued before any of its consumers, so the CTA waits for its + # round trips in parallel: the slot-independent operands (new-token projections, conv + # weights, per-head constants, output-norm operands) under the slot load, the + # slot-dependent ones (S0 TMA, committed conv windows, every replay entry) as soon as the + # slot arrives; only the smem placement of the new tokens and accepted drafts waits for + # the accepted count. Each lane moves four consecutive channels (8 B for bf16, 16 B for + # fp32) per row, held as raw words until the row's smem position is known. + ch = tidx % 128 + quad = tidx // 128 + k4 = 4 * in_warp_tid + # (1) slot-independent: conv weights ([ch][w] rows of 16 B: threads [0, 384) one row each), + # dt_bias / output-norm weight / A_log (threads [384, 512)), new token j = warp + # (warps < NEW_TOK) + r_w = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Int32) + if quad == 0: + ld_row4(w_q, (hk_off + ch, 0), r_w, 0) + if quad == 1: + ld_row4(w_k, (hk_off + ch, 0), r_w, 0) + if quad == 2: + ld_row4(w_v, (hv_off + ch, 0), r_w, 0) + r_c = cute.make_rmem_tensor(cute.make_layout((3,)), cutlass.Float32) + if quad == 3: + r_c[0] = cutlass.Float32(dt_bias[hk_off + ch]) + r_c[1] = cutlass.Float32(onorm_weight[ch]) + r_c[2] = cutlass.Float32(A_log[i_h]) + r_new = cute.make_rmem_tensor(cute.make_layout((5 * 4,)), cutlass.Int32) + r_newb = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Float32) + tok_row = i_n * NEW_TOK + warp_idx + if warp_idx < NEW_TOK: + if cutlass.const_expr(USE_FLAT_LAYOUT): + ld_row4(x_q, (0, tok_row, hk_off + k4), r_new, 0) + ld_row4(x_k, (0, tok_row, hk_off + k4), r_new, 4) + ld_row4(x_v, (0, tok_row, hv_off + k4), r_new, 8) + else: + ld_row4(x_q, (0, tok_row, i_h, k4), r_new, 0) + ld_row4(x_k, (0, tok_row, i_h, k4), r_new, 4) + ld_row4(x_v, (0, tok_row, i_hv, k4), r_new, 8) + ld_row4(g, (0, tok_row, i_hv, k4), r_new, 12) + ld_row4(onorm_g, (0, tok_row, i_hv, k4), r_new, 16) + r_newb[0] = cutlass.Float32(beta[0, tok_row, i_hv]) + if warp_idx == 0: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mbar1, 1) + cute.arch.mbarrier_init(mbar2, 1) + cute.arch.mbarrier_init(mbar2b, 1) + cute.arch.mbarrier_init(mbar_ld, 1) + cute.arch.mbarrier_init_fence() + cute.arch.barrier() + # V2 state views may place many layer pages between adjacent slots. Widen + # the slot before indexing: slot * state stride can exceed INT32_MAX. + slot = Int64(slot32) + h0_idx = slot * HV + i_hv + # (2) slot-dependent: replay entry r = warp - REPLAY_WARP0 (every pending entry, accepted or + # not), the committed conv windows (TMA warp: three columns of q, k and v) + rep = warp_idx - REPLAY_WARP0 + r_rep = cute.make_rmem_tensor(cute.make_layout((6 * 4,)), cutlass.Int32) + r_repb = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Float32) + if (warp_idx >= REPLAY_WARP0) & (rep < NUM_SPEC): + ld_row4(k_cache, (slot, rep, hk_off + k4), r_rep, 0) + ld_row4(g_cache, (slot, rep, hk_off + k4), r_rep, 4) + ld_row4(v_cache, (slot, rep, hv_off + k4), r_rep, 8) + ld_row4(cs_q, (slot, hk_off + k4, KERNEL_WIDTH - 1 + rep), r_rep, 12) + ld_row4(cs_k, (slot, hk_off + k4, KERNEL_WIDTH - 1 + rep), r_rep, 16) + ld_row4(cs_v, (slot, hv_off + k4, KERNEL_WIDTH - 1 + rep), r_rep, 20) + r_repb[0] = cutlass.Float32(beta_cache[slot, rep, i_hv]) + r_win = cute.make_rmem_tensor(cute.make_layout((3 * (KERNEL_WIDTH - 1) * 4,)), cutlass.Int32) + if warp_idx == TMA_WARP: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + ld_row4(cs_q, (slot, hk_off + k4, w), r_win, 4 * w) + ld_row4(cs_k, (slot, hk_off + k4, w), r_win, 4 * (KERNEL_WIDTH - 1 + w)) + ld_row4(cs_v, (slot, hv_off + k4, w), r_win, 4 * (2 * (KERNEL_WIDTH - 1) + w)) + commit_len = num_accepted_tokens[slot] + # ---- TMA bulk load of S0 into the boxes; completes under the token staging / + # precompute and is consumed (mbar_ld) at the start of the chunk stage. + # Issued by the whole of one warp (a TMA load issued from a single thread never + # completes on this DSL build; the copy elects its issuing lane internally). + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_ld, V * K * (STATE_DTYPE.width // 8)) + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sTl_h = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + sl_h * (V * 32), box_layout.inner, dtype=STATE_DTYPE + ), + box_layout.outer, + ) + gS_h_all = cute.local_tile(mS_tma, (V, BOX_COLS, 1), (0, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gS_h = gS_h_all[(None, None, 0, h0_idx)] + else: + gS_h = gS_h_all[(None, None, 0, (i_hv, slot))] + tSl_h, tGl_h = cpasync.tma_partition( + tma_atom_l, + 0, + cute.make_layout(1), + cute.group_modes(sTl_h, 0, 2), + cute.group_modes(gS_h, 0, 2), + ) + if warp_idx == TMA_WARP: + cute.copy(tma_atom_l, tGl_h, tSl_h, tma_bar_ptr=mbar_ld) + # TMEM allocation overlaps the in-flight loads; the pointer is retrieved after the + # precompute barrier (first TMEM use is the S0 staging in the chunk stage). + if warp_idx == 0: + cute.arch.alloc_tmem(TMEM_COLS, tmem_holder) + # (3) placement: the slot-independent operands first, the token rows once the accepted count + # is in. Smem rows keep the [row][channel] fp32 layout of the precompute. + if quad < 3: + wv = unpack4(w_q.element_type, r_w, 0) + for w in cutlass.range_constexpr(KERNEL_WIDTH): + sConvW[quad * (KERNEL_WIDTH * K) + w * K + ch] = wv[w] + if quad == 3: + sDtBias[ch] = r_c[0] + sOnW[ch] = r_c[1] + if ch == 0: + sAlog[0] = r_c[2] + if warp_idx < NEW_TOK: + st_shared_f32x4(saddr(sOnG, warp_idx * V + k4), *unpack4(onorm_g.element_type, r_new, 16)) + if warp_idx == TMA_WARP: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + st_shared_f32x4(saddr(sXQ, w * K + k4), *unpack4(CONV_DTYPE, r_win, 4 * w)) + st_shared_f32x4( + saddr(sXK, w * K + k4), *unpack4(CONV_DTYPE, r_win, 4 * (KERNEL_WIDTH - 1 + w)) + ) + st_shared_f32x4( + saddr(sXV, w * V + k4), + *unpack4(CONV_DTYPE, r_win, 4 * (2 * (KERNEL_WIDTH - 1) + w)), + ) + # Only NUM_SPEC drafts can be pending from the previous round. Clamp + # so a malformed count cannot drive T_loop past the t_max-sized SMEM + # buffers or the num_spec extents of the replay caches. + if commit_len > NUM_SPEC: + commit_len = cutlass.Int32(NUM_SPEC) + T_loop = commit_len + NEW_TOK + if warp_idx < NEW_TOK: + t_new = commit_len + warp_idx + st_shared_f32x4(saddr(sXQ, (t_new + 3) * K + k4), *unpack4(x_q.element_type, r_new, 0)) + st_shared_f32x4(saddr(sXK, (t_new + 3) * K + k4), *unpack4(x_k.element_type, r_new, 4)) + st_shared_f32x4(saddr(sXV, (t_new + 3) * V + k4), *unpack4(x_v.element_type, r_new, 8)) + st_shared_f32x4(saddr(sGin, t_new * K + k4), *unpack4(g.element_type, r_new, 12)) + if in_warp_tid == 0: + sBetaRaw[t_new] = r_newb[0] + if (warp_idx >= REPLAY_WARP0) & (rep < commit_len): + # replayed drafts produce no output: their q rows are unused + z = cutlass.Float32(0.0) + st_shared_f32x4(saddr(sQ, rep * K + k4), z, z, z, z) + st_shared_f32x4(saddr(sK, rep * K + k4), *unpack4(K_CACHE_DTYPE, r_rep, 0)) + st_shared_f32x4(saddr(sG, rep * K + k4), *unpack4(g_cache.element_type, r_rep, 4)) + st_shared_f32x4(saddr(sVall, rep * V + k4), *unpack4(V_CACHE_DTYPE, r_rep, 8)) + st_shared_f32x4(saddr(sXQ, (rep + 3) * K + k4), *unpack4(CONV_DTYPE, r_rep, 12)) + st_shared_f32x4(saddr(sXK, (rep + 3) * K + k4), *unpack4(CONV_DTYPE, r_rep, 16)) + st_shared_f32x4(saddr(sXV, (rep + 3) * V + k4), *unpack4(CONV_DTYPE, r_rep, 20)) + if in_warp_tid == 0: + sBeta[rep] = r_repb[0] + # Padding rows [T_loop, t_max) of the gates, values and betas are zeroed so the chunk + # stage's r_Y / cumulative-gate / V' loops are branch-free. + if tidx < K: + for _z in cutlass.range(T_loop, t_max, 1): + sG[_z, tidx] = cutlass.Float32(0.0) + sVall[_z * V + tidx] = cutlass.Float32(0.0) + if tidx == 0: + sBeta[_z] = cutlass.Float32(0.0) + wg = warp_idx // 4 + wg_tidx = tidx % 128 + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,)), cutlass.Float32) + r_state = cute.make_rmem_tensor(cute.make_layout((32,)), cutlass.Float32) + cute.arch.barrier() + + # ---- per-token-parallel precompute (smem only; the new tokens' cache and conv-window writes + # are issued from smem while the tensor cores run GEMM1/GEMM3): warp j < NEW_TOK computes + # the q and k conv + SiLU + L2 norm of new token j (the two chains interleaved), warp 8 + j + # its gate, beta and v conv. + half = NUM_THREADS // 64 + role = warp_idx // half + j = warp_idx % half + if j < NEW_TOK: + t = commit_len + j + if role == 0: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_conv_0 = cutlass.Float32(0.0) + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + r_conv_0 += sXQ[(t + w) * K + k_idx] * sConvW[w * K + k_idx] + r_conv = sXK[t * K + k_idx] * sConvW[k_weight_base + 0 * K + k_idx] + for w in cutlass.range(1, KERNEL_WIDTH, unroll_full=True): + r_conv += sXK[(t + w) * K + k_idx] * sConvW[k_weight_base + w * K + k_idx] + e0 = cute.math.exp(-r_conv_0, fastmath=True) + r_q[i] = r_conv_0 * cute.arch.rcp_approx(cutlass.Float32(1.0) + e0) + r_conv = r_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_conv, fastmath=True) + ) + r_k[i] = r_conv + sum_q = 0.0 + sum_k = 0.0 + for i in cutlass.range(vec_size, unroll_full=True): + sum_q += r_q[i] * r_q[i] + sum_k += r_k[i] * r_k[i] + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly( + sum_q, offset=offset, mask=-1, mask_and_clamp=31 + ) + sum_k += cute.arch.shuffle_sync_bfly( + sum_k, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_q_scaled = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + sQ[t, k_idx] = r_q[i] * rnorm_q_scaled + sK[t, k_idx] = r_k[i] * rnorm_k + else: + r_exp_A = cute.math.exp(sAlog[0], fastmath=True) + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_g_raw = sGin[t * K + k_idx] + r_g_raw = r_g_raw + sDtBias[k_idx] + exp_A_x = r_exp_A * r_g_raw + sigmoid_val = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-exp_A_x, fastmath=True) + ) + sG[t, k_idx] = lower_bound * sigmoid_val + r_beta_t = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-sBetaRaw[t], fastmath=True) + ) + if in_warp_tid == 0: + sBeta[t] = r_beta_t + for i in cutlass.range(vec_size, unroll_full=True): + v_idx = i * 32 + in_warp_tid + _v_conv = 0.0 + for w in cutlass.range(KERNEL_WIDTH, unroll_full=True): + _v_conv += sXV[(t + w) * V + v_idx] * sConvW[v_weight_base + w * V + v_idx] + _v_conv = _v_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_v_conv, fastmath=True) + ) + sVall[t * V + v_idx] = _v_conv + cute.arch.barrier() + + # ================================================================= + # Chunk (WY-form) recurrence on tensor cores. + # ================================================================= + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Float32, 16, tmem_holder) + pool = cutlass.utils.TmemBufferPool(tmem_ptr, TMEM_COLS) + tD4_frag = tiled_mma4.make_fragment_C(tiled_mma4.partition_shape_C((V, T_PAD))) + tD4 = pool.allocate_tensor(tD4_frag.layout, cutlass.Float32) # outputs (columns 0..15) + tD2h_frag = tiled_mma2h.make_fragment_C(tiled_mma2h.partition_shape_C((V, K // 2))) + tD2h0 = pool.allocate_tensor(tD2h_frag.layout, cutlass.Float32) # GEMM2 half 0: k cols 0..63 + tD2h1 = pool.allocate_tensor(tD2h_frag.layout, cutlass.Float32) # GEMM2 half 1: k cols 64..127 + tD2q_frag = tiled_mma2q.make_fragment_C(tiled_mma2q.partition_shape_C((V, T_PAD + K // 2))) + tD2q = cute.make_tensor(tD4.iterator, tD2q_frag.layout) # [D4 | dS half 0] accumulator view + tA_frag = tiled_mma1.make_fragment_A(tiled_mma1.partition_shape_A((V, K))) + tAhi = pool.allocate_tensor(tA_frag.layout, cutlass.Float32) + tAlo = pool.allocate_tensor(tA_frag.layout, cutlass.Float32) + tD1_frag = tiled_mma1.make_fragment_C(tiled_mma1.partition_shape_C((V, N1))) + tD1 = pool.allocate_tensor(tD1_frag.layout, cutlass.Float32) + tD3_frag = tiled_mma3.make_fragment_C(tiled_mma3.partition_shape_C((4 * T_PAD, T_PAD))) + tD3 = pool.allocate_tensor(tD3_frag.layout, cutlass.Float32) + tD2 = cute.make_tensor(tD2h0.iterator, tA_frag.layout) # A-view layout (all 128 columns) + + st32 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + ld16 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + ld32 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x32), cutlass.Float32) + thr_stA = tcgen05.make_tmem_copy(st32, tD2).get_slice(wg_tidx) + tD2_st = thr_stA.partition_D(tD2) + rA_shape = thr_stA.partition_S(tD2).shape + tAhi_st = thr_stA.partition_D(tAhi) + tAlo_st = thr_stA.partition_D(tAlo) + thr_ldD2 = tcgen05.make_tmem_copy(ld32, tD2).get_slice(wg_tidx) + tD2_ld = thr_ldD2.partition_S(tD2) + rS = cute.make_rmem_tensor(thr_ldD2.partition_D(tD2).shape[:-1], cutlass.Float32) + rS_flat = cute.make_tensor(rS.iterator, cute.make_layout((32,))) + # TMA store partitions for the NUM_SLABS boxes of this CTA's committed-state tile. + tS_list = [] + tG_list = [] + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sT_h = cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + sl_h * ((V // SPLIT_V) * 32), + box_layout_t.inner, + dtype=STATE_DTYPE, + ), + box_layout_t.outer, + ) + gT_h_all = cute.local_tile(mT_tma, (V // SPLIT_V, BOX_COLS, 1), (vh, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gT_h = gT_h_all[(None, None, 0, h0_idx)] + else: + gT_h = gT_h_all[(None, None, 0, (i_hv, slot))] + tS_h, tG_h = cpasync.tma_partition( + tma_atom_t, + 0, + cute.make_layout(1), + cute.group_modes(sT_h, 0, 2), + cute.group_modes(gT_h, 0, 2), + ) + tS_list.append(tS_h) + tG_list.append(tG_h) + r_hi = cute.make_rmem_tensor(rA_shape[:-1], cutlass.Float32) + r_lo = cute.make_rmem_tensor(rA_shape[:-1], cutlass.Float32) + r_hi_flat = cute.make_tensor(r_hi.iterator, cute.make_layout((32,))) + r_lo_flat = cute.make_tensor(r_lo.iterator, cute.make_layout((32,))) + + # ---- cumulative gates per k column (thread kq, quarter q owns tokens t = q mod 4): + # G_t -> sE, the T x T coupling factors, and the operand rows written with closed-form + # 128B-swizzle offsets (lanes = consecutive k: conflict-free). The GEMM3 rows (A3 = + # [Kf; Qf hi; Qf lo], B3 = Kb hi / lo) lie clear of the S0 boxes and are written first, + # so GEMM3 runs while S0 moves into TMEM; the GEMM1 rows (K~ hi / lo, Q~) stay in + # registers until the S0 boxes have been read. + kq = tidx % K + quarter = tidx // K + t_mid = T_loop // 2 + r_G = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + g_run = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t < t_max: + g_run = g_run + sG[t, kq] + r_G[t] = g_run + g_mid = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t == t_mid: + g_mid = r_G[t] + # Factors of the T x T coupling: e^{G_t - G_mid} (t side) and e^{G_mid - G_s} (s side), + # clamped at e^87 so a chunk whose decay exceeds the fp32 range yields ~0 instead of + # inf/NaN (the affected products are < e^{-|g|} and negligible either way). + exp_clamp = cutlass.Float32(87.0) + ka = kq // 32 + kc = kq % 32 + b1_base = ka * (3 * T_PAD * 32) + a3_base = ka * (4 * T_PAD * 32) + b3_base = ka * (T_PAD * 32) + b1k_base = ka * (T_PAD * 32) + # GEMM1 rows of this thread's tokens t = quarter + 4 i: (K~ hi, Q~, K~ lo) per token + r_b1 = cute.make_rmem_tensor(cute.make_layout((3 * (T_PAD // 4),)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + kt_hi = cutlass.Float32(0.0) + kt_lo = cutlass.Float32(0.0) + qt_r = cutlass.Float32(0.0) + kf = cutlass.Float32(0.0) + qf_hi = cutlass.Float32(0.0) + qf_lo = cutlass.Float32(0.0) + kb_hi = cutlass.Float32(0.0) + kb_lo = cutlass.Float32(0.0) + if t < T_loop: + d_t = r_G[t] - g_mid + e_t = cute.math.exp(-cute.arch.fmax(-d_t, -exp_clamp), fastmath=True) + e_bt = cute.math.exp(-cute.arch.fmax(d_t, -exp_clamp), fastmath=True) + eg_t = cute.math.exp(r_G[t], fastmath=True) + sE[t * K + kq] = r_G[t] + if t == commit_len: + sEc[kq] = eg_t + k_t = sK[t, kq] + q_t = sQ[t, kq] + kt = k_t * eg_t + kt_hi = round_tf32(kt) + qt_r = round_tf32(q_t * eg_t) + kf = round_tf32(k_t * e_t) + qf = q_t * e_t + qf_hi = round_tf32(qf) + kb = k_t * e_bt + kb_hi = round_tf32(kb) + kt_lo = kt - kt_hi + qf_lo = qf - qf_hi + kb_lo = kb - kb_hi + # rows t, 16+t, 32+t share t % 8 -> the same 16-byte-chunk XOR + xc = kc ^ (4 * (t % 8)) + sA3f[a3_base + 32 * t + xc] = kf + sA3f[a3_base + 32 * (T_PAD + t) + xc] = qf_hi + sA3f[a3_base + 32 * (2 * T_PAD + t) + xc] = qf_lo + sB3hif[b3_base + 32 * t + xc] = kb_hi + sB3lof[b3_base + 32 * t + xc] = kb_lo + r_b1[3 * (t // 4)] = kt_hi + r_b1[3 * (t // 4) + 1] = qt_r + r_b1[3 * (t // 4) + 2] = kt_lo + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + # ---- GEMM3: D3 = A3 Kb_hi^T + A3 Kb_lo^T (warp 0; completion is tracked with GEMM1's) + tCrA3 = tiled_mma3.make_fragment_A( + cute.make_tensor( + cute.recast_ptr(sA3.iterator, smem_layout_a3.inner, dtype=cutlass.TFloat32), + sA3.layout, + ) + ) + tCrB3hi = tiled_mma3.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB3hi.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3hi.layout, + ) + ) + tCrB3lo = tiled_mma3.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB3lo.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3lo.layout, + ) + ) + if warp_idx == 0: + tiled_mma3.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, tD3, tCrA3[(None, None, kb_i, 0)], tCrB3hi[(None, None, kb_i, 0)], tD3 + ) + tiled_mma3.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, tD3, tCrA3[(None, None, kb_i, 0)], tCrB3lo[(None, None, kb_i, 0)], tD3 + ) + + # ---- S0 row v = wg_tidx, columns [32 wg, +32), from the TMA-loaded swizzled box + # (16 B chunks: conflict-free LDS.128 under the 128B swizzle) into TMEM (hi / lo) + cute.arch.mbarrier_wait(mbar_ld, 0) + # In-place update of a cluster pair: the peer CTA must not store its state half before our + # S0 load completed, and CTA 0 must not overwrite the replay entries and conv windows before + # the peer's load phase read them (those values have been consumed by now) -> arrive here; + # CTA 0 waits before its replay-cache writes, CTA 1 before its state store. + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_arrive_relaxed() + vx = 4 * (wg_tidx % 8) + if cutlass.const_expr(STATE_BF16): + # bf16 box: 64 columns per 128 B row, so this warpgroup's 32 columns are the 16 B + # chunks [4 (wg % 2), +4) of box wg // 2. A bf16 value is exact in tf32: no lo part, + # and the S0_lo sweep of GEMM1 is skipped. + for c in cutlass.range(4, unroll_full=True): + chunk = (4 * (wg % 2) + c) ^ (wg_tidx % 8) + off_l = (wg // 2) * (V * 64) + 64 * wg_tidx + 8 * chunk + for i in cutlass.range(8, unroll_full=True): + r_state[8 * c + i] = cutlass.Float32(sBoxh[off_l + i]) + for i in cutlass.range(32, unroll_full=True): + r_hi_flat[i] = r_state[i] + for jj in cutlass.range_constexpr(4): + if wg == jj: + cute.copy(st32, r_hi, tAhi_st[(None, None, None, jj)]) + else: + for c in cutlass.range(8, unroll_full=True): + off_l = wg * (V * 32) + 32 * wg_tidx + ((4 * c) ^ vx) + for i in cutlass.range(4, unroll_full=True): + r_state[4 * c + i] = sBoxf[off_l + i] + for i in cutlass.range(32, unroll_full=True): + s_hi = round_tf32(r_state[i]) + r_hi_flat[i] = s_hi + r_lo_flat[i] = r_state[i] - s_hi + for jj in cutlass.range_constexpr(4): + if wg == jj: + cute.copy(st32, r_hi, tAhi_st[(None, None, None, jj)]) + cute.copy(st32, r_lo, tAlo_st[(None, None, None, jj)]) + cute.arch.fence_view_async_tmem_store() + # the S0 boxes have been read: the GEMM1 operand rows go over them + cute.arch.barrier() + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + xc = kc ^ (4 * (t % 8)) + sB1f[b1_base + 32 * t + xc] = r_b1[3 * (t // 4)] + if cutlass.const_expr(S0_LO): + sB1kf[b1k_base + 32 * t + xc] = r_b1[3 * (t // 4)] + sB1f[b1_base + 32 * (T_PAD + t) + xc] = r_b1[3 * (t // 4) + 1] + sB1f[b1_base + 32 * (2 * T_PAD + t) + xc] = r_b1[3 * (t // 4) + 2] + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + + # ---- operand fragments of GEMM1 (D1 = S0_hi B1^T + S0_lo K~hi^T) + tCrB1 = tiled_mma1.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB1.iterator, smem_layout_b1.inner, dtype=cutlass.TFloat32), + sB1.layout, + ) + ) + tCrB1k = tiled_mma1k.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB1k.iterator, smem_layout_b1k.inner, dtype=cutlass.TFloat32), + sB1k.layout, + ) + ) + tD1k = cute.make_tensor( + tD1.iterator, + tiled_mma1k.make_fragment_C(tiled_mma1k.partition_shape_C((V, T_PAD))).layout, + ) + tAhi_tf = cute.make_tensor(cute.recast_ptr(tAhi.iterator, dtype=cutlass.TFloat32), tAhi.layout) + tAlo_tf = cute.make_tensor(cute.recast_ptr(tAlo.iterator, dtype=cutlass.TFloat32), tAlo.layout) + # ---- warp 0 issues GEMM1 right away; warps 1..15 stage Kbar^T (sKbT[k, s], lanes = + # consecutive k: conflict-free), seed the GEMM2 accumulator, build the B2 tiles and write + # the new tokens' cache rows while the tensor cores run. + if warp_idx > 0: + for i in cutlass.range( + (K * T_PAD + NUM_THREADS - 33) // (NUM_THREADS - 32), unroll_full=True + ): + e_kb = i * (NUM_THREADS - 32) + (tidx - 32) + if e_kb < K * T_PAD: + k_kb = e_kb % K + s_idx = e_kb // K + b_val = cutlass.Float32(0.0) + if s_idx <= commit_len: + b_val = sK[s_idx, k_kb] * cute.math.exp( + sE[commit_len * K + k_kb] - sE[s_idx * K + k_kb], fastmath=True + ) + sKbT[k_kb * TP_STRIDE + s_idx] = b_val + if warp_idx == 0: + tiled_mma1.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma1, tD1, tAhi_tf[(None, None, kb_i)], tCrB1[(None, None, kb_i, 0)], tD1 + ) + tiled_mma1.set(tcgen05.Field.ACCUMULATE, True) + if cutlass.const_expr(S0_LO): + tiled_mma1k.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma1k, + tD1k, + tAlo_tf[(None, None, kb_i)], + tCrB1k[(None, None, kb_i, 0)], + tD1k, + ) + # tracks GEMM3 and GEMM1 + with cute.arch.elect_one(): + tcgen05.commit(mbar1) + else: + # sKbT complete before the B2 tiles read it + cute.arch.barrier(barrier_id=5, number_of_threads=NUM_THREADS - 32) + # ---- GEMM2 accumulator seed (overlaps the GEMM1/GEMM3 execution) + # D4 (output-combination) columns start from zero: GEMM2 half 0 accumulates into [D4 | dS] + st16 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + thr_stD4 = tcgen05.make_tmem_copy(st16, tD4).get_slice(wg_tidx) + tD4_st = thr_stD4.partition_D(tD4) + r_zero = cute.make_rmem_tensor(thr_stD4.partition_S(tD4).shape, cutlass.Float32) + for i in cutlass.range(cute.size(r_zero), unroll_full=True): + r_zero[i] = cutlass.Float32(0.0) + if wg == 0: + cute.copy(st16, r_zero, tD4_st) + # only this CTA's state rows (TMEM lanes [vh * V/2, +V/2)) are seeded and stored + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + for i in cutlass.range(32, unroll_full=True): + r_hi_flat[i] = r_state[i] * sEc[wg * 32 + i] + for jj in cutlass.range_constexpr(4): + if wg == jj: + cute.copy(st32, r_hi, tD2_st[(None, None, None, jj)]) + cute.arch.fence_view_async_tmem_store() + # ---- B2 tiles (k rows x s cols, K-major SW64) from sKbT while the MMAs run: + # lanes 0..15 -> row 2w, lanes 16..31 -> row 2w+1, four 32-row rounds + s_col = in_warp_tid % T_PAD + if warp_idx > 0: + for r in cutlass.range(4, unroll_full=True): + k_row = 32 * r + 2 * warp_idx + in_warp_tid // T_PAD + b_val2 = sKbT[k_row * TP_STRIDE + s_col] + b_hi = round_tf32(b_val2) + o_b2 = swz64_k16(k_row, s_col) + sB2hif[o_b2] = b_hi + sB2lof[o_b2] = b_val2 - b_hi + # warp 0 is issuing the MMAs: warps 1..4 cover its rows 32 r + {0, 1} + if warp_idx <= 4: + k_row0 = 32 * (warp_idx - 1) + in_warp_tid // T_PAD + b_val0 = sKbT[k_row0 * TP_STRIDE + s_col] + b_hi0 = round_tf32(b_val0) + o_b20 = swz64_k16(k_row0, s_col) + sB2hif[o_b20] = b_hi0 + sB2lof[o_b20] = b_val0 - b_hi0 + # ---- the new tokens' replay-cache rows and conv-window commits, from smem: warp 1 + p + # writes draft p's k / gate / v cache rows, its raw conv columns and beta, warps + # REPLAY_WARP0 + {0, 1, 2} commit the q / k / v conv windows (the raw inputs after the + # golden token). One CTA of a pair writes them. + if vh == 0: + if cutlass.const_expr(SPLIT_V == 2): + cute.arch.cluster_wait() + if (warp_idx >= 1) & (warp_idx <= NUM_SPEC): + p = warp_idx - 1 + t_c = commit_len + 1 + p + st_row4(k_cache, (slot, p, hk_off + k4), ld_shared_f32x4(saddr(sK, t_c * K + k4))) + st_row4(g_cache, (slot, p, hk_off + k4), ld_shared_f32x4(saddr(sG, t_c * K + k4))) + st_row4(v_cache, (slot, p, hv_off + k4), ld_shared_f32x4(saddr(sVall, t_c * V + k4))) + st_row4( + cs_q, + (slot, hk_off + k4, KERNEL_WIDTH - 1 + p), + ld_shared_f32x4(saddr(sXQ, (t_c + 3) * K + k4)), + ) + st_row4( + cs_k, + (slot, hk_off + k4, KERNEL_WIDTH - 1 + p), + ld_shared_f32x4(saddr(sXK, (t_c + 3) * K + k4)), + ) + st_row4( + cs_v, + (slot, hv_off + k4, KERNEL_WIDTH - 1 + p), + ld_shared_f32x4(saddr(sXV, (t_c + 3) * V + k4)), + ) + if in_warp_tid == 0: + beta_cache[slot, p, i_hv] = sBeta[t_c].to(BETA_CACHE_DTYPE) + if warp_idx == REPLAY_WARP0: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + st_row4( + cs_q, + (slot, hk_off + k4, w), + ld_shared_f32x4(saddr(sXQ, (commit_len + 1 + w) * K + k4)), + ) + if warp_idx == REPLAY_WARP0 + 1: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + st_row4( + cs_k, + (slot, hk_off + k4, w), + ld_shared_f32x4(saddr(sXK, (commit_len + 1 + w) * K + k4)), + ) + if warp_idx == REPLAY_WARP0 + 2: + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + st_row4( + cs_v, + (slot, hv_off + k4, w), + ld_shared_f32x4(saddr(sXV, (commit_len + 1 + w) * V + k4)), + ) + # ---- wait for GEMM1/GEMM3 + cute.arch.mbarrier_wait(mbar1, 0) + tcgen05_fence_after_thread_sync() + thr_ld1 = tcgen05.make_tmem_copy(ld16, tD1).get_slice(wg_tidx) + rD1 = cute.make_rmem_tensor(thr_ld1.partition_D(tD1).shape, cutlass.Float32) + cute.copy(ld16, thr_ld1.partition_S(tD1), rD1) + cute.arch.fence_view_async_tmem_load() + rD1_flat = cute.make_tensor(rD1.iterator, cute.make_layout((N1,))) + v_row = wg_tidx + # warpgroup 0: D3 (64 x 16) -> smem [L; B_hi; B_lo; -] + ld3 = cute.make_copy_atom( + tcgen05.Ld16x128bOp(tcgen05.Repetition(T_PAD * 32 // 128)), cutlass.Float32 + ) + thr_ld3 = tcgen05.make_tmem_copy(ld3, tD3).get_slice(wg_tidx) + tD3_ld = thr_ld3.partition_S(tD3) + sLB3 = cute.make_tensor( + sLB.iterator, + cute.make_layout(((4 * T_PAD, T_PAD), 1, 1), stride=((T_PAD, 1), 0, 0)), + ) + tD3_sm = thr_ld3.partition_D(sLB3) + rD3 = cute.make_rmem_tensor(tD3_sm.shape, cutlass.Float32) + if wg == 0: + cute.copy(ld3, tD3_ld, rD3) + cute.arch.fence_view_async_tmem_load() + cute.autovec_copy(rD3, tD3_sm) + cute.arch.barrier() + # ---- warpgroup 0: V' by forward substitution (row v per thread, four accumulators); + # warpgroups 1-2: output-combination operand Bcoef[t, s] = B_hi + B_lo (rows t < T_loop, + # s <= t; zero elsewhere), one element per thread, K-major SW64 16 x 16 tiles (hi / lo) + # that head the GEMM2 half-0 B tile + if wg == 0: + r_Y = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for s in cutlass.range(T_PAD, unroll_full=True): + y = cutlass.Float32(0.0) + if s < t_max: + u0 = rD1_flat[s] + rD1_flat[2 * T_PAD + s] + y = sBeta[s] * (sVall[s * V + v_row] - u0) + r_Y[s] = y + r_V = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + vp = cutlass.Float32(0.0) + if t < t_max: + acc0 = cutlass.Float32(0.0) + acc1 = cutlass.Float32(0.0) + acc2 = cutlass.Float32(0.0) + acc3 = cutlass.Float32(0.0) + for s in range(t): + if s % 4 == 0: + acc0 = acc0 + sLB[t, s] * r_V[s] + elif s % 4 == 1: + acc1 = acc1 + sLB[t, s] * r_V[s] + elif s % 4 == 2: + acc2 = acc2 + sLB[t, s] * r_V[s] + else: + acc3 = acc3 + sLB[t, s] * r_V[s] + vp = r_Y[t] - sBeta[t] * ((acc0 + acc1) + (acc2 + acc3)) + r_V[t] = vp + for t in cutlass.range(T_PAD, unroll_full=True): + sVpT[v_row * TP_STRIDE + t] = r_V[t] + else: + if tidx < 128 + T_PAD * T_PAD: + t4 = (tidx - 128) // T_PAD + s4 = (tidx - 128) % T_PAD + bc = cutlass.Float32(0.0) + if (t4 < T_loop) & (s4 <= t4): + bc = sLB[T_PAD + t4, s4] + sLB[2 * T_PAD + t4, s4] + bc_hi = round_tf32(bc) + o_b4 = swz64_k16(t4, s4) + sB4hif[o_b4] = bc_hi + sB4lof[o_b4] = bc - bc_hi + cute.arch.barrier() + # ---- A2 tiles (v rows x s cols) from sVpT, same 16-lanes-per-row mapping; s > c -> 0 + for r in cutlass.range(4, unroll_full=True): + v_row2 = 32 * r + 2 * warp_idx + in_warp_tid // T_PAD + a_val = sVpT[v_row2 * TP_STRIDE + s_col] + a_hi = round_tf32(a_val) + o_a2 = swz64_k16(v_row2, s_col) + sA2hif[o_a2] = a_hi + sA2lof[o_a2] = a_val - a_hi + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + sA2hi_tf = cute.make_tensor( + cute.recast_ptr(sA2hi.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), + sA2hi.layout, + ) + sA2lo_tf = cute.make_tensor( + cute.recast_ptr(sA2lo.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), + sA2lo.layout, + ) + tCrA2hi = tiled_mma2h.make_fragment_A(sA2hi_tf) + tCrA2lo = tiled_mma2h.make_fragment_A(sA2lo_tf) + # B2 halves: k rows 0..63 / 64..127 of the SW64 K-major tiles (a 4 KB shift keeps the swizzle) + tCrB2hi_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B2HI + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.TFloat32, + ), + smem_layout_b2h.outer, + ) + ) + tCrB2lo_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B2LO + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.TFloat32, + ), + smem_layout_b2h.outer, + ) + ) + tCrA2hi_q = tiled_mma2q.make_fragment_A(sA2hi_tf) + tCrA2lo_q = tiled_mma2q.make_fragment_A(sA2lo_tf) + tCrB2q_hi = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B4HI, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + smem_layout_b2q.outer, + ) + ) + tCrB2q_lo = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sScratch.iterator + OFF_B4LO, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + smem_layout_b2q.outer, + ) + ) + if warp_idx == 0: + # GEMM2 half 0 fused with the output combination: [D4 | dS_0..63] += A2 x [Bcoef; Kbar_0..63] + tiled_mma2q.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_lo[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2lo_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2) + tiled_mma2h.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2lo_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2lo[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2b) + + # ---- outputs of the new tokens: o_t = O0_t + D4[v, t] (columns 0..15 of GEMM2 half 0), + # warpgroup w writes new tokens j = w mod 4, bf16-rounded, into sVall rows j + cute.arch.mbarrier_wait(mbar2, 0) + tcgen05_fence_after_thread_sync() + thr_ld4 = tcgen05.make_tmem_copy(ld16, tD4).get_slice(wg_tidx) + rD4 = cute.make_rmem_tensor(thr_ld4.partition_D(tD4).shape, cutlass.Float32) + cute.copy(ld16, thr_ld4.partition_S(tD4), rD4) + cute.arch.fence_view_async_tmem_load() + rD4_flat = cute.make_tensor(rD4.iterator, cute.make_layout((T_PAD,))) + for j_new in cutlass.range(NEW_TOK, unroll_full=True): + if wg == j_new % 4: + t_out = commit_len + j_new + o_val = cutlass.Float32(0.0) + for tt in cutlass.range(T_PAD, unroll_full=True): + if tt == t_out: + o_val = rD1_flat[T_PAD + tt] + rD4_flat[tt] + sVall[j_new * V + v_row] = cutlass.Float32(cutlass.BFloat16(o_val)) + # ---- committed-state store through TMA: D2 slab -> swizzled smem box (row per thread, + # 16 B chunks land in distinct banks under the 128B swizzle) -> one bulk tensor store + if cutlass.const_expr(SPLIT_V == 2): + if vh == 1: + cute.arch.cluster_wait() + # All 16 warps produced disjoint value rows above; the barrier makes them visible to the + # norm epilogue, which assigns one complete token to each warp. + cute.arch.barrier() + # The state-owning warps of warpgroups 0/1 stage k columns 0..63 (GEMM2 half 0, mbar2), those + # of warpgroups 2/3 columns 64..127 (half 1, mbar2b) into the swizzled smem boxes; the first + # staging thread of each pair issues the TMA stores once its half is visible (named barriers + # 6 / 7) and waits for the bulk groups at the very end. The warps without state rows run the + # norm epilogue. + if wg < 2: + cute.arch.mbarrier_wait(mbar2, 0) + else: + cute.arch.mbarrier_wait(mbar2b, 0) + tcgen05_fence_after_thread_sync() + for sl in cutlass.range_constexpr(4): + if wg == sl: + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + cute.copy(ld32, tD2_ld[(None, None, None, sl)], rS) + cute.arch.fence_view_async_tmem_load() + if cutlass.const_expr(STATE_BF16): + # bf16 box: chunks [4 (sl % 2), +4) of box sl // 2, round to nearest + for c in cutlass.range(4, unroll_full=True): + chunk = (4 * (sl % 2) + c) ^ (wg_tidx % 8) + off_b = ( + (sl // 2) * ((V // SPLIT_V) * 64) + + 64 * (wg_tidx - (V // SPLIT_V) * vh) + + 8 * chunk + ) + for i in cutlass.range(8, unroll_full=True): + sBoxh[off_b + i] = cutlass.BFloat16(rS_flat[8 * c + i]) + else: + for c in cutlass.range(8, unroll_full=True): + off_b = ( + sl * ((V // SPLIT_V) * 32) + + 32 * (wg_tidx - (V // SPLIT_V) * vh) + + ((4 * c) ^ vx) + ) + for i in cutlass.range(4, unroll_full=True): + sBoxf[off_b + i] = rS_flat[4 * c + i] + # last TMEM reads of the kernel: warps 1..15 signal barrier 8, warp 0 frees TMEM below + tcgen05_fence_before_thread_sync() + if warp_idx > 0: + cute.arch.barrier_arrive(barrier_id=8, number_of_threads=NUM_THREADS) + cute.arch.fence_view_async_shared() + st_tid0 = (V // 2) * vh + if cutlass.const_expr(SPLIT_V == 1) or ((warp_idx % 4) // 2 == vh): + if wg < 2: + cute.arch.barrier(barrier_id=6, number_of_threads=256 // SPLIT_V) + if tidx == st_tid0: + for s in cutlass.range_constexpr(NUM_SLABS // 2): + cute.copy(tma_atom_t, tS_list[s], tG_list[s]) + cute.arch.cp_async_bulk_commit_group() + else: + cute.arch.barrier(barrier_id=7, number_of_threads=256 // SPLIT_V) + if tidx == 256 + st_tid0: + for s in cutlass.range_constexpr(NUM_SLABS // 2, NUM_SLABS): + cute.copy(tma_atom_t, tS_list[s], tG_list[s]) + cute.arch.cp_async_bulk_commit_group() + if warp_idx == 0: + # all TMEM reads are done (barrier 8): free the columns while the store drains + cute.arch.barrier(barrier_id=8, number_of_threads=NUM_THREADS) + tcgen05_fence_after_thread_sync() + cute.arch.relinquish_tmem_alloc_permit() + cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS) + # ---- gated RMSNorm epilogue (CTA 0 of a pair) on the warps (warp % 4) in {2, 3}: one token + # per warp, p_v = o_v * w_v * sigmoid(g_v), out = bf16(p_v * rsqrt(mean(o_v^2) + eps)) + token_idx = (warp_idx // 4) * 2 + (warp_idx % 4 - 2) + if (vh == 0) & ((warp_idx % 4) // 2 == 1): + while token_idx < NEW_TOK: + pre_values = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) + rms_partial = cutlass.Float32(0.0) + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + value = sVall[token_idx * V + v_idx] + gate = sOnG[token_idx * V + v_idx] + sigmoid_gate = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) + ) + pre_values[i] = value * sOnW[v_idx] * sigmoid_gate + rms_partial += value * value + for offset in [16, 8, 4, 2, 1]: + rms_partial += cute.arch.shuffle_sync_bfly( + rms_partial, offset=offset, mask=-1, mask_and_clamp=31 + ) + inv_rms = cute.math.rsqrt( + rms_partial / cutlass.Float32(V) + cutlass.Float32(onorm_eps), + fastmath=True, + ) + output_row = i_n * NEW_TOK + token_idx + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + o[0, output_row, i_hv, v_idx] = cutlass.BFloat16(pre_values[i] * inv_rms) + token_idx = token_idx + NUM_THREADS // 32 + # the state stores must have read their smem boxes before the CTA exits + if (tidx == st_tid0) | (tidx == 256 + st_tid0): + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py index c67ed3e2e55a..ac124d733866 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py @@ -50,11 +50,18 @@ Kernel shape contract: ``K == V == 128``, conv width ``W == 4``, ``HV == H``, one V pass (TILE_V=128, 512 threads). ``num_spec`` is a compile-time constant per cache allocation. + +Kernel selection (``_select_verify_kernel``): calls on the runtime verify +contract (packed token rows, fused gated RMSNorm with a bf16 output, 3xTF32 +recurrence, vector-aligned operands) run ``kda_decode_mtp_v2_kernel`` +(``kda_mtp_decode_v2.py``: the same arithmetic in the same order, bitwise +identical results, shorter critical path); every other variant runs +``kda_decode_mtp_kernel``. """ import os from math import gcd -from typing import Optional, Tuple +from typing import Callable, Literal, NamedTuple, Optional, Tuple import torch @@ -71,12 +78,14 @@ from cutlass.cute.runtime import from_dlpack from ..attention.kernels.blackwell.kimi_k3_kda.kda_mtp_decode import ( + MAX_NUM_SPEC, NUM_THREADS, SPLIT_MAX_TILES, T_PAD, TILE_K, kda_decode_mtp_kernel, ) + from ..attention.kernels.blackwell.kimi_k3_kda.kda_mtp_decode_v2 import kda_decode_mtp_v2_kernel else: raise ImportError("Kimi K3 KDA MTP decode requires NVIDIA CUTLASS DSL") @@ -86,6 +95,148 @@ if IS_CUTLASS_DSL_AVAILABLE: + class _VerifyMmaAtoms(NamedTuple): + """tcgen05 MMA atoms and operand smem layouts shared by the verify kernels.""" + + tiled_mma1: "cute.TiledMma" + smem_layout_b1: "cute.ComposedLayout" + tiled_mma_f32a: "cute.TiledMma" + tiled_mma1k: "cute.TiledMma" + smem_layout_b1k: "cute.ComposedLayout" + tiled_mma2h: "cute.TiledMma" + smem_layout_b2h: "cute.ComposedLayout" + smem_layout_a2: "cute.ComposedLayout" + tiled_mma2q: "cute.TiledMma" + smem_layout_b2q: "cute.ComposedLayout" + tiled_mma3: "cute.TiledMma" + smem_layout_a3: "cute.ComposedLayout" + smem_layout_b3: "cute.ComposedLayout" + tiled_mma4: "cute.TiledMma" + + def _verify_mma_atoms(op_dtype, V: int, K: int, n1: int) -> _VerifyMmaAtoms: + """The verify MMAs for ``op_dtype`` operands (fp32 = 3xTF32, or bf16) with fp32 + accumulators. + + GEMM1: M = V rows of the state (TMEM), N = ``n1`` operand rows ([K~; Q~; K~lo] for + tf32). ``tiled_mma_f32a`` is the fp32 A-fragment layout of the (V x K) state in TMEM + (the GEMM2 accumulator seed / state store views); for 3xTF32 it is GEMM1's atom. + GEMM1b: the S0_lo sweep only needs K~hi (N = T_PAD). GEMM2: (V x T_PAD) x (T_PAD x K), + both operands in smem, issued as two N = K/2 column halves so the state store can start + on half 0; half 0 also carries the output-combination GEMM (N = T_PAD (Bcoef) + K/2). + GEMM3: T x T coupling matrices, A3 = [Kf; Qf] padded to 64 rows, B3 = Kb (16 rows). + ``tiled_mma4`` only provides the accumulator layout of the output combination + D4[v, t] = sum_s V'[s, v] Bcoef[t, s], the first T_PAD columns of GEMM2 half 0. + """ + f32 = cutlass.Float32 + + def mma(dtype, shape_mn, source): + return sm100_utils.make_trivial_tiled_mma( + dtype, + OperandMajorMode.K, + OperandMajorMode.K, + f32, + tcgen05.CtaGroup.ONE, + shape_mn, + source, + ) + + tmem, smem = tcgen05.OperandSource.TMEM, tcgen05.OperandSource.SMEM + tiled_mma_f32a = mma(f32, (V, 3 * T_PAD), tmem) + if op_dtype is f32 and n1 == 3 * T_PAD: + tiled_mma1 = tiled_mma_f32a + else: + tiled_mma1 = mma(op_dtype, (V, n1), tmem) + tiled_mma1k = mma(f32, (V, T_PAD), tmem) + tiled_mma2h = mma(op_dtype, (V, K // 2), smem) + tiled_mma2q = mma(op_dtype, (V, T_PAD + K // 2), smem) + tiled_mma3 = mma(op_dtype, (4 * T_PAD, T_PAD), smem) + return _VerifyMmaAtoms( + tiled_mma1=tiled_mma1, + smem_layout_b1=sm100_utils.make_smem_layout_b(tiled_mma1, (V, n1, K), op_dtype, 1), + tiled_mma_f32a=tiled_mma_f32a, + tiled_mma1k=tiled_mma1k, + smem_layout_b1k=sm100_utils.make_smem_layout_b(tiled_mma1k, (V, T_PAD, K), f32, 1), + tiled_mma2h=tiled_mma2h, + smem_layout_b2h=sm100_utils.make_smem_layout_b( + tiled_mma2h, (V, K // 2, T_PAD), op_dtype, 1 + ), + smem_layout_a2=sm100_utils.make_smem_layout_a( + tiled_mma2h, (V, K // 2, T_PAD), op_dtype, 1 + ), + tiled_mma2q=tiled_mma2q, + smem_layout_b2q=sm100_utils.make_smem_layout_b( + tiled_mma2q, (V, T_PAD + K // 2, T_PAD), op_dtype, 1 + ), + tiled_mma3=tiled_mma3, + smem_layout_a3=sm100_utils.make_smem_layout_a( + tiled_mma3, (4 * T_PAD, T_PAD, K), op_dtype, 1 + ), + smem_layout_b3=sm100_utils.make_smem_layout_b( + tiled_mma3, (4 * T_PAD, T_PAD, K), op_dtype, 1 + ), + tiled_mma4=mma(f32, (V, T_PAD), smem), + ) + + class _StateTmaAtoms(NamedTuple): + """TMA load of the initial state S0 and bulk store of the committed state.""" + + tma_atom_l: "cpasync.CopyAtom" + mS_tma: "cute.Tensor" + box_layout: "cute.ComposedLayout" + tma_atom_t: "cpasync.CopyAtom" + mT_tma: "cute.Tensor" + box_layout_t: "cute.ComposedLayout" + + def _state_tiles(state: cute.Tensor, V: int, K: int) -> cute.Tensor: + """(V, K, tiles) view of a V-first state pool with its own strides (the per-layer state + cache may be a strided slice of a larger buffer); the tile mode is the pool index (flat + ``[pool, V, K]`` layout) or the (head, slot) pair (``[pool, H, V, K]``).""" + stride = state.layout.stride + if cute.rank(state.shape) == 3: + return cute.make_tensor( + state.iterator, + cute.make_layout((V, K, state.shape[0]), stride=(stride[1], stride[2], stride[0])), + ) + return cute.make_tensor( + state.iterator, + cute.make_layout( + (V, K, (state.shape[1], state.shape[0])), + stride=(stride[2], stride[3], (stride[1], stride[0])), + ), + ) + + def _state_tma_atoms( + h0: cute.Tensor, ht: cute.Tensor, V: int, K: int, split_v: int + ) -> _StateTmaAtoms: + """S0 is loaded from ``h0`` and the committed state stored to ``ht`` in 128B-swizzled + boxes of 32 fp32 or 64 bf16 state columns; each CTA of a ``split_v`` group stores its + V / split_v rows.""" + box_cols = 128 // (h0.element_type.width // 8) + box_atom = sm100_utils.make_smem_layout_atom( + tcgen05.SmemLayoutAtomKind.K_SW128, h0.element_type + ) + box_layout = cute.tile_to_shape(box_atom, (V, box_cols), order=(0, 1)) + box_layout_t = cute.tile_to_shape(box_atom, (V // split_v, box_cols), order=(0, 1)) + tma_atom_l, mS_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileG2SOp(), _state_tiles(h0, V, K), box_layout, (V, box_cols) + ) + tma_atom_t, mT_tma = cpasync.make_tiled_tma_atom( + cpasync.CopyBulkTensorTileS2GOp(), + _state_tiles(ht, V, K), + box_layout_t, + (V // split_v, box_cols), + ) + return _StateTmaAtoms(tma_atom_l, mS_tma, box_layout, tma_atom_t, mT_tma, box_layout_t) + + def _launch_verify(launched, HV: int, N, split_v: int, stream) -> None: + if split_v == 2: + # cluster pair: the two CTAs of a (head, request) must be co-scheduled (in-place state) + launched.launch( + grid=(HV, N, 2), block=[NUM_THREADS, 1, 1], cluster=(1, 1, 2), stream=stream + ) + else: + launched.launch(grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream) + @cute.jit def _run_kda_decode_mtp( h0: cute.Tensor, @@ -147,159 +298,13 @@ def _run_kda_decode_mtp( t_max = 2 * NUM_SPEC + 1 smem_qk_layout = cute.make_layout((t_max, K), stride=(K, 1)) # BF16_MMA: bf16 tensor-core operands ([K~; Q~] only, K = 16 per instruction). - OP_DTYPE = cutlass.BFloat16 if BF16_MMA else cutlass.Float32 - N1 = 2 * T_PAD if BF16_MMA else 3 * T_PAD - # GEMM1: M = V rows of the state (TMEM), N = operand rows ([K~; Q~; K~lo] for tf32). - tiled_mma1 = sm100_utils.make_trivial_tiled_mma( - OP_DTYPE, - OperandMajorMode.K, - OperandMajorMode.K, - cutlass.Float32, - tcgen05.CtaGroup.ONE, - (V, N1), - tcgen05.OperandSource.TMEM, - ) - smem_layout_b1 = sm100_utils.make_smem_layout_b(tiled_mma1, (V, N1, K), OP_DTYPE, 1) - # fp32 A-fragment layout of the (V x K) state in TMEM: the GEMM2 accumulator seed / state - # store views use it in every mode. - tiled_mma_f32a = sm100_utils.make_trivial_tiled_mma( - cutlass.Float32, - OperandMajorMode.K, - OperandMajorMode.K, - cutlass.Float32, - tcgen05.CtaGroup.ONE, - (V, 3 * T_PAD), - tcgen05.OperandSource.TMEM, - ) - # GEMM1b: the S0_lo sweep only needs K~hi (N = T_PAD) -> a third of the N=48 work. - tiled_mma1k = sm100_utils.make_trivial_tiled_mma( - cutlass.Float32, - OperandMajorMode.K, - OperandMajorMode.K, - cutlass.Float32, - tcgen05.CtaGroup.ONE, - (V, T_PAD), - tcgen05.OperandSource.TMEM, - ) - smem_layout_b1k = sm100_utils.make_smem_layout_b( - tiled_mma1k, (V, T_PAD, K), cutlass.Float32, 1 - ) - # TMA bulk store of the committed state: (V, K, tiles) view of ht, one 128B-swizzled - # (V/SPLIT_V x 32 fp32 | 64 bf16) box per 128 B-wide column slab. The view takes ht's own strides (the - # per-layer state cache is a strided slice of a larger buffer); the tile mode is the pool - # index (flat layout) or the (head, slot) pair. - if cutlass.const_expr(cute.rank(ht.shape) == 3): - mT_vkl = cute.make_tensor( - ht.iterator, - cute.make_layout( - (V, K, ht.shape[0]), - stride=(ht.layout.stride[1], ht.layout.stride[2], ht.layout.stride[0]), - ), - ) - else: - mT_vkl = cute.make_tensor( - ht.iterator, - cute.make_layout( - (V, K, (ht.shape[1], ht.shape[0])), - stride=( - ht.layout.stride[2], - ht.layout.stride[3], - (ht.layout.stride[1], ht.layout.stride[0]), - ), - ), - ) - # 128 B swizzled box rows: 32 fp32 or 64 bf16 state columns per box - box_cols = 128 // (h0.element_type.width // 8) - box_atom = sm100_utils.make_smem_layout_atom( - tcgen05.SmemLayoutAtomKind.K_SW128, h0.element_type - ) - box_layout = cute.tile_to_shape(box_atom, (V, box_cols), order=(0, 1)) - # committed-state store: each CTA writes its V/SPLIT_V rows through (V/SPLIT_V x box_cols) boxes - box_layout_t = cute.tile_to_shape(box_atom, (V // SPLIT_V, box_cols), order=(0, 1)) - tma_atom_t, mT_tma = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileS2GOp(), mT_vkl, box_layout_t, (V // SPLIT_V, box_cols) - ) - # TMA bulk load of the initial state S0 into the same box layout (issued at kernel start). - if cutlass.const_expr(cute.rank(h0.shape) == 3): - mS_vkl = cute.make_tensor( - h0.iterator, - cute.make_layout( - (V, K, h0.shape[0]), - stride=(h0.layout.stride[1], h0.layout.stride[2], h0.layout.stride[0]), - ), - ) - else: - mS_vkl = cute.make_tensor( - h0.iterator, - cute.make_layout( - (V, K, (h0.shape[1], h0.shape[0])), - stride=( - h0.layout.stride[2], - h0.layout.stride[3], - (h0.layout.stride[1], h0.layout.stride[0]), - ), - ), - ) - tma_atom_l, mS_tma = cpasync.make_tiled_tma_atom( - cpasync.CopyBulkTensorTileG2SOp(), mS_vkl, box_layout, (V, box_cols) - ) - # GEMM2: (V x T_PAD) x (T_PAD x K), both operands in smem, tf32; issued as two N = K/2 - # column halves so the state store can start on half 0. - tiled_mma2h = sm100_utils.make_trivial_tiled_mma( - OP_DTYPE, - OperandMajorMode.K, - OperandMajorMode.K, - cutlass.Float32, - tcgen05.CtaGroup.ONE, - (V, K // 2), - tcgen05.OperandSource.SMEM, - ) - smem_layout_b2h = sm100_utils.make_smem_layout_b( - tiled_mma2h, (V, K // 2, T_PAD), OP_DTYPE, 1 - ) - smem_layout_a2 = sm100_utils.make_smem_layout_a( - tiled_mma2h, (V, K // 2, T_PAD), OP_DTYPE, 1 - ) - # GEMM2 half 0 also carries the output-combination GEMM: N = T_PAD (Bcoef) + K/2 (Kbar rows 0..63). - tiled_mma2q = sm100_utils.make_trivial_tiled_mma( - OP_DTYPE, - OperandMajorMode.K, - OperandMajorMode.K, - cutlass.Float32, - tcgen05.CtaGroup.ONE, - (V, T_PAD + K // 2), - tcgen05.OperandSource.SMEM, - ) - smem_layout_b2q = sm100_utils.make_smem_layout_b( - tiled_mma2q, (V, T_PAD + K // 2, T_PAD), OP_DTYPE, 1 - ) - # GEMM3: T x T coupling matrices, A3 = [Kf; Qf] padded to 64 rows, B3 = Kb (16 rows). - tiled_mma3 = sm100_utils.make_trivial_tiled_mma( - OP_DTYPE, - OperandMajorMode.K, - OperandMajorMode.K, - cutlass.Float32, - tcgen05.CtaGroup.ONE, - (4 * T_PAD, T_PAD), - tcgen05.OperandSource.SMEM, - ) - smem_layout_a3 = sm100_utils.make_smem_layout_a( - tiled_mma3, (4 * T_PAD, T_PAD, K), OP_DTYPE, 1 - ) - smem_layout_b3 = sm100_utils.make_smem_layout_b( - tiled_mma3, (4 * T_PAD, T_PAD, K), OP_DTYPE, 1 - ) - # Output combination D4[v, t] = sum_s V'[s, v] Bcoef[t, s]: computed as the first 16 columns - # of GEMM2 half 0; this tiled MMA only provides the accumulator layout of those columns. - tiled_mma4 = sm100_utils.make_trivial_tiled_mma( - cutlass.Float32, - OperandMajorMode.K, - OperandMajorMode.K, - cutlass.Float32, - tcgen05.CtaGroup.ONE, - (V, T_PAD), - tcgen05.OperandSource.SMEM, + mma = _verify_mma_atoms( + cutlass.BFloat16 if BF16_MMA else cutlass.Float32, + V, + K, + 2 * T_PAD if BF16_MMA else 3 * T_PAD, ) + state_tma = _state_tma_atoms(h0, ht, V, K, SPLIT_V) launched = kda_decode_mtp_kernel( h0, x_q, @@ -349,37 +354,130 @@ def _run_kda_decode_mtp( QUANTIZE_OUTPUT, stage_timing, PROFILE_STAGES, - tiled_mma1, - tiled_mma3, - tiled_mma1k, - smem_layout_b1, - smem_layout_a2, - smem_layout_a3, - smem_layout_b3, - smem_layout_b1k, - tiled_mma4, - tma_atom_t, - mT_tma, - box_layout, - tma_atom_l, - mS_tma, - tiled_mma2h, - smem_layout_b2h, - box_layout_t, - tiled_mma2q, - smem_layout_b2q, + mma.tiled_mma1, + mma.tiled_mma3, + mma.tiled_mma1k, + mma.smem_layout_b1, + mma.smem_layout_a2, + mma.smem_layout_a3, + mma.smem_layout_b3, + mma.smem_layout_b1k, + mma.tiled_mma4, + state_tma.tma_atom_t, + state_tma.mT_tma, + state_tma.box_layout, + state_tma.tma_atom_l, + state_tma.mS_tma, + mma.tiled_mma2h, + mma.smem_layout_b2h, + state_tma.box_layout_t, + mma.tiled_mma2q, + mma.smem_layout_b2q, SPLIT_V, BF16_MMA, - tiled_mma_f32a, + mma.tiled_mma_f32a, REPLAY_ONLY, ) - if cutlass.const_expr(SPLIT_V == 2): - # cluster pair: the two CTAs of a (head, request) must be co-scheduled (in-place state) - launched.launch( - grid=(HV, N, 2), block=[NUM_THREADS, 1, 1], cluster=(1, 1, 2), stream=stream - ) - else: - launched.launch(grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream) + _launch_verify(launched, HV, N, SPLIT_V, stream) + + @cute.jit + def _run_kda_decode_mtp_v2( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + o: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + ssm_state_indices: cute.Tensor, + num_accepted_tokens: cute.Tensor, + N: cutlass.Int32, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + SPLIT_V: cutlass.Constexpr[int], + stream: cuda.CUstream, + ): + """Host side of ``kda_decode_mtp_v2_kernel`` (packed rows, fused gated RMSNorm, bf16 + output, 3xTF32); the state is committed in place.""" + mma = _verify_mma_atoms(cutlass.Float32, V, K, 3 * T_PAD) + state_tma = _state_tma_atoms(h0, h0, V, K, SPLIT_V) + launched = kda_decode_mtp_v2_kernel( + h0, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + A_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + o, + k_cache, + g_cache, + v_cache, + beta_cache, + ssm_state_indices, + num_accepted_tokens, + cute.make_layout((2 * NUM_SPEC + 1, K), stride=(K, 1)), + mma.tiled_mma1, + mma.tiled_mma3, + mma.tiled_mma1k, + mma.smem_layout_b1, + mma.smem_layout_a2, + mma.smem_layout_a3, + mma.smem_layout_b3, + mma.smem_layout_b1k, + mma.tiled_mma4, + state_tma.tma_atom_t, + state_tma.mT_tma, + state_tma.box_layout, + state_tma.tma_atom_l, + state_tma.mS_tma, + mma.tiled_mma2h, + mma.smem_layout_b2h, + state_tma.box_layout_t, + mma.tiled_mma2q, + mma.smem_layout_b2q, + scale, + HV, + K, + V, + NUM_SPEC, + KERNEL_WIDTH, + lower_bound, + onorm_eps, + USE_FLAT_LAYOUT, + SPLIT_V, + ) + _launch_verify(launched, HV, N, SPLIT_V, stream) def _require_stride_layout( @@ -601,6 +699,64 @@ def _layout_key(tensor: torch.Tensor, dynamic_layout: bool = False, *, assumed_a return (tensor.dtype, shape, stride, _fits_32bit_stride(tensor), assumed_align) +def _vec4_rows(tensor: torch.Tensor, dim: int = -1) -> bool: + """True when every group of four consecutive elements along the unit-stride dimension `dim` + that starts at a multiple of four is one aligned 8 B (bf16) or 16 B (fp32) access: the v2 + kernel moves its per-token rows with such vector loads and stores.""" + if tensor.dtype not in (torch.bfloat16, torch.float32) or tensor.stride(dim) != 1: + return False + if tensor.data_ptr() % (4 * tensor.element_size()): + return False + return all( + stride % 4 == 0 for d, stride in enumerate(tensor.stride()) if d != dim % tensor.ndim + ) + + +def _v2_operands_ok(*, w_q, w_k, w_v, rows, conv_caches) -> bool: + """Operand layouts / dtypes the v2 kernel's vector accesses support (see `_vec4_rows`): + token rows and replay caches vectorized along their last dimension, the dim-contiguous conv + caches along the channel dimension, conv weights as contiguous [channel][4] rows.""" + weights_ok = all( + w.dtype == w_q.dtype and w.is_contiguous() and w.shape[1] == 4 and _vec4_rows(w) + for w in (w_q, w_k, w_v) + ) + return ( + weights_ok + and all(_vec4_rows(t) for t in rows) + and all(_vec4_rows(t, dim=1) for t in conv_caches) + ) + + +def _select_verify_kernel( + *, + packed_token_layout: bool, + fuse_output_norm: bool, + quantize_output: bool, + bf16_math: bool, + replay_only: bool, + num_spec: int, + v2_operands_ok: Callable[[], bool], +) -> Literal["v1", "v2"]: + """The verify kernel a call runs (see the module docstring). + + ``v2`` (``kda_decode_mtp_v2_kernel``) serves the runtime verify contract: packed token rows, + the fused gated RMSNorm with a bf16 output, the 3xTF32 recurrence and vector-aligned operands + (``v2_operands_ok``, evaluated last). ``v1`` (``kda_decode_mtp_kernel``) serves everything + else; both compute bitwise the same results where both apply. + """ + if ( + packed_token_layout + and fuse_output_norm + and not quantize_output + and not bf16_math + and not replay_only + and num_spec <= MAX_NUM_SPEC + and v2_operands_ok() + ): + return "v2" + return "v1" + + # (device_index, enabled) -> persistent int32 [1] control tensor. Keys are # plain values (not tensor identities), so entries never go stale. _precompute_control_cache = {} @@ -887,6 +1043,84 @@ def kda_mtp_decode_impl( split_v = 2 if HV * N <= SPLIT_MAX_TILES else 1 bf16_math = _bf16_math_enabled() + kernel = _select_verify_kernel( + packed_token_layout=packed_token_layout, + fuse_output_norm=fuse_output_norm, + quantize_output=quantize_output, + bf16_math=bf16_math, + replay_only=replay_only, + num_spec=num_spec, + v2_operands_ok=lambda: _v2_operands_ok( + w_q=w_q, + w_k=w_k, + w_v=w_v, + rows=(x_q_arg, x_k_arg, x_v_arg, g, onorm_g, k_cache, g_cache, v_cache), + conv_caches=(cs_q, cs_k, cs_v), + ), + ) + if kernel == "v2": + v2_args = ( + (h0_arg, 16), + (x_q_arg, 16), + (x_k_arg, 16), + (x_v_arg, 16), + (w_q, 16), + (w_k, 16), + (w_v, 16), + (cs_q, 16), + (cs_k, 16), + (cs_v, 16), + (A_log, 16), + (g, 16), + (dt_bias, 16), + (beta, 16), + (onorm_g, 16), + (onorm_weight, 16), + (out, 16), + (k_cache, 16), + (g_cache, 16), + (v_cache, 16), + (beta_cache, beta_cache_assumed_align), + (ssm_state_indices, 4), + (num_accepted_tokens, 4), + ) + v2_key = ( + "v2", + tuple(_layout_key(t, dynamic_layout=True, assumed_align=a) for t, a in v2_args), + scale, + HV, + K, + V_dim, + num_spec, + W, + float(lower_bound), + float(onorm_eps), + use_flat_layout, + split_v, + ) + if v2_key not in _compiled_cache: + logger.info( + f"kda_mtp_decode: compiling v2 variant H={HV} num_spec={num_spec} split_v={split_v}" + ) + _compiled_cache[v2_key] = cute.compile( + _run_kda_decode_mtp_v2, + *[_dlpack_arg(t, assumed_align=a) for t, a in v2_args], + N, + scale=scale, + HV=HV, + K=K, + V=V_dim, + NUM_SPEC=num_spec, + KERNEL_WIDTH=W, + lower_bound=float(lower_bound), + onorm_eps=float(onorm_eps), + USE_FLAT_LAYOUT=use_flat_layout, + SPLIT_V=split_v, + stream=stream, + ) + _compiled_cache[v2_key](*[_dlpack_arg(t, assumed_align=a) for t, a in v2_args], N, stream) + return out + key = ( x_q.dtype, scale, diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py index aa936b063399..bca3d6206507 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py @@ -1321,6 +1321,130 @@ def small(*shape, scale): assert not conv_pool[others].any() +def _packed_fused_norm_inputs(B, H, num_spec, state_dtype): + """Runtime-contract inputs for the packed / fused-norm path: section-strided q / k / v / + output-gate views of one fused projection, the bf16 beta tail of [f_a | beta], replay caches + and conv windows chained from a CPU-golden round 1, mixed acceptance 0..num_spec and + non-identity state slots.""" + data, _ = _round2_inputs(B, H, num_spec) + packed, _ = _pack_new_token_projection_views(data) + slots = torch.arange(B - 1, -1, -1, dtype=torch.int32, device="cuda") + packed["ssm_state_indices"] = slots + for name in ( + "initial_state_cute", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + slot_order = torch.empty_like(packed[name]) + slot_order[slots.long()] = packed[name] + packed[name] = slot_order + pool_accepted = torch.zeros(B, dtype=torch.int32, device="cuda") + pool_accepted[slots.long()] = data["num_accepted_tokens"] + packed["num_accepted_tokens"] = pool_accepted + packed["initial_state_cute"] = packed["initial_state_cute"].to(state_dtype) + packed["onorm_weight"] = torch.linspace( + 0.75, 1.25, data["V"], dtype=torch.float32, device="cuda" + ) + packed["onorm_eps"] = 1e-5 + return packed + + +_VERIFY_RESULTS = ( + "out", + "state_v_first", + "k_cache", + "g_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", +) + + +def _run_with_verify_kernel(monkeypatch, kernel, run): + """``run()`` with the verify kernel forced to ``kernel``; asserts the selection would have + picked it on its own unless ``kernel`` is ``"v1"`` (the reference).""" + from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops as mtp_ops + + select = mtp_ops._select_verify_kernel + selected = [] + + def forced(**kwargs): + selected.append(select(**kwargs)) + return kernel + + with monkeypatch.context() as patch: + patch.setattr(mtp_ops, "_select_verify_kernel", forced) + result = run() + assert selected, "the call did not reach the verify kernel selection" + if kernel != "v1": + assert selected == [kernel], f"selection picked {selected}, not {kernel}" + return result + + +@pytest.mark.parametrize( + "B,H,num_spec", + [(1, 6, 6), (4, 6, 7), (16, 6, 6), (32, 6, 7), (128, 6, 6), (8, 12, 7)], + ids=lambda v: str(v), +) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +@pytest.mark.parametrize("strided_state", [False, True], ids=["pool", "strided"]) +def test_v2_kernel_matches_v1_bitwise(B, H, num_spec, state_dtype, strided_state, monkeypatch): + """The latency-scheduled verify kernel (packed rows + fused gated RMSNorm) performs the same + arithmetic as ``kda_decode_mtp_kernel`` in the same order: the outputs, committed state, + conv pools and replay caches must match bit for bit. Covers the cluster-pair + (HV * B <= 64) and single-CTA grids, fp32 / bf16 state pools and a state that is a strided + slice of a larger buffer (non-flat TMA view).""" + packed = _packed_fused_norm_inputs(B, H, num_spec, state_dtype) + + def run(): + override = None + if strided_state: + shape = tuple(packed["initial_state_cute"].shape) + big = torch.zeros((shape[0], 3) + shape[1:], dtype=state_dtype, device="cuda") + override = big[:, 1] + return cute_run( + packed, packed_token_layout=True, fuse_output_norm=True, state_override=override + ) + + ref = _run_with_verify_kernel(monkeypatch, "v1", run) + new = _run_with_verify_kernel(monkeypatch, "v2", run) + for name in _VERIFY_RESULTS: + assert torch.equal(new[name], ref[name]), f"{name}: v2 differs from v1" + + +def test_verify_kernel_selection(): + """v2 serves only the runtime verify contract; every other variant stays on v1.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import _select_verify_kernel + + contract = dict( + packed_token_layout=True, + fuse_output_norm=True, + quantize_output=False, + bf16_math=False, + replay_only=False, + num_spec=7, + v2_operands_ok=lambda: True, + ) + assert _select_verify_kernel(**contract) == "v2" + for name, value in ( + ("packed_token_layout", False), + ("fuse_output_norm", False), + ("quantize_output", True), + ("bf16_math", True), + ("replay_only", True), + ("num_spec", 8), + ("v2_operands_ok", lambda: False), + ): + assert _select_verify_kernel(**{**contract, name: value}) == "v1", name + + if __name__ == "__main__": import sys From 3f05bfacf31a737968dc5589b9a0138d55fd52f7 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:11:28 -0700 Subject: [PATCH 3/8] [None][perf] Kimi K3 KDA MTP verify: two-CTAs-per-SM kernel for multi-wave grids Add kda_decode_mtp_2cta_kernel: the v2 verify arithmetic in v2's order with a footprint that lets two (head, request) tiles share an SM - 256 tensor-memory columns (S0 hi / lo as sequential GEMM1 A operands, GEMM2 seeded in the same columns), <= 114 KB shared memory (one 64 KB region reused across phases, raw conv inputs kept bf16 and staged through cp.async), 64 registers per thread (min_blocks_per_mp=2) and an early release of the tensor-memory allocation permit. Named barrier ids are compile-time constants <= 4 so that two CTAs can also co-reside where an SM limits a CTA's named barriers. fp32 and bf16 recurrent states are supported (a bf16 S0 row stays packed in registers). Selection: _select_verify_kernel runs it for calls on the v2 verify contract whose grid has more (head, request) tiles than the device has SMs, when the operands have the dtypes it stages (bf16 projections and conv caches, fp32 replay gates, fp32/bf16 state); single-wave grids keep the v2 kernel (faster for a tile alone on an SM at 3-5 drafts with a bf16 state and at every draft count with an fp32 state, see the previous commit). Outputs, committed state, conv windows and replay caches are bitwise identical to kda_decode_mtp_kernel / the v2 kernel. The shared v2 inline-asm helpers the kernel imports are now public (asm_words, st_asm). The two-CTA kernel imports kda_mtp_decode.py's shared helpers and v2's vector helpers, its host function uses the shared MMA and TMA helpers, and the v2 and two-CTA launches share one compile-and-launch path. Evidence: GB300 (a34dcd2207, 5ed174863c, 6158a944b6): verify b128 MTP6 bf16 state 45.2 -> 32.9 us per call; KDA decode chain b128 82.9 -> 75.0 us; two tiles per SM finish ~20% (fp32 state) to ~27% (bf16 state) sooner than two waves of v2 tiles. GB300 job 1368796 (cold L2): multi-wave grids 18-36% (bf16 state) / 10-23% (fp32 state) faster than v2, e.g. bf16 dep16 b8 mtp7 52.24 -> 33.24 us. Rubin (with the MXFP8 epilogue of the next commit), dep16 per call: b8 mtp7 30.61 -> 19.59 us, b16 mtp5 53.71 -> 34.82 us; dep16 b8 mtp7 gen-only 32.38 -> 31.37 ms/step (-3.1%, jobs 792567/792568). Platform: S (SM100/103 BF16-output verify with more tiles than SMs, e.g. dep16; on SM107 the target configurations use the MXFP8 output, next commit). Tests: test_kda_mtp_decode_cute_parity.py::test_two_cta_kernel_matches_v1_bitwise (single- and multi-wave grids, 2..7 drafts, fp32/bf16 state, K3 shape 8 x 96 heads), ::test_verify_kernel_selection; test_kimi_kda_fused_verify_parity.py::test_fused_verify_multi_wave_runs_two_cta_kernel (mixer, two rounds with mixed acceptance, bitwise vs the one-CTA kernel). GB300 (OCI AGA) job 1387423. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../kimi_k3_kda/kda_mtp_decode_2cta.py | 1196 +++++++++++++++++ .../kimi_k3_kda/kda_mtp_decode_v2.py | 16 +- .../cute_dsl_kimi_k3_kda_mtp_ops.py | 228 +++- .../test_kda_mtp_decode_cute_parity.py | 62 +- .../test_kimi_kda_fused_verify_parity.py | 67 + 5 files changed, 1509 insertions(+), 60 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_2cta.py diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_2cta.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_2cta.py new file mode 100644 index 000000000000..ce1b734f165e --- /dev/null +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_2cta.py @@ -0,0 +1,1196 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# 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. +"""KDA multi-token verify sized for two CTAs per SM (large batches). + +Computes exactly what ``kda_mtp_decode_v2.kda_decode_mtp_v2_kernel`` computes on its runtime +contract (packed rows, fused gated RMSNorm with a bf16 output, 3xTF32 chunk recurrence, fp32 or +bf16 state, bf16 projections and conv caches, fp32 replay gates, one CTA per (head, request) +tile, forget gate from the projection) with the same arithmetic in the same order, so outputs, +committed state, conv windows and replay caches are bitwise identical. A tile still takes one +512-thread CTA; the footprint shrinks so that two tiles share an SM when the grid has more tiles +than SMs: + +* **Tensor memory: 256 columns.** S0 enters GEMM1 as two sequential A operands in the same 128 + columns (S0 hi for GEMM1, then S0 lo for the K~ sweep, issued in v2's order), and the GEMM2 + accumulator (S0 * decay) is seeded into those columns once GEMM1 is done. D4 | D2 | D1 | D3 + = 208 columns. A bf16 state is exact in tf32, so (as in v2) it has no lo part and GEMM1 is + the hi pass alone. +* **Shared memory: <= 114 KB** (113.7 KB at seven drafts). One 64 KB region holds, in turn, the + S0 TMA boxes (S0 is moved into registers and its hi part into tensor memory right after the + precompute), the GEMM3 tiles with the cumulative gates and e^{G_c}, the GEMM1 / Kbar / Bcoef + tiles once GEMM3 is done, V' over the GEMM1 tiles once GEMM1 is done, and finally the + committed-state TMA boxes. The raw conv inputs stay bf16 (the committed window and replay + rows land through ``cp.async``; the new tokens are indexed by token, not position, so their + placement does not wait for the accepted count), q is staged for the new tokens only, the + gate warps keep their token's raw gate, beta, dt_bias and A_log in registers (gate of four + consecutive channels per lane), the epilogue keeps its output-norm operands in registers, and + the GEMM3 result reuses the conv-weight words. Kbar and V' go straight into their swizzled + tiles from the row-per-thread producers. +* **Registers: 64 per thread** (``min_blocks_per_mp=2``); D1 is read from tensor memory in the + 16-column pieces each consumer needs, and a bf16 S0 row stays in registers as packed bf16x2 + words. The tensor-memory allocation permit is released right after the allocation so that the + SM's other CTA can allocate. +* **Named barriers: ids <= 4, all compile-time constants** (see ``BAR_WG0``), which SM107 needs + to co-schedule two CTAs. + +Alone on an SM a tile is not faster than in the v2 kernel (the precompute reads bf16 and GEMM1's +two A operands are serialized), while two tiles per SM finish sooner than two waves of v2 tiles, +so the kernel pays off for grids with more tiles than SMs. +""" + +import cutlass +import cutlass.cute as cute +from cutlass._mlir.dialects import llvm +from cutlass.cute.nvgpu import cpasync, tcgen05 +from cutlass.cute.typing import Int64 +from cutlass.cutlass_dsl import T, dsl_user_op + +from .kda_mtp_decode import ( + NUM_THREADS, + T_PAD, + TILE_K, + cvt_bf16x2, + round_tf32, + swz64_k16, + tcgen05_fence_after_thread_sync, + tcgen05_fence_before_thread_sync, +) +from .kda_mtp_decode_v2 import ( + asm_words, + gaddr, + ld_row4, + saddr, + st_asm, + st_global_b32x2, + st_row4, + st_shared_f32x4, + unpack4, +) +from .kda_mtp_decode_v2 import ld_shared_f32x4 as _ld_shared_f32x4 + +TMEM_COLS = 256 +# The 64 KB operand region (floats), used in turn as +# P0: the S0 TMA boxes (128 B swizzled rows: four 128 x 32 fp32 or two 128 x 64 bf16 boxes), +# until S0 is in registers +# P1: A3 = [Kf; Qf hi; Qf lo; -] 64 x 128, B3 hi / lo 16 x 128 (GEMM3), the cumulative gates +# and e^{G_c} (both kept through P2 / P3) +# P2 (once GEMM3 is done): B1 = [K~hi; Q~; K~lo] 48 x 128, B1k = K~hi 16 x 128, then +# [Bcoef hi; Kbar hi] and [Bcoef lo; Kbar lo] (16 + 128) x 16 each (GEMM2 B tiles) +# P3 (once GEMM1 is done): A2 hi / lo = V'^T 128 x 16 over B1 +# P4 (once GEMM2 is done): the committed-state TMA boxes +OFF_A3 = 0 +OFF_B3HI = 8192 +OFF_B3LO = 10240 +OFF_B1 = 0 +OFF_B1K = 6144 +OFF_B4HI = 8192 +OFF_B2HI = 8448 +OFF_B4LO = 10496 +OFF_B2LO = 10752 +OFF_E = 12800 +OFF_EC = 14720 +OFF_A2HI = 0 +OFF_A2LO = 2048 +REGION_FLOATS = 16384 +# Named barrier ids (0 is the CTA-wide barrier), all compile-time constants. An SM reserves +# (max id + 1) named barriers per CTA and SM107 co-schedules two CTAs only up to id 4, so the ids +# are reused across phases that CTA-wide barriers separate: BAR_WG0 + wg (1..4) are the +# per-warpgroup GEMM3 barriers, the store / TMEM-release barriers of the epilogue take ids 1..3 +# again. +BAR_WG0 = 1 +BAR_STORE_LO = 1 # committed-state boxes staged, threads 0..255 +BAR_STORE_HI = 2 # committed-state boxes staged, threads 256..511 +BAR_TMEM_FREE = 3 # last TMEM reads done -> warp 0 frees the columns + + +@dsl_user_op +def _cp_async_8(smem_addr, gmem_addr, *, loc=None, ip=None): + st_asm( + "cp.async.ca.shared.global [$0], [$1], 8;", + "r,l", + [ + cutlass.Int32(smem_addr).ir_value(loc=loc, ip=ip), + Int64(gmem_addr).ir_value(loc=loc, ip=ip), + ], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _cp_async_16(smem_addr, gmem_addr, *, loc=None, ip=None): + st_asm( + "cp.async.cg.shared.global [$0], [$1], 16;", + "r,l", + [ + cutlass.Int32(smem_addr).ir_value(loc=loc, ip=ip), + Int64(gmem_addr).ir_value(loc=loc, ip=ip), + ], + loc=loc, + ip=ip, + ) + + +@dsl_user_op +def _cp_async_commit(*, loc=None, ip=None): + st_asm("cp.async.commit_group;", "", [], loc=loc, ip=ip) + + +@dsl_user_op +def _cp_async_wait_all(*, loc=None, ip=None): + st_asm("cp.async.wait_group 0;", "", [], loc=loc, ip=ip) + + +@dsl_user_op +def _tmem_ld_x1(taddr, *, loc=None, ip=None) -> cutlass.Float32: + """One 32-bit tensor-memory word per lane (``tcgen05.ld 32x32b.x1``) at ``taddr``.""" + return cutlass.Float32( + llvm.inline_asm( + T.f32(), + [cutlass.Int32(taddr).ir_value(loc=loc, ip=ip)], + "tcgen05.ld.sync.aligned.32x32b.x1.b32 {$0}, [$1];", + "=f,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +def _ld_shared_b32x2(addr): + return asm_words( + 2, + "ld.shared.v2.b32 {$0, $1}, [$2];", + "=r,=r,r", + [cutlass.Int32(addr).ir_value()], + loc=None, + ip=None, + ) + + +def _st_shared_b32x2(addr, a, b): + st_asm( + "st.shared.v2.b32 [$0], {$1, $2};", + "r,r,r", + [cutlass.Int32(addr).ir_value(), cutlass.Int32(a).ir_value(), cutlass.Int32(b).ir_value()], + ) + + +def _s0_value(r_state, r_state_w, idx, packed): + """fp32 value of S0 column ``idx`` of this thread's row: ``r_state[idx]``, or for a packed bf16 + state the bf16 value ``idx`` of the bf16x2 words ``r_state_w`` (even values in the low half).""" + if cutlass.const_expr(packed): + w = r_state_w[idx // 2] + if cutlass.const_expr(idx % 2 == 0): + w = w << 16 + else: + w = w & cutlass.Int32(-65536) + return w.bitcast(cutlass.Float32) + return r_state[idx] + + +def _load_new_token_qkv( + x_q, x_k, x_v, tok_row, hk_off, hv_off, k4, r_new, use_flat_layout, i_h, i_hv +): + """v2's new-token q / k / v loads (raw bf16 words, four channels per lane).""" + if cutlass.const_expr(use_flat_layout): + ld_row4(x_q, (0, tok_row, hk_off + k4), r_new, 0) + ld_row4(x_k, (0, tok_row, hk_off + k4), r_new, 4) + ld_row4(x_v, (0, tok_row, hv_off + k4), r_new, 8) + else: + ld_row4(x_q, (0, tok_row, i_h, k4), r_new, 0) + ld_row4(x_k, (0, tok_row, i_h, k4), r_new, 4) + ld_row4(x_v, (0, tok_row, i_hv, k4), r_new, 8) + + +@cute.kernel +def kda_decode_mtp_2cta_kernel( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + o: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + ssm_state_indices: cute.Tensor, + num_accepted_tokens: cute.Tensor, + tiled_mma1: cute.TiledMma, + tiled_mma3: cute.TiledMma, + tiled_mma1k: cute.TiledMma, + smem_layout_b1: cute.ComposedLayout, + smem_layout_a2: cute.ComposedLayout, + smem_layout_a3: cute.ComposedLayout, + smem_layout_b3: cute.ComposedLayout, + smem_layout_b1k: cute.ComposedLayout, + tiled_mma4: cute.TiledMma, + tiled_mma2h: cute.TiledMma, + smem_layout_b2h: cute.ComposedLayout, + tiled_mma2q: cute.TiledMma, + smem_layout_b2q: cute.ComposedLayout, + tma_atom_l: cute.CopyAtom, + mS_tma: cute.Tensor, + box_layout: cute.ComposedLayout, + tma_atom_t: cute.CopyAtom, + mT_tma: cute.Tensor, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], +): + """See the module docstring. Grid ``(HV, N, 1)``, one (head, request) tile per CTA.""" + tidx, _, _ = cute.arch.thread_idx() + in_warp_tid = tidx % 32 + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + i_hv, i_n, _ = cute.arch.block_idx() + i_h = i_hv + slot32 = ssm_state_indices[i_n] + hk_off = i_h * K + hv_off = i_hv * V + N1 = 3 * T_PAD + K_CACHE_DTYPE = k_cache.element_type + V_CACHE_DTYPE = v_cache.element_type + BETA_CACHE_DTYPE = beta_cache.element_type + t_max = 2 * NUM_SPEC + 1 + NEW_TOK = 1 + NUM_SPEC + NUM_OLD = KERNEL_WIDTH - 1 + NUM_SPEC # committed window + replay rows of the raw staging + vec_size = TILE_K // 32 + REPLAY_WARP0 = NUM_THREADS // 64 + TMA_WARP = NUM_THREADS // 32 - 1 + assert NUM_SPEC <= 7 and t_max <= T_PAD + assert NEW_TOK <= REPLAY_WARP0 and REPLAY_WARP0 + NUM_SPEC <= TMA_WARP + assert K == 128 and V == 128 and KERNEL_WIDTH == 4 and NUM_THREADS == 512 + STATE_DTYPE = h0.element_type + STATE_BF16 = STATE_DTYPE == cutlass.BFloat16 + assert STATE_BF16 or STATE_DTYPE == cutlass.Float32, "fp32 or bf16 recurrent state" + # a bf16 state is exact in tf32: no S0 lo part (v2's S0_LO) + S0_LO = not STATE_BF16 + BOX_COLS = 128 // (STATE_DTYPE.width // 8) # state columns per 128 B swizzled box row + NUM_SLABS = K // BOX_COLS + assert g_cache.element_type == cutlass.Float32, "fp32 replay gates (cp.async into sG)" + assert cs_q.element_type == cutlass.BFloat16 and x_q.element_type == cutlass.BFloat16 + k_weight_base = KERNEL_WIDTH * K + v_weight_base = 2 * KERNEL_WIDTH * K + conv_weight_elems = 2 * KERNEL_WIDTH * K + KERNEL_WIDTH * V + + smem = cutlass.utils.SmemAllocator() + sQn = smem.allocate_tensor(cutlass.Float32, cute.make_layout((NEW_TOK * K,)), 16) + sK = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max, K), stride=(K, 1)), 16) + sG = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max, K), stride=(K, 1)), 16) + sVall = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max * V,)), 16) + # Conv weights during the precompute; the GEMM3 result [L; B hi; B lo; -] (64 x 16) after. + sConvW = smem.allocate_tensor(cutlass.Float32, cute.make_layout((conv_weight_elems,)), 16) + sLB = cute.make_tensor(sConvW.iterator, cute.make_layout((4 * T_PAD, T_PAD), stride=(T_PAD, 1))) + # Raw conv inputs (bf16): [window; replayed drafts] and the new tokens, per q / k / v. + # rows [0, NUM_OLD): window + replay rows, then the new tokens' rows + sXq = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(((NUM_OLD + NEW_TOK) * K,)), 16) + sXk = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(((NUM_OLD + NEW_TOK) * K,)), 16) + sXv = smem.allocate_tensor(cutlass.BFloat16, cute.make_layout(((NUM_OLD + NEW_TOK) * V,)), 16) + sXrq, sXrk, sXrv = sXq, sXk, sXv + sXnq = cute.make_tensor(sXq.iterator + NUM_OLD * K, cute.make_layout((NEW_TOK * K,))) + sXnk = cute.make_tensor(sXk.iterator + NUM_OLD * K, cute.make_layout((NEW_TOK * K,))) + sXnv = cute.make_tensor(sXv.iterator + NUM_OLD * V, cute.make_layout((NEW_TOK * V,))) + sBeta = smem.allocate_tensor(cutlass.Float32, cute.make_layout((t_max,)), 16) + tmem_holder = smem.allocate_array(cutlass.Int32, 1) + mbar1a = smem.allocate_array(cutlass.Int64, 1) + mbar1 = smem.allocate_array(cutlass.Int64, 1) + mbar3 = smem.allocate_array(cutlass.Int64, 1) + mbar2 = smem.allocate_array(cutlass.Int64, 1) + mbar2b = smem.allocate_array(cutlass.Int64, 1) + mbar_ld = smem.allocate_array(cutlass.Int64, 1) + sRegion = smem.allocate_tensor(cutlass.Float32, cute.make_layout((REGION_FLOATS,)), 1024) + + def region_tile(off, layout): + return cute.make_tensor( + cute.recast_ptr(sRegion.iterator + off, layout.inner, dtype=cutlass.Float32), + layout.outer, + ) + + sA3 = region_tile(OFF_A3, smem_layout_a3) + sB3hi = region_tile(OFF_B3HI, smem_layout_b3) + sB3lo = region_tile(OFF_B3LO, smem_layout_b3) + sB1 = region_tile(OFF_B1, smem_layout_b1) + sB1k = region_tile(OFF_B1K, smem_layout_b1k) + sA2hi = region_tile(OFF_A2HI, smem_layout_a2) + sA2lo = region_tile(OFF_A2LO, smem_layout_a2) + # Flat views for closed-form swizzled element writes. + sA3f = cute.make_tensor(sRegion.iterator + OFF_A3, cute.make_layout((64 * 128,))) + sB3hif = cute.make_tensor(sRegion.iterator + OFF_B3HI, cute.make_layout((16 * 128,))) + sB3lof = cute.make_tensor(sRegion.iterator + OFF_B3LO, cute.make_layout((16 * 128,))) + sB1f = cute.make_tensor(sRegion.iterator + OFF_B1, cute.make_layout((48 * 128,))) + sB1kf = cute.make_tensor(sRegion.iterator + OFF_B1K, cute.make_layout((T_PAD * 128,))) + sB4hif = cute.make_tensor(sRegion.iterator + OFF_B4HI, cute.make_layout((T_PAD * T_PAD,))) + sB4lof = cute.make_tensor(sRegion.iterator + OFF_B4LO, cute.make_layout((T_PAD * T_PAD,))) + # Cumulative gates G_t (P1 / P2) and the S0 / committed-state boxes (P0 / P4). + sE = cute.make_tensor(sRegion.iterator + OFF_E, cute.make_layout((t_max * K,))) + sEc = cute.make_tensor(sRegion.iterator + OFF_EC, cute.make_layout((K,))) + sBoxf = cute.make_tensor(sRegion.iterator, cute.make_layout((V * K,))) + sBoxh = cute.make_tensor( + cute.recast_ptr(sRegion.iterator, dtype=cutlass.BFloat16), cute.make_layout((V * K,)) + ) + + # ---- load phase (v2's schedule): slot-independent loads under the slot load; then the + # slot-dependent ones. Window and replay raw conv rows and replay gates go to smem through + # cp.async (no registers); replay k / v / beta (converted on placement) and the new + # tokens through registers. + ch = tidx % 128 + quad = tidx // 128 + k4 = 4 * in_warp_tid + r_w = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Int32) + if quad == 0: + ld_row4(w_q, (hk_off + ch, 0), r_w, 0) + if quad == 1: + ld_row4(w_k, (hk_off + ch, 0), r_w, 0) + if quad == 2: + ld_row4(w_v, (hv_off + ch, 0), r_w, 0) + # gate warps (REPLAY_WARP0 + j): this lane's dt_bias channels [4 l, 4 l + 4) and A_log + j_gate = warp_idx - REPLAY_WARP0 + tok_row_g = i_n * NEW_TOK + j_gate + r_dtb = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Int32) + r_gate = cute.make_rmem_tensor(cute.make_layout((4,)), cutlass.Int32) + r_gb = cute.make_rmem_tensor(cute.make_layout((2,)), cutlass.Float32) + if (warp_idx >= REPLAY_WARP0) & (j_gate < NEW_TOK): + ld_row4(dt_bias, (hk_off + k4,), r_dtb, 0) + r_gb[1] = cutlass.Float32(A_log[i_h]) + r_new = cute.make_rmem_tensor(cute.make_layout((3 * 4,)), cutlass.Int32) + tok_row = i_n * NEW_TOK + warp_idx + if warp_idx < NEW_TOK: + _load_new_token_qkv( + x_q, x_k, x_v, tok_row, hk_off, hv_off, k4, r_new, USE_FLAT_LAYOUT, i_h, i_hv + ) + if (warp_idx >= REPLAY_WARP0) & (j_gate < NEW_TOK): + ld_row4(g, (0, tok_row_g, i_hv, k4), r_gate, 0) + r_gb[0] = cutlass.Float32(beta[0, tok_row_g, i_hv]) + # The TMA warp, the first to use an mbarrier, initializes them; every other warp first touches + # one after the placement barrier. + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_init(mbar1a, 1) + cute.arch.mbarrier_init(mbar1, 1) + cute.arch.mbarrier_init(mbar3, 1) + cute.arch.mbarrier_init(mbar2, 1) + cute.arch.mbarrier_init(mbar2b, 1) + cute.arch.mbarrier_init(mbar_ld, 1) + cute.arch.mbarrier_init_fence() + cute.arch.fence_view_async_shared() + slot = Int64(slot32) + h0_idx = slot * HV + i_hv + # byte step between consecutive window positions (mode 2) of the conv states + cs_step = [cute.crd2idx((0, 0, 1), t.layout) * 2 for t in (cs_q, cs_k, cs_v)] + rep = warp_idx - REPLAY_WARP0 + r_rep = cute.make_rmem_tensor(cute.make_layout((2 * 4,)), cutlass.Int32) + r_repb = cute.make_rmem_tensor(cute.make_layout((1,)), cutlass.Float32) + if (warp_idx >= REPLAY_WARP0) & (rep < NUM_SPEC): + ld_row4(k_cache, (slot, rep, hk_off + k4), r_rep, 0) + ld_row4(v_cache, (slot, rep, hv_off + k4), r_rep, 4) + r_repb[0] = cutlass.Float32(beta_cache[slot, rep, i_hv]) + # every pending entry lands in its row; rows from the accepted count on are overwritten + # by the new tokens' precompute before they are read + _cp_async_16(saddr(sG, rep * K + k4), gaddr(g_cache, (slot, rep, hk_off + k4))) + rr = KERNEL_WIDTH - 1 + rep + _cp_async_8(saddr(sXrq, rr * K + k4), gaddr(cs_q, (slot, hk_off + k4, rr))) + _cp_async_8(saddr(sXrk, rr * K + k4), gaddr(cs_k, (slot, hk_off + k4, rr))) + _cp_async_8(saddr(sXrv, rr * V + k4), gaddr(cs_v, (slot, hv_off + k4, rr))) + _cp_async_commit() + if warp_idx == TMA_WARP: + # committed window rows 0 .. KERNEL_WIDTH - 2: one base per conv state + wq_base = gaddr(cs_q, (slot, hk_off + k4, 0)) + wk_base = gaddr(cs_k, (slot, hk_off + k4, 0)) + wv_base = gaddr(cs_v, (slot, hv_off + k4, 0)) + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + _cp_async_8(saddr(sXrq, w * K + k4), wq_base + w * cs_step[0]) + _cp_async_8(saddr(sXrk, w * K + k4), wk_base + w * cs_step[1]) + _cp_async_8(saddr(sXrv, w * V + k4), wv_base + w * cs_step[2]) + _cp_async_commit() + commit_len = num_accepted_tokens[slot] + # S0 into the region's boxes (consumed right after the precompute) + if warp_idx == TMA_WARP: + with cute.arch.elect_one(): + cute.arch.mbarrier_arrive_and_expect_tx(mbar_ld, V * K * (STATE_DTYPE.width // 8)) + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sTl_h = cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + sl_h * (V * 32), box_layout.inner, dtype=STATE_DTYPE + ), + box_layout.outer, + ) + gS_h_all = cute.local_tile(mS_tma, (V, BOX_COLS, 1), (0, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gS_h = gS_h_all[(None, None, 0, h0_idx)] + else: + gS_h = gS_h_all[(None, None, 0, (i_hv, slot))] + tSl_h, tGl_h = cpasync.tma_partition( + tma_atom_l, + 0, + cute.make_layout(1), + cute.group_modes(sTl_h, 0, 2), + cute.group_modes(gS_h, 0, 2), + ) + if warp_idx == TMA_WARP: + cute.copy(tma_atom_l, tGl_h, tSl_h, tma_bar_ptr=mbar_ld) + if warp_idx == 0: + cute.arch.alloc_tmem(TMEM_COLS, tmem_holder) + # the SM's other CTA allocates its 256 columns as soon as we hold ours + cute.arch.relinquish_tmem_alloc_permit() + # ---- placement + if quad < 3: + wv = unpack4(w_q.element_type, r_w, 0) + for w in cutlass.range_constexpr(KERNEL_WIDTH): + sConvW[quad * (KERNEL_WIDTH * K) + w * K + ch] = wv[w] + if commit_len > NUM_SPEC: + commit_len = cutlass.Int32(NUM_SPEC) + T_loop = commit_len + NEW_TOK + if warp_idx < NEW_TOK: + # raw rows of new token j stay bf16 and are indexed by j + _st_shared_b32x2(saddr(sXnq, warp_idx * K + k4), r_new[0], r_new[1]) + _st_shared_b32x2(saddr(sXnk, warp_idx * K + k4), r_new[4], r_new[5]) + _st_shared_b32x2(saddr(sXnv, warp_idx * V + k4), r_new[8], r_new[9]) + if (warp_idx >= REPLAY_WARP0) & (rep < commit_len): + st_shared_f32x4(saddr(sK, rep * K + k4), *unpack4(K_CACHE_DTYPE, r_rep, 0)) + st_shared_f32x4(saddr(sVall, rep * V + k4), *unpack4(V_CACHE_DTYPE, r_rep, 4)) + if in_warp_tid == 0: + sBeta[rep] = r_repb[0] + if ((warp_idx >= REPLAY_WARP0) & (rep < NUM_SPEC)) | (warp_idx == TMA_WARP): + _cp_async_wait_all() + if tidx < K: + for _z in cutlass.range(T_loop, t_max, 1): + sG[_z, tidx] = cutlass.Float32(0.0) + sVall[_z * V + tidx] = cutlass.Float32(0.0) + if tidx == 0: + sBeta[_z] = cutlass.Float32(0.0) + wg = warp_idx // 4 + wg_tidx = tidx % 128 + r_q = cute.make_rmem_tensor(cute.make_layout((vec_size,)), cutlass.Float32) + r_k = cute.make_rmem_tensor(cute.make_layout((vec_size,)), cutlass.Float32) + cute.arch.barrier() + + # ---- precompute (v2's arithmetic; raw rows read as bf16 from the split staging) + half = NUM_THREADS // 64 + role = warp_idx // half + j = warp_idx % half + if j < NEW_TOK: + t = commit_len + j + # raw row of tap w of token j: j + w - 3 >= 0 -> new token row, else window / replay row + x_row = [] + for w in cutlass.range_constexpr(KERNEL_WIDTH): + r_new_idx = j + w - (KERNEL_WIDTH - 1) + old_row = commit_len + j + w + is_new = cutlass.Int32(r_new_idx >= 0) + x_row.append(old_row + is_new * (NUM_OLD + r_new_idx - old_row)) + if role == 0: + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + r_conv_0 = cutlass.Float32(0.0) + for w in cutlass.range_constexpr(KERNEL_WIDTH): + r_conv_0 += cutlass.Float32(sXq[x_row[w] * K + k_idx]) * sConvW[w * K + k_idx] + r_conv = ( + cutlass.Float32(sXk[x_row[0] * K + k_idx]) + * sConvW[k_weight_base + 0 * K + k_idx] + ) + for w in cutlass.range_constexpr(1, KERNEL_WIDTH): + r_conv += ( + cutlass.Float32(sXk[x_row[w] * K + k_idx]) + * sConvW[k_weight_base + w * K + k_idx] + ) + e0 = cute.math.exp(-r_conv_0, fastmath=True) + r_q[i] = r_conv_0 * cute.arch.rcp_approx(cutlass.Float32(1.0) + e0) + r_conv = r_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_conv, fastmath=True) + ) + r_k[i] = r_conv + sum_q = 0.0 + sum_k = 0.0 + for i in cutlass.range(vec_size, unroll_full=True): + sum_q += r_q[i] * r_q[i] + sum_k += r_k[i] * r_k[i] + for offset in [16, 8, 4, 2, 1]: + sum_q += cute.arch.shuffle_sync_bfly( + sum_q, offset=offset, mask=-1, mask_and_clamp=31 + ) + sum_k += cute.arch.shuffle_sync_bfly( + sum_k, offset=offset, mask=-1, mask_and_clamp=31 + ) + rnorm_q_scaled = cute.math.rsqrt(sum_q + 1e-06, fastmath=True) * scale + rnorm_k = cute.math.rsqrt(sum_k + 1e-06, fastmath=True) + for i in cutlass.range(vec_size, unroll_full=True): + k_idx = i * 32 + in_warp_tid + sQn[j * K + k_idx] = r_q[i] * rnorm_q_scaled + sK[t, k_idx] = r_k[i] * rnorm_k + else: + # gate of channels [4 l, 4 l + 4) from this warp's registers (v2's per-channel math) + r_exp_A = cute.math.exp(r_gb[1], fastmath=True) + g_raw = unpack4(g.element_type, r_gate, 0) + dtb = unpack4(dt_bias.element_type, r_dtb, 0) + gates = [] + for i in cutlass.range_constexpr(4): + r_g_raw = g_raw[i] + r_g_raw = r_g_raw + dtb[i] + exp_A_x = r_exp_A * r_g_raw + sigmoid_val = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-exp_A_x, fastmath=True) + ) + gates.append(lower_bound * sigmoid_val) + st_shared_f32x4(saddr(sG, t * K + k4), *gates) + r_beta_t = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-r_gb[0], fastmath=True) + ) + if in_warp_tid == 0: + sBeta[t] = r_beta_t + for i in cutlass.range(vec_size, unroll_full=True): + v_idx = i * 32 + in_warp_tid + _v_conv = 0.0 + for w in cutlass.range_constexpr(KERNEL_WIDTH): + _v_conv += ( + cutlass.Float32(sXv[x_row[w] * V + v_idx]) + * sConvW[v_weight_base + w * V + v_idx] + ) + _v_conv = _v_conv * cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-_v_conv, fastmath=True) + ) + sVall[t * V + v_idx] = _v_conv + + # ================================================================= + # Chunk (WY-form) recurrence on tensor cores. + # ================================================================= + tmem_ptr = cute.arch.retrieve_tmem_ptr(cutlass.Float32, 16, tmem_holder) + pool = cutlass.utils.TmemBufferPool(tmem_ptr, TMEM_COLS) + tD4_frag = tiled_mma4.make_fragment_C(tiled_mma4.partition_shape_C((V, T_PAD))) + tD4 = pool.allocate_tensor(tD4_frag.layout, cutlass.Float32) + tD2h_frag = tiled_mma2h.make_fragment_C(tiled_mma2h.partition_shape_C((V, K // 2))) + tD2h0 = pool.allocate_tensor(tD2h_frag.layout, cutlass.Float32) + tD2h1 = pool.allocate_tensor(tD2h_frag.layout, cutlass.Float32) + tD2q_frag = tiled_mma2q.make_fragment_C(tiled_mma2q.partition_shape_C((V, T_PAD + K // 2))) + tD2q = cute.make_tensor(tD4.iterator, tD2q_frag.layout) + tA_frag = tiled_mma1.make_fragment_A(tiled_mma1.partition_shape_A((V, K))) + # S0 hi, then S0 lo, then the GEMM2 accumulator: the same 128 columns + tX = cute.make_tensor(tD2h0.iterator, tA_frag.layout) + tD1_frag = tiled_mma1.make_fragment_C(tiled_mma1.partition_shape_C((V, N1))) + tD1 = pool.allocate_tensor(tD1_frag.layout, cutlass.Float32) + tD3_frag = tiled_mma3.make_fragment_C(tiled_mma3.partition_shape_C((4 * T_PAD, T_PAD))) + tD3 = pool.allocate_tensor(tD3_frag.layout, cutlass.Float32) + + st8 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x8), cutlass.Float32) + ld8 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x8), cutlass.Float32) + ld16 = cute.make_copy_atom(tcgen05.Ld32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + thr_st8 = tcgen05.make_tmem_copy(st8, tX).get_slice(wg_tidx) + tX_st8 = thr_st8.partition_D(tX) + r8 = cute.make_rmem_tensor(thr_st8.partition_S(tX).shape[:-1], cutlass.Float32) + thr_ld8 = tcgen05.make_tmem_copy(ld8, tX).get_slice(wg_tidx) + tX_ld8 = thr_ld8.partition_S(tX) + tS_list = [] + tG_list = [] + for sl_h in cutlass.range_constexpr(NUM_SLABS): + sT_h = cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + sl_h * (V * 32), box_layout.inner, dtype=STATE_DTYPE + ), + box_layout.outer, + ) + gT_h_all = cute.local_tile(mT_tma, (V, BOX_COLS, 1), (0, sl_h, None)) + if cutlass.const_expr(USE_FLAT_LAYOUT): + gT_h = gT_h_all[(None, None, 0, h0_idx)] + else: + gT_h = gT_h_all[(None, None, 0, (i_hv, slot))] + tS_h, tG_h = cpasync.tma_partition( + tma_atom_t, + 0, + cute.make_layout(1), + cute.group_modes(sT_h, 0, 2), + cute.group_modes(gT_h, 0, 2), + ) + tS_list.append(tS_h) + tG_list.append(tG_h) + + # ---- S0 row v = wg_tidx, columns [32 wg, +32): from its TMA box (v2's conflict-free 16 B + # chunks) into registers, its hi part into X; the boxes are free after the barrier, which + # also publishes the precompute + v_row = wg_tidx + vx = 4 * (wg_tidx % 8) + r_state = cute.make_rmem_tensor(cute.make_layout((32,)), cutlass.Float32) + # a bf16 state stays packed: bf16x2 words, value 2 i in the low half of word i + r_state_w = cute.make_rmem_tensor(cute.make_layout((16,)), cutlass.Int32) + cute.arch.mbarrier_wait(mbar_ld, 0) + if cutlass.const_expr(STATE_BF16): + # bf16 box: 64 columns per 128 B row, so this warpgroup's 32 columns are the 16 B + # chunks [4 (wg % 2), +4) of box wg // 2 (v2), kept as raw bf16x2 words + for c in cutlass.range_constexpr(4): + chunk = (4 * (wg % 2) + c) ^ (wg_tidx % 8) + off_l = (wg // 2) * (V * 64) + 64 * wg_tidx + 8 * chunk + words = _ld_shared_f32x4(saddr(sBoxh, off_l)) + for i in cutlass.range_constexpr(4): + r_state_w[4 * c + i] = words[i].bitcast(cutlass.Int32) + else: + for c in cutlass.range(8, unroll_full=True): + off_l = wg * (V * 32) + 32 * wg_tidx + ((4 * c) ^ vx) + for i in cutlass.range(4, unroll_full=True): + r_state[4 * c + i] = sBoxf[off_l + i] + for jj in cutlass.range_constexpr(4): + if wg == jj: + for c in cutlass.range_constexpr(4): + for i in cutlass.range_constexpr(8): + if cutlass.const_expr(S0_LO): + r8[i] = round_tf32(r_state[8 * c + i]) + else: + r8[i] = _s0_value(r_state, r_state_w, 8 * c + i, STATE_BF16) + cute.copy(st8, r8, tX_st8[(None, None, None, 4 * jj + c)]) + cute.arch.fence_view_async_tmem_store() + cute.arch.barrier() + + # ---- cumulative gates (v2), q of the replayed drafts is zero + kq = tidx % K + quarter = tidx // K + t_mid = T_loop // 2 + r_G = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + g_run = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t < t_max: + g_run = g_run + sG[t, kq] + r_G[t] = g_run + g_mid = cutlass.Float32(0.0) + for t in cutlass.range(T_PAD, unroll_full=True): + if t == t_mid: + g_mid = r_G[t] + exp_clamp = cutlass.Float32(87.0) + ka = kq // 32 + kc = kq % 32 + b1_base = ka * (3 * T_PAD * 32) + a3_base = ka * (4 * T_PAD * 32) + b3_base = ka * (T_PAD * 32) + b1k_base = ka * (T_PAD * 32) + r_b1 = cute.make_rmem_tensor(cute.make_layout((3 * (T_PAD // 4),)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + kt_hi = cutlass.Float32(0.0) + kt_lo = cutlass.Float32(0.0) + qt_r = cutlass.Float32(0.0) + kf = cutlass.Float32(0.0) + qf_hi = cutlass.Float32(0.0) + qf_lo = cutlass.Float32(0.0) + kb_hi = cutlass.Float32(0.0) + kb_lo = cutlass.Float32(0.0) + if t < T_loop: + d_t = r_G[t] - g_mid + e_t = cute.math.exp(-cute.arch.fmax(-d_t, -exp_clamp), fastmath=True) + e_bt = cute.math.exp(-cute.arch.fmax(d_t, -exp_clamp), fastmath=True) + eg_t = cute.math.exp(r_G[t], fastmath=True) + sE[t * K + kq] = r_G[t] + if t == commit_len: + sEc[kq] = eg_t + k_t = sK[t, kq] + q_t = cutlass.Float32(0.0) + if t >= commit_len: + q_t = sQn[(t - commit_len) * K + kq] + kt = k_t * eg_t + kt_hi = round_tf32(kt) + qt_r = round_tf32(q_t * eg_t) + kf = round_tf32(k_t * e_t) + qf = q_t * e_t + qf_hi = round_tf32(qf) + kb = k_t * e_bt + kb_hi = round_tf32(kb) + kt_lo = kt - kt_hi + qf_lo = qf - qf_hi + kb_lo = kb - kb_hi + xc = kc ^ (4 * (t % 8)) + sA3f[a3_base + 32 * t + xc] = kf + sA3f[a3_base + 32 * (T_PAD + t) + xc] = qf_hi + sA3f[a3_base + 32 * (2 * T_PAD + t) + xc] = qf_lo + sB3hif[b3_base + 32 * t + xc] = kb_hi + sB3lof[b3_base + 32 * t + xc] = kb_lo + r_b1[3 * (t // 4)] = kt_hi + r_b1[3 * (t // 4) + 1] = qt_r + r_b1[3 * (t // 4) + 2] = kt_lo + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + tCrA3 = tiled_mma3.make_fragment_A( + cute.make_tensor( + cute.recast_ptr(sA3.iterator, smem_layout_a3.inner, dtype=cutlass.TFloat32), + sA3.layout, + ) + ) + tCrB3hi = tiled_mma3.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB3hi.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3hi.layout, + ) + ) + tCrB3lo = tiled_mma3.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB3lo.iterator, smem_layout_b3.inner, dtype=cutlass.TFloat32), + sB3lo.layout, + ) + ) + if warp_idx == 0: + tiled_mma3.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, tD3, tCrA3[(None, None, kb_i, 0)], tCrB3hi[(None, None, kb_i, 0)], tD3 + ) + tiled_mma3.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma3, tD3, tCrA3[(None, None, kb_i, 0)], tCrB3lo[(None, None, kb_i, 0)], tD3 + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar3) + + # GEMM1 operand rows over the GEMM3 tiles once GEMM3 has read them; D3 -> smem. One warp per + # warpgroup polls the mbarrier, the others wait at the warpgroup's barrier. + if warp_idx % 4 == 1: + cute.arch.mbarrier_wait(mbar3, 0) + # one compile-time barrier id per warpgroup: a barrier id held in a register makes the CTA + # reserve all 16 named barriers, which leaves room for one CTA per SM on SM107 + for jj in cutlass.range_constexpr(4): + if wg == jj: + cute.arch.barrier(barrier_id=BAR_WG0 + jj, number_of_threads=128) + tcgen05_fence_after_thread_sync() + for t in cutlass.range(T_PAD, unroll_full=True): + if quarter == t % 4: + xc = kc ^ (4 * (t % 8)) + sB1f[b1_base + 32 * t + xc] = r_b1[3 * (t // 4)] + if cutlass.const_expr(S0_LO): + sB1kf[b1k_base + 32 * t + xc] = r_b1[3 * (t // 4)] + sB1f[b1_base + 32 * (T_PAD + t) + xc] = r_b1[3 * (t // 4) + 1] + sB1f[b1_base + 32 * (2 * T_PAD + t) + xc] = r_b1[3 * (t // 4) + 2] + ld3 = cute.make_copy_atom( + tcgen05.Ld16x128bOp(tcgen05.Repetition(T_PAD * 32 // 128)), cutlass.Float32 + ) + thr_ld3 = tcgen05.make_tmem_copy(ld3, tD3).get_slice(wg_tidx) + tD3_ld = thr_ld3.partition_S(tD3) + sLB3 = cute.make_tensor( + sLB.iterator, + cute.make_layout(((4 * T_PAD, T_PAD), 1, 1), stride=((T_PAD, 1), 0, 0)), + ) + tD3_sm = thr_ld3.partition_D(sLB3) + rD3 = cute.make_rmem_tensor(tD3_sm.shape, cutlass.Float32) + if wg == 0: + cute.copy(ld3, tD3_ld, rD3) + cute.arch.fence_view_async_tmem_load() + cute.autovec_copy(rD3, tD3_sm) + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + tCrB1 = tiled_mma1.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB1.iterator, smem_layout_b1.inner, dtype=cutlass.TFloat32), + sB1.layout, + ) + ) + tCrB1k = tiled_mma1k.make_fragment_B( + cute.make_tensor( + cute.recast_ptr(sB1k.iterator, smem_layout_b1k.inner, dtype=cutlass.TFloat32), + sB1k.layout, + ) + ) + tD1k = cute.make_tensor( + tD1.iterator, + tiled_mma1k.make_fragment_C(tiled_mma1k.partition_shape_C((V, T_PAD))).layout, + ) + tX_tf = cute.make_tensor(cute.recast_ptr(tX.iterator, dtype=cutlass.TFloat32), tX.layout) + if warp_idx == 0: + tiled_mma1.set(tcgen05.Field.ACCUMULATE, False) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm(tiled_mma1, tD1, tX_tf[(None, None, kb_i)], tCrB1[(None, None, kb_i, 0)], tD1) + tiled_mma1.set(tcgen05.Field.ACCUMULATE, True) + with cute.arch.elect_one(): + if cutlass.const_expr(S0_LO): + tcgen05.commit(mbar1a) + else: + tcgen05.commit(mbar1) + # ---- under GEMM1: D4 zero, Bcoef tiles (warpgroups 1..2, from sLB), Kbar tiles, the new + # tokens' replay-cache rows and conv-window commits + st16 = cute.make_copy_atom(tcgen05.St32x32bOp(tcgen05.Repetition.x16), cutlass.Float32) + thr_stD4 = tcgen05.make_tmem_copy(st16, tD4).get_slice(wg_tidx) + tD4_st = thr_stD4.partition_D(tD4) + r_zero = cute.make_rmem_tensor(thr_stD4.partition_S(tD4).shape, cutlass.Float32) + for i in cutlass.range(cute.size(r_zero), unroll_full=True): + r_zero[i] = cutlass.Float32(0.0) + if wg == 0: + cute.copy(st16, r_zero, tD4_st) + if (tidx >= 128) & (tidx < 128 + T_PAD * T_PAD): + t4 = (tidx - 128) // T_PAD + s4 = (tidx - 128) % T_PAD + bc = cutlass.Float32(0.0) + if (t4 < T_loop) & (s4 <= t4): + bc = sLB[T_PAD + t4, s4] + sLB[2 * T_PAD + t4, s4] + bc_hi = round_tf32(bc) + o_b4 = swz64_k16(t4, s4) + sB4hif[o_b4] = bc_hi + sB4lof[o_b4] = bc - bc_hi + # Kbar[k, s] = K[s, k] e^{G_c[k] - G_s[k]} (s <= c, else 0): thread (k = kq, quarter) owns + # columns s = 4 quarter + {0..3}, i.e. one 16 B chunk of row k of the SW64 hi / lo tiles + sB2hi_row = saddr(sRegion, OFF_B2HI) + sB2lo_row = saddr(sRegion, OFF_B2LO) + # (quarters whose first column is past the committed count hold zeros only) + kb_0 = cutlass.Float32(0.0) + kb_1 = cutlass.Float32(0.0) + kb_2 = cutlass.Float32(0.0) + kb_3 = cutlass.Float32(0.0) + s_q = 4 * quarter + if s_q <= commit_len: + e_c = sE[commit_len * K + kq] + kb_0 = sK[s_q, kq] * cute.math.exp(e_c - sE[s_q * K + kq], fastmath=True) + if s_q + 1 <= commit_len: + kb_1 = sK[s_q + 1, kq] * cute.math.exp(e_c - sE[(s_q + 1) * K + kq], fastmath=True) + if s_q + 2 <= commit_len: + kb_2 = sK[s_q + 2, kq] * cute.math.exp(e_c - sE[(s_q + 2) * K + kq], fastmath=True) + if s_q + 3 <= commit_len: + kb_3 = sK[s_q + 3, kq] * cute.math.exp(e_c - sE[(s_q + 3) * K + kq], fastmath=True) + kb_vals = [kb_0, kb_1, kb_2, kb_3] + kb_his = [round_tf32(v) for v in kb_vals] + o_b2 = swz64_k16(kq, 4 * quarter) + st_shared_f32x4(sB2hi_row + 4 * o_b2, *kb_his) + st_shared_f32x4(sB2lo_row + 4 * o_b2, *[kb_vals[i] - kb_his[i] for i in range(4)]) + if (warp_idx >= 1) & (warp_idx <= NUM_SPEC): + p = warp_idx - 1 + t_c = commit_len + 1 + p + st_row4(k_cache, (slot, p, hk_off + k4), _ld_shared_f32x4(saddr(sK, t_c * K + k4))) + st_row4(g_cache, (slot, p, hk_off + k4), _ld_shared_f32x4(saddr(sG, t_c * K + k4))) + st_row4(v_cache, (slot, p, hv_off + k4), _ld_shared_f32x4(saddr(sVall, t_c * V + k4))) + # draft p is new token 1 + p + wq = _ld_shared_b32x2(saddr(sXnq, (1 + p) * K + k4)) + st_global_b32x2(gaddr(cs_q, (slot, hk_off + k4, KERNEL_WIDTH - 1 + p)), wq[0], wq[1]) + wk = _ld_shared_b32x2(saddr(sXnk, (1 + p) * K + k4)) + st_global_b32x2(gaddr(cs_k, (slot, hk_off + k4, KERNEL_WIDTH - 1 + p)), wk[0], wk[1]) + wvv = _ld_shared_b32x2(saddr(sXnv, (1 + p) * V + k4)) + st_global_b32x2(gaddr(cs_v, (slot, hv_off + k4, KERNEL_WIDTH - 1 + p)), wvv[0], wvv[1]) + if in_warp_tid == 0: + beta_cache[slot, p, i_hv] = sBeta[t_c].to(BETA_CACHE_DTYPE) + # committed window: raw rows commit_len + 1 + w (w = 2 is new token 0) + for sec in cutlass.range_constexpr(3): + if warp_idx == REPLAY_WARP0 + sec: + cs_sec = (cs_q, cs_k, cs_v)[sec] + old_sec = (sXrq, sXrk, sXrv)[sec] + new_sec = (sXnq, sXnk, sXnv)[sec] + off_sec = hv_off if sec == 2 else hk_off + sec_base = gaddr(cs_sec, (slot, off_sec + k4, 0)) + for w in cutlass.range_constexpr(KERNEL_WIDTH - 1): + if cutlass.const_expr(w + 1 >= KERNEL_WIDTH - 1): + src = saddr(new_sec, (w + 1 - (KERNEL_WIDTH - 1)) * K + k4) + else: + src = saddr(old_sec, (commit_len + 1 + w) * K + k4) + ww = _ld_shared_b32x2(src) + st_global_b32x2(sec_base + w * cs_step[sec], ww[0], ww[1]) + # ---- GEMM1's S0 lo sweep once the hi pass has read X + if cutlass.const_expr(S0_LO): + cute.arch.mbarrier_wait(mbar1a, 0) + tcgen05_fence_after_thread_sync() + for jj in cutlass.range_constexpr(4): + if wg == jj: + for c in cutlass.range_constexpr(4): + for i in cutlass.range_constexpr(8): + s_hi = round_tf32(r_state[8 * c + i]) + r8[i] = r_state[8 * c + i] - s_hi + cute.copy(st8, r8, tX_st8[(None, None, None, 4 * jj + c)]) + cute.arch.fence_view_async_tmem_store() + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + if warp_idx == 0: + tiled_mma1k.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(K // 8): + cute.gemm( + tiled_mma1k, + tD1k, + tX_tf[(None, None, kb_i)], + tCrB1k[(None, None, kb_i, 0)], + tD1k, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar1) + # ---- GEMM2 accumulator seed S0 * e^{G_c} into X once GEMM1 is done + cute.arch.mbarrier_wait(mbar1, 0) + tcgen05_fence_after_thread_sync() + for jj in cutlass.range_constexpr(4): + if wg == jj: + for c in cutlass.range_constexpr(4): + for i in cutlass.range_constexpr(8): + r8[i] = ( + _s0_value(r_state, r_state_w, 8 * c + i, STATE_BF16) + * sEc[jj * 32 + 8 * c + i] + ) + cute.copy(st8, r8, tX_st8[(None, None, None, 4 * jj + c)]) + cute.arch.fence_view_async_tmem_store() + # ---- warpgroup 0: V' by forward substitution (row v per thread) straight into the A2 tiles + thr_ld1 = tcgen05.make_tmem_copy(ld16, tD1).get_slice(wg_tidx) + tD1_ld = thr_ld1.partition_S(tD1) + # 16-column chunks of D1 are mode 1 of the ld16 partition: [K~hi | Q~ | K~lo] + rD1_shape = thr_ld1.partition_D(tD1)[(None, 0, None, None)].shape + rD1p = cute.make_rmem_tensor(rD1_shape, cutlass.Float32) + rD1q = cute.make_rmem_tensor(rD1_shape, cutlass.Float32) + rD1p_flat = cute.make_tensor(rD1p.iterator, cute.make_layout((T_PAD,))) + rD1q_flat = cute.make_tensor(rD1q.iterator, cute.make_layout((T_PAD,))) + if wg == 0: + cute.copy(ld16, tD1_ld[(None, 0, None, None)], rD1p) + cute.copy(ld16, tD1_ld[(None, 2, None, None)], rD1q) + cute.arch.fence_view_async_tmem_load() + r_Y = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for s in cutlass.range(T_PAD, unroll_full=True): + y = cutlass.Float32(0.0) + if s < t_max: + u0 = rD1p_flat[s] + rD1q_flat[s] + y = sBeta[s] * (sVall[s * V + v_row] - u0) + r_Y[s] = y + r_V = cute.make_rmem_tensor(cute.make_layout((T_PAD,)), cutlass.Float32) + for t in cutlass.range(T_PAD, unroll_full=True): + vp = cutlass.Float32(0.0) + if t < t_max: + acc0 = cutlass.Float32(0.0) + acc1 = cutlass.Float32(0.0) + acc2 = cutlass.Float32(0.0) + acc3 = cutlass.Float32(0.0) + for s in range(t): + if s % 4 == 0: + acc0 = acc0 + sLB[t, s] * r_V[s] + elif s % 4 == 1: + acc1 = acc1 + sLB[t, s] * r_V[s] + elif s % 4 == 2: + acc2 = acc2 + sLB[t, s] * r_V[s] + else: + acc3 = acc3 + sLB[t, s] * r_V[s] + vp = r_Y[t] - sBeta[t] * ((acc0 + acc1) + (acc2 + acc3)) + r_V[t] = vp + sA2hi_row = saddr(sRegion, OFF_A2HI) + sA2lo_row = saddr(sRegion, OFF_A2LO) + for c in cutlass.range_constexpr(4): + his = [round_tf32(r_V[4 * c + i]) for i in range(4)] + o_a2 = swz64_k16(v_row, 4 * c) + st_shared_f32x4(sA2hi_row + 4 * o_a2, *his) + st_shared_f32x4(sA2lo_row + 4 * o_a2, *[r_V[4 * c + i] - his[i] for i in range(4)]) + cute.arch.fence_view_async_shared() + tcgen05_fence_before_thread_sync() + cute.arch.barrier() + tcgen05_fence_after_thread_sync() + sA2hi_tf = cute.make_tensor( + cute.recast_ptr(sA2hi.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), sA2hi.layout + ) + sA2lo_tf = cute.make_tensor( + cute.recast_ptr(sA2lo.iterator, smem_layout_a2.inner, dtype=cutlass.TFloat32), sA2lo.layout + ) + tCrA2hi = tiled_mma2h.make_fragment_A(sA2hi_tf) + tCrA2lo = tiled_mma2h.make_fragment_A(sA2lo_tf) + tCrB2hi_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + OFF_B2HI + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.TFloat32, + ), + smem_layout_b2h.outer, + ) + ) + tCrB2lo_h1 = tiled_mma2h.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + OFF_B2LO + 64 * T_PAD, + smem_layout_b2h.inner, + dtype=cutlass.TFloat32, + ), + smem_layout_b2h.outer, + ) + ) + tCrA2hi_q = tiled_mma2q.make_fragment_A(sA2hi_tf) + tCrA2lo_q = tiled_mma2q.make_fragment_A(sA2lo_tf) + tCrB2q_hi = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + OFF_B4HI, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + smem_layout_b2q.outer, + ) + ) + tCrB2q_lo = tiled_mma2q.make_fragment_B( + cute.make_tensor( + cute.recast_ptr( + sRegion.iterator + OFF_B4LO, smem_layout_b2q.inner, dtype=cutlass.TFloat32 + ), + smem_layout_b2q.outer, + ) + ) + if warp_idx == 0: + tiled_mma2q.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2hi_q[(None, None, kb_i, 0)], + tCrB2q_lo[(None, None, kb_i, 0)], + tD2q, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2q, + tD2q, + tCrA2lo_q[(None, None, kb_i, 0)], + tCrB2q_hi[(None, None, kb_i, 0)], + tD2q, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2) + tiled_mma2h.set(tcgen05.Field.ACCUMULATE, True) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2hi[(None, None, kb_i, 0)], + tCrB2lo_h1[(None, None, kb_i, 0)], + tD2h1, + ) + for kb_i in cutlass.range_constexpr(T_PAD // 8): + cute.gemm( + tiled_mma2h, + tD2h1, + tCrA2lo[(None, None, kb_i, 0)], + tCrB2hi_h1[(None, None, kb_i, 0)], + tD2h1, + ) + with cute.arch.elect_one(): + tcgen05.commit(mbar2b) + # output-gate operand of the epilogue (warps 2, 3 of each warpgroup: one token each) + token_idx = (warp_idx // 4) * 2 + (warp_idx % 4 - 2) + r_og = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) + r_ow = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) + if ((warp_idx % 4) // 2 == 1) & (token_idx < NEW_TOK): + for i in cutlass.range(V // 32, unroll_full=True): + r_og[i] = cutlass.Float32( + onorm_g[0, i_n * NEW_TOK + token_idx, i_hv, i * 32 + in_warp_tid] + ) + r_ow[i] = cutlass.Float32(onorm_weight[i * 32 + in_warp_tid]) + # ---- outputs of the new tokens: o_t = Q-part of D1 + D4 (v2), bf16-rounded into sVall rows; + # warpgroup wg reads column t of both for its tokens wg and wg + 4 (lane = row v) + cute.arch.mbarrier_wait(mbar2, 0) + tcgen05_fence_after_thread_sync() + lane_hi = (warp_idx % 4) << 21 + d4_col = cutlass.Int32(tD4.iterator.toint()) + lane_hi + d1q_col = cutlass.Int32(tD1.iterator.toint()) + lane_hi + T_PAD + for jw in cutlass.range_constexpr((NEW_TOK + 3) // 4): + j_out = wg + 4 * jw + if j_out < NEW_TOK: + t_col = commit_len + j_out + o_q = _tmem_ld_x1(d1q_col + t_col) + o_d = _tmem_ld_x1(d4_col + t_col) + cute.arch.fence_view_async_tmem_load() + sVall[j_out * V + v_row] = cutlass.Float32(cutlass.BFloat16(o_q + o_d)) + cute.arch.barrier() + # ---- committed state: warpgroup wg stages columns [32 wg, +32) into box wg of the region + # (free once GEMM2 is done), then one thread per half issues the TMA stores (v2) + cute.arch.mbarrier_wait(mbar2b, 0) + tcgen05_fence_after_thread_sync() + for jj in cutlass.range_constexpr(4): + if wg == jj: + for cc in cutlass.range_constexpr(4): + cute.copy(ld8, tX_ld8[(None, None, None, 4 * jj + cc)], r8) + cute.arch.fence_view_async_tmem_load() + if cutlass.const_expr(STATE_BF16): + # bf16 box jj // 2, 16 B chunk 4 (jj % 2) + cc, round to nearest (v2's + # cvt.rn per value, two values per cvt.rn.bf16x2) + chunk = (4 * (jj % 2) + cc) ^ (wg_tidx % 8) + off_b = (jj // 2) * (V * 64) + 64 * wg_tidx + 8 * chunk + st_shared_f32x4( + saddr(sBoxh, off_b), + *[cvt_bf16x2(r8[2 * i + 1], r8[2 * i]) for i in range(4)], + ) + else: + for h in cutlass.range_constexpr(2): + c = 2 * cc + h + off_b = jj * (V * 32) + 32 * wg_tidx + ((4 * c) ^ vx) + for i in cutlass.range_constexpr(4): + sBoxf[off_b + i] = r8[4 * h + i] + tcgen05_fence_before_thread_sync() + if warp_idx > 0: + cute.arch.barrier_arrive(barrier_id=BAR_TMEM_FREE, number_of_threads=NUM_THREADS) + cute.arch.fence_view_async_shared() + if wg < 2: + cute.arch.barrier(barrier_id=BAR_STORE_LO, number_of_threads=256) + if tidx == 0: + for sl in cutlass.range_constexpr(NUM_SLABS // 2): + cute.copy(tma_atom_t, tS_list[sl], tG_list[sl]) + cute.arch.cp_async_bulk_commit_group() + else: + cute.arch.barrier(barrier_id=BAR_STORE_HI, number_of_threads=256) + if tidx == 256: + for sl in cutlass.range_constexpr(NUM_SLABS // 2, NUM_SLABS): + cute.copy(tma_atom_t, tS_list[sl], tG_list[sl]) + cute.arch.cp_async_bulk_commit_group() + if warp_idx == 0: + cute.arch.barrier(barrier_id=BAR_TMEM_FREE, number_of_threads=NUM_THREADS) + tcgen05_fence_after_thread_sync() + cute.arch.dealloc_tmem(tmem_ptr, TMEM_COLS) + # ---- gated RMSNorm epilogue (v2) + if ((warp_idx % 4) // 2 == 1) & (token_idx < NEW_TOK): + pre_values = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) + rms_partial = cutlass.Float32(0.0) + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + value = sVall[token_idx * V + v_idx] + gate = r_og[i] + sigmoid_gate = cute.arch.rcp_approx( + cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) + ) + pre_values[i] = value * r_ow[i] * sigmoid_gate + rms_partial += value * value + for offset in [16, 8, 4, 2, 1]: + rms_partial += cute.arch.shuffle_sync_bfly( + rms_partial, offset=offset, mask=-1, mask_and_clamp=31 + ) + inv_rms = cute.math.rsqrt( + rms_partial / cutlass.Float32(V) + cutlass.Float32(onorm_eps), fastmath=True + ) + output_row = i_n * NEW_TOK + token_idx + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + o[0, output_row, i_hv, v_idx] = cutlass.BFloat16(pre_values[i] * inv_rms) + # the state stores must have read their smem boxes before the CTA exits + if (tidx == 0) | (tidx == 256): + cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_v2.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_v2.py index dcb520c8c67b..fe72a75d98f6 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_v2.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_v2.py @@ -75,7 +75,7 @@ OFF_B3LO = 26624 # 16 x 128 Kb lo -def _asm_words(n, asm, constraints, operands, *, loc, ip): +def asm_words(n, asm, constraints, operands, *, loc, ip): res = llvm.inline_asm( llvm.StructType.get_literal([T.i32()] * n), operands, @@ -94,7 +94,7 @@ def _asm_words(n, asm, constraints, operands, *, loc, ip): def ld_global_b32x2(addr: Int64, *, loc=None, ip=None): """Two 32-bit words (8 B, 8-byte aligned) from global memory, unconverted: the load has no consumer until its value is placed, so a warp keeps all of its loads in flight.""" - return _asm_words( + return asm_words( 2, "ld.global.v2.b32 {$0, $1}, [$2];", "=r,=r,l", @@ -107,7 +107,7 @@ def ld_global_b32x2(addr: Int64, *, loc=None, ip=None): @dsl_user_op def ld_global_b32x4(addr: Int64, *, loc=None, ip=None): """Four 32-bit words (16 B, 16-byte aligned) from global memory, unconverted.""" - return _asm_words( + return asm_words( 4, "ld.global.v4.b32 {$0, $1, $2, $3}, [$4];", "=r,=r,=r,=r,l", @@ -120,7 +120,7 @@ def ld_global_b32x4(addr: Int64, *, loc=None, ip=None): @dsl_user_op def ld_shared_f32x4(addr: cutlass.Int32, *, loc=None, ip=None): """Four consecutive fp32 (16-byte aligned) from shared memory.""" - words = _asm_words( + words = asm_words( 4, "ld.shared.v4.b32 {$0, $1, $2, $3}, [$4];", "=r,=r,=r,=r,r", @@ -132,7 +132,7 @@ def ld_shared_f32x4(addr: cutlass.Int32, *, loc=None, ip=None): @dsl_user_op -def _st_asm(asm, constraints, operands, *, loc=None, ip=None): +def st_asm(asm, constraints, operands, *, loc=None, ip=None): llvm.inline_asm( None, operands, @@ -147,7 +147,7 @@ def _st_asm(asm, constraints, operands, *, loc=None, ip=None): def st_global_b32x2(addr, a, b): - _st_asm( + st_asm( "st.global.v2.b32 [$0], {$1, $2};", "l,r,r", [Int64(addr).ir_value(), cutlass.Int32(a).ir_value(), cutlass.Int32(b).ir_value()], @@ -155,7 +155,7 @@ def st_global_b32x2(addr, a, b): def st_global_b32x4(addr, a, b, c, d): - _st_asm( + st_asm( "st.global.v4.b32 [$0], {$1, $2, $3, $4};", "l,r,r,r,r", [Int64(addr).ir_value()] + [cutlass.Int32(x).ir_value() for x in (a, b, c, d)], @@ -163,7 +163,7 @@ def st_global_b32x4(addr, a, b, c, d): def st_shared_f32x4(addr, a, b, c, d): - _st_asm( + st_asm( "st.shared.v4.f32 [$0], {$1, $2, $3, $4};", "r,f,f,f,f", [cutlass.Int32(addr).ir_value()] + [cutlass.Float32(x).ir_value() for x in (a, b, c, d)], diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py index ac124d733866..d12712349ef0 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py @@ -55,7 +55,9 @@ contract (packed token rows, fused gated RMSNorm with a bf16 output, 3xTF32 recurrence, vector-aligned operands) run ``kda_decode_mtp_v2_kernel`` (``kda_mtp_decode_v2.py``: the same arithmetic in the same order, bitwise -identical results, shorter critical path); every other variant runs +identical results, shorter critical path), or with more ``(head, request)`` +tiles than SMs ``kda_decode_mtp_2cta_kernel`` (``kda_mtp_decode_2cta.py``: the +same arithmetic sized for two CTAs per SM); every other variant runs ``kda_decode_mtp_kernel``. """ @@ -85,12 +87,17 @@ TILE_K, kda_decode_mtp_kernel, ) + from ..attention.kernels.blackwell.kimi_k3_kda.kda_mtp_decode_2cta import ( + kda_decode_mtp_2cta_kernel, + ) from ..attention.kernels.blackwell.kimi_k3_kda.kda_mtp_decode_v2 import kda_decode_mtp_v2_kernel else: raise ImportError("Kimi K3 KDA MTP decode requires NVIDIA CUTLASS DSL") # One V pass: the kernel runs 16 warps x 8 rows over the full 128-row state. _TILE_V = 128 +# CTAs per SM kda_decode_mtp_2cta_kernel is compiled for (64 registers per thread). +_TWO_CTA_BLOCKS_PER_SM = 2 if IS_CUTLASS_DSL_AVAILABLE: @@ -228,14 +235,16 @@ def _state_tma_atoms( ) return _StateTmaAtoms(tma_atom_l, mS_tma, box_layout, tma_atom_t, mT_tma, box_layout_t) - def _launch_verify(launched, HV: int, N, split_v: int, stream) -> None: + def _launch_verify(launched, HV: int, N, split_v: int, stream, **launch_kwargs) -> None: if split_v == 2: # cluster pair: the two CTAs of a (head, request) must be co-scheduled (in-place state) launched.launch( grid=(HV, N, 2), block=[NUM_THREADS, 1, 1], cluster=(1, 1, 2), stream=stream ) else: - launched.launch(grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream) + launched.launch( + grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream, **launch_kwargs + ) @cute.jit def _run_kda_decode_mtp( @@ -479,6 +488,101 @@ def _run_kda_decode_mtp_v2( ) _launch_verify(launched, HV, N, SPLIT_V, stream) + @cute.jit + def _run_kda_decode_mtp_2cta( + h0: cute.Tensor, + x_q: cute.Tensor, + x_k: cute.Tensor, + x_v: cute.Tensor, + w_q: cute.Tensor, + w_k: cute.Tensor, + w_v: cute.Tensor, + cs_q: cute.Tensor, + cs_k: cute.Tensor, + cs_v: cute.Tensor, + A_log: cute.Tensor, + g: cute.Tensor, + dt_bias: cute.Tensor, + beta: cute.Tensor, + onorm_g: cute.Tensor, + onorm_weight: cute.Tensor, + o: cute.Tensor, + k_cache: cute.Tensor, + g_cache: cute.Tensor, + v_cache: cute.Tensor, + beta_cache: cute.Tensor, + ssm_state_indices: cute.Tensor, + num_accepted_tokens: cute.Tensor, + N: cutlass.Int32, + scale: cutlass.Constexpr[float], + HV: cutlass.Constexpr[int], + K: cutlass.Constexpr[int], + V: cutlass.Constexpr[int], + NUM_SPEC: cutlass.Constexpr[int], + KERNEL_WIDTH: cutlass.Constexpr[int], + lower_bound: cutlass.Constexpr[float], + onorm_eps: cutlass.Constexpr[float], + USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + stream: cuda.CUstream, + ): + """Host side of ``kda_decode_mtp_2cta_kernel`` (fp32 or bf16 state), launched for two + CTAs per SM (64 registers per thread).""" + mma = _verify_mma_atoms(cutlass.Float32, V, K, 3 * T_PAD) + state_tma = _state_tma_atoms(h0, h0, V, K, 1) + launched = kda_decode_mtp_2cta_kernel( + h0, + x_q, + x_k, + x_v, + w_q, + w_k, + w_v, + cs_q, + cs_k, + cs_v, + A_log, + g, + dt_bias, + beta, + onorm_g, + onorm_weight, + o, + k_cache, + g_cache, + v_cache, + beta_cache, + ssm_state_indices, + num_accepted_tokens, + mma.tiled_mma1, + mma.tiled_mma3, + mma.tiled_mma1k, + mma.smem_layout_b1, + mma.smem_layout_a2, + mma.smem_layout_a3, + mma.smem_layout_b3, + mma.smem_layout_b1k, + mma.tiled_mma4, + mma.tiled_mma2h, + mma.smem_layout_b2h, + mma.tiled_mma2q, + mma.smem_layout_b2q, + state_tma.tma_atom_l, + state_tma.mS_tma, + state_tma.box_layout, + state_tma.tma_atom_t, + state_tma.mT_tma, + scale, + HV, + K, + V, + NUM_SPEC, + KERNEL_WIDTH, + lower_bound, + onorm_eps, + USE_FLAT_LAYOUT, + ) + _launch_verify(launched, HV, N, 1, stream, min_blocks_per_mp=_TWO_CTA_BLOCKS_PER_SM) + def _require_stride_layout( *, @@ -727,6 +831,27 @@ def _v2_operands_ok(*, w_q, w_k, w_v, rows, conv_caches) -> bool: ) +def _two_cta_operands_ok(*, state, rows, conv_caches, g_cache) -> bool: + """Operand dtypes of ``kda_decode_mtp_2cta_kernel``: fp32 or bf16 state, bf16 new-token + projections and conv caches (staged as bf16), fp32 replay gates.""" + return ( + state.dtype in (torch.float32, torch.bfloat16) + and all(t.dtype == torch.bfloat16 for t in (*rows, *conv_caches)) + and g_cache.dtype == torch.float32 + ) + + +_num_sms_cache = {} + + +def _num_sms(device: torch.device) -> int: + dev = torch.device(device) + index = dev.index if dev.index is not None else torch.cuda.current_device() + if index not in _num_sms_cache: + _num_sms_cache[index] = torch.cuda.get_device_properties(index).multi_processor_count + return _num_sms_cache[index] + + def _select_verify_kernel( *, packed_token_layout: bool, @@ -735,16 +860,22 @@ def _select_verify_kernel( bf16_math: bool, replay_only: bool, num_spec: int, + num_tiles: int, + num_sms: int, v2_operands_ok: Callable[[], bool], -) -> Literal["v1", "v2"]: + two_cta_operands_ok: Callable[[], bool], +) -> Literal["v1", "v2", "two_cta"]: """The verify kernel a call runs (see the module docstring). - ``v2`` (``kda_decode_mtp_v2_kernel``) serves the runtime verify contract: packed token rows, - the fused gated RMSNorm with a bf16 output, the 3xTF32 recurrence and vector-aligned operands - (``v2_operands_ok``, evaluated last). ``v1`` (``kda_decode_mtp_kernel``) serves everything - else; both compute bitwise the same results where both apply. + The runtime verify contract is packed token rows, the fused gated RMSNorm with a bf16 + output, the 3xTF32 recurrence and vector-aligned operands (``v2_operands_ok``). On it, + ``two_cta`` (``kda_decode_mtp_2cta_kernel``, two CTAs per SM) serves grids of more + ``(head, request)`` tiles than SMs with the operand dtypes it stages + (``two_cta_operands_ok``), and ``v2`` (``kda_decode_mtp_v2_kernel``) serves the rest. ``v1`` + (``kda_decode_mtp_kernel``) serves every other variant. All compute bitwise the same results + where they apply. The operand checks are evaluated last. """ - if ( + if not ( packed_token_layout and fuse_output_norm and not quantize_output @@ -753,8 +884,10 @@ def _select_verify_kernel( and num_spec <= MAX_NUM_SPEC and v2_operands_ok() ): - return "v2" - return "v1" + return "v1" + if num_tiles > num_sms and two_cta_operands_ok(): + return "two_cta" + return "v2" # (device_index, enabled) -> persistent int32 [1] control tensor. Keys are @@ -1050,6 +1183,8 @@ def kda_mtp_decode_impl( bf16_math=bf16_math, replay_only=replay_only, num_spec=num_spec, + num_tiles=HV * N, + num_sms=_num_sms(x_q.device), v2_operands_ok=lambda: _v2_operands_ok( w_q=w_q, w_k=w_k, @@ -1057,9 +1192,16 @@ def kda_mtp_decode_impl( rows=(x_q_arg, x_k_arg, x_v_arg, g, onorm_g, k_cache, g_cache, v_cache), conv_caches=(cs_q, cs_k, cs_v), ), + two_cta_operands_ok=lambda: _two_cta_operands_ok( + state=h0_arg, + rows=(x_q, x_k, x_v), + conv_caches=(cs_q, cs_k, cs_v), + g_cache=g_cache, + ), ) - if kernel == "v2": - v2_args = ( + if kernel != "v1": + # Shared operands of kda_decode_mtp_v2_kernel and kda_decode_mtp_2cta_kernel. + args = [ (h0_arg, 16), (x_q_arg, 16), (x_k_arg, 16), @@ -1083,42 +1225,40 @@ def kda_mtp_decode_impl( (beta_cache, beta_cache_assumed_align), (ssm_state_indices, 4), (num_accepted_tokens, 4), + ] + if kernel == "two_cta": + runner = _run_kda_decode_mtp_2cta + variant = {} + else: + runner = _run_kda_decode_mtp_v2 + variant = {"SPLIT_V": split_v} + constexprs = dict( + scale=scale, + HV=HV, + K=K, + V=V_dim, + NUM_SPEC=num_spec, + KERNEL_WIDTH=W, + lower_bound=float(lower_bound), + onorm_eps=float(onorm_eps), + USE_FLAT_LAYOUT=use_flat_layout, + **variant, ) - v2_key = ( - "v2", - tuple(_layout_key(t, dynamic_layout=True, assumed_align=a) for t, a in v2_args), - scale, - HV, - K, - V_dim, - num_spec, - W, - float(lower_bound), - float(onorm_eps), - use_flat_layout, - split_v, + key = ( + kernel, + tuple(_layout_key(t, dynamic_layout=True, assumed_align=a) for t, a in args), + tuple(sorted(constexprs.items())), ) - if v2_key not in _compiled_cache: - logger.info( - f"kda_mtp_decode: compiling v2 variant H={HV} num_spec={num_spec} split_v={split_v}" - ) - _compiled_cache[v2_key] = cute.compile( - _run_kda_decode_mtp_v2, - *[_dlpack_arg(t, assumed_align=a) for t, a in v2_args], + if key not in _compiled_cache: + logger.info(f"kda_mtp_decode: compiling {kernel} variant {constexprs}") + _compiled_cache[key] = cute.compile( + runner, + *[_dlpack_arg(t, assumed_align=a) for t, a in args], N, - scale=scale, - HV=HV, - K=K, - V=V_dim, - NUM_SPEC=num_spec, - KERNEL_WIDTH=W, - lower_bound=float(lower_bound), - onorm_eps=float(onorm_eps), - USE_FLAT_LAYOUT=use_flat_layout, - SPLIT_V=split_v, + **constexprs, stream=stream, ) - _compiled_cache[v2_key](*[_dlpack_arg(t, assumed_align=a) for t, a in v2_args], N, stream) + _compiled_cache[key](*[_dlpack_arg(t, assumed_align=a) for t, a in args], N, stream) return out key = ( diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py index bca3d6206507..174d59e3994a 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py @@ -1324,9 +1324,17 @@ def small(*shape, scale): def _packed_fused_norm_inputs(B, H, num_spec, state_dtype): """Runtime-contract inputs for the packed / fused-norm path: section-strided q / k / v / output-gate views of one fused projection, the bf16 beta tail of [f_a | beta], replay caches - and conv windows chained from a CPU-golden round 1, mixed acceptance 0..num_spec and + and bf16 conv windows chained from a CPU-golden round 1, mixed acceptance 0..num_spec and non-identity state slots.""" data, _ = _round2_inputs(B, H, num_spec) + for name in ("cs_q", "cs_k", "cs_v"): + # The runtime's conv caches are bf16; round 1 wrote bf16 values into them. + src = data[name] + dst = torch.empty( + src.shape[0], src.shape[2], src.shape[1], dtype=torch.bfloat16, device=src.device + ).transpose(1, 2) + dst.copy_(src) + data[name] = dst packed, _ = _pack_new_token_projection_views(data) slots = torch.arange(B - 1, -1, -1, dtype=torch.int32, device="cuda") packed["ssm_state_indices"] = slots @@ -1367,9 +1375,9 @@ def _packed_fused_norm_inputs(B, H, num_spec, state_dtype): ) -def _run_with_verify_kernel(monkeypatch, kernel, run): - """``run()`` with the verify kernel forced to ``kernel``; asserts the selection would have - picked it on its own unless ``kernel`` is ``"v1"`` (the reference).""" +def _run_with_verify_kernel(monkeypatch, kernel, run, expected_selection=None): + """``run()`` with the verify kernel forced to ``kernel``; with ``expected_selection``, also + asserts the kernel ``_select_verify_kernel`` picks for the call on its own.""" from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops as mtp_ops select = mtp_ops._select_verify_kernel @@ -1383,11 +1391,15 @@ def forced(**kwargs): patch.setattr(mtp_ops, "_select_verify_kernel", forced) result = run() assert selected, "the call did not reach the verify kernel selection" - if kernel != "v1": - assert selected == [kernel], f"selection picked {selected}, not {kernel}" + if expected_selection is not None: + assert selected == [expected_selection], f"selection picked {selected}" return result +def _multi_wave(B, H): + return B * H > torch.cuda.get_device_properties(0).multi_processor_count + + @pytest.mark.parametrize( "B,H,num_spec", [(1, 6, 6), (4, 6, 7), (16, 6, 6), (32, 6, 7), (128, 6, 6), (8, 12, 7)], @@ -1414,13 +1426,37 @@ def run(): ) ref = _run_with_verify_kernel(monkeypatch, "v1", run) - new = _run_with_verify_kernel(monkeypatch, "v2", run) + expected = "two_cta" if _multi_wave(B, H) else "v2" + new = _run_with_verify_kernel(monkeypatch, "v2", run, expected_selection=expected) for name in _VERIFY_RESULTS: assert torch.equal(new[name], ref[name]), f"{name}: v2 differs from v1" +@pytest.mark.parametrize( + "B,H,num_spec", + [(26, 6, 2), (32, 6, 6), (32, 6, 7), (64, 6, 5), (128, 6, 6), (4, 6, 7), (8, 96, 7)], + ids=lambda v: str(v), +) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +def test_two_cta_kernel_matches_v1_bitwise(B, H, num_spec, state_dtype, monkeypatch): + """The two-CTAs-per-SM verify kernel performs v2's arithmetic in v2's order: outputs, + committed state, conv windows and replay caches match ``kda_decode_mtp_kernel`` bit for bit, + for single- and multi-wave grids (selected for multi-wave grids only).""" + packed = _packed_fused_norm_inputs(B, H, num_spec, state_dtype) + + def run(): + return cute_run(packed, packed_token_layout=True, fuse_output_norm=True) + + ref = _run_with_verify_kernel(monkeypatch, "v1", run) + expected = "two_cta" if _multi_wave(B, H) else "v2" + new = _run_with_verify_kernel(monkeypatch, "two_cta", run, expected_selection=expected) + for name in _VERIFY_RESULTS: + assert torch.equal(new[name], ref[name]), f"{name}: two-CTA kernel differs from v1" + + def test_verify_kernel_selection(): - """v2 serves only the runtime verify contract; every other variant stays on v1.""" + """v2 and the two-CTA kernel serve only the runtime verify contract, the two-CTA kernel + for grids of more tiles than SMs; every other variant stays on v1.""" from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import _select_verify_kernel contract = dict( @@ -1430,9 +1466,19 @@ def test_verify_kernel_selection(): bf16_math=False, replay_only=False, num_spec=7, + num_tiles=148, + num_sms=148, v2_operands_ok=lambda: True, + two_cta_operands_ok=lambda: True, ) assert _select_verify_kernel(**contract) == "v2" + assert _select_verify_kernel(**{**contract, "num_tiles": 149}) == "two_cta" + assert ( + _select_verify_kernel( + **{**contract, "num_tiles": 149, "two_cta_operands_ok": lambda: False} + ) + == "v2" + ) for name, value in ( ("packed_token_layout", False), ("fuse_output_norm", False), diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py index 00ada502421e..e159c69db5bf 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py @@ -492,6 +492,73 @@ def test_fused_verify_falls_back_without_output_norm_weights(): assert _rep("fallback vs fused epilogue", fallback_out, fused_out) +@torch.no_grad() +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +def test_fused_verify_multi_wave_runs_two_cta_kernel(monkeypatch, state_dtype): + """A verify grid of more (head, request) tiles than SMs runs the two-CTAs-per-SM kernel, with + bitwise the outputs, state and replay caches of ``kda_decode_mtp_kernel``, over two rounds + with mixed acceptance.""" + from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_kda_mtp_ops as mtp_ops + + num_sms = torch.cuda.get_device_properties(0).multi_processor_count + batch = num_sms // H + 4 + steps = M + 1 + runtime = _make_runtime(seed=1) + runtime.finalize_decode_weights() + conv_pool, ssm_pool = _make_pools(batch, seed=2) + cache = _make_fused_layer_cache(batch, conv_pool) + for name in ("kda_conv_q", "kda_conv_k", "kda_conv_v"): + # The runtime's conv caches are bf16 (dim-contiguous). + src = getattr(cache, name) + dst = torch.empty_like(src.transpose(-1, -2), dtype=torch.bfloat16).transpose(-1, -2) + dst.copy_(src) + setattr(cache, name, dst) + worlds = [] + for _ in range(2): + worlds.append( + ( + ssm_pool.to(state_dtype, copy=True), + SimpleNamespace( + **{ + key: value.clone() if isinstance(value, torch.Tensor) else value + for key, value in vars(cache).items() + } + ), + ) + ) + select = mtp_ops._select_verify_kernel + selected = [] + + def v1_reference(**kwargs): + selected.append(select(**kwargs)) + return "v1" + + slots = torch.arange(batch, dtype=torch.int32, device="cuda") + gen = torch.Generator(device="cuda").manual_seed(3) + for rnd in range(2): + x = (torch.randn(batch * steps, HIDDEN, generator=gen, device="cuda") * 0.5).bfloat16() + with monkeypatch.context() as patch: + patch.setattr(mtp_ops, "_select_verify_kernel", v1_reference) + ref = runtime.forward_verify_fused(x, steps, worlds[0][1], worlds[0][0], slots) + out = runtime.forward_verify_fused(x, steps, worlds[1][1], worlds[1][0], slots) + assert selected[-1] == "two_cta" + assert torch.equal(out, ref), f"round {rnd} out" + assert torch.equal(worlds[1][0], worlds[0][0]), f"round {rnd} state" + for key in ( + "kda_conv_q", + "kda_conv_k", + "kda_conv_v", + "kda_k_cache", + "kda_g_cache", + "kda_v_cache", + "kda_beta_cache", + ): + assert torch.equal(getattr(worlds[1][1], key), getattr(worlds[0][1], key)), key + accept = torch.randint(0, M + 1, (batch,), generator=gen, device="cuda", dtype=torch.int32) + for _, world_cache in worlds: + world_cache.prev_num_accepted_tokens.copy_(accept) + + class _PrequantizedProjectionStub(torch.nn.Module): """o_proj stand-in that records the prequantized pair it receives.""" From 65f1c278a5c822df5e572865e65ba4041b73da64 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:13:21 -0700 Subject: [PATCH 4/8] [None][perf] Kimi K3 KDA MTP verify: MXFP8 output on the two-CTAs-per-SM kernel On SM107 the verify call hands o_proj the MXFP8 operand of the Rubin CuTe GEMM (quantize_output), which only kda_decode_mtp_kernel produced: one 512-thread CTA per SM, and the v2 kernel has no MXFP8 epilogue. - kda_mtp_decode_2cta.py: an MXFP8_OUT epilogue that applies kda_decode_mtp_kernel's quantization (bf16-rounded normalized values, UE8M0 scale of bf16(max |p| * inv_rms) / 448, E4M3 values, the scale byte four times in the R128c4 layout). No value-producing instruction of the BF16 mode changes. - _select_verify_kernel: MXFP8-output calls on the verify contract run that kernel, one CTA per (head, request) tile at every grid size (also where kda_decode_mtp_kernel used cluster pairs). The contract (packed rows, fused norm, 3xTF32) now requires its operands (vector-aligned rows and conv caches, bf16 projections and conv caches, fp32 replay gates; ValueError otherwise) instead of falling back to kda_decode_mtp_kernel, which keeps the replay-only commit, the shifted-metadata layout, the raw output and the BF16 recurrence and is the bitwise reference of the v2 and two-CTA kernels. The MXFP8 operand is consumed only by the Rubin o_proj GEMM (supports_prequantized_input), so on SM100/103 production verify calls never request it. Bitwise vs kda_decode_mtp_kernel: E4M3 values, valid-row scales, committed state, conv windows and replay caches. GB300 test exemption: effective only on SM107 (only Rubin requests the MXFP8 verify output). Rubin evidence (hecate R200, 212 SMs, bf16 state, CUDA graph of 69 calls, us per call, default kernel -> two-CTA kernel, L0 792445): - dep16 (96 heads/rank): b1 mtp7 7.97 -> 7.57, b4 mtp7 16.57 -> 11.07, b8 mtp7 30.61 -> 19.59 (-0.76 ms/step), b16 mtp7 59.39 -> 37.26, b16 mtp5 53.71 -> 34.82 (-1.30 ms/step), b32 mtp5 101.5 -> 66.3 - helix8 (12 heads/rank): b8 mtp7 7.95 -> 7.55, b32 mtp5 15.15 -> 10.65 - helix16 (6 heads/rank): b8 mtp7 7.59 -> 7.06 - dep16 b8 mtp7 gen-only, same image: 32.38 -> 31.37 ms/step (-3.1%), 118.6 -> 122.4 tok/s/user (jobs 792567/792568); outputs, state and caches bitwise identical to the previous path on 18 decode shapes (L0 792445, 792447). GB300 job 1368796 (cold L2, bf16 state): two-CTA MXFP8 vs kda_decode_mtp_kernel 8.41 -> 7.68 us (H=6 b1 mtp3) .. 45.29 -> 29.13 us (dep16 b8 mtp3). Platform: R. Tests: test_kda_mtp_decode_cute_parity.py::test_two_cta_mxfp8_output_matches_v1_bitwise (helix and dep16 head counts, 3..7 drafts, fp32/bf16 state; the kernel is arch-generic, so it runs on SM100/103 too) and ::test_verify_kernel_selection; test_kimi_kda_fused_verify_parity.py uses the bf16 conv caches the runtime allocates. GB300 (OCI AGA) jobs 1369212, 1387340: all pass. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../kimi_k3_kda/kda_mtp_decode_2cta.py | 66 +++++++++++-- .../cute_dsl_kimi_k3_kda_mtp_ops.py | 99 ++++++++++--------- .../test_kda_mtp_decode_cute_parity.py | 64 +++++++++--- .../test_kimi_kda_fused_verify_parity.py | 11 +-- 4 files changed, 167 insertions(+), 73 deletions(-) diff --git a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_2cta.py b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_2cta.py index ce1b734f165e..3b20983b079c 100644 --- a/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_2cta.py +++ b/tensorrt_llm/_torch/attention/kernels/blackwell/kimi_k3_kda/kda_mtp_decode_2cta.py @@ -44,9 +44,15 @@ * **Named barriers: ids <= 4, all compile-time constants** (see ``BAR_WG0``), which SM107 needs to co-schedule two CTAs. -Alone on an SM a tile is not faster than in the v2 kernel (the precompute reads bf16 and GEMM1's -two A operands are serialized), while two tiles per SM finish sooner than two waves of v2 tiles, -so the kernel pays off for grids with more tiles than SMs. +The epilogue writes the BF16 output or, with ``MXFP8_OUT``, its MXFP8 quantization with the +scale layout and the arithmetic of ``kda_decode_mtp_kernel``'s ``QUANTIZE_OUTPUT`` epilogue (the +operand of the Rubin o_proj GEMM). + +A tile alone on an SM is in general not faster than in the v2 kernel (the precompute reads bf16 +and GEMM1's two A operands are serialized), while two tiles per SM finish sooner than two waves +of v2 tiles, so the kernel pays off for grids with more tiles than SMs. With the MXFP8 output a +tile alone on an SM also finishes sooner than in ``kda_decode_mtp_kernel`` (one CTA or a cluster +pair). """ import cutlass @@ -241,6 +247,8 @@ def kda_decode_mtp_2cta_kernel( beta_cache: cute.Tensor, ssm_state_indices: cute.Tensor, num_accepted_tokens: cute.Tensor, + o8: cute.Tensor, + osf: cute.Tensor, tiled_mma1: cute.TiledMma, tiled_mma3: cute.TiledMma, tiled_mma1k: cute.TiledMma, @@ -268,8 +276,14 @@ def kda_decode_mtp_2cta_kernel( lower_bound: cutlass.Constexpr[float], onorm_eps: cutlass.Constexpr[float], USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + MXFP8_OUT: cutlass.Constexpr[bool], ): - """See the module docstring. Grid ``(HV, N, 1)``, one (head, request) tile per CTA.""" + """See the module docstring. Grid ``(HV, N, 1)``, one (head, request) tile per CTA. + + ``MXFP8_OUT``: instead of the BF16 output ``o``, write the MXFP8 operand of the Rubin CuTe + o_proj GEMM with the arithmetic of ``kda_decode_mtp_kernel``'s ``QUANTIZE_OUTPUT`` epilogue: + ``o8`` (E4M3 ``[T, HV * V]``) and ``osf`` (uint8, UE8M0 1x128 scales in the R128c4 layout: + four copies per (row, head)). Without it ``o8`` and ``osf`` are not accessed.""" tidx, _, _ = cute.arch.thread_idx() in_warp_tid = tidx % 32 warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -1171,6 +1185,7 @@ def region_tile(off, layout): if ((warp_idx % 4) // 2 == 1) & (token_idx < NEW_TOK): pre_values = cute.make_rmem_tensor(cute.make_layout((V // 32,)), cutlass.Float32) rms_partial = cutlass.Float32(0.0) + pre_max = cutlass.Float32(0.0) for i in cutlass.range(V // 32, unroll_full=True): v_idx = i * 32 + in_warp_tid value = sVall[token_idx * V + v_idx] @@ -1178,19 +1193,54 @@ def region_tile(off, layout): sigmoid_gate = cute.arch.rcp_approx( cutlass.Float32(1.0) + cute.math.exp(-gate, fastmath=True) ) - pre_values[i] = value * r_ow[i] * sigmoid_gate + p_v = value * r_ow[i] * sigmoid_gate + pre_values[i] = p_v rms_partial += value * value + if cutlass.const_expr(MXFP8_OUT): + pre_max = cute.arch.fmax(pre_max, cute.arch.fmax(p_v, -p_v)) for offset in [16, 8, 4, 2, 1]: rms_partial += cute.arch.shuffle_sync_bfly( rms_partial, offset=offset, mask=-1, mask_and_clamp=31 ) + if cutlass.const_expr(MXFP8_OUT): + pre_max = cute.arch.fmax( + pre_max, + cute.arch.shuffle_sync_bfly(pre_max, offset=offset, mask=-1, mask_and_clamp=31), + ) inv_rms = cute.math.rsqrt( rms_partial / cutlass.Float32(V) + cutlass.Float32(onorm_eps), fastmath=True ) output_row = i_n * NEW_TOK + token_idx - for i in cutlass.range(V // 32, unroll_full=True): - v_idx = i * 32 + in_warp_tid - o[0, output_row, i_hv, v_idx] = cutlass.BFloat16(pre_values[i] * inv_rms) + if cutlass.const_expr(MXFP8_OUT): + # kda_decode_mtp_kernel's QUANTIZE_OUTPUT epilogue: BF16-rounded normalized values, + # UE8M0 scale of bf16(max |p_v| * inv_rms) / 448 (bf16 rounding is monotonic), E4M3 + # values, the scale byte replicated four times in the R128c4 layout + output_amax = cutlass.Float32(cutlass.BFloat16(pre_max * inv_rms)) + output_amax = cute.arch.fmax(output_amax, cutlass.Float32(1.0e-10)) + output_sf = (output_amax * cutlass.Float32(1.0 / 448.0)).to(cutlass.Float8E8M0FNU) + quant_scale = cute.arch.rcp_approx(cutlass.Float32(output_sf)) + for i in cutlass.range(V // 32, unroll_full=True): + normalized = cutlass.Float32(cutlass.BFloat16(pre_values[i] * inv_rms)) + o8[output_row, hv_off + i * 32 + in_warp_tid] = cutlass.Float8E4M3FN( + normalized * quant_scale + ) + if in_warp_tid == 0: + scale_bytes = cute.make_tensor( + cute.recast_ptr(osf.iterator, dtype=cutlass.Float8E8M0FNU), + cute.make_layout(osf.shape), + ) + scale_offset = ( + (output_row // 128) * HV * 512 + + i_hv * 512 + + (output_row % 32) * 16 + + ((output_row % 128) // 32) * 4 + ) + for replica in cutlass.range_constexpr(4): + scale_bytes[scale_offset + replica] = output_sf + else: + for i in cutlass.range(V // 32, unroll_full=True): + v_idx = i * 32 + in_warp_tid + o[0, output_row, i_hv, v_idx] = cutlass.BFloat16(pre_values[i] * inv_rms) # the state stores must have read their smem boxes before the CTA exits if (tidx == 0) | (tidx == 256): cute.arch.cp_async_bulk_wait_group(0, read=True) diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py index d12712349ef0..9484e7acd5bd 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_kimi_k3_kda_mtp_ops.py @@ -51,14 +51,20 @@ ``HV == H``, one V pass (TILE_V=128, 512 threads). ``num_spec`` is a compile-time constant per cache allocation. -Kernel selection (``_select_verify_kernel``): calls on the runtime verify -contract (packed token rows, fused gated RMSNorm with a bf16 output, 3xTF32 -recurrence, vector-aligned operands) run ``kda_decode_mtp_v2_kernel`` -(``kda_mtp_decode_v2.py``: the same arithmetic in the same order, bitwise -identical results, shorter critical path), or with more ``(head, request)`` -tiles than SMs ``kda_decode_mtp_2cta_kernel`` (``kda_mtp_decode_2cta.py``: the -same arithmetic sized for two CTAs per SM); every other variant runs -``kda_decode_mtp_kernel``. +Kernel selection (``_select_verify_kernel``). The runtime verify contract is +packed token rows, the fused gated RMSNorm and the 3xTF32 recurrence (at most +``MAX_NUM_SPEC`` drafts, vector-aligned bf16 operands). The three kernels have +the same arithmetic in the same order and bitwise identical results: + +* ``kda_decode_mtp_kernel``: every call off the contract (replay-only commits, + the shifted-metadata layout, the raw output, the BF16 recurrence of + ``TRTLLM_KDA_MTP_BF16_MATH=1``). Its MXFP8 epilogue serves only that BF16 + recurrence and is the bitwise reference of the two-CTA kernel's. +* ``kda_decode_mtp_2cta_kernel`` (``kda_mtp_decode_2cta.py``, two CTAs per + SM): contract calls with more ``(head, request)`` tiles than SMs or with the + MXFP8 output. +* ``kda_decode_mtp_v2_kernel`` (``kda_mtp_decode_v2.py``, one CTA or a cluster + pair per tile): the other contract calls (BF16 output, one wave). """ import os @@ -513,6 +519,8 @@ def _run_kda_decode_mtp_2cta( beta_cache: cute.Tensor, ssm_state_indices: cute.Tensor, num_accepted_tokens: cute.Tensor, + o8: cute.Tensor, + osf: cute.Tensor, N: cutlass.Int32, scale: cutlass.Constexpr[float], HV: cutlass.Constexpr[int], @@ -523,6 +531,7 @@ def _run_kda_decode_mtp_2cta( lower_bound: cutlass.Constexpr[float], onorm_eps: cutlass.Constexpr[float], USE_FLAT_LAYOUT: cutlass.Constexpr[bool], + MXFP8_OUT: cutlass.Constexpr[bool], stream: cuda.CUstream, ): """Host side of ``kda_decode_mtp_2cta_kernel`` (fp32 or bf16 state), launched for two @@ -553,6 +562,8 @@ def _run_kda_decode_mtp_2cta( beta_cache, ssm_state_indices, num_accepted_tokens, + o8, + osf, mma.tiled_mma1, mma.tiled_mma3, mma.tiled_mma1k, @@ -580,6 +591,7 @@ def _run_kda_decode_mtp_2cta( lower_bound, onorm_eps, USE_FLAT_LAYOUT, + MXFP8_OUT, ) _launch_verify(launched, HV, N, 1, stream, min_blocks_per_mp=_TWO_CTA_BLOCKS_PER_SM) @@ -816,10 +828,15 @@ def _vec4_rows(tensor: torch.Tensor, dim: int = -1) -> bool: ) -def _v2_operands_ok(*, w_q, w_k, w_v, rows, conv_caches) -> bool: - """Operand layouts / dtypes the v2 kernel's vector accesses support (see `_vec4_rows`): - token rows and replay caches vectorized along their last dimension, the dim-contiguous conv - caches along the channel dimension, conv weights as contiguous [channel][4] rows.""" +def _verify_operands_ok(*, w_q, w_k, w_v, state, projections, rows, conv_caches, g_cache) -> bool: + """Operands of the runtime verify contract that ``kda_decode_mtp_v2_kernel`` and + ``kda_decode_mtp_2cta_kernel`` take. + + Layouts for their vector accesses (see `_vec4_rows`): token rows and replay caches + vectorized along their last dimension, the dim-contiguous conv caches along the channel + dimension, conv weights as contiguous [channel][4] rows. Dtypes: fp32 or bf16 state, bf16 + new-token projections and conv caches, fp32 replay gates. + """ weights_ok = all( w.dtype == w_q.dtype and w.is_contiguous() and w.shape[1] == 4 and _vec4_rows(w) for w in (w_q, w_k, w_v) @@ -828,15 +845,8 @@ def _v2_operands_ok(*, w_q, w_k, w_v, rows, conv_caches) -> bool: weights_ok and all(_vec4_rows(t) for t in rows) and all(_vec4_rows(t, dim=1) for t in conv_caches) - ) - - -def _two_cta_operands_ok(*, state, rows, conv_caches, g_cache) -> bool: - """Operand dtypes of ``kda_decode_mtp_2cta_kernel``: fp32 or bf16 state, bf16 new-token - projections and conv caches (staged as bf16), fp32 replay gates.""" - return ( - state.dtype in (torch.float32, torch.bfloat16) - and all(t.dtype == torch.bfloat16 for t in (*rows, *conv_caches)) + and state.dtype in (torch.float32, torch.bfloat16) + and all(t.dtype == torch.bfloat16 for t in (*projections, *conv_caches)) and g_cache.dtype == torch.float32 ) @@ -862,30 +872,26 @@ def _select_verify_kernel( num_spec: int, num_tiles: int, num_sms: int, - v2_operands_ok: Callable[[], bool], - two_cta_operands_ok: Callable[[], bool], + operands_ok: Callable[[], bool], ) -> Literal["v1", "v2", "two_cta"]: - """The verify kernel a call runs (see the module docstring). - - The runtime verify contract is packed token rows, the fused gated RMSNorm with a bf16 - output, the 3xTF32 recurrence and vector-aligned operands (``v2_operands_ok``). On it, - ``two_cta`` (``kda_decode_mtp_2cta_kernel``, two CTAs per SM) serves grids of more - ``(head, request)`` tiles than SMs with the operand dtypes it stages - (``two_cta_operands_ok``), and ``v2`` (``kda_decode_mtp_v2_kernel``) serves the rest. ``v1`` - (``kda_decode_mtp_kernel``) serves every other variant. All compute bitwise the same results - where they apply. The operand checks are evaluated last. - """ + """The verify kernel a call runs, by the rule in the module docstring. A call on the + contract whose operands fail ``operands_ok`` (`_verify_operands_ok`, evaluated last) is + rejected.""" if not ( packed_token_layout and fuse_output_norm - and not quantize_output and not bf16_math and not replay_only and num_spec <= MAX_NUM_SPEC - and v2_operands_ok() ): return "v1" - if num_tiles > num_sms and two_cta_operands_ok(): + if not operands_ok(): + raise ValueError( + "KDA MTP verify: the packed, fused-norm verify takes 16-byte-vectorizable bf16 " + "projections and conv caches, contiguous conv weights, fp32 replay gates and an fp32 " + "or bf16 state" + ) + if quantize_output or num_tiles > num_sms: return "two_cta" return "v2" @@ -991,8 +997,10 @@ def kda_mtp_decode_impl( or packed uniformly by request. w_q/w_k/w_v: conv weights ``[H*128, W]`` fp32, width-contiguous. cs_q/cs_k/cs_v: extended conv caches ``[pool, H*128, >= W-1+M]`` - fp32, dim-contiguous. Columns ``[0, W-1)`` are the committed - window; tail columns hold raw pending-draft inputs. Mutated. + of raw projection inputs, bf16 (the runtime's caches; the packed + fused-norm verify requires bf16), dim-contiguous. Columns + ``[0, W-1)`` are the committed window; tail columns hold raw + pending-draft inputs. Mutated. g, beta: raw gate ``[1, T, H, 128]`` and beta ``[1, T, H]`` bf16. A_log, dt_bias: fp32 ``[H]`` / ``[H*128]``. recurrent_state: pool ``[pool, H, V, K]`` fp32 or bf16, **V-first** layout @@ -1185,16 +1193,13 @@ def kda_mtp_decode_impl( num_spec=num_spec, num_tiles=HV * N, num_sms=_num_sms(x_q.device), - v2_operands_ok=lambda: _v2_operands_ok( + operands_ok=lambda: _verify_operands_ok( w_q=w_q, w_k=w_k, w_v=w_v, - rows=(x_q_arg, x_k_arg, x_v_arg, g, onorm_g, k_cache, g_cache, v_cache), - conv_caches=(cs_q, cs_k, cs_v), - ), - two_cta_operands_ok=lambda: _two_cta_operands_ok( state=h0_arg, - rows=(x_q, x_k, x_v), + projections=(x_q, x_k, x_v), + rows=(x_q_arg, x_k_arg, x_v_arg, g, onorm_g, k_cache, g_cache, v_cache), conv_caches=(cs_q, cs_k, cs_v), g_cache=g_cache, ), @@ -1228,7 +1233,9 @@ def kda_mtp_decode_impl( ] if kernel == "two_cta": runner = _run_kda_decode_mtp_2cta - variant = {} + # E4M3 values and R128c4 UE8M0 scales of the MXFP8 output (placeholders without it) + args += [(out, 16), (output_scale if quantize_output else out, 16)] + variant = {"MXFP8_OUT": bool(quantize_output)} else: runner = _run_kda_decode_mtp_v2 variant = {"SPLIT_V": split_v} @@ -1259,6 +1266,8 @@ def kda_mtp_decode_impl( stream=stream, ) _compiled_cache[key](*[_dlpack_arg(t, assumed_align=a) for t, a in args], N, stream) + if quantize_output: + return out, output_scale return out key = ( diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py index 174d59e3994a..4dc1662a2797 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py @@ -1454,9 +1454,51 @@ def run(): assert torch.equal(new[name], ref[name]), f"{name}: two-CTA kernel differs from v1" +def _valid_mxfp8_scales(output_scale, rows, heads): + """The R128c4 UE8M0 scale bytes of the ``rows`` output rows (four copies per (row, head)); + the padding rows of the layout are not written.""" + r = torch.arange(rows, device=output_scale.device) + h = torch.arange(heads, device=output_scale.device) + base = ( + (r // 128)[:, None] * heads * 512 + + h[None, :] * 512 + + ((r % 32) * 16 + ((r % 128) // 32) * 4)[:, None] + ) + return output_scale[base[..., None] + torch.arange(4, device=output_scale.device)] + + +@pytest.mark.parametrize( + "B,H,num_spec", + [(1, 96, 7), (8, 96, 7), (16, 96, 5), (8, 12, 7), (32, 12, 5), (8, 6, 7), (4, 96, 3)], + ids=lambda v: str(v), +) +@pytest.mark.parametrize("state_dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +def test_two_cta_mxfp8_output_matches_v1_bitwise(B, H, num_spec, state_dtype, monkeypatch): + """MXFP8-output verify calls run the two-CTA kernel at every grid size, with the E4M3 values, + valid-row scales, committed state, conv windows and replay caches of + ``kda_decode_mtp_kernel``'s quantization epilogue bit for bit (helix and dep16 head counts).""" + packed = _packed_fused_norm_inputs(B, H, num_spec, state_dtype) + + def run(): + return cute_run(packed, packed_token_layout=True, quantize_output=True) + + ref = _run_with_verify_kernel(monkeypatch, "v1", run) + new = _run_with_verify_kernel(monkeypatch, "two_cta", run, expected_selection="two_cta") + rows = B * (num_spec + 1) + assert new["out"].dtype == torch.float8_e4m3fn + assert torch.equal(new["out"].view(torch.uint8), ref["out"].view(torch.uint8)) + assert torch.equal( + _valid_mxfp8_scales(new["output_scale"], rows, H), + _valid_mxfp8_scales(ref["output_scale"], rows, H), + ) + for name in _VERIFY_RESULTS[1:]: + assert torch.equal(new[name], ref[name]), f"{name}: two-CTA kernel differs from v1" + + def test_verify_kernel_selection(): """v2 and the two-CTA kernel serve only the runtime verify contract, the two-CTA kernel - for grids of more tiles than SMs; every other variant stays on v1.""" + for MXFP8 output and for grids of more tiles than SMs; the contract requires its operands; + every other variant stays on v1.""" from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import _select_verify_kernel contract = dict( @@ -1468,27 +1510,25 @@ def test_verify_kernel_selection(): num_spec=7, num_tiles=148, num_sms=148, - v2_operands_ok=lambda: True, - two_cta_operands_ok=lambda: True, + operands_ok=lambda: True, ) assert _select_verify_kernel(**contract) == "v2" assert _select_verify_kernel(**{**contract, "num_tiles": 149}) == "two_cta" - assert ( - _select_verify_kernel( - **{**contract, "num_tiles": 149, "two_cta_operands_ok": lambda: False} - ) - == "v2" - ) + # MXFP8 output: the two-CTA kernel at every grid size + assert _select_verify_kernel(**{**contract, "quantize_output": True}) == "two_cta" + for quantize_output in (False, True): + with pytest.raises(ValueError, match="packed, fused-norm verify"): + _select_verify_kernel( + **{**contract, "quantize_output": quantize_output, "operands_ok": lambda: False} + ) for name, value in ( ("packed_token_layout", False), ("fuse_output_norm", False), - ("quantize_output", True), ("bf16_math", True), ("replay_only", True), ("num_spec", 8), - ("v2_operands_ok", lambda: False), ): - assert _select_verify_kernel(**{**contract, name: value}) == "v1", name + assert _select_verify_kernel(**{**contract, name: value, "operands_ok": None}) == "v1", name if __name__ == "__main__": diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py index e159c69db5bf..4298c655c834 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_fused_verify_parity.py @@ -154,8 +154,9 @@ def _make_fused_layer_cache(B, conv_pool, *, num_heads=H): S = W - 1 + M def _conv_cache(section): - cache = torch.zeros(B, S, d, device="cuda", dtype=torch.float32).transpose(-1, -2) - cache[:, :, : W - 1] = conv_pool[:, section * d : (section + 1) * d].float() + # Raw bf16 projection inputs, dim-contiguous. + cache = torch.zeros(B, S, d, device="cuda", dtype=torch.bfloat16).transpose(-1, -2) + cache[:, :, : W - 1] = conv_pool[:, section * d : (section + 1) * d] return cache return SimpleNamespace( @@ -507,12 +508,6 @@ def test_fused_verify_multi_wave_runs_two_cta_kernel(monkeypatch, state_dtype): runtime.finalize_decode_weights() conv_pool, ssm_pool = _make_pools(batch, seed=2) cache = _make_fused_layer_cache(batch, conv_pool) - for name in ("kda_conv_q", "kda_conv_k", "kda_conv_v"): - # The runtime's conv caches are bf16 (dim-contiguous). - src = getattr(cache, name) - dst = torch.empty_like(src.transpose(-1, -2), dtype=torch.bfloat16).transpose(-1, -2) - dst.copy_(src) - setattr(cache, name, dst) worlds = [] for _ in range(2): worlds.append( From 0c47ec51997b0011f9ecedd792120d04d8d68fef Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:02:04 -0700 Subject: [PATCH 5/8] [None][perf] DWDP: cross-step cyclic expert prefetch The DWDP expert prefetch now follows one cyclic schedule over the MoE layers of consecutive forward passes (PrefetchSchedule in modules/dwdp/prefetch_schedule.py) instead of restarting at every pass. MoE order o uses slot o % num_slots; a layer's copy is issued once every earlier layer is issued and its slot was released by the MoE that last read it. The release (consume event) is recorded right after the MoE is enqueued, as before, so each MoE layer is gathered once per pass, and the last MoE layers of a pass already gather the first layers of the next pass: the gather runs through the host gap between executor steps and start_forward enqueues nothing in steady state. This is the default and only protocol; the bound weights are the same, only the copy timing changes. The DwdpManager entry points are renamed to match the weight manager: start_forward / finish_layer (were prefetch_first_layers / record_compute_and_prefetch_next). Correctness: - A copy waits for the consume event of its slot; a bind waits for the prefetch event of its slot, whose last copy is the bound layer's because no other copy is issued into the slot until that layer finishes. - A pass without start_forward (model warmup, the KV-estimation executor) issues its copies on the first bind; the per-step protocol bound whatever the slot held there. - A pass that stops inside the MoE sequence or binds out of order restarts the sequence at that layer (warning at restart counts 1, 2, 4, ...); if a bound layer was abandoned (exception inside its MoE) the device is synchronized before its slot is reused. - After every control action (update_weights, sleep/wakeup) PyExecutor calls DwdpManager.invalidate_prefetch(); the next pass gathers its first layers again. - DWDP now rejects cuda_graph_config at executor creation: a replayed graph would bind slots the host schedule does not track (every DWDP config and test already sets cuda_graph_config=None). Base already gathers each layer once per pass and releases the slot after the MoE (#18644); this adds the cross-step part of the reference's prefetch, combined with the scheduled prefetch of e47e99fa94 (dedup / early-consume / late-release variants and their env switches are not carried: base behaviour already is the dedup + early-release one). Evidence: Rubin (hecate R200), ctx-only DWDP16 conc 128, cross-step 2 slots together with the FP8-KV MLA context and kv_b items: +5.2% req/s (jobs 798859 vs 799151). Pipeline model of the R200 ctx nsys trace (job 797928): cross-step 2 slots -1.1% full step / -1.3% all steps on top of the dedup. GB300 replay (e47e99fa94, real weight manager + SM gather, 15.6 GB per layer): +cross-step -0.1..-0.5% per trace on top of dedup + early release. Tests (GB300, NSC LHR, base image 23bec710e5 + this tree): job 316326 at this commit: tests/unittest/_torch/modules/dwdp 190 passed. New: test_dwdp_prefetch_schedule.py (issue order for 1..12 and 92 layers x 2..4 slots, wrap points, warmup/out-of-order/truncated/abandoned passes, invalidation, randomized RAW/WAR slot model), test_dwdp_prefetch_pipeline.py (event order with mocked streams/events), test_dwdp_prefetch_gpu.py (real streams and copies, delayed MoE or delayed copies, copy engine and SM fabric gather from peer GPUs: every bound weight exact, one copy per layer per pass, stale first layers without invalidation, missing consume detected). A probe in the same image confirms that a stream wait on an event recorded outside a CUDA graph capture fails inside the capture, so DWDP was not capturable before this change either. test_dwdp_prefetch_gpu.py is listed in l0_b200.yml and l0_gb300_multi_gpus.yml (the peer-GPU cases need 2 GPUs); the CPU tests run in l0_cpu through unittest/_torch/modules. OCI AGA job 1369418: test_dwdp_prefetch_gpu.py on 4 GPUs 24 passed. Platform: N (architecture independent). Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../_torch/modules/dwdp/prefetch_schedule.py | 161 ++++++++++ .../_torch/modules/dwdp/weight_manager.py | 181 +++++++---- .../moe/fused_moe/MOE_DEVELOPER_GUIDE.md | 11 +- .../_torch/moe/fused_moe/configurable_moe.py | 2 +- tensorrt_llm/_torch/pyexecutor/dwdp.py | 51 +-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 6 +- .../_torch/pyexecutor/py_executor_creator.py | 7 + .../test_lists/test-db/l0_b200.yml | 3 + .../test-db/l0_gb300_multi_gpus.yml | 3 + .../_torch/modules/dwdp/test_dwdp_manager.py | 112 ++----- .../modules/dwdp/test_dwdp_prefetch_gpu.py | 302 ++++++++++++++++++ .../dwdp/test_dwdp_prefetch_pipeline.py | 165 +++++++--- .../dwdp/test_dwdp_prefetch_schedule.py | 226 +++++++++++++ 13 files changed, 1000 insertions(+), 230 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/dwdp/prefetch_schedule.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py create mode 100644 tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_schedule.py diff --git a/tensorrt_llm/_torch/modules/dwdp/prefetch_schedule.py b/tensorrt_llm/_torch/modules/dwdp/prefetch_schedule.py new file mode 100644 index 000000000000..5207fc55ce6d --- /dev/null +++ b/tensorrt_llm/_torch/modules/dwdp/prefetch_schedule.py @@ -0,0 +1,161 @@ +# 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. +"""Issue order of the DWDP expert prefetch. + +The MoE layers of consecutive forward passes form one cyclic sequence of positions: position +``p`` is MoE order ``p % num_layers`` and uses the prefetch slot ``order % num_slots`` (the slot +its composite VA maps). A position is issued (its copy enqueued) once every earlier position is +issued and its slot no longer holds a position that has not finished. A slot is released right +after the MoE that read it is enqueued, so each MoE layer is gathered once per pass, the copy into +a slot overlaps everything up to the next MoE that reads the slot, and the last MoE layers of a +pass already gather the first layers of the next pass, through the gap between executor steps. + +``PrefetchSchedule`` only does the host bookkeeping. Each MoE layer goes through ``bind`` (before +its MoE reads the slot) and ``finish`` (right after the MoE is enqueued); a pass starts with +``start_forward``. Every call returns a ``ScheduleStep`` that the caller enqueues in field order. +A pass that stops inside the sequence, or a bind out of sequence, restarts the sequence at that +layer; a pass without ``start_forward`` (e.g. a warmup forward) issues its copies on the first +bind. +""" + +from __future__ import annotations + +from typing import List, NamedTuple, Optional + + +class ScheduleStep(NamedTuple): + """What the caller enqueues for one schedule call, in field order. + + Attributes: + synchronize: Synchronize the device first: a layer bound by an unfinished pass may still + be reading its slot, and its release was never recorded. + release: Slot whose consume event is recorded now on the compute stream, if any. + prefetch: MoE orders whose copies are enqueued now, in issue order. + """ + + synchronize: bool + release: Optional[int] + prefetch: List[int] + + +class PrefetchSchedule: + """Host-side issue order of the DWDP expert prefetch over consecutive forward passes. + + Args: + num_layers: Number of MoE layers of one forward pass. + num_slots: Number of prefetch buffer slots; MoE order ``o`` uses slot ``o % num_slots``. + """ + + def __init__(self, num_layers: int, num_slots: int) -> None: + if num_layers <= 0: + raise ValueError(f"num_layers must be positive, got {num_layers}") + if num_slots < 2: + raise ValueError(f"num_slots must be at least 2, got {num_slots}") + self._num_layers = num_layers + self._num_slots = num_slots + # Next position to bind and number of positions issued so far. + self._next_bind = 0 + self._issued = 0 + # Position holding each slot until it finishes (None: free). + self._slot_owner: List[Optional[int]] = [None] * num_slots + # Position bound and not finished yet. + self._bound: Optional[int] = None + self._num_restarts = 0 + + @property + def num_restarts(self) -> int: + """How often a pass broke the sequence (stopped inside it or bound out of order).""" + return self._num_restarts + + def slot(self, order: int) -> int: + """Prefetch slot of MoE order ``order``.""" + return order % self._num_slots + + def start_forward(self) -> ScheduleStep: + """Start a forward pass: issue its first layers unless they are already in flight.""" + if self._bound is None and self._next_bind % self._num_layers == 0: + return ScheduleStep(False, None, self._issue()) + self._num_restarts += 1 + synchronize = self._restart(0) + return ScheduleStep(synchronize, None, self._issue()) + + def bind(self, order: int) -> ScheduleStep: + """Bind MoE order ``order``; its MoE reads the slot only after the returned step. + + A bind out of sequence restarts the sequence at ``order``. + + Raises: + ValueError: If ``order`` is out of range. + """ + self._check_order(order) + synchronize = False + if self._bound is not None or order != self._next_bind % self._num_layers: + self._num_restarts += 1 + synchronize = self._restart(order) + prefetch = self._issue() + if self._issued <= self._next_bind: + raise RuntimeError( + f"DWDP prefetch schedule: MoE order {order} bound before its copy was issued" + ) + self._bound = self._next_bind + return ScheduleStep(synchronize, None, prefetch) + + def finish(self, order: int) -> ScheduleStep: + """Finish the bound MoE order ``order`` (its MoE has just been enqueued). + + Raises: + ValueError: If ``order`` is out of range. + RuntimeError: If ``order`` is not the bound layer. + """ + self._check_order(order) + if self._bound is None or order != self._bound % self._num_layers: + bound = None if self._bound is None else self._bound % self._num_layers + raise RuntimeError( + f"DWDP prefetch schedule: finish of MoE order {order} without a matching bind " + f"(bound={bound})" + ) + self._next_bind = self._bound + 1 + self._bound = None + slot = self.slot(order) + self._slot_owner[slot] = None + return ScheduleStep(False, slot, self._issue()) + + def invalidate(self) -> ScheduleStep: + """Forget every issued prefetch; the next pass issues its first layers again.""" + return ScheduleStep(self._restart(0), None, []) + + def _check_order(self, order: int) -> None: + if not 0 <= order < self._num_layers: + raise ValueError(f"MoE order {order} out of range [0, {self._num_layers})") + + def _restart(self, order: int) -> bool: + """Restart the sequence at ``order``; returns whether a bound layer was abandoned.""" + abandoned = self._bound is not None + self._slot_owner = [None] * self._num_slots + self._bound = None + self._next_bind = order + self._issued = order + return abandoned + + def _issue(self) -> List[int]: + orders: List[int] = [] + while True: + order = self._issued % self._num_layers + slot = self.slot(order) + if self._slot_owner[slot] is not None: + return orders + self._slot_owner[slot] = self._issued + self._issued += 1 + orders.append(order) diff --git a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py index 2ec645f18998..5d15da642c5e 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py @@ -18,24 +18,27 @@ The WeightManager is the runtime scheduler for DWDP expert weight prefetch. It runs INDEPENDENTLY per instance (no cross-instance sync at runtime) and manages: - - Async P2P copies from peer MNNVL handles into WeightBuffer's remote slices - - Bidirectional event sync (prefetch_events, consume_events) with zero CPU blocking + - Async P2P copies from peer MNNVL handles into WeightBuffer's remote slices, + in the order of a ``PrefetchSchedule`` over consecutive forward passes + - Per-slot event sync (prefetch_events, consume_events) with zero CPU blocking - param.data binding for MoE forward pass Design invariants: - - NEVER calls torch.cuda.synchronize() — all sync via stream.wait_event() + - Never calls torch.cuda.synchronize() in steady state — all sync via + stream.wait_event(); only a forward pass abandoned inside a MoE layer is + recovered with a device synchronization - NEVER allocates/frees memory during inference — all buffers pre-allocated - Only copies remote experts — local experts are zero-copy MNNVL - - Double buffer with bidirectional events prevents RAW and WAR hazards + - Per-slot prefetch/consume events prevent RAW and WAR hazards -Timeline visualization: - compute_stream: [forward L3] --wait(prefetch[0])--> [forward L4] --wait(prefetch[1])--> [forward L5] - copy_stream: --wait(consume[0])--> [P2P L4] --wait(consume[1])--> [P2P L5] +Timeline visualization (two slots, MoE layers L3, L4, L5, L6): + compute_stream: --wait(prefetch[0])--> [MoE L3] rec(consume[0]) [attention] + --wait(prefetch[1])--> [MoE L4] rec(consume[1]) [attention] ... + copy_stream: --wait(consume[0])--> [P2P L5] --wait(consume[1])--> [P2P L6] """ from __future__ import annotations -import bisect import logging import traceback from typing import Dict, List, Optional, Tuple @@ -49,6 +52,7 @@ logger = logging.getLogger(__name__) from . import fabric_gather +from .prefetch_schedule import PrefetchSchedule, ScheduleStep from .specs import PeerRanges, lookup_owner from .weight_buffer import WeightBuffer @@ -69,13 +73,23 @@ # number of blocks each group gets. _GatherPlan = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int] +# Number of prefetch buffer slots (``WeightBuffer`` double buffer). +_NUM_SLOTS = 2 + class DWDPWeightManager: """Runtime P2P copy orchestrator for DWDP expert weight prefetch. Manages async P2P copies from peer MNNVL handles into WeightBuffer's - remote slices using a double-buffered event protocol. All synchronization - is done via CUDA events and stream.wait_event() — the CPU is never blocked. + remote slices. A ``PrefetchSchedule`` decides when each copy is issued; + every slot has a prefetch (RAW) and a consume (WAR) event. All + synchronization is done via CUDA events and stream.wait_event() — the CPU + is never blocked. + + Runtime protocol: ``start_forward`` before a forward pass, then per MoE + layer ``wait_and_bind`` before its MoE and ``finish_layer`` right after + the MoE is enqueued; ``invalidate_prefetch`` after the source expert + weights changed. Attributes: weight_buffer: WeightBuffer with composite VA layout for all MoE layers. @@ -101,6 +115,9 @@ class DWDPWeightManager: "_batched_copy_plans", "_use_gather", "_gather_plans", + "_layer_order", + "_schedule", + "_reported_restarts", "_released", ) @@ -117,9 +134,9 @@ def __init__( ) -> None: """Initialize the DWDPWeightManager. - Creates a dedicated copy stream and pre-allocates all CUDA events for - the double-buffer protocol. Both consume_events are recorded immediately - so the first prefetch does not stall waiting for a never-recorded event. + Creates a dedicated 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. Args: weight_buffer: WeightBuffer with composite VA layout for all MoE layers. @@ -146,7 +163,8 @@ def __init__( ``data_ptr()`` are stable after initialization. Raises: - ValueError: If dwdp_rank is out of range or dwdp_size is invalid. + ValueError: If dwdp_rank is out of range, dwdp_size is invalid or + there are no MoE layers. """ if dwdp_size <= 0: raise ValueError(f"dwdp_size must be positive, got {dwdp_size}") @@ -173,13 +191,22 @@ def __init__( device = torch.device("cuda", weight_buffer.device_id) self._copy_stream = torch.cuda.Stream(device=device) - # Double-buffer events: one per buffer slot (0 and 1). + num_slots = _NUM_SLOTS + self._layer_order = {layer: order for order, layer in enumerate(self._moe_layer_indices)} + self._schedule = PrefetchSchedule(len(self._moe_layer_indices), num_slots) + self._reported_restarts = 0 + + # Buffer-slot events: one pair per slot. # prefetch_events[i]: recorded on copy_stream after P2P copy into slot i. # Waited on by compute_stream (RAW — must finish copy before read). - # consume_events[i]: recorded on compute_stream after forward using slot i. + # consume_events[i]: recorded on compute_stream after the MoE using slot i. # Waited on by copy_stream (WAR — must finish read before overwrite). - self._prefetch_events: List[torch.cuda.Event] = [torch.cuda.Event() for _ in range(2)] - self._consume_events: List[torch.cuda.Event] = [torch.cuda.Event() for _ in range(2)] + self._prefetch_events: List[torch.cuda.Event] = [ + torch.cuda.Event() for _ in range(num_slots) + ] + self._consume_events: List[torch.cuda.Event] = [ + torch.cuda.Event() for _ in range(num_slots) + ] self._transport = None # Initialize consume_events as "already signaled" so the first prefetch @@ -250,43 +277,12 @@ def is_moe_layer(self, layer_idx: int) -> bool: """ return layer_idx in self._moe_layer_set - def next_moe_layer(self, layer_idx: int) -> Optional[int]: - """Return the next MoE layer index after layer_idx. - - Uses binary search for O(log n) lookup. - - Args: - layer_idx: Current layer index. - - Returns: - The next MoE layer index strictly greater than layer_idx, - or None if layer_idx is at or past the last MoE layer. - """ - if not self._moe_layer_indices: - return None - pos = bisect.bisect_right(self._moe_layer_indices, layer_idx) - if pos < len(self._moe_layer_indices): - return self._moe_layer_indices[pos] - return None - - def first_moe_layer(self) -> int: - """Return the first MoE layer index. - - Returns: - The smallest MoE layer index. - - Raises: - IndexError: If there are no MoE layers. - """ - if not self._moe_layer_indices: - raise IndexError("No MoE layers configured") - return self._moe_layer_indices[0] - def prefetch_layer(self, layer_idx: int) -> None: """Kick off async P2P copies for all remote expert slices of a layer. - Switches to the copy_stream, waits for the consume_event on the target - buffer slot (WAR hazard — ensures compute has finished reading the slot), + Called for the copies the prefetch schedule issues. Switches to the + copy_stream, waits for the consume_event on the target buffer slot + (WAR hazard — ensures compute has finished reading the slot), then issues P2P copies from peer MNNVL views into the WeightBuffer's remote slices. Finally, records a prefetch_event so the compute stream can wait for the copy to complete. @@ -522,15 +518,25 @@ def _prefetch_layer_gather(self, layer_idx: int) -> None: dst_ptrs, src_ptrs, lines, row_starts, blocks, fabric_gather.THREADS ) + def start_forward(self) -> None: + """Prefetch the first MoE layers of the coming forward pass. + + Enqueues nothing when the previous forward pass already prefetched + them. + """ + self._apply(self._schedule.start_forward()) + def wait_and_bind(self, backend_module: torch.nn.Module, layer_idx: int) -> None: """Wait for prefetch completion and bind weight tensors to the module. - On the compute stream: + Issues the layer's copy first if nothing did (a forward pass without + ``start_forward``, or one that restarts the prefetch sequence). Then, + on the compute stream: 1. Wait for prefetch_events[buf_idx] (RAW — copy must finish before read). 2. Bind each weight's full tensor (local zero-copy + remote P2P data) as the param.data of the backend module's corresponding attribute. - The matching WAR signal is deliberately recorded by ``record_compute`` + The matching WAR signal is deliberately recorded by ``finish_layer`` after the MoE kernel has been enqueued. Recording it here would place the event before the kernel on the compute stream and could let the copy stream overwrite a slot while that kernel is still reading it. @@ -545,10 +551,7 @@ def wait_and_bind(self, backend_module: torch.nn.Module, layer_idx: int) -> None KeyError: If layer_idx is not a valid MoE layer. AttributeError: If backend_module lacks a required weight attribute. """ - if layer_idx not in self._moe_layer_set: - raise KeyError( - f"Layer {layer_idx} is not a MoE layer. Valid layers: {self._moe_layer_indices}" - ) + self._apply(self._schedule.bind(self._order_of(layer_idx))) buf_idx = self._weight_buffer.buffer_index_for_layer(layer_idx) compute_stream = torch.cuda.current_stream( @@ -556,7 +559,9 @@ def wait_and_bind(self, backend_module: torch.nn.Module, layer_idx: int) -> None ) # RAW hazard: wait until P2P copy into this slot is complete - # before the forward pass reads the data. + # before the forward pass reads the data. The schedule issues no + # other copy into the slot until this layer finishes, so the last + # recorded prefetch event of the slot is this layer's. compute_stream.wait_event(self._prefetch_events[buf_idx]) # Bind full [num_experts, ...] tensors to the backend module's parameters. @@ -569,29 +574,65 @@ def wait_and_bind(self, backend_module: torch.nn.Module, layer_idx: int) -> None logger.debug(f"[DWDPWeightManager] Bound weights: layer={layer_idx}, buf_slot={buf_idx}") - def record_compute(self, layer_idx: int) -> None: - """Record that the current layer has finished consuming its buffer slot. + def finish_layer(self, layer_idx: int) -> None: + """Release the layer's buffer slot and issue the copies this unblocks. This method must be called after the layer's MoE kernel is enqueued on the current compute stream. CUDA stream ordering then guarantees that the consume event cannot complete until the kernel has stopped reading - the slot. The copy stream waits on this event before reusing the slot - for the layer two positions ahead in the MoE sequence. + the slot. The copy stream waits on this event before reusing the slot; + the last MoE layers of a forward pass issue the copies of the first + layers of the next one. Args: layer_idx: MoE layer whose compute was just enqueued. + + Raises: + KeyError: If layer_idx is not a valid MoE layer. + RuntimeError: If layer_idx is not the bound layer. """ - if layer_idx not in self._moe_layer_set: + self._apply(self._schedule.finish(self._order_of(layer_idx))) + + def invalidate_prefetch(self) -> None: + """Drop the copies issued ahead of the next forward pass. + + Call after the source expert weights changed; the next forward pass + gathers its first layers again. + """ + self._apply(self._schedule.invalidate()) + + def _order_of(self, layer_idx: int) -> int: + order = self._layer_order.get(layer_idx) + if order is None: raise KeyError( f"Layer {layer_idx} is not a MoE layer. Valid layers: {self._moe_layer_indices}" ) + return order - buf_idx = self._weight_buffer.buffer_index_for_layer(layer_idx) - compute_stream = torch.cuda.current_stream( - torch.device("cuda", self._weight_buffer.device_id) - ) - self._consume_events[buf_idx].record(compute_stream) - logger.debug(f"[DWDPWeightManager] Compute recorded: layer={layer_idx}, buf_slot={buf_idx}") + def _apply(self, step: ScheduleStep) -> None: + """Enqueue one schedule step: device sync if needed, slot release, then copies.""" + if step.synchronize: + logger.warning( + "[DWDPWeightManager] Prefetch sequence restarted after an unfinished forward " + "pass; synchronizing the device before reusing its buffer slots" + ) + torch.cuda.synchronize(self._weight_buffer.device_id) + if step.release is not None: + compute_stream = torch.cuda.current_stream( + torch.device("cuda", self._weight_buffer.device_id) + ) + self._consume_events[step.release].record(compute_stream) + restarts = self._schedule.num_restarts + if restarts != self._reported_restarts: + self._reported_restarts = restarts + # Power-of-two counts only: a model that keeps breaking the sequence still shows up. + if restarts & (restarts - 1) == 0: + logger.warning( + f"[DWDPWeightManager] Prefetch sequence restarted ({restarts} restart(s) so " + f"far): a forward pass did not run the MoE layers in order" + ) + for order in step.prefetch: + self.prefetch_layer(self._moe_layer_indices[order]) def release(self) -> None: """Synchronize outstanding work and release all owned VMM resources.""" 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 57063ebb6368..999626e8e30f 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/moe/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -53,7 +53,7 @@ def forward_impl(self, x, router_logits, ...): self.dwdp_manager.wait_and_bind(self.backend, self.layer_idx) outputs = self.scheduler.forward(x, router_logits, ...) if self.enable_dwdp: - self.dwdp_manager.record_compute_and_prefetch_next(self.layer_idx) + self.dwdp_manager.finish_layer(self.layer_idx) self.repeat_idx = (self.repeat_idx + 1) % self.repeat_count return outputs ``` @@ -79,12 +79,17 @@ Executor construction → after weight loading: create VMM transport + composite VA + page pool PyExecutor forward start - → DwdpManager.prefetch_first_layers() for the first two MoE layers + → DwdpManager.start_forward(): first MoE layers, unless the previous + pass already prefetched them → ConfigurableMoE.wait_and_bind(composite VA) → MoEScheduler.forward(... ordinary CuTeDSL kernel ...) → main stream joins scheduler auxiliary work → DwdpManager records current-slot consumption - → prefetch the MoE layer two positions ahead + → prefetch the next MoE layer in the cyclic sequence whose slot is now free + (the last layers of a pass prefetch the first layers of the next pass) + +Control action (weight update, sleep/wakeup) + → DwdpManager.invalidate_prefetch() Executor shutdown or construction rollback → detach model references diff --git a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py index 892c8e15c570..1e213d342e69 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py +++ b/tensorrt_llm/_torch/moe/fused_moe/configurable_moe.py @@ -736,7 +736,7 @@ def forward_impl( # Owned at the wrapper because schedulers must not run it twice (external-comm # might enter via single- or multi-chunk paths). if self.enable_dwdp: - self.dwdp_manager.record_compute_and_prefetch_next(self.layer_idx) + self.dwdp_manager.finish_layer(self.layer_idx) # EPLB repeat counter: advance once per forward, regardless of chunk count. # Schedulers are forbidden from rotating ``repeat_idx`` themselves. diff --git a/tensorrt_llm/_torch/pyexecutor/dwdp.py b/tensorrt_llm/_torch/pyexecutor/dwdp.py index 2bcbebd92e3f..476c40d7112d 100644 --- a/tensorrt_llm/_torch/pyexecutor/dwdp.py +++ b/tensorrt_llm/_torch/pyexecutor/dwdp.py @@ -21,11 +21,11 @@ a TCPStore rendezvous for independently launched workers * Layer index registration (SSOT: ``add_layer`` → ``_registered_layers``) plus runtime entry points that forward to ``DWDPWeightManager`` - (``prefetch_first_layers``, ``wait_and_bind``, - ``record_compute_and_prefetch_next``) + (``start_forward``, ``wait_and_bind``, + ``finish_layer``, ``invalidate_prefetch``) The heavy lifting — MNNVL handle allocation, composite VA layout, -double-buffer scheduling — lives in ``tensorrt_llm._torch.modules.dwdp`` +prefetch scheduling — lives in ``tensorrt_llm._torch.modules.dwdp`` and is wired up by ``setup_dwdp()`` during ``setup(model)``. """ @@ -76,8 +76,8 @@ class DwdpManager: 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 (``prefetch_first_layers``, - ``wait_and_bind``, ``record_compute_and_prefetch_next``) forward to the + After ``setup()`` returns, the runtime methods (``start_forward``, + ``wait_and_bind``, ``finish_layer``) forward to the ``DWDPWeightManager`` instance produced by ``setup_dwdp``. """ @@ -291,15 +291,15 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: # Runtime entry points (forwarded to DWDPWeightManager) # ------------------------------------------------------------------ - def prefetch_first_layers(self) -> None: - """Warm-up: kick off prefetch for the first MoE layers (double-buffer depth).""" + 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). + """ if self._weight_manager is None: raise RuntimeError("DwdpManager.setup() has not been called yet") - first = self._weight_manager.first_moe_layer() - self._weight_manager.prefetch_layer(first) - second = self._weight_manager.next_moe_layer(first) - if second is not None: - self._weight_manager.prefetch_layer(second) + self._weight_manager.start_forward() def wait_and_bind(self, backend_module: nn.Module, layer_idx: int) -> None: """Wait for prefetch and bind full-tensor views onto backend weights.""" @@ -307,20 +307,21 @@ def wait_and_bind(self, backend_module: nn.Module, layer_idx: int) -> None: raise RuntimeError("DwdpManager.setup() has not been called yet") self._weight_manager.wait_and_bind(backend_module, layer_idx) - def record_compute_and_prefetch_next(self, layer_idx: int) -> None: - """Record compute and maintain the two-layer prefetch window. + def finish_layer(self, layer_idx: int) -> None: + """Release the layer's buffer slot and prefetch the layers it unblocks. - ``prefetch_first_layers`` warms the first two MoE layers. After each - layer's kernel has been enqueued, record consumption of its current - buffer slot and prefetch the layer two positions ahead. This avoids - duplicate prefetch of the second layer and prevents reuse of a slot - before the current kernel has finished reading it. + Called after the layer's MoE kernel has been enqueued, so the slot's + consume event covers that kernel. """ if self._weight_manager is None: raise RuntimeError("DwdpManager.setup() has not been called yet") - self._weight_manager.record_compute(layer_idx) - next_idx = self._weight_manager.next_moe_layer(layer_idx) - if next_idx is not None: - next_next_idx = self._weight_manager.next_moe_layer(next_idx) - if next_next_idx is not None: - self._weight_manager.prefetch_layer(next_next_idx) + self._weight_manager.finish_layer(layer_idx) + + def invalidate_prefetch(self) -> None: + """Drop the prefetches issued ahead of the next forward pass. + + Call after the expert weights may have changed (e.g. a weight update); + the next forward pass gathers its first layers again. + """ + if self._weight_manager is not None: + self._weight_manager.invalidate_prefetch() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e0502b4cd181..5a5b810339b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5365,7 +5365,7 @@ def _executor_loop(self) -> None: gpu_forward_start, gpu_forward_end) as fwd_timing: if self.dwdp_manager is not None: - self.dwdp_manager.prefetch_first_layers() + self.dwdp_manager.start_forward() batch_outputs = self._forward_step(scheduled_batch) # None: _forward_step hit an exception and @@ -5651,6 +5651,10 @@ def _handle_control_request(self): self.control_action_done.wait() self.control_action_done.clear() self._active_control_id = None + if self.dwdp_manager is not None: + # The action may have changed the expert weights that the last + # forward pass already prefetched for the next one. + self.dwdp_manager.invalidate_prefetch() logger.debug("[control_action] control request finished") def _sync_and_process_resource_governor_queue(self): diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 5d6c7e1a7a43..f0e465944c5a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -761,6 +761,13 @@ def _create_py_executor( "tp_size=1 (disaggregated context worker) or " "enable_attention_dp=True (aggregated serving), but got " f"tp_size={mapping.tp_size}, dp_size={mapping.dp_size}.") + # The host issues the expert prefetch as the MoE layers run, and the + # last layers of a pass already prefetch the next pass; a replayed + # CUDA graph would bind buffer slots the host schedule does not track. + if llm_args.cuda_graph_config is not None: + raise ValueError( + "DWDP does not support CUDA graphs; set cuda_graph_config=None." + ) dwdp_manager = DwdpManager(config=llm_args.dwdp_config, dist=dist, mapping=mapping) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index b597efc96bc0..7b1f78d23052 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -199,6 +199,9 @@ l0_b200: - unittest/_torch/modules/test_low_m_gemm.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules + # 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 # ------------- 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 fe2c8c5c0271..cd7500c03d99 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 @@ -55,6 +55,9 @@ l0_gb300_multi_gpus: - unittest/_torch/modules/test_group_rmn_norm.py - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules + # DWDP expert prefetch on real streams and copies, incl. the gather from + # peer GPUs. + - unittest/_torch/modules/dwdp/test_dwdp_prefetch_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_dwdp_manager.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py index 9aaf54ac6463..1cca0d9a9957 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py @@ -357,42 +357,11 @@ def test_setup_forwards_layer_indices(self, mock_comm_world, _rank, mock_setup): # Runtime forwards require setup() # ------------------------------------------------------------------ - def test_prefetch_first_layers_forwards(self, mock_comm_world, _rank, mock_setup): + def _set_up_manager(self, mock_comm_world, mock_setup): mock_comm_world.Create_group.return_value = MagicMock() mock_comm_world.group.Incl.return_value = MagicMock() mock_comm_world.Get_size.return_value = 1024 fake_wm = MagicMock() - fake_wm.first_moe_layer.return_value = 3 - fake_wm.next_moe_layer.return_value = 5 - mock_setup.return_value = fake_wm - mgr = DwdpManager( - config=_make_config(), - dist=MagicMock(spec=MPIDist), - mapping=_make_mapping(), - ) - with patch( - "tensorrt_llm._torch.pyexecutor.dwdp.torch.cuda.current_device", - return_value=0, - ): - mgr.setup(MagicMock()) - mgr.prefetch_first_layers() - # First MoE layer 3 + ping-pong depth -> also prefetch next (5) - self.assertEqual( - fake_wm.prefetch_layer.call_args_list, - [ - unittest.mock.call(3), - unittest.mock.call(5), - ], - ) - - def test_record_compute_and_prefetch_next_schedules_next_next( - self, mock_comm_world, _rank, mock_setup - ): - mock_comm_world.Create_group.return_value = MagicMock() - mock_comm_world.group.Incl.return_value = MagicMock() - mock_comm_world.Get_size.return_value = 1024 - fake_wm = MagicMock() - fake_wm.next_moe_layer.side_effect = lambda layer_idx: {7: 9, 9: 11}.get(layer_idx) mock_setup.return_value = fake_wm mgr = DwdpManager( config=_make_config(), @@ -404,76 +373,33 @@ def test_record_compute_and_prefetch_next_schedules_next_next( return_value=0, ): mgr.setup(MagicMock()) + return mgr, fake_wm - mgr.record_compute_and_prefetch_next(7) - fake_wm.record_compute.assert_called_once_with(7) - self.assertEqual( - fake_wm.next_moe_layer.call_args_list, - [unittest.mock.call(7), unittest.mock.call(9)], - ) - fake_wm.prefetch_layer.assert_called_once_with(11) + def test_start_forward_forwards(self, mock_comm_world, _rank, mock_setup): + mgr, fake_wm = self._set_up_manager(mock_comm_world, mock_setup) + mgr.start_forward() + fake_wm.start_forward.assert_called_once_with() - def test_record_compute_and_prefetch_next_last_layer_is_noop( - self, mock_comm_world, _rank, mock_setup - ): - mock_comm_world.Create_group.return_value = MagicMock() - mock_comm_world.group.Incl.return_value = MagicMock() - mock_comm_world.Get_size.return_value = 1024 - fake_wm = MagicMock() - fake_wm.next_moe_layer.return_value = None - mock_setup.return_value = fake_wm - mgr = DwdpManager( - config=_make_config(), - dist=MagicMock(spec=MPIDist), - mapping=_make_mapping(), - ) - with patch( - "tensorrt_llm._torch.pyexecutor.dwdp.torch.cuda.current_device", - return_value=0, - ): - mgr.setup(MagicMock()) + def test_finish_layer_forwards(self, mock_comm_world, _rank, mock_setup): + mgr, fake_wm = self._set_up_manager(mock_comm_world, mock_setup) + mgr.finish_layer(7) + fake_wm.finish_layer.assert_called_once_with(7) - mgr.record_compute_and_prefetch_next(60) - fake_wm.record_compute.assert_called_once_with(60) - fake_wm.prefetch_layer.assert_not_called() + def test_invalidate_prefetch_forwards(self, mock_comm_world, _rank, mock_setup): + mgr, fake_wm = self._set_up_manager(mock_comm_world, mock_setup) + mgr.invalidate_prefetch() + fake_wm.invalidate_prefetch.assert_called_once_with() - def test_prefetch_pipeline_schedules_each_layer_once(self, mock_comm_world, _rank, mock_setup): - mock_comm_world.Create_group.return_value = MagicMock() - mock_comm_world.group.Incl.return_value = MagicMock() - mock_comm_world.Get_size.return_value = 1024 - fake_wm = MagicMock() - fake_wm.first_moe_layer.return_value = 3 - fake_wm.next_moe_layer.side_effect = lambda layer_idx: { - 3: 5, - 5: 7, - 7: 9, - 9: None, - }.get(layer_idx) - mock_setup.return_value = fake_wm + def test_invalidate_prefetch_before_setup_is_noop(self, mock_comm_world, _rank, _setup): + _configure_mock_comm_world(mock_comm_world) mgr = DwdpManager( config=_make_config(), dist=MagicMock(spec=MPIDist), mapping=_make_mapping(), ) - with patch( - "tensorrt_llm._torch.pyexecutor.dwdp.torch.cuda.current_device", - return_value=0, - ): - mgr.setup(MagicMock()) - - mgr.prefetch_first_layers() - for layer_idx in (3, 5, 7, 9): - mgr.record_compute_and_prefetch_next(layer_idx) - - self.assertEqual( - fake_wm.prefetch_layer.call_args_list, - [ - unittest.mock.call(3), - unittest.mock.call(5), - unittest.mock.call(7), - unittest.mock.call(9), - ], - ) + mgr.invalidate_prefetch() + with self.assertRaises(RuntimeError): + mgr.start_forward() def test_wait_and_bind_forwards(self, mock_comm_world, _rank, mock_setup): mock_comm_world.Create_group.return_value = MagicMock() diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py new file mode 100644 index 000000000000..ad3b3624ed37 --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py @@ -0,0 +1,302 @@ +# 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. +"""DWDP expert prefetch on real streams, events and copies. + +A stand-in for ``WeightBuffer`` backs every layer's remote experts with its slot's storage, and +the peers' experts are plain tensors (on this GPU, or on the other GPUs of the node for the +``trtllm::dwdp_fabric_gather`` kernel). Each "MoE" multiplies activations of a per-pass token +count with the bound remote experts after a delay; delaying the MoE exposes a missing WAR wait +(the next copy overwrites the slot first), delaying the copies exposes a missing RAW wait. Every +output is checked against the source weights of its pass. +""" + +import time +from typing import Dict, List, Optional, Tuple +from unittest.mock import patch + +import pytest +import torch + +from tensorrt_llm._torch.modules.dwdp import fabric_gather +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") + +NUM_EXPERTS = 8 # experts per rank +DIM = 4096 +DTYPE = torch.float64 +DELAY_CYCLES = 400_000 + + +class _SlotBuffer: + """Stand-in for ``WeightBuffer`` of DWDP rank 0: 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, num_ranks: int) -> None: + self.device_id = torch.cuda.current_device() + self.num_slots = num_slots + self._layers = layers + self.slots = [ + torch.full((num_ranks * NUM_EXPERTS, DIM), -1.0, dtype=DTYPE, 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=DTYPE, device="cuda"), requires_grad=False) + + +class _AbortPass(Exception): + pass + + +class _Harness: + """Drives the facade the way the executor loop and ``ConfigurableMoE`` do.""" + + def __init__( + self, layers: List[int], num_slots: int, peer_devices: Optional[List[int]] = None + ) -> None: + self.layers = layers + peers = [torch.cuda.current_device()] if peer_devices is None else peer_devices + self.num_ranks = len(peers) + 1 + self.buffer = _SlotBuffer(layers, num_slots, self.num_ranks) + # src[layer][r - 1]: the experts of DWDP rank r (r >= 1) on its device. + self.src: Dict[int, List[torch.Tensor]] = {} + for layer in layers: + self.src[layer] = [ + torch.randint(0, 64, (NUM_EXPERTS, DIM), device=f"cuda:{dev}", dtype=torch.int32) + .to(DTYPE) + .contiguous() + for dev in peers + ] + peer_views = { + (rank, layer, "w"): self.src[layer][rank - 1] + for layer in layers + for rank in range(1, self.num_ranks) + } + self.wm = DWDPWeightManager( + weight_buffer=self.buffer, + peer_views=peer_views, + peer_ranges=[ + (rank * NUM_EXPERTS, (rank + 1) * NUM_EXPERTS) for rank in range(self.num_ranks) + ], + moe_layer_indices=layers, + weight_names=["w"], + dwdp_rank=0, + dwdp_size=self.num_ranks, + ) + self.mgr = object.__new__(DwdpManager) + self.mgr._weight_manager = self.wm + self.backend = _Backend() + self.compute_stream = torch.cuda.Stream() + self.copies: List[int] = [] + self.results: List[Tuple[int, torch.Tensor, torch.Tensor, torch.Tensor]] = [] + + def snapshot(self) -> Dict[int, torch.Tensor]: + """Remote experts of every layer as rank 0 must see them, on this GPU.""" + return { + layer: torch.cat([t.to(torch.cuda.current_device()) for t in tensors]) + for layer, tensors in self.src.items() + } + + def run_pass( + self, + tokens: int, + moe_delay: int, + expected: Dict[int, torch.Tensor], + start: bool = True, + abort_after_bind: Optional[int] = None, + ) -> None: + x = torch.randint(0, 4, (tokens, DIM), device="cuda", dtype=torch.int32).to(DTYPE) + if start: + self.mgr.start_forward() + self.compute_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.compute_stream): + for layer in self.layers: + self.mgr.wait_and_bind(self.backend, layer) + if moe_delay: + torch.cuda._sleep(moe_delay) + out = x @ self.backend.w[NUM_EXPERTS:].T + self.results.append((layer, x, out, expected[layer])) + if layer == abort_after_bind: + raise _AbortPass() + self.mgr.finish_layer(layer) + torch.cuda.current_stream().wait_stream(self.compute_stream) + + def mismatches(self) -> List[int]: + torch.cuda.synchronize() + bad = [ + layer for layer, x, out, weights in self.results if not torch.equal(out, x @ weights.T) + ] + self.results.clear() + return bad + + +def _logged(method_name: str, delay: int, log: List[int]): + """Wrap a copy method: first stall the copy stream for ``delay`` cycles, log the layer.""" + original = getattr(DWDPWeightManager, method_name) + + def copy(self, layer_idx): + if delay: + torch.cuda._sleep(delay) + log.append(layer_idx) + original(self, layer_idx) + + return copy + + +def _drive(h: _Harness, moe_delay: int, num_slots: int) -> None: + """Warmup pass, steps of varying size with idle gaps, an aborted pass and a weight update.""" + layers = h.layers + weights = h.snapshot() + # Warmup forward that never calls start_forward; its last layers already gather the + # first layers of the next pass. + h.run_pass(5, moe_delay, weights, start=False) + assert h.mismatches() == [] + assert h.copies[: len(layers)] == layers + assert h.copies[len(layers) :] == layers[: min(num_slots, len(layers))] + + # Steps with different token counts, idle gaps and unrelated GPU work in between. + for step, tokens in enumerate([16, 3, 64, 1]): + del h.copies[:] + h.run_pass(tokens, moe_delay, weights) + assert h.mismatches() == [], f"step {step}" + # Every layer is gathered exactly once per pass. + assert sorted(h.copies) == sorted(layers), f"step {step}" + time.sleep(0.01) + torch.cuda._sleep(DELAY_CYCLES) + + # A pass that raises inside MoE(layer 3): the next pass restarts the sequence. + with pytest.raises(_AbortPass): + h.run_pass(8, moe_delay, weights, abort_after_bind=layers[2]) + torch.cuda.current_stream().wait_stream(h.compute_stream) + assert h.mismatches() == [] + h.run_pass(8, moe_delay, weights) + assert h.mismatches() == [] + + # Weight update between steps (e.g. update_weights): the copies issued by the last pass hold + # the old weights until invalidate_prefetch drops them. + torch.cuda.synchronize() + for tensors in h.src.values(): + for t in tensors: + t.add_(1000.0) + weights = h.snapshot() + h.mgr.invalidate_prefetch() + h.run_pass(7, moe_delay, weights) + assert h.mismatches() == [] + h.run_pass(2, moe_delay, weights) + assert h.mismatches() == [] + + +@pytest.mark.parametrize("num_layers", [4, 5]) +@pytest.mark.parametrize("num_slots", [2]) +@pytest.mark.parametrize("slow", ["moe", "copy"]) +def test_bound_weights_across_passes(monkeypatch, num_layers, num_slots, slow): + monkeypatch.delenv("TRTLLM_DWDP_PREFETCH_GATHER", raising=False) + layers = [2 * i + 1 for i in range(num_layers)] # non-contiguous decoder indices + h = _Harness(layers, num_slots) + moe_delay = DELAY_CYCLES if slow == "moe" else 0 + copy_delay = DELAY_CYCLES if slow == "copy" else 0 + copy = _logged("_prefetch_layer_per_slice", copy_delay, h.copies) + with patch.object(DWDPWeightManager, "_prefetch_layer_per_slice", copy): + _drive(h, moe_delay, num_slots) + + +def test_without_invalidate_first_layers_are_stale(monkeypatch): + """The prefetch issued at the end of a pass holds the weights of that moment.""" + monkeypatch.delenv("TRTLLM_DWDP_PREFETCH_GATHER", raising=False) + layers = [1, 2, 3, 4] + h = _Harness(layers, 2) + weights = h.snapshot() + h.run_pass(4, 0, weights) + assert h.mismatches() == [] + torch.cuda.synchronize() + for tensors in h.src.values(): + for t in tensors: + t.add_(1000.0) + h.run_pass(4, 0, h.snapshot()) + assert h.mismatches() == layers[:2] + + +def test_missing_war_wait_is_detected(monkeypatch): + """Sensitivity check: without the consume events the next copies overwrite slots in use.""" + monkeypatch.delenv("TRTLLM_DWDP_PREFETCH_GATHER", raising=False) + layers = [1, 2, 3, 4, 5, 6] + h = _Harness(layers, 2) + weights = h.snapshot() + original = DWDPWeightManager._apply + + def without_consume(self, step): + original(self, step._replace(release=None)) + + with patch.object(DWDPWeightManager, "_apply", without_consume): + for _ in range(3): + h.run_pass(4, DELAY_CYCLES, weights) + assert h.mismatches() != [] + + +def _enable_peer_access(devices: List[int]) -> None: + from cuda.bindings import runtime as cudart + + current = torch.cuda.current_device() + for dev in devices: + if not torch.cuda.can_device_access_peer(current, dev): + pytest.skip(f"GPU {current} cannot access GPU {dev}") + (err,) = cudart.cudaDeviceEnablePeerAccess(dev, 0) + if err not in ( + cudart.cudaError_t.cudaSuccess, + cudart.cudaError_t.cudaErrorPeerAccessAlreadyEnabled, + ): + pytest.skip(f"cudaDeviceEnablePeerAccess({dev}) failed: {err}") + + +@pytest.mark.parametrize("num_slots", [2]) +@pytest.mark.parametrize("slow", ["moe", "copy"]) +def test_fabric_gather_from_peer_gpus(monkeypatch, num_slots, slow): + """The production transport: the SM gather kernel pulls the remote experts from peer GPUs.""" + if torch.cuda.device_count() < 2: + pytest.skip("needs at least 2 GPUs") + monkeypatch.setenv("TRTLLM_DWDP_PREFETCH_GATHER", "1") + # The gather sets these for the CuTeDSL kernels unless pinned; pin them to keep the env clean. + monkeypatch.setenv("TRTLLM_CUTEDSL_DENSE_GEMM_SCHEDULER", "clc_dynamic") + monkeypatch.setenv("TLLM_ROUTING_DISABLE_COOP_LAUNCH", "1") + if fabric_gather.gather_op() is None: + pytest.skip("trtllm::dwdp_fabric_gather is not registered") + torch.cuda.set_device(0) + peers = list(range(1, min(torch.cuda.device_count(), 4))) + _enable_peer_access(peers) + layers = [2 * i + 1 for i in range(5)] + h = _Harness(layers, num_slots, peer_devices=peers) + assert h.wm._use_gather + moe_delay = DELAY_CYCLES if slow == "moe" else 0 + copy = _logged("_prefetch_layer_gather", DELAY_CYCLES if slow == "copy" else 0, h.copies) + with patch.object(DWDPWeightManager, "_prefetch_layer_gather", copy): + _drive(h, moe_delay, num_slots) 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 2d5219e351e2..f51ece0ac753 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_pipeline.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_pipeline.py @@ -12,37 +12,60 @@ # 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-mocked tests for the DWDP double-buffer event protocol.""" +"""CPU-mocked tests for the DWDP prefetch event protocol.""" from contextlib import nullcontext +from typing import List, Tuple from unittest.mock import MagicMock, patch +import pytest import torch import torch.nn as nn from tensorrt_llm._torch.modules.dwdp.weight_manager import DWDPWeightManager +pytestmark = pytest.mark.cpu_only + +LAYERS = [3, 5, 7] + + +class _Event: + def __init__(self, name: str, log: List[Tuple[str, ...]]) -> None: + self.name = name + self._log = log + + def record(self, stream) -> None: + self._log.append(("record", self.name, stream.name)) + + +class _Stream: + def __init__(self, name: str, log: List[Tuple[str, ...]]) -> None: + self.name = name + self._log = log + + def wait_event(self, event: _Event) -> None: + self._log.append(("wait", self.name, event.name)) + class TestDWDPPrefetchPipeline: def setup_method(self): + self.log: List[Tuple[str, ...]] = [] self.weight_buffer = MagicMock() self.weight_buffer.device_id = 0 self.weight_buffer.local_start = 0 self.weight_buffer.local_end = 2 - self.weight_buffer.buffer_index_for_layer.side_effect = { - 3: 0, - 5: 1, - 7: 0, - }.__getitem__ - - self.copy_stream = MagicMock(name="copy_stream") - self.compute_stream = MagicMock(name="compute_stream") - self.events = [ - MagicMock(name="prefetch_0"), - MagicMock(name="prefetch_1"), - MagicMock(name="consume_0"), - MagicMock(name="consume_1"), - ] + self.weight_buffer.buffer_index_for_layer.side_effect = lambda layer: ( + LAYERS.index(layer) % 2 + ) + self.weight_buffer.get_full_tensor.side_effect = lambda layer, name: torch.tensor( + [float(layer)] + ) + + self.copy_stream = _Stream("copy", self.log) + self.compute_stream = _Stream("compute", self.log) + events = iter( + [_Event(name, self.log) for name in ("prefetch0", "prefetch1", "consume0", "consume1")] + ) self.patches = [ patch( "tensorrt_llm._torch.modules.dwdp.weight_manager.torch.cuda.Stream", @@ -50,7 +73,7 @@ def setup_method(self): ), patch( "tensorrt_llm._torch.modules.dwdp.weight_manager.torch.cuda.Event", - side_effect=self.events, + side_effect=lambda: next(events), ), patch( "tensorrt_llm._torch.modules.dwdp.weight_manager.torch.cuda.current_stream", @@ -60,6 +83,15 @@ def setup_method(self): "tensorrt_llm._torch.modules.dwdp.weight_manager.torch.cuda.stream", return_value=nullcontext(), ), + patch( + "tensorrt_llm._torch.modules.dwdp.weight_manager.torch.cuda.synchronize", + side_effect=lambda device: self.log.append(("synchronize",)), + ), + patch.object( + DWDPWeightManager, + "_prefetch_layer_per_slice", + lambda _self, layer: self.log.append(("copy", layer)), + ), ] for cuda_patch in self.patches: cuda_patch.start() @@ -68,38 +100,97 @@ def setup_method(self): weight_buffer=self.weight_buffer, peer_views={}, peer_ranges=[(0, 2), (2, 4)], - moe_layer_indices=[3, 5, 7], + moe_layer_indices=LAYERS, weight_names=["gate_up_proj", "down_proj"], dwdp_rank=0, dwdp_size=2, ) - for event in self.events: - event.reset_mock() + self.backend = nn.Module() + self.backend.gate_up_proj = nn.Parameter(torch.empty(1), requires_grad=False) + self.backend.down_proj = nn.Parameter(torch.empty(1), requires_grad=False) + self.log.clear() def teardown_method(self): self.manager._released = True for cuda_patch in reversed(self.patches): cuda_patch.stop() - def test_wait_and_bind_waits_for_copy_without_signalling_consumption(self): - full_tensors = { - "gate_up_proj": torch.tensor([1.0]), - "down_proj": torch.tensor([2.0]), - } - self.weight_buffer.get_full_tensor.side_effect = lambda _layer_idx, name: full_tensors[name] - backend = nn.Module() - backend.gate_up_proj = nn.Parameter(torch.empty(1), requires_grad=False) - backend.down_proj = nn.Parameter(torch.empty(1), requires_grad=False) + @staticmethod + def _copy(layer: int, slot: int) -> List[Tuple[str, ...]]: + return [ + ("wait", "copy", f"consume{slot}"), + ("copy", layer), + ("record", f"prefetch{slot}", "copy"), + ] + + def _bind(self, layer: int) -> None: + self.manager.wait_and_bind(self.backend, layer) + assert self.backend.gate_up_proj.data.item() == float(layer) + assert self.backend.down_proj.data.item() == float(layer) + + def test_event_order_across_passes(self): + self.manager.start_forward() + assert self.log == self._copy(3, 0) + self._copy(5, 1) + + # Binding waits for the copy and does not release the slot. + self.log.clear() + self._bind(3) + assert self.log == [("wait", "compute", "prefetch0")] + + # The release is recorded after the MoE was enqueued, then the copy two layers ahead + # waits for it. + self.log.clear() + self.manager.finish_layer(3) + assert self.log == [("record", "consume0", "compute")] + self._copy(7, 0) + + self.log.clear() + self._bind(5) + self.manager.finish_layer(5) + # Layer 3 of the next pass needs slot 0, still held by layer 7. + assert self.log == [ + ("wait", "compute", "prefetch1"), + ("record", "consume1", "compute"), + ] + + # The last MoE layer of the pass prefetches the first layers of the next one. + self.log.clear() + self._bind(7) + self.manager.finish_layer(7) + assert self.log == [ + ("wait", "compute", "prefetch0"), + ("record", "consume0", "compute"), + ] + self._copy(3, 0) + self._copy(5, 1) + + self.log.clear() + self.manager.start_forward() + assert self.log == [] + self._bind(3) + assert self.log == [("wait", "compute", "prefetch0")] - self.manager.wait_and_bind(backend, 3) + def test_bind_without_start_forward_issues_its_copy_first(self): + self._bind(3) + assert self.log == self._copy(3, 0) + self._copy(5, 1) + [("wait", "compute", "prefetch0")] - self.compute_stream.wait_event.assert_called_once_with(self.events[0]) - self.events[2].record.assert_not_called() - assert backend.gate_up_proj.data.data_ptr() == full_tensors["gate_up_proj"].data_ptr() - assert backend.down_proj.data.data_ptr() == full_tensors["down_proj"].data_ptr() + def test_abandoned_pass_synchronizes_before_reusing_slots(self): + self.manager.start_forward() + self._bind(3) + self.log.clear() + self.manager.start_forward() + assert self.log == [("synchronize",)] + self._copy(3, 0) + self._copy(5, 1) - def test_record_compute_signals_the_current_slot(self): - self.manager.record_compute(3) + def test_invalidate_regathers_the_first_layers(self): + self.manager.start_forward() + for layer in LAYERS: + self._bind(layer) + self.manager.finish_layer(layer) + self.log.clear() + self.manager.invalidate_prefetch() + assert self.log == [] + self.manager.start_forward() + assert self.log == self._copy(3, 0) + self._copy(5, 1) - self.events[2].record.assert_called_once_with(self.compute_stream) - self.events[3].record.assert_not_called() + def test_unknown_layer_raises(self): + with pytest.raises(KeyError): + self.manager.wait_and_bind(self.backend, 4) + with pytest.raises(KeyError): + self.manager.finish_layer(4) diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_schedule.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_schedule.py new file mode 100644 index 000000000000..bfae959efb6d --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_schedule.py @@ -0,0 +1,226 @@ +# 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. +"""Host-side issue order of the DWDP expert prefetch (``PrefetchSchedule``).""" + +import random +from typing import Dict, List, Optional + +import pytest + +from tensorrt_llm._torch.modules.dwdp.prefetch_schedule import PrefetchSchedule, ScheduleStep + +pytestmark = pytest.mark.cpu_only + + +def _run_pass(schedule: PrefetchSchedule, num_layers: int) -> None: + """One complete forward pass in MoE order.""" + schedule.start_forward() + for order in range(num_layers): + schedule.bind(order) + assert schedule.finish(order).release == schedule.slot(order) + + +@pytest.mark.parametrize("num_slots", [2, 3, 4]) +@pytest.mark.parametrize("num_layers", list(range(1, 13)) + [92]) +def test_each_layer_is_gathered_once_per_pass(num_layers, num_slots): + schedule = PrefetchSchedule(num_layers, num_slots) + first = schedule.start_forward() + assert first.prefetch == list(range(min(num_layers, num_slots))) + issued = list(first.prefetch) + for order in range(num_layers): + issued += schedule.bind(order).prefetch + step = schedule.finish(order) + assert step.release == order % num_slots + issued += step.prefetch + for _ in range(3): + # Steady state: the previous pass already issued the first layers of this one. + assert schedule.start_forward().prefetch == [] + for order in range(num_layers): + assert schedule.bind(order).prefetch == [] + issued += schedule.finish(order).prefetch + # One cyclic sequence: every pass gathers each layer once, in MoE order. + assert issued[: 4 * num_layers] == list(range(num_layers)) * 4 + assert schedule.num_restarts == 0 + + +@pytest.mark.parametrize( + "num_slots, wrap", + [ + (2, {}), + (4, {}), + # 92 = 30 * 3 + 2: orders 90 and 91 hold slots 0 and 1, so the next pass's order 0 waits + # for MoE(90) and orders 1 and 2 for MoE(91). + (3, {89: [], 90: [0], 91: [1, 2]}), + ], +) +def test_finish_issues_the_layer_num_slots_ahead_and_wraps(num_slots, wrap): + num_layers = 92 + schedule = PrefetchSchedule(num_layers, num_slots) + _run_pass(schedule, num_layers) + assert schedule.start_forward().prefetch == [] + for order in range(num_layers): + assert schedule.bind(order).prefetch == [] + expected = wrap.get(order, [(order + num_slots) % num_layers]) + assert schedule.finish(order).prefetch == expected + + +def test_two_slots_wrap_explicitly(): + # 5 layers, 2 slots: slots 0 1 0 1 0 | 0 1 0 1 0. + schedule = PrefetchSchedule(5, 2) + assert schedule.start_forward().prefetch == [0, 1] + expected_after_finish = {0: [2], 1: [3], 2: [4], 3: [], 4: [0, 1]} + for order in range(5): + assert schedule.bind(order).prefetch == [] + assert schedule.finish(order).prefetch == expected_after_finish[order] + assert schedule.start_forward().prefetch == [] + # 4 layers, 3 slots: slots 0 1 2 0 | 0 1 2 0. + schedule = PrefetchSchedule(4, 3) + assert schedule.start_forward().prefetch == [0, 1, 2] + expected_after_finish = {0: [3], 1: [], 2: [], 3: [0, 1, 2]} + for order in range(4): + schedule.bind(order) + assert schedule.finish(order).prefetch == expected_after_finish[order] + assert schedule.start_forward().prefetch == [] + + +def test_pass_without_start_forward_issues_on_first_bind(): + schedule = PrefetchSchedule(4, 2) + step = schedule.bind(0) + assert step == ScheduleStep(False, None, [0, 1]) + schedule.finish(0) + assert schedule.num_restarts == 0 + + +def test_out_of_order_bind_restarts_at_that_layer(): + schedule = PrefetchSchedule(6, 2) + _ = schedule.start_forward() + schedule.bind(0) + schedule.finish(0) + step = schedule.bind(3) + assert step == ScheduleStep(False, None, [3, 4]) + assert schedule.num_restarts == 1 + assert schedule.finish(3).prefetch == [5] + + +def test_pass_stopped_inside_the_sequence_restarts_at_start_forward(): + schedule = PrefetchSchedule(6, 2) + schedule.start_forward() + for order in range(3): + schedule.bind(order) + schedule.finish(order) + step = schedule.start_forward() + assert step == ScheduleStep(False, None, [0, 1]) + assert schedule.num_restarts == 1 + + +def test_abandoned_bind_requests_a_device_sync(): + schedule = PrefetchSchedule(4, 2) + schedule.start_forward() + schedule.bind(0) + # The MoE of layer 0 raised before finish: its slot may still be read. + step = schedule.start_forward() + assert step == ScheduleStep(True, None, [0, 1]) + schedule.bind(0) + assert schedule.bind(0).synchronize + + +def test_finish_without_matching_bind_raises(): + schedule = PrefetchSchedule(4, 2) + with pytest.raises(RuntimeError): + schedule.finish(0) + schedule.bind(0) + with pytest.raises(RuntimeError): + schedule.finish(1) + with pytest.raises(ValueError): + schedule.bind(4) + + +def test_invalidate_reissues_the_first_layers(): + schedule = PrefetchSchedule(4, 2) + _run_pass(schedule, 4) + assert schedule.invalidate() == ScheduleStep(False, None, []) + assert schedule.start_forward().prefetch == [0, 1] + assert schedule.num_restarts == 0 + schedule.bind(0) + assert schedule.invalidate().synchronize + + +def test_constructor_validates_arguments(): + with pytest.raises(ValueError): + PrefetchSchedule(0, 2) + with pytest.raises(ValueError): + PrefetchSchedule(4, 1) + + +class _SlotModel: + """Device-side model of the slots: checks RAW at every bind and WAR at every copy. + + A copy into a slot may only start after every MoE that read the slot released it (WAR), unless + the caller synchronized the device; a bind must find its own layer as the last copy into the + slot (RAW). + """ + + def __init__(self, num_layers: int, num_slots: int) -> None: + self.num_layers = num_layers + self.num_slots = num_slots + self.last_copy: List[Optional[int]] = [None] * num_slots + # Readers of each slot whose release was not recorded yet. + self.unreleased: List[set] = [set() for _ in range(num_slots)] + self.copies_per_pass: Dict[int, int] = {} + + def apply(self, step: ScheduleStep) -> None: + if step.synchronize: + for readers in self.unreleased: + readers.clear() + if step.release is not None: + self.unreleased[step.release].clear() + for order in step.prefetch: + slot = order % self.num_slots + assert not self.unreleased[slot], f"WAR: copy of {order} into slot {slot} in use" + self.last_copy[slot] = order + self.copies_per_pass[order] = self.copies_per_pass.get(order, 0) + 1 + + def read(self, order: int) -> None: + slot = order % self.num_slots + assert self.last_copy[slot] == order, f"RAW: slot {slot} holds {self.last_copy[slot]}" + self.unreleased[slot].add(order) + + +@pytest.mark.parametrize("seed", range(20)) +def test_randomized_passes_keep_raw_and_war(seed): + rng = random.Random(seed) + num_layers = rng.randint(1, 9) + num_slots = rng.randint(2, 4) + schedule = PrefetchSchedule(num_layers, num_slots) + model = _SlotModel(num_layers, num_slots) + for _ in range(40): + kind = rng.random() + if rng.random() < 0.8: + model.apply(schedule.start_forward()) + if kind < 0.1: + model.apply(schedule.invalidate()) + continue + stop = num_layers + if kind < 0.3: + stop = rng.randint(0, num_layers) + orders = list(range(stop)) + if kind > 0.95 and orders: + orders = sorted(rng.sample(orders, rng.randint(1, len(orders)))) + for order in orders: + model.apply(schedule.bind(order)) + model.read(order) + if kind < 0.2 and order == orders[-1]: + break # The MoE raised: no finish. + model.apply(schedule.finish(order)) From 9f963a3e7ea20a60fc2a517b0ba771163217b4e2 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Wed, 21 Oct 2026 04:05:09 -0700 Subject: [PATCH 6/8] [None][perf] DWDP: configurable number of expert prefetch slots DwdpConfig.num_prefetch_slots (default 2, >= 2) sets the number of remote prefetch slots of the DWDP weight buffer: MoE order o maps its remote experts to slot o % num_prefetch_slots, a table the weight buffer builds once and the prefetch schedule is constructed from; the page pool holds one pool per slot (sized by the largest layer of the slot) and the weight manager keeps one prefetch and one consume event per slot. With the cyclic prefetch schedule, N slots let the gather run N - 1 MoE layers ahead of the compute, so it also runs ahead through compute-bound attention layers instead of re-synchronizing with the compute at each of them. Each slot costs about one layer of remote experts (15.6 GB per slot for Kimi-K3 at DWDP16); the weight buffer is allocated before the KV-cache estimation, so a third slot shrinks the KV pool rather than running out of memory. That trade-off depends on the deployment (memory, prefix-reuse rate), hence a config field; the default keeps the current two slots and lives only in DwdpConfig. A MoE layer that runs several times in a row within a pass (the one-model MTP-Eagle layer, reused for each of the max_draft_len draft steps) is declared once (DwdpManager.add_repeated_layer, called by MTPForCausalLM) and keeps its slot after each finish. The next call that is not a repeat releases it, so every repeat binds the copy still resident instead of restarting the prefetch sequence and gathering the layer again; the cross-step copy into that slot is then issued at the next start_forward. WeightBuffer rejects fewer than two slots with a ValueError. Evidence: Rubin (hecate R200), ctx-only DWDP16 conc 128, cross-step prefetch with the FP8-KV MLA context and kv_b items: 2 slots +5.2% req/s, 3 slots +10.3% req/s (jobs 798859 vs 799151 / 799150). Pipeline model of the R200 ctx nsys trace (job 797928): 3 slots -2.8% full step / -3.3% all steps vs -1.1% / -1.3% for 2 slots. GB300 replay (e47e99fa94): 3 slots -0.3..-1.0% per trace on top of the cross-step prefetch. The new field needs telemetry/privacy CODEOWNER approval (tensorrt_llm/usage/llm_args_golden_manifest.json). Tests (GB300, NSC LHR, base image 23bec710e5 + this tree, job 316218): tests/unittest/_torch/modules/dwdp 207 passed, 1 skipped (a two-slot-only case in the three-slot class); test_trtllm_serve_config_schemas.py + test_llmapi_config_telemetry_docs.py 104 passed; generate_llm_args_golden_manifest.py --check and generate_trtllm_serve_schemas.py (job 316285) match the committed files; executor control-request tests 117 passed. New: test_dwdp_prefetch_slots.py (config default/validation (Pydantic ge=2), slot = MoE order % N, per-slot sizing, slots come from the given mapping, a non-MoE layer has no slot), three-slot event order, GPU tests with 2 and 3 slots. OCI AGA jobs 1369171 (modules/dwdp 206 passed) and 1369418 (GPU tests on 4 GPUs, 2/3 slots). Repeated layers and the slot minimum: test_dwdp_prefetch_schedule.py (repeated-layer pattern N, N, N; randomized RAW/WAR with repeats), test_dwdp_prefetch_pipeline.py (event order with a repeated last layer), test_dwdp_manager.py (repeated layers forwarded to setup_dwdp), test_dwdp_prefetch_slots.py (fewer than two slots rejected); OCI AGA job 1387406. Platform: N (architecture independent). Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- docs/source/developer-guide/moe-perf-guide.md | 2 +- docs/source/developer-guide/telemetry.md | 1 + .../_torch/models/modeling_speculative.py | 5 + tensorrt_llm/_torch/modules/dwdp/__init__.py | 2 +- tensorrt_llm/_torch/modules/dwdp/page_pool.py | 44 ++++--- .../_torch/modules/dwdp/prefetch_schedule.py | 85 +++++++++---- tensorrt_llm/_torch/modules/dwdp/setup.py | 18 ++- tensorrt_llm/_torch/modules/dwdp/vmm.py | 4 +- .../_torch/modules/dwdp/weight_buffer.py | 66 ++++++---- .../_torch/modules/dwdp/weight_manager.py | 32 +++-- tensorrt_llm/_torch/pyexecutor/dwdp.py | 16 +++ tensorrt_llm/llmapi/llm_args.py | 10 +- .../schemas/trtllm-serve-config.schema.json | 9 +- .../trtllm-serve-disagg-config.schema.json | 9 +- .../usage/llm_args_golden_manifest.json | 5 + .../_torch/modules/dwdp/test_dwdp_manager.py | 23 +++- .../modules/dwdp/test_dwdp_prefetch_gpu.py | 4 +- .../dwdp/test_dwdp_prefetch_pipeline.py | 105 ++++++++++++++-- .../dwdp/test_dwdp_prefetch_schedule.py | 115 +++++++++++++++--- .../modules/dwdp/test_dwdp_prefetch_slots.py | 85 +++++++++++++ .../modules/dwdp/test_dwdp_setup_lifecycle.py | 1 + 21 files changed, 531 insertions(+), 110 deletions(-) create mode 100644 tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_slots.py diff --git a/docs/source/developer-guide/moe-perf-guide.md b/docs/source/developer-guide/moe-perf-guide.md index 95b6a60baa06..8fd16236c2e4 100644 --- a/docs/source/developer-guide/moe-perf-guide.md +++ b/docs/source/developer-guide/moe-perf-guide.md @@ -402,7 +402,7 @@ Usage: [`examples/wide_ep/ep_load_balancer/README.md`](https://github.com/NVIDIA [Distributed Weight Data Parallelism (DWDP)](https://doi.org/10.48550/arXiv.2604.01621) targets layer-wise synchronization and token-exchange overhead in the prefill phase, where differences in sequence lengths and KV-cache hits can leave workers waiting for the slowest rank. DWDP keeps a resident subset of experts on each worker and asynchronously prefetches the remaining expert weights from its NVLink peers. Each worker computes on its own tokens, replacing synchronized token dispatch/combine with weight transfers that overlap with computation. -DWDP can be applied to the prefill phase of a serving system when the context batch is large enough for computation to hide the weight prefetch. Decode typically has too little computation to hide these transfers, so generation keeps the normal EP dispatch. Experimental cost of DWDP is a TTFT regression, from a lower service rate on the context stage, plus extra prefetch memory: each rank stores only its local experts, but adds a double-buffered region of about two layers of remote experts. DWDP cannot run on the same MoE path as EPLB. A disaggregated deployment can use DWDP on context servers and EPLB on generation servers. +DWDP can be applied to the prefill phase of a serving system when the context batch is large enough for computation to hide the weight prefetch. Decode typically has too little computation to hide these transfers, so generation keeps the normal EP dispatch. Experimental cost of DWDP is a TTFT regression, from a lower service rate on the context stage, plus extra prefetch memory: each rank stores only its local experts, but adds `num_prefetch_slots` prefetch slots (two by default) of about one layer of remote experts each; more slots let the prefetch run further ahead of the compute. DWDP cannot run on the same MoE path as EPLB. A disaggregated deployment can use DWDP on context servers and EPLB on generation servers. Usage: [`examples/dwdp/README.md`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/dwdp/README.md) and [blog19](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/blogs/tech_blog/blog19_DWDP_Distributed_Weight_Data_Parallelism_for_High_Performance_LLM_Inference_on_NVL72.md). diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 6f8a5f0cb7ca..cc668e2e1647 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/docs/source/developer-guide/telemetry.md @@ -85,6 +85,7 @@ unset or when the safety sanitizer rejects the runtime value. | `dwdp_config.num_experts_per_worker` | `int` | `value` | | | `dwdp_config.num_groups` | `int` | `value` | | | `dwdp_config.num_prefetch_experts` | `int` | `value` | | +| `dwdp_config.num_prefetch_slots` | `int` | `value` | | | `enable_attention_dp` | `bool` | `value` | | | `enable_autotuner` | `bool` | `value` | | | `enable_chunked_prefill` | `bool` | `value` | | diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index dd365170b5f6..ad02c4d606e6 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -29,6 +29,7 @@ WeightsLoadingConfig) from ..modules.rms_norm import RMSNorm from ..moe.fused_moe import moe_load_balancer_set_repeated_for_next_layer +from ..pyexecutor.dwdp import get_global_dwdp_manager from ..pyexecutor.guided_decoder import CapturableGuidedDecoder from ..speculative import (SpecMetadata, get_spec_worker, should_use_separate_draft_kv_cache) @@ -1438,6 +1439,10 @@ def __init__( model.aux_stream_dict) for layer_idx in range(mtp_num_layers) ]) + dwdp_manager = get_global_dwdp_manager() + if dwdp_manager is not None and mtp_repeat_count > 1: + for layer_idx in range(mtp_num_layers): + dwdp_manager.add_repeated_layer(layer_idx + start_layer_idx) self.owns_lm_head = lm_head_quant_config is not None if self.owns_lm_head: if getattr(model_config.pretrained_config, "tie_word_embeddings", diff --git a/tensorrt_llm/_torch/modules/dwdp/__init__.py b/tensorrt_llm/_torch/modules/dwdp/__init__.py index f4c64dfa49b6..b029c68a35f3 100644 --- a/tensorrt_llm/_torch/modules/dwdp/__init__.py +++ b/tensorrt_llm/_torch/modules/dwdp/__init__.py @@ -16,7 +16,7 @@ """DWDP (Distributed Weight Data Parallelism) VA-based infrastructure. MNNVL-based physical memory allocation, composite VA layout, page-aligned -mapping, and double buffer lifecycle management for DWDP expert weights. +mapping, and prefetch slot lifecycle management for DWDP expert weights. Key design principles: - No D2D copy: local experts are always zero-copy via cuMemMap diff --git a/tensorrt_llm/_torch/modules/dwdp/page_pool.py b/tensorrt_llm/_torch/modules/dwdp/page_pool.py index a89badeba87c..84c09788a68c 100644 --- a/tensorrt_llm/_torch/modules/dwdp/page_pool.py +++ b/tensorrt_llm/_torch/modules/dwdp/page_pool.py @@ -13,15 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Page pool for DWDP remote double buffer. +"""Page pool for the DWDP remote prefetch slots. The page pool manages per-page fabric handles for the remote regions of the -composite buffer. Two pools (slot 0, slot 1) support double buffering. +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: - Layer 3 (buf_idx=0): pool[0] pages 0..K for pre/post regions - - Layer 5 (buf_idx=0): SAME pool[0] pages (double buffer reuse) + - Layer 5 (buf_idx=0): SAME pool[0] pages (slot reuse) - Layer 4 (buf_idx=1): pool[1] pages 0..M This naturally handles heterogeneous layers: each layer maps exactly the pages @@ -50,15 +51,15 @@ class PagePool: - """Pool of per-page fabric handles for remote double buffer. + """Pool of per-page fabric handles for the remote prefetch slots. - Two pools (slot 0, slot 1) for double buffering. Each pool is a list + One pool per prefetch slot. Each pool is a list of page-sized fabric handles. Different layers map these pages at different VA positions within their composite VA. 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 (double buffer reuse) + - Layer 5 (buf_idx=0): SAME pool[0] pages (slot reuse) - Layer 4 (buf_idx=1): pool[1] pages 0..M Pool page size vs VMM granularity: @@ -101,8 +102,8 @@ def __init__( """Create a page pool with the specified slot sizes. Args: - slot_sizes: [max_remote_bytes_slot0, max_remote_bytes_slot1] - Computed as max over all layers assigned to each slot. + slot_sizes: Per-slot ``max_remote_bytes``, computed as max over + all layers assigned to each slot. device_id: CUDA device ordinal. granularity: VMM page granularity in bytes. If None, queries the device. @@ -180,8 +181,8 @@ def create( """Factory method to create a page pool. Args: - slot_sizes: [max_remote_bytes_slot0, max_remote_bytes_slot1] - Computed as max over all layers assigned to each slot. + slot_sizes: Per-slot ``max_remote_bytes``, computed as max over + all layers assigned to each slot. device_id: CUDA device ordinal. page_size: Pool page handle size in bytes. If None, uses default (8 * granularity = 16MB on GB200). @@ -208,14 +209,14 @@ def page_size(self) -> int: @property def num_slots(self) -> int: - """Get number of double buffer slots.""" + """Get number of prefetch slots.""" return len(self._page_handles) def num_pages(self, slot: int) -> int: """Get number of pages in a slot. Args: - slot: Double buffer slot (0 or 1). + slot: Prefetch slot index. Returns: Number of pages in the slot. @@ -226,7 +227,7 @@ def slot_size(self, slot: int) -> int: """Get the size of a slot in bytes. Args: - slot: Double buffer slot (0 or 1). + slot: Prefetch slot index. Returns: Size of the slot in bytes. @@ -237,7 +238,7 @@ def get_page_handle(self, slot: int, page_idx: int) -> int: """Get handle for a specific page. Args: - slot: Double buffer slot (0 or 1). + slot: Prefetch slot index. page_idx: Page index within the slot. Returns: @@ -276,7 +277,7 @@ def map_pages( number of cuMemMap calls by 8x. Args: - slot: Double buffer slot (0 or 1). + slot: Prefetch slot index. va_start: Page-aligned VA position to start mapping. size: Total bytes to map (must be a multiple of page_size; callers should use pool_granularity-aligned pre/post sizes). @@ -346,7 +347,7 @@ def release(self) -> None: f"[PagePool] Failed to release page {page_idx} in slot {slot_idx}: {e}" ) - self._page_handles = [[], []] + self._page_handles = [[] for _ in self._slot_pages] logger.debug("[PagePool] Released all page handles") def __del__(self): @@ -369,6 +370,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): def compute_slot_sizes( layouts: dict[int, dict[str, "PageAlignedLayout"]], # noqa: F821 buffer_slot_assignments: dict[int, int], + num_slots: int, ) -> List[int]: """Compute the required slot sizes for a set of layouts. @@ -377,16 +379,18 @@ def compute_slot_sizes( Args: layouts: Mapping from layer_idx to {weight_name: PageAlignedLayout}. - buffer_slot_assignments: Mapping from layer_idx to slot (0 or 1). + buffer_slot_assignments: Slot of every layer in ``layouts``, in + ``[0, num_slots)`` (``WeightBuffer.buffer_index_for_layer``). + num_slots: Number of prefetch slots. Returns: - [slot0_size, slot1_size] in bytes. + Per-slot sizes in bytes, ``num_slots`` entries. """ - slot_sizes = [0, 0] + slot_sizes = [0] * num_slots for layer_idx, weight_layouts in layouts.items(): - slot = buffer_slot_assignments.get(layer_idx, layer_idx % 2) + slot = buffer_slot_assignments[layer_idx] # Sum remote sizes across all weight names within a layer. # Each weight name needs its own pool pages (they hold different data diff --git a/tensorrt_llm/_torch/modules/dwdp/prefetch_schedule.py b/tensorrt_llm/_torch/modules/dwdp/prefetch_schedule.py index 5207fc55ce6d..5bdd2277bde8 100644 --- a/tensorrt_llm/_torch/modules/dwdp/prefetch_schedule.py +++ b/tensorrt_llm/_torch/modules/dwdp/prefetch_schedule.py @@ -15,12 +15,13 @@ """Issue order of the DWDP expert prefetch. The MoE layers of consecutive forward passes form one cyclic sequence of positions: position -``p`` is MoE order ``p % num_layers`` and uses the prefetch slot ``order % num_slots`` (the slot -its composite VA maps). A position is issued (its copy enqueued) once every earlier position is -issued and its slot no longer holds a position that has not finished. A slot is released right -after the MoE that read it is enqueued, so each MoE layer is gathered once per pass, the copy into -a slot overlaps everything up to the next MoE that reads the slot, and the last MoE layers of a -pass already gather the first layers of the next pass, through the gap between executor steps. +``p`` is MoE order ``p % num_layers`` and uses that order's prefetch slot (the slot its composite +VA maps, ``WeightBuffer.buffer_index_for_layer``). A position is issued (its copy enqueued) once +every earlier position is issued and its slot no longer holds a position that has not finished. +A slot is released right after the MoE that read it is enqueued, so each MoE layer is gathered +once per pass, the copy into a slot overlaps everything up to the next MoE that reads the slot, +and the last MoE layers of a pass already gather the first layers of the next pass, through the +gap between executor steps. ``PrefetchSchedule`` only does the host bookkeeping. Each MoE layer goes through ``bind`` (before its MoE reads the slot) and ``finish`` (right after the MoE is enqueued); a pass starts with @@ -28,11 +29,16 @@ A pass that stops inside the sequence, or a bind out of sequence, restarts the sequence at that layer; a pass without ``start_forward`` (e.g. a warmup forward) issues its copies on the first bind. + +A repeated MoE order runs several times in a row within a pass (the one-model MTP-Eagle layer, +reused for every draft step). Its slot is held after each finish and released only by the next +call that is not a bind of the same order, so every repeat binds the copy that is still in the +slot instead of gathering the layer again. """ from __future__ import annotations -from typing import List, NamedTuple, Optional +from typing import Iterable, List, NamedTuple, Optional, Sequence class ScheduleStep(NamedTuple): @@ -54,17 +60,24 @@ class PrefetchSchedule: """Host-side issue order of the DWDP expert prefetch over consecutive forward passes. Args: - num_layers: Number of MoE layers of one forward pass. - num_slots: Number of prefetch buffer slots; MoE order ``o`` uses slot ``o % num_slots``. + slots: Prefetch slot of each MoE order of a forward pass, as the weight buffer maps it. + repeated_orders: MoE orders that may be bound several times in a row within a pass. + + Raises: + ValueError: If ``slots`` is empty or negative, or a repeated order is out of range. """ - def __init__(self, num_layers: int, num_slots: int) -> None: - if num_layers <= 0: - raise ValueError(f"num_layers must be positive, got {num_layers}") - if num_slots < 2: - raise ValueError(f"num_slots must be at least 2, got {num_slots}") - self._num_layers = num_layers - self._num_slots = num_slots + def __init__(self, slots: Sequence[int], repeated_orders: Iterable[int] = ()) -> None: + if not slots: + raise ValueError("a prefetch schedule needs at least one MoE layer") + if min(slots) < 0: + raise ValueError(f"negative prefetch slot in {list(slots)}") + self._slots = list(slots) + self._num_layers = len(self._slots) + self._repeated = frozenset(repeated_orders) + for order in self._repeated: + self._check_order(order) + num_slots = max(self._slots) + 1 # Next position to bind and number of positions issued so far. self._next_bind = 0 self._issued = 0 @@ -72,6 +85,8 @@ def __init__(self, num_layers: int, num_slots: int) -> None: self._slot_owner: List[Optional[int]] = [None] * num_slots # Position bound and not finished yet. self._bound: Optional[int] = None + # Finished position of a repeated order whose slot is still held for a repeat. + self._held: Optional[int] = None self._num_restarts = 0 @property @@ -81,25 +96,35 @@ def num_restarts(self) -> int: def slot(self, order: int) -> int: """Prefetch slot of MoE order ``order``.""" - return order % self._num_slots + return self._slots[order] def start_forward(self) -> ScheduleStep: """Start a forward pass: issue its first layers unless they are already in flight.""" + release = self._release_held() if self._bound is None and self._next_bind % self._num_layers == 0: - return ScheduleStep(False, None, self._issue()) + return ScheduleStep(False, release, self._issue()) self._num_restarts += 1 synchronize = self._restart(0) - return ScheduleStep(synchronize, None, self._issue()) + return ScheduleStep(synchronize, release, self._issue()) def bind(self, order: int) -> ScheduleStep: """Bind MoE order ``order``; its MoE reads the slot only after the returned step. - A bind out of sequence restarts the sequence at ``order``. + A bind out of sequence restarts the sequence at ``order``. A repeated order bound again + right after its finish reads the copy still held in its slot. Raises: ValueError: If ``order`` is out of range. """ self._check_order(order) + if ( + self._bound is None + and self._held is not None + and order == self._held % self._num_layers + ): + self._bound = self._held + return ScheduleStep(False, None, []) + release = self._release_held() synchronize = False if self._bound is not None or order != self._next_bind % self._num_layers: self._num_restarts += 1 @@ -110,7 +135,7 @@ def bind(self, order: int) -> ScheduleStep: f"DWDP prefetch schedule: MoE order {order} bound before its copy was issued" ) self._bound = self._next_bind - return ScheduleStep(synchronize, None, prefetch) + return ScheduleStep(synchronize, release, prefetch) def finish(self, order: int) -> ScheduleStep: """Finish the bound MoE order ``order`` (its MoE has just been enqueued). @@ -127,6 +152,10 @@ def finish(self, order: int) -> ScheduleStep: f"(bound={bound})" ) self._next_bind = self._bound + 1 + if order in self._repeated: + self._held = self._bound + self._bound = None + return ScheduleStep(False, None, []) self._bound = None slot = self.slot(order) self._slot_owner[slot] = None @@ -134,16 +163,26 @@ def finish(self, order: int) -> ScheduleStep: def invalidate(self) -> ScheduleStep: """Forget every issued prefetch; the next pass issues its first layers again.""" - return ScheduleStep(self._restart(0), None, []) + release = self._release_held() + return ScheduleStep(self._restart(0), release, []) def _check_order(self, order: int) -> None: if not 0 <= order < self._num_layers: raise ValueError(f"MoE order {order} out of range [0, {self._num_layers})") + def _release_held(self) -> Optional[int]: + """Release the slot held for a repeat, if any; returns the slot.""" + if self._held is None: + return None + slot = self.slot(self._held % self._num_layers) + self._slot_owner[slot] = None + self._held = None + return slot + def _restart(self, order: int) -> bool: """Restart the sequence at ``order``; returns whether a bound layer was abandoned.""" abandoned = self._bound is not None - self._slot_owner = [None] * self._num_slots + self._slot_owner = [None] * len(self._slot_owner) self._bound = None self._next_bind = order self._issued = order diff --git a/tensorrt_llm/_torch/modules/dwdp/setup.py b/tensorrt_llm/_torch/modules/dwdp/setup.py index b15f4330e69f..b8203c2e6e0b 100644 --- a/tensorrt_llm/_torch/modules/dwdp/setup.py +++ b/tensorrt_llm/_torch/modules/dwdp/setup.py @@ -27,7 +27,7 @@ local weights into fabric memory and freeing the originals. Handles are exchanged via the DWDP MPI communicator. 5. Constructs a WeightBuffer with composite VA layout (zero-copy local - + page-pool-backed remote double buffer). + + 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 @@ -42,7 +42,7 @@ from __future__ import annotations import traceback -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Sequence, Tuple import torch from torch import nn @@ -93,7 +93,9 @@ def setup_dwdp( layer_indices: List[int], num_experts_per_worker: int, num_prefetch_experts: int, + num_prefetch_slots: int, contention_opt: bool = False, + repeated_layer_indices: Sequence[int] = (), ) -> Optional[DWDPWeightManager]: """Set up DWDP for a model after weight loading. @@ -124,6 +126,12 @@ def setup_dwdp( redundancy: 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 + remote experts of one MoE layer (``DwdpConfig.num_prefetch_slots``). + contention_opt: Copy-engine prefetch as one interleaved + ``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``). Returns: A ready-to-use DWDPWeightManager if DWDP is enabled, ``None`` otherwise. @@ -235,11 +243,14 @@ def setup_dwdp( local_end=local_end, dwdp_size=mapping.dwdp_size, device_id=device_id, + num_slots=num_prefetch_slots, ) except BaseException: _release_during_rollback(transport, "transport") raise - logger.info("[DWDP Setup] Weight buffer created successfully.") + logger.info( + f"[DWDP Setup] Weight buffer created successfully ({num_prefetch_slots} prefetch slots)." + ) # 6. Fill edge bytes logger.info("[DWDP Setup] Filling page-alignment edge bytes...") @@ -270,6 +281,7 @@ def setup_dwdp( dwdp_rank=mapping.dwdp_rank, dwdp_size=mapping.dwdp_size, contention_opt=contention_opt, + repeated_layer_indices=repeated_layer_indices, ) except BaseException: _release_during_rollback(weight_buffer, "weight buffer") diff --git a/tensorrt_llm/_torch/modules/dwdp/vmm.py b/tensorrt_llm/_torch/modules/dwdp/vmm.py index b472152d0bb5..95dfaa7dc82e 100644 --- a/tensorrt_llm/_torch/modules/dwdp/vmm.py +++ b/tensorrt_llm/_torch/modules/dwdp/vmm.py @@ -236,8 +236,8 @@ def create_local_handle(size: int, device_id: int) -> int: Uses CU_MEM_HANDLE_TYPE_NONE — the handle cannot be exported to other processes but avoids consuming NVLink fabric routing table entries. - Use this for buffers that only need local GPU access (e.g., page pool - double buffers that are written via P2P copy and read locally). + Use this for buffers that only need local GPU access (e.g., the page pool's + prefetch slots that are written via P2P copy and read locally). Args: size: Size in bytes (must be aligned to granularity). diff --git a/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py b/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py index bad47c0d8425..42ac02afb167 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_buffer.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""WeightBuffer: Composite VA layout with zero-copy local + page-pool-backed remote double buffer. +"""WeightBuffer: Composite VA layout with zero-copy local + page-pool-backed remote prefetch slots. The WeightBuffer is the central component of DWDP weight management. It creates a composite virtual address space that seamlessly combines: @@ -82,7 +82,7 @@ class LayerBufferState: class WeightBuffer: - """Composite VA layout with zero-copy local + page-pool-backed remote double buffer. + """Composite VA layout with zero-copy local + page-pool-backed remote prefetch slots. The WeightBuffer does NOT copy any data. It only sets up VA mappings: - Local region: cuMemMap to MNNVL handle (zero-copy, no D2D) @@ -108,6 +108,8 @@ class WeightBuffer: "_page_pool", "_layer_states", "_moe_layer_indices", + "_num_slots", + "_slot_of_layer", "_released", ) @@ -119,6 +121,7 @@ def __init__( local_end: int, dwdp_size: int, device_id: int, + num_slots: int, granularity: Optional[int] = None, pool_page_size: Optional[int] = None, ): @@ -131,12 +134,21 @@ def __init__( local_end: Last local expert index (exclusive). dwdp_size: Number of DWDP ranks. device_id: CUDA device ordinal. + num_slots: Number of remote prefetch slots + (``DwdpConfig.num_prefetch_slots``); MoE order ``o`` maps its + remote experts to slot ``o % num_slots``. granularity: VMM page granularity. If None, queries device. pool_page_size: Pool page handle size in bytes. If None, defaults to ``PagePool.DEFAULT_PAGE_SIZE_MULTIPLIER * granularity`` (16MB on GB200). Must be a multiple of granularity. + + Raises: + ValueError: If ``num_slots`` is less than 2. """ + if num_slots < 2: + raise ValueError(f"num_slots must be at least 2, got {num_slots}") + self._num_slots = num_slots self._layer_weight_specs = layer_weight_specs self._handles = handles self._local_start = local_start @@ -157,6 +169,9 @@ def __init__( self._page_pool: Optional[PagePool] = None self._layer_states: Dict[int, LayerBufferState] = {} self._moe_layer_indices: List[int] = sorted(layer_weight_specs.keys()) + self._slot_of_layer: Dict[int, int] = { + layer_idx: order % num_slots for order, layer_idx in enumerate(self._moe_layer_indices) + } self._released = False @classmethod @@ -168,13 +183,14 @@ def create( local_end: int, dwdp_size: int, device_id: int, + num_slots: int, ) -> WeightBuffer: """Create all VA mappings and remote page pool. 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 remote double buffer + 2. Allocate page pool for the remote prefetch slots 3. Create per-layer composite VA mappings 4. Create tensor views @@ -186,6 +202,8 @@ def create( local_end: Last local expert index (exclusive). dwdp_size: Number of DWDP ranks. device_id: CUDA device ordinal. + num_slots: Number of remote prefetch slots, each holding the + remote experts of one MoE layer. Returns: Ready-to-use WeightBuffer with all VA mappings established. @@ -204,6 +222,7 @@ def create( local_end=local_end, dwdp_size=dwdp_size, device_id=device_id, + num_slots=num_slots, ) try: @@ -214,11 +233,7 @@ def create( # Step 2: Compute slot sizes and create page pool. # Because layouts already have pre/post sizes aligned to # pool_page_size, the slot sizes are automatically aligned too. - buffer_slot_assignments = { - layer_idx: buffer.buffer_index_for_layer(layer_idx) - for layer_idx in layer_weight_specs.keys() - } - slot_sizes = compute_slot_sizes(layouts, buffer_slot_assignments) + slot_sizes = compute_slot_sizes(layouts, buffer._slot_of_layer, num_slots) buffer._page_pool = PagePool.create( slot_sizes, device_id, page_size=buffer._pool_page_size @@ -621,22 +636,27 @@ def device_id(self) -> int: """CUDA device ordinal.""" return self._device_id + @property + def num_slots(self) -> int: + """Number of remote prefetch slots.""" + return self._num_slots + def buffer_index_for_layer(self, layer_idx: int) -> int: - """Double buffer slot: 0 or 1 based on MoE order (odd/even). + """Prefetch slot of a MoE layer: its MoE order modulo ``num_slots``. + + This is the only definition of the slot mapping; the weight manager's + prefetch schedule is built from it. Args: - layer_idx: Layer index. + layer_idx: MoE layer index. Returns: - Buffer slot index (0 or 1). - """ - # Find position of this layer in the MoE layer sequence - if layer_idx in self._moe_layer_indices: - moe_order = self._moe_layer_indices.index(layer_idx) - else: - moe_order = layer_idx + Buffer slot index in ``[0, num_slots)``. - return moe_order % 2 + Raises: + KeyError: If ``layer_idx`` is not a MoE layer of this buffer. + """ + return self._slot_of_layer[layer_idx] def weight_names(self, layer_idx: int) -> List[str]: """Get weight names for a specific layer. @@ -733,10 +753,12 @@ def debug_info(self) -> Dict: if self._page_pool is not None: info["page_pool"] = { "page_size": self._page_pool.page_size, - "slot0_pages": self._page_pool.num_pages(0), - "slot1_pages": self._page_pool.num_pages(1), - "slot0_size": self._page_pool.slot_size(0), - "slot1_size": self._page_pool.slot_size(1), + "slot_pages": [ + self._page_pool.num_pages(slot) for slot in range(self._page_pool.num_slots) + ], + "slot_sizes": [ + self._page_pool.slot_size(slot) for slot in range(self._page_pool.num_slots) + ], } return info diff --git a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py index 5d15da642c5e..a918e2e2c7af 100644 --- a/tensorrt_llm/_torch/modules/dwdp/weight_manager.py +++ b/tensorrt_llm/_torch/modules/dwdp/weight_manager.py @@ -41,7 +41,7 @@ import logging import traceback -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Optional, Sequence, Tuple import torch @@ -73,9 +73,6 @@ # number of blocks each group gets. _GatherPlan = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int] -# Number of prefetch buffer slots (``WeightBuffer`` double buffer). -_NUM_SLOTS = 2 - class DWDPWeightManager: """Runtime P2P copy orchestrator for DWDP expert weight prefetch. @@ -131,6 +128,7 @@ def __init__( dwdp_rank: int, dwdp_size: int, contention_opt: bool = False, + repeated_layer_indices: Sequence[int] = (), ) -> None: """Initialize the DWDPWeightManager. @@ -161,10 +159,13 @@ def __init__( to reduce single-NVLink-link contention. Plans are cached per layer since peer ``data_ptr()`` and destination ``data_ptr()`` are stable after initialization. + 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``). Raises: - ValueError: If dwdp_rank is out of range, dwdp_size is invalid or - there are no MoE layers. + 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. """ if dwdp_size <= 0: raise ValueError(f"dwdp_size must be positive, got {dwdp_size}") @@ -191,9 +192,17 @@ def __init__( device = torch.device("cuda", weight_buffer.device_id) self._copy_stream = torch.cuda.Stream(device=device) - num_slots = _NUM_SLOTS + num_slots = weight_buffer.num_slots self._layer_order = {layer: order for order, layer in enumerate(self._moe_layer_indices)} - self._schedule = PrefetchSchedule(len(self._moe_layer_indices), num_slots) + # The weight buffer owns the slot mapping; the schedule and the event + # indices below both read it from there. + unknown = set(repeated_layer_indices) - self._moe_layer_set + if unknown: + raise ValueError(f"repeated layers {sorted(unknown)} are not DWDP MoE layers") + self._schedule = PrefetchSchedule( + [weight_buffer.buffer_index_for_layer(layer) for layer in self._moe_layer_indices], + repeated_orders=[self._layer_order[layer] for layer in repeated_layer_indices], + ) self._reported_restarts = 0 # Buffer-slot events: one pair per slot. @@ -307,7 +316,7 @@ def prefetch_layer(self, layer_idx: int) -> None: f"Layer {layer_idx} is not a MoE layer. Valid layers: {self._moe_layer_indices}" ) - buf_idx = self._weight_buffer.buffer_index_for_layer(layer_idx) + buf_idx = self._schedule.slot(self._order_of(layer_idx)) with torch.cuda.stream(self._copy_stream): # WAR hazard: wait until compute finishes reading this buffer slot @@ -551,9 +560,10 @@ def wait_and_bind(self, backend_module: torch.nn.Module, layer_idx: int) -> None KeyError: If layer_idx is not a valid MoE layer. AttributeError: If backend_module lacks a required weight attribute. """ - self._apply(self._schedule.bind(self._order_of(layer_idx))) + order = self._order_of(layer_idx) + self._apply(self._schedule.bind(order)) - buf_idx = self._weight_buffer.buffer_index_for_layer(layer_idx) + buf_idx = self._schedule.slot(order) compute_stream = torch.cuda.current_stream( torch.device("cuda", self._weight_buffer.device_id) ) diff --git a/tensorrt_llm/_torch/pyexecutor/dwdp.py b/tensorrt_llm/_torch/pyexecutor/dwdp.py index 476c40d7112d..0497b1299119 100644 --- a/tensorrt_llm/_torch/pyexecutor/dwdp.py +++ b/tensorrt_llm/_torch/pyexecutor/dwdp.py @@ -166,6 +166,8 @@ def __init__( # SSOT for MoE layer indices; populated by configurable_moe.add_layer() self._registered_layers: List[int] = [] + # Layers that run several times in a row per pass; see add_repeated_layer() + self._repeated_layers: List[int] = [] # Set by setup(model); None until then self._weight_manager: Optional[DWDPWeightManager] = None @@ -260,6 +262,16 @@ def add_layer(self, layer_idx: int) -> None: return self._registered_layers.append(layer_idx) + def add_repeated_layer(self, layer_idx: int) -> None: + """Declare that layer ``layer_idx`` runs several times in a row per forward pass. + + Called by models that reuse a decoder layer, e.g. the one-model MTP-Eagle layer that runs + once per draft step. Its repeats bind the expert copy still in the prefetch slot instead + of gathering the layer again. Layers whose MoE does not use DWDP are ignored. + """ + if layer_idx not in self._repeated_layers: + self._repeated_layers.append(layer_idx) + def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: """Run the VA initialization pipeline on ``model``. @@ -280,6 +292,10 @@ def setup(self, model: nn.Module) -> Optional[DWDPWeightManager]: 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 + ], ) self._model_ref = weakref.ref(model) if self._store_rendezvous: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e02c42ae24f5..514aa4e26f05 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5704,7 +5704,7 @@ class DwdpConfig(StrictBaseModel): Each worker holds a subset of experts locally and asynchronously prefetches the remaining experts from peer workers via CUDA VMM + MNNVL fabric handles (composite virtual-address layout with zero-copy local mapping and - double-buffered P2P remote regions), enabling fully asynchronous execution + P2P remote prefetch slots), enabling fully asynchronous execution across ranks without synchronization barriers. Currently supported with the CuteDSL MoE backend and NVFP4 quantization @@ -5726,6 +5726,14 @@ class DwdpConfig(StrictBaseModel): "Enable limited-contention prefetch optimization. Uses batched " "memcpy with round-robin slice ordering across peers and a 2 MiB " "slice granularity to reduce NVLink contention.") + num_prefetch_slots: int = Field( + default=2, + ge=2, + description= + "The number of expert prefetch buffer slots. Each slot holds the " + "remote experts of one MoE layer; more slots let the prefetch run " + "further ahead of the MoE compute, at the cost of device memory that " + "is then not available to the KV cache.") class BaseLlmArgs(StrictBaseModel): diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 3db659f73b23..b2ca2567a6fe 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -1836,7 +1836,7 @@ }, "DwdpConfig": { "additionalProperties": false, - "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\ndouble-buffered P2P remote regions), 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).", + "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": { "default": false, @@ -1867,6 +1867,13 @@ "description": "The number of prefetch experts per worker.", "title": "Num Prefetch Experts", "type": "integer" + }, + "num_prefetch_slots": { + "default": 2, + "description": "The number of expert prefetch buffer slots. Each slot holds the remote experts of one MoE layer; more slots let the prefetch run further ahead of the MoE compute, at the cost of device memory that is then not available to the KV cache.", + "minimum": 2, + "title": "Num Prefetch Slots", + "type": "integer" } }, "title": "DwdpConfig", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 371b27ff1e08..704cffba5538 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -3400,7 +3400,7 @@ }, "DwdpConfig": { "additionalProperties": false, - "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\ndouble-buffered P2P remote regions), 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).", + "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": { "default": false, @@ -3431,6 +3431,13 @@ "description": "The number of prefetch experts per worker.", "title": "Num Prefetch Experts", "type": "integer" + }, + "num_prefetch_slots": { + "default": 2, + "description": "The number of expert prefetch buffer slots. Each slot holds the remote experts of one MoE layer; more slots let the prefetch run further ahead of the MoE compute, at the cost of device memory that is then not available to the KV cache.", + "minimum": 2, + "title": "Num Prefetch Slots", + "type": "integer" } }, "title": "DwdpConfig", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f940f2a80937..30acb66df517 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -310,6 +310,11 @@ "kind": "value", "path": "dwdp_config.num_prefetch_experts" }, + { + "capture_policy": "int", + "kind": "value", + "path": "dwdp_config.num_prefetch_slots" + }, { "capture_policy": "bool", "kind": "value", diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py index 1cca0d9a9957..31217886e537 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_manager.py @@ -350,14 +350,33 @@ def test_setup_forwards_layer_indices(self, mock_comm_world, _rank, mock_setup): # Config-driven expert range fields are forwarded from DwdpConfig. self.assertEqual(kwargs["num_experts_per_worker"], 4) self.assertEqual(kwargs["num_prefetch_experts"], 4) + self.assertEqual(kwargs["num_prefetch_slots"], 2) self.assertIs(result, fake_weight_manager) self.assertIs(mgr._weight_manager, fake_weight_manager) + self.assertEqual(kwargs["repeated_layer_indices"], []) + + 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.""" + mgr, _ = self._set_up_manager( + mock_comm_world, + _setup, + before_setup=lambda m: ( + m.add_layer(3), + m.add_layer(61), + m.add_repeated_layer(61), + m.add_repeated_layer(61), + m.add_repeated_layer(62), + ), + ) + _, kwargs = _setup.call_args + self.assertEqual(kwargs["layer_indices"], [3, 61]) + self.assertEqual(kwargs["repeated_layer_indices"], [61]) # ------------------------------------------------------------------ # Runtime forwards require setup() # ------------------------------------------------------------------ - def _set_up_manager(self, mock_comm_world, mock_setup): + def _set_up_manager(self, mock_comm_world, mock_setup, before_setup=None): mock_comm_world.Create_group.return_value = MagicMock() mock_comm_world.group.Incl.return_value = MagicMock() mock_comm_world.Get_size.return_value = 1024 @@ -368,6 +387,8 @@ def _set_up_manager(self, mock_comm_world, mock_setup): dist=MagicMock(spec=MPIDist), mapping=_make_mapping(), ) + if before_setup is not None: + before_setup(mgr) with patch( "tensorrt_llm._torch.pyexecutor.dwdp.torch.cuda.current_device", return_value=0, diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py index ad3b3624ed37..96f21c46f79b 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_gpu.py @@ -217,7 +217,7 @@ def _drive(h: _Harness, moe_delay: int, num_slots: int) -> None: @pytest.mark.parametrize("num_layers", [4, 5]) -@pytest.mark.parametrize("num_slots", [2]) +@pytest.mark.parametrize("num_slots", [2, 3]) @pytest.mark.parametrize("slow", ["moe", "copy"]) def test_bound_weights_across_passes(monkeypatch, num_layers, num_slots, slow): monkeypatch.delenv("TRTLLM_DWDP_PREFETCH_GATHER", raising=False) @@ -278,7 +278,7 @@ def _enable_peer_access(devices: List[int]) -> None: pytest.skip(f"cudaDeviceEnablePeerAccess({dev}) failed: {err}") -@pytest.mark.parametrize("num_slots", [2]) +@pytest.mark.parametrize("num_slots", [2, 3]) @pytest.mark.parametrize("slow", ["moe", "copy"]) def test_fabric_gather_from_peer_gpus(monkeypatch, num_slots, slow): """The production transport: the SM gather kernel pulls the remote experts from peer GPUs.""" 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 f51ece0ac753..c8f121c26ada 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_pipeline.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_pipeline.py @@ -27,6 +27,7 @@ pytestmark = pytest.mark.cpu_only LAYERS = [3, 5, 7] +NUM_SLOTS = 2 class _Event: @@ -48,14 +49,18 @@ def wait_event(self, event: _Event) -> None: class TestDWDPPrefetchPipeline: + num_slots = NUM_SLOTS + repeated_layers: Tuple[int, ...] = () + def setup_method(self): self.log: List[Tuple[str, ...]] = [] self.weight_buffer = MagicMock() self.weight_buffer.device_id = 0 self.weight_buffer.local_start = 0 self.weight_buffer.local_end = 2 + self.weight_buffer.num_slots = self.num_slots self.weight_buffer.buffer_index_for_layer.side_effect = lambda layer: ( - LAYERS.index(layer) % 2 + LAYERS.index(layer) % self.num_slots ) self.weight_buffer.get_full_tensor.side_effect = lambda layer, name: torch.tensor( [float(layer)] @@ -63,9 +68,9 @@ def setup_method(self): self.copy_stream = _Stream("copy", self.log) self.compute_stream = _Stream("compute", self.log) - events = iter( - [_Event(name, self.log) for name in ("prefetch0", "prefetch1", "consume0", "consume1")] - ) + 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]) self.patches = [ patch( "tensorrt_llm._torch.modules.dwdp.weight_manager.torch.cuda.Stream", @@ -104,6 +109,7 @@ def setup_method(self): weight_names=["gate_up_proj", "down_proj"], dwdp_rank=0, dwdp_size=2, + repeated_layer_indices=self.repeated_layers, ) self.backend = nn.Module() self.backend.gate_up_proj = nn.Parameter(torch.empty(1), requires_grad=False) @@ -123,12 +129,18 @@ def _copy(layer: int, slot: int) -> List[Tuple[str, ...]]: ("record", f"prefetch{slot}", "copy"), ] + def _first_layers(self) -> List[Tuple[str, ...]]: + """Copies of the first MoE layers, one per slot.""" + return [entry for slot in range(self.num_slots) for entry in self._copy(LAYERS[slot], slot)] + def _bind(self, layer: int) -> None: self.manager.wait_and_bind(self.backend, layer) assert self.backend.gate_up_proj.data.item() == float(layer) assert self.backend.down_proj.data.item() == float(layer) def test_event_order_across_passes(self): + if self.num_slots != 2: + pytest.skip("two-slot sequence") self.manager.start_forward() assert self.log == self._copy(3, 0) + self._copy(5, 1) @@ -169,14 +181,14 @@ def test_event_order_across_passes(self): def test_bind_without_start_forward_issues_its_copy_first(self): self._bind(3) - assert self.log == self._copy(3, 0) + self._copy(5, 1) + [("wait", "compute", "prefetch0")] + assert self.log == self._first_layers() + [("wait", "compute", "prefetch0")] def test_abandoned_pass_synchronizes_before_reusing_slots(self): self.manager.start_forward() self._bind(3) self.log.clear() self.manager.start_forward() - assert self.log == [("synchronize",)] + self._copy(3, 0) + self._copy(5, 1) + assert self.log == [("synchronize",)] + self._first_layers() def test_invalidate_regathers_the_first_layers(self): self.manager.start_forward() @@ -187,10 +199,89 @@ def test_invalidate_regathers_the_first_layers(self): self.manager.invalidate_prefetch() assert self.log == [] self.manager.start_forward() - assert self.log == self._copy(3, 0) + self._copy(5, 1) + assert self.log == self._first_layers() def test_unknown_layer_raises(self): with pytest.raises(KeyError): self.manager.wait_and_bind(self.backend, 4) with pytest.raises(KeyError): self.manager.finish_layer(4) + + +class TestDWDPPrefetchPipelineRepeatedLayer(TestDWDPPrefetchPipeline): + """The last MoE layer runs three times per pass, like the one-model MTP-Eagle layer.""" + + repeated_layers = (7,) + + def _pass(self, repeats: int = 3) -> None: + for layer in LAYERS[:-1] + [LAYERS[-1]] * repeats: + self._bind(layer) + self.manager.finish_layer(layer) + + def test_event_order_across_passes(self): + self.manager.start_forward() + self._bind(3) + self.manager.finish_layer(3) + self._bind(5) + self.manager.finish_layer(5) + self.log.clear() + + # Every use of layer 7 binds the copy in slot 0; the slot stays held. + for _ in range(3): + self._bind(7) + self.manager.finish_layer(7) + assert self.log == [("wait", "compute", "prefetch0")] * 3 + + # The next pass releases it after the last use and gathers its first layers once. + self.log.clear() + self.manager.start_forward() + assert self.log == [("record", "consume0", "compute")] + self._first_layers() + self.log.clear() + self._pass() + self.manager.start_forward() + copies = [entry for entry in self.log if entry[0] == "copy"] + assert copies == [("copy", 7), ("copy", 3), ("copy", 5)] + assert self.manager._schedule.num_restarts == 0 + + def test_invalidate_regathers_the_first_layers(self): + self.manager.start_forward() + self._pass() + self.log.clear() + self.manager.invalidate_prefetch() + assert self.log == [("record", "consume0", "compute")] + self.log.clear() + self.manager.start_forward() + assert self.log == self._first_layers() + + def test_repeated_layer_must_be_a_dwdp_layer(self): + with pytest.raises(ValueError): + DWDPWeightManager( + 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, + repeated_layer_indices=[4], + ) + + +class TestDWDPPrefetchPipelineThreeSlots(TestDWDPPrefetchPipeline): + num_slots = 3 + + def test_three_slots_prefetch_a_whole_pass_ahead(self): + self.manager.start_forward() + assert self.log == self._copy(3, 0) + self._copy(5, 1) + self._copy(7, 2) + for layer, slot in zip(LAYERS, range(3)): + self.log.clear() + self._bind(layer) + self.manager.finish_layer(layer) + # Each finish frees the slot of the same layer in the next pass. + assert self.log == [ + ("wait", "compute", f"prefetch{slot}"), + ("record", f"consume{slot}", "compute"), + ] + self._copy(layer, slot) + self.log.clear() + self.manager.start_forward() + assert self.log == [] diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_schedule.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_schedule.py index bfae959efb6d..b82a7638bd4d 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_schedule.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_schedule.py @@ -24,6 +24,11 @@ pytestmark = pytest.mark.cpu_only +def _schedule(num_layers: int, num_slots: int, repeated=()) -> PrefetchSchedule: + """The weight buffer's mapping: MoE order ``o`` uses slot ``o % num_slots``.""" + return PrefetchSchedule([order % num_slots for order in range(num_layers)], repeated) + + def _run_pass(schedule: PrefetchSchedule, num_layers: int) -> None: """One complete forward pass in MoE order.""" schedule.start_forward() @@ -35,7 +40,7 @@ def _run_pass(schedule: PrefetchSchedule, num_layers: int) -> None: @pytest.mark.parametrize("num_slots", [2, 3, 4]) @pytest.mark.parametrize("num_layers", list(range(1, 13)) + [92]) def test_each_layer_is_gathered_once_per_pass(num_layers, num_slots): - schedule = PrefetchSchedule(num_layers, num_slots) + schedule = _schedule(num_layers, num_slots) first = schedule.start_forward() assert first.prefetch == list(range(min(num_layers, num_slots))) issued = list(first.prefetch) @@ -67,7 +72,7 @@ def test_each_layer_is_gathered_once_per_pass(num_layers, num_slots): ) def test_finish_issues_the_layer_num_slots_ahead_and_wraps(num_slots, wrap): num_layers = 92 - schedule = PrefetchSchedule(num_layers, num_slots) + schedule = _schedule(num_layers, num_slots) _run_pass(schedule, num_layers) assert schedule.start_forward().prefetch == [] for order in range(num_layers): @@ -78,7 +83,7 @@ def test_finish_issues_the_layer_num_slots_ahead_and_wraps(num_slots, wrap): def test_two_slots_wrap_explicitly(): # 5 layers, 2 slots: slots 0 1 0 1 0 | 0 1 0 1 0. - schedule = PrefetchSchedule(5, 2) + schedule = _schedule(5, 2) assert schedule.start_forward().prefetch == [0, 1] expected_after_finish = {0: [2], 1: [3], 2: [4], 3: [], 4: [0, 1]} for order in range(5): @@ -86,7 +91,7 @@ def test_two_slots_wrap_explicitly(): assert schedule.finish(order).prefetch == expected_after_finish[order] assert schedule.start_forward().prefetch == [] # 4 layers, 3 slots: slots 0 1 2 0 | 0 1 2 0. - schedule = PrefetchSchedule(4, 3) + schedule = _schedule(4, 3) assert schedule.start_forward().prefetch == [0, 1, 2] expected_after_finish = {0: [3], 1: [], 2: [], 3: [0, 1, 2]} for order in range(4): @@ -96,7 +101,7 @@ def test_two_slots_wrap_explicitly(): def test_pass_without_start_forward_issues_on_first_bind(): - schedule = PrefetchSchedule(4, 2) + schedule = _schedule(4, 2) step = schedule.bind(0) assert step == ScheduleStep(False, None, [0, 1]) schedule.finish(0) @@ -104,7 +109,7 @@ def test_pass_without_start_forward_issues_on_first_bind(): def test_out_of_order_bind_restarts_at_that_layer(): - schedule = PrefetchSchedule(6, 2) + schedule = _schedule(6, 2) _ = schedule.start_forward() schedule.bind(0) schedule.finish(0) @@ -115,7 +120,7 @@ def test_out_of_order_bind_restarts_at_that_layer(): def test_pass_stopped_inside_the_sequence_restarts_at_start_forward(): - schedule = PrefetchSchedule(6, 2) + schedule = _schedule(6, 2) schedule.start_forward() for order in range(3): schedule.bind(order) @@ -126,7 +131,7 @@ def test_pass_stopped_inside_the_sequence_restarts_at_start_forward(): def test_abandoned_bind_requests_a_device_sync(): - schedule = PrefetchSchedule(4, 2) + schedule = _schedule(4, 2) schedule.start_forward() schedule.bind(0) # The MoE of layer 0 raised before finish: its slot may still be read. @@ -137,7 +142,7 @@ def test_abandoned_bind_requests_a_device_sync(): def test_finish_without_matching_bind_raises(): - schedule = PrefetchSchedule(4, 2) + schedule = _schedule(4, 2) with pytest.raises(RuntimeError): schedule.finish(0) schedule.bind(0) @@ -148,7 +153,7 @@ def test_finish_without_matching_bind_raises(): def test_invalidate_reissues_the_first_layers(): - schedule = PrefetchSchedule(4, 2) + schedule = _schedule(4, 2) _run_pass(schedule, 4) assert schedule.invalidate() == ScheduleStep(False, None, []) assert schedule.start_forward().prefetch == [0, 1] @@ -159,9 +164,87 @@ def test_invalidate_reissues_the_first_layers(): def test_constructor_validates_arguments(): with pytest.raises(ValueError): - PrefetchSchedule(0, 2) + PrefetchSchedule([]) with pytest.raises(ValueError): - PrefetchSchedule(4, 1) + PrefetchSchedule([0, -1]) + with pytest.raises(ValueError): + PrefetchSchedule([0, 1], repeated_orders=[2]) + + +@pytest.mark.parametrize("num_slots", [2, 3]) +@pytest.mark.parametrize("repeats", [1, 2, 4]) +def test_repeated_layer_rebinds_its_resident_copy(num_slots, repeats): + """MTP-Eagle pattern ``..., N, N, N``: every repeat reads the copy still in the slot.""" + num_layers = 5 + last = num_layers - 1 + schedule = _schedule(num_layers, num_slots, repeated=[last]) + model = _SlotModel(num_layers, num_slots) + copies = [] + for _ in range(4): + step = schedule.start_forward() + model.apply(step) + copies += step.prefetch + for order in list(range(num_layers)) + [last] * (repeats - 1): + step = schedule.bind(order) + model.apply(step) + model.read(order) + copies += step.prefetch + step = schedule.finish(order) + if order == last: + # The slot stays held: no release and no copy until a different layer runs. + assert step == ScheduleStep(False, None, []) + model.apply(step) + copies += step.prefetch + assert copies[: 4 * num_layers] == list(range(num_layers)) * 4 + assert schedule.num_restarts == 0 + + +def test_held_slot_is_released_by_the_next_call(): + # 5 layers, 2 slots: order 4 holds slot 0 after it finishes. + schedule = _schedule(5, 2, repeated=[4]) + schedule.start_forward() + for order in range(5): + schedule.bind(order) + schedule.finish(order) + assert schedule.bind(4) == ScheduleStep(False, None, []) + assert schedule.finish(4) == ScheduleStep(False, None, []) + # The next pass releases slot 0 before it reissues the slot. + assert schedule.start_forward() == ScheduleStep(False, 0, [0, 1]) + # A bind of the next layer releases it too (pass without start_forward). + for order in range(5): + schedule.bind(order) + schedule.finish(order) + assert schedule.bind(0) == ScheduleStep(False, 0, [0, 1]) + schedule.finish(0) + # An abandoned repeat still synchronizes before its slot is reused. + for order in range(1, 5): + schedule.bind(order) + schedule.finish(order) + schedule.bind(4) + assert schedule.start_forward() == ScheduleStep(True, 0, [0, 1]) + assert schedule.num_restarts == 1 + + +def test_undeclared_repeat_restarts(): + """Without the declaration a repeat is out of sequence and gathers the layer again.""" + schedule = _schedule(5, 2) + schedule.start_forward() + for order in range(5): + schedule.bind(order) + schedule.finish(order) + assert schedule.bind(4).prefetch == [4] + assert schedule.num_restarts == 1 + + +def test_slots_come_from_the_given_mapping(): + """The schedule uses the slot table it is given, not a formula of its own.""" + slots = [1, 0, 2, 1, 0, 2] + schedule = PrefetchSchedule(slots) + assert [schedule.slot(order) for order in range(len(slots))] == slots + assert schedule.start_forward().prefetch == [0, 1, 2] + for order in range(len(slots)): + schedule.bind(order) + assert schedule.finish(order).release == slots[order] class _SlotModel: @@ -198,12 +281,14 @@ def read(self, order: int) -> None: self.unreleased[slot].add(order) -@pytest.mark.parametrize("seed", range(20)) +@pytest.mark.parametrize("seed", range(40)) def test_randomized_passes_keep_raw_and_war(seed): rng = random.Random(seed) num_layers = rng.randint(1, 9) num_slots = rng.randint(2, 4) - schedule = PrefetchSchedule(num_layers, num_slots) + # Half of the seeds repeat the last layer of each pass, as one-model MTP-Eagle does. + repeated = [num_layers - 1] if seed % 2 else [] + schedule = _schedule(num_layers, num_slots, repeated) model = _SlotModel(num_layers, num_slots) for _ in range(40): kind = rng.random() @@ -218,6 +303,8 @@ def test_randomized_passes_keep_raw_and_war(seed): orders = list(range(stop)) if kind > 0.95 and orders: orders = sorted(rng.sample(orders, rng.randint(1, len(orders)))) + if repeated and orders and orders[-1] == repeated[0]: + orders += repeated * rng.randint(0, 3) for order in orders: model.apply(schedule.bind(order)) model.read(order) diff --git a/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_slots.py b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_slots.py new file mode 100644 index 000000000000..37aff81d7a9e --- /dev/null +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_prefetch_slots.py @@ -0,0 +1,85 @@ +# 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. +"""Number of DWDP prefetch slots: config, slot assignment and slot sizing.""" + +from types import SimpleNamespace + +import pytest +from pydantic import ValidationError + +from tensorrt_llm._torch.modules.dwdp.page_pool import compute_slot_sizes +from tensorrt_llm._torch.modules.dwdp.weight_buffer import WeightBuffer +from tensorrt_llm.llmapi.llm_args import DwdpConfig + +pytestmark = pytest.mark.cpu_only + +LAYERS = [3, 4, 6, 7, 9] + + +def _buffer(num_slots: int) -> WeightBuffer: + # The constructor only records the layout inputs; ``create`` maps memory. + return WeightBuffer( + layer_weight_specs={layer: {} for layer in LAYERS}, + handles=None, + local_start=0, + local_end=4, + dwdp_size=2, + device_id=0, + granularity=1 << 21, + num_slots=num_slots, + ) + + +def test_config_defaults_to_two_slots_and_rejects_fewer(): + assert DwdpConfig().num_prefetch_slots == 2 + assert DwdpConfig(num_prefetch_slots=3).num_prefetch_slots == 3 + with pytest.raises(ValidationError): + DwdpConfig(num_prefetch_slots=1) + + +@pytest.mark.parametrize("num_slots", [2, 3, 4]) +def test_slot_is_moe_order_modulo_num_slots(num_slots): + buffer = _buffer(num_slots) + assert buffer.num_slots == num_slots + assert [buffer.buffer_index_for_layer(layer) for layer in LAYERS] == [ + order % num_slots for order in range(len(LAYERS)) + ] + + +def test_weight_buffer_rejects_fewer_than_two_slots(): + with pytest.raises(ValueError): + _buffer(1) + + +def test_slot_of_a_non_moe_layer_raises(): + with pytest.raises(KeyError): + _buffer(2).buffer_index_for_layer(5) + + +def test_slot_size_is_the_largest_layer_of_the_slot(): + def layout(pre: int, post: int) -> SimpleNamespace: + return SimpleNamespace(pre_size=pre, post_size=post) + + layouts = { + 3: {"w1": layout(10, 0), "w2": layout(5, 5)}, + 4: {"w1": layout(30, 0)}, + 6: {"w1": layout(0, 7)}, + 7: {"w1": layout(1, 1)}, + 9: {"w1": layout(0, 50)}, + } + assignments = {layer: order % 3 for order, layer in enumerate(LAYERS)} + assert compute_slot_sizes(layouts, assignments, 3) == [20, 50, 7] + assignments = {layer: order % 2 for order, layer in enumerate(LAYERS)} + assert compute_slot_sizes(layouts, assignments, 2) == [50, 30] 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 36303c165de1..1ae2e167e0da 100644 --- a/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py +++ b/tests/unittest/_torch/modules/dwdp/test_dwdp_setup_lifecycle.py @@ -88,6 +88,7 @@ def _run_setup(): layer_indices=[3], num_experts_per_worker=2, num_prefetch_experts=2, + num_prefetch_slots=2, ) From ed8079272b526a2c3c219c8db4a6ed424a1db81b Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Tue, 27 Oct 2026 19:25:36 -0700 Subject: [PATCH 7/8] [None][perf] DWDP: validate the SM reserve against the device SM count set_sm_reserve now rejects a reserve that would leave no SM on the current device, so a bad TRTLLM_CUTEDSL_SM_RESERVE fails at DwdpManager setup instead of at the first kda_prefill launch. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- tensorrt_llm/_torch/cute_dsl_utils.py | 15 ++++++++++++++- .../_torch/modules/dwdp/test_fabric_gather.py | 10 ++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/cute_dsl_utils.py b/tensorrt_llm/_torch/cute_dsl_utils.py index e6baaf9f0777..b99362b2cffb 100644 --- a/tensorrt_llm/_torch/cute_dsl_utils.py +++ b/tensorrt_llm/_torch/cute_dsl_utils.py @@ -91,10 +91,23 @@ def install_cutlass_dsl_compatibility() -> None: def set_sm_reserve(reserve: int) -> None: - """Keep ``reserve`` SMs out of persistent grids from now on (0 = none).""" + """Keep ``reserve`` SMs out of persistent grids from now on (0 = none). + + A nonzero ``reserve`` must leave at least one SM of the current device. + """ global _sm_reserve if reserve < 0: raise ValueError(f"SM reserve must be non-negative, got {reserve}") + if reserve: + import torch + + if torch.cuda.is_available(): + sm_count = torch.cuda.get_device_properties( + torch.cuda.current_device()).multi_processor_count + if reserve >= sm_count: + raise ValueError( + f"SM reserve {reserve} must leave at least one of " + f"{sm_count} SMs") if reserve and reserve != _sm_reserve: logger.info( f"[cute_dsl_utils] persistent grids launch with {reserve} SMs reserved" diff --git a/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py b/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py index 5b7f9827c9d6..f7cf0ff253fc 100644 --- a/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py +++ b/tests/unittest/_torch/modules/dwdp/test_fabric_gather.py @@ -121,6 +121,16 @@ def test_sm_reserve_is_set_explicitly(): cute_dsl_utils.set_sm_reserve(-1) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") +def test_sm_reserve_must_leave_an_sm(): + from tensorrt_llm._torch import cute_dsl_utils + + sm_count = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count + with pytest.raises(ValueError, match="must leave at least one"): + cute_dsl_utils.set_sm_reserve(sm_count) + assert cute_dsl_utils.sm_reserve() == 0 + + @pytest.mark.parametrize( "full, cluster_size, reserve, expected", [(148, 1, 30, 118), (74, 2, 30, 59), (2, 2, 30, 1), (0, 2, 30, 0), (148, 1, 0, 148)], From c3442391a4c0b2141d2ba06b830f7243ba9d8fca Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Wed, 28 Oct 2026 01:06:05 -0700 Subject: [PATCH 8/8] [None][test] Stub dwdp_manager in the control-request hang-detector test _handle_control_request now reads self.dwdp_manager to invalidate the DWDP prefetch after a control action. The test builds PyExecutor with __new__, so set the attribute like the other minimal executor stubs do, and also check that a configured manager has its prefetch invalidated. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../unittest/_torch/executor/test_hang_detector_kill.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/executor/test_hang_detector_kill.py b/tests/unittest/_torch/executor/test_hang_detector_kill.py index 542b204399af..aa3e55e12cb8 100644 --- a/tests/unittest/_torch/executor/test_hang_detector_kill.py +++ b/tests/unittest/_torch/executor/test_hang_detector_kill.py @@ -25,7 +25,7 @@ import time import types from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest @@ -203,7 +203,8 @@ def failing_provider(): assert events[-1] == ("detected", None) -def test_control_request_wait_pauses_detector(): +@pytest.mark.parametrize("with_dwdp", [False, True], ids=["no_dwdp", "dwdp"]) +def test_control_request_wait_pauses_detector(with_dwdp): fired = [] executor = PyExecutor.__new__(PyExecutor) executor.control_requests = [SimpleNamespace(control_requires_drain=False, control_id=None)] @@ -214,6 +215,7 @@ def test_control_request_wait_pauses_detector(): executor.dist = SimpleNamespace(pp_size=1, world_size=1) executor.disable_overlap_scheduler = True executor.enable_attention_dp = False + executor.dwdp_manager = Mock() if with_dwdp else None executor.hang_detector = HangDetector(timeout=1, on_detected=lambda: fired.append(1)) def complete_control_action(): @@ -236,6 +238,9 @@ def complete_control_action(): assert fired == [] assert executor.control_requests == [] assert not executor.control_action_done.is_set() + if with_dwdp: + # The control action may have changed expert weights already prefetched for the next step. + executor.dwdp_manager.invalidate_prefetch.assert_called_once_with() def test_propagate_hard_kill_self_sigkills_without_mpi():