From 1c1d49d2913db6b091d309df43a39db7a26b4b99 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 30 Sep 2026 19:00:00 -0700 Subject: [PATCH 01/67] [None][perf] Cut host work between speculative decoding steps One-model speculative decoding spends host time between two decode-step graph replays on small kernel launches and host-to-device copies. - On SM 100, a CUDA graph decode step of an engine whose attention metadata is a plain TrtllmAttentionMetadata writes its per-step inputs (overlap gathers, positions, prompt and KV lengths, KV block offsets) with one CuTe DSL kernel launch (StepInputStage). Eager steps on SM 100 run the overlap gathers as one kernel. - On SM 100, the speculative sampler moves the forward's outputs into its slot stores with one CuTe DSL kernel (SlotScatter). - The sampler's host copies of its stores run on its D2H side stream, and one-model speculative sampling runs on the execution stream. - CUDA graphs of a speculative engine read input ids and positions from the engine's buffers instead of copying them in at each replay. - Copies whose values did not change are skipped (gather ids, slot tables, previous-batch indices), as are the Philox seed and offset uploads of all-greedy batches and the unused num_accepted_draft_tokens upload. Other devices, and calls the kernels do not cover, keep the torch path. Signed-off-by: Vasanth Sabavat --- .../_torch/attention/backends/trtllm.py | 60 +- .../spec_step_copies/__init__.py | 18 + .../cute_dsl_kernels/spec_step_copies/op.py | 433 +++++++++ .../spec_step_copies_kernel.py | 362 ++++++++ .../_torch/pyexecutor/cuda_graph_runner.py | 42 +- .../kv_cache/kv_cache_manager_v2.py | 32 + .../_torch/pyexecutor/model_engine.py | 219 +++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 13 +- tensorrt_llm/_torch/speculative/interface.py | 7 +- .../_torch/speculative/spec_sampler_base.py | 145 ++- tensorrt_llm/_utils.py | 28 + .../test_lists/test-db/l0_b200.yml | 4 + .../cute_dsl_kernels/test_spec_step_copies.py | 822 ++++++++++++++++++ .../test_copy_to_device_if_changed.py | 119 +++ .../test_cuda_graph_capture_replay.py | 79 ++ .../executor/test_pytorch_model_engine.py | 128 +++ tests/unittest/_torch/helpers.py | 13 +- .../hw_agnostic/test_rng_window_counter.py | 38 +- 18 files changed, 2400 insertions(+), 162 deletions(-) create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/__init__.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py create mode 100644 tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py create mode 100644 tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py create mode 100644 tests/unittest/_torch/executor/test_copy_to_device_if_changed.py diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 6f792d64dd5a..3b835cfc3ef0 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -25,6 +25,7 @@ if TYPE_CHECKING: from tensorrt_llm.mapping import Mapping + from ...cute_dsl_kernels.spec_step_copies.op import StepInputStage from ...model_config import ModelConfig from ...speculative.interface import SpecMetadata from ...speculative.spec_tree_manager import SpecTreeManager @@ -123,6 +124,10 @@ class TrtllmAttentionMetadata(AttentionMetadata): workspace: Optional[torch.Tensor] = None cuda_graph_workspace: Optional[torch.Tensor] = None workspace_reclaimable: bool = field(default=True, init=False) + # Set by the model engine around prepare() on a CUDA graph decode step: + # the per-step host-to-device copies are staged on it (a StepInputStage) + # and written by its commit, instead of being copied here. + h2d_stage: Optional["StepInputStage"] = field(default=None, init=False) # TrtllmAttention needs to know the beam width to access to the cache indirection buffer, # when beam search is enabled. @@ -873,6 +878,29 @@ def restore_after_draft_forward(self, saved_state: dict | None) -> None: """Restore backend state modified for draft-forward execution.""" return None + def _copy_block_offsets(self, manager, dst: torch.Tensor, + max_blocks: Optional[int]) -> None: + """``manager``'s block-offset copy for this batch, staged on + ``h2d_stage`` when the engine set one and the manager can stage it + (``stage_batch_block_offsets``).""" + stage_copy = getattr(manager, "stage_batch_block_offsets", None) + if self.h2d_stage is not None and stage_copy is not None and stage_copy( + self.h2d_stage, dst, self.request_ids, self.beam_width, + self.num_contexts, self.num_seqs): + return + manager.copy_batch_block_offsets(dst, + self.request_ids, + self.beam_width, + self.num_contexts, + self.num_seqs, + max_blocks=max_blocks) + + def _copy_step_input(self, dst: torch.Tensor, src: torch.Tensor) -> None: + """``dst.copy_(src)`` for a per-step input, staged on ``h2d_stage`` + when the engine set one.""" + if self.h2d_stage is None or not self.h2d_stage.copy(dst, src): + dst.copy_(maybe_pin_memory(src), non_blocking=True) + def prepare(self) -> None: super().prepare() # Recomputed on first use this iteration; see mla_prepare_scheduler_buffers. @@ -901,8 +929,8 @@ def prepare(self) -> None: device='cpu', ) self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens) - self.prompt_lens_cuda[:self.num_seqs].copy_( - self.prompt_lens_cpu[:self.num_seqs], non_blocking=True) + self._copy_step_input(self.prompt_lens_cuda[:self.num_seqs], + self.prompt_lens_cpu[:self.num_seqs]) # number of tokens in the kv cache for each sequence in the batch cached_token_lens = torch.tensor( @@ -955,9 +983,8 @@ def prepare(self) -> None: # the sequence length including the cached tokens and the input tokens. self.kv_lens[:self.num_seqs].copy_( kv_lens + self.kv_cache_params.num_extra_kv_tokens) - self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory( - kv_lens[:self.num_seqs]), - non_blocking=True) + self._copy_step_input(self.kv_lens_cuda[:self.num_seqs], + kv_lens[:self.num_seqs]) # total kv lens for context requests and generation requests, without extra tokens self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item() self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self. @@ -995,24 +1022,14 @@ def prepare(self) -> None: if not spec_active and self.kv_cache_manager.tokens_per_block: max_blocks = ceil_div(max_kv_len, self.kv_cache_manager.tokens_per_block) - self.kv_cache_manager.copy_batch_block_offsets( - self.kv_cache_block_offsets, - self.request_ids, - self.beam_width, - self.num_contexts, - self.num_seqs, - max_blocks=max_blocks) + self._copy_block_offsets(self.kv_cache_manager, + self.kv_cache_block_offsets, max_blocks) # Also prepare draft KV cache block offsets if draft_kv_cache_manager exists if self.draft_kv_cache_manager is not None: - # Use the wrapper method which works for both V1 and V2 - self.draft_kv_cache_manager.copy_batch_block_offsets( - self.draft_kv_cache_block_offsets, - self.request_ids, - self.beam_width, - self.num_contexts, - self.num_seqs, - max_blocks=max_blocks) + self._copy_block_offsets(self.draft_kv_cache_manager, + self.draft_kv_cache_block_offsets, + max_blocks) # Don't pass self.kv_lens as kv_lens here because it includes extra # tokens. Use the actual KV length (without extra tokens) for @@ -1319,6 +1336,9 @@ def prepare_context_mla_with_cached_kv(self, self.max_ctx_cached_token_len = 0 self.max_ctx_kv_len = 0 self.max_ctx_seq_len = 0 + # The indptrs below are read only by the context MLA kernels, and + # every batch with context requests rewrites them. + return torch.cumsum(cached_token_lens[:self.num_contexts], dim=0, dtype=torch.int64, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/__init__.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/__init__.py new file mode 100644 index 000000000000..459e4e88bd94 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The one-model speculative decoding step's eager copy passes as one CuTe DSL kernel each. + +Import :mod:`.op` for the host side; nothing is imported eagerly here so that the CuTe DSL dependency stays optional. +""" diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py new file mode 100644 index 000000000000..c0b2e56c3658 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py @@ -0,0 +1,433 @@ +# 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 of the one-model speculative decoding step's copy kernels (``spec_step_copies_kernel``). + +``SlotScatter`` moves a step's per-row outputs into the sampler's slot stores (one launch for the sampler's four +``index_copy_``). ``StepInputStage`` writes a decode step's per-step inputs with one launch: the overlap scheduler's +gathers from those stores, host-to-device copies whose values the kernel reads from a pinned host record, and KV cache +block-offset copies. Each kernel takes every size and address as a launch argument, so it is compiled once per process, +on the first launch, which must happen outside CUDA-graph capture. The kernels are called through TVM-FFI: a launch +costs a few microseconds of host time. + +Callers use them only where ``is_supported()`` holds. A call whose arguments are outside what the kernel covers launches +and stages nothing and returns False; the caller then does that work itself. +""" + +from __future__ import annotations + +from typing import Any, ClassVar + +import torch + +from ...._utils import get_sm_version, prefer_pinned +from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + +def is_supported() -> bool: + """Whether the kernels run on the current device: SM 100 with the CuTe DSL installed.""" + return IS_CUTLASS_DSL_AVAILABLE and torch.cuda.is_available() and get_sm_version() == 100 + + +def _is_i32(t: torch.Tensor) -> bool: + """A contiguous CUDA int32 tensor.""" + return t.is_cuda and t.dtype == torch.int32 and t.is_contiguous() + + +def _is_pinned_i32(t: torch.Tensor) -> bool: + """A contiguous pinned host int32 tensor.""" + return not t.is_cuda and t.dtype == torch.int32 and t.is_contiguous() and t.is_pinned() + + +def _compile(name: str, args: tuple): + """Compile kernel ``name`` for its scalar ``args`` with TVM-FFI; a call then takes ``args`` and the CUDA stream + handle to launch on.""" + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"spec_step_copies: the {name} kernel must run once outside CUDA-graph capture first " + "(it compiles on its first launch)." + ) + import cutlass.cute as cute + + from . import spec_step_copies_kernel as kernel + + return cute.compile( + getattr(kernel, name), *args, cute.runtime.make_fake_stream(), options="--enable-tvm-ffi" + ) + + +def _gather_shapes_ok( + store_next_new_tokens: torch.Tensor, + store_next_draft: torch.Tensor, + store_lens: torch.Tensor, + rows: int, + tokens_per_row: int, + draft_width: int, +) -> bool: + """The overlap gather's stores are [width, slots, 1], [slots, draft width] and [slots], wide enough for a row.""" + num_slots = store_lens.numel() + return ( + store_lens.dim() == 1 + and store_next_new_tokens.dim() == 3 + and store_next_new_tokens.shape[1:] == (num_slots, 1) + and store_next_draft.dim() == 2 + and store_next_draft.shape[0] == num_slots + and 0 < tokens_per_row <= store_next_new_tokens.shape[0] + and 0 <= draft_width < tokens_per_row + and draft_width <= store_next_draft.shape[1] + and rows >= 0 + ) + + +class SlotScatter: + """The speculative sampler's store update as one kernel. + + ``store[..., slots[r]] = outputs[row_begin + r]`` for every row ``r``, each row padded with zeros or cut to its + store's width. + """ + + _kernel: ClassVar[Any] = None + + def scatter( + self, + outputs: dict[str, torch.Tensor], + row_begin: int, + rows: int, + slots: torch.Tensor, + store_new_tokens: torch.Tensor, + store_next_new_tokens: torch.Tensor, + store_lens: torch.Tensor, + store_next_draft: torch.Tensor, + ) -> bool: + """Launch the store update on the current stream. + + Args: + outputs: The forward's ``new_tokens`` [rows_total, a], ``next_new_tokens`` [rows_total, b], + ``new_tokens_lens`` [rows_total] and ``next_draft_tokens`` [rows_total, c], int32. + row_begin: The first row of ``outputs`` to move. + rows: The number of rows to move. + slots: The distinct slot of each moved row, int32 [>= rows] on the device. + store_new_tokens: int32 [width, num_slots, 1]. + store_next_new_tokens: int32 [width, num_slots, 1]. + store_lens: int32 [num_slots]. + store_next_draft: int32 [num_slots, width]. + + Returns: + False, launching nothing, when an argument is outside what the kernel covers (a non-contiguous or + non-int32 tensor, a store of another layout, fewer output rows or slots than ``rows``); True otherwise. + """ + if rows == 0: + return True + new_tokens = outputs["new_tokens"] + next_new_tokens = outputs["next_new_tokens"] + lens = outputs["new_tokens_lens"] + next_draft = outputs["next_draft_tokens"] + num_slots = store_lens.numel() + tensors = ( + new_tokens, + next_new_tokens, + lens, + next_draft, + slots, + store_new_tokens, + store_next_new_tokens, + store_lens, + store_next_draft, + ) + if not ( + all(_is_i32(t) for t in tensors) + and new_tokens.dim() == next_new_tokens.dim() == next_draft.dim() == 2 + and lens.dim() == store_lens.dim() == 1 + and min(t.shape[0] for t in (new_tokens, next_new_tokens, lens, next_draft)) + >= row_begin + rows + and row_begin >= 0 + and slots.numel() >= rows + and store_new_tokens.dim() == store_next_new_tokens.dim() == 3 + and store_new_tokens.shape[1:] == (num_slots, 1) + and store_next_new_tokens.shape[1:] == (num_slots, 1) + and store_next_draft.dim() == 2 + and store_next_draft.shape[0] == num_slots + ): + return False + columns = max( + store_new_tokens.shape[0], store_next_new_tokens.shape[0], store_next_draft.shape[1], 1 + ) + args = ( + new_tokens.data_ptr(), + next_new_tokens.data_ptr(), + lens.data_ptr(), + next_draft.data_ptr(), + slots.data_ptr(), + store_new_tokens.data_ptr(), + store_next_new_tokens.data_ptr(), + store_lens.data_ptr(), + store_next_draft.data_ptr(), + rows, + row_begin, + columns, + num_slots, + store_new_tokens.shape[0], + new_tokens.shape[1], + store_next_new_tokens.shape[0], + next_new_tokens.shape[1], + store_next_draft.shape[1], + next_draft.shape[1], + ) + if SlotScatter._kernel is None: + SlotScatter._kernel = _compile("scatter", args) + SlotScatter._kernel(*args, torch.cuda.current_stream().cuda_stream) + return True + + +class StepInputStage: + """One decode step's per-step device inputs, written by one ``stage_kernel`` launch. + + A step calls ``begin()``, then stages the overlap gathers (``gather``), host-to-device copies whose values go into a + pinned host record that the kernel reads in place (``copy``, up to ``max_copies``) and KV cache block-offset copies + (``block_copy``, up to ``max_block_copies``), then ``commit()`` launches one kernel on the current stream that + performs all of them. The record is reused every step: the first ``copy`` after a commit that staged copies waits + for that commit's kernel (in the overlap loop it has run by then: the host is at most one step ahead). + + Args: + capacity: The record's size in int32 values, the most a step's copies stage in total. + """ + + _kernel: ClassVar[Any] = None + + def __init__(self, capacity: int) -> None: + from . import spec_step_copies_kernel as kernel + + self.max_copies = kernel.STAGED_COPIES + self.max_block_copies = kernel.BLOCK_COPIES + self.record = torch.empty((capacity,), dtype=torch.int32, pin_memory=prefer_pinned()) + self._copies: list[tuple[int, int, int]] = [] + self._block_copies: list[tuple[int, ...]] = [] + self._gather: tuple[int, ...] | None = None + self._used = 0 + self._done: torch.cuda.Event | None = None + self._record_free = True + + def begin(self) -> None: + """Start a step: drop anything staged by a step that did not commit (its preparation raised).""" + self._copies = [] + self._block_copies = [] + self._gather = None + self._used = 0 + + def gather( + self, + store_next_new_tokens: torch.Tensor, + store_next_draft: torch.Tensor, + store_lens: torch.Tensor, + slots: torch.Tensor, + pos_indices: torch.Tensor, + rows: int, + tokens_per_row: int, + draft_width: int, + input_ids: torch.Tensor, + input_begin: int, + draft_tokens: torch.Tensor, + draft_begin: int, + pos_offsets: torch.Tensor, + pos_begin: int, + kv_offsets: torch.Tensor, + kv_begin: int, + ) -> bool: + """Stage the overlap scheduler's gathers of the rows whose request ran in the previous step. + + For ``r < rows`` with ``s = slots[r]`` and ``j < tokens_per_row``: + ``input_ids[input_begin + r * tokens_per_row + j] = store_next_new_tokens[j, s, 0]``, + ``pos_offsets[pos_begin + r * tokens_per_row + j] = store_lens[pos_indices[r * tokens_per_row + j]]``, + ``draft_tokens[draft_begin + r * draft_width + j] = store_next_draft[s, j]`` for ``j < draft_width`` and + ``kv_offsets[kv_begin + r] = store_lens[s] - tokens_per_row``. Every tensor is int32 on the device; the + stores are [width, slots, 1], [slots, width] and [slots]. + + Returns: + False, staging nothing, when an argument is outside what the kernel covers or a gather is already + staged; True otherwise. + """ + tensors = ( + store_next_new_tokens, + store_next_draft, + store_lens, + slots, + pos_indices, + input_ids, + draft_tokens, + pos_offsets, + kv_offsets, + ) + if not ( + self._gather is None + and all(_is_i32(t) for t in tensors) + and _gather_shapes_ok( + store_next_new_tokens, + store_next_draft, + store_lens, + rows, + tokens_per_row, + draft_width, + ) + and slots.numel() >= rows + and pos_indices.numel() >= rows * tokens_per_row + and input_ids.numel() >= input_begin + rows * tokens_per_row + and draft_tokens.numel() >= draft_begin + rows * draft_width + and pos_offsets.numel() >= pos_begin + rows * tokens_per_row + and kv_offsets.numel() >= kv_begin + rows + and min(input_begin, draft_begin, pos_begin, kv_begin) >= 0 + ): + return False + if rows == 0: + return True + self._gather = ( + store_next_new_tokens.data_ptr(), + store_next_draft.data_ptr(), + store_lens.data_ptr(), + slots.data_ptr(), + pos_indices.data_ptr(), + input_ids.data_ptr(), + draft_tokens.data_ptr(), + pos_offsets.data_ptr(), + kv_offsets.data_ptr(), + rows, + tokens_per_row, + draft_width, + store_lens.numel(), + store_next_draft.shape[1], + input_begin, + draft_begin, + pos_begin, + kv_begin, + ) + return True + + def copy(self, dst: torch.Tensor, values: torch.Tensor) -> bool: + """Stage ``dst[:n] = values`` for the ``n`` int32 host ``values``; ``dst`` is a 1-D int32 device tensor. + + Returns: + False, staging nothing, when ``dst`` is not a contiguous 1-D CUDA int32 tensor of at least ``n`` + values, the step already staged ``max_copies`` copies or the record cannot hold the values; True + otherwise. + """ + host = values.reshape(-1) + n = host.numel() + if not ( + _is_i32(dst) + and dst.dim() == 1 + and n <= dst.numel() + and host.dtype == torch.int32 + and not host.is_cuda + and len(self._copies) < self.max_copies + and self._used + n <= self.record.numel() + and self.record.is_pinned() + ): + return False + if n == 0: + return True + if not self._record_free: + self._done.synchronize() + self._record_free = True + self.record[self._used : self._used + n].copy_(host) + self._copies.append((dst.data_ptr(), self._used, n)) + self._used += n + return True + + def block_copy( + self, + offsets: torch.Tensor, + table: torch.Tensor, + copy_index: torch.Tensor, + index_scales: torch.Tensor, + kv_offset: torch.Tensor, + ) -> bool: + """Stage a KV cache manager's block-offset copy (``copy_batch_block_offsets_to_device``). + + Args: + offsets: The device block offsets, int32 [pools, seqs_cap, 2, blocks]. + table: The pinned host page table, int32 [pools, table_seqs, 2, blocks]. + copy_index: The table row of each of the step's sequences, pinned host int32 [seqs]. + index_scales: The per-pool page index scale, pinned host int32 [pools]. + kv_offset: The per-pool V offset, pinned host int32 [pools]. + + For every pool and sequence the K row is ``index_scale * page`` and the V row ``index_scale * page + + kv_offset`` over the table row ``copy_index[sequence]``; a bad page (-1) gives 0. + + Returns: + False, staging nothing, when an argument is outside what the kernel covers or the step already staged + ``max_block_copies`` block copies; True otherwise. + """ + if not ( + _is_i32(offsets) + and offsets.dim() == 4 + and table.dim() == 4 + and all(_is_pinned_i32(t) for t in (table, copy_index, index_scales, kv_offset)) + and len(self._block_copies) < self.max_block_copies + ): + return False + pools, table_seqs, kv, blocks = table.shape + seqs = copy_index.numel() + if not ( + kv == 2 + and offsets.shape[0] >= pools + and offsets.shape[1] >= seqs + and offsets.shape[2] == 2 + and offsets.shape[3] == blocks + and index_scales.numel() >= pools + and kv_offset.numel() >= pools + ): + return False + if pools * seqs * blocks == 0: + return True + self._block_copies.append( + ( + table.data_ptr(), + offsets.data_ptr(), + copy_index.data_ptr(), + index_scales.data_ptr(), + kv_offset.data_ptr(), + pools, + table_seqs, + offsets.shape[1], + blocks, + seqs, + ) + ) + return True + + def commit(self) -> None: + """Launch the staged gathers and copies as one kernel on the current stream, then start the next staging.""" + if self._gather is None and not self._copies and not self._block_copies: + return + gather = self._gather or (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0) + copies = self._copies + [(0, 0, 0)] * (self.max_copies - len(self._copies)) + blocks = self._block_copies + [(0, 0, 0, 0, 0, 0, 0, 0, 0, 0)] * ( + self.max_block_copies - len(self._block_copies) + ) + args = ( + *gather, + self.record.data_ptr(), + *(c[0] for c in copies), + *(c[1] for c in copies), + *(c[2] for c in copies), + *(v for b in blocks for v in b), + ) + if StepInputStage._kernel is None: + StepInputStage._kernel = _compile("stage", args) + StepInputStage._kernel(*args, torch.cuda.current_stream().cuda_stream) + if self._copies: + if self._done is None: + self._done = torch.cuda.Event() + self._done.record() + self._record_free = False + self.begin() diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py new file mode 100644 index 000000000000..9780847687b4 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py @@ -0,0 +1,362 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The one-model speculative decoding step's eager copy passes, one kernel each. + +* ``scatter_kernel``: the sampler moves the forward's per-row outputs into its slot-indexed stores (``index_copy_`` + of new tokens, next new tokens, accepted lengths and next draft tokens, each padded or cut to its store width). +* ``stage_kernel``: one launch writes a decode step's per-step inputs: the overlap scheduler's gathers from those + stores for the rows whose request ran in the previous step (``index_select`` into input ids and draft tokens by + slot, into the position offsets by the per-token index list, and into the KV-length offsets by slot, minus the + tokens per step), up to four host-to-device copies whose values it reads in place from a pinned host record, and up + to two KV cache block-offset copies. + +Every buffer is int32 and passed as a raw address, so a call passes integers only. Row and element offsets are +arguments. One thread per element. +""" + +from __future__ import annotations + +import cuda.bindings.driver as cuda_driver +import cutlass +import cutlass.cute as cute + +THREADS = 128 + + +def _i32_at(address, index): + """A one-element int32 tensor at ``address + 4 * index`` (trace-time helper; global or mapped host memory).""" + ptr = cute.make_ptr( + cutlass.Int32, + address + cutlass.Int64(index) * cutlass.Int64(4), + cute.AddressSpace.gmem, + assumed_align=4, + ) + return cute.make_tensor(ptr, cute.make_layout((1,))) + + +@cute.kernel +def scatter_kernel( + out_new_tokens: cutlass.Int64, # int32 [rows_total, out_new_width], the forward's accepted tokens + out_next_new_tokens: cutlass.Int64, # int32 [rows_total, out_next_width] + out_lens: cutlass.Int64, # int32 [rows_total] + out_next_draft: cutlass.Int64, # int32 [rows_total, out_draft_width] + slots: cutlass.Int64, # int32 [rows] + store_new_tokens: cutlass.Int64, # int32 [new_width, num_slots] (token-major) + store_next_new_tokens: cutlass.Int64, # int32 [next_width, num_slots] + store_lens: cutlass.Int64, # int32 [num_slots] + store_next_draft: cutlass.Int64, # int32 [num_slots, draft_width] + rows: cutlass.Int32, + row_begin: cutlass.Int32, + columns: cutlass.Int32, + num_slots: cutlass.Int32, + new_width: cutlass.Int32, + out_new_width: cutlass.Int32, + next_width: cutlass.Int32, + out_next_width: cutlass.Int32, + draft_width: cutlass.Int32, + out_draft_width: cutlass.Int32, +): + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + t = bidx * cutlass.Int32(THREADS) + tidx + if t < rows * columns: + row = t // columns + col = t - row * columns + src_row = row_begin + row + slot = _i32_at(slots, row)[0] + zero = cutlass.Int32(0) + if col < new_width: + value = zero + if col < out_new_width: + value = _i32_at(out_new_tokens, src_row * out_new_width + col)[0] + _i32_at(store_new_tokens, col * num_slots + slot)[0] = value + if col < next_width: + value = zero + if col < out_next_width: + value = _i32_at(out_next_new_tokens, src_row * out_next_width + col)[0] + _i32_at(store_next_new_tokens, col * num_slots + slot)[0] = value + if col < draft_width: + value = zero + if col < out_draft_width: + value = _i32_at(out_next_draft, src_row * out_draft_width + col)[0] + _i32_at(store_next_draft, slot * draft_width + col)[0] = value + if col == zero: + _i32_at(store_lens, slot)[0] = _i32_at(out_lens, src_row)[0] + + +@cute.jit +def scatter( + out_new_tokens: cutlass.Int64, + out_next_new_tokens: cutlass.Int64, + out_lens: cutlass.Int64, + out_next_draft: cutlass.Int64, + slots: cutlass.Int64, + store_new_tokens: cutlass.Int64, + store_next_new_tokens: cutlass.Int64, + store_lens: cutlass.Int64, + store_next_draft: cutlass.Int64, + rows: cutlass.Int32, + row_begin: cutlass.Int32, + columns: cutlass.Int32, + num_slots: cutlass.Int32, + new_width: cutlass.Int32, + out_new_width: cutlass.Int32, + next_width: cutlass.Int32, + out_next_width: cutlass.Int32, + draft_width: cutlass.Int32, + out_draft_width: cutlass.Int32, + stream: cuda_driver.CUstream, +) -> None: + scatter_kernel( + out_new_tokens, out_next_new_tokens, out_lens, out_next_draft, slots, store_new_tokens, + store_next_new_tokens, store_lens, store_next_draft, rows, row_begin, columns, num_slots, new_width, + out_new_width, next_width, out_next_width, draft_width, out_draft_width, + ).launch( + grid=[(rows * columns + THREADS - 1) // THREADS, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) # fmt: skip + + +STAGED_COPIES = 4 +BLOCK_COPIES = 2 + + +@cute.kernel +def stage_kernel( + store_next_new_tokens: cutlass.Int64, # int32 [next_width, num_slots] (token-major) + store_next_draft: cutlass.Int64, # int32 [num_slots, draft_stride] + store_lens: cutlass.Int64, # int32 [num_slots] + slots: cutlass.Int64, # int32 [rows] + pos_indices: cutlass.Int64, # int32 [rows * tokens_per_row] + input_ids: cutlass.Int64, + draft_tokens: cutlass.Int64, + pos_offsets: cutlass.Int64, + kv_offsets: cutlass.Int64, + rows: cutlass.Int32, + tokens_per_row: cutlass.Int32, + draft_width: cutlass.Int32, + num_slots: cutlass.Int32, + draft_stride: cutlass.Int32, + input_begin: cutlass.Int32, + draft_begin: cutlass.Int32, + pos_begin: cutlass.Int32, + kv_begin: cutlass.Int32, + record: cutlass.Int64, # int32 values in pinned host memory, read in place + dst0: cutlass.Int64, + dst1: cutlass.Int64, + dst2: cutlass.Int64, + dst3: cutlass.Int64, + off0: cutlass.Int32, + off1: cutlass.Int32, + off2: cutlass.Int32, + off3: cutlass.Int32, + n0: cutlass.Int32, + n1: cutlass.Int32, + n2: cutlass.Int32, + n3: cutlass.Int32, + table0: cutlass.Int64, # block copy 0: int32 [pools, table_seqs, 2, blocks] in pinned host memory + offsets0: cutlass.Int64, # int32 [pools, offsets_seqs, 2, blocks] (device) + copy_index0: cutlass.Int64, # int32 [seqs], pinned host + index_scales0: cutlass.Int64, # int32 [pools], pinned host + kv_offset0: cutlass.Int64, # int32 [pools], pinned host + pools0: cutlass.Int32, + table_seqs0: cutlass.Int32, + offsets_seqs0: cutlass.Int32, + blocks0: cutlass.Int32, + seqs0: cutlass.Int32, + table1: cutlass.Int64, + offsets1: cutlass.Int64, + copy_index1: cutlass.Int64, + index_scales1: cutlass.Int64, + kv_offset1: cutlass.Int64, + pools1: cutlass.Int32, + table_seqs1: cutlass.Int32, + offsets_seqs1: cutlass.Int32, + blocks1: cutlass.Int32, + seqs1: cutlass.Int32, +): + """The overlap gathers for threads [0, rows * tokens_per_row) (for row r, slot s = slots[r] and token j: + ``input_ids[input_begin + r * tokens_per_row + j] = store_next_new_tokens[j, s]``, ``pos_offsets[pos_begin + r * + tokens_per_row + j] = store_lens[pos_indices[r * tokens_per_row + j]]``, ``draft_tokens[draft_begin + r * + draft_width + j] = store_next_draft[s, j]`` for j < draft_width, ``kv_offsets[kv_begin + r] = store_lens[s] - + tokens_per_row``); then up to four staged copies ``dst_i[j] = record[off_i + j]`` for j < n_i; then up to two KV + block-offset copies, each the + ``copyBatchBlockOffsetsToDeviceKernel`` of kvCacheManagerV2Utils.cu (the K and V rows of every (pool, sequence) + from the pinned host table's row ``copy_index[sequence]``: ``index_scale * page``, and ``+ kv_offset`` for V; a + bad page (-1) gives 0). One thread per int32 (per K / V pair for the block copies).""" + tidx, _, _ = cute.arch.thread_idx() + bidx, _, _ = cute.arch.block_idx() + t = bidx * cutlass.Int32(THREADS) + tidx + gathered = rows * tokens_per_row + if t < gathered: + row = t // tokens_per_row + col = t - row * tokens_per_row + slot = _i32_at(slots, row)[0] + _i32_at(input_ids, input_begin + t)[0] = _i32_at( + store_next_new_tokens, col * num_slots + slot + )[0] + _i32_at(pos_offsets, pos_begin + t)[0] = _i32_at(store_lens, _i32_at(pos_indices, t)[0])[0] + if col < draft_width: + _i32_at(draft_tokens, draft_begin + row * draft_width + col)[0] = _i32_at( + store_next_draft, slot * draft_stride + col + )[0] + if col == cutlass.Int32(0): + _i32_at(kv_offsets, kv_begin + row)[0] = _i32_at(store_lens, slot)[0] - tokens_per_row + else: + u = t - gathered + dst = cutlass.Int64(0) + src = cutlass.Int32(0) + j = cutlass.Int32(0) + active = cutlass.Int32(0) + if u < n0: + dst = dst0 + src = off0 + u + j = u + active = cutlass.Int32(1) + elif u < n0 + n1: + dst = dst1 + src = off1 + u - n0 + j = u - n0 + active = cutlass.Int32(1) + elif u < n0 + n1 + n2: + dst = dst2 + src = off2 + u - n0 - n1 + j = u - n0 - n1 + active = cutlass.Int32(1) + elif u < n0 + n1 + n2 + n3: + dst = dst3 + src = off3 + u - n0 - n1 - n2 + j = u - n0 - n1 - n2 + active = cutlass.Int32(1) + if active == cutlass.Int32(1): + _i32_at(dst, j)[0] = _i32_at(record, src)[0] + v = u - n0 - n1 - n2 - n3 + per0 = pools0 * seqs0 * blocks0 + per1 = pools1 * seqs1 * blocks1 + if v >= cutlass.Int32(0): + if v < per0 + per1: + table = table0 + offsets = offsets0 + copy_index = copy_index0 + index_scales = index_scales0 + kv_offset = kv_offset0 + table_seqs = table_seqs0 + offsets_seqs = offsets_seqs0 + blocks = blocks0 + seqs = seqs0 + w = v + if v >= per0: + table = table1 + offsets = offsets1 + copy_index = copy_index1 + index_scales = index_scales1 + kv_offset = kv_offset1 + table_seqs = table_seqs1 + offsets_seqs = offsets_seqs1 + blocks = blocks1 + seqs = seqs1 + w = v - per0 + pool = w // (seqs * blocks) + rest = w - pool * seqs * blocks + seq = rest // blocks + block = rest - seq * blocks + row = pool * table_seqs + _i32_at(copy_index, seq)[0] + page = _i32_at(table, row * cutlass.Int32(2) * blocks + block)[0] + key = cutlass.Int32(0) + value = cutlass.Int32(0) + if page != cutlass.Int32(-1): + key = _i32_at(index_scales, pool)[0] * page + value = key + _i32_at(kv_offset, pool)[0] + out = (pool * offsets_seqs + seq) * cutlass.Int32(2) * blocks + block + _i32_at(offsets, out)[0] = key + _i32_at(offsets, out + blocks)[0] = value + + +@cute.jit +def stage( + store_next_new_tokens: cutlass.Int64, + store_next_draft: cutlass.Int64, + store_lens: cutlass.Int64, + slots: cutlass.Int64, + pos_indices: cutlass.Int64, + input_ids: cutlass.Int64, + draft_tokens: cutlass.Int64, + pos_offsets: cutlass.Int64, + kv_offsets: cutlass.Int64, + rows: cutlass.Int32, + tokens_per_row: cutlass.Int32, + draft_width: cutlass.Int32, + num_slots: cutlass.Int32, + draft_stride: cutlass.Int32, + input_begin: cutlass.Int32, + draft_begin: cutlass.Int32, + pos_begin: cutlass.Int32, + kv_begin: cutlass.Int32, + record: cutlass.Int64, + dst0: cutlass.Int64, + dst1: cutlass.Int64, + dst2: cutlass.Int64, + dst3: cutlass.Int64, + off0: cutlass.Int32, + off1: cutlass.Int32, + off2: cutlass.Int32, + off3: cutlass.Int32, + n0: cutlass.Int32, + n1: cutlass.Int32, + n2: cutlass.Int32, + n3: cutlass.Int32, + table0: cutlass.Int64, + offsets0: cutlass.Int64, + copy_index0: cutlass.Int64, + index_scales0: cutlass.Int64, + kv_offset0: cutlass.Int64, + pools0: cutlass.Int32, + table_seqs0: cutlass.Int32, + offsets_seqs0: cutlass.Int32, + blocks0: cutlass.Int32, + seqs0: cutlass.Int32, + table1: cutlass.Int64, + offsets1: cutlass.Int64, + copy_index1: cutlass.Int64, + index_scales1: cutlass.Int64, + kv_offset1: cutlass.Int64, + pools1: cutlass.Int32, + table_seqs1: cutlass.Int32, + offsets_seqs1: cutlass.Int32, + blocks1: cutlass.Int32, + seqs1: cutlass.Int32, + stream: cuda_driver.CUstream, +) -> None: + total = ( + rows * tokens_per_row + + n0 + + n1 + + n2 + + n3 + + pools0 * seqs0 * blocks0 + + pools1 * seqs1 * blocks1 + ) + stage_kernel( + store_next_new_tokens, store_next_draft, store_lens, slots, pos_indices, input_ids, draft_tokens, pos_offsets, + kv_offsets, rows, tokens_per_row, draft_width, num_slots, draft_stride, input_begin, draft_begin, pos_begin, + kv_begin, record, dst0, dst1, dst2, dst3, off0, off1, off2, off3, n0, n1, n2, n3, table0, offsets0, + copy_index0, index_scales0, kv_offset0, pools0, table_seqs0, offsets_seqs0, blocks0, seqs0, table1, offsets1, + copy_index1, index_scales1, kv_offset1, pools1, table_seqs1, offsets_seqs1, blocks1, seqs1, + ).launch( + grid=[(total + THREADS - 1) // THREADS, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 4614477976e0..dc78b8a9f967 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -136,6 +136,13 @@ class CUDAGraphRunnerConfig: sparse_attention_config: Optional[BaseSparseAttentionConfig] = None enable_encoder_decoder_mixed_cuda_graph: bool = False enable_in_graph_sampling: bool = False + static_input_ids: Optional[torch.Tensor] = None + """The engine's own int32 input-id buffer. When given (with + ``static_position_ids``), the graphs' static inputs are views of the + engine's buffers: every write the engine makes is already the graphs' + input, and replay copies nothing. For engines that rewrite every token's + input each step.""" + static_position_ids: Optional[torch.Tensor] = None class CUDAGraphRunner: @@ -203,6 +210,25 @@ def _create_shared_static_tensors(self): max_total_tokens = self.config.max_num_tokens max_total_tokens = min(max_total_tokens, self.config.max_num_tokens) + if self.config.static_input_ids is not None: + input_ids = self.config.static_input_ids + position_ids = self.config.static_position_ids + if (self.config.use_mrope or position_ids is None + or input_ids.dtype != torch.int32 + or position_ids.dtype != torch.int32 or input_ids.dim() != 1 + or position_ids.dim() != 1 or not input_ids.is_contiguous() + or not position_ids.is_contiguous() + or min(input_ids.numel(), + position_ids.numel()) < max_total_tokens): + raise ValueError( + "CUDA graph static inputs from the engine's buffers need " + "two contiguous 1-D int32 buffers of at least " + f"{max_total_tokens} tokens and no MRoPE.") + self.shared_static_tensors = { + "input_ids": input_ids[:max_total_tokens], + "position_ids": position_ids[:max_total_tokens].unsqueeze(0), + } + return self.shared_static_tensors = { "input_ids": torch.ones((max_total_tokens, ), device="cuda", dtype=torch.int32), @@ -732,10 +758,12 @@ def capture(self, # Do not keep the eager result live from this runner across graph # setup/capture; release its reference before entering. output = None + # A captured forward does not run, so there is no in-place input + # change to undo (postprocess_fn is for the warmup forwards above). + # Undoing one here would shift the static inputs, which are the + # engine's own buffers when static_input_ids is set. with torch.cuda.graph(graph, pool=self.memory_pool): output = forward_fn(capture_inputs) - if postprocess_fn is not None: - postprocess_fn(capture_inputs) _restore_spec_decode_capture_state(attn_metadata, saved_kv_lens_cuda) @@ -765,7 +793,11 @@ def replay(self, key: KeyType, f"replay() got {seqlen} tokens for key {key}, but the graph " f"was captured for {expected_num_tokens} tokens. A shorter " "input_ids leaves the tail of the static input buffer stale.") - static_tensors["input_ids"][:seqlen].copy_(input_ids) + # With the engine's buffers as the static inputs (static_input_ids) + # the source is the static view itself, and there is nothing to copy. + static_input_ids = static_tensors["input_ids"][:seqlen] + if input_ids.data_ptr() != static_input_ids.data_ptr(): + static_input_ids.copy_(input_ids) position_ids = current_inputs["position_ids"] if self.config.use_mrope: @@ -812,7 +844,9 @@ def replay(self, key: KeyType, f"for key {key}, but expected {expected_position_ids_shape}. " "torch.Tensor.copy_() silently broadcasts mismatched shapes, " "which would corrupt the static input buffer.") - static_tensors["position_ids"][:, :seqlen].copy_(position_ids) + static_position_ids = static_tensors["position_ids"][:, :seqlen] + if position_ids.data_ptr() != static_position_ids.data_ptr(): + static_position_ids.copy_(position_ids) num_encoder_tokens = key.num_encoder_tokens if num_encoder_tokens: diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index ed0cc2a89afe..cb880fe0f69b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -5762,6 +5762,38 @@ def update_resources( ) self._allocated_draft_lens.pop(req.py_request_id, None) + def stage_batch_block_offsets( + self, + stage, + dst_tensor: torch.Tensor, + request_ids: List[int], + beam_width: int, + num_contexts: int, + num_seqs: int, + ) -> bool: + """Stage ``copy_batch_block_offsets`` on a StepInputStage instead of launching its kernel. + + Returns False, staging nothing, where that copy is not this class's single-table copy on the current + stream: the per-layer page-table layout, a subclass that overrides ``copy_batch_block_offsets``, a manager + whose stream is not the current one, or a copy the stage does not cover. The caller then copies as usual. + """ + if ( + self._use_per_layer_page_tables + or type(self).copy_batch_block_offsets is not KVCacheManagerV2.copy_batch_block_offsets + or self._stream != torch.cuda.current_stream() + ): + return False + assert beam_width == 1, "beam_width must be 1 for KVCacheManagerV2" + copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width) + assert copy_idx.shape[0] == num_seqs + return stage.block_copy( + dst_tensor, + self.host_kv_cache_block_offsets, + copy_idx, + self.index_scales, + self.kv_offset, + ) + def copy_batch_block_offsets( self, dst_tensor: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2e93692a3947..b67f11f8789a 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -19,8 +19,9 @@ from tensorrt_llm._torch.peft.lora.config import LoraConfig from tensorrt_llm._torch.peft.lora.manager import LoraModelConfig from tensorrt_llm._torch.pyexecutor.warmup_timer import _WarmupTimer -from tensorrt_llm._utils import (global_mpi_rank, maybe_pin_memory, nvtx_range, - prefer_pinned, release_gc) +from tensorrt_llm._utils import (copy_to_device_if_changed, global_mpi_rank, + maybe_pin_memory, nvtx_range, prefer_pinned, + release_gc) from tensorrt_llm.bindings.internal import \ batch_manager as batch_manager_bindings from tensorrt_llm.inputs.multimodal import (MultimodalParams, @@ -48,6 +49,7 @@ from ..autotuner import AutoTuner, autotune from ..compilation.backend import Backend from ..compilation.utils import capture_piecewise_cuda_graph +from ..cute_dsl_kernels.spec_step_copies import op as spec_step_copies from ..distributed import Distributed from ..distributed.communicator import init_pp_comm from ..memory_buffer_utils import clear_memory_buffers, with_shared_pool @@ -756,8 +758,6 @@ def __init__( self.gather_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') - self.num_accepted_draft_tokens_cuda = torch.empty( - (self.batch_size, ), dtype=torch.int, device='cuda') self.previous_pos_indices_cuda = torch.empty( (self.max_num_tokens, ), dtype=torch.int, device='cuda') self.previous_pos_id_offsets_cuda = torch.zeros( @@ -811,6 +811,11 @@ def __init__( dtype=torch.int, device='cuda') self._encoder_decoder_staged_request_ids: Optional[List[int]] = None + # Host copies of what previous_batch_indices_cuda and + # previous_pos_indices_cuda last received: a decode step whose batch is + # unchanged skips both copies. + self._staged_previous_batch_indices: Optional[List[int]] = None + self._staged_previous_pos_indices: Optional[List[int]] = None if self._fallback_to_engine: self.input_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, @@ -876,6 +881,12 @@ def __init__( self.kv_cache_dtype_byte_size = self.get_kv_cache_dtype_byte_size() self._prepare_inputs_event: Optional[torch.cuda.Event] = None + # Where the step-copy kernels run: writes a speculative decode step's + # per-step inputs with one launch. Positions (up to max_num_tokens) and + # two per-sequence lengths fit its record. + self._step_input_stage: Optional[spec_step_copies.StepInputStage] = ( + spec_step_copies.StepInputStage(3 * self.max_num_tokens) if + self.is_spec_decode and spec_step_copies.is_supported() else None) # Cache for enc-dec cross-attention stable generation steps. # Populated on the first CUDA-graph generation step; cleared whenever @@ -918,6 +929,13 @@ def _initialize_cuda_graph_runner(self) -> Optional[CUDAGraphRunner]: enable_encoder_decoder_mixed_cuda_graph=( enable_encoder_decoder_mixed_cuda_graph), enable_in_graph_sampling=self.enable_in_graph_sampling, + # A speculative engine rewrites every token's input id and position + # each step (its prepare never takes the steady fast path), so the + # graphs read those straight from the engine's buffers. + static_input_ids=(self.input_ids_cuda if self.is_spec_decode + and not self.use_mrope else None), + static_position_ids=(self.position_ids_cuda if self.is_spec_decode + and not self.use_mrope else None), ) return CUDAGraphRunner(config) @@ -4364,6 +4382,7 @@ def _prepare_encoder_decoder_inputs_fast( [:num_previous_batch_requests], non_blocking=True) self._encoder_decoder_staged_request_ids = staged_request_ids + self._staged_previous_batch_indices = None generation_begin = num_context_tokens generation_end = generation_begin + num_previous_batch_requests torch.index_select( @@ -4736,7 +4755,6 @@ def _prepare_tp_inputs( # permanently reads back a zero delta. mrope_dummy_seq_slot = get_mrope_dummy_seq_slot(self.max_num_tokens, self.mapping.pp_size) - num_accepted_draft_tokens = [] # per request is_enc_dec = self._is_encoder_decoder_model() cross_encoder_hidden_states: List[torch.Tensor] = [] cross_encoder_seq_lens: List[int] = [ @@ -4806,7 +4824,6 @@ def append_cross_attention_state(request: LlmRequest, gather_ids.append(len(input_ids) - 1) sequence_lengths.append(len(prompt_tokens)) - num_accepted_draft_tokens.append(len(prompt_tokens) - 1) prompt_lengths.append(len(prompt_tokens)) past_seen_token_num = begin_compute num_cached_tokens_per_seq.append(past_seen_token_num - @@ -5064,7 +5081,6 @@ def _helix_pack_extend(request, group: int) -> int: prompt_lengths.append(request.py_prompt_len) sequence_lengths.append(1 + num_draft_tokens) - num_accepted_draft_tokens.append(num_draft_tokens) gather_ids.extend( list( range(len(position_ids), @@ -5099,8 +5115,6 @@ def _helix_pack_extend(request, group: int) -> int: request.py_batch_idx = request.py_seq_slot sequence_lengths.append(runtime_tokens_per_gen_step) - num_accepted_draft_tokens.append( - request.py_num_accepted_draft_tokens) past_seen_token_num = request.max_beam_num_tokens - 1 draft_lens.append(runtime_draft_token_buffer_width) @@ -5160,8 +5174,6 @@ def _helix_pack_extend(request, group: int) -> int: gather_ids.append( len(input_ids) - 1 - (self.original_max_draft_len - request.py_num_accepted_draft_tokens)) - num_accepted_draft_tokens.append( - request.py_num_accepted_draft_tokens) sequence_lengths.append(1 + self.original_max_draft_len) prompt_lengths.append(request.py_prompt_len) @@ -5231,7 +5243,6 @@ def _helix_pack_extend(request, group: int) -> int: # overhead (saves ~3 append calls per request). draft_lens.extend([0] * (_n_gen * beam_width)) sequence_lengths.extend([1] * (_n_gen * beam_width)) - num_accepted_draft_tokens.extend([0] * (_n_gen * beam_width)) for request in generation_requests: request_ids.append(request.py_request_id) @@ -5403,18 +5414,36 @@ def _helix_pack_extend(request, group: int) -> int: previous_batch_len = len(previous_batch_indices) def previous_seq_slots_device(): - previous_batch_indices_host = torch.tensor( - previous_batch_indices, - dtype=torch.int, - pin_memory=prefer_pinned()) previous_slots = self.previous_batch_indices_cuda[: previous_batch_len] - previous_slots.copy_(previous_batch_indices_host, non_blocking=True) + if previous_batch_indices != self._staged_previous_batch_indices: + previous_batch_indices_host = torch.tensor( + previous_batch_indices, + dtype=torch.int, + pin_memory=prefer_pinned()) + previous_slots.copy_(previous_batch_indices_host, + non_blocking=True) + self._staged_previous_batch_indices = list( + previous_batch_indices) return previous_slots num_tokens = len(input_ids) num_draft_tokens = len(draft_tokens) total_num_tokens = len(position_ids) + # Where the step-copy kernels run, the overlap gathers below are one + # kernel. On a CUDA graph step whose attention metadata is a plain + # TrtllmAttentionMetadata (its prepare() reads none of these inputs on + # the device), the gathers, the positions and the metadata's prompt + # lengths, KV lengths and block offsets are staged instead and written + # by one launch right after attn_metadata.prepare(). + step_copies = (self._step_input_stage + if self.enable_spec_decode else None) + if step_copies is not None: + step_copies.begin() + stage = (step_copies if step_copies is not None + and attn_metadata.is_cuda_graph and not self.use_mrope + and type(attn_metadata) is TrtllmAttentionMetadata + and attn_metadata.fp4_mla_state is None else None) assert total_num_tokens <= self.max_num_tokens, ( f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self.max_num_tokens})" ) @@ -5431,53 +5460,38 @@ def previous_seq_slots_device(): pin_memory=prefer_pinned()) self.draft_tokens_cuda[:len(draft_tokens)].copy_(draft_tokens, non_blocking=True) - if self.is_spec_decode and len(num_accepted_draft_tokens) > 0: - num_accepted_draft_tokens = torch.tensor(num_accepted_draft_tokens, - dtype=torch.int, - pin_memory=prefer_pinned()) - self.num_accepted_draft_tokens_cuda[:len( - num_accepted_draft_tokens)].copy_(num_accepted_draft_tokens, - non_blocking=True) if next_draft_tokens_device is not None: - # Initialize these two values to zeros - self.previous_pos_id_offsets_cuda *= 0 - self.previous_kv_lens_offsets_cuda *= 0 runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step( self.runtime_draft_len) runtime_draft_token_buffer_width = runtime_tokens_per_gen_step - 1 + num_extend_reqeust_wo_dummy = len(extend_requests) - len( + extend_dummy_requests) + # The offsets of requests without a previous batch and of dummy + # requests must read 0. On a CUDA graph step (no token padding) + # where every row is a request with a previous batch, the gathers + # below overwrite all rows _preprocess_inputs reads, so the buffers + # need no zeroing first. + if not (attn_metadata.is_cuda_graph and previous_batch_len + == num_extend_reqeust_wo_dummy and not extend_dummy_requests + and not generation_requests and not first_draft_requests + and scheduled_requests.num_context_requests == 0): + self.previous_pos_id_offsets_cuda.zero_() + self.previous_kv_lens_offsets_cuda.zero_() if previous_batch_len > 0: previous_slots = previous_seq_slots_device() - # previous input ids previous_batch_tokens = (previous_batch_len * runtime_tokens_per_gen_step) - new_tokens = new_tokens_device.transpose( - 0, - 1)[previous_slots, :runtime_tokens_per_gen_step].flatten() - self.input_ids_cuda[num_tokens:num_tokens + - previous_batch_tokens].copy_( - new_tokens, non_blocking=True) - - # previous draft tokens - previous_batch_draft_tokens = (previous_batch_len * - runtime_draft_token_buffer_width) - if runtime_draft_token_buffer_width > 0: - self.draft_tokens_cuda[ - num_draft_tokens:num_draft_tokens + - previous_batch_draft_tokens].copy_( - next_draft_tokens_device[ - previous_slots, : - runtime_draft_token_buffer_width].flatten(), - non_blocking=True) - # prepare data for the preprocess inputs - kv_len_offsets_device = (new_tokens_lens_device - - runtime_tokens_per_gen_step) - previous_pos_indices_host = torch.tensor( - previous_pos_indices, - dtype=torch.int, - pin_memory=prefer_pinned()) - self.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( - previous_pos_indices_host, non_blocking=True) + if previous_pos_indices != self._staged_previous_pos_indices: + previous_pos_indices_host = torch.tensor( + previous_pos_indices, + dtype=torch.int, + pin_memory=prefer_pinned()) + self.previous_pos_indices_cuda[ + 0:previous_batch_tokens].copy_( + previous_pos_indices_host, non_blocking=True) + self._staged_previous_pos_indices = list( + previous_pos_indices) # The order of requests in a batch: [context requests, generation requests] # generation requests: ['requests that do not have previous batch', 'requests that already have previous batch', 'dummy requests'] @@ -5487,27 +5501,60 @@ def previous_seq_slots_device(): # Therefore, both of self.previous_pos_id_offsets_cuda and self.previous_kv_lens_offsets_cuda are also 3 segments. # For 1) 'requests that do not have previous batch': disable overlap scheduler or the first step in the generation server of disaggregated serving. # Set these requests' previous_pos_id_offsets and previous_kv_lens_offsets to '0' to skip the value changes in _preprocess_inputs. - # Already set to '0' during initialization. + # Zeroed above. # For 2) 'requests that already have previous batch': enable overlap scheduler. - # Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device and kv_len_offsets_device. + # Set their previous_pos_id_offsets and previous_kv_lens_offsets according to new_tokens_lens_device. # For 3) 'dummy requests': pad dummy requests for CUDA graph or attention dp. - # Already set to '0' during initialization. - - num_extend_reqeust_wo_dummy = len(extend_requests) - len( - extend_dummy_requests) - self.previous_pos_id_offsets_cuda[ - (num_extend_reqeust_wo_dummy - previous_batch_len) * - runtime_tokens_per_gen_step:num_extend_reqeust_wo_dummy * - runtime_tokens_per_gen_step].copy_( - new_tokens_lens_device[self.previous_pos_indices_cuda[ - 0:previous_batch_tokens]], - non_blocking=True) - - self.previous_kv_lens_offsets_cuda[ - num_extend_reqeust_wo_dummy - - previous_batch_len:num_extend_reqeust_wo_dummy].copy_( - kv_len_offsets_device[previous_slots], - non_blocking=True) + # Zeroed above. + previous_begin = (num_extend_reqeust_wo_dummy - + previous_batch_len) + if step_copies is not None and step_copies.gather( + new_tokens_device, next_draft_tokens_device, + new_tokens_lens_device, + self.previous_batch_indices_cuda, + self.previous_pos_indices_cuda, previous_batch_len, + runtime_tokens_per_gen_step, + runtime_draft_token_buffer_width, self.input_ids_cuda, + num_tokens, self.draft_tokens_cuda, num_draft_tokens, + self.previous_pos_id_offsets_cuda, + previous_begin * runtime_tokens_per_gen_step, + self.previous_kv_lens_offsets_cuda, previous_begin): + if stage is None: + step_copies.commit() + else: + # previous input ids + new_tokens = new_tokens_device.transpose(0, 1)[ + previous_slots, :runtime_tokens_per_gen_step].flatten() + self.input_ids_cuda[num_tokens:num_tokens + + previous_batch_tokens].copy_( + new_tokens, non_blocking=True) + + # previous draft tokens + previous_batch_draft_tokens = ( + previous_batch_len * runtime_draft_token_buffer_width) + if runtime_draft_token_buffer_width > 0: + self.draft_tokens_cuda[ + num_draft_tokens:num_draft_tokens + + previous_batch_draft_tokens].copy_( + next_draft_tokens_device[ + previous_slots, : + runtime_draft_token_buffer_width].flatten(), + non_blocking=True) + # prepare data for the preprocess inputs + kv_len_offsets_device = (new_tokens_lens_device - + runtime_tokens_per_gen_step) + self.previous_pos_id_offsets_cuda[ + previous_begin * runtime_tokens_per_gen_step: + num_extend_reqeust_wo_dummy * + runtime_tokens_per_gen_step].copy_( + new_tokens_lens_device[ + self.previous_pos_indices_cuda[ + 0:previous_batch_tokens]], + non_blocking=True) + self.previous_kv_lens_offsets_cuda[ + previous_begin:num_extend_reqeust_wo_dummy].copy_( + kv_len_offsets_device[previous_slots], + non_blocking=True) elif new_tokens_device is not None: seq_slots_device = previous_seq_slots_device() @@ -5572,16 +5619,18 @@ def previous_seq_slots_device(): final_position_ids = self.mrope_position_ids_cuda[:, :, : total_num_tokens] else: - self.position_ids_cuda[:total_num_tokens].copy_(host_position_ids, - non_blocking=True) + if stage is None or not stage.copy( + self.position_ids_cuda[:total_num_tokens], + host_position_ids): + self.position_ids_cuda[:total_num_tokens].copy_( + host_position_ids, non_blocking=True) final_position_ids = self.position_ids_cuda[: total_num_tokens].unsqueeze( 0) if self.enable_spec_decode: - self.gather_ids_cuda[:len(gather_ids)].copy_(torch.tensor( - gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), - non_blocking=True) + copy_to_device_if_changed(self.gather_ids_cuda, + torch.tensor(gather_ids, dtype=torch.int)) if self.mapping.has_cp_helix(): # A non-None owned-count list is what arms @@ -5657,7 +5706,15 @@ def previous_seq_slots_device(): # pre-prepare counts so the steady-gen recording below stores values # that the per-step prepare() can re-clamp from scratch. num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) - attn_metadata.prepare() + if stage is None: + attn_metadata.prepare() + else: + attn_metadata.h2d_stage = stage + try: + attn_metadata.prepare() + finally: + attn_metadata.h2d_stage = None + stage.commit() cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( cross_encoder_hidden_states, cross_encoder_seq_lens, @@ -5775,8 +5832,6 @@ def previous_seq_slots_device(): # num_generations / num_tokens / seq_lens are set above, before the # attention-DP allgather that must agree with prepare(). spec_metadata.host_position_ids = host_position_ids - spec_metadata.num_accepted_draft_tokens = self.num_accepted_draft_tokens_cuda[:len( - num_accepted_draft_tokens)] if context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead( context_prompt_lookahead) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index acd9da97a83d..5e4770d63b8b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -72,7 +72,7 @@ from ..modules.decoder_layer import DecoderLayer from ..moe.expert_statistic import ExpertStatistic from ..speculative.drafter import Drafter -from ..speculative.spec_sampler_base import SampleStateTensorsSpec +from ..speculative.spec_sampler_base import SampleStateTensorsSpec, SpecSampler from ..speculative.speculation_gate import SpeculationGate from ..speculative.utils import update_draft_len from .adp_iter_stats import ADPIterStatsBuffer @@ -5420,7 +5420,16 @@ def _executor_loop_overlap(self): if can_queue: guided_decoder_failed_requests = None - with self.perf_manager.record_perf_events( + # The one-model speculative sampler only moves the forward's + # outputs into its stores, and the next forward reads those + # on the execution stream. Sampling on that stream too keeps + # the chain between two forwards on one stream (None: the + # current stream). + sample_stream = (self.execution_stream if isinstance( + self.sampler, SpecSampler) else None) + with torch.cuda.stream( + sample_stream + ), self.perf_manager.record_perf_events( None, gpu_sample_end) as sample_timing: with self._step_scope(scheduled_batch, phase="sample"): if self.guided_decoder is not None: diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 37992e2f8835..d6d6864f48b4 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -509,8 +509,6 @@ class SpecMetadata: host_position_ids: Optional[torch.Tensor] = None # The gather ids for logits. gather_ids: Optional[torch.Tensor] = None - # The number of accepted draft tokens for each request. - num_accepted_draft_tokens: Optional[torch.Tensor] = None # The number of tokens for speculative model/layer num_tokens: int = 0 # The number of tokens for speculative model/layer of different rank @@ -863,6 +861,11 @@ def _upload(dst: torch.Tensor, values: list[int]) -> None: pin_memory=prefer_pinned()), non_blocking=True) + # An all-greedy batch runs the argmax graph, which reads none of these + # buffers. The offset windows above still advance, so skipping the + # copies changes no later batch's streams. + if self.is_all_greedy_sample: + return _upload(self.request_seeds, request_seeds) _upload(self.request_offsets, request_offsets) _upload(self.seeds, flat_seeds) diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 030fb18f4fd9..8cd5ad530b76 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -27,7 +27,9 @@ import torch +from ..._utils import copy_to_device_if_changed from ...sampling_params import SamplingParams +from ..cute_dsl_kernels.spec_step_copies import op as spec_step_copies from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..pyexecutor.resource_manager import BaseResourceManager from ..pyexecutor.sampler import ( @@ -301,6 +303,16 @@ def __init__( # tree (K=6, T=60), MTP dynamic tree, PARD (T=2K-1) and the linear # modes -- none exceed it. self.max_accepted_path_len = args.max_draft_len + 1 + # Where the step-copy kernels run, one kernel moves a step's outputs + # into the stores below, reading each row's slot from _slot_table. + self._store_scatter: Optional[spec_step_copies.SlotScatter] = None + self._slot_table: Optional[torch.Tensor] = None + if spec_step_copies.is_supported(): + self._store_scatter = spec_step_copies.SlotScatter() + self._slot_table = torch.zeros((seq_slots,), dtype=torch.int32, device="cuda") + # Recorded after the last step's host copies of the stores, which read + # them on the D2H side stream. + self._store_copies_done: Optional[torch.cuda.Event] = None self.store = self.Store( new_tokens=int_tensor((self.max_accepted_path_len, seq_slots, self.max_beam_width)), next_new_tokens=int_tensor( @@ -371,6 +383,52 @@ def update_requests( req.py_rewind_len = runtime_draft_len - req.py_num_accepted_draft_tokens self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len) + def _scatter_to_stores( + self, outputs: dict[str, torch.Tensor], num_skip: int, slots: list[int] + ) -> None: + """Move rows ``num_skip`` onward of the forward's outputs into the stores at ``slots``.""" + num_sampling_requests = len(slots) + slots_device = torch.as_tensor(slots, dtype=torch.long) + slots_device = slots_device.to(device="cuda", non_blocking=True) + + o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests] + o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests] + o_next_draft_tokens = outputs["next_draft_tokens"][ + num_skip : num_skip + num_sampling_requests + ] + o_next_new_tokens = outputs["next_new_tokens"][num_skip : num_skip + num_sampling_requests] + + # Pad or truncate to match fixed-size store buffers for index_copy_. + # The worker output width tracks runtime_draft_len, which dynamic draft + # length shrinks below the statically allocated store width. + new_tokens_width = self.store.new_tokens.shape[0] + next_new_tokens_width = self.store.next_new_tokens.shape[0] + draft_tokens_width = self.store.next_draft_tokens.shape[1] + if o_new_tokens.shape[1] < new_tokens_width: + o_new_tokens = torch.nn.functional.pad( + o_new_tokens, (0, new_tokens_width - o_new_tokens.shape[1]) + ) + elif o_new_tokens.shape[1] > new_tokens_width: + o_new_tokens = o_new_tokens[:, :new_tokens_width] + if o_next_draft_tokens.shape[1] < draft_tokens_width: + o_next_draft_tokens = torch.nn.functional.pad( + o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]) + ) + elif o_next_draft_tokens.shape[1] > draft_tokens_width: + o_next_draft_tokens = o_next_draft_tokens[:, :draft_tokens_width] + if o_next_new_tokens.shape[1] < next_new_tokens_width: + o_next_new_tokens = torch.nn.functional.pad( + o_next_new_tokens, (0, next_new_tokens_width - o_next_new_tokens.shape[1]) + ) + elif o_next_new_tokens.shape[1] > next_new_tokens_width: + o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width] + + # Use index_copy_ for efficient copying (slots are unique) + self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots_device, o_new_tokens) + self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots_device, o_next_new_tokens) + self.store.new_tokens_lens.index_copy_(0, slots_device, o_new_tokens_lens) + self.store.next_draft_tokens.index_copy_(0, slots_device, o_next_draft_tokens) + def sample_async( self, scheduled_requests: ScheduledRequests, @@ -396,7 +454,6 @@ def sample_async( num_skip = len(scheduled_requests.context_requests_chunking) finished_context_requests = scheduled_requests.context_requests_last_chunk sampling_requests = finished_context_requests + scheduled_requests.generation_requests - num_sampling_requests = len(sampling_requests) # Snapshot each request's draft count for THIS step before # _add_dummy_draft_tokens below installs placeholder drafts on @@ -413,47 +470,27 @@ def sample_async( for r in sampling_requests ] - slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long) - slots = slots.to(device="cuda", non_blocking=True) - - o_new_tokens = outputs["new_tokens"][num_skip : num_skip + num_sampling_requests] - o_new_tokens_lens = outputs["new_tokens_lens"][num_skip : num_skip + num_sampling_requests] - o_next_draft_tokens = outputs["next_draft_tokens"][ - num_skip : num_skip + num_sampling_requests - ] - o_next_new_tokens = outputs["next_new_tokens"][num_skip : num_skip + num_sampling_requests] - runtime_draft_len = o_next_draft_tokens.shape[1] - - # Pad or truncate to match fixed-size store buffers for index_copy_. - # The worker output width tracks runtime_draft_len, which dynamic draft - # length shrinks below the statically allocated store width. - new_tokens_width = self.store.new_tokens.shape[0] - next_new_tokens_width = self.store.next_new_tokens.shape[0] - draft_tokens_width = self.store.next_draft_tokens.shape[1] - if o_new_tokens.shape[1] < new_tokens_width: - o_new_tokens = torch.nn.functional.pad( - o_new_tokens, (0, new_tokens_width - o_new_tokens.shape[1]) + runtime_draft_len = outputs["next_draft_tokens"].shape[1] + if self._store_copies_done is not None: + # The previous step's host copies may still be reading the stores. + torch.cuda.current_stream().wait_event(self._store_copies_done) + slots = [r.py_seq_slot for r in sampling_requests] + scattered = False + if self._store_scatter is not None: + # A step whose batch is unchanged leaves the slot table as it is. + copy_to_device_if_changed(self._slot_table, torch.tensor(slots, dtype=torch.int32)) + scattered = self._store_scatter.scatter( + outputs, + num_skip, + len(slots), + self._slot_table, + self.store.new_tokens, + self.store.next_new_tokens, + self.store.new_tokens_lens, + self.store.next_draft_tokens, ) - elif o_new_tokens.shape[1] > new_tokens_width: - o_new_tokens = o_new_tokens[:, :new_tokens_width] - if o_next_draft_tokens.shape[1] < draft_tokens_width: - o_next_draft_tokens = torch.nn.functional.pad( - o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]) - ) - elif o_next_draft_tokens.shape[1] > draft_tokens_width: - o_next_draft_tokens = o_next_draft_tokens[:, :draft_tokens_width] - if o_next_new_tokens.shape[1] < next_new_tokens_width: - o_next_new_tokens = torch.nn.functional.pad( - o_next_new_tokens, (0, next_new_tokens_width - o_next_new_tokens.shape[1]) - ) - elif o_next_new_tokens.shape[1] > next_new_tokens_width: - o_next_new_tokens = o_next_new_tokens[:, :next_new_tokens_width] - - # Use index_copy_ for efficient copying (slots are unique) - self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens) - self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens) - self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens) - self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens) + if not scattered: + self._scatter_to_stores(outputs, num_skip, slots) # Create sample state with async D2H copy device_tensors = SampleStateTensorsSpec( @@ -462,12 +499,26 @@ def sample_async( next_draft_tokens=self.store.next_draft_tokens, ) - host_tensors = SampleStateTensorsSpec( - new_tokens=self._copy_to_host(self.store.new_tokens), - new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), - next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), - ) - sampler_event = self._record_sampler_event() + if self._async_worker_active(): + host_tensors = SampleStateTensorsSpec( + new_tokens=self._copy_to_host(self.store.new_tokens), + new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), + next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), + ) + sampler_event = self._record_sampler_event() + else: + # The next step's input preparation waits on this stream and only + # needs the stores, so the host copies run on the side stream. + # update_requests syncs their event before reading them, and the + # next step's store update waits for it on this stream. + with self._make_side_stream_copier() as copier: + host_tensors = SampleStateTensorsSpec( + new_tokens=copier.stage_copy_to_host(self.store.new_tokens), + new_tokens_lens=copier.stage_copy_to_host(self.store.new_tokens_lens), + next_draft_tokens=copier.stage_copy_to_host(self.store.next_draft_tokens), + ) + self._store_copies_done = copier.event + sampler_event = self._record_sampler_event(side_stream_event=copier.event) # Add dummy draft tokens to context requests for KV cache preparation for request in finished_context_requests: diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index c875715c4501..b45bb486b740 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -1336,6 +1336,34 @@ def maybe_pin_memory(tensor: torch.Tensor) -> torch.Tensor: return tensor +def copy_to_device_if_changed(dst: torch.Tensor, host: torch.Tensor) -> None: + """Copy ``host`` into the leading elements of the device buffer ``dst`` + unless they already hold it. + + ``dst`` is the whole buffer (not a slice), and the values it last received + are kept on it, so a reallocated buffer starts over. Only for per-step + metadata buffers that nothing else writes: a decode step whose batch is + unchanged then enqueues no copy. "Nothing else" includes other tensor + objects over the same memory, such as the views a shared buffer pool hands + to each metadata object. The copy itself goes through fresh pinned staging, + so ``host`` may be reused right away. + """ + flat = host.reshape(-1) + n = flat.numel() + last = getattr(dst, "_last_host_values", None) + if last is not None and last.numel() >= n and torch.equal(last[:n], flat): + return + staging = torch.empty_like(flat, device="cpu", pin_memory=prefer_pinned()) + staging.copy_(flat) + dst.view(-1)[:n].copy_(staging, non_blocking=True) + if last is None or last.numel() <= n: + dst._last_host_values = flat.clone() + else: + last = last.clone() + last[:n] = flat + dst._last_host_values = last + + def async_tensor_h2d(data, dtype: torch.dtype, device: Union[str, torch.device]) -> torch.Tensor: """Build a CPU tensor from `data` and ship it to `device` with a diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 2ca039299f9c..3a7ca13438e7 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -238,6 +238,10 @@ l0_b200: # accuracy deviation found so far is Blackwell-only), so it is not # hw-agnostic and needs its own Blackwell run. - unittest/_torch/speculative/test_fused_sampling_op.py + # One-model speculative decoding step-copy kernels and the engine's staged + # decode-step inputs (SM 100 only; they skip on every other list). + - unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py + - unittest/_torch/executor/test_pytorch_model_engine.py -k "test_staged_spec_decode_graph_step" - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/thop/serial diff --git a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py new file mode 100644 index 000000000000..10408b60106d --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py @@ -0,0 +1,822 @@ +# 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. +"""``spec_step_copies`` (the one-model speculative decoding step's copy kernels) against the torch ops each replaces, +bit for bit (int32 copies), at every split of R = 1 .. 8 requests x T = 1, 2, 4, 8 tokens, on buffers sized like an +engine's for max batch 8 and max draft length 7 (sampler stores [8, slots, 1], [8, slots, 1], [slots], [slots, 7] over +8 or 16 slots; draft-token buffer 56, KV-length offsets 8, per-token buffers max_num_tokens). + +* ``SlotScatter.scatter`` vs SpecSampler's torch store update (each output padded with zeros or cut to its store's + width, then four ``index_copy_`` by slot): outputs narrower than (dynamic draft length), as wide as and wider than the + stores, and mixed per field; row_begin 0 and > 0; shuffled distinct slot tables. +* ``StepInputStage.gather`` committed on its own vs PyTorchModelEngine._prepare_tp_inputs's torch overlap gathers (the + stores by slot into the input ids and draft tokens, the lengths by the per-token index list into the position + offsets, ``new_tokens_lens - T`` by slot into the KV-length offsets): the engine's offsets (the requests without a + previous batch first), zero and arbitrary offsets, a draft width below T - 1, the engine's and arbitrary index lists. +* ``StepInputStage`` vs those gathers plus what it stages instead of copying: ``dst.copy_(pinned, non_blocking=True)`` + per staged copy (positions, prompt and KV lengths, a fourth) and the KV cache manager's block-offset copy + (``copy_batch_block_offsets_to_device``, target and draft managers, bad pages included): the engine's step, all four + copy slots at offset destinations, the gathers alone, the copies alone. +* The calls the kernels do not cover launch nothing and return False, so the caller keeps its torch path. + +Every buffer starts random and is compared whole (what an op must not touch included), every case runs twice (reruns +bit-identical), and one CUDA graph per split family (R = 1 .. 8 at one T) captures each step's scatter, gather and +staged commit; it is replayed with every input (and the staged values in the stages' pinned records) rewritten in +place, bit-identical to the eager ops and to the torch ops on every replay. + +Table: ``python3 test_spec_step_copies.py report``. Timing: ``python3 test_spec_step_copies.py time`` (CUDA graphs of +back-to-back calls, median us per call over 15 replays, each kernel vs the torch ops it replaces, at every split). +""" + +import statistics +import sys + +import pytest +import torch + +MAX_BATCH = 8 # max_batch_size +MAX_DRAFT = 7 # max_draft_len (linear speculation: max_total_draft_tokens == max_draft_len) +STORE_WIDTH = MAX_DRAFT + 1 # the sampler's new_tokens / next_new_tokens stores +MAX_NUM_TOKENS = 8192 # the engine's per-token buffers (input ids, positions, previous_*) +DRAFT_BUFFER = MAX_DRAFT * MAX_BATCH # draft_tokens_cuda: max_draft_loop_tokens * batch_size +STAGE_CAPACITY = 3 * MAX_NUM_TOKENS # PyTorchModelEngine._get_step_input_stage +BLOCKS = 40 # KV blocks per sequence (copy_batch_block_offsets_to_device needs a multiple of 4) +TABLE_SEQS = 24 # rows of a KV cache manager's pinned host block-offset table +POOLS = (2, 1) # pools of the target and the draft KV cache managers +COPY_NAMES = ("position_ids", "prompt_lens", "kv_lens", "extra") # staged copies, in order + +ROWS = list(range(1, MAX_BATCH + 1)) +TOKENS = [1, 2, 4, 8] +SPLITS = [(r, t) for t in TOKENS for r in ROWS] + + +def _op(): + from tensorrt_llm._torch.cute_dsl_kernels.spec_step_copies import op + + return op + + +def _supported() -> bool: + return torch.cuda.is_available() and _op().is_supported() + + +pytestmark = pytest.mark.skipif( + not _supported(), reason="the step-copy kernels run on SM 100 with the CuTe DSL" +) + + +def _pinned_ok() -> bool: + """StepInputStage reads a pinned host record in place; it refuses to run where pinned memory is not preferred.""" + from tensorrt_llm._utils import prefer_pinned + + return prefer_pinned() + + +_eager = {} + + +def eager_stage(): + """The StepInputStage of the eager runs (never captured: a captured commit's stage cannot stage again).""" + if "stage" not in _eager: + _eager["stage"] = _op().StepInputStage(STAGE_CAPACITY) + return _eager["stage"] + + +def kernel_scatter(args) -> None: + """``SlotScatter.scatter`` on ``args``, which it must cover.""" + assert _op().SlotScatter().scatter(*args), "scatter declined a covered call" + + +def kernel_gather(stage, args) -> None: + """The overlap gathers on ``args`` (which they must cover) committed on their own through ``stage``.""" + stage.begin() + assert stage.gather(*args), "gather declined a covered call" + stage.commit() + + +# ---------------------------------------------------------------------------------------------------------------- +# The torch ops each kernel replaces +# ---------------------------------------------------------------------------------------------------------------- + + +def torch_scatter( + outputs, + num_skip, + num_sampling_requests, + slot_table, + store_new_tokens, + store_next_new_tokens, + store_new_tokens_lens, + store_next_draft_tokens, +): + """SpecSampler's torch store update (SlotScatter.scatter's arguments; the slots as the long tensor it builds from + the requests' seq slots).""" + slots = slot_table[:num_sampling_requests].long() + end = num_skip + num_sampling_requests + o_new_tokens = outputs["new_tokens"][num_skip:end] + o_new_tokens_lens = outputs["new_tokens_lens"][num_skip:end] + o_next_draft_tokens = outputs["next_draft_tokens"][num_skip:end] + o_next_new_tokens = outputs["next_new_tokens"][num_skip:end] + + def fit(t, width): # pad with zeros or truncate to the store width + if t.shape[1] < width: + return torch.nn.functional.pad(t, (0, width - t.shape[1])) + return t[:, :width] + + o_new_tokens = fit(o_new_tokens, store_new_tokens.shape[0]) + o_next_draft_tokens = fit(o_next_draft_tokens, store_next_draft_tokens.shape[1]) + o_next_new_tokens = fit(o_next_new_tokens, store_next_new_tokens.shape[0]) + store_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens) + store_next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens) + store_new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens) + store_next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens) + + +def torch_gather( + new_tokens_device, + next_draft_tokens_device, + new_tokens_lens_device, + previous_batch_indices_cuda, + previous_pos_indices_cuda, + previous_batch_len, + runtime_tokens_per_gen_step, + runtime_draft_token_buffer_width, + input_ids_cuda, + num_tokens, + draft_tokens_cuda, + num_draft_tokens, + previous_pos_id_offsets_cuda, + pos_begin, + previous_kv_lens_offsets_cuda, + kv_begin, +): + """PyTorchModelEngine._prepare_tp_inputs's torch overlap gathers (StepInputStage.gather's arguments).""" + tokens = runtime_tokens_per_gen_step + width = runtime_draft_token_buffer_width + previous_slots = previous_batch_indices_cuda[:previous_batch_len] + previous_batch_tokens = previous_batch_len * tokens + new_tokens = new_tokens_device.transpose(0, 1)[previous_slots, :tokens].flatten() + input_ids_cuda[num_tokens : num_tokens + previous_batch_tokens].copy_( + new_tokens, non_blocking=True + ) + previous_batch_draft_tokens = previous_batch_len * width + if width > 0: + draft_tokens_cuda[num_draft_tokens : num_draft_tokens + previous_batch_draft_tokens].copy_( + next_draft_tokens_device[previous_slots, :width].flatten(), non_blocking=True + ) + kv_len_offsets_device = new_tokens_lens_device - tokens + previous_pos_id_offsets_cuda[pos_begin : pos_begin + previous_batch_tokens].copy_( + new_tokens_lens_device[previous_pos_indices_cuda[0:previous_batch_tokens]], + non_blocking=True, + ) + previous_kv_lens_offsets_cuda[kv_begin : kv_begin + previous_batch_len].copy_( + kv_len_offsets_device[previous_slots], non_blocking=True + ) + + +def torch_block_copy(offsets, table, copy_index, index_scales, kv_offset): + """KVCacheManagerV2.copy_batch_block_offsets's copy (StepInputStage.block_copy's arguments), current stream.""" + from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( + copy_batch_block_offsets_to_device, + ) + + copy_batch_block_offsets_to_device( + table, offsets, copy_index, index_scales, kv_offset, torch.cuda.current_stream().cuda_stream + ) + + +# ---------------------------------------------------------------------------------------------------------------- +# One decode step's buffers +# ---------------------------------------------------------------------------------------------------------------- + + +def rand_i32(shape, g, low=-(1 << 30), high=1 << 30): + """Random host int32 (``g``: a CPU generator).""" + return torch.randint(low, high, tuple(shape), generator=g, dtype=torch.int32) + + +def shuffled_distinct(count, total, g): + """``count`` distinct indices below ``total`` (>= 2) in a shuffled order: never ascending and never index r at + position r, so a kernel that takes a row for its slot (or ignores the order) fails.""" + while True: + picked = torch.randperm(total, generator=g)[:count].tolist() + if all(p != r for r, p in enumerate(picked)) and (count < 2 or picked != sorted(picked)): + return picked + + +def engine_begins(rows, tokens): + """The engine's gather offsets (input ids, draft tokens, position offsets, KV-length offsets) when the batch's + other MAX_BATCH - rows generation requests (no previous batch; T tokens, T - 1 draft tokens each) come first.""" + first = MAX_BATCH - rows + return first * tokens, first * (tokens - 1), first * tokens, first + + +class Step: + """One decode step of ``rows`` requests x ``tokens`` tokens on engine-sized int32 buffers, all random so that a + stray or a missing write shows (index buffers hold valid indices past their used part): + + * the forward's outputs: new_tokens [N, a], next_new_tokens [N, b], new_tokens_lens [N], next_draft_tokens [N, c] + (default a = b = T, c = T - 1); + * the sampler's slot table [slots] and stores new_tokens / next_new_tokens [8, slots, 1], new_tokens_lens [slots], + next_draft_tokens [slots, 7]; + * the engine's previous_batch_indices / previous_pos_indices, and two output sets ("a": the gathers committed on + their own, "b": the gathers in a full staged step) of input_ids, draft_tokens [56], previous_pos_id_offsets and + previous_kv_lens_offsets [8] (grown only when an offset needs it); + * the staged copies' destinations (positions, prompt and KV lengths [slots], a fourth buffer) and host values; + * the target and draft KV cache managers' device block offsets [pools, slots, 2, 40] and pinned host tables + [pools, 24, 2, 40] (a fifth of the pages bad, -1), copy indices, index scales and KV offsets. + + ``chain``: the gather reads the scatter's slots in another order (the next step's inputs from this step's stores). + """ + + def __init__( + self, + rows, + tokens, + seed, + *, + num_slots=MAX_BATCH, + widths=None, + row_begin=0, + out_rows=None, + begins=(0, 0, 0, 0), + draft_width=None, + pos_list="engine", + chain=False, + copy_at=(0, 0, 0, 0), + ): + self.rows, self.tokens, self.num_slots = rows, tokens, num_slots + self.widths = widths or (tokens, tokens, tokens - 1) + self.row_begin, self.begins = row_begin, begins + self.draft_width = tokens - 1 if draft_width is None else draft_width + self.pos_list, self.chain, self.copy_at = pos_list, chain, copy_at + self.g = torch.Generator().manual_seed(seed) + n_out = max(MAX_BATCH, row_begin + rows) if out_rows is None else out_rows + assert n_out >= row_begin + rows + n_draft = max(DRAFT_BUFFER, begins[1] + rows * self.draft_width) + n_kv = max(MAX_BATCH, begins[3] + rows) + a, b, c = self.widths + device = { + "out_new": (n_out, a), + "out_next": (n_out, b), + "out_lens": (n_out,), + "out_draft": (n_out, c), + "slot_table": (num_slots,), + "st_new": (STORE_WIDTH, num_slots, 1), + "st_next": (STORE_WIDTH, num_slots, 1), + "st_lens": (num_slots,), + "st_draft": (num_slots, MAX_DRAFT), + "prev_slots": (MAX_NUM_TOKENS,), + "prev_pos": (MAX_NUM_TOKENS,), + "position_ids": (MAX_NUM_TOKENS,), + "prompt_lens": (num_slots,), + "kv_lens": (num_slots,), + "extra": (MAX_NUM_TOKENS,), + } + host = {} + for s in "ab": + device[f"input_ids_{s}"] = (MAX_NUM_TOKENS,) + device[f"draft_{s}"] = (n_draft,) + device[f"pos_off_{s}"] = (MAX_NUM_TOKENS,) + device[f"kv_off_{s}"] = (n_kv,) + for m, pools in zip("td", POOLS): + device[f"blk_{m}"] = (pools, num_slots, 2, BLOCKS) + host[f"tab_{m}"] = (pools, TABLE_SEQS, 2, BLOCKS) + host[f"cidx_{m}"] = (rows,) + host[f"scale_{m}"] = (pools,) + host[f"kvoff_{m}"] = (pools,) + self.bufs = {k: torch.empty(v, dtype=torch.int32, device="cuda") for k, v in device.items()} + for k, v in host.items(): + self.bufs[k] = torch.empty(v, dtype=torch.int32, pin_memory=True) + self.rewrite() + + def rewrite(self) -> None: + """New random contents for every buffer and host value, written in place (a captured graph keeps the + addresses).""" + g, rows, tokens, num_slots = self.g, self.rows, self.tokens, self.num_slots + b = self.bufs + for t in b.values(): + if t.is_cuda: + t.copy_(rand_i32(t.shape, g)) + slots = shuffled_distinct(rows, num_slots, g) + table = rand_i32((num_slots,), g, 0, num_slots) + table[:rows] = torch.tensor(slots, dtype=torch.int32) + b["slot_table"].copy_(table) + previous = slots[1:] + slots[:1] if self.chain else shuffled_distinct(rows, num_slots, g) + previous = torch.tensor(previous, dtype=torch.int32) + index = rand_i32((MAX_NUM_TOKENS,), g, 0, num_slots) + index[:rows] = previous + b["prev_slots"].copy_(index) + per_token = rand_i32((MAX_NUM_TOKENS,), g, 0, num_slots) + if self.pos_list == "engine": # each row's slot, once per token + per_token[: rows * tokens] = previous.repeat_interleave(tokens) + b["prev_pos"].copy_(per_token) + self.values = { + "position_ids": rand_i32((rows * tokens,), g, 0, 1 << 20), + "prompt_lens": rand_i32((rows,), g, 1, 1 << 20), + "kv_lens": rand_i32((rows,), g, 1, 1 << 20), + "extra": rand_i32((rows,), g), + } + self.pinned = {k: v.pin_memory() for k, v in self.values.items()} + for m in "td": + pages = rand_i32(b[f"tab_{m}"].shape, g, 0, 1 << 20) + pages[torch.rand(pages.shape, generator=g) < 0.2] = -1 + b[f"tab_{m}"].copy_(pages) + b[f"cidx_{m}"].copy_(torch.tensor(shuffled_distinct(rows, TABLE_SEQS, g))) + b[f"scale_{m}"].copy_(rand_i32(b[f"scale_{m}"].shape, g, 1, 8)) + b[f"kvoff_{m}"].copy_(rand_i32(b[f"kvoff_{m}"].shape, g, 0, 1 << 16)) + + def outputs(self, b): + return { + "new_tokens": b["out_new"], + "next_new_tokens": b["out_next"], + "new_tokens_lens": b["out_lens"], + "next_draft_tokens": b["out_draft"], + } + + def scatter_args(self, b): + """SlotScatter.scatter's arguments on the buffers ``b``.""" + return (self.outputs(b), self.row_begin, self.rows, b["slot_table"], b["st_new"], b["st_next"], + b["st_lens"], b["st_draft"]) # fmt: skip + + def gather_args(self, b, out): + """StepInputStage.gather's arguments on the buffers ``b``, into output set ``out``.""" + ib, db, pb, kb = self.begins + return (b["st_next"], b["st_draft"], b["st_lens"], b["prev_slots"], b["prev_pos"], self.rows, self.tokens, + self.draft_width, b[f"input_ids_{out}"], ib, b[f"draft_{out}"], db, b[f"pos_off_{out}"], pb, + b[f"kv_off_{out}"], kb) # fmt: skip + + def copies(self, b, count): + """The first ``count`` staged copies: (destination view, host value name).""" + return [ + (b[name][at : at + self.values[name].numel()], name) + for name, at in zip(COPY_NAMES[:count], self.copy_at) + ] + + def block_copies(self, b): + """StepInputStage.block_copy's arguments for the target and the draft KV cache managers.""" + return [ + (b[f"blk_{m}"], b[f"tab_{m}"], b[f"cidx_{m}"], b[f"scale_{m}"], b[f"kvoff_{m}"]) + for m in "td" + ] + + def run_stage(self, b, stage, gather=True, copies=3, blocks=True): + """The step's inputs through ``stage`` (a StepInputStage): one commit.""" + stage.begin() + if gather: + assert stage.gather(*self.gather_args(b, "b")), "gather declined" + for dst, name in self.copies(b, copies): + assert stage.copy(dst, self.values[name]), "copy declined" + if blocks: + for args in self.block_copies(b): + assert stage.block_copy(*args), "block copy declined" + stage.commit() + + def torch_stage(self, b, gather=True, copies=3, blocks=True): + """What the stage replaces: the gathers, a non-blocking copy from pinned memory per staged copy, and the KV + cache managers' block-offset copies.""" + if gather: + torch_gather(*self.gather_args(b, "b")) + for dst, name in self.copies(b, copies): + dst.copy_(self.pinned[name], non_blocking=True) + if blocks: + for args in self.block_copies(b): + torch_block_copy(*args) + + +def clone_bufs(bufs): + """A copy of every buffer (pinned host buffers stay pinned).""" + out = {} + for k, t in bufs.items(): + if t.is_cuda: + out[k] = t.clone() + else: + out[k] = torch.empty(t.shape, dtype=t.dtype, pin_memory=t.is_pinned()).copy_(t) + return out + + +def mismatches(x, y): + """Names of the buffers whose contents differ.""" + return [k for k in x if not torch.equal(x[k], y[k])] + + +def compare(bufs, run_torch, run_kernel): + """The buffers where the kernel differs from the torch ops and where a rerun differs from the kernel, each run on + its own copy of ``bufs``.""" + want, got, again = clone_bufs(bufs), clone_bufs(bufs), clone_bufs(bufs) + run_torch(want) + run_kernel(got) + run_kernel(again) + torch.cuda.synchronize() + return dict(torch=mismatches(want, got), rerun=mismatches(got, again)) + + +def result(case, **bad): + """A table row; ``bad`` maps each comparison (torch, eager, rerun) to the buffers where it differed.""" + failed = {k: v for k, v in bad.items() if v} + if not failed: + return dict(case=case, identical="yes (= " + ", = ".join(bad) + ")", ok=True) + detail = "; ".join(f"!= {k}: {', '.join(v)}" for k, v in failed.items()) + return dict(case=case, identical=f"no: {detail}", ok=False) + + +def _seed(op_index, rows, tokens, case): + return 100000 * op_index + 1000 * rows + 10 * tokens + case + + +# ---------------------------------------------------------------------------------------------------------------- +# The cases of one split +# ---------------------------------------------------------------------------------------------------------------- + + +def scatter_cases(rows, tokens): + """(label, Step arguments): the forward's outputs narrower than or as wide as the stores (dynamic draft length; + the engine's batch of 8 rows, the first 8 - R skipped), wider than every store, and mixed per field.""" + t = tokens + fit = "pad" if t < STORE_WIDTH else "exact" + return [ + (f"{fit}: widths ({t}, {t}, {t - 1}) -> (8, 8, 7), 8 slots, row_begin {MAX_BATCH - rows} of 8 rows", + dict(num_slots=MAX_BATCH, row_begin=MAX_BATCH - rows, out_rows=MAX_BATCH)), + (f"cut: widths ({t + 8}, {t + 8}, {t + 7}) -> (8, 8, 7), 16 slots, row_begin 0", + dict(num_slots=2 * MAX_BATCH, widths=(t + 8, t + 8, t + 7), out_rows=rows + 1)), + (f"mixed: widths ({t + 8}, {t}, {t + 7}) -> (8, 8, 7), 16 slots, row_begin 3", + dict(num_slots=2 * MAX_BATCH, widths=(t + 8, t, t + 7), row_begin=3, out_rows=rows + 4)), + ] # fmt: skip + + +def gather_cases(rows, tokens): + """(label, Step arguments): the engine's layout, zero offsets, and arbitrary offsets with a narrower draft width + and an arbitrary per-token index list.""" + begins = engine_begins(rows, tokens) + return [ + (f"engine: offsets {begins}, draft width {tokens - 1}, per-token list = slots x T, 8 slots", + dict(num_slots=MAX_BATCH, begins=begins)), + (f"offsets 0, draft width {tokens - 1}, per-token list = slots x T, 16 slots", + dict(num_slots=2 * MAX_BATCH)), + (f"offsets (37, 5, 11, 3), draft width {(tokens - 1) // 2}, arbitrary per-token list, 16 slots", + dict(num_slots=2 * MAX_BATCH, begins=(37, 5, 11, 3), draft_width=(tokens - 1) // 2, pos_list="any")), + ] # fmt: skip + + +def stage_cases(rows, tokens): + """(label, Step arguments, run_stage arguments).""" + begins = engine_begins(rows, tokens) + return [ + (f"engine step: gathers at {begins}, positions / prompt / KV lengths, target + draft block offsets", + dict(num_slots=MAX_BATCH, begins=begins), dict()), + ("4 copies to offset destinations, gathers at (37, 5, 11, 3), block offsets, 16 slots", + dict(num_slots=2 * MAX_BATCH, begins=(37, 5, 11, 3), pos_list="any", copy_at=(7, 1, 2, 3)), + dict(copies=4)), + ("gathers only, 16 slots", dict(num_slots=2 * MAX_BATCH), dict(copies=0, blocks=False)), + ("copies + block offsets only (no previous batch), 16 slots", + dict(num_slots=2 * MAX_BATCH, copy_at=(5, 3, 0, 0)), dict(gather=False)), + ] # fmt: skip + + +def measure_scatter(rows, tokens): + out = [] + for i, (case, kw) in enumerate(scatter_cases(rows, tokens)): + st = Step(rows, tokens, _seed(1, rows, tokens, i), **kw) + bad = compare( + st.bufs, + lambda b: torch_scatter(*st.scatter_args(b)), + lambda b: kernel_scatter(st.scatter_args(b)), + ) + out.append(result(case, **bad)) + return out + + +def measure_gather(rows, tokens): + stage = eager_stage() + out = [] + for i, (case, kw) in enumerate(gather_cases(rows, tokens)): + st = Step(rows, tokens, _seed(2, rows, tokens, i), **kw) + bad = compare( + st.bufs, + lambda b: torch_gather(*st.gather_args(b, "a")), + lambda b: kernel_gather(stage, st.gather_args(b, "a")), + ) + out.append(result(case, **bad)) + return out + + +def measure_stage(rows, tokens): + stage = eager_stage() + out = [] + for i, (case, kw, run_kw) in enumerate(stage_cases(rows, tokens)): + st = Step(rows, tokens, _seed(3, rows, tokens, i), **kw) + bad = compare( + st.bufs, + lambda b: st.torch_stage(b, **run_kw), + lambda b: st.run_stage(b, stage, **run_kw), + ) + out.append(result(case, **bad)) + return out + + +def measure_graph(tokens, replays=6): + """One CUDA graph for the split family R = 1 .. 8 at ``tokens``: each step's scatter, its gather (from the stores + the scatter wrote, as the next step reads them) and its staged commit, captured once and replayed with every + input rewritten in place; each replay against the eager ops and the torch ops on the same inputs, then a replay + of the last inputs against the first.""" + op = _op() + with_stage = _pinned_ok() + gather_stage = op.StepInputStage(STAGE_CAPACITY) # gathers only: its record is never read + steps = [ + Step(r, tokens, _seed(4, r, tokens, 0), num_slots=2 * MAX_BATCH, row_begin=MAX_BATCH - r, + out_rows=MAX_BATCH, begins=engine_begins(r, tokens), chain=True) + for r in ROWS + ] # fmt: skip + stage = eager_stage() if with_stage else None + # A captured commit reads its stage's record at every replay, so each step keeps its own stage. + captured = [op.StepInputStage(STAGE_CAPACITY) if with_stage else None for _ in steps] + + def run_kernels(st, b, step_stage): + kernel_scatter(st.scatter_args(b)) + kernel_gather(gather_stage, st.gather_args(b, "a")) + if step_stage is not None: + st.run_stage(b, step_stage) + + def run_torch(st, b): + torch_scatter(*st.scatter_args(b)) + torch_gather(*st.gather_args(b, "a")) + if with_stage: + st.torch_stage(b) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + for st in steps: # the kernels compile on their first call, outside capture + run_kernels(st, clone_bufs(st.bufs), stage) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for st, step_stage in zip(steps, captured): + run_kernels(st, st.bufs, step_stage) + torch.cuda.synchronize() + # The staged values sit in each captured stage's pinned record in staging order; a replay reads them in place. + spans = [] + for st, step_stage in zip(steps, captured): + at, step_spans = 0, [] + for _, name in st.copies(st.bufs, 3) if with_stage else []: + n = st.values[name].numel() + assert torch.equal(step_stage.record[at : at + n], st.values[name]), "record layout" + step_spans.append((at, n, name)) + at += n + spans.append(step_spans) + what = "scatter, gather" + (", staged commit" if with_stage else "") + out = [] + before = None + for rep in range(replays): + for st, step_stage, step_spans in zip(steps, captured, spans): + st.rewrite() + for at, n, name in step_spans: + step_stage.record[at : at + n].copy_(st.values[name]) + before = [clone_bufs(st.bufs) for st in steps] + graph.replay() + torch.cuda.synchronize() + bad = dict(torch=[], eager=[]) + for st, start in zip(steps, before): + want, got = clone_bufs(start), clone_bufs(start) + run_torch(st, want) + run_kernels(st, got, stage) + torch.cuda.synchronize() + bad["torch"] += [f"{st.rows}x{tokens} {k}" for k in mismatches(want, st.bufs)] + bad["eager"] += [f"{st.rows}x{tokens} {k}" for k in mismatches(got, st.bufs)] + out.append(result(f"replay {rep}: {what}, every input rewritten", **bad)) + after = [clone_bufs(st.bufs) for st in steps] + for st, start in zip(steps, before): + for k, t in st.bufs.items(): + t.copy_(start[k]) + graph.replay() + torch.cuda.synchronize() + rerun = [ + f"{st.rows}x{tokens} {k}" for st, a in zip(steps, after) for k in mismatches(a, st.bufs) + ] + out.append(result("the last inputs replayed again", rerun=rerun)) + return [dict(family=f"R x {tokens}", **r) for r in out] + + +# ---------------------------------------------------------------------------------------------------------------- +# Tests +# ---------------------------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize("rows,tokens", SPLITS, ids=[f"{r}x{t}" for r, t in SPLITS]) +def test_scatter(rows, tokens): + bad = [r for r in measure_scatter(rows, tokens) if not r["ok"]] + assert not bad, bad + + +@pytest.mark.parametrize("rows,tokens", SPLITS, ids=[f"{r}x{t}" for r, t in SPLITS]) +def test_gather(rows, tokens): + bad = [r for r in measure_gather(rows, tokens) if not r["ok"]] + assert not bad, bad + + +@pytest.mark.parametrize("rows,tokens", SPLITS, ids=[f"{r}x{t}" for r, t in SPLITS]) +def test_stage(rows, tokens): + if not _pinned_ok(): + pytest.skip("StepInputStage needs pinned host memory (not preferred here)") + bad = [r for r in measure_stage(rows, tokens) if not r["ok"]] + assert not bad, bad + + +@pytest.mark.parametrize("tokens", TOKENS, ids=[f"Rx{t}" for t in TOKENS]) +def test_graph_replay(tokens): + bad = [r for r in measure_graph(tokens) if not r["ok"]] + assert not bad, bad + + +def test_contract(): + """rows == 0 launches nothing; every call the kernels do not cover returns False and launches nothing, so the + caller keeps its torch path; StepInputStage's per-step limits (max_copies = 4 copies, max_block_copies = 2 block + copies, one gather) decline the call past them.""" + op = _op() + assert op.is_supported() + st = Step(MAX_BATCH, STORE_WIDTH, _seed(6, 0, 0, 0)) + b = clone_bufs(st.bufs) + before = clone_bufs(b) + scatter = list(st.scatter_args(b)) + gather = list(st.gather_args(b, "a")) + stage = op.StepInputStage(STAGE_CAPACITY) + assert op.SlotScatter().scatter(*scatter[:2], 0, *scatter[3:]) + stage.begin() + assert stage.gather(*gather[:5], 0, *gather[6:]) + stage.commit() + torch.cuda.synchronize() + assert not mismatches(before, b), "rows == 0 wrote" + # Argument 7 (draft width) == tokens, argument 6 (tokens) wider than the 8-wide store, argument 15 (KV begin) + # ending past the engine's 8 KV-length offsets, and an int64 slot table (argument 3). + for index, value in ( + (7, STORE_WIDTH), + (6, STORE_WIDTH + 1), + (15, 1), + (3, b["prev_slots"].long()), + ): + args = list(gather) + args[index] = value + stage.begin() + assert not stage.gather(*args), f"gather argument {index} = {value!r} was staged" + # Outputs of 8 rows with rows 1 .. 8 asked, an int64 slot table, a non-contiguous output. + assert not op.SlotScatter().scatter(scatter[0], 1, *scatter[2:]) + assert not op.SlotScatter().scatter(*scatter[:3], b["slot_table"].long(), *scatter[4:]) + strided = dict( + scatter[0], next_draft_tokens=torch.empty_like(b["out_draft"]).t().contiguous().t() + ) + assert not op.SlotScatter().scatter(strided, *scatter[1:]) + if _pinned_ok(): + stage.begin() + for _ in range(stage.max_copies): + assert stage.copy(b["extra"][:2], torch.tensor([1, 2], dtype=torch.int32)) + assert not stage.copy(b["extra"][:2], torch.tensor([1, 2], dtype=torch.int32)) + stage.begin() + assert not stage.copy(b["extra"][:2], torch.tensor([1, 2])) # int64 values + assert not stage.copy( + b["out_new"], torch.tensor([1, 2], dtype=torch.int32) + ) # 2-D destination + assert stage.gather(*gather) + assert not stage.gather(*gather) + blocks = st.block_copies(b) + for args in blocks: + assert stage.block_copy(*args) + assert not stage.block_copy(*blocks[0]) + stage.begin() # drops the staged step: nothing was launched + unpinned = blocks[0][1].clone() # a pageable host table + assert not unpinned.is_pinned() + assert not stage.block_copy(blocks[0][0], unpinned, *blocks[0][2:]) + stage.begin() + torch.cuda.synchronize() + assert not mismatches(before, b), "a declined call wrote" + + +# ---------------------------------------------------------------------------------------------------------------- +# Error table (python3 test_spec_step_copies.py report) and timing (python3 test_spec_step_copies.py time) +# ---------------------------------------------------------------------------------------------------------------- + + +def report() -> int: + print(f"{torch.cuda.get_device_name()}") + ok_all = True + sections = [("SlotScatter.scatter", measure_scatter), ("StepInputStage.gather", measure_gather)] + if _pinned_ok(): + sections.append(("StepInputStage", measure_stage)) + else: + print("StepInputStage: skipped (pinned host memory not preferred here)") + for name, measure in sections: + print(f"\n## {name}\n") + print("| split | case | identical | result |") + print("| :-- | :-- | :-- | :-- |") + for rows, tokens in SPLITS: + for r in measure(rows, tokens): + ok_all &= r["ok"] + print(f"| {rows}x{tokens} | {r['case']} | {r['identical']} | {'PASS' if r['ok'] else 'FAIL'} |", + flush=True) # fmt: skip + print("\n## CUDA graph replays (one capture per family R = 1 .. 8)\n") + print("| family | case | identical | result |") + print("| :-- | :-- | :-- | :-- |") + for tokens in TOKENS: + for r in measure_graph(tokens): + ok_all &= r["ok"] + print(f"| {r['family']} | {r['case']} | {r['identical']} | {'PASS' if r['ok'] else 'FAIL'} |", + flush=True) # fmt: skip + print("\nALL PASS" if ok_all else "\nFAIL") + return 0 if ok_all else 1 + + +def time_graph(body, calls, replays=15): + """Per-call us of a CUDA graph of ``calls`` back-to-back ``body(i)``: median (min, max) over ``replays`` replays. + ``body(-1)`` runs first, outside capture.""" + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + body(-1) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for i in range(calls): + body(i) + torch.cuda.synchronize() + for _ in range(3): + graph.replay() + torch.cuda.synchronize() + per_call = [] + for _ in range(replays): + start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + start.record() + graph.replay() + end.record() + torch.cuda.synchronize() + per_call.append(start.elapsed_time(end) * 1e3 / calls) + return statistics.median(per_call), min(per_call), max(per_call) + + +class StageArm: + """Back-to-back staged commits: fresh stages for every capture (a captured commit's stage cannot stage again), + and one of its own for the warm-up call.""" + + def __init__(self, step, bufs, calls): + self.step, self.bufs, self.calls = step, bufs, calls + self.warm = _op().StepInputStage(STAGE_CAPACITY) + self.fresh = [] + + def __call__(self, i): + if i < 0: + self.fresh = [_op().StepInputStage(STAGE_CAPACITY) for _ in range(self.calls)] + self.step.run_stage(self.bufs, self.warm) + else: + self.step.run_stage(self.bufs, self.fresh[i]) + + +def timing() -> None: + op = _op() + calls = 32 + with_stage = _pinned_ok() + print(f"{torch.cuda.get_device_name()}; graphs of {calls} back-to-back calls on one step's buffers (the engine's " + "layout), 15 replays: median (min-max) us per call") # fmt: skip + print("| split | scatter: torch | scatter: kernel | gather: torch | gather: kernel | stage: torch " + "| stage: kernel |") # fmt: skip + print("| :-- | --: | --: | --: | --: | --: | --: |") + for rows, tokens in SPLITS: + st = Step(rows, tokens, _seed(5, rows, tokens, 0), row_begin=MAX_BATCH - rows, out_rows=MAX_BATCH, + begins=engine_begins(rows, tokens)) # fmt: skip + b = st.bufs + gather_stage = op.StepInputStage(STAGE_CAPACITY) + arms = [ + lambda i: torch_scatter(*st.scatter_args(b)), + lambda i: kernel_scatter(st.scatter_args(b)), + lambda i: torch_gather(*st.gather_args(b, "a")), + lambda i: kernel_gather(gather_stage, st.gather_args(b, "a")), + ] + if with_stage: + arms += [lambda i: st.torch_stage(b), StageArm(st, b, calls)] + res = [[] for _ in arms] + for rep in range(3): # alternating order + order = range(len(arms)) if rep % 2 == 0 else reversed(range(len(arms))) + for a in order: + res[a].append(time_graph(arms[a], calls)) + cells = [] + for timings in res: + meds = sorted(x[0] for x in timings) + lo, hi = min(x[1] for x in timings), max(x[2] for x in timings) + cells.append(f"{meds[1]:.2f} ({lo:.2f}-{hi:.2f})") + cells += ["n/a"] * (6 - len(cells)) + print(f"| {rows}x{tokens} | " + " | ".join(cells) + " |", flush=True) + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "time": + timing() + elif len(sys.argv) > 1 and sys.argv[1] == "report": + sys.exit(report()) + else: + sys.exit(pytest.main([__file__, "-q", "-p", "no:cacheprovider", *sys.argv[1:]])) diff --git a/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py b/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py new file mode 100644 index 000000000000..c4305439b5a4 --- /dev/null +++ b/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py @@ -0,0 +1,119 @@ +# 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. +"""``copy_to_device_if_changed``: which calls copy and what the buffer then holds. + +The skip decision does not depend on the device, so the buffer here is a host tensor: a skipped copy shows as a +sentinel written behind the function's back that survives the call. +""" + +import pytest +import torch + +from tensorrt_llm._utils import copy_to_device_if_changed + +pytestmark = pytest.mark.cpu_only + +SENTINEL = -7 + + +def _host(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int32) + + +def _tamper(dst: torch.Tensor) -> None: + """Overwrite the buffer without the function knowing, so only a real copy restores it.""" + dst.fill_(SENTINEL) + + +def test_first_call_copies() -> None: + dst = torch.zeros(8, dtype=torch.int32) + copy_to_device_if_changed(dst, _host([3, 1, 2])) + assert dst.tolist() == [3, 1, 2, 0, 0, 0, 0, 0] + + +def test_unchanged_values_skip_the_copy() -> None: + dst = torch.zeros(8, dtype=torch.int32) + copy_to_device_if_changed(dst, _host([3, 1, 2])) + _tamper(dst) + copy_to_device_if_changed(dst, _host([3, 1, 2])) + assert dst.eq(SENTINEL).all() + + +def test_changed_values_copy() -> None: + dst = torch.zeros(8, dtype=torch.int32) + copy_to_device_if_changed(dst, _host([3, 1, 2])) + _tamper(dst) + copy_to_device_if_changed(dst, _host([3, 1, 4])) + assert dst[:3].tolist() == [3, 1, 4] + assert dst[3:].eq(SENTINEL).all(), "only the leading elements are written" + + +def test_prefix_of_the_last_values_skips_the_copy() -> None: + dst = torch.zeros(8, dtype=torch.int32) + copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) + _tamper(dst) + copy_to_device_if_changed(dst, _host([5, 6])) + assert dst.eq(SENTINEL).all() + + +def test_longer_values_copy_and_extend_what_is_kept() -> None: + dst = torch.zeros(8, dtype=torch.int32) + copy_to_device_if_changed(dst, _host([5, 6])) + copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) + assert dst[:4].tolist() == [5, 6, 7, 8] + _tamper(dst) + copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) + assert dst.eq(SENTINEL).all() + + +def test_shorter_changed_values_keep_the_tail() -> None: + # A shorter copy rewrites only its leading elements; the buffer's tail still holds the earlier values, and what + # the function keeps says so. + dst = torch.zeros(8, dtype=torch.int32) + copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) + copy_to_device_if_changed(dst, _host([9, 9])) + assert dst[:4].tolist() == [9, 9, 7, 8] + _tamper(dst) + copy_to_device_if_changed(dst, _host([9, 9, 7, 8])) + assert dst.eq(SENTINEL).all() + copy_to_device_if_changed(dst, _host([9, 9, 7, 1])) + assert dst[:4].tolist() == [9, 9, 7, 1] + + +def test_host_values_may_be_reused_right_away() -> None: + dst = torch.zeros(4, dtype=torch.int32) + host = _host([1, 2, 3]) + copy_to_device_if_changed(dst, host) + host.fill_(0) + assert dst[:3].tolist() == [1, 2, 3] + copy_to_device_if_changed(dst, _host([1, 2, 3])) + assert dst[:3].tolist() == [1, 2, 3] + + +def test_a_new_buffer_starts_over() -> None: + first = torch.zeros(4, dtype=torch.int32) + copy_to_device_if_changed(first, _host([1, 2])) + second = torch.full((4,), SENTINEL, dtype=torch.int32) + copy_to_device_if_changed(second, _host([1, 2])) + assert second[:2].tolist() == [1, 2] + + +def test_two_dimensional_values_are_flattened() -> None: + dst = torch.zeros(6, dtype=torch.int32) + copy_to_device_if_changed(dst, torch.tensor([[1, 2], [3, 4]], dtype=torch.int32)) + assert dst[:4].tolist() == [1, 2, 3, 4] + _tamper(dst) + copy_to_device_if_changed(dst, _host([1, 2, 3, 4])) + assert dst.eq(SENTINEL).all() diff --git a/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py b/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py index 6c41f7821351..6b0a1b1d75d4 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py @@ -330,3 +330,82 @@ def forward_fn(inputs): logits_eager = forward_fn(self._make_inputs(attn_metadata, num_tokens, batch_size, value=2)) torch.testing.assert_close(logits_cuda_graph, logits_eager) + + +class TestEngineBuffersAsStaticInputs: + """With static_input_ids / static_position_ids, the graphs' static inputs are views of the engine's own buffers: + the engine's writes are the graphs' inputs, and replay copies an input only when it is another tensor.""" + + def test_graph_reads_the_engine_buffers(self): + batch_size = 1 + engine_input_ids = torch.zeros((8,), device="cuda", dtype=torch.int32) + engine_position_ids = torch.zeros((8,), device="cuda", dtype=torch.int32) + runner = create_mock_cuda_graph_runner( + batch_size, + max_num_tokens=8, + static_input_ids=engine_input_ids, + static_position_ids=engine_position_ids, + ) + assert runner.shared_static_tensors["input_ids"].data_ptr() == engine_input_ids.data_ptr() + assert ( + runner.shared_static_tensors["position_ids"].data_ptr() + == engine_position_ids.data_ptr() + ) + key = KeyType(batch_size=batch_size, draft_len=0, is_first_draft=False) + num_tokens = runner._get_num_tokens_for_key(key) + attn_metadata = object() + + def forward_fn(inputs): + return inputs["input_ids"] * 1000 + inputs["position_ids"][0] + + def engine_inputs(): + return { + "attn_metadata": attn_metadata, + "input_ids": engine_input_ids[:num_tokens], + "position_ids": engine_position_ids[:num_tokens].unsqueeze(0), + } + + engine_input_ids.fill_(3) + engine_position_ids.fill_(4) + runner.capture(key, forward_fn, engine_inputs()) + # The captured forward did not run, so the engine's buffers are as they were. + assert engine_input_ids.tolist() == [3] * 8 + assert engine_position_ids.tolist() == [4] * 8 + + engine_input_ids[:num_tokens] = 7 + engine_position_ids[:num_tokens] = 9 + output = runner.replay(key, engine_inputs()) + assert output.tolist() == [7009] * num_tokens + + # Another tensor as the input is copied into the engine's buffers, which the graph reads. + other = { + "attn_metadata": attn_metadata, + "input_ids": torch.full((num_tokens,), 5, device="cuda", dtype=torch.int32), + "position_ids": torch.full((1, num_tokens), 6, device="cuda", dtype=torch.int32), + } + output = runner.replay(key, other) + assert output.tolist() == [5006] * num_tokens + assert engine_input_ids[:num_tokens].tolist() == [5] * num_tokens + + @pytest.mark.parametrize( + "input_ids, position_ids, use_mrope", + [ + pytest.param(torch.int64, torch.int32, False, id="int64-input-ids"), + pytest.param(torch.int32, None, False, id="no-position-ids"), + pytest.param(torch.int32, torch.int32, True, id="mrope"), + pytest.param("short", torch.int32, False, id="short"), + ], + ) + def test_unusable_engine_buffers_are_rejected(self, input_ids, position_ids, use_mrope): + size = 0 if input_ids == "short" else 8 # the graphs need one token here + dtype = torch.int32 if input_ids == "short" else input_ids + with pytest.raises(ValueError): + create_mock_cuda_graph_runner( + 1, + use_mrope=use_mrope, + max_num_tokens=8, + static_input_ids=torch.zeros((size,), device="cuda", dtype=dtype), + static_position_ids=None + if position_ids is None + else torch.zeros((8,), device="cuda", dtype=position_ids), + ) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index d86a355f6d21..e9b99e12dd82 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1575,6 +1575,134 @@ def test_promoted_context_precedes_speculative_overlap_generation( [generation.py_seq_slot], 0) kv_cache_manager.shutdown() + def _check_staged_spec_decode_graph_step( + self, use_kv_cache_manager_v2: bool) -> None: + """A CUDA graph decode step of a speculative engine whose step inputs + are staged on the StepInputStage (one launch after the attention + metadata's prepare) writes what the torch path writes.""" + from tensorrt_llm._torch.attention.backends.trtllm import \ + TrtllmAttentionMetadata + from tensorrt_llm._torch.cute_dsl_kernels.spec_step_copies import \ + op as spec_step_copies + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import \ + KVCacheManagerV2 + from tensorrt_llm.llmapi.llm_args import \ + KvCacheConfig as LlmKvCacheConfig + if not spec_step_copies.is_supported(): + self.skipTest("the step-copy kernels run on SM 100") + max_draft_len = 3 + tokens_per_step = max_draft_len + 1 + model_engine, kv_cache_manager = create_model_engine_and_kvcache( + spec_config=SADecodingConfig(max_draft_len=max_draft_len)) + if use_kv_cache_manager_v2: + kv_cache_manager.shutdown() + kv_cache_manager = KVCacheManagerV2( + LlmKvCacheConfig(max_tokens=512, enable_block_reuse=False), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=1, + num_kv_heads=model_engine.model.config.num_key_value_heads, + head_dim=model_engine.model.config.head_dim, + tokens_per_block=4, + max_seq_len=256, + max_batch_size=8, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=tensorrt_llm.bindings.DataType.HALF) + model_engine.runtime_draft_len = max_draft_len + stage = model_engine._step_input_stage + self.assertIsNotNone(stage) + resource_manager = ResourceManager( + {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) + + # Three requests that ran in the previous step, in slots out of order. + slots = [5, 2, 7] + requests = kv_cache_manager.add_dummy_requests( + [1, 2, 3], + token_nums=[10, 17, 33], + is_gen=True, + max_num_draft_tokens=max_draft_len) + for request, slot in zip(requests, slots): + request.is_dummy_request = False + request.py_seq_slot = slot + batch = ScheduledRequests() + batch.generation_requests = requests + attn_metadata = model_engine._set_up_attn_metadata(kv_cache_manager) + self.assertIs(type(attn_metadata), TrtllmAttentionMetadata) + graph_metadata = attn_metadata.create_cuda_graph_metadata( + len(requests), False, max_draft_len) + spec_metadata = Mock( + _force_non_greedy_for_capture=False, + context_prompt_lookahead_tokens=None, + ) + + # The sampler's slot stores as the previous step left them. + num_slots = 8 + generator = torch.Generator(device="cuda").manual_seed(0) + previous = SampleStateTensorsSpec( + new_tokens=torch.randint(0, + 1000, (tokens_per_step, num_slots, 1), + generator=generator, + dtype=torch.int32, + device="cuda"), + new_tokens_lens=torch.randint(1, + tokens_per_step + 1, (num_slots, ), + generator=generator, + dtype=torch.int32, + device="cuda"), + next_draft_tokens=torch.randint(0, + 1000, (num_slots, max_draft_len), + generator=generator, + dtype=torch.int32, + device="cuda"), + ) + + def step_inputs(step_input_stage): + """The buffers one step writes, from a filler they all start at.""" + model_engine._step_input_stage = step_input_stage + for request, slot in zip(requests, slots): + request.py_batch_idx = slot + written = (model_engine.input_ids_cuda, + model_engine.position_ids_cuda, + model_engine.draft_tokens_cuda, + model_engine.previous_pos_id_offsets_cuda, + model_engine.previous_kv_lens_offsets_cuda, + graph_metadata.prompt_lens_cuda, + graph_metadata.kv_lens_cuda, + graph_metadata.kv_cache_block_offsets) + for buffer in written: + buffer.fill_(-3) + model_engine._prepare_tp_inputs(scheduled_requests=batch, + kv_cache_manager=kv_cache_manager, + attn_metadata=graph_metadata, + spec_metadata=spec_metadata, + new_tensors_device=previous, + resource_manager=resource_manager) + torch.cuda.synchronize() + return [buffer.clone() for buffer in written] + + torch_path = step_inputs(None) + with patch.object(stage, "commit", wraps=stage.commit) as commit, \ + patch.object(stage, "block_copy", + wraps=stage.block_copy) as block_copy: + staged = step_inputs(stage) + commit.assert_called_once() + self.assertEqual(block_copy.call_count, + 1 if use_kv_cache_manager_v2 else 0) + for expected, actual in zip(torch_path, staged): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + # The previous rows' input ids came from the stores, by slot. + num_tokens = len(requests) * tokens_per_step + expected_input_ids = previous.new_tokens[:, slots, 0].t().reshape(-1) + self.assertEqual(staged[0][:num_tokens].tolist(), + expected_input_ids.tolist()) + kv_cache_manager.shutdown() + + def test_staged_spec_decode_graph_step_matches_torch_path(self) -> None: + self._check_staged_spec_decode_graph_step(use_kv_cache_manager_v2=False) + + def test_staged_spec_decode_graph_step_matches_torch_path_kv_v2( + self) -> None: + self._check_staged_spec_decode_graph_step(use_kv_cache_manager_v2=True) + def test_multimodal_encoder_max_seq_len(self) -> None: class CapturingEncoder(torch.nn.Module, MultimodalEncoderMixin): diff --git a/tests/unittest/_torch/helpers.py b/tests/unittest/_torch/helpers.py index c2adb2b9b68d..7350a10e0dda 100644 --- a/tests/unittest/_torch/helpers.py +++ b/tests/unittest/_torch/helpers.py @@ -236,9 +236,12 @@ def block_scale_gemm(mat_a: torch.Tensor, mat_scale_a: torch.Tensor, return results.view_as(x) -def create_mock_cuda_graph_runner(batch_size: int, - use_mrope: bool = False, - max_num_tokens: int = 1): +def create_mock_cuda_graph_runner( + batch_size: int, + use_mrope: bool = False, + max_num_tokens: int = 1, + static_input_ids: Optional[torch.Tensor] = None, + static_position_ids: Optional[torch.Tensor] = None): config = CUDAGraphRunnerConfig( use_cuda_graph=True, cuda_graph_padding_enabled=False, @@ -256,7 +259,9 @@ def create_mock_cuda_graph_runner(batch_size: int, is_encoder_decoder=False, mapping=Mapping(), dist=None, - kv_cache_manager_key=ResourceManagerType.KV_CACHE_MANAGER) + kv_cache_manager_key=ResourceManagerType.KV_CACHE_MANAGER, + static_input_ids=static_input_ids, + static_position_ids=static_position_ids) return CUDAGraphRunner(config) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_rng_window_counter.py b/tests/unittest/_torch/speculative/hw_agnostic/test_rng_window_counter.py index ae735562317f..23bab3ac0aec 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_rng_window_counter.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_rng_window_counter.py @@ -20,7 +20,9 @@ import types from typing import Optional -from tensorrt_llm._torch.speculative.interface import SpecMetadata +import torch + +from tensorrt_llm._torch.speculative.interface import DEFAULT_SAMPLING_SEED, SpecMetadata MAX_DRAFT_LEN = 3 WINDOW = MAX_DRAFT_LEN + 1 @@ -180,3 +182,37 @@ def test_graph_copy_shares_the_counters() -> None: assert _offsets(meta, [_request(1)]) == [0] assert _offsets(graph_meta, [_request(1)]) == [WINDOW] assert _offsets(meta, [_request(1)]) == [2 * WINDOW] + + +# --- all-greedy batches ------------------------------------------------------ + + +def _populated(meta: SpecMetadata, requests: list[types.SimpleNamespace]) -> list[int]: + """Run _populate_request_rng_state on CPU buffers and return request_offsets.""" + for request in requests: + request.sampling_config = types.SimpleNamespace(seed=request.seed) + normalized = [(0.0, 0, 1.0, 0.0, WINDOW) for _ in requests] + meta._populate_request_rng_state(requests, normalized) + return meta.request_offsets[: len(requests)].tolist() + + +def test_all_greedy_batch_skips_the_copies_but_advances_the_windows() -> None: + # The argmax graph reads none of the Philox buffers, so an all-greedy batch + # leaves them as they are; its windows are still taken, so the next sampled + # batch gets the same offsets it would have had. + meta = _meta() + sentinel = -5 + meta.temperatures = torch.ones(8 * WINDOW) + meta.request_seeds = torch.full((8,), sentinel, dtype=torch.int64) + meta.request_offsets = torch.full((8,), sentinel, dtype=torch.int64) + meta.seeds = torch.full((8 * WINDOW,), sentinel, dtype=torch.int64) + meta.offsets = torch.full((8 * WINDOW,), sentinel, dtype=torch.int64) + + meta.is_all_greedy_sample = True + assert _populated(meta, [_request(0, seed=7), _request(1)]) == [sentinel, sentinel] + assert meta.seeds.eq(sentinel).all() and meta.offsets.eq(sentinel).all() + + meta.is_all_greedy_sample = False + assert _populated(meta, [_request(0, seed=7), _request(1)]) == [WINDOW, WINDOW] + assert meta.request_seeds[:2].tolist() == [7, DEFAULT_SAMPLING_SEED] + assert meta.offsets[: 2 * WINDOW].tolist() == [WINDOW] * (2 * WINDOW) From f6cb23b6fec9022728bc2277cddfe781a8bf0579 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 1 Oct 2026 00:08:13 -0700 Subject: [PATCH 02/67] [None][feat] KDA verify: optional per-token states in the V2 hybrid cache manager MambaHybridCacheManagerV2 takes kda_token_states (default off). With the KDA replay caches it then also keeps the fp32 state after every draft of the last verify round per slot (kda_state_tok, [layers, slots, num_spec, H, V, K]), passes the layer's view in the speculative state and moves it with the slot. A verify kernel can start the next round from the state after the accepted drafts instead of replaying them. The cache factory enables it when the model's backbone sets kda_token_states. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/pyexecutor/_util.py | 8 ++ .../kv_cache/mamba_cache_manager.py | 55 +++++++- .../kv_cache/test_mamba_cache_manager.py | 127 +++++++++++++++++- 3 files changed, 185 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index f830d3c23191..2028050aa034 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3048,6 +3048,14 @@ def _create_kv_cache_manager( if is_kda_mtp_verify_available(): kda_extra_kwargs["kda_replay_num_spec"] = ( spec_config.tokens_per_gen_step - 1) + # A model whose KDA verify starts from the state after the + # accepted drafts (instead of replaying them) asks for the + # state after every draft with `kda_token_states`. + backbone = getattr(getattr(model_engine, "model", None), + "model", None) + if (issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2) + and getattr(backbone, "kda_token_states", False)): + kda_extra_kwargs["kda_token_states"] = True if is_glm5_next: # The manager places an indexer buffer on every attention layer. from ..attention.backends.sparse.glm_kpool import \ diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index ad417618950b..ed5cef76dd35 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -330,6 +330,14 @@ def use_kda_replay_update(self) -> bool: """Whether KDA fused verification owns per-slot replay caches.""" return getattr(self, "_use_kda_replay_update", False) + @property + def keeps_kda_token_states(self) -> bool: + """Whether the KDA replay caches come with the state after every + draft of the last verify round (``kda_state_tok``), so that a verify + kernel can start the next round from the state after the accepted + drafts instead of replaying them.""" + return False + @abstractmethod def get_conv_states(self, layer_idx: int) -> torch.Tensor: """Return conv states for specific layer. @@ -438,6 +446,14 @@ class SpeculativeState(State): kda_qkg_cache: torch.Tensor | None = None kda_v_cache: torch.Tensor | None = None kda_beta_cache: torch.Tensor | None = None + # Optional: the state after every draft of the last verify round, + # [slots, num_spec, H, V, K] fp32 (the SSM pool holds the state + # after the round's first, non-draft token). A verify kernel starts + # the next round from entry n - 1 when the round accepted n > 0 + # drafts (prev_num_accepted_tokens) instead of replaying them from + # the caches above; a slot reset for a new request has n = 0, so a + # previous owner's entries are never read. + kda_state_tok: torch.Tensor | None = None # Replay path: compact double-buffered cache # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). @@ -3024,6 +3040,7 @@ def __init__( conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", kda_replay_num_spec: Optional[int] = None, qwen4_exp_ple_cache_params: Optional["Qwen4ExpPLECacheParams"] = None, + kda_token_states: bool = False, **kwargs, ) -> None: if conv_state_layout not in ("x_b_c", "q_k_v"): @@ -3040,6 +3057,9 @@ def __init__( self._kda_replay_num_spec = kda_replay_num_spec self._use_kda_replay_update = kda_replay_num_spec is not None + # Only with the KDA replay caches: also keep the state after every + # verify token (kda_state_tok). + self._kda_token_states = kda_token_states and self._use_kda_replay_update if self._use_kda_replay_update: if use_replay_state_update: raise ValueError( @@ -3420,6 +3440,10 @@ def get_disagg_role_layouts(self) -> Dict[DataRole, RoleLayout]: def use_gdn_cached_replay_all_layer_commit(self) -> bool: return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) + @property + def keeps_kda_token_states(self) -> bool: + return getattr(self, "kda_state_tok", None) is not None + def _setup_mtp_intermediate_states(self, spec_config, max_batch_size: int) -> None: if not self.use_kda_replay_update: @@ -3450,6 +3474,7 @@ def _allocate_pool_replay_buffers( self.kda_qkg_cache = None self.kda_v_cache = None self.kda_beta_cache = None + self.kda_state_tok = None if (not self.use_kda_replay_update or self.local_num_mamba_layers == 0): return allocated @@ -3504,8 +3529,17 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: self.ssm_state_shape[0]), device, ) + if self._kda_token_states: + self.kda_state_tok = torch.zeros( + (self.local_num_mamba_layers, cache_size, num_spec, + *self.ssm_state_shape), + dtype=torch.float32, + device=device, + ) + per_token = (", with per-token states" + if self.kda_state_tok is not None else "") logger.info("Mamba Cache (kda-replay) is allocated for " - f"{cache_size} state slots") + f"{cache_size} state slots{per_token}") return True def _commit_gdn_cached_replay_history_layers( @@ -3801,6 +3835,12 @@ def _max_resident_sequences(self) -> int: def _mamba_state_bytes_per_slot(self) -> int: base_bytes = self.local_num_mamba_layers * (self.ssm_bytes + self.conv_bytes) + if getattr(self, "_kda_token_states", False): + # kda_state_tok: an fp32 SSM state per draft of every slot. + base_bytes += (self.local_num_mamba_layers * + self._kda_replay_num_spec * + math.prod(self.ssm_state_shape) * + torch.float32.itemsize) local_ple_layers = sum(layer_id in self.pp_layers for layer_id in self._ple_layer_ids) if local_ple_layers == 0: @@ -4285,16 +4325,20 @@ def _relocate_kda_replay_slots(self, old_slots: List[int], destination_slots = torch.tensor([new for _, new in moves], dtype=torch.long, device=device) - replay_buffers = ( + replay_buffers = [ self.kda_conv_q, self.kda_conv_k, self.kda_conv_v, self.kda_qkg_cache, self.kda_v_cache, self.kda_beta_cache, - ) + ] + assert all(replay_buffer is not None + for replay_buffer in replay_buffers) + # The per-token verify states belong to the slot as well. + if self.kda_state_tok is not None: + replay_buffers.append(self.kda_state_tok) for replay_buffer in replay_buffers: - assert replay_buffer is not None replay_buffer.index_copy_( 1, destination_slots, @@ -4498,6 +4542,8 @@ def mamba_layer_cache( "kda_v_cache": self.kda_v_cache[layer_offset], "kda_beta_cache": self.kda_beta_cache[layer_offset], } + if self.kda_state_tok is not None: + spec_kwargs["kda_state_tok"] = self.kda_state_tok[layer_offset] if self.mamba_ssm_rand_seed is not None: spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed return PythonMambaCacheManager.SpeculativeState( @@ -4781,4 +4827,5 @@ def shutdown(self): self.kda_qkg_cache = None self.kda_v_cache = None self.kda_beta_cache = None + self.kda_state_tok = None super().shutdown() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 1f920c7b718c..0079ea8632a2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Regression tests for Python, Cpp, and V2 Mamba cache managers.""" +import math import os from types import SimpleNamespace from unittest.mock import MagicMock @@ -300,12 +301,16 @@ def _kimi_model_config() -> SimpleNamespace: quant_config=None, sparse_attention_config=None, get_num_mamba_layers=lambda: 2, + # get_kv_cache_manager_cls reads the mapping and spec config (helix check). + mapping=None, + spec_config=None, ) def _capture_kimi_v2_manager_ctor( monkeypatch: pytest.MonkeyPatch, spec_config=None, + model_engine=None, ) -> tuple[tuple, dict]: """Route a Kimi config through _create_kv_cache_manager with an explicit V2 manager and capture the constructor arguments.""" @@ -324,7 +329,7 @@ def __init__(self, *args: object, **kwargs: object) -> None: assert get_kv_cache_manager_cls(model_config, kv_cache_config) is MambaHybridCacheManagerV2 _create_kv_cache_manager( - model_engine=None, + model_engine=model_engine, kv_cache_manager_cls=RecordingV2Manager, mapping=Mapping(world_size=1, tp_size=1, pp_size=1), kv_cache_config=kv_cache_config, @@ -380,6 +385,33 @@ def test_kimi_explicit_v2_manager_enables_kda_replay( assert kwargs["kda_replay_num_spec"] == spec_config.tokens_per_gen_step - 1 assert kwargs["conv_state_layout"] == "q_k_v" + assert "kda_token_states" not in kwargs + + +@pytest.mark.parametrize("backbone_asks", [True, False]) +def test_kimi_explicit_v2_manager_kda_token_states_follow_the_model( + monkeypatch: pytest.MonkeyPatch, + backbone_asks: bool, +) -> None: + """The factory asks the V2 manager for per-token KDA verify states only when + the model's backbone sets ``kda_token_states``.""" + monkeypatch.setattr( + "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", + lambda: True, + ) + backbone = SimpleNamespace(kda_token_states=True) if backbone_asks else SimpleNamespace() + model_engine = SimpleNamespace( + model=SimpleNamespace(model=backbone), + _max_cuda_graph_batch_size=4, + is_multimodal=True, + ) + + _, kwargs = _capture_kimi_v2_manager_ctor( + monkeypatch, MTPDecodingConfig(max_draft_len=3), model_engine + ) + + assert kwargs["kda_replay_num_spec"] == 3 + assert kwargs.get("kda_token_states", False) is backbone_asks def test_kimi_explicit_v2_manager_uses_qkv_convolution_layout( @@ -2353,6 +2385,7 @@ def _build_v2_hybrid_with_mamba_layer( mamba_n_groups=1, mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, + kda_token_states=False, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2407,6 +2440,7 @@ def _build_v2_hybrid_with_mamba_layer( dtype=dtype, conv_state_layout=conv_state_layout, kda_replay_num_spec=kda_replay_num_spec, + kda_token_states=kda_token_states, ) @@ -3629,6 +3663,43 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): ) assert layer_cache.intermediate_ssm is None assert layer_cache.intermediate_conv_window is None + assert layer_cache.kda_state_tok is None + assert not mgr.keeps_kda_token_states + finally: + mgr.shutdown() + + +@skip_no_cuda +@pytest.mark.parametrize("kda_replay_num_spec", [2, None]) +def test_v2_kda_token_states_allocated_with_the_replay_caches(kda_replay_num_spec): + """With the replay caches, ``kda_token_states`` adds the fp32 state after every + draft of each slot, per layer; without them it allocates nothing.""" + mgr = _build_v2_hybrid_with_mamba_layer( + max_batch_size=4, + num_mamba_layers=2, + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_d_conv=5, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=kda_replay_num_spec, + kda_token_states=True, + ) + try: + if kda_replay_num_spec is None: + assert getattr(mgr, "kda_state_tok", None) is None + assert not mgr.keeps_kda_token_states + return + assert mgr.keeps_kda_token_states + for layer_idx in range(2): + layer_cache = mgr.mamba_layer_cache(layer_idx) + cache_size = layer_cache.temporal.shape[0] + states = layer_cache.kda_state_tok + assert states.shape == (cache_size, 2, *layer_cache.temporal.shape[1:]) + assert states.dtype is torch.float32 + assert states.data_ptr() == mgr.kda_state_tok[layer_idx].data_ptr() + assert not states.any() finally: mgr.shutdown() @@ -3785,6 +3856,60 @@ def test_v2_kda_replay_relocates_live_slot_history(): mgr.shutdown() +@skip_no_cuda +def test_v2_kda_token_states_relocate_with_their_slot(): + mgr = _build_v2_hybrid_with_mamba_layer( + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_n_groups=4, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, + kda_token_states=True, + ) + try: + states = mgr.kda_state_tok + assert states is not None + states.zero_() + states[:, 0].fill_(1.0) + states[:, 1].fill_(2.0) + source_zero, source_one = states[:, 0].clone(), states[:, 1].clone() + + mgr._relocate_kda_replay_slots([0, 1], [1, 2]) + + torch.testing.assert_close(states[:, 1], source_zero, rtol=0, atol=0) + torch.testing.assert_close(states[:, 2], source_one, rtol=0, atol=0) + finally: + mgr.shutdown() + # Released with the other replay buffers. + assert mgr.kda_state_tok is None + + +@skip_no_cuda +def test_v2_kda_token_states_count_in_the_per_slot_budget(): + """The capacity math sees the per-token states: a slot costs one more fp32 SSM state per draft and layer.""" + kwargs = dict( + num_mamba_layers=2, + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_d_conv=5, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, + ) + plain = _build_v2_hybrid_with_mamba_layer(**kwargs) + try: + plain_bytes = plain._mamba_state_bytes_per_slot() + finally: + plain.shutdown() + with_states = _build_v2_hybrid_with_mamba_layer(kda_token_states=True, **kwargs) + try: + per_token = 2 * 2 * math.prod(with_states.ssm_state_shape) * 4 + assert with_states._mamba_state_bytes_per_slot() == plain_bytes + per_token + finally: + with_states.shutdown() + + @skip_no_cuda def test_v2_kda_state_index_setup_relocates_generation_history(): class StateCache: From b44c537918a606c01f5187dfae94aeb5d1fbd2c6 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 15:41:56 -0700 Subject: [PATCH 03/67] [None][fix] spec_step_copies: store zeros past a row's accepted tokens The one-model worker's acceptance returns new_tokens as torch.empty [N, K + 1]. A context row writes only column 0 and accepts one token; a generation row writes every column. The scatter kernel copied every column of every row into the sampler's new_tokens store, so for context rows it read columns 1..K, which were never written; compute-sanitizer initcheck reports those reads. update_requests reads the store only up to new_tokens_lens, so no output changed. The scatter kernel now reads a row's new tokens only below its new_tokens_lens and stores zeros past them. A thread loads the row's length once; the column 0 thread's store_lens write reuses that load. The sampler's torch store update, used where the kernel does not run, is unchanged: nothing reads past new_tokens_lens. test_spec_step_copies: the reference zeroes the same columns, and the Step cases draw accepted lengths from 0 to past the output width. test_scatter_zeros_past_accepted_length covers a mixed batch laid out as the worker writes it, with a skipped chunked-context row. "poisoned" fills the never-written columns with a sentinel and fails on the old kernel; "never_written" leaves them unwritten for compute-sanitizer initcheck. Signed-off-by: Vasanth Sabavat --- .../cute_dsl_kernels/spec_step_copies/op.py | 3 +- .../spec_step_copies_kernel.py | 10 ++- .../cute_dsl_kernels/test_spec_step_copies.py | 64 ++++++++++++++++++- 3 files changed, 71 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py index c0b2e56c3658..9bc81fef63bb 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py @@ -94,7 +94,8 @@ class SlotScatter: """The speculative sampler's store update as one kernel. ``store[..., slots[r]] = outputs[row_begin + r]`` for every row ``r``, each row padded with zeros or cut to its - store's width. + store's width. A row's ``new_tokens`` at or past its ``new_tokens_lens`` are stored as zeros: the forward writes + only column 0 of a context row, and readers of the store stop at that length. """ _kernel: ClassVar[Any] = None diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py index 9780847687b4..269da57f8d01 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py @@ -15,7 +15,8 @@ """The one-model speculative decoding step's eager copy passes, one kernel each. * ``scatter_kernel``: the sampler moves the forward's per-row outputs into its slot-indexed stores (``index_copy_`` - of new tokens, next new tokens, accepted lengths and next draft tokens, each padded or cut to its store width). + of new tokens, next new tokens, accepted lengths and next draft tokens, each padded or cut to its store width; + a row's new tokens past its accepted length are written as zeros). * ``stage_kernel``: one launch writes a decode step's per-step inputs: the overlap scheduler's gathers from those stores for the rows whose request ran in the previous step (``index_select`` into input ids and draft tokens by slot, into the position offsets by the per-token index list, and into the KV-length offsets by slot, minus the @@ -76,11 +77,14 @@ def scatter_kernel( col = t - row * columns src_row = row_begin + row slot = _i32_at(slots, row)[0] + accepted = _i32_at(out_lens, src_row)[0] zero = cutlass.Int32(0) if col < new_width: + # Only the row's accepted tokens are read: past them a context row's forward output is never written. value = zero if col < out_new_width: - value = _i32_at(out_new_tokens, src_row * out_new_width + col)[0] + if col < accepted: + value = _i32_at(out_new_tokens, src_row * out_new_width + col)[0] _i32_at(store_new_tokens, col * num_slots + slot)[0] = value if col < next_width: value = zero @@ -93,7 +97,7 @@ def scatter_kernel( value = _i32_at(out_next_draft, src_row * out_draft_width + col)[0] _i32_at(store_next_draft, slot * draft_width + col)[0] = value if col == zero: - _i32_at(store_lens, slot)[0] = _i32_at(out_lens, src_row)[0] + _i32_at(store_lens, slot)[0] = accepted @cute.jit diff --git a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py index 10408b60106d..a28ad5b659b1 100644 --- a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py +++ b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py @@ -18,8 +18,11 @@ 8 or 16 slots; draft-token buffer 56, KV-length offsets 8, per-token buffers max_num_tokens). * ``SlotScatter.scatter`` vs SpecSampler's torch store update (each output padded with zeros or cut to its store's - width, then four ``index_copy_`` by slot): outputs narrower than (dynamic draft length), as wide as and wider than the - stores, and mixed per field; row_begin 0 and > 0; shuffled distinct slot tables. + width, the new tokens at or past a row's accepted length zeroed, then four ``index_copy_`` by slot): outputs + narrower than (dynamic draft length), as wide as and wider than the stores, and mixed per field; row_begin 0 and + > 0; shuffled distinct slot tables; accepted lengths from 0 to past the output width. A mixed context / generation + batch laid out as the one-model worker writes it (a context row's tokens past column 0 never written) stores zeros + there. * ``StepInputStage.gather`` committed on its own vs PyTorchModelEngine._prepare_tp_inputs's torch overlap gathers (the stores by slot into the input ids and draft tokens, the lengths by the per-token index list into the position offsets, ``new_tokens_lens - T`` by slot into the KV-length offsets): the engine's offsets (the requests without a @@ -135,6 +138,8 @@ def fit(t, width): # pad with zeros or truncate to the store width return t[:, :width] o_new_tokens = fit(o_new_tokens, store_new_tokens.shape[0]) + columns = torch.arange(o_new_tokens.shape[1], device=o_new_tokens.device) + o_new_tokens = torch.where(columns < o_new_tokens_lens[:, None], o_new_tokens, 0) o_next_draft_tokens = fit(o_next_draft_tokens, store_next_draft_tokens.shape[1]) o_next_new_tokens = fit(o_next_new_tokens, store_next_new_tokens.shape[0]) store_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens) @@ -309,6 +314,8 @@ def rewrite(self) -> None: for t in b.values(): if t.is_cuda: t.copy_(rand_i32(t.shape, g)) + # Accepted lengths from none to past the new-token width, so every row cuts its new tokens somewhere. + b["out_lens"].copy_(rand_i32(b["out_lens"].shape, g, 0, self.widths[0] + 2)) slots = shuffled_distinct(rows, num_slots, g) table = rand_i32((num_slots,), g, 0, num_slots) table[:rows] = torch.tensor(slots, dtype=torch.int32) @@ -638,6 +645,59 @@ def test_graph_replay(tokens): assert not bad, bad +POISON = 0x5EEDF00D # what the never-written columns hold in the poisoned case + + +@pytest.mark.parametrize("poison", [True, False], ids=["poisoned", "never_written"]) +def test_scatter_zeros_past_accepted_length(poison): + """A mixed batch as the one-model worker's acceptance leaves its outputs (``new_tokens`` is ``torch.empty`` + [N, K + 1]; a context row writes its first token only and accepts 1, a generation row writes every column): the + store holds each row's accepted tokens, then zeros, never a never-written column. ``poisoned``: those columns hold + POISON. ``never_written``: they keep the allocation's contents, so compute-sanitizer initcheck reports any read of + them. Row 0 is a context row whose chunk is not its last (row_begin 1 skips it).""" + op = _op() + g = torch.Generator().manual_seed(49) + skipped, num_contexts, accepted = 1, 3, [1, 3, MAX_DRAFT + 1, 5] + n = skipped + num_contexts + len(accepted) + rows, row_begin, num_slots = n - skipped, skipped, 2 * MAX_BATCH + width = MAX_DRAFT + 1 + contexts = skipped + num_contexts + tokens = rand_i32((n, width), g, 0, 1 << 20) + lens = torch.tensor([1] * contexts + accepted, dtype=torch.int32) + new_tokens = torch.empty((n, width), dtype=torch.int32, device="cuda") + if poison: + new_tokens.fill_(POISON) + new_tokens[:contexts, 0] = tokens[:contexts, 0].cuda() + new_tokens[contexts:] = tokens[contexts:].cuda() + outputs = { + "new_tokens": new_tokens, + "new_tokens_lens": lens.cuda(), + "next_new_tokens": rand_i32((n, width), g).cuda(), + "next_draft_tokens": rand_i32((n, MAX_DRAFT), g).cuda(), + } + slots = shuffled_distinct(rows, num_slots, g) + slot_table = torch.tensor(slots + [0] * (num_slots - rows), dtype=torch.int32, device="cuda") + stores = [ + torch.full(shape, -1, dtype=torch.int32, device="cuda") + for shape in ( + (STORE_WIDTH, num_slots, 1), + (STORE_WIDTH, num_slots, 1), + (num_slots,), + (num_slots, MAX_DRAFT), + ) + ] + assert op.SlotScatter().scatter(outputs, row_begin, rows, slot_table, *stores) + store_new, store_lens = stores[0][:, :, 0].T.cpu(), stores[2].cpu() + want = torch.full((num_slots, STORE_WIDTH), -1, dtype=torch.int32) + for r, s in enumerate(slots): + src = row_begin + r + want[s] = 0 + want[s, : lens[src]] = tokens[src, : lens[src]] + assert store_lens[s] == lens[src], (r, s) + assert not (store_new == POISON).any(), "a never-written column reached the store" + assert torch.equal(store_new, want), (store_new, want) + + def test_contract(): """rows == 0 launches nothing; every call the kernels do not cover returns False and launches nothing, so the caller keeps its torch path; StepInputStage's per-step limits (max_copies = 4 copies, max_block_copies = 2 block From 92d67f38d78dde4f756a661e9cc919ad2ed1bfed Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 15:42:23 -0700 Subject: [PATCH 04/67] [None][test] spec_step_copies tests: side stream waits for the current stream measure_graph (test_graph_replay) and time_graph build their inputs on the current stream, then make their first eager call on a fresh side stream without ordering it after that work. Under the stream-ordered cudaMallocAsync allocator the call can write memory before its allocation in the side stream's order; with any allocator it can read inputs whose initialization has not finished. The side stream now waits for the current stream first. Signed-off-by: Vasanth Sabavat --- tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py index a28ad5b659b1..a2a44221d1da 100644 --- a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py +++ b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py @@ -562,6 +562,7 @@ def run_torch(st, b): st.torch_stage(b) stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for st in steps: # the kernels compile on their first call, outside capture run_kernels(st, clone_bufs(st.bufs), stage) @@ -798,6 +799,7 @@ def time_graph(body, calls, replays=15): """Per-call us of a CUDA graph of ``calls`` back-to-back ``body(i)``: median (min, max) over ``replays`` replays. ``body(-1)`` runs first, outside capture.""" stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): body(-1) torch.cuda.synchronize() From f3cf23f5fd63b007e3a1b0ffabdebec8dfa2719a Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 15:44:12 -0700 Subject: [PATCH 05/67] [None][fix] KV cache manager V2: keep the SSM pool at its live floor with per-token KDA states With per-token KDA states, kda_state_tok holds num_spec recurrent states per SSM state slot, outside the cache quota, addressed by the request's SSM slot (_allocate_pool_replay_buffers). V2 sizes the SSM pool group by the typical step's ratio, so its slot count, and kda_state_tok with it, grows with the quota. Without block reuse a slot past the min-slots floor (resident sequences plus reserved dummies) never holds a request, so those slots only add memory. On the Kimi K3 stack (7 DSpark drafts, TP16, max_seq_len 4096, free_gpu_memory_fraction 0.25) the pool got 183 slots: 32.4 GiB of kda_state_tok beside a 13.2 GiB quota. It ran out of device memory above fraction 0.28 at max_seq_len 4096, and above 0.16 at 512. With per-token KDA states and block reuse off, the typical step's request capacity is now raised until its attention pages alone take the quota. The ratio then gives the SSM pool group fewer slots than its floor, so the min-slots constraint sets its size and attention gets the rest. Without per-token states (replay caches only), or with block reuse, the ratio sizes the pool as before. test_v2_kda_token_states_keep_ssm_pool_at_live_floor builds the V2 storage for a 128 MiB quota with 2 MiB state slots: with per-token states the SSM pool keeps its floor (plus one for max_util_for_resume) and attention takes the rest of the quota; without them, or with block reuse, the pool keeps its ratio share. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 30 ++++++++ .../kv_cache/test_mamba_cache_manager.py | 72 +++++++++++++++++++ 2 files changed, 102 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index ed5cef76dd35..6b60514f8dfa 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3985,6 +3985,32 @@ def _minimum_live_gpu_quota(self) -> int: state_quota + attention_block_quota, ) + def _ssm_pool_at_live_floor(self, kv_cache_config: KvCacheConfig) -> bool: + """Whether the SSM pool group keeps only its min-slots floor (one slot + per resident sequence plus the reserved dummies) instead of its share + of the typical step's ratio. + + With per-token KDA states, ``kda_state_tok`` takes num_spec states per + SSM slot outside the cache quota (``_allocate_pool_replay_buffers``), + addressed by the request's SSM slot. Without block reuse no SSM slot + holds anything but a live request or a reserved dummy, so slots past + the floor are never used and would only add that memory. + """ + return (getattr(self, "_kda_token_states", False) + and self.local_num_mamba_layers > 0 + and not kv_cache_config.enable_block_reuse + and self._attention_cache_bytes_per_token() > 0) + + def _quota_filling_request_capacity(self, gpu_quota: int) -> int: + """A typical request capacity whose resident requests' attention pages + alone take ``gpu_quota``. The typical step's ratio then gives the SSM + pool group fewer slots than its floor, so the min-slots constraint + sets its size and attention gets the rest of the quota.""" + bytes_per_token = (self._max_resident_sequences() * + self._attention_cache_bytes_per_token()) + # KVCacheDesc.capacity is a 32-bit int. + return min(-(-gpu_quota // bytes_per_token), 1 << 30) + def _build_cache_config( self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: kv_cache_config = self.kv_cache_config @@ -4030,6 +4056,10 @@ def _build_cache_config( if config.initial_pool_ratio is None: typical_capacity = self._get_typical_request_capacity( kv_cache_config) + if self._ssm_pool_at_live_floor(kv_cache_config): + typical_capacity = max( + typical_capacity, + self._quota_filling_request_capacity(gpu_quota)) request_descs = self._typical_request_descs(typical_capacity, kv_cache_config) typical_step = BatchDesc(request_descs * diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 0079ea8632a2..df05a69a015e 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3747,6 +3747,78 @@ def test_v2_kda_replay_validates_configuration( ) +@pytest.mark.parametrize( + ("kda_token_states", "enable_block_reuse", "at_floor"), + [(True, False, True), (False, False, False), (True, True, False)], +) +def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( + kda_token_states, enable_block_reuse, at_floor +): + """The per-token KDA states take memory per SSM slot outside the cache quota. Without block reuse the SSM pool + keeps only its live floor and attention gets the rest of the quota; otherwise (replay caches only, or block reuse) + the typical step's ratio sizes it.""" + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._generation_kv_capacity_headroom = 1 + mgr._has_cp_helix = False + mgr.kv_cache_type = CacheTypeCpp.SELF + mgr.head_dim_per_layer = [64, 64] + mgr.pp_layers = [0, 1] + mgr._mamba_layer_mask = [True, False] + # One 2 MiB grain per state slot, so the pool's slot count follows its grains (a Kimi K3 slot holds 27 MB). + mgr.ssm_bytes = 2 << 20 + mgr.conv_bytes = 64 << 10 + mgr.max_attention_window_vec = [128, 128] + mgr.max_batch_size = 2 + mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) + mgr.max_seq_len = 128 + mgr.max_num_tokens = 128 + mgr.tokens_per_block = 32 + mgr.num_local_layers = 2 + mgr.local_num_mamba_layers = 1 + mgr._num_reserved_dummy_slots = 1 + mgr.dtype = DataType.HALF + mgr.enable_swa_scratch_reuse = False + mgr.enable_stats = False + mgr.num_extra_kv_tokens = 0 + mgr.get_layer_bytes_per_token = lambda **kwargs: 8 + # Layer 1's 256-byte pages of 32 tokens (_base_attention_layer_configs; layer 0 becomes the SSM layer). + mgr._attention_cache_bytes_per_token = lambda: 8 + mgr._use_kda_replay_update = True + mgr._kda_token_states = kda_token_states + # The per-token states count in each slot's state bytes (_mamba_state_bytes_per_slot): num_spec fp32 states of + # the SSM state shape, 2 MiB each here. + mgr._kda_replay_num_spec = 2 + mgr.ssm_state_shape = [8, 256, 256] + mgr.kv_cache_config = KvCacheConfig( + avg_seq_len=64, + enable_block_reuse=enable_block_reuse, + enable_partial_reuse=False, + ) + base_config = KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=128 << 20)], + layers=_base_attention_layer_configs(2), + ) + runtime_manager = RuntimeKVCacheManager(mgr._build_cache_config(base_config)) + try: + slots = {} + for stats in runtime_manager.get_storage_statistics(): + sizes = stats.slot_sizes if hasattr(stats, "slot_sizes") else stats.slot_size + role = "ssm" if mgr.ssm_bytes in [int(s) for s in sizes] else "attention" + slots[role] = int(stats.total) + finally: + runtime_manager.shutdown() + + floor = mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots + if at_floor: + # The floor's min-slots constraint is scaled by 1 / max_util_for_resume (0.97). + assert floor <= slots["ssm"] <= floor + 1 + # Every other 2 MiB grain goes to attention's 256-byte pages. + assert slots["attention"] >= ((128 << 20) // (2 << 20) - (floor + 2)) * ((2 << 20) // 256) + else: + assert slots["ssm"] > 10 * floor + + def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: mgr = object.__new__(MambaCacheManager) mgr._impl = SimpleNamespace(use_kda_replay_update=True) From 67d93e9115247f2d26523b0646cec93877840445 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 10:52:37 -0700 Subject: [PATCH 06/67] [None][fix] multi_stream_utils: record the calling stream on fn1's outputs maybe_execute_in_parallel runs fn1 on aux_stream and returns its outputs to the caller, which uses them on its own stream after the event join. Those tensors are allocated in aux_stream order. When the caller drops them, the caching allocator returns their blocks to aux_stream's pool at once and cudaMallocAsync frees them in aux_stream order, so later aux_stream work can take and overwrite the memory while the calling stream's reads are still queued. compute-sanitizer memcheck with cudaMallocAsync reports this as use-after-free on the outputs Kimi K3's KDA mixer and MoE shared expert return this way. The CUDA tensors fn1 returns, directly or in nested tuples, lists and dict values, are now recorded on the calling stream after the join. cudaMallocAsync raises in record_stream on memory it did not allocate (a buffer from outside PyTorch, wrapped as a tensor); only that error is caught, and the tensor is skipped, as the native allocator does. Nothing is recorded under CUDA graph capture or torch.compile tracing. In a capture the caching allocator would hold every recorded block freed during the capture until it ends, so each call's outputs would take new graph pool memory, and under cudaMallocAsync freeing a recorded output inside the capture fails it with "capturing stream has unjoined work". A captured graph relies instead on later aux_stream work waiting on the calling stream, as the next call's event0.wait() does. Dynamo traces neither the capture query nor record_stream, so a compiled caller keeps the previous behavior and still traces fullgraph. test_multi_stream_utils.py reads fn1's output on the calling stream behind a sleep, drops it, and fills a new aux_stream allocation of the same size at once; the read must still see fn1's values. Before, it saw the fill under both allocators. The file also covers a buffer from cudaMalloc, a capture whose outputs outlive it, outputs dropped inside a capture, and a fullgraph torch.compile of a caller. It reruns itself under PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync and joins the l0_h100 and l0_b200 pre-merge module tests. Signed-off-by: Vasanth Sabavat --- .../_torch/modules/multi_stream_utils.py | 45 +++- .../test_lists/test-db/l0_b200.yml | 1 + .../test_lists/test-db/l0_h100.yml | 1 + .../_torch/modules/test_multi_stream_utils.py | 208 ++++++++++++++++++ 4 files changed, 254 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/modules/test_multi_stream_utils.py diff --git a/tensorrt_llm/_torch/modules/multi_stream_utils.py b/tensorrt_llm/_torch/modules/multi_stream_utils.py index f3e60a208646..c8b3c3af0ee3 100644 --- a/tensorrt_llm/_torch/modules/multi_stream_utils.py +++ b/tensorrt_llm/_torch/modules/multi_stream_utils.py @@ -32,6 +32,29 @@ def with_multi_stream(enable: bool): set_do_multi_stream(prev_do_multi_stream) +def _record_stream(obj: Any, stream: torch.cuda.Stream) -> None: + """Record ``stream`` on every CUDA tensor in ``obj``: a tensor, or tuples, + lists and dict values of them, nested. Other objects are not walked.""" + if isinstance(obj, torch.Tensor): + if obj.is_cuda: + try: + obj.record_stream(stream) + except RuntimeError as e: + # The cudaMallocAsync backend asserts on memory it did not + # allocate (a buffer from outside PyTorch, wrapped as a + # tensor). It never frees that memory: nothing to record. + if (torch.cuda.memory.get_allocator_backend() + != "cudaMallocAsync" + or "ptr not found in ptr_info" not in str(e)): + raise + elif isinstance(obj, (tuple, list)): + for item in obj: + _record_stream(item, stream) + elif isinstance(obj, dict): + for item in obj.values(): + _record_stream(item, stream) + + def maybe_execute_in_parallel( fn0: Callable, fn1: Callable, @@ -46,6 +69,15 @@ def maybe_execute_in_parallel( use case. For simplicity, fn0 and fn1 do not support inputs. + fn1()'s outputs are allocated in aux_stream order and used on the calling + stream. Outside CUDA graph capture and torch.compile tracing, the CUDA + tensors fn1() returns (directly, or in nested tuples, lists and dict + values; not tensors held by other objects) are recorded on the calling + stream, so their memory outlives the caller's use. During capture nothing + is recorded: a captured graph relies on every later use of aux_stream + first waiting for the calling stream, as the next call's event0.wait() + does. Other work captured on aux_stream must do the same. + Args: fn0 (Callable): callable for the default stream fn1 (Callable): callable for the second stream, aux_stream @@ -60,7 +92,7 @@ def maybe_execute_in_parallel( leave this as False since custom ops are opaque to the compiler. Returns: - tuple[Any, Any]: the return values of fn0() and fn1() + tuple[Any, Any]: the return values of fn0() and fn1(). """ multi_stream = (do_multi_stream() and aux_stream is not None and @@ -75,6 +107,17 @@ def maybe_execute_in_parallel( result1 = fn1() event1.record() event1.wait() + # Without the record, freeing fn1's outputs hands their memory to later + # aux_stream work (caching allocator) or frees it in aux_stream order + # (cudaMallocAsync) while the calling stream's reads may still be + # queued. Under graph capture the caching allocator would hold every + # recorded block freed in the capture until it ends, and under + # cudaMallocAsync freeing a recorded output inside the capture fails it + # ("capturing stream has unjoined work"). Dynamo traces neither the + # capture query nor record_stream, so compiled callers skip both. + if not (torch.compiler.is_compiling() + or torch.cuda.is_current_stream_capturing()): + _record_stream(result1, torch.cuda.current_stream()) else: result0 = fn0() result1 = fn1() diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 2ca039299f9c..c3ff228e1c35 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -120,6 +120,7 @@ l0_b200: - unittest/llmapi/test_deepseek_v4_tokenizer.py - unittest/_torch/modules/test_mhc.py - unittest/_torch/modules/test_engram.py + - unittest/_torch/modules/test_multi_stream_utils.py # ------------- Kimi K3 (KimiLinear) unit tests --------------- - unittest/_torch/modeling/test_kimi_k3_checkpoint.py TIMEOUT (15) - unittest/_torch/modeling/test_kimi_k3_mla.py TIMEOUT (15) diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 31659d64e83c..1deaf60afd8e 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -39,6 +39,7 @@ l0_h100: - unittest/_torch/modules/mamba - unittest/_torch/modules/tests_lora_modules - unittest/_torch/modules/test_w4a16_nvfp4_linear.py + - unittest/_torch/modules/test_multi_stream_utils.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/test_multi_stream_utils.py b/tests/unittest/_torch/modules/test_multi_stream_utils.py new file mode 100644 index 000000000000..91166f14134c --- /dev/null +++ b/tests/unittest/_torch/modules/test_multi_stream_utils.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""maybe_execute_in_parallel: the tensors fn1 returns stay valid for the calling stream. + +fn1 runs on aux_stream, so its outputs are allocated in aux_stream order, while the caller +reads them on its own stream. Once the caller drops them, the allocator must not give their +memory to new aux_stream work until the caller's queued reads have run: the caching allocator +returns a freed block to the pool of the stream it was allocated on, and cudaMallocAsync frees +it in that stream's order. Under CUDA graph capture the outputs are not recorded; dropping them +inside a capture must neither grow the graph pool nor fail the capture. A compiled caller must +still trace as one graph. The tests need no sanitizer. They run under the process's allocator backend, and +test_this_file_under_cuda_malloc_async reruns them under cudaMallocAsync. +""" + +import os +import subprocess +import sys + +import pytest +import torch + +from tensorrt_llm._torch.modules.multi_stream_utils import ( + maybe_execute_in_parallel, + with_multi_stream, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") + +# About 25 ms at 2 GHz: the calling stream is still sleeping when the host has dropped fn1's +# output and queued the next aux_stream fill, which takes microseconds. +_SLEEP_CYCLES = 50_000_000 + + +def _wrap(out: torch.Tensor, shape: str): + if shape == "tuple_list": + return (None, [out]) + if shape == "dict": + return {"out": out, "none": None} + return out + + +def _unwrap(result, shape: str) -> torch.Tensor: + if shape == "tuple_list": + return result[1][0] + if shape == "dict": + return result["out"] + return result + + +@pytest.mark.parametrize("shape", ["tensor", "tuple_list", "dict"]) +def test_aux_outputs_outlive_the_calling_streams_reads(shape: str) -> None: + numel = 1 << 20 + aux_stream = torch.cuda.Stream() + event0, event1 = torch.cuda.Event(), torch.cuda.Event() + + def fn1(): + return _wrap(torch.full((numel,), 1.0, device="cuda"), shape) + + torch.cuda.synchronize() + # A free block of this size left in aux_stream's pool by earlier work could take the fill + # below instead of fn1's block, and hide a missing record. + torch.cuda.empty_cache() + with with_multi_stream(True): + _, result1 = maybe_execute_in_parallel(lambda: None, fn1, event0, event1, aux_stream) + out = _unwrap(result1, shape) + torch.cuda._sleep(_SLEEP_CYCLES) + # Queued behind the sleep: this read of fn1's output runs after everything below. + copy = out.clone() + del out, result1 + # The next aux_stream allocation of the same size, filled at once (aux_stream is idle). + # Without the record it takes fn1's freed memory and overwrites it before the read above. + with torch.cuda.stream(aux_stream): + torch.full((numel,), 2.0, device="cuda") + torch.cuda.synchronize() + assert torch.equal(copy, torch.ones_like(copy)) + + +def test_aux_output_over_memory_torch_did_not_allocate() -> None: + # The native allocator skips such memory in record_stream; cudaMallocAsync raises on it. + cudart = pytest.importorskip("cuda.bindings.runtime") + numel = 1024 + err, ptr = cudart.cudaMalloc(numel * 4) + assert err == cudart.cudaError_t.cudaSuccess + + class _Buffer: + def __init__(self): + self.__cuda_array_interface__ = { + "shape": (numel,), + "typestr": " None: + aux_stream = torch.cuda.Stream() + event0, event1 = torch.cuda.Event(), torch.cuda.Event() + x = torch.randn(4096, device="cuda") + + def fn0(): + return x * 2 + + def fn1(): + return x + 1, [x - 1], {"y": x * 3} + + graph = torch.cuda.CUDAGraph() + with with_multi_stream(True): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + maybe_execute_in_parallel(fn0, fn1, event0, event1, aux_stream) + torch.cuda.current_stream().wait_stream(side) + with torch.cuda.graph(graph): + r0, (a, [b], c) = maybe_execute_in_parallel(fn0, fn1, event0, event1, aux_stream) + out = r0 + a + b + c["y"] + for _ in range(3): + x.copy_(torch.randn_like(x)) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, x * 2 + (x + 1) + (x - 1) + x * 3) + + +def test_aux_outputs_dropped_inside_graph_capture() -> None: + layers, numel = 8, 1 << 20 + aux_stream = torch.cuda.Stream() + event0, event1 = torch.cuda.Event(), torch.cuda.Event() + # Values below 2^24: the fp32 sums are exact. + x = (torch.arange(numel, device="cuda") % 4096).to(torch.float32) + addresses = [] + + def layer_stack(): + acc = torch.zeros_like(x) + for layer in range(layers): + _, out = maybe_execute_in_parallel( + lambda: None, lambda layer=layer: x + layer, event0, event1, aux_stream + ) + addresses.append(out.data_ptr()) + acc = acc + out + del out + # cudaMallocAsync frees the outputs in aux_stream order: the capture must join it. + torch.cuda.current_stream().wait_stream(aux_stream) + return acc + + graph = torch.cuda.CUDAGraph() + with with_multi_stream(True): + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + layer_stack() + torch.cuda.current_stream().wait_stream(side) + addresses.clear() + with torch.cuda.graph(graph): + acc = layer_stack() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(acc, x * layers + sum(range(layers))) + if torch.cuda.memory.get_allocator_backend() == "native": + # A recorded block freed during capture is held until the capture ends. + assert len(set(addresses)) == 1, f"{len(set(addresses))} blocks for {layers} outputs" + + +def test_compiled_caller_traces_fullgraph() -> None: + # Dynamo traces neither the capture query nor record_stream, so a compiled caller must take + # the path without the record and still give one graph. + aux_stream = torch.cuda.Stream() + event0, event1 = torch.cuda.Event(), torch.cuda.Event() + + def caller(x): + a, b = maybe_execute_in_parallel(lambda: x * 2, lambda: x + 1, event0, event1, aux_stream) + return a + b + + torch._dynamo.reset() + compiled = torch.compile(caller, backend="eager", fullgraph=True) + x = torch.randn(16, device="cuda") + with with_multi_stream(True): + out = compiled(x) + torch.cuda.synchronize() + torch.testing.assert_close(out, 3 * x + 1) + + +def test_this_file_under_cuda_malloc_async() -> None: + if torch.cuda.memory.get_allocator_backend() == "cudaMallocAsync": + pytest.skip("this process already runs under cudaMallocAsync") + env = dict(os.environ, PYTORCH_CUDA_ALLOC_CONF="backend:cudaMallocAsync") + result = subprocess.run( + [sys.executable, "-m", "pytest", "-q", "-p", "no:cacheprovider", __file__], + capture_output=True, + text=True, + env=env, + timeout=1200, + ) + assert result.returncode == 0, result.stdout[-6000:] + result.stderr[-3000:] From 5e16d82b3171d04bc37c76dc17d372d25d82e025 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 15:47:02 -0700 Subject: [PATCH 07/67] [None][fix] MNNVL all-reduce: keep a grown workspace's predecessor alive for captured graphs Growing the MNNVL all-reduce workspace replaced it and freed the previous buffers, pointer tables and flags while CUDA graphs captured before the growth still launched on them. The replaced workspaces now stay alive with the process, and creating or growing a workspace during graph capture raises instead of failing inside the driver. test_mnnvl_workspace_growth_keeps_captured_graphs captures a one-token all-reduce, grows the workspace with an eager two-shot call, replays the graph and checks growth under capture is refused. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/distributed/ops.py | 14 +++ .../_torch/multi_gpu/test_mnnvl_allreduce.py | 107 ++++++++++++++++++ 2 files changed, 121 insertions(+) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 6d20a9aa6394..31d06468876d 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -287,6 +287,10 @@ def _get_or_scale_allreduce_mnnvl_workspace( if mapping not in allreduce_mnnvl_workspaces or allreduce_mnnvl_workspaces[ mapping]["buffer_size_bytes"] < (buffer_size_bytes or 0): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "MNNVL all-reduce workspace creation or growth during CUDA graph capture: run each shape once " + "outside capture first") # Initial buffer to be large enough to support 1024 tokens * 8192 hidden_dim init_buffer_size_bytes = max(1024 * 8192 * elem_size, buffer_size_bytes or 0) @@ -363,6 +367,11 @@ def _get_or_scale_allreduce_mnnvl_workspace( _initialize_allreduce_mnnvl_protocol(candidate_workspace) # Hand ownership of the communicator to the workspace. pending_comms.pop(mapping, None) + previous_workspace = allreduce_mnnvl_workspaces.get(mapping) + if previous_workspace is not None: + # CUDA graphs captured before this growth keep launching on the previous buffers and flags. + MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.setdefault( + mapping, []).append(previous_workspace) allreduce_mnnvl_workspaces[mapping] = candidate_workspace return allreduce_mnnvl_workspaces[mapping] @@ -754,6 +763,11 @@ class MNNVLAllReduce(nn.Module): allreduce_mnnvl_workspaces: typing.ClassVar[dict[Mapping, _MnnvlWorkspace]] = {} + # Workspaces a larger one replaced. CUDA graphs captured before the growth still launch on their buffers and + # flags, so they stay alive with the process. The checkpoint hooks cover only the current workspace. + allreduce_mnnvl_retired_workspaces: typing.ClassVar[dict[ + Mapping, list[_MnnvlWorkspace]]] = {} + # Communicators split for a mapping whose workspace construction has not # succeeded yet. Ownership moves to the workspace once it is published, so # an entry here is never reachable from allreduce_mnnvl_workspaces. diff --git a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py index f0abca2492bf..1a19977e4ad7 100644 --- a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py +++ b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py @@ -334,6 +334,96 @@ def mnnvl_checkpoint_worker_env(_: int): } +def mnnvl_growth_graph_forward(tensor_parallel_size: int, + tensor_parallel_rank: int) -> bool: + """A graph captured on the MNNVL workspace replays correctly after an eager call grew the workspace, and growth + during capture is refused.""" + env_names = ("TLLM_TEST_MNNVL", "TRTLLM_FORCE_MNNVL_AR") + previous_env = { + name: (name in os.environ, os.environ.get(name)) + for name in env_names + } + tensor_parallel_rank = tensorrt_llm.mpi_rank() + torch.cuda.set_device(tensor_parallel_rank) + os.environ["TLLM_TEST_MNNVL"] = "1" + os.environ["TRTLLM_FORCE_MNNVL_AR"] = "1" + mapping = None + try: + MPI.COMM_WORLD.barrier() + mapping = Mapping( + world_size=tensor_parallel_size, + tp_size=tensor_parallel_size, + rank=tensor_parallel_rank, + ) + MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None) + MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.pop(mapping, None) + gc.collect() + MPI.COMM_WORLD.barrier() + expected = tensor_parallel_size * (tensor_parallel_size + 1) // 2 + with torch.inference_mode(): + allreduce = AllReduce( + mapping=mapping, + strategy=AllReduceStrategy.MNNVL, + dtype=torch.bfloat16, + ) + input_ = torch.full((1, 7168), + tensor_parallel_rank + 1, + dtype=torch.bfloat16, + device="cuda") + allreduce(input_) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = allreduce(input_) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, + torch.full_like(output, expected)) + + before = MNNVLAllReduce.allreduce_mnnvl_workspaces[mapping] + # Two-shot footprint 2 * 4096 * 7168 * 2 B: well past the initial 16 MiB per Lamport buffer. + big = torch.ones((4096, 7168), dtype=torch.bfloat16, device="cuda") + torch.testing.assert_close( + allreduce(big), + torch.full_like(big, tensor_parallel_size)) + after = MNNVLAllReduce.allreduce_mnnvl_workspaces[mapping] + assert after["buffer_size_bytes"] > before["buffer_size_bytes"] + assert MNNVLAllReduce.allreduce_mnnvl_retired_workspaces[ + mapping] == [before] + assert before["handle"].is_mapped() + del big + + for step in range(3): + input_.fill_(tensor_parallel_rank + 2 + step) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + output, + torch.full_like(output, + expected + tensor_parallel_size * + (1 + step))) + + huge = torch.ones((8192, 7168), + dtype=torch.bfloat16, + device="cuda") + with pytest.raises(RuntimeError, + match="during CUDA graph capture"): + with torch.cuda.graph(torch.cuda.CUDAGraph()): + allreduce(huge) + return True + finally: + if mapping is not None: + MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None) + MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.pop( + mapping, None) + gc.collect() + for name, (was_present, value) in previous_env.items(): + if was_present: + os.environ[name] = value + else: + os.environ.pop(name, None) + + def mnnvl_checkpoint_rejects_wrong_membership(world_size: int, world_rank: int) -> bool: env_names = ("TLLM_TEST_MNNVL", "TRTLLM_FORCE_MNNVL_AR") @@ -705,6 +795,23 @@ def test_mnnvl_checkpoint_rejects_wrong_group_membership( assert all(results) +@pytest.mark.skipif( + platform.machine().lower() != "aarch64" or torch.cuda.device_count() < 2 + or not MnnvlMemory.supports_mnnvl(), + reason="requires at least two GB200 GPUs with fabric-backed MNNVL", +) +@pytest.mark.parametrize("mpi_pool_executor", [2], indirect=True) +def test_mnnvl_workspace_growth_keeps_captured_graphs( + mpi_pool_executor) -> None: + tensor_parallel_size = mpi_pool_executor.num_workers + results = mpi_pool_executor.map( + mnnvl_growth_graph_forward, + [tensor_parallel_size] * tensor_parallel_size, + range(tensor_parallel_size), + ) + assert all(results) + + def _make_quant_scale(reference_norm: torch.Tensor, fusion_op: AllReduceFusionOp) -> torch.Tensor: amax = reference_norm.abs().max().float().clamp_min(1e-6) From f0ff9a83ff2ef63e700b669ea7f0b40a9117cec6 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 1 Oct 2026 22:11:00 -0700 Subject: [PATCH 08/67] MNNVL all-reduce: run the workspace-growth graph test in the GB200 multi-GPU list test_mnnvl_workspace_growth_keeps_captured_graphs needs fabric-backed MNNVL on two GB200 GPUs, like the checkpoint graph tests next to it. Signed-off-by: Vasanth Sabavat --- tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 25a83b2d12f5..5102abe33232 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -50,6 +50,7 @@ l0_gb200_multi_gpus: - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_preserves_cuda_graph_addresses - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_checkpoint_rejects_wrong_group_membership + - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_growth_keeps_captured_graphs - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_engine_checkpoint_coordination - unittest/_torch/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_failure_is_collective_and_bounded From 551ca164479d6294dc1bbf8dbb27f1370f254225 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 15:47:47 -0700 Subject: [PATCH 09/67] [None][chore] MNNVL all-reduce test: yapf the workspace-growth test The workspace-growth test came from a branch whose yapf settings differ from main's; format it as main's pre-commit does. No behaviour change. Signed-off-by: Vasanth Sabavat --- .../_torch/multi_gpu/test_mnnvl_allreduce.py | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py index 1a19977e4ad7..4efd124b4820 100644 --- a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py +++ b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py @@ -384,8 +384,7 @@ def mnnvl_growth_graph_forward(tensor_parallel_size: int, # Two-shot footprint 2 * 4096 * 7168 * 2 B: well past the initial 16 MiB per Lamport buffer. big = torch.ones((4096, 7168), dtype=torch.bfloat16, device="cuda") torch.testing.assert_close( - allreduce(big), - torch.full_like(big, tensor_parallel_size)) + allreduce(big), torch.full_like(big, tensor_parallel_size)) after = MNNVLAllReduce.allreduce_mnnvl_workspaces[mapping] assert after["buffer_size_bytes"] > before["buffer_size_bytes"] assert MNNVLAllReduce.allreduce_mnnvl_retired_workspaces[ @@ -399,23 +398,18 @@ def mnnvl_growth_graph_forward(tensor_parallel_size: int, torch.cuda.synchronize() torch.testing.assert_close( output, - torch.full_like(output, - expected + tensor_parallel_size * - (1 + step))) - - huge = torch.ones((8192, 7168), - dtype=torch.bfloat16, - device="cuda") - with pytest.raises(RuntimeError, - match="during CUDA graph capture"): + torch.full_like( + output, expected + tensor_parallel_size * (1 + step))) + + huge = torch.ones((8192, 7168), dtype=torch.bfloat16, device="cuda") + with pytest.raises(RuntimeError, match="during CUDA graph capture"): with torch.cuda.graph(torch.cuda.CUDAGraph()): allreduce(huge) return True finally: if mapping is not None: MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None) - MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.pop( - mapping, None) + MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.pop(mapping, None) gc.collect() for name, (was_present, value) in previous_env.items(): if was_present: From df9d9bfc92aaed597c53d5c0a93299f6b1014b68 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 15:48:31 -0700 Subject: [PATCH 10/67] [None][test] MNNVL all-reduce: one-rank test of the workspace's capture guard Creating or growing the MNNVL all-reduce workspace under CUDA-graph capture raises before the collective allocation. Its only test was the last block of the 2-rank growth test, which runs on GB200 multi-GPU stages only. test_mnnvl_workspace_creation_refuses_graph_capture creates the workspace for a one-rank group inside a capture and expects the guard's RuntimeError. The guard runs before any MNNVL call, so one GPU suffices; without it the call enters the workspace construction under capture and fails inside it with another error. The test joins l0_b200's single-GPU PyTorch block. Signed-off-by: Vasanth Sabavat --- .../integration/test_lists/test-db/l0_b200.yml | 3 +++ .../_torch/multi_gpu/test_mnnvl_allreduce.py | 17 ++++++++++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 2ca039299f9c..aba43618bb35 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -148,6 +148,9 @@ l0_b200: - unittest/_torch/modeling/test_kimi_linear_checkpoint.py # CPU self-test of the Kimi K3 disagg parity harness comparison logic. - test_kimi_k3_specdec.py::test_kimi_k3_disagg_parity_selftest + # ------------- distributed --------------- + # One process: the MNNVL workspace's capture guard raises before any MNNVL call. + - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_creation_refuses_graph_capture # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_fused_add_rms_norm_quant.py - unittest/_torch/modules/test_fused_activation_quant.py diff --git a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py index 4efd124b4820..ce053dba16e4 100644 --- a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py +++ b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py @@ -29,7 +29,8 @@ from tensorrt_llm._mnnvl_utils import MnnvlMemory from tensorrt_llm._torch.distributed import (AllReduce, AllReduceFusionOp, AllReduceParams) -from tensorrt_llm._torch.distributed.ops import MNNVLAllReduce +from tensorrt_llm._torch.distributed.ops import ( + MNNVLAllReduce, get_or_scale_allreduce_mnnvl_workspace) from tensorrt_llm.functional import AllReduceStrategy from tensorrt_llm.mapping import Mapping @@ -806,6 +807,20 @@ def test_mnnvl_workspace_growth_keeps_captured_graphs( assert all(results) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a GPU") +def test_mnnvl_workspace_creation_refuses_graph_capture() -> None: + """Creating or growing the MNNVL workspace allocates collectively, so under + CUDA-graph capture it raises before entering the allocation. One process, + a group of one: the guard runs before any MNNVL call.""" + mapping = Mapping(world_size=1, rank=0, tp_size=1) + MNNVLAllReduce.allreduce_mnnvl_workspaces.pop(mapping, None) + graph, stream = torch.cuda.CUDAGraph(), torch.cuda.Stream() + with torch.cuda.graph(graph, stream=stream): + with pytest.raises(RuntimeError, match="during CUDA graph capture"): + get_or_scale_allreduce_mnnvl_workspace(mapping, torch.bfloat16) + assert mapping not in MNNVLAllReduce.allreduce_mnnvl_workspaces + + def _make_quant_scale(reference_norm: torch.Tensor, fusion_op: AllReduceFusionOp) -> torch.Tensor: amax = reference_norm.abs().max().float().clamp_min(1e-6) From fd1bddcaf67cf4bc90a526df5eb82fd1dcd7656a Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 1 Oct 2026 21:41:57 -0700 Subject: [PATCH 11/67] MNNVL all-reduce: count Lamport arrivals per CTA in the RMSNorm-fused kernels oneshotAllreduceFusionKernel and rmsNormLamport used LamportFlags' cluster arrival, in which every thread arrives on the cluster barrier but only cluster rank 0's first warp waits, and then arrived on the cluster barrier again in the RMSNorm epilogue's cluster reduction. Each thread must arrive and wait once per phase: a second arrival counted toward the first phase lets the reduction read a peer CTA's partial sum before it is written, and counts the cluster in before a late CTA has read the buffer flags. Both kernels now count arrivals per CTA, as the attn_res kernel already does; their cluster barrier is left to the reduction. Signed-off-by: Vasanth Sabavat --- cpp/tensorrt_llm/common/lamportUtils.cuh | 2 ++ .../communicationKernels/mnnvlAllreduceKernels.cu | 9 ++++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/cpp/tensorrt_llm/common/lamportUtils.cuh b/cpp/tensorrt_llm/common/lamportUtils.cuh index 60639e2f4870..28a06d1d062a 100644 --- a/cpp/tensorrt_llm/common/lamportUtils.cuh +++ b/cpp/tensorrt_llm/common/lamportUtils.cuh @@ -261,6 +261,8 @@ public: if constexpr (UseCGA) { #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + // Every thread arrives on the cluster barrier but only rank 0's first warp waits, so a kernel using this + // must not arrive on the cluster barrier again (each thread arrives and waits once per phase). cg::cluster_group cluster = cg::this_cluster(); __cluster_barrier_arrive(); if (cluster.block_rank() == 0 && threadIdx.x < kWARP_SIZE) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu index a8ff3dd431a8..876836c81f7e 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu @@ -597,8 +597,9 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed int threadOffset = token * params.tokenDim + packedIdx * kELTS_PER_THREAD; #endif - // We only use 1 stage for the oneshot allreduce - LamportFlags flag(params.bufferFlags, 1); + // We only use 1 stage for the oneshot allreduce. Arrivals are counted per CTA: the RMSNorm epilogue's cluster + // reduction below must be each thread's only phase of the cluster barrier. + LamportFlags flag(params.bufferFlags, 1); T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, 0)); T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0)); bool const inBounds = packedIdx * kELTS_PER_THREAD < params.tokenDim; @@ -1016,7 +1017,9 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam T* smemResidual = reinterpret_cast(&smem[smemBufferSize]); T* smemGamma = reinterpret_cast(&smem[2 * smemBufferSize]); - LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + // Arrivals are counted per CTA: with UseCGA the RMSNorm cluster reduction below must be each thread's only phase + // of the cluster barrier. + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); T* input = reinterpret_cast( flag.getCurLamportBuf(reinterpret_cast(params.bufferInputPtr), MNNVLTwoShotStage::BROADCAST)); From 7a7ccf360fe011b7e9089099771365b51e7aed7a Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 00:03:12 -0700 Subject: [PATCH 12/67] MNNVL all-reduce: start every CTA of the cluster before the RMSNorm cluster reduction writes The RESIDUAL_RMS_NORM epilogues of oneshotAllreduceFusionKernel and rmsNormLamport write each CTA's partial sum into the other CTAs' shared memory through the cluster mapping before any cluster barrier wait, so a target CTA may not have started yet; racecheck reports every such write once the kernels' arrival counting no longer stalls it. A cluster.sync() before the writes makes every CTA of the cluster start first, as the attn_res kernel's early arrive / wait does. Signed-off-by: Vasanth Sabavat --- .../kernels/communicationKernels/mnnvlAllreduceKernels.cu | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu index 876836c81f7e..e60a2f233c15 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu @@ -711,6 +711,8 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed fullSum = 0.F; // Need to reduce over the entire cluster int const blockRank = cluster.block_rank(); + // Every CTA of the cluster must have started before another CTA writes its shared memory. + cluster.sync(); if (threadIdx.x < numBlocks) { cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; @@ -1140,6 +1142,8 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam fullSum = 0.F; // Need to reduce over the entire cluster int const blockRank = cluster.block_rank(); + // Every CTA of the cluster must have started before another CTA writes its shared memory. + cluster.sync(); if (threadIdx.x < numBlocks) { cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; From f4672116460e12b094bb9548e0b34d628c74e77d Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 1 Oct 2026 21:10:23 -0700 Subject: [PATCH 13/67] [None][fix] Kimi K3 attn_res_fwd online kernel: order the ws_stats reads before the next chunk's writes attn_res_fwd_online_v2_kernel reuses one ws_stats buffer every chunk of candidates: lane 0 of each consumer warp writes its row, a named barrier orders the writes before every consumer's cross-warp reads, but nothing ordered those reads before the next chunk's writes. With more than one chunk (N > 4 at H 7168) a warp that finished reading could overwrite its row while a slower warp still read it, mixing the next chunk's statistics into this chunk's logits. Under compute-sanitizer racecheck, which serializes warps, every such call returned a wrong mixture (relative error 0.6-2.1 of max |ref|); natively it was not observed. Add the second barrier the persistent fork already has. Also order every lane's reads of a chunk's slots before lane 0 releases them to the producer (__syncwarp before the bar_consumed arrive), in online_v2 and in the persistent fork. Signed-off-by: Vasanth Sabavat --- .../kernels/kimiK3AttnRes/attnResFwd.cu | 15 +++++++++++++-- .../kimiK3AttnRes/attnResFwdPersistentFused.cu | 2 ++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index e714466dbb3a..952b54da16ab 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -411,9 +411,14 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c float2 f[2] = {__bfloat1622float2(v2[0]), __bfloat1622float2(v2[1])}; if constexpr (FULL_N12) { - if (n == AN - 1 && lane == 0) + if (n == AN - 1) { - cute::arrive_barrier(plan.bar_consumed[chunk_slot]); + // Every lane's reads of the chunk's slots before lane 0 releases them. + __syncwarp(); + if (lane == 0) + { + cute::arrive_barrier(plan.bar_consumed[chunk_slot]); + } } } tmem_st_32dp32bNx<4>( @@ -492,6 +497,8 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c } if constexpr (!FULL_N12) { + // Every lane's reads of the chunk's slots before lane 0 releases them to the producer. + __syncwarp(); if (lane == 0) { cute::arrive_barrier(plan.bar_consumed[chunk_slot]); @@ -551,6 +558,10 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c { cross_warp_tail(lane); } + // ws_stats is one buffer reused every chunk. The barrier above orders this chunk's writes before the + // reads; this one orders the reads before the next chunk's writes, so a warp that has finished + // reading cannot overwrite its row while a slower warp still reads it. + cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 0); float logit_n[N_CHUNK]; #pragma unroll for (int n = 0; n < N_CHUNK; n++) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu index 7245f724134a..32bfae257368 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu @@ -650,6 +650,8 @@ __global__ void __launch_bounds__(BLK, 1) default: __builtin_unreachable(); } } + // Every lane's reads of the chunk's slots before lane 0 releases them to the producer. + __syncwarp(); if (lane == 0) { mbarrier_arrive(plan.bar_consumed[chunk_slot]); From 26cb3c37cd70812e6d74040138709fc3b51cbe52 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 02:28:16 -0700 Subject: [PATCH 14/67] [None][fix] Kimi K3 attn_res: cross-proxy fence before a consumer releases a TMA-filled slot The online_v2, N = 1 tile and persistent fused kernels read their V (and delta) slots with generic-proxy shared loads and release them with an mbarrier arrive; the producer then refills them with cp.async.bulk (async proxy). Accesses to one location through two proxies need a cross-proxy fence: fence.proxy.async.shared::cta in every reading thread before the release. In the current SASS each release already issues after an instruction that consumes the last load, so this makes that ordering independent of instruction scheduling. Outputs are bit-identical. Signed-off-by: Vasanth Sabavat --- cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu | 9 +++++++++ .../kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu | 3 +++ 2 files changed, 12 insertions(+) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index 952b54da16ab..39a14df806b6 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -413,6 +413,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c { if (n == AN - 1) { + // The producer refills the slots through the async proxy (TMA): a + // cross-proxy fence orders this lane's generic-proxy reads before it. + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); // Every lane's reads of the chunk's slots before lane 0 releases them. __syncwarp(); if (lane == 0) @@ -497,6 +500,9 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c } if constexpr (!FULL_N12) { + // The producer refills the slots through the async proxy (TMA): a cross-proxy fence orders this + // lane's generic-proxy reads before it. + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); // Every lane's reads of the chunk's slots before lane 0 releases them to the producer. __syncwarp(); if (lane == 0) @@ -983,6 +989,9 @@ __global__ void __launch_bounds__(BLK, 1) } cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 1); } + // The producer refills the tile through the async proxy (TMA): a cross-proxy fence orders this thread's + // generic-proxy reads before it. + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); cute::arrive_barrier(plan.bar_consumed[slot]); } } diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu index 32bfae257368..ee734add0661 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu @@ -650,6 +650,9 @@ __global__ void __launch_bounds__(BLK, 1) default: __builtin_unreachable(); } } + // The producer refills the slots through the async proxy (TMA): a cross-proxy fence orders this lane's + // generic-proxy reads before it. + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); // Every lane's reads of the chunk's slots before lane 0 releases them to the producer. __syncwarp(); if (lane == 0) From 05e9e4f1f25555d9a91e89f95f344a07fa4d17d0 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 11:49:29 -0700 Subject: [PATCH 15/67] [None][test] Kimi K3 attn_res: check the cross-proxy fence before each slot release test_attn_res_proxy_fences.py reads attnResFwd.cu and attnResFwdPersistentFused.cu and checks every consumer release of a cp.async.bulk-filled slot (an arrive on bar_consumed). Each must follow fence.proxy.async.shared::cta, with only the warp sync and the lane-0 guard between the two. The fence orders the consumers' generic-proxy reads of the slot before the producer's async-proxy refill. Without it the order rests on instruction scheduling, which no test of values can see. The check fails on the sources from before the fence was added (four releases) and passes on the current ones. CPU only. Signed-off-by: Vasanth Sabavat --- .../test_attn_res_proxy_fences.py | 87 +++++++++++++++++++ 1 file changed, 87 insertions(+) create mode 100644 tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_proxy_fences.py diff --git a/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_proxy_fences.py b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_proxy_fences.py new file mode 100644 index 000000000000..b1a83d52af7f --- /dev/null +++ b/tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_proxy_fences.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Source check of the cross-proxy fences in the Kimi K3 attention-residual kernels (PTX ISA, memory consistency +model, proxies). + +The consumers read the V (and delta) slots with generic-proxy shared loads and release each slot with an mbarrier +arrive on ``bar_consumed``. The producer then refills the slot with ``cp.async.bulk``, an async-proxy write. Accesses +to one location through two proxies need a cross-proxy fence: every reading thread issues +``fence.proxy.async.shared::cta`` after its last read and before the release. Without it the order rests on how the +compiler schedules the instructions. No test of values can see that, so the kernels' sources are read. + + pytest test_attn_res_proxy_fences.py +""" + +import pathlib +import re + +import pytest + +pytestmark = pytest.mark.cpu_only + +_SRC = pathlib.Path(__file__).resolve().parents[5] / "cpp/tensorrt_llm/kernels/kimiK3AttnRes" + +# The kernels whose consumers release cp.async.bulk-filled slots, with the number of releases each file has. +RELEASES = {"attnResFwd.cu": 3, "attnResFwdPersistentFused.cu": 1} + +RELEASE = re.compile(r"\b(cute::arrive_barrier|mbarrier_arrive)\(\s*plan\.bar_consumed\b") +FENCE = 'asm volatile("fence.proxy.async.shared::cta;" ::: "memory");' +# What may sit between the fence and the release: the warp sync that gathers every lane's reads before lane 0 releases, +# and the lane-0 guard. Any other statement could hold a read after the fence. +BETWEEN = re.compile(r"^(__syncwarp\(\);|if \(lane == 0\)|\{)$") + + +def _code(line): + return line.split("//", 1)[0].strip() + + +def violations(lines): + """(line number, release) of every ``bar_consumed`` release whose closest preceding statement, past the warp sync + and the lane-0 guard, is not the cross-proxy fence.""" + out = [] + for i, line in enumerate(lines): + if not RELEASE.search(_code(line)): + continue + for j in range(i - 1, -1, -1): + code = _code(lines[j]) + if not code or BETWEEN.match(code): + continue + if code != FENCE: + out.append((i + 1, line.strip())) + break + else: + out.append((i + 1, line.strip())) + return out + + +@pytest.mark.parametrize("name", sorted(RELEASES)) +def test_slot_release_follows_cross_proxy_fence(name): + lines = (_SRC / name).read_text().splitlines() + releases = [i for i, line in enumerate(lines) if RELEASE.search(_code(line))] + assert len(releases) == RELEASES[name], ( + f"{name}: {len(releases)} bar_consumed releases, expected {RELEASES[name]} " + "(update RELEASES if the kernels changed)" + ) + assert any("cp.async.bulk.shared::cta.global" in line for line in lines), ( + f"{name}: no cp.async.bulk refill; the fence may no longer be needed" + ) + assert violations(lines) == [], ( + f"{name}: slot releases without fence.proxy.async.shared::cta right before them: {violations(lines)}" + ) + + +def test_check_flags_a_release_without_the_fence(): + """The check itself: the release pattern without the fence, as before the fix, is reported.""" + lines = [ + "// Every lane's reads of the chunk's slots before lane 0 releases them.", + "__syncwarp();", + "if (lane == 0)", + "{", + " cute::arrive_barrier(plan.bar_consumed[chunk_slot]);", + "}", + ] + assert violations(lines) == [(5, "cute::arrive_barrier(plan.bar_consumed[chunk_slot]);")] + assert violations([FENCE] + lines) == [] + assert violations([FENCE, "float x = buf[0];"] + lines) == [ + (7, "cute::arrive_barrier(plan.bar_consumed[chunk_slot]);") + ] From 1e0dceb74e5e18d92294c6a62965c7043d87d3dd Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 1 Oct 2026 21:28:59 -0700 Subject: [PATCH 16/67] DFlash: start the dummy slot at context length 0 every step CUDA-graph padding rows and warmup dummies share the drafter's dummy slot. Their accepted tokens grow its context length like any request's, but nothing reset it, while their page-table rows are the padding request's pages with every other entry mapped to page 0, another request's page. Once the dummy length passed the padding request's own pages, the padding rows' context K / V (k3_ctx_kv, or the torch path) landed on that page: a live request's drafter context overwritten on every padded step, costing acceptance. prepare() now writes 0 to the dummy slot every step, with the evicted slots. test_dflash_dummy_slot.py: padded steps keep the dummy slot at 0 and the real slots untouched; an evicted request and the dummy reset in one step. It fails without the reset. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/speculative/dflash.py | 4 + .../hw_agnostic/test_dflash_dummy_slot.py | 83 +++++++++++++++++++ 2 files changed, 87 insertions(+) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index d8097dffbd58..d760e15fbc78 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -551,6 +551,10 @@ def prepare(self): evicted[slot] = 0 worker._req_ctx_pos.pop(rid, None) worker._free_slots.append(slot) + # The dummy slot starts every step empty. Its rows (CUDA-graph padding, warmup dummies) add their accepted + # tokens to its length like any request, and their table rows are the padding request's pages, the rest + # mapped to page 0 (another request's). Left growing, a padding row's context K / V would land there. + evicted[worker._dummy_slot] = 0 worker._write_ctx_len(evicted) # A disagg generation worker receives prompt KV instead of diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py new file mode 100644 index 000000000000..37d837aa7c28 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -0,0 +1,83 @@ +# 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. +"""DFlash's dummy slot (CUDA-graph padding and warmup dummies) starts every step at context length 0. + +Padding rows add their accepted tokens to the dummy slot's length like any request, while their page-table rows are +the padding request's pages with the rest mapped to page 0, which another request owns. A dummy length that kept +growing would put the padding rows' context K / V on that page. ``DFlashSpecMetadata.prepare`` runs before every step +(eager or replayed graph); after it the dummy slot's length is 0 and every real slot's is untouched. +""" + +from functools import partial +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="the slot lengths live on the GPU" +) + +FLOOR = 1 << 40 # request ids from here on are CUDA-graph padding dummies + + +def _worker(lengths): + w = SimpleNamespace( + _ctx_buf_inited=True, + _req_to_slot={11: 0, 12: 1, 13: 2}, + _req_ctx_pos={}, + _free_slots=[3], + _graph_dummy_id_floor=FLOOR, + _dummy_slot=len(lengths) - 1, + _ctx_len=torch.tensor(lengths, dtype=torch.long, device="cuda"), + _ctx_len_host=list(lengths), + _batch_to_slot=torch.zeros(8, dtype=torch.long, device="cuda"), + ) + w._write_ctx_len = partial(DFlashWorker._write_ctx_len, w) + w._assign_slot = lambda rid, *a, **k: None + return w + + +def _prepare(worker, request_ids, num_generations): + meta = SimpleNamespace( + request_ids=request_ids, + num_generations=num_generations, + batch_indices_cuda=torch.zeros(8, dtype=torch.int, device="cuda"), + _dflash_worker=worker, + ) + DFlashSpecMetadata.prepare(meta) + torch.cuda.synchronize() + + +def test_padding_steps_keep_the_dummy_slot_empty(): + worker = _worker([40, 300, 7, 0, 0]) + padded = [11, 12, 13, FLOOR + 1, FLOOR + 1] # three requests in a graph of five + for step in range(4): + # What the step's acceptance does to the slots of its rows (k3_ctx_kv / the torch path): + accepted. + _prepare(worker, padded, num_generations=5) + assert worker._ctx_len.tolist() == [40 + 8 * step, 300 + 8 * step, 7 + 8 * step, 0, 0], step + assert worker._batch_to_slot[:5].tolist() == [0, 1, 2, 4, 4] + assert worker._ctx_len_host[4] == 0 + worker._ctx_len[[0, 1, 2]] += 8 + worker._ctx_len[4] += 2 * 8 # both padding rows land on the dummy slot + + +def test_evicted_request_and_dummy_reset_together(): + worker = _worker([40, 300, 7, 0, 96]) + _prepare(worker, [11, 13, FLOOR + 1], num_generations=3) # request 12 left; one padding row + assert worker._ctx_len.tolist() == [40, 0, 7, 0, 0] + assert 1 in worker._free_slots and 12 not in worker._req_to_slot From adb66182fadee7007a9d4bb9f133e783eaed2ca4 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 20:15:52 -0700 Subject: [PATCH 17/67] [https://nvbugs/6852473][fix] Skip zero-row launches in fp4_quantize, mxfp8_quantize and noaux_tc_op invokeFP4Quantization and invokeMxFP8Quantization launch one block per row (up to an occupancy cap), invokeNoAuxTc one block per token. With no rows the grid has zero blocks: the launch fails with cudaErrorInvalidValue, nothing checks it, and the error stays pending until the next launch that reads the last error raises it, usually an unrelated PyTorch kernel. The three launchers now return before the launch when there are no rows (there is nothing to write; the swizzled scale-factor layout pads zero rows to zero), and their cudaLaunchKernelEx calls are checked. noaux_tc_op still rejects an unsupported configuration at zero tokens. That pending error is what failed the three modeling_v2 catalog tests that run right after a zero-row call in their files; their waivers are removed. Tests: zero-row fp4_quantize (bf16 and fp8 input, both scale-factor layouts) and mxfp8_quantize (linear and swizzled), and zero-token noaux_tc_op (single and multi group), each followed by a normal call that must match the one before it. Signed-off-by: Vasanth Sabavat --- cpp/tensorrt_llm/kernels/noAuxTcKernels.cu | 10 +++++-- cpp/tensorrt_llm/kernels/quantization.cu | 22 +++++++++++---- tests/integration/test_lists/waives.txt | 3 -- .../thop/parallel/test_fp4_gemm_quantize.py | 28 +++++++++++++++++++ .../_torch/thop/parallel/test_fp8_quantize.py | 25 +++++++++++++++++ .../parallel_hw_agnostic/test_noaux_tc.py | 27 ++++++++++++++++++ 6 files changed, 105 insertions(+), 10 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu b/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu index 986e030c71c0..d3041739d959 100644 --- a/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu +++ b/cpp/tensorrt_llm/kernels/noAuxTcKernels.cu @@ -289,6 +289,12 @@ void invokeNoAuxTc(InputT* scores, BiasT* bias, OutputT* topk_values, IdxT* topk if (is_single_group || is_multi_group) { + // No tokens: there is nothing to write, and a grid of zero blocks is an invalid launch. + if (num_tokens == 0) + { + return; + } + cudaLaunchConfig_t config; auto* kernel_instance = &deepseek_v3_topk_kernel; int num_threads = NumDeepseekExperts; @@ -344,8 +350,8 @@ void invokeNoAuxTc(InputT* scores, BiasT* bias, OutputT* topk_values, IdxT* topk config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values, topk_indices, bias, num_tokens, n_group, - topk_group, topk, num_experts, num_experts / n_group, routed_scaling_factor); + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel_instance, scores, topk_values, topk_indices, bias, + num_tokens, n_group, topk_group, topk, num_experts, num_experts / n_group, routed_scaling_factor)); sync_check_cuda_error(stream); } else diff --git a/cpp/tensorrt_llm/kernels/quantization.cu b/cpp/tensorrt_llm/kernels/quantization.cu index 1daa4a397120..3d47be9e2e74 100644 --- a/cpp/tensorrt_llm/kernels/quantization.cu +++ b/cpp/tensorrt_llm/kernels/quantization.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. + * Copyright (c) 2019-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. @@ -136,6 +136,12 @@ template void invokeFP4Quantization(int b, int m, int n, T const* input, float const* SFScale, int64_t* output, int32_t* SFOuput, bool useUE8M0, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream) { + // No rows: there is nothing to write, and a grid of zero blocks is an invalid launch. + if (m == 0) + { + return; + } + #ifdef ENABLE_FP8 if constexpr (std::is_same_v) { @@ -181,8 +187,8 @@ void invokeFP4Quantization(int b, int m, int n, T const* input, float const* SFS attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, kernel_instance, b, m, n, n, input, SFScale, reinterpret_cast(output), - reinterpret_cast(SFOuput), layout); + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, kernel_instance, b, m, n, n, input, SFScale, + reinterpret_cast(output), reinterpret_cast(SFOuput), layout)); } } @@ -197,6 +203,12 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, static_assert(SF_OUTPUT_VEC_SIZE == 32 || SF_OUTPUT_VEC_SIZE == SF_VEC_SIZE, "MXFP8 output SF vector size must be 32 or match the quantization SF vector size."); + // No rows: there is nothing to write, and a grid of zero blocks is an invalid launch. + if (m == 0) + { + return; + } + // Grid, Block size. // Each thread converts 8 values. dim3 block(std::min(int(padded_n / CVT_ELTS_PER_THREAD), 512)); @@ -217,10 +229,10 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input, attrs[0].val.programmaticStreamSerializationAllowed = tensorrt_llm::common::getEnvEnablePDL(); config.numAttrs = 1; config.attrs = attrs; - cudaLaunchKernelEx(&config, + TLLM_CUDA_CHECK(cudaLaunchKernelEx(&config, quantize_with_block_size, b, m, n, padded_n, input, nullptr, reinterpret_cast(output), reinterpret_cast(SFOuput), - layout); + layout)); } //////////////////////////////////////////////////////////////////////////////////////////////////// diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index b60dc375273d..e0c192222196 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -284,9 +284,6 @@ unittest/_torch/modeling/test_modeling_nemotron_h_multimodal.py::test_nemotron_n unittest/_torch/modeling_v2 --ignore=unittest/_torch/modeling_v2/comm SKIP (https://nvbugs/6853741) unittest/_torch/modeling_v2/comm/test_modeling_v2_allgather_op_matrix.py::test_allgather_op_matrix SKIP (https://nvbugs/6843474) unittest/_torch/modeling_v2/comm/test_modeling_v2_reducescatter_op_matrix.py::test_reducescatter_op_matrix SKIP (https://nvbugs/6843474) -unittest/_torch/modeling_v2/moe/test_modeling_v2_noaux_tc_op.py::test_input_not_mutated_and_deterministic SKIP (https://nvbugs/6852473) -unittest/_torch/modeling_v2/quantization/test_modeling_v2_fp4_quantize.py::test_input_not_mutated_and_deterministic SKIP (https://nvbugs/6852473) -unittest/_torch/modeling_v2/quantization/test_modeling_v2_mxfp8_quantize.py::test_input_not_mutated SKIP (https://nvbugs/6852473) unittest/_torch/modules/mamba/test_flashinfer_gdn_decode.py SKIP (https://nvbugs/6842831) unittest/_torch/modules/tests_lora_modules/test_nemotron_h_lora_sanity.py::TestNemotronHLoRA::test_lora_pp2_sanity SKIP (https://nvbugs/6428124) unittest/_torch/moe/test_moe_backend.py::test_moe_backend[act=Relu2-e60_k4_h2048_i1408-seq=8-dtype=torch.bfloat16-backend=TRTLLM-quant=NVFP4-routing=Renormalize] SKIP (https://nvbugs/5989912) diff --git a/tests/unittest/_torch/thop/parallel/test_fp4_gemm_quantize.py b/tests/unittest/_torch/thop/parallel/test_fp4_gemm_quantize.py index d8e7e8ab8a26..7cfff778fc29 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp4_gemm_quantize.py +++ b/tests/unittest/_torch/thop/parallel/test_fp4_gemm_quantize.py @@ -272,6 +272,34 @@ def test_fp4_quantize_torch_fp8(self, m, k, dtype, use_ue8m0, # The gap is too large for ue8m0, so we just make sure that it runs self.assertTrue(torch.allclose(a_pt, aq_fp32, atol=1, rtol=0)) + @parameterized.expand(list([[torch.bfloat16, True], [torch.bfloat16, False], + [torch.float8_e4m3fn, True], + [torch.float8_e4m3fn, False]]), + name_func=unittest_name_func) + @skip_pre_blackwell_unittest + def test_fp4_quantize_zero_rows(self, dtype, is_sf_swizzled_layout): + """No rows: empty outputs, and no launch error left for the next call.""" + k = 512 + a = torch.randn([4, k], dtype=torch.float32) + a_global_sf = ((448 * 6) / a.abs().max().float()).cuda() + a = a.to(dtype).cuda() + ref_fp4, ref_sf = torch.ops.trtllm.fp4_quantize(a, a_global_sf, 16, + False, + is_sf_swizzled_layout) + + empty = torch.empty([0, k], dtype=dtype, device="cuda") + empty_fp4, empty_sf = torch.ops.trtllm.fp4_quantize( + empty, a_global_sf, 16, False, is_sf_swizzled_layout) + self.assertEqual(tuple(empty_fp4.shape), (0, k // 2)) + self.assertEqual(empty_sf.numel(), 0) + + # A launch error left pending by the empty call is raised by the next + # kernel launch, here torch.equal's. + a_fp4, a_sf = torch.ops.trtllm.fp4_quantize(a, a_global_sf, 16, False, + is_sf_swizzled_layout) + self.assertTrue(torch.equal(a_fp4, ref_fp4)) + self.assertTrue(torch.equal(a_sf, ref_sf)) + class TestProfiling(unittest.TestCase): diff --git a/tests/unittest/_torch/thop/parallel/test_fp8_quantize.py b/tests/unittest/_torch/thop/parallel/test_fp8_quantize.py index 6964284c941e..2c75c5c9467a 100644 --- a/tests/unittest/_torch/thop/parallel/test_fp8_quantize.py +++ b/tests/unittest/_torch/thop/parallel/test_fp8_quantize.py @@ -202,6 +202,31 @@ def test_mxfp8_quantize_alignment_torch_device(m, k, dtype, a.cpu().to(torch.float32), 8, 0, 0.999) +@pytest.mark.parametrize("is_sf_swizzled_layout, alignment", [(False, 512), + (True, 32)]) +@skip_pre_blackwell_unittest +def test_mxfp8_quantize_zero_rows(is_sf_swizzled_layout, alignment): + """No rows: empty outputs, and no launch error left for the next call.""" + k = 2880 + torch.random.manual_seed(0) + a = torch.randn([4, k], dtype=torch.bfloat16, device="cuda") + ref_fp8, ref_sf = torch.ops.trtllm.mxfp8_quantize(a, is_sf_swizzled_layout, + alignment) + + empty = torch.empty([0, k], dtype=torch.bfloat16, device="cuda") + empty_fp8, empty_sf = torch.ops.trtllm.mxfp8_quantize( + empty, is_sf_swizzled_layout, alignment) + assert empty_fp8.shape == (0, ref_fp8.shape[1]) + assert empty_sf.numel() == 0 + + # A launch error left pending by the empty call is raised by the next + # kernel launch, here torch.equal's. + a_fp8, a_sf = torch.ops.trtllm.mxfp8_quantize(a, is_sf_swizzled_layout, + alignment) + assert torch.equal(a_fp8.view(torch.uint8), ref_fp8.view(torch.uint8)) + assert torch.equal(a_sf, ref_sf) + + def _run_megamoe_prepare(hidden_states, token_selected_experts, token_final_scales): m, k = hidden_states.shape diff --git a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_noaux_tc.py b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_noaux_tc.py index a2f49cb3516d..ce131855c593 100644 --- a/tests/unittest/_torch/thop/parallel_hw_agnostic/test_noaux_tc.py +++ b/tests/unittest/_torch/thop/parallel_hw_agnostic/test_noaux_tc.py @@ -80,3 +80,30 @@ def test_noaux_tc_run(seq_len, num_experts, n_group, topk_group, top_k, dtype): torch.testing.assert_close( sorted_selected_values, ref_sorted_selected_values, rtol=0.01, atol=0.01 ) + + +@pytest.mark.parametrize( + "num_experts, n_group, topk_group, top_k", + [ + (72, 1, 1, 6), + (256, 8, 4, 8), + ], +) +def test_noaux_tc_zero_tokens(num_experts, n_group, topk_group, top_k): + """No tokens: empty outputs, and no launch error left for the next call.""" + torch.manual_seed(24) + scores = torch.randn((8, num_experts), dtype=torch.bfloat16, device="cuda") + bias = torch.randn((num_experts,), dtype=torch.float32, device="cuda") + args = (n_group, topk_group, top_k, 2.5) + ref_values, ref_indices = torch.ops.trtllm.noaux_tc_op(scores, bias, *args) + + empty = torch.empty((0, num_experts), dtype=torch.bfloat16, device="cuda") + values, indices = torch.ops.trtllm.noaux_tc_op(empty, bias, *args) + assert values.shape == (0, top_k) + assert indices.shape == (0, top_k) + + # A launch error left pending by the empty call is raised by the next kernel launch, here + # torch.equal's. + values, indices = torch.ops.trtllm.noaux_tc_op(scores, bias, *args) + assert torch.equal(values, ref_values) + assert torch.equal(indices, ref_indices) From 5e238440d749d3c822e670069d87f313cf1e6a38 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 20:45:26 -0700 Subject: [PATCH 18/67] [None][fix] kda_prefill: return a new output tensor on every call trtllm::kda_prefill returned a view of a buffer it reuses: the output was one of the intermediate tensors _get_buffers caches per input shape, so every call with the same shapes wrote the same storage, and a caller's output changed under it at the next such call. The op's fake declares a new tensor (v.new_empty(v.shape)), so nothing orders the readers of one call's output before the next call. Allocate the output on every call, with the varlen path's one chunk of slack past the last token like the cached buffers, and wrap it for K4 there; the cache no longer holds it. Test: a second call with the same shapes and other inputs must leave the first call's output unchanged (varlen, and one padded equal-length row). Signed-off-by: Vasanth Sabavat --- .../custom_ops/cute_dsl_kimi_k3_custom_ops.py | 24 ++++--- .../modules/kimi_kda/test_kda_prefill_op.py | 68 +++++++++++++++++++ 2 files changed, 83 insertions(+), 9 deletions(-) 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 171cd4ceb615..fee6b354cc69 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 @@ -406,7 +406,6 @@ def _with_slack(ctor, *shape, dtype): gk_last_exp = torch.empty(B, NT, H, K_dim, device=dev, dtype=torch.float32) A_qk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k) A_kk = _with_slack(torch.zeros, B, T_alloc, H, BT, dtype=dtype_k) - O_flat = _with_slack(torch.empty, B, T_alloc, H, V_dim, dtype=dtype_k) cu_eqlen = torch.arange(0, (B + 1) * T, T, dtype=torch.int32, device=dev) # co_eqlen is consumed only on the eqlen path (varlen K4 derives its # chunk offsets from cu_seqlens). Eqlen inputs are pre-padded to a @@ -423,7 +422,6 @@ def _with_slack(ctor, *shape, dtype): KS_flat = k_scaled.reshape(T_total, H, K_dim) QS_flat = q_scaled.reshape(T_total, H, K_dim) KG_flat = kg.reshape(T_total, H, K_dim) - O_token = O_flat.reshape(T_total, H, V_dim) gk_flat = gk_last_exp.reshape(-1, H, K_dim) def _wrap(t, etype): @@ -436,7 +434,6 @@ def _wrap(t, etype): ks_ct = _wrap(KS_flat, bf16) # raw k_scaled (beta absorbed in akk_inv) qs_ct = _wrap(QS_flat, bf16) kg_ct = _wrap(KG_flat, bf16) - o_ct = _wrap(O_token, bf16) gk_ct = _wrap(gk_flat, fp32) cu_eqlen_ct = from_dlpack(cu_eqlen, assumed_align=4).mark_layout_dynamic() cu_eqlen_ct.element_type = cutlass.Int32 @@ -458,7 +455,6 @@ def _wrap(t, etype): ks_ct=ks_ct, qs_ct=qs_ct, kg_ct=kg_ct, - o_ct=o_ct, gk_ct=gk_ct, cu_eqlen_ct=cu_eqlen_ct, co_eqlen_ct=co_eqlen_ct, @@ -480,7 +476,6 @@ def _wrap(t, etype): gk_last_exp, A_qk, A_kk, - O_flat, cu_eqlen, co_eqlen, cute_wrappers, @@ -491,8 +486,20 @@ def _wrap(t, etype): return _buf_cache[key] +def _new_output(dev, dtype, B, T, H, V_dim, BT, varlen=False): + """The prefill's output ``o`` [B, T, H, V] and K4's wrapper of it, new for every call: ``o`` belongs to the + caller, so no later call may write it. On the varlen path the buffer has one chunk of slack past T, like the + scratch.""" + T_alloc = T + BT if varlen else T + o = torch.empty(B, T_alloc, H, V_dim, device=dev, dtype=dtype)[:, :T] + o_ct = from_dlpack(o.reshape(B * T, H, V_dim), assumed_align=16).mark_layout_dynamic() + o_ct.element_type = cutlass.BFloat16 + return o, o_ct + + def _launch_k4_persistent( cute_wrappers, + o_ct, v_beta, state_pool, state_indices, @@ -528,7 +535,6 @@ def _launch_k4_persistent( q_ct = cute_wrappers["qs_ct"] aqc_ct = cute_wrappers["aqc_ct"] kg_ct = cute_wrappers["kg_ct"] - o_ct = cute_wrappers["o_ct"] gk_ct = cute_wrappers["gk_ct"] # v is a per-call activation — wrap fresh every call (never cache; see @@ -1024,11 +1030,11 @@ def _chunk_kda_fwd( gk_last_exp, A_qk, A_kk, - O_flat, cu_eqlen, co_eqlen, cute_wrappers, ) = _get_buffers(device, k.dtype, B, T, H, K, V_dim, NT, N_seqs, BT, varlen=is_varlen) + o, o_ct = _new_output(device, k.dtype, B, T, H, V_dim, BT, varlen=is_varlen) # Beta is fused entirely in akk_inv kernel epilogue (post-inv column-scale). # No host v*beta and no K1 k_scaled*beta any more. @@ -1072,6 +1078,7 @@ def _chunk_kda_fwd( _launch_k4_persistent( cute_wrappers, + o_ct, v, state_pool, state_indices, @@ -1082,10 +1089,9 @@ def _chunk_kda_fwd( V_dim=V_dim, ) - o = O_flat if needs_eqlen_pad: # Caller called with original T = real_T; their downstream code expects - # outputs at that shape. Slice the padded scratch tail back off. + # outputs at that shape. Slice the padded tail back off. o = o[:, :real_T] return o 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 bf141d7a6306..94faa0d97f9c 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 @@ -598,6 +598,74 @@ def test_kda_prefill_op_rejects_misaligned_state_indices() -> None: ) +@torch.no_grad() +@pytest.mark.parametrize( + "sequence_lengths", + [[100, 200], None], + ids=["varlen", "equal-length-padded"], +) +def test_kda_prefill_op_output_survives_the_next_call(sequence_lengths) -> None: + """The op's output belongs to the caller, as its fake (a new tensor) declares. + + A second call with the same shapes and other inputs must leave the first + call's output as it was. ``None`` runs one 300-token row without + cu_seqlens, which the op pads to a chunk multiple and slices back. + """ + from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops # noqa: F401 + from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices + + heads = 2 + tokens = sum(sequence_lengths) if sequence_lengths else 300 + rows = len(sequence_lengths) if sequence_lengths else 1 + state_pool = torch.zeros(rows, heads, HEAD_DIM, HEAD_DIM, dtype=torch.float32, device="cuda") + state_indices = torch.arange(rows, dtype=torch.int32, device="cuda") + cu_seqlens = None + chunk_indices = None + if sequence_lengths: + cu_seqlens = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + chunk_indices = prepare_chunk_indices(cu_seqlens, 64) + + def call(seed: int) -> torch.Tensor: + gen = torch.Generator(device="cuda").manual_seed(seed) + shape = (1, tokens, heads, HEAD_DIM) + + def unit() -> torch.Tensor: + t = torch.randn(*shape, generator=gen, device="cuda") + return (t / t.norm(dim=-1, keepdim=True)).to(torch.bfloat16) + + state_pool.zero_() + return torch.ops.trtllm.kda_prefill( + q=unit(), + k=unit(), + v=(torch.randn(*shape, generator=gen, device="cuda") * 0.5).to(torch.bfloat16), + g=torch.randn(*shape, generator=gen, device="cuda").to(torch.bfloat16), + beta=torch.randn(1, tokens, heads, generator=gen, device="cuda"), + state_pool=state_pool, + state_indices=state_indices, + scale=HEAD_DIM**-0.5, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + safe_gate=True, + lower_bound=-5.0, + use_gate_in_kernel=True, + use_beta_sigmoid_in_kernel=True, + A_log=torch.zeros(heads, device="cuda"), + dt_bias=torch.zeros(heads * HEAD_DIM, device="cuda"), + ) + + first = call(seed=0) + kept = first.clone() + second = call(seed=1) + torch.cuda.synchronize() + assert first.shape == second.shape == (1, tokens, heads, HEAD_DIM) + assert not torch.equal(second, kept), "the two calls must compute different outputs" + assert torch.equal(first, kept), "a later call wrote an earlier call's output" + + @torch.no_grad() def test_kda_prefill_op_partial_final_chunk_large_batch(): """Regression: varlen batches whose FINAL chunk is partial. From 936f38fbe75d546456355909d07261135d8fc454 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 20:53:38 -0700 Subject: [PATCH 19/67] [None][fix] Kimi Linear / Kimi K3: make fla-core a runtime dependency The KDA layers (tensorrt_llm/_torch/modules/kimi_kda) import flash-linear-attention's ShortConvolution when they load and run its chunked KDA for prefill batches below four 64-token chunks, but flash-linear-attention is listed only in requirements-dev.txt. CI installs that file, so every test passes, while a plain pip install of the wheel, and the release container, which installs the wheel with its runtime requirements only, cannot load either model: No module named 'fla'. Require fla-core, the distribution that provides fla.modules and fla.ops, at the version requirements-dev.txt already installs through flash-linear-attention 0.5.2. Its only base dependency, einops, is a runtime requirement already. A CPU test checks that every package the KDA layers import when they load is a runtime requirement. Signed-off-by: Vasanth Sabavat --- requirements.txt | 1 + .../test_kimi_kda_runtime_dependencies.py | 64 +++++++++++++++++++ 2 files changed, 65 insertions(+) create mode 100644 tests/unittest/others/test_kimi_kda_runtime_dependencies.py diff --git a/requirements.txt b/requirements.txt index 63c1972c7c35..81a35db18be0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -68,6 +68,7 @@ ordered-set peft>=0.18.1,<0.19.0 patchelf einops +fla-core==0.5.2 # the Kimi Linear / Kimi K3 KDA layers import fla.modules and fla.ops when they load flashinfer-python==0.6.18 xgrammar==0.1.32 llguidance==0.7.29 diff --git a/tests/unittest/others/test_kimi_kda_runtime_dependencies.py b/tests/unittest/others/test_kimi_kda_runtime_dependencies.py new file mode 100644 index 000000000000..73add6c0ce5a --- /dev/null +++ b/tests/unittest/others/test_kimi_kda_runtime_dependencies.py @@ -0,0 +1,64 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The packages the Kimi KDA layers import when they load must be runtime requirements. + +CI installs requirements-dev.txt, so a package listed only there passes every test, yet it is +missing from a plain `pip install tensorrt_llm` and from the release container, where loading +Kimi Linear or Kimi K3 then fails. +""" + +import ast +from importlib import metadata +from pathlib import Path + +import pytest +from packaging.requirements import InvalidRequirement, Requirement +from packaging.utils import canonicalize_name + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parents[3] +KDA_PACKAGE = REPO_ROOT / "tensorrt_llm" / "_torch" / "modules" / "kimi_kda" + + +def _runtime_requirements() -> set: + names = set() + for line in (REPO_ROOT / "requirements.txt").read_text().splitlines(): + # Trailing comments are not PEP 508; option lines (-c, --extra-index-url) name no package. + line = line.split("#", 1)[0].strip() + if not line or line.startswith("-"): + continue + try: + names.add(canonicalize_name(Requirement(line).name)) + except InvalidRequirement: + continue + return names + + +def _load_time_imports(path: Path) -> set: + """Top-level packages of a module's module-level absolute imports.""" + names = set() + for node in ast.parse(path.read_text()).body: + if isinstance(node, ast.Import): + names.update(alias.name.split(".")[0] for alias in node.names) + elif isinstance(node, ast.ImportFrom) and node.level == 0 and node.module: + names.add(node.module.split(".")[0]) + return names + + +def test_kda_load_time_imports_are_runtime_requirements() -> None: + requirements = _runtime_requirements() + providers = metadata.packages_distributions() + # The KDA layers import `fla` at load; with no distribution to map it to, the check is vacuous. + assert providers.get("fla"), ( + "flash-linear-attention is not installed: `fla` cannot be mapped to a distribution" + ) + missing = [] + for path in sorted(KDA_PACKAGE.glob("*.py")): + for name in sorted(_load_time_imports(path) - {"tensorrt_llm"}): + distributions = {canonicalize_name(d) for d in providers.get(name, ())} + if distributions and not distributions & requirements: + missing.append( + f"{path.name} imports `{name}`, provided by {', '.join(sorted(distributions))}" + ) + assert not missing, "not runtime requirements (requirements.txt):\n" + "\n".join(missing) From 0102a35eeaf53279a1a2cb8d283fd3ad1e38bc29 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 23:31:37 -0700 Subject: [PATCH 20/67] [None][fix] MNNVL all-reduce: keep a replaced workspace only if a capture used it Growth retired every replaced workspace for the life of the process, including ones no CUDA graph had captured. Growth rounds to 8 MiB, so a run of rising eager sizes could keep many dead workspaces. The getter now marks a workspace it hands out during capture, and a growth keeps the previous workspace only when it is marked. An unmarked one is released at the replacement, after every rank's device synchronize in the protocol reset. The growth test checks that a second growth, past a workspace no capture used, leaves only the captured one in the retired list and that the graph still replays. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/distributed/ops.py | 28 ++++++++++++++----- .../_torch/multi_gpu/test_mnnvl_allreduce.py | 21 ++++++++++++-- 2 files changed, 40 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/distributed/ops.py b/tensorrt_llm/_torch/distributed/ops.py index 31d06468876d..e6f764938b95 100644 --- a/tensorrt_llm/_torch/distributed/ops.py +++ b/tensorrt_llm/_torch/distributed/ops.py @@ -99,6 +99,9 @@ class _MnnvlWorkspace(TypedDict): # An MPI communicator under MPI, the TP ProcessGroup under a non-MPI orchestrator (Ray). # None between checkpoint_prepare() and a successful checkpoint_restore(). comm: Optional[Union[_MpiCommProtocol, "torch.distributed.ProcessGroup"]] + # Set once the workspace is handed out during CUDA graph capture: the captured graphs launch on its buffers and + # flags, so a larger workspace that replaces it must keep it alive. + captured: bool def get_allreduce_workspace(mapping: Mapping) -> torch.LongTensor: @@ -256,12 +259,20 @@ def get_or_scale_allreduce_mnnvl_workspace( if not workspace["handle"].is_mapped(): raise RuntimeError("MNNVL workspace handles are not attached") if workspace["buffer_size_bytes"] >= (buffer_size_bytes or 0): - return workspace + return _mark_if_capturing(workspace) workspace_lock = MNNVLAllReduce._get_allreduce_mnnvl_workspace_lock(mapping) with workspace_lock: - return _get_or_scale_allreduce_mnnvl_workspace(mapping, dtype, - buffer_size_bytes) + return _mark_if_capturing( + _get_or_scale_allreduce_mnnvl_workspace(mapping, dtype, + buffer_size_bytes)) + + +def _mark_if_capturing(workspace: _MnnvlWorkspace) -> _MnnvlWorkspace: + """Record that a CUDA graph being captured launches on this workspace.""" + if not workspace["captured"] and torch.cuda.is_current_stream_capturing(): + workspace["captured"] = True + return workspace def _get_or_scale_allreduce_mnnvl_workspace( @@ -351,6 +362,7 @@ def _get_or_scale_allreduce_mnnvl_workspace( "buffer_flags": buffer_flags, "buffer_size_bytes": buffer_size_bytes, "comm": comm, + "captured": False, } except Exception as error: candidate_error = error @@ -368,8 +380,10 @@ def _get_or_scale_allreduce_mnnvl_workspace( # Hand ownership of the communicator to the workspace. pending_comms.pop(mapping, None) previous_workspace = allreduce_mnnvl_workspaces.get(mapping) - if previous_workspace is not None: - # CUDA graphs captured before this growth keep launching on the previous buffers and flags. + if previous_workspace is not None and previous_workspace["captured"]: + # CUDA graphs captured on the previous workspace keep launching on its buffers and flags. One that no + # capture used is released by the replacement below: every rank had synchronized its device before the + # protocol reset above returned, so no kernel still uses it. MNNVLAllReduce.allreduce_mnnvl_retired_workspaces.setdefault( mapping, []).append(previous_workspace) allreduce_mnnvl_workspaces[mapping] = candidate_workspace @@ -763,8 +777,8 @@ class MNNVLAllReduce(nn.Module): allreduce_mnnvl_workspaces: typing.ClassVar[dict[Mapping, _MnnvlWorkspace]] = {} - # Workspaces a larger one replaced. CUDA graphs captured before the growth still launch on their buffers and - # flags, so they stay alive with the process. The checkpoint hooks cover only the current workspace. + # Workspaces a larger one replaced after a CUDA graph capture used them. Those graphs still launch on their buffers + # and flags, so they stay alive with the process. The checkpoint hooks cover only the current workspace. allreduce_mnnvl_retired_workspaces: typing.ClassVar[dict[ Mapping, list[_MnnvlWorkspace]]] = {} diff --git a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py index ce053dba16e4..71cba1c68c3c 100644 --- a/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py +++ b/tests/unittest/_torch/multi_gpu/test_mnnvl_allreduce.py @@ -337,8 +337,8 @@ def mnnvl_checkpoint_worker_env(_: int): def mnnvl_growth_graph_forward(tensor_parallel_size: int, tensor_parallel_rank: int) -> bool: - """A graph captured on the MNNVL workspace replays correctly after an eager call grew the workspace, and growth - during capture is refused.""" + """A graph captured on the MNNVL workspace replays correctly after an eager call grew the workspace, a replaced + workspace that no capture used is not kept, and growth during capture is refused.""" env_names = ("TLLM_TEST_MNNVL", "TRTLLM_FORCE_MNNVL_AR") previous_env = { name: (name in os.environ, os.environ.get(name)) @@ -402,6 +402,23 @@ def mnnvl_growth_graph_forward(tensor_parallel_size: int, torch.full_like( output, expected + tensor_parallel_size * (1 + step))) + # No capture used `after`: growing past it keeps only `before`, which the graph still launches on. + wider = torch.ones((6144, 7168), + dtype=torch.bfloat16, + device="cuda") + torch.testing.assert_close( + allreduce(wider), torch.full_like(wider, tensor_parallel_size)) + assert MNNVLAllReduce.allreduce_mnnvl_workspaces[mapping][ + "buffer_size_bytes"] > after["buffer_size_bytes"] + assert MNNVLAllReduce.allreduce_mnnvl_retired_workspaces[ + mapping] == [before] + del wider, after + input_.fill_(tensor_parallel_rank + 1) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(output, + torch.full_like(output, expected)) + huge = torch.ones((8192, 7168), dtype=torch.bfloat16, device="cuda") with pytest.raises(RuntimeError, match="during CUDA graph capture"): with torch.cuda.graph(torch.cuda.CUDAGraph()): From ba20fa17fbebbdca0a400a4cb2dbd1b006394f06 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 2 Oct 2026 23:55:45 -0700 Subject: [PATCH 21/67] [None][fix] DSpark GQA drafter: keep the checkpoint's trained mask-token embedding GQADSparkForCausalLM takes the target's embedding in load_weights_from_target_model, so the DFlash worker built every masked block slot from the target's row for the mask token, which the target never trained. A drafter checkpoint that ships its own embedding has trained that row. load_weights now keeps it as mask_token_embedding, and the worker's noise block (dflash_noise_block_embedding, the same computation moved into a function) uses it for every masked slot when the drafter has one for the mask id in use; slot 0 stays the bonus token's shared lookup. Every rank still embeds the mask id through embed_tokens.forward, as before. Drafts change; under greedy verification the target's tokens do not. Plain DFlash and MLA DSpark drafters set no mask_token_embedding and are unchanged. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/models/modeling_dspark.py | 15 +++++++ tensorrt_llm/_torch/speculative/dflash.py | 43 ++++++++++++++---- .../test_kimi_k3_dspark_semantics.py | 44 +++++++++++++++++++ 3 files changed, 94 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 10aa541a6ff4..1855c56b4312 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2673,8 +2673,23 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """Take the DSpark head weights, then hand the rest to DFlash.""" weights, _ = self._take_dspark_head_weights(weights) + self._keep_trained_mask_embedding(weights) return super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + def _keep_trained_mask_embedding(self, weights: Dict) -> None: + """Keep the checkpoint's own embedding row for the mask token. + + This drafter takes the target's embedding (``load_weights_from_target_model``). A drafter checkpoint that ships + an embedding has trained the mask token's row, which the target's embedding does not have: the block decode + reads ``mask_token_embedding`` for every masked slot instead of the shared lookup. + """ + names = [k for k in ("embed_tokens.weight", "model.embed_tokens.weight") if k in weights] + if not names: + return + weight = weights[names[0]] + row = weight[self.mask_token_id : self.mask_token_id + 1] + self.mask_token_embedding = torch.as_tensor(row).reshape(-1).to("cuda") + class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): """DSpark drafter on an MLA-shaped backbone, from a standalone checkpoint. diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index d8097dffbd58..02be1e4f6aba 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -432,6 +432,32 @@ def dflash_draft_slot_ids( return (request_bases.unsqueeze(1) + first_slot + offsets.unsqueeze(0)).flatten() +def dflash_noise_block_embedding( + embed_tokens: nn.Module, + bonus: torch.Tensor, + mask_token_id: int, + block_size: int, + trained_mask_embedding: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """The drafter's input block per request, [len(bonus), block_size, hidden]: slot 0 the bonus token's embedding, + every other slot the mask token's. + + The mask row is the drafter's own trained one when it has one (``trained_mask_embedding``); otherwise the shared + lookup's. Both rows go through ``embed_tokens.forward`` (NOT ``.weight[...]``) so TP-sharded vocabs mask out + ranks that don't own the token id and all-reduce, the mask row included, so every rank makes the same call. + """ + num_gens = bonus.shape[0] + mask_tok = torch.full((1,), int(mask_token_id), dtype=torch.long, device=bonus.device) + combined_embed = embed_tokens(torch.cat([bonus, mask_tok], dim=0)) + embed_bonus = combined_embed[:num_gens] + embed_mask = combined_embed[num_gens] + if trained_mask_embedding is not None: + embed_mask = trained_mask_embedding.to(embed_mask.dtype) + noise_embed_2d = embed_mask.expand(num_gens, block_size, -1).clone() + noise_embed_2d[:, 0, :] = embed_bonus + return noise_embed_2d + + def dflash_position_ceiling(max_ctx: int, block_size: int, max_draft_len: int) -> int: """Positions a drafter that indexes absolute positions must be able to encode. @@ -1995,14 +2021,15 @@ def prepare_1st_drafter_inputs( query_position_ids = ctx_len_now.unsqueeze(1) + j_block.unsqueeze(0) ctx_position_ids = ctx_len_gen.unsqueeze(1) + offsets_kp1.unsqueeze(0) - # Go through embed_tokens.forward (NOT .weight[...]) so TP-sharded - # vocabs mask out ranks that don't own the token id and all-reduce. - mask_tok = torch.full((1,), int(mask_token_id), dtype=torch.long, device="cuda") - combined_embed = embed_tokens(torch.cat([bonus, mask_tok], dim=0)) - embed_bonus = combined_embed[:num_gens] - embed_mask = combined_embed[num_gens] - noise_embed_2d = embed_mask.expand(num_gens, query_tokens_per_req, -1).clone() - noise_embed_2d[:, 0, :] = embed_bonus + # The drafter's own trained mask row, if it kept one for the mask id in use. + trained_mask_embedding = ( + getattr(draft_model, "mask_token_embedding", None) + if mask_token_id == getattr(draft_model, "mask_token_id", None) + else None + ) + noise_embed_2d = dflash_noise_block_embedding( + embed_tokens, bonus, mask_token_id, query_tokens_per_req, trained_mask_embedding + ) # Accumulate new accepted features into context buffers if has_target_features: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index 70a6688c4fa2..0939bfdca109 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -454,6 +454,50 @@ def test_dspark_drafter_loads_head_weights_and_parses_config(): assert drafter.confidence_proj_bias is not None +@needs_gpu +def test_gqa_dspark_keeps_its_trained_mask_row_over_the_target_embedding(): + """The GQA drafter takes the target's embedding (load_weights_from_target_model), but a checkpoint that ships its + own embedding trained the mask token's row, which the target's embedding does not have. The drafter keeps that row + as ``mask_token_embedding``, and the worker's noise block puts it in every masked slot; slot 0 stays the bonus + token's embedding from the shared lookup. + """ + h = TINY["hidden_size"] + weights = _tiny_weights() + g = torch.Generator().manual_seed(23) + weights["embed_tokens.weight"] = (torch.randn(VOCAB, h, generator=g) * 0.05).to(torch.bfloat16) + drafter = _build_drafter(True, weights) + target_embed = torch.nn.Embedding(VOCAB, h).to("cuda", torch.bfloat16) + target = SimpleNamespace( + model=SimpleNamespace(embed_tokens=target_embed), + lm_head=torch.nn.Linear(h, VOCAB, bias=False), + ) + drafter.load_weights_from_target_model(target) + + mask_id = drafter.mask_token_id + trained = weights["embed_tokens.weight"][mask_id] + assert not torch.equal(target_embed.weight[mask_id].detach().cpu(), trained) + torch.testing.assert_close(drafter.mask_token_embedding.cpu(), trained, rtol=0, atol=0) + + from tensorrt_llm._torch.speculative.dflash import dflash_noise_block_embedding + + bonus = torch.tensor([3, 5], dtype=torch.long, device="cuda") + block = 4 + with torch.no_grad(): + noise = dflash_noise_block_embedding( + drafter.draft_model_full.model.embed_tokens, + bonus, + mask_id, + block, + drafter.mask_token_embedding, + ) + assert tuple(noise.shape) == (2, block, h) + torch.testing.assert_close( + noise[:, 0].cpu(), target_embed.weight[bonus].detach().cpu(), rtol=0, atol=0 + ) + for j in range(1, block): + torch.testing.assert_close(noise[:, j].cpu(), trained.expand(2, -1), rtol=0, atol=0) + + @needs_gpu def test_published_drafter_spelling_activates_the_heads(): """Both public K3 DSpark checkpoints load with their heads live. From e30477fb61ec921830fe84c81702979a66473e27 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Sat, 3 Oct 2026 00:44:20 -0700 Subject: [PATCH 22/67] [None][fix] KV cache manager V2: take the per-token KDA states out of the GPU quota at the SSM live floor With per-token KDA states and block reuse off, the SSM pool is held at its live floor and the typical step's attention pages fill the GPU quota. kda_state_tok, num_spec fp32 states per SSM slot, is allocated beside the pools afterwards. So the pools took the whole quota and the states came on top of it: initialization could run out of device memory although _minimum_live_gpu_quota() passed. At the floor the SSM pool keeps ceil(floor / max_util_for_resume) slots: the floor's min-slots constraint, as the runtime scales it. _build_cache_config now takes the token states of that many slots out of the GPU tier quota before the pools are sized, and logs the split. The runtime rounds each pool up to whole allocation grains, which can add SSM slots when a slot is smaller than a grain; their states are not reserved. The quota for max_tokens and the minimum quota count the token states of all those slots (_state_quota), not only the floor's. Otherwise the attention pages could get less than max_tokens needs: on Kimi K3 (DSpark with 7 drafts, TP16, max_batch_size 8: a floor of 9 slots, 10 kept) the estimation quota for max_tokens 16896 would leave the KV pool about 13.5K tokens. With it the quota is 2.56 GiB and the pool about 20.6K tokens (arithmetic from the logged 2.38 GiB quota at 27648 bytes per token). test_v2_kda_token_states_keep_ssm_pool_at_live_floor now also checks that the pools and the token states of every SSM slot fit in the quota, with attention taking the rest of it to within one grain. test_v2_kda_token_states_reserve_quota_before_sizing checks the reserved bytes, and that the minimum quota and the quota for max_tokens exceed those of the same manager without token states by the states of every kept slot. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 84 +++++++++++++++---- .../kv_cache/test_mamba_cache_manager.py | 74 +++++++++++++--- 2 files changed, 132 insertions(+), 26 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 6b60514f8dfa..1f9f6b0d04c7 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -52,7 +52,9 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, BatchDesc, BufferConfig, - DataRole, KVCacheDesc) + DataRole, + GpuCacheTierConfig, + KVCacheDesc) from tensorrt_llm.runtime.kv_cache_manager_v2 import \ KVCacheManagerConfig as KVCacheManagerConfigPy from tensorrt_llm.runtime.kv_cache_manager_v2 import (LayerId, PageIndexMode, @@ -3530,6 +3532,8 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: device, ) if self._kda_token_states: + # Beside the V2 pools. With the SSM pool at its live floor, + # _build_cache_config took these bytes out of the GPU quota. self.kda_state_tok = torch.zeros( (self.local_num_mamba_layers, cache_size, num_spec, *self.ssm_state_shape), @@ -3832,15 +3836,18 @@ def _get_pool_roles(self, def _max_resident_sequences(self) -> int: return self.max_batch_size * self.mapping.pp_size + def _kda_token_state_bytes_per_slot(self) -> int: + """Bytes of ``kda_state_tok`` per SSM slot: an fp32 SSM state per + draft of every local Mamba layer, or 0 without per-token states.""" + if not getattr(self, "_kda_token_states", False): + return 0 + return (self.local_num_mamba_layers * self._kda_replay_num_spec * + math.prod(self.ssm_state_shape) * torch.float32.itemsize) + def _mamba_state_bytes_per_slot(self) -> int: - base_bytes = self.local_num_mamba_layers * (self.ssm_bytes + - self.conv_bytes) - if getattr(self, "_kda_token_states", False): - # kda_state_tok: an fp32 SSM state per draft of every slot. - base_bytes += (self.local_num_mamba_layers * - self._kda_replay_num_spec * - math.prod(self.ssm_state_shape) * - torch.float32.itemsize) + base_bytes = (self.local_num_mamba_layers * + (self.ssm_bytes + self.conv_bytes) + + self._kda_token_state_bytes_per_slot()) local_ple_layers = sum(layer_id in self.pp_layers for layer_id in self._ple_layer_ids) if local_ple_layers == 0: @@ -3936,6 +3943,15 @@ def _get_typical_request_capacity( "pool ratio.") return fallback_capacity + def _state_quota(self, state_slots: int) -> int: + """Quota for ``state_slots`` recurrent-state slots, plus the token + states of the further slots the live floor keeps + (``_live_floor_token_state_slots``).""" + extra_slots = self._live_floor_token_state_slots( + self.kv_cache_config.max_util_for_resume) - state_slots + return (state_slots * self._mamba_state_bytes_per_slot() + + max(0, extra_slots) * self._kda_token_state_bytes_per_slot()) + def _get_quota_from_max_tokens(self, max_tokens: int) -> int: attention_quota = super()._get_quota_from_max_tokens(max_tokens) num_request_lineages = self._max_resident_sequences() @@ -3943,7 +3959,7 @@ def _get_quota_from_max_tokens(self, max_tokens: int) -> int: max_tokens, self.kv_cache_config) state_slots = (num_request_lineages + self._num_reserved_dummy_slots + snapshot_slots) - state_quota = state_slots * self._mamba_state_bytes_per_slot() + state_quota = self._state_quota(state_slots) # Once the plan contains any non-live SSM capacity, reserve one partial # attention page per request lineage. This remains conservative when # the plan contains fewer than one non-live slot per lineage. @@ -3979,7 +3995,7 @@ def _minimum_live_gpu_quota(self) -> int: self.tokens_per_block) num_state_slots = (self._max_resident_sequences() + self._num_reserved_dummy_slots) - state_quota = num_state_slots * self._mamba_state_bytes_per_slot() + state_quota = self._state_quota(num_state_slots) return max( self._get_quota_from_max_tokens(0), state_quota + attention_block_quota, @@ -3991,16 +4007,36 @@ def _ssm_pool_at_live_floor(self, kv_cache_config: KvCacheConfig) -> bool: of the typical step's ratio. With per-token KDA states, ``kda_state_tok`` takes num_spec states per - SSM slot outside the cache quota (``_allocate_pool_replay_buffers``), + SSM slot beside the V2 pools (``_allocate_pool_replay_buffers``), addressed by the request's SSM slot. Without block reuse no SSM slot holds anything but a live request or a reserved dummy, so slots past - the floor are never used and would only add that memory. + the floor are never used and would only add that memory. At the floor + the pool's slot count is known before sizing, so + ``_build_cache_config`` takes those states' bytes out of the GPU quota. """ return (getattr(self, "_kda_token_states", False) and self.local_num_mamba_layers > 0 and not kv_cache_config.enable_block_reuse and self._attention_cache_bytes_per_token() > 0) + def _live_floor_token_state_slots(self, max_util_for_resume: float) -> int: + """SSM slots whose ``kda_state_tok`` entries come out of the GPU quota. + + With the SSM pool at its live floor and no ``pool_ratio``, that is + every slot the pool keeps: the floor's min-slots constraint as the + runtime scales it by 1 / max_util_for_resume. (Rounding each pool up + to whole allocation grains can add slots when they are smaller than a + grain.) ``_build_cache_config`` takes their token states out of the + GPU quota before the pools are sized. Otherwise 0. + """ + kv_cache_config = self.kv_cache_config + if (kv_cache_config.pool_ratio is not None + or not self._ssm_pool_at_live_floor(kv_cache_config)): + return 0 + floor_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + return math.ceil(floor_slots / max_util_for_resume) + def _quota_filling_request_capacity(self, gpu_quota: int) -> int: """A typical request capacity whose resident requests' attention pages alone take ``gpu_quota``. The typical step's ratio then gives the SSM @@ -4022,6 +4058,25 @@ def _build_cache_config( "The V2 Mamba GPU cache quota is too small for live recurrent " f"states and attention pages: got {gpu_quota} bytes, need at " f"least {minimum_live_quota} bytes.") + ssm_pool_at_live_floor = (config.initial_pool_ratio is None and + self._ssm_pool_at_live_floor(kv_cache_config)) + if ssm_pool_at_live_floor: + # kda_state_tok is allocated beside the pools, one entry per SSM + # slot: take those states' bytes out of the GPU quota before the + # pools are sized. + token_state_slots = self._live_floor_token_state_slots( + config.max_util_for_resume) + token_state_quota = (token_state_slots * + self._kda_token_state_bytes_per_slot()) + gpu_quota -= token_state_quota + cache_tiers = [ + GpuCacheTierConfig(quota=gpu_quota), *cache_tiers[1:] + ] + logger.info( + "KV cache manager v2: the per-token KDA states of " + f"{token_state_slots} SSM slots take " + f"{token_state_quota / (1 << 30)}GiB of the device quota; the " + f"pools get {gpu_quota / (1 << 30)}GiB") # _build_base_config already constructed every attention layer, # including dtype-specific scale and subclass-provided side buffers. # Preserve those configs and replace only the local Mamba layers. @@ -4056,7 +4111,7 @@ def _build_cache_config( if config.initial_pool_ratio is None: typical_capacity = self._get_typical_request_capacity( kv_cache_config) - if self._ssm_pool_at_live_floor(kv_cache_config): + if ssm_pool_at_live_floor: typical_capacity = max( typical_capacity, self._quota_filling_request_capacity(gpu_quota)) @@ -4086,6 +4141,7 @@ def _build_cache_config( ] return replace( config, + cache_tiers=cache_tiers, layers=layers, typical_step=typical_step, constraints=constraints, diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index df05a69a015e..87cadcbe3bad 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -78,6 +78,7 @@ ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + GPU_LEVEL, AttentionLayerConfig, BatchDesc, BufferConfig, @@ -3747,16 +3748,9 @@ def test_v2_kda_replay_validates_configuration( ) -@pytest.mark.parametrize( - ("kda_token_states", "enable_block_reuse", "at_floor"), - [(True, False, True), (False, False, False), (True, True, False)], -) -def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( - kda_token_states, enable_block_reuse, at_floor -): - """The per-token KDA states take memory per SSM slot outside the cache quota. Without block reuse the SSM pool - keeps only its live floor and attention gets the rest of the quota; otherwise (replay caches only, or block reuse) - the typical step's ratio sizes it.""" +def _v2_kda_token_state_manager(kda_token_states, enable_block_reuse): + """A V2 hybrid manager stub for _build_cache_config: one KDA layer with the replay caches and one attention + layer.""" mgr = object.__new__(MambaHybridCacheManagerV2) mgr._generation_kv_capacity_headroom = 1 mgr._has_cp_helix = False @@ -3794,6 +3788,20 @@ def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( enable_block_reuse=enable_block_reuse, enable_partial_reuse=False, ) + return mgr + + +@pytest.mark.parametrize( + ("kda_token_states", "enable_block_reuse", "at_floor"), + [(True, False, True), (False, False, False), (True, True, False)], +) +def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( + kda_token_states, enable_block_reuse, at_floor +): + """The per-token KDA states take memory per SSM slot beside the V2 pools. Without block reuse the SSM pool keeps + only its live floor, the states' bytes come out of the quota, and attention gets the rest; otherwise (replay caches + only, or block reuse) the typical step's ratio sizes the pool.""" + mgr = _v2_kda_token_state_manager(kda_token_states, enable_block_reuse) base_config = KVCacheManagerConfig( tokens_per_block=32, cache_tiers=[GpuCacheTierConfig(quota=128 << 20)], @@ -3806,6 +3814,7 @@ def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( sizes = stats.slot_sizes if hasattr(stats, "slot_sizes") else stats.slot_size role = "ssm" if mgr.ssm_bytes in [int(s) for s in sizes] else "attention" slots[role] = int(stats.total) + pool_bytes = runtime_manager.get_quota(GPU_LEVEL) finally: runtime_manager.shutdown() @@ -3813,12 +3822,53 @@ def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( if at_floor: # The floor's min-slots constraint is scaled by 1 / max_util_for_resume (0.97). assert floor <= slots["ssm"] <= floor + 1 - # Every other 2 MiB grain goes to attention's 256-byte pages. - assert slots["attention"] >= ((128 << 20) // (2 << 20) - (floor + 2)) * ((2 << 20) // 256) + # kda_state_tok, allocated beside the pools: (layers, SSM slots, num_spec, *ssm_state_shape) fp32. + token_state_bytes = ( + mgr.local_num_mamba_layers + * slots["ssm"] + * mgr._kda_replay_num_spec + * math.prod(mgr.ssm_state_shape) + * 4 + ) + # The pools and the states fit in the quota, and attention takes the rest of it, to within one 2 MiB grain. + assert (128 << 20) - (2 << 20) < pool_bytes + token_state_bytes <= 128 << 20 else: assert slots["ssm"] > 10 * floor +def test_v2_kda_token_states_reserve_quota_before_sizing(): + """At the live floor the token states of every slot the SSM pool keeps come out of the GPU quota before the pools + are sized, and the minimum quota and the quota for max_tokens count them all.""" + mgr = _v2_kda_token_state_manager(kda_token_states=True, enable_block_reuse=False) + plain = _v2_kda_token_state_manager(kda_token_states=False, enable_block_reuse=False) + assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 3 + # The pool keeps ceil(3 / max_util_for_resume) = 4 slots, each with 2 fp32 token states of 2 MiB; the quotas + # without token states differ by exactly those. + reserved = 4 * 2 * (2 << 20) + minimum = mgr._minimum_live_gpu_quota() + assert minimum - plain._minimum_live_gpu_quota() == reserved + for max_tokens in (0, 4096): + assert ( + mgr._get_quota_from_max_tokens(max_tokens) + - plain._get_quota_from_max_tokens(max_tokens) + == reserved + ) + + def build(quota): + return mgr._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + layers=_base_attention_layer_configs(2), + ) + ) + + with pytest.raises(ValueError, match="too small for live recurrent states"): + build(minimum - 1) + assert build(minimum).cache_tiers[0].quota == minimum - reserved + assert build(128 << 20).cache_tiers[0].quota == (128 << 20) - reserved + + def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: mgr = object.__new__(MambaCacheManager) mgr._impl = SimpleNamespace(use_kda_replay_update=True) From ea84b2378a9dd9f5cb64a3c69c4b1e6ce9726239 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Sat, 3 Oct 2026 05:11:08 -0700 Subject: [PATCH 23/67] [None][perf] Mamba state metadata: skip unchanged per-step uploads A decode step of a hybrid Mamba model uploaded its recurrent-state metadata on every step, even when the batch had not changed: the V2 manager's state indices (prepare_resources), the dummy-request mask (twice per step on the V2 manager: from prepare_resources and from get_state_indices) and Mamba2Metadata's copy of the state indices. Each is a small host-to-device copy on the execution stream between two CUDA graph replays, and the manager's two first allocate fresh pinned staging. Each of these buffers has one writer, so they now go through copy_to_device_if_changed, as the speculative engine's per-step inputs already do: a step whose values did not change enqueues no copy. test_v2_state_index_setup_skips_unchanged_uploads and test_prepare_skips_an_unchanged_state_index_upload overwrite the device buffers after a first step, then check that an unchanged step leaves them alone and that a changed step uploads again. Signed-off-by: Vasanth Sabavat --- .../_torch/modules/mamba/mamba2_metadata.py | 6 +-- .../kv_cache/mamba_cache_manager.py | 16 +++---- .../kv_cache/test_mamba_cache_manager.py | 48 +++++++++++++++++++ .../modules/mamba/test_mamba2_metadata.py | 38 +++++++++++++++ 4 files changed, 95 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 9b6545cfc58f..dd11a72d6cc5 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -24,7 +24,7 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ CUDA_GRAPH_DUMMY_REQUEST_ID -from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm._utils import copy_to_device_if_changed, prefer_pinned REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0 REPLAY_WORK_CACHE_SLOT = 1 @@ -500,8 +500,8 @@ def prepare(self, attn_metadata: AttentionMetadata): self.state_indices_cpu[:batch_size].copy_( torch.as_tensor(indices, dtype=self.state_indices_cpu.dtype)) - self.state_indices[:batch_size].copy_( - self.state_indices_cpu[:batch_size], non_blocking=True) + copy_to_device_if_changed(self.state_indices, + self.state_indices_cpu[:batch_size]) self._prepare_replay_work_items(kv_cache_manager, batch_size, num_contexts) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index ad417618950b..075bd85fbc13 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -43,8 +43,8 @@ PoolConfiguration, get_pp_layers) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, - nvtx_range, prefer_pinned, - torch_dtype_to_binding) + copy_to_device_if_changed, nvtx_range, + prefer_pinned, torch_dtype_to_binding) from tensorrt_llm.bindings.internal.batch_manager import ( LinearAttentionMetadata, LinearCacheType) from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -1470,10 +1470,8 @@ def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: if n > 0: self._dummy_request_mask_host[:n].copy_( torch.tensor(is_dummy, dtype=torch.bool)) - mask_staging = torch.empty_like(self._dummy_request_mask_host, - pin_memory=prefer_pinned()) - mask_staging.copy_(self._dummy_request_mask_host) - self._dummy_request_mask.copy_(mask_staging, non_blocking=True) + copy_to_device_if_changed(self._dummy_request_mask, + self._dummy_request_mask_host) def _reset_context_mamba_slots(self, num_contexts: int) -> None: if num_contexts == 0: @@ -4387,10 +4385,8 @@ def _setup_state_indices(self, f"request {req.py_request_id}") self._host_state_indices[i] = base_index - idx_staging = torch.empty_like(self._host_state_indices, - pin_memory=prefer_pinned()) - idx_staging.copy_(self._host_state_indices) - self.cuda_state_indices.copy_(idx_staging, non_blocking=True) + copy_to_device_if_changed(self.cuda_state_indices, + self._host_state_indices) is_dummy = [req.is_dummy for req in requests] self._refresh_dummy_request_mask(is_dummy) state_values = self._host_state_indices[:n].tolist() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 1f920c7b718c..2aabe73ff7b7 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3840,6 +3840,54 @@ def close(self) -> None: mgr.shutdown() +@skip_no_cuda +def test_v2_state_index_setup_skips_unchanged_uploads(): + """A step whose state indices and dummy flags did not change enqueues no upload of either: the device buffers + keep values the host never sent. A changed step uploads both again.""" + + class StateCache: + def __init__(self, slot: int) -> None: + self.slot = slot + + def get_ssm_block_base_index(self, _layer_group_id) -> int: + return self.slot + + def close(self) -> None: + pass + + mgr = _build_v2_hybrid_with_mamba_layer( + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_n_groups=4, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, + ) + live = SimpleNamespace(py_request_id=101, is_dummy=False) + padding = SimpleNamespace(py_request_id=202, is_dummy=True) + mgr.kv_cache_map.update({101: StateCache(0), 202: StateCache(1)}) + try: + mgr._setup_state_indices([live, padding]) + torch.cuda.synchronize() + assert mgr.cuda_state_indices[:2].tolist() == [0, 1] + assert mgr._dummy_request_mask[:2].tolist() == [False, True] + + mgr.cuda_state_indices.fill_(-1) + mgr._dummy_request_mask.fill_(True) + mgr._setup_state_indices([live, padding]) + torch.cuda.synchronize() + assert mgr.cuda_state_indices[:2].tolist() == [-1, -1] + assert mgr._dummy_request_mask[:2].tolist() == [True, True] + + mgr.kv_cache_map[202].slot = 2 + padding.is_dummy = False + mgr._setup_state_indices([live, padding]) + torch.cuda.synchronize() + assert mgr.cuda_state_indices[:2].tolist() == [0, 2] + assert mgr._dummy_request_mask[:2].tolist() == [False, False] + finally: + mgr.shutdown() + + def test_v2_kda_replay_seeds_disaggregated_generation_slots(): mgr = object.__new__(MambaHybridCacheManagerV2) mgr._use_kda_replay_update = True diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 56fb548b8793..16958c073f84 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -138,6 +138,44 @@ def test_prepare_handles_tensor_cached_tokens(self): assert metadata.chunk_indices is not None assert metadata.chunk_offsets is not None + def test_prepare_skips_an_unchanged_state_index_upload(self): + """A decode step whose state indices (a list from the cache manager) did not change enqueues no upload: the + device buffer keeps values the host never sent. A changed step uploads again.""" + + class ListCacheManager: + def __init__(self): + self.state_indices = [3, 1] + + def get_state_indices(self, request_ids, is_padding): + return self.state_indices[: len(request_ids)] + + manager = ListCacheManager() + metadata = Mamba2Metadata(max_batch_size=2, chunk_size=8) + seq_lens = torch.tensor([1, 1], dtype=torch.int) + attn_metadata = SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=0, + num_ctx_tokens=0, + kv_cache_manager=manager, + request_ids=[10, 11], + kv_cache_params=None, + ) + + metadata.prepare(attn_metadata) + torch.cuda.synchronize() + assert metadata.state_indices[:2].tolist() == [3, 1] + + metadata.state_indices.fill_(-1) + metadata.prepare(attn_metadata) + torch.cuda.synchronize() + assert metadata.state_indices[:2].tolist() == [-1, -1] + + manager.state_indices = [3, 2] + metadata.prepare(attn_metadata) + torch.cuda.synchronize() + assert metadata.state_indices[:2].tolist() == [3, 2] + def test_prepare_replay_work_items_write_first(self): class ReplayCacheManager: use_replay_state_update = True From 546e567f7b3962a4c68791f7b01bd12960a371a6 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Sat, 3 Oct 2026 05:38:20 -0700 Subject: [PATCH 24/67] [None][perf] DFlash: skip unchanged slot-buffer copies on decode steps DFlashSpecMetadata.prepare copied batch_indices_cuda and the worker's _batch_to_slot to the device on every step. Both now go through copy_to_device_if_changed, so a decode step whose batch is unchanged enqueues no copy. The worker's post-prefill rebuild of _batch_to_slot moves into _rebuild_batch_to_slot and goes through the same helper: the helper keeps the values it last copied on the buffer, so a plain copy there would let a later step that maps back to those values skip its copy and leave the rebuild's slots in place. test_dflash_copy_skip.py (cpu_only) checks that both buffers always hold the latest mapping, including a step whose mapping returns to an earlier one after a rebuild, and that an unchanged step copies nothing. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/speculative/dflash.py | 33 +++--- .../hw_agnostic/test_dflash_copy_skip.py | 111 ++++++++++++++++++ 2 files changed, 125 insertions(+), 19 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index d8097dffbd58..e0bf55de16b3 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -22,7 +22,7 @@ from torch import nn from tensorrt_llm._torch.custom_ops import inplace_slice_copy -from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm._utils import copy_to_device_if_changed, prefer_pinned from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -535,10 +535,7 @@ def prepare(self): assert self.request_ids is not None num_seqs = len(self.request_ids) - batch_indices = torch.arange( - num_seqs, dtype=torch.int, device="cpu", pin_memory=prefer_pinned() - ) - self.batch_indices_cuda[:num_seqs].copy_(batch_indices, non_blocking=True) + copy_to_device_if_changed(self.batch_indices_cuda, torch.arange(num_seqs, dtype=torch.int)) # Update slot mapping for DFlash context buffers worker = getattr(self, "_dflash_worker", None) @@ -576,10 +573,8 @@ def prepare(self): mapping = torch.tensor( [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids], dtype=torch.long, - device="cpu", - pin_memory=prefer_pinned(), ) - worker._batch_to_slot[:num_seqs].copy_(mapping, non_blocking=True) + copy_to_device_if_changed(worker._batch_to_slot, mapping) def is_layer_capture(self, layer_id: int) -> bool: return layer_id in self._capture_layer_set @@ -1341,6 +1336,16 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): attn_metadata.kv_lens_cuda[nc:bs] -= self._kv_rewind_amount attn_metadata.kv_lens_cuda[nc:bs].clamp_(min=0) + def _rebuild_batch_to_slot(self, request_ids: List[int]) -> None: + """Point _batch_to_slot at each request's slot after prefill assigned new ones. + + Like DFlashSpecMetadata.prepare, this writes through copy_to_device_if_changed, + which keeps the values it last copied on the buffer: every write of the buffer + must go through it, or a later unchanged batch skips its copy over other values. + """ + mapping = [self._req_to_slot.get(rid, self._dummy_slot) for rid in request_ids] + copy_to_device_if_changed(self._batch_to_slot, torch.tensor(mapping, dtype=torch.long)) + def _write_ctx_len(self, updates: dict[int, int]) -> None: """Apply a slot-to-length mapping to _ctx_len in one async scatter.""" if not updates: @@ -1636,17 +1641,7 @@ def _forward_impl( ) # Rebuild batch_to_slot after prefill assigns new slots if self._ctx_buf_inited and spec_metadata.request_ids: - num_seqs = len(spec_metadata.request_ids) - mapping = [ - self._req_to_slot.get(rid, self._dummy_slot) - for rid in spec_metadata.request_ids - ] - self._batch_to_slot[:num_seqs].copy_( - torch.tensor( - mapping, dtype=torch.long, device="cpu", pin_memory=prefer_pinned() - ), - non_blocking=True, - ) + self._rebuild_batch_to_slot(spec_metadata.request_ids) inputs = self.prepare_1st_drafter_inputs( input_ids=input_ids, diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py new file mode 100644 index 000000000000..e58a89378a6c --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py @@ -0,0 +1,111 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the DFlash worker's per-step slot buffers and copy_to_device_if_changed. + +DFlashSpecMetadata.prepare writes batch_indices_cuda and the worker's _batch_to_slot, and the +worker's post-prefill rebuild writes _batch_to_slot again. Both buffers go through +copy_to_device_if_changed, which skips a copy whose values equal the last ones it copied. The +tests check that each buffer always holds the latest mapping, including a step whose mapping +returns to an earlier one after the rebuild wrote another, and that an unchanged step enqueues +no copy. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker + +pytestmark = pytest.mark.cpu_only + +DUMMY_SLOT = 0 + + +def _worker(req_to_slot): + worker = SimpleNamespace( + _ctx_buf_inited=True, + _req_to_slot=dict(req_to_slot), + _req_ctx_pos={}, + _free_slots=[], + _graph_dummy_id_floor=1 << 30, + _dummy_slot=DUMMY_SLOT, + _batch_to_slot=torch.zeros(8, dtype=torch.long), + _write_ctx_len=lambda updates: None, + ) + # Generation requests are already assigned in these tests; prepare assigns nothing. + worker._assign_slot = lambda rid: worker._req_to_slot.get(rid) + return worker + + +def _metadata(worker, request_ids, num_generations): + return SimpleNamespace( + request_ids=list(request_ids), + num_generations=num_generations, + batch_indices_cuda=torch.full((8,), -1, dtype=torch.int), + _dflash_worker=worker, + ) + + +def _prepare(metadata, request_ids, num_generations): + metadata.request_ids = list(request_ids) + metadata.num_generations = num_generations + DFlashSpecMetadata.prepare(metadata) + + +def _rebuild(worker, request_ids): + DFlashWorker._rebuild_batch_to_slot(worker, list(request_ids)) + + +def test_mapping_that_returns_after_a_rebuild_is_written_again(): + # Request 5 decodes on slot 1. Request 8 arrives as a context request, so prepare maps it to + # the dummy slot; prefill then assigns it slot 4 and the rebuild writes [4, 1]. + worker = _worker({5: 1}) + metadata = _metadata(worker, [8, 5], num_generations=1) + _prepare(metadata, [8, 5], num_generations=1) + assert worker._batch_to_slot[:2].tolist() == [DUMMY_SLOT, 1] + worker._req_to_slot[8] = 4 + _rebuild(worker, [8, 5]) + assert worker._batch_to_slot[:2].tolist() == [4, 1] + # Request 8 is gone and context request 9 takes its row: prepare maps [dummy, 1] again, the + # values it copied two writes ago. The buffer must hold them, not the rebuild's [4, 1]. + del worker._req_to_slot[8] + _prepare(metadata, [9, 5], num_generations=1) + assert worker._batch_to_slot[:2].tolist() == [DUMMY_SLOT, 1] + + +def test_rebuild_then_prepare_back_to_back_keep_the_latest_mapping(): + worker = _worker({5: 1, 6: 2}) + metadata = _metadata(worker, [5, 6], num_generations=2) + for request_ids, slots in (([5, 6], [1, 2]), ([6, 5], [2, 1]), ([5, 6], [1, 2])): + _rebuild(worker, request_ids) + assert worker._batch_to_slot[:2].tolist() == slots + _prepare(metadata, request_ids, num_generations=2) + assert worker._batch_to_slot[:2].tolist() == slots + + +def test_batch_indices_hold_the_batch_rows(): + worker = _worker({5: 1, 6: 2, 7: 3}) + metadata = _metadata(worker, [5], num_generations=1) + for request_ids in ([5], [5, 6, 7], [5, 6], [5, 6, 7]): + _prepare(metadata, request_ids, num_generations=len(request_ids)) + n = len(request_ids) + assert metadata.batch_indices_cuda[:n].tolist() == list(range(n)) + + +def test_an_unchanged_step_enqueues_no_copy(monkeypatch): + worker = _worker({5: 1, 6: 2}) + metadata = _metadata(worker, [5, 6], num_generations=2) + _prepare(metadata, [5, 6], num_generations=2) + copies = [] + original = torch.Tensor.copy_ + + def counted(dst, src, *args, **kwargs): + copies.append(tuple(dst.shape)) + return original(dst, src, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "copy_", counted) + _prepare(metadata, [5, 6], num_generations=2) + monkeypatch.undo() + assert copies == [] + assert worker._batch_to_slot[:2].tolist() == [1, 2] From a9938c2a8a9fdb576358791c1372894bad408941 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Sat, 3 Oct 2026 17:01:55 -0700 Subject: [PATCH 25/67] [None][fix] KV cache manager V2: count the live-floor KDA token-state slots as the runtime does The runtime holds max_util_for_resume as a C++ float, and _build_cache_config passed that value to _live_floor_token_state_slots, while _state_quota (behind _minimum_live_gpu_quota and _get_quota_from_max_tokens) passed the Python double. When the floor divided by the ratio is a whole number, the two ceilings differ by one slot: at 0.95, a floor of 19 slots keeps 21 slots (19 / float32(0.95) is just above 20), while the double gave 20. The minimum and the quota for max_tokens then counted one slot's token states fewer than the build reserves. _live_floor_token_state_slots now rounds the ratio through float32, as kv_cache_manager_v2.py does for its own floor. The new test checks that the minimum quota and the build reserve the same 21 slots at a floor of 19 and a ratio of 0.95. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 6 +++- .../kv_cache/test_mamba_cache_manager.py | 29 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 1f9f6b0d04c7..693dedc7c2ed 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -20,6 +20,7 @@ from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, Optional, Protocol, Sequence, Tuple, Union, cast) +import numpy as np import torch import triton import triton.language as tl @@ -4035,7 +4036,10 @@ def _live_floor_token_state_slots(self, max_util_for_resume: float) -> int: return 0 floor_slots = (self._max_resident_sequences() + self._num_reserved_dummy_slots) - return math.ceil(floor_slots / max_util_for_resume) + # The runtime holds max_util_for_resume as a C++ float; round the + # Python value the same way, or the ceiling can differ by one slot + # when floor_slots / max_util_for_resume is a whole number. + return math.ceil(floor_slots / float(np.float32(max_util_for_resume))) def _quota_filling_request_capacity(self, gpu_quota: int) -> int: """A typical request capacity whose resident requests' attention pages diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 87cadcbe3bad..e2ba109a90fc 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3869,6 +3869,35 @@ def build(quota): assert build(128 << 20).cache_tiers[0].quota == (128 << 20) - reserved +def test_v2_kda_token_states_reserve_the_slots_the_runtime_keeps(): + """The quota checks and the build count the same slots when floor / max_util_for_resume is a whole number in + exact arithmetic: the runtime holds the ratio as a C++ float, so 19 / 0.95 keeps 21 slots, not 20.""" + mgr = _v2_kda_token_state_manager(kda_token_states=True, enable_block_reuse=False) + plain = _v2_kda_token_state_manager(kda_token_states=False, enable_block_reuse=False) + for m in (mgr, plain): + m.max_batch_size = 18 + m.kv_cache_config = KvCacheConfig( + avg_seq_len=64, + enable_block_reuse=False, + enable_partial_reuse=False, + max_util_for_resume=0.95, + ) + assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 19 + reserved = 21 * 2 * (2 << 20) + minimum = mgr._minimum_live_gpu_quota() + assert minimum - plain._minimum_live_gpu_quota() == reserved + quota = minimum + (64 << 20) + built = mgr._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=quota)], + layers=_base_attention_layer_configs(2), + max_util_for_resume=0.95, + ) + ) + assert quota - built.cache_tiers[0].quota == reserved + + def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: mgr = object.__new__(MambaCacheManager) mgr._impl = SimpleNamespace(use_kda_replay_update=True) From 54b87a271a06cbb7688eb67fa5d039ac8d812e95 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Sat, 3 Oct 2026 17:02:19 -0700 Subject: [PATCH 26/67] [None][fix] KV cache manager V2: refuse per-token KDA states off the SSM live floor kda_state_tok is allocated for every slot the SSM pool keeps, beside the V2 pools. Its bytes come out of the GPU quota only at the live floor (block reuse off, no pool_ratio, attention layers on the rank), where the slot count is known before the pools are sized. With block reuse, a pool_ratio, or no attention layers, the ratio-sized pool kept many more slots and kda_state_tok took their states outside the quota. The K3 targets only refused block reuse at their first forward, after the manager had allocated. _build_cache_config now raises ValueError for per-token states off the live floor. test_v2_kda_token_states_refuse_a_pool_off_the_live_floor covers the three cases. The ratio-sizing case of test_v2_kda_token_states_keep_ssm_pool_at_live_floor now runs without the states, since that combination is refused. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 12 ++++++++ .../kv_cache/test_mamba_cache_manager.py | 30 +++++++++++++++++-- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 693dedc7c2ed..6fb06cb63aac 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -4054,6 +4054,18 @@ def _quota_filling_request_capacity(self, gpu_quota: int) -> int: def _build_cache_config( self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: kv_cache_config = self.kv_cache_config + if (getattr(self, "_kda_token_states", False) + and self.local_num_mamba_layers > 0 + and (config.initial_pool_ratio is not None + or not self._ssm_pool_at_live_floor(kv_cache_config))): + # kda_state_tok is allocated for every slot the SSM pool keeps; + # only at the live floor is that count known before sizing, so + # that its bytes can come out of the quota. + raise ValueError( + "Per-token KDA states (kda_token_states) need the SSM pool at " + "its live floor: kv_cache_config.enable_block_reuse off, " + "kv_cache_config.pool_ratio unset, and attention layers on " + "this rank.") cache_tiers = config.cache_tiers gpu_quota = cache_tiers[0].quota minimum_live_quota = self._minimum_live_gpu_quota() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index e2ba109a90fc..b4805d3ed744 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3793,14 +3793,14 @@ def _v2_kda_token_state_manager(kda_token_states, enable_block_reuse): @pytest.mark.parametrize( ("kda_token_states", "enable_block_reuse", "at_floor"), - [(True, False, True), (False, False, False), (True, True, False)], + [(True, False, True), (False, False, False), (False, True, False)], ) def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( kda_token_states, enable_block_reuse, at_floor ): """The per-token KDA states take memory per SSM slot beside the V2 pools. Without block reuse the SSM pool keeps - only its live floor, the states' bytes come out of the quota, and attention gets the rest; otherwise (replay caches - only, or block reuse) the typical step's ratio sizes the pool.""" + only its live floor, the states' bytes come out of the quota, and attention gets the rest; without the states + (replay caches only, with or without block reuse) the typical step's ratio sizes the pool.""" mgr = _v2_kda_token_state_manager(kda_token_states, enable_block_reuse) base_config = KVCacheManagerConfig( tokens_per_block=32, @@ -3836,6 +3836,30 @@ def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( assert slots["ssm"] > 10 * floor +@pytest.mark.parametrize("off_the_floor", ["block_reuse", "pool_ratio", "no_attention"]) +def test_v2_kda_token_states_refuse_a_pool_off_the_live_floor(off_the_floor): + """kda_state_tok is allocated for every slot the SSM pool keeps, so with block reuse, a pool_ratio, or no + attention layers on the rank (the SSM pool sized by its ratio, its slot count unknown before sizing) the build + refuses the per-token states instead of allocating them outside the quota.""" + mgr = _v2_kda_token_state_manager( + kda_token_states=True, enable_block_reuse=off_the_floor == "block_reuse" + ) + pool_ratio = None + if off_the_floor == "pool_ratio": + pool_ratio = [0.5, 0.5] + elif off_the_floor == "no_attention": + mgr._attention_cache_bytes_per_token = lambda: 0 + with pytest.raises(ValueError, match="need the SSM pool at its live floor"): + mgr._build_cache_config( + KVCacheManagerConfig( + tokens_per_block=32, + cache_tiers=[GpuCacheTierConfig(quota=128 << 20)], + layers=_base_attention_layer_configs(2), + initial_pool_ratio=pool_ratio, + ) + ) + + def test_v2_kda_token_states_reserve_quota_before_sizing(): """At the live floor the token states of every slot the SSM pool keeps come out of the GPU quota before the pools are sized, and the minimum quota and the quota for max_tokens count them all.""" From 1e6eb678c9d87ceec81d92c0ea266b5cf08ec272 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Sat, 3 Oct 2026 18:04:24 -0700 Subject: [PATCH 27/67] [None][fix] KV cache manager V2: keep a suspended request's KDA replay state MambaHybridCacheManagerV2 keeps the KDA replay state (prev_num_accepted_tokens and the six replay caches) per SSM slot, outside the V2 pools, and remembers each request's last slot until free_resources. A V2 suspend (the scheduler's eviction and self-eviction, the pool rebalance) only suspends the pages. Under pressure the host tier takes them, the slot goes to another request, and the suspended request resumes in whichever slot is free. Its next _setup_state_indices then relocated the old slot's contents, by then the other request's replay state, or, back in its old slot, kept them: the next verify replayed another request's drafts. suspend_request now copies the slot's replay state to host memory and forgets the slot; _setup_state_indices writes it into the slot the request resumes in. update_resources gives the host drafter's acceptance of a request suspended after its batch was scheduled to the saved state instead of the old slot. Relocation and the copy share one list of the per-slot buffers. The new tests resume a request in a new slot and in its old one and record a late acceptance on a manager stub, and on the V2 runtime let new caches take a suspended request's slot through the host tier before it resumes. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 108 ++++++++-- .../kv_cache/test_mamba_cache_manager.py | 200 ++++++++++++++++++ 2 files changed, 286 insertions(+), 22 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index ad417618950b..778832696355 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3235,6 +3235,10 @@ def __init__( } self._request_id_to_state_index = {} self._request_id_to_is_dummy = {} + # KDA replay state of suspended requests (suspend_request), and the + # draft counts update_resources accepts for them after their suspend. + self._suspended_kda_replay_state: Dict[int, List[torch.Tensor]] = {} + self._suspended_kda_replay_accepted: Dict[int, int] = {} state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) @@ -4268,6 +4272,18 @@ def _reset_kda_replay_slots( if slots.numel() > 0: self.prev_num_accepted_tokens[slots] = 0 + def _kda_replay_state_buffers(self) -> List[Tuple[torch.Tensor, int]]: + """The per-slot KDA replay state kept outside the V2 pools, as + (buffer, slot dimension) pairs.""" + assert self.prev_num_accepted_tokens is not None + buffers = [(self.prev_num_accepted_tokens, 0)] + for replay_buffer in (self.kda_conv_q, self.kda_conv_k, self.kda_conv_v, + self.kda_qkg_cache, self.kda_v_cache, + self.kda_beta_cache): + assert replay_buffer is not None + buffers.append((replay_buffer, 1)) + return buffers + def _relocate_kda_replay_slots(self, old_slots: List[int], new_slots: List[int]) -> None: """Move replay history when V2 remaps a live request's SSM slot.""" @@ -4285,26 +4301,43 @@ def _relocate_kda_replay_slots(self, old_slots: List[int], destination_slots = torch.tensor([new for _, new in moves], dtype=torch.long, device=device) - replay_buffers = ( - self.kda_conv_q, - self.kda_conv_k, - self.kda_conv_v, - self.kda_qkg_cache, - self.kda_v_cache, - self.kda_beta_cache, - ) - for replay_buffer in replay_buffers: - assert replay_buffer is not None - replay_buffer.index_copy_( - 1, - destination_slots, - replay_buffer.index_select(1, source_slots), - ) - self.prev_num_accepted_tokens.index_copy_( - 0, - destination_slots, - self.prev_num_accepted_tokens.index_select(0, source_slots), - ) + for buffer, dim in self._kda_replay_state_buffers(): + buffer.index_copy_(dim, destination_slots, + buffer.index_select(dim, source_slots)) + + def suspend_request(self, req: LlmRequest) -> None: + """Suspend the request's KV cache and keep its KDA replay state. + + The replay state is indexed by SSM slot, outside the V2 pools. While + the request is suspended, its pages can move to a lower tier and its + slot can go to another request; it resumes in whichever slot is free. + The slot's replay state is copied out here, and the request's next + ``_setup_state_indices`` writes it into the slot it resumes in. + """ + request_id = req.py_request_id + kv_cache = self.kv_cache_map.get(request_id) + slot = self._request_id_to_state_index.get(request_id, -1) + if (self.use_kda_replay_update and slot >= 0 and kv_cache is not None + and kv_cache.is_active): + saved = [] + for buffer, dim in self._kda_replay_state_buffers(): + state = buffer.select(dim, slot) + host = torch.empty(state.shape, + dtype=state.dtype, + pin_memory=prefer_pinned()) + host.copy_(state, non_blocking=True) + saved.append(host) + self._suspended_kda_replay_state[request_id] = saved + del self._request_id_to_state_index[request_id] + super().suspend_request(req) + + def _restore_kda_replay_state(self, request_id: int, slot: int, + saved: List[torch.Tensor]) -> None: + for (buffer, dim), host in zip(self._kda_replay_state_buffers(), saved): + buffer.select(dim, slot).copy_(host, non_blocking=True) + accepted = self._suspended_kda_replay_accepted.pop(request_id, None) + if accepted is not None: + self.prev_num_accepted_tokens[slot] = accepted def add_dummy_requests( self, @@ -4345,6 +4378,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self._request_id_to_state_index.pop(request.py_request_id, None) self._request_id_to_is_dummy.pop(request.py_request_id, None) self._branch_snapshot_points.pop(request.py_request_id, None) + self._suspended_kda_replay_state.pop(request.py_request_id, None) + self._suspended_kda_replay_accepted.pop(request.py_request_id, None) super().free_resources(request, pin_on_release) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -4399,15 +4434,33 @@ def _setup_state_indices(self, self._request_id_to_is_dummy[req.py_request_id] = dummy if self.use_kda_replay_update: + # A suspended request that comes back as a context request starts + # over; one that comes back generating takes its replay state. + for request in requests[:num_contexts]: + self._suspended_kda_replay_state.pop(request.py_request_id, + None) + self._suspended_kda_replay_accepted.pop(request.py_request_id, + None) + saved_states = [ + self._suspended_kda_replay_state.pop(request.py_request_id, + None) + for request in replay_requests + ] new_state_values = [ self._request_id_to_state_index[request.py_request_id] for request in replay_requests ] self._relocate_kda_replay_slots(old_state_values, new_state_values) self._reset_kda_replay_slots([ - new for old, new in zip(old_state_values, new_state_values) - if old < 0 + new for old, new, saved in zip(old_state_values, + new_state_values, saved_states) + if old < 0 and saved is None ]) + for request, new, saved in zip(replay_requests, new_state_values, + saved_states): + if saved is not None: + self._restore_kda_replay_state(request.py_request_id, new, + saved) def _record_kda_replay_acceptance( self, @@ -4532,6 +4585,15 @@ def update_resources( "Mixed drafted/undrafted generation batch is not supported " "for KDA replay bookkeeping") + for request in generation_requests: + # Suspended after this batch was scheduled: its saved replay state + # predates this acceptance, so the acceptance goes with it. + request_id = request.py_request_id + if (request_id in self._suspended_kda_replay_state and + not self._request_id_to_is_dummy.get(request_id, False)): + self._suspended_kda_replay_accepted[request_id] = max( + 0, request.py_num_accepted_draft_tokens) + active_requests = [ request for request in generation_requests if request.py_request_id in self._request_id_to_state_index @@ -4781,4 +4843,6 @@ def shutdown(self): self.kda_qkg_cache = None self.kda_v_cache = None self.kda_beta_cache = None + self._suspended_kda_replay_state.clear() + self._suspended_kda_replay_accepted.clear() super().shutdown() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 1f920c7b718c..b4a5e7cdcf5a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3719,6 +3719,8 @@ def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): mgr._record_kda_replay_in_update_resources = True mgr._request_id_to_state_index = {101: 1, 303: 2} mgr._request_id_to_is_dummy = {101: False, 303: True} + mgr._suspended_kda_replay_state = {} + mgr._suspended_kda_replay_accepted = {} mgr.prev_num_accepted_tokens = torch.tensor([7, 8, 9], dtype=torch.int32) base_update = MagicMock() monkeypatch.setattr(KVCacheManagerV2, "update_resources", base_update) @@ -3745,6 +3747,204 @@ def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): assert mgr.prev_num_accepted_tokens.tolist() == [7, 2, 9] +_KDA_REPLAY_BUFFERS = ( + "kda_conv_q", + "kda_conv_k", + "kda_conv_v", + "kda_qkg_cache", + "kda_v_cache", + "kda_beta_cache", +) + + +class _SlotCache: + """A V2 KV cache stand-in that reports one SSM slot and can be suspended.""" + + def __init__(self, slot: int) -> None: + self.slot = slot + self.is_active = True + + def get_ssm_block_base_index(self, _layer_group_id) -> int: + return self.slot + + def suspend(self) -> None: + self.is_active = False + + +def _v2_kda_replay_stub(num_slots=4): + """A V2 hybrid manager stub with CPU KDA replay buffers, for the state-index setup that prepare_resources runs.""" + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr.local_num_mamba_layers = 1 + mgr._use_kda_replay_update = True + mgr._use_replay_state_update = False + mgr._request_id_to_state_index = {} + mgr._request_id_to_is_dummy = {} + mgr._suspended_kda_replay_state = {} + mgr._suspended_kda_replay_accepted = {} + mgr._host_state_indices = torch.zeros(num_slots, dtype=torch.int32) + mgr.cuda_state_indices = torch.zeros(num_slots, dtype=torch.int32) + mgr._dummy_request_mask = None + mgr.mamba_ssm_rand_seed = None + mgr.kv_cache_map = {} + mgr.ssm_layer_group_id = 0 + mgr.prev_num_accepted_tokens = torch.zeros(num_slots, dtype=torch.int32) + for name in _KDA_REPLAY_BUFFERS: + setattr(mgr, name, torch.zeros(1, num_slots, 2, 4)) + return mgr + + +def _write_kda_replay_slot(mgr, slot, accepted, value): + """What a verify step leaves in its request's slot for the next one.""" + mgr.prev_num_accepted_tokens[slot] = accepted + for name in _KDA_REPLAY_BUFFERS: + getattr(mgr, name)[:, slot] = value + + +def _kda_replay_slot(mgr, slot): + return int(mgr.prev_num_accepted_tokens[slot]), [ + float(getattr(mgr, name)[0, slot].flatten()[0]) for name in _KDA_REPLAY_BUFFERS + ] + + +@pytest.mark.parametrize("resume_slot", [1, 0], ids=["new_slot", "old_slot"]) +def test_v2_kda_replay_state_follows_a_suspended_request(resume_slot): + """While a request is suspended its SSM page can move to a lower tier and its slot can go to another request. It + resumes in a new slot, or in its old one after that request has finished, with its own replay state either way.""" + mgr = _v2_kda_replay_stub() + first = SimpleNamespace(py_request_id=1, is_dummy=False) + second = SimpleNamespace(py_request_id=2, is_dummy=False) + mgr.kv_cache_map[1] = _SlotCache(0) + mgr._setup_state_indices([first], num_contexts=0) + _write_kda_replay_slot(mgr, 0, accepted=2, value=1.0) + mgr.suspend_request(first) + assert not mgr.kv_cache_map[1].is_active + + mgr.kv_cache_map[2] = _SlotCache(0) + mgr._setup_state_indices([second], num_contexts=1) + mgr._reset_context_mamba_slots(1) + mgr._setup_state_indices([second], num_contexts=0) + _write_kda_replay_slot(mgr, 0, accepted=3, value=7.0) + resumed = [first] + if resume_slot == 0: + del mgr.kv_cache_map[2] + del mgr._request_id_to_state_index[2] + else: + resumed.append(second) + mgr.kv_cache_map[1].slot = resume_slot + mgr.kv_cache_map[1].is_active = True + mgr._setup_state_indices(resumed, num_contexts=0) + + assert _kda_replay_slot(mgr, resume_slot) == (2, [1.0] * len(_KDA_REPLAY_BUFFERS)) + if resume_slot != 0: + assert _kda_replay_slot(mgr, 0) == (3, [7.0] * len(_KDA_REPLAY_BUFFERS)) + assert not mgr._suspended_kda_replay_state + + +def test_v2_kda_replay_late_acceptance_follows_a_suspended_request(monkeypatch): + """The host drafter's acceptance is recorded by update_resources after the next batch is scheduled. A request + suspended in between takes that acceptance with its saved replay state instead of writing its old slot.""" + mgr = _v2_kda_replay_stub() + mgr._record_kda_replay_in_update_resources = True + monkeypatch.setattr(KVCacheManagerV2, "update_resources", MagicMock()) + request = SimpleNamespace( + py_request_id=1, is_dummy=False, py_draft_tokens=[11, 12], py_num_accepted_draft_tokens=0 + ) + mgr.kv_cache_map[1] = _SlotCache(0) + mgr._setup_state_indices([request], num_contexts=0) + _write_kda_replay_slot(mgr, 0, accepted=1, value=1.0) + mgr.suspend_request(request) + mgr.kv_cache_map[2] = _SlotCache(0) + _write_kda_replay_slot(mgr, 0, accepted=3, value=7.0) + request.py_num_accepted_draft_tokens = 2 + mgr.update_resources(SimpleNamespace(generation_requests=[request])) + + assert int(mgr.prev_num_accepted_tokens[0]) == 3 + mgr.kv_cache_map[1].slot = 2 + mgr.kv_cache_map[1].is_active = True + mgr._setup_state_indices([request], num_contexts=0) + assert _kda_replay_slot(mgr, 2) == (2, [1.0] * len(_KDA_REPLAY_BUFFERS)) + + +def _v2_kda_replay_manager_with_large_slots(): + """A real V2 hybrid manager at Kimi K3's per-rank KDA shapes (6 heads, K = V = 128, conv width 4) on 6 layers. An + SSM slot (~2.25 MiB) is larger than the 2 MiB allocation grain, so the SSM pool holds a handful of slots and runs + out before the attention pages do. max_util_for_resume 1.0 lets resume admit caches until the pool is full.""" + num_mamba_layers = 6 + return MambaHybridCacheManagerV2( + mamba_d_state=128, + mamba_d_conv=4, + mamba_num_heads=6, + mamba_n_groups=6, + mamba_head_dim=128, + mamba_num_layers=num_mamba_layers, + mamba_layer_mask=[True] * num_mamba_layers + [False], + mamba_cache_dtype=torch.bfloat16, + mamba_ssm_cache_dtype=torch.float32, + kv_cache_config=KvCacheConfig( + max_tokens=4096, enable_block_reuse=False, max_util_for_resume=1.0 + ), + kv_cache_type=CacheTypeCpp.SELF, + num_layers=1, + num_kv_heads=4, + head_dim=64, + tokens_per_block=32, + max_seq_len=128, + max_batch_size=4, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + spec_config=MTPDecodingConfig(max_draft_len=2), + layer_mask=[False] * num_mamba_layers + [True], + vocab_size=1024, + conv_state_layout="q_k_v", + kda_replay_num_spec=2, + ) + + +@skip_no_cuda +def test_v2_kda_replay_state_follows_a_suspended_request_on_the_runtime(): + """The same on the V2 runtime: once new caches fill the SSM pool, the suspended request's SSM page moves to the + host tier and one of them takes its slot; the request resumes in another slot with its own replay state.""" + mgr = _v2_kda_replay_manager_with_large_slots() + stream = mgr._stream.cuda_stream + others = [] + request = None + try: + (request,) = mgr.add_dummy_requests( + [1], token_nums=[8], is_gen=True, max_num_draft_tokens=2 + ) + slot = mgr._request_id_to_state_index[1] + _write_kda_replay_slot(mgr, slot, accepted=2, value=1.0) + own = _kda_replay_slot(mgr, slot) + mgr.suspend_request(request) + holder = None + for _ in range(mgr.prev_num_accepted_tokens.shape[0] + 1): + cache = mgr.impl.create_kv_cache() + assert cache.resume(stream) + others.append(cache) + assert cache.resize(1) + if cache.get_ssm_block_base_index(mgr.ssm_layer_group_id) == slot: + holder = cache + break + assert holder is not None, "no cache took the suspended request's SSM slot" + _write_kda_replay_slot(mgr, slot, accepted=3, value=7.0) + for cache in others: + if cache is not holder: + cache.close() + others = [holder] + assert mgr.resume_request(request) + mgr._setup_state_indices([request], num_contexts=0) + new_slot = mgr._request_id_to_state_index[1] + + assert new_slot != slot + assert _kda_replay_slot(mgr, new_slot) == own + finally: + for cache in others: + cache.close() + if request is not None: + mgr.free_resources(request) + mgr.shutdown() + + @skip_no_cuda def test_v2_kda_replay_relocates_live_slot_history(): mgr = _build_v2_hybrid_with_mamba_layer( From 10fd5cf9f52d06c37934722bef98a844a92b94ed Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Sat, 3 Oct 2026 18:29:36 -0700 Subject: [PATCH 28/67] [None][fix] KV cache manager V2: refuse a KV pool rebalance with the KDA replay caches The KDA replay caches and prev_num_accepted_tokens are allocated once, one entry per slot the SSM pool keeps at construction (get_page_index_upper_bound). With enable_kv_pool_rebalance, adjust() resizes the pool groups to the auto-tuner's target ratio, which gives the SSM pool one slot per typical request, so a pool built smaller than that grows past the buffers. A request in a slot past them makes the prev_num_accepted_tokens[slot] = 0 reset fail, and the KDA verify kernels index the replay entries by slot without a bound. Nothing refused the combination: _util.py refuses a rebalance only with block reuse, and the executor's rebalance gate skips only two-model drafters, so one-model speculative decoding on KDA layers reaches it. _build_cache_config now raises ValueError for the KDA replay path with enable_kv_pool_rebalance, as KvCacheConfig does for fp8_ds_mla's fixed pool views. test_v2_kda_replay_refuses_kv_pool_rebalance covers the refusal and the control without the replay caches. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 12 +++++++++++ .../kv_cache/test_mamba_cache_manager.py | 21 +++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 778832696355..bd299fb07164 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3952,6 +3952,18 @@ def _minimum_live_gpu_quota(self) -> int: def _build_cache_config( self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: kv_cache_config = self.kv_cache_config + if (getattr(self, "_use_kda_replay_update", False) + and self.local_num_mamba_layers > 0 + and kv_cache_config.enable_kv_pool_rebalance): + # The KDA replay caches and prev_num_accepted_tokens hold one + # entry per SSM slot the pool keeps at construction; a rebalance + # can grow the SSM pool past them, and the verify kernels index + # them by slot without a bound. + raise ValueError( + "The KDA replay caches (speculative decoding on KDA layers) " + "are sized once from the SSM pool's slot count, and a KV pool " + "rebalance can grow that pool past them: set " + "kv_cache_config.enable_kv_pool_rebalance=False.") cache_tiers = config.cache_tiers gpu_quota = cache_tiers[0].quota minimum_live_quota = self._minimum_live_gpu_quota() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index b4a5e7cdcf5a..0b379e3a2877 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2353,6 +2353,7 @@ def _build_v2_hybrid_with_mamba_layer( mamba_n_groups=1, mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, + enable_kv_pool_rebalance=False, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2380,6 +2381,7 @@ def _build_v2_hybrid_with_mamba_layer( enable_branch_snapshot=enable_branch_snapshot, ), dtype=kv_cache_dtype, + enable_kv_pool_rebalance=enable_kv_pool_rebalance, ) return MambaHybridCacheManagerV2( mamba_d_state=8, @@ -3676,6 +3678,25 @@ def test_v2_kda_replay_validates_configuration( ) +@pytest.mark.parametrize("kda_replay", [True, False], ids=["kda_replay", "no_replay"]) +def test_v2_kda_replay_refuses_kv_pool_rebalance(kda_replay): + """The KDA replay caches and prev_num_accepted_tokens hold one entry per SSM slot the pool keeps at construction, + and a KV pool rebalance can grow that pool past them, so the manager refuses enable_kv_pool_rebalance with them; + without them the rebalance stays allowed.""" + if not kda_replay: + _build_v2_hybrid_with_mamba_layer(enable_kv_pool_rebalance=True).shutdown() + return + with pytest.raises(ValueError, match="enable_kv_pool_rebalance=False"): + _build_v2_hybrid_with_mamba_layer( + spec_config=MTPDecodingConfig(max_draft_len=2), + conv_state_layout="q_k_v", + mamba_n_groups=4, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, + enable_kv_pool_rebalance=True, + ) + + def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: mgr = object.__new__(MambaCacheManager) mgr._impl = SimpleNamespace(use_kda_replay_update=True) From 50ac9d64e4db6a88b4c1531f727e31585dae56ce Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Sat, 3 Oct 2026 22:57:01 -0700 Subject: [PATCH 29/67] [None][fix] kda_prefill: release evicted scratch with its cache entry _get_buffers keeps the prefill scratch for at most _BUF_CACHE_MAX_ENTRIES token-count shapes and evicts the oldest entry on a new shape. The K123 and akk_inv launches wrapped that scratch with _ct_cached, whose id-keyed _input_wrap_cache holds each wrapper until the wrapped tensor is garbage collected. The wrapper's DLPack capsule keeps the tensor alive, so that never happens: every evicted entry's scratch stayed allocated. A server whose prefill batches keep changing token count grows by one entry's scratch per eviction (GSM8K on Kimi K3: about 500 MiB per 200 requests on every rank). Cache the scratch wrappers in the scratch's own _buf_cache entry instead (_scratch_ct), so evicting the entry releases the scratch and its wrappers together; the per-call from_dlpack savings stay. _ct_cached and _input_wrap_cache have no other users and are removed. Test: test_evicted_prefill_scratch_is_released runs cyclic passes over more token-count shapes than the cache holds. After every later pass the cache holds at most _BUF_CACHE_MAX_ENTRIES entries, the allocated bytes are back at their value after the first pass, and the tensors of an entry cached at the end of the first pass, which the next pass evicts, are freed. Signed-off-by: Vasanth Sabavat --- .../custom_ops/cute_dsl_kimi_k3_custom_ops.py | 66 +++++++++---------- .../kimi_kda/test_kda_cache_soundness.py | 52 +++++++++++++++ 2 files changed, 83 insertions(+), 35 deletions(-) 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 a7aacaa3d8ce..8b8a9d09300e 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 @@ -164,13 +164,6 @@ def _cute_int_type(dtype): # single-seq cu_seqlens. _varlen_single_seqlen_cache = {} -# id(tensor) -> cute_wrapper. The wrappers themselves are stateless views -# over the tensor's storage, so they remain valid as long as the tensor's -# data pointer / shape / strides don't change. Callers that reuse the same -# tensor objects across iterations (typical benchmark pattern) hit the -# cache; per-call fresh tensors (the executor runtime pattern) rebuild. -_input_wrap_cache = {} - def _prune_on_gc(cache, key, *keyobjs): """Drop ``cache[key]`` when any of ``keyobjs`` is garbage-collected. @@ -187,24 +180,24 @@ def _prune_on_gc(cache, key, *keyobjs): weakref.finalize(o, cache.pop, key, None) -def _ct_cached(t, etype): - """`_ct(t, etype)` with id(t)-based cache. Returns the same cute wrapper - for repeated calls with the same tensor object, avoiding per-call - `from_dlpack` overhead (~5-10us each). +def _scratch_ct(cute_wrappers, name, t, etype): + """`_ct(t, etype)` for the ``_get_buffers`` scratch tensor ``t``, cached + under ``name`` in ``cute_wrappers``, the wrapper dict of the same + ``_buf_cache`` entry, so a launch skips its ``from_dlpack`` (~5-10us each). - ONLY use for tensors with process-long lifetime (module params, the - module-level scratch from ``_get_buffers``): the cached wrapper pins the - tensor's storage, so the weakref pruning never fires for the keyed - object and a per-call activation would be pinned forever (~100MB/call - leak in the executor runtime). Per-call tensors must use plain ``_ct``. + A CuTe wrapper keeps the tensor it wraps alive (its DLPack capsule holds a + reference), so a wrapper is cached only beside its tensor: evicting the + entry from ``_buf_cache`` releases the scratch and its wrappers together. + Per-call tensors are never cached; they are wrapped with plain ``_ct``. + Without an entry (``cute_wrappers`` is None) ``t`` is wrapped afresh. """ - key = (id(t), etype) - w = _input_wrap_cache.get(key) - if w is None: - w = _ct(t, etype) - _input_wrap_cache[key] = w - _prune_on_gc(_input_wrap_cache, key, t) - return w + if cute_wrappers is None: + return _ct(t, etype) + key = ("scratch_ct", name, etype) + cached = cute_wrappers.get(key) + if cached is None or cached[0] is not t: + cached = cute_wrappers[key] = (t, _ct(t, etype)) + return cached[1] # Cache for dt_bias `.float().contiguous().view(H, K)` + cute wrapper. @@ -532,7 +525,8 @@ def _launch_k4_persistent( gk_ct = cute_wrappers["gk_ct"] # v is a per-call activation — wrap fresh every call (never cache; see - # _ct_cached docstring). + # _scratch_ct: only the _get_buffers scratch keeps cached wrappers, in its + # own _buf_cache entry). v_view = v_beta.reshape(-1, H, V_dim) if v_beta.dim() == 4 else v_beta v_ct = from_dlpack(v_view, assumed_align=16).mark_layout_dynamic() v_ct.element_type = bf16 @@ -717,9 +711,9 @@ def _launch_fused_k123_inv( # Inputs are guaranteed contiguous by upstream linear projections. # A_log is fp32 model param; .float() is no-op when dtype already matches. - # q/k/g/beta/cu/ci are per-call activations: plain _ct, never cached - # (see _ct_cached docstring). The _get_buffers scratch below is - # module-persistent, so caching its wrappers is safe and worthwhile. + # q/k/g/beta/cu/ci are per-call activations: plain _ct, never cached. + # The _get_buffers scratch below is wrapped once per _buf_cache entry, and + # its wrappers live in that entry's cute_wrappers (see _scratch_ct). q_ct = _ct(q, cutlass.BFloat16) k_ct = _ct(k, cutlass.BFloat16) g_ct = _ct(g, cutlass.BFloat16) @@ -731,14 +725,16 @@ def _launch_fused_k123_inv( else: raise ValueError(f"Kimi K3 KDA prefill beta must be float32 or bfloat16, got {beta.dtype}") beta_ct = _ct(beta, beta_etype) - beta_activated_ct = _ct_cached(beta_activated, cutlass.BFloat16) + beta_activated_ct = _scratch_ct( + cute_wrappers, "beta_activated", beta_activated, cutlass.BFloat16 + ) - ks_ct = _ct_cached(k_scaled, cutlass.BFloat16) - kg_ct = _ct_cached(kg, cutlass.BFloat16) - qs_ct = _ct_cached(q_scaled, cutlass.BFloat16) - gk_ct = _ct_cached(gk_last_exp, cutlass.Float32) - aqk_ct = _ct_cached(A_qk, cutlass.BFloat16) - akk_ct = _ct_cached(A_kk_inv, cutlass.BFloat16) + ks_ct = _scratch_ct(cute_wrappers, "k_scaled", k_scaled, cutlass.BFloat16) + kg_ct = _scratch_ct(cute_wrappers, "kg", kg, cutlass.BFloat16) + qs_ct = _scratch_ct(cute_wrappers, "q_scaled", q_scaled, cutlass.BFloat16) + gk_ct = _scratch_ct(cute_wrappers, "gk_last_exp", gk_last_exp, cutlass.Float32) + aqk_ct = _scratch_ct(cute_wrappers, "A_qk", A_qk, cutlass.BFloat16) + akk_ct = _scratch_ct(cute_wrappers, "A_kk", A_kk_inv, cutlass.BFloat16) if is_varlen: cu_ct = _ct(cu_seqlens, _cute_int_type(cu_seqlens.dtype)) @@ -824,7 +820,7 @@ def _launch_fused_k123_inv( else: akk_beta_etype = cutlass.BFloat16 akk_beta_ct = ( - _ct_cached(beta_for_akk, akk_beta_etype) + _scratch_ct(cute_wrappers, "beta_activated", beta_for_akk, akk_beta_etype) if use_beta_sigmoid_in_kernel else _ct(beta_for_akk, akk_beta_etype) ) diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py b/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py index aa04f8e196a6..d042bc61469d 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py @@ -25,9 +25,17 @@ against FLA. ``test_repeated_single_sequence_metadata_matches_fla`` additionally covers the Phase 2.1 poisoning case (same cu_seqlens object, two calls). + +3. Bounded scratch. The prefill scratch cache ``_buf_cache`` keeps at most + ``_BUF_CACHE_MAX_ENTRIES`` token-count shapes, and the runtime calls with a + new token count for most prefill batches, so entries are evicted all the + time. ``test_evicted_prefill_scratch_is_released`` checks that an evicted + entry's scratch is freed: nothing else, such as a cached CuTe wrapper, may + keep it alive. """ import gc +import weakref import pytest import torch @@ -145,6 +153,10 @@ def _allocate_with_recycled_id(target_id: int, make, attempts: int = 512): return None +def _tensor_weakrefs(values) -> list[weakref.ref]: + return [weakref.ref(value) for value in values if isinstance(value, torch.Tensor)] + + @pytest.mark.parametrize("regime", ["eqlen", "varlen"]) @torch.no_grad() def test_indexed_prefill_uses_current_stream( @@ -246,6 +258,46 @@ def test_recycled_cu_seqlens_id_matches_fla( assert module._varlen_pure_cache.get(id(cu_seqlens)) is False +@torch.no_grad() +def test_evicted_prefill_scratch_is_released( + dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch], + gate_params: tuple[torch.Tensor, torch.Tensor], +) -> None: + """Scratch evicted from the bounded prefill buffer cache is freed.""" + optimized, _ = dispatch_pair + module = _op_module() + # Every call has its own token count, and one pass has more shapes than the + # cache holds, so the cyclic passes below miss and evict on every call. + # Two sequences per call keep the batch at four or more chunks. + first_lengths = [300 + 64 * i for i in range(module._BUF_CACHE_MAX_ENTRIES + 4)] + + def run_pass() -> None: + for i, first_length in enumerate(first_lengths): + sequence_lengths = [first_length, 200] + inputs = _make_inputs(sum(sequence_lengths), seed=500 + i) + run_indexed_prefill(optimized, gate_params, inputs, _make_cu_seqlens(sequence_lengths)) + torch.cuda.synchronize() + _flush_tensor_cache_pins() + gc.collect() + + # The first pass compiles the kernels and leaves the cache holding the last + # shapes of the cycle; every later pass evicts all of those entries and + # ends with the same shapes cached again. + run_pass() + baseline = torch.cuda.memory_allocated() + evicted_scratch = _tensor_weakrefs(next(iter(module._buf_cache.values()))) + assert evicted_scratch + for iteration in range(2): + run_pass() + assert len(module._buf_cache) <= module._BUF_CACHE_MAX_ENTRIES + grown = torch.cuda.memory_allocated() - baseline + assert grown <= 16 << 20, ( + f"pass {iteration}: {grown / (1 << 20):.1f} MiB more allocated after the evictions" + ) + alive = sum(ref() is not None for ref in evicted_scratch) + assert alive == 0, f"pass {iteration}: {alive} tensors of an evicted entry are still alive" + + @torch.no_grad() def test_repeated_single_sequence_metadata_matches_fla( dispatch_pair: tuple[KDAKernelDispatch, KDAKernelDispatch], From a0509b335194704ebe193e39855b537855c78ea0 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Mon, 5 Oct 2026 17:15:48 -0700 Subject: [PATCH 30/67] [None][fix] LoRA overlap: leave the output record to maybe_execute_in_parallel forward_with_base recorded the calling stream on each LoRA output of the overlapped path. maybe_execute_in_parallel now records fn1's outputs itself, outside CUDA graph capture and torch.compile tracing, so there the LoRA record only repeated it. The overlapped path runs only inside the CUDA graph runner, which turns multi-stream on for its eager warm-up steps and for the capture. During the capture the LoRA record still ran, which is what the helper avoids: the caching allocator holds a recorded block freed in the capture until the capture ends, so each overlapped LoRA call took new graph pool memory, and under cudaMallocAsync a recorded output freed inside the capture breaks it. maybe_execute_in_parallel is now the only place the rule is applied. lora_aux_stream is used only through it, so in a capture each overlapped call's fn1 waits on the calling stream (event0) before it allocates, after the previous call's outputs were read. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/peft/lora/layer.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tensorrt_llm/_torch/peft/lora/layer.py b/tensorrt_llm/_torch/peft/lora/layer.py index 4d7e87de1935..fd64e6e2efc4 100644 --- a/tensorrt_llm/_torch/peft/lora/layer.py +++ b/tensorrt_llm/_torch/peft/lora/layer.py @@ -308,8 +308,6 @@ def lora_forward() -> tuple[torch.Tensor | None, ...]: for lora_output in lora_outputs: if not isinstance(lora_output, torch.Tensor): continue - if execute_in_parallel: - lora_output.record_stream(torch.cuda.current_stream()) base_output = add_lora_result(base_output, lora_output) return base_output From 77a450e895b85e2a7a1c66bc2c4a946ef1655607 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Mon, 5 Oct 2026 17:15:51 -0700 Subject: [PATCH 31/67] [None][doc] multi_stream_utils: note the compiled path and the error's source Two notes in maybe_execute_in_parallel, no behavior change: - The docstring says that callers should not record fn1's outputs themselves, and that a caller torch.compile traces with disable_on_compile=False runs fn1 on aux_stream without the record, outside CUDA graphs too, so it must order the reuse of fn1's outputs itself. - The comment on the caught cudaMallocAsync error names where its text comes from: the TORCH_INTERNAL_ASSERT in CudaMallocAsyncAllocator::recordStream (c10/cuda/CUDAMallocAsyncAllocator.cpp, torch 2.14.0a0+4fdf77b940.nv26.08), to be re-checked on torch upgrades. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/modules/multi_stream_utils.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/multi_stream_utils.py b/tensorrt_llm/_torch/modules/multi_stream_utils.py index c8b3c3af0ee3..ce96e4221dcd 100644 --- a/tensorrt_llm/_torch/modules/multi_stream_utils.py +++ b/tensorrt_llm/_torch/modules/multi_stream_utils.py @@ -43,6 +43,10 @@ def _record_stream(obj: Any, stream: torch.cuda.Stream) -> None: # The cudaMallocAsync backend asserts on memory it did not # allocate (a buffer from outside PyTorch, wrapped as a # tensor). It never frees that memory: nothing to record. + # The text is the TORCH_INTERNAL_ASSERT in + # CudaMallocAsyncAllocator::recordStream + # (c10/cuda/CUDAMallocAsyncAllocator.cpp, torch + # 2.14.0a0+4fdf77b940.nv26.08); re-check it on torch upgrades. if (torch.cuda.memory.get_allocator_backend() != "cudaMallocAsync" or "ptr not found in ptr_info" not in str(e)): @@ -76,7 +80,12 @@ def maybe_execute_in_parallel( stream, so their memory outlives the caller's use. During capture nothing is recorded: a captured graph relies on every later use of aux_stream first waiting for the calling stream, as the next call's event0.wait() - does. Other work captured on aux_stream must do the same. + does. Other work captured on aux_stream must do the same. Callers should + not record fn1()'s outputs themselves. + + A caller that torch.compile traces with disable_on_compile=False runs + fn1() on aux_stream without this record, outside CUDA graphs too, so it + must order the reuse of fn1()'s outputs itself. Args: fn0 (Callable): callable for the default stream From c3ced8d8c13f4fb0f609f8e1802d2e9fc6334c7a Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Mon, 5 Oct 2026 17:15:52 -0700 Subject: [PATCH 32/67] [None][test] test_multi_stream_utils: bound the capture's block count - test_aux_outputs_dropped_inside_graph_capture asserted that every output in the capture reused one block. That is the native allocator's choice today, not this code's contract. It now asserts fewer blocks than outputs: recording in the capture would hold every output's block until the capture ends. - test_aux_output_over_memory_torch_did_not_allocate imports cuda.bindings instead of importorskip-ing it. cuda-python is a TensorRT-LLM requirement (requirements.txt), so the check of the cudaMallocAsync error text always runs. Signed-off-by: Vasanth Sabavat --- .../_torch/modules/test_multi_stream_utils.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/modules/test_multi_stream_utils.py b/tests/unittest/_torch/modules/test_multi_stream_utils.py index 91166f14134c..b691408e739e 100644 --- a/tests/unittest/_torch/modules/test_multi_stream_utils.py +++ b/tests/unittest/_torch/modules/test_multi_stream_utils.py @@ -77,7 +77,10 @@ def fn1(): def test_aux_output_over_memory_torch_did_not_allocate() -> None: # The native allocator skips such memory in record_stream; cudaMallocAsync raises on it. - cudart = pytest.importorskip("cuda.bindings.runtime") + # A plain import, not importorskip: cuda-python is a TensorRT-LLM requirement, so this test + # always runs and catches a change in cudaMallocAsync's error text. + from cuda.bindings import runtime as cudart + numel = 1024 err, ptr = cudart.cudaMalloc(numel * 4) assert err == cudart.cudaError_t.cudaSuccess @@ -171,8 +174,10 @@ def layer_stack(): torch.cuda.synchronize() assert torch.equal(acc, x * layers + sum(range(layers))) if torch.cuda.memory.get_allocator_backend() == "native": - # A recorded block freed during capture is held until the capture ends. - assert len(set(addresses)) == 1, f"{len(set(addresses))} blocks for {layers} outputs" + # A recorded block freed during capture is held until the capture ends, so recording + # would give every output its own block. Which free block an allocation reuses is up to + # the allocator; the bound does not depend on it. + assert len(set(addresses)) < layers, f"{len(set(addresses))} blocks for {layers} outputs" def test_compiled_caller_traces_fullgraph() -> None: From f39f875fa8c0b7a8bba5897273b566d21d66c7dc Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Mon, 5 Oct 2026 21:51:57 -0700 Subject: [PATCH 33/67] [None][perf] Kimi K3 attn_res_fwd online kernel: double-buffer ws_stats instead of a second barrier attn_res_fwd_online_v2_kernel ordered each chunk's cross-warp ws_stats reads before the next chunk's writes with a second NamedBarrier per chunk. Instead, double-buffer ws_stats by chunk parity: chunk gci writes and reads ws_stats[gci & 1]. A warp writes chunk gci + 2 into the same buffer only after it passes chunk gci + 1's barrier, which every warp reaches only after reading chunk gci, so the reads need no barrier of their own. gci keeps counting across token blocks, so the parity holds there too. This removes one bar.sync per chunk for CONSUMER_WARPS * NC * 8 bytes of shared memory. The N == 1 tile kernel shares FwdSmemPlan; it uses buffer 0 and keeps its own second barrier. The __syncwarp() before the lane-0 bar_consumed arrive and the cross-proxy fences stay. Signed-off-by: Vasanth Sabavat --- .../kernels/kimiK3AttnRes/attnResFwd.cu | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu index 39a14df806b6..9f0521b84f18 100644 --- a/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu +++ b/cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu @@ -217,7 +217,10 @@ struct FwdSmemPlan { alignas(16) uint64_t bar_ready[CHUNK_DEPTH]; alignas(16) uint64_t bar_consumed[CHUNK_DEPTH]; - alignas(16) float2 ws_stats[CONSUMER_WARPS][NC]; + // online_v2 writes chunk gci's per-warp statistics to ws_stats[gci & 1]. A warp writes chunk gci + 2 to the same + // buffer only after it passes chunk gci + 1's barrier, and every warp reads chunk gci before it arrives there, so + // the reads need no barrier of their own. The N == 1 tile kernel uses buffer 0 and orders its reads itself. + alignas(16) float2 ws_stats[2][CONSUMER_WARPS][NC]; alignas(16) float logits_all[N_MAX]; uint32_t tmem_base; }; @@ -538,7 +541,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c #pragma unroll for (int n = 0; n < N_CHUNK; n++) { - plan.ws_stats[comp_wid][n] = reduce_pair[n]; + plan.ws_stats[gci & 1][comp_wid][n] = reduce_pair[n]; } } cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 0); @@ -551,7 +554,7 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c #pragma unroll for (int w = 0; w < CONSUMER_WARPS; w++) { - totals = float2_add(totals, plan.ws_stats[w][n]); + totals = float2_add(totals, plan.ws_stats[gci & 1][w][n]); } local_rsig = rsqrtf(totals.x / H + eps_cache); local_logit = totals.y * local_rsig; @@ -564,10 +567,6 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c { cross_warp_tail(lane); } - // ws_stats is one buffer reused every chunk. The barrier above orders this chunk's writes before the - // reads; this one orders the reads before the next chunk's writes, so a warp that has finished - // reading cannot overwrite its row while a slower warp still reads it. - cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 0); float logit_n[N_CHUNK]; #pragma unroll for (int n = 0; n < N_CHUNK; n++) @@ -968,7 +967,7 @@ __global__ void __launch_bounds__(BLK, 1) } if (lane == 0) { - plan.ws_stats[comp_wid][0] = make_float2(sq_local, dot_local); + plan.ws_stats[0][comp_wid][0] = make_float2(sq_local, dot_local); } cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 0); @@ -978,7 +977,7 @@ __global__ void __launch_bounds__(BLK, 1) #pragma unroll for (int w = 0; w < CONSUMER_WARPS; w++) { - totals = float2_add(totals, plan.ws_stats[w][0]); + totals = float2_add(totals, plan.ws_stats[0][w][0]); } float rs = rsqrtf(totals.x / H + rms_eps); rsigma_out[tb] = rs; From b2a6af468072050b516e04211986dd6043c7e1b0 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 1 Oct 2026 23:00:49 -0700 Subject: [PATCH 34/67] Kimi K3 KDA prefill: K1 warps finish reading sPartialLast before the next chunk fused_kernel123's K1 warps read their scanned prefix row and the chunk total from sPartialLast after the scan barrier, and the next chunk's partial sums overwrite sPartialLast before any K1 barrier (racecheck: write at fused_k123.py:1268 against reads at :1293 / :1299). A K1 warp that reached its reads late used the next chunk's sums: with one K1 warp stalled 20 us after the scan barrier, every case of test_fused_prefill_beta_sigmoid_matches_unfused_kernel changed (output max |d| up to 7e14). One K1 barrier after the reads orders them; outputs are bit-identical to before. test_fused_prefill_does_not_depend_on_k1_warp_timing runs a copy of the kernel with that stall and requires the unpatched kernel's bits; it fails without the barrier. Signed-off-by: Vasanth Sabavat --- .../blackwell/kimi_k3_kda/fused_k123.py | 4 + .../modules/kimi_kda/test_kda_prefill_op.py | 125 ++++++++++++++++++ 2 files changed, 129 insertions(+) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py index 0db49ce5a164..99f0cd3f535d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/fused_k123.py @@ -1298,6 +1298,10 @@ def fused_kernel123( for vi in cutlass.range_constexpr(VEC): rPrefix[vi] = sPartialLast[warp_row_group - 1, col_base + vi] + # The next chunk's partial sums overwrite sPartialLast, so every + # K1 warp must have read its prefix and the chunk total first. + k1_internal_barrier() + # ---- Pass 2a: ONLY cumsum + write csGcum (critical path, minimal work) ---- for vi in cutlass.range_constexpr(VEC): rAcc[vi] = rPrefix[vi] 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 c56c246b0bf3..f78d3722a8fb 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 @@ -2,7 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 """Parity tests for the optimized Kimi K3 KDA prefill op.""" +import importlib.util +import sys from collections.abc import Mapping +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -490,6 +493,128 @@ def randn( _assert_close(state_fused, state_unfused) +# fused_k123's K1 warps meet at this barrier after the cross-warp scan, then +# read the scanned prefix rows and the chunk total from sPartialLast. +_K1_SCAN_BARRIER = " k1_internal_barrier()\n" +_K1_SCAN_READS = ( + _K1_SCAN_BARRIER + + "\n" + + " for vi in cutlass.range_constexpr(VEC):\n" + + " rGkLast[vi] = sPartialLast[K1_ROW_GROUPS - 1, col_base + vi]\n" +) +# K1 warp 1 reads the prefix row 0 and the chunk total. +_LATE_K1_WARP = ( + " if k1_warp == 1:\n" + " late_t = _test_globaltimer()\n" + " late_end = late_t + cutlass.Int64(20000)\n" + " while late_t < late_end:\n" + " late_t = _test_globaltimer()\n" +) +_GLOBALTIMER = """ + +@dsl_user_op +def _test_globaltimer(*, loc=None, ip=None): + return cutlass.Int64( + llvm.inline_asm( + T.i64(), + [], + "mov.u64 $0, %globaltimer;", + "=l", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) +""" + + +@torch.no_grad() +@pytest.mark.parametrize( + "sequence_lengths", + [None, [100, 257, 63]], + ids=["padded_eqlen", "varlen"], +) +def test_fused_prefill_does_not_depend_on_k1_warp_timing( + sequence_lengths: list[int] | None, tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A K1 warp that reads the scanned partial sums late gets the same bits. + + A copy of fused_k123 stalls K1 warp 1 for 20 us right after the scan + barrier, so the other K1 warps run ahead into the next chunk before it + reads sPartialLast. Outputs and states must equal the unpatched kernel's. + """ + from tensorrt_llm._torch.custom_ops import cute_dsl_kimi_k3_custom_ops as kimi_k3_ops + from tensorrt_llm._torch.cute_dsl_kernels.blackwell.kimi_k3_kda import fused_k123 + + source = Path(fused_k123.__file__).read_text() + assert source.count(_K1_SCAN_READS) == 1, "fused_k123's K1 scan reads moved" + late_source = source.replace( + _K1_SCAN_READS, _K1_SCAN_BARRIER + _LATE_K1_WARP + _K1_SCAN_READS[len(_K1_SCAN_BARRIER) :] + ) + late_path = tmp_path / "fused_k123_late_k1_warp.py" + late_path.write_text(late_source + _GLOBALTIMER) + spec = importlib.util.spec_from_file_location(late_path.stem, late_path) + late_module = importlib.util.module_from_spec(spec) + monkeypatch.setitem(sys.modules, late_path.stem, late_module) + spec.loader.exec_module(late_module) + + torch.manual_seed(10) + if sequence_lengths is None: + sequence_length = 300 + num_sequences = 1 + cu_seqlens = None + else: + sequence_length = sum(sequence_lengths) + num_sequences = len(sequence_lengths) + cu_seqlens = torch.tensor( + [0, *torch.tensor(sequence_lengths).cumsum(0).tolist()], + dtype=torch.long, + device="cuda", + ) + + def randn(*shape: int, dtype: torch.dtype = torch.bfloat16, scale: float = 0.05): + return (torch.randn(*shape, dtype=torch.float32, device="cuda") * scale).to(dtype) + + def l2norm(x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.normalize(x.float(), dim=-1).to(torch.bfloat16) + + shape = (1, sequence_length, NUM_HEADS, HEAD_DIM) + inputs = dict( + q=l2norm(randn(*shape)), + k=l2norm(randn(*shape)), + v=randn(*shape), + g=randn(*shape), + beta=randn(1, sequence_length, NUM_HEADS, dtype=torch.float32), + scale=HEAD_DIM**-0.5, + state_indices=torch.arange(num_sequences, dtype=torch.int32, device="cuda"), + safe_gate=True, + lower_bound=-5.0, + use_gate_in_kernel=True, + A_log=torch.randn(NUM_HEADS, dtype=torch.float32, device="cuda") * 0.1, + dt_bias=torch.randn(NUM_HEADS * HEAD_DIM, dtype=torch.float32, device="cuda") * 0.1, + cu_seqlens=cu_seqlens, + use_beta_sigmoid_in_kernel=True, + ) + initial_state = randn( + num_sequences, NUM_HEADS, HEAD_DIM, HEAD_DIM, dtype=torch.float32, scale=0.01 + ) + + def run() -> tuple[torch.Tensor, torch.Tensor]: + state = initial_state.clone() + output = torch.ops.trtllm.kda_prefill(state_pool=state, **inputs).clone() + return output, state + + expected_output, expected_state = run() + monkeypatch.setattr(kimi_k3_ops, "_fused_make_host", late_module.make_host_function) + monkeypatch.setattr(kimi_k3_ops, "_fused_k123_cache", {}) + output, state = run() + + torch.testing.assert_close(output, expected_output, rtol=0, atol=0) + torch.testing.assert_close(state, expected_state, rtol=0, atol=0) + + @torch.no_grad() def test_kda_mixer_empty_prefill(): """An empty token payload still publishes the conv state for its slot.""" From 2e598e10cb797f5b3297c145810f3a43a15c8738 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 1 Oct 2026 00:52:42 -0700 Subject: [PATCH 35/67] [None][feat] One-model speculative decoding: the worker produces the target logits SpecDecOneEngineForCausalLM asks its spec worker for the target logits (SpecWorkerBase.target_logits) instead of calling the logits processor itself. The default is the logits processor's fp32 [rows, vocab], so no worker changes behaviour. A worker whose acceptance reads another layout can override it; when its forward then returns this TP rank's vocabulary shard as "logits", it sets "logits_vocab_shard", and the engine all-gathers the logits before running logits post-processors (only when a scheduled request has one). Signed-off-by: Vasanth Sabavat --- .../_torch/models/modeling_speculative.py | 8 +- .../_torch/pyexecutor/model_engine.py | 24 +++ tensorrt_llm/_torch/speculative/interface.py | 19 ++ .../test_spec_worker_target_logits.py | 179 ++++++++++++++++++ 4 files changed, 227 insertions(+), 3 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_spec_worker_target_logits.py diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 38788d8f06b4..ba54c65d95d1 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1840,12 +1840,14 @@ def forward( hidden_states = hidden_states[:attn_metadata.num_tokens] if self.spec_worker is not None: - # get logits - logits = self.logits_processor.forward( + # The target logits, in the layout the worker's acceptance reads. + logits = self.spec_worker.target_logits( hidden_states[spec_metadata.gather_ids], self.lm_head, + self.logits_processor, attn_metadata, - True, + spec_metadata, + self.draft_model, ) # VLM wrappers (e.g. Qwen3VLModelBase) replace input_ids with diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 2429877c9cb4..06b544b0d9b4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -6496,6 +6496,30 @@ def _execute_logit_post_processors(self, return logits_tensor = outputs["logits"] + if outputs.get("logits_vocab_shard", False): + # The spec worker kept the logits vocabulary-sharded (this TP + # rank's columns). A post-processor sees the whole vocabulary, so + # gather them first, only when a request has one. Every TP rank + # schedules the same requests, so all of them take the gather; + # under attention DP they would not. + if self.mapping.enable_attention_dp: + raise RuntimeError( + "vocabulary-sharded spec-worker logits need every TP rank " + "to schedule the same requests, which attention DP does " + "not") + if not any( + getattr(request, "py_logits_post_processors", None) + for request in scheduled_requests.all_requests()): + return + from ..distributed import allgather + logits_tensor = allgather(logits_tensor, self.mapping, + dim=-1).float() + # Drop the vocabulary's TP padding, as the LM head's own gather + # does. + vocab_size = getattr(getattr(self.model, "lm_head", None), + "num_embeddings", None) + if vocab_size is not None: + logits_tensor = logits_tensor[..., :vocab_size] logits_row_offset = 0 request_groups = ( diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 85d3fd1dd829..db98d830cac5 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -1648,6 +1648,25 @@ def forward(self, *args, **kwargs): def _forward_impl(self, *args, **kwargs): """Worker-specific forward logic, called by SpecWorkerBase.forward.""" + def target_logits(self, hidden_states: torch.Tensor, lm_head, + logits_processor, attn_metadata, spec_metadata, + draft_model) -> torch.Tensor: + """The target logits of the gathered rows ``hidden_states``, which + the model hands to ``forward`` as ``logits``: fp32 [rows, vocab]. + + A worker whose acceptance reads another layout overrides this. If its + ``forward`` then returns this TP rank's vocabulary shard as + ``"logits"``, it also returns ``"logits_vocab_shard": True``, and the + engine gathers the logits (on every TP rank together, so not under + attention DP) before any logits post-processor sees them; the + post-processors get a read-only view, since the acceptance already + ran inside ``forward``. Such a worker applies a guided decoder's + bitmask to its shard at the shard's column offset, or rejects guided + decoding. + """ + return logits_processor.forward(hidden_states, lm_head, attn_metadata, + True) + def _ensure_spec_dec_state_restored(self, attn_metadata, spec_metadata): """Restore attn-metadata spec-dec state if a failure skipped it. diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_spec_worker_target_logits.py b/tests/unittest/_torch/speculative/hw_agnostic/test_spec_worker_target_logits.py new file mode 100644 index 000000000000..c3a541a91654 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_spec_worker_target_logits.py @@ -0,0 +1,179 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The one-model spec worker's target-logits hook and the engine gather of vocabulary-sharded logits.""" + +from types import SimpleNamespace + +import pytest +import torch + +import tensorrt_llm._torch.distributed as distributed +from tensorrt_llm._torch.models.modeling_speculative import SpecDecOneEngineForCausalLM +from tensorrt_llm._torch.pyexecutor.model_engine import PyTorchModelEngine +from tensorrt_llm._torch.speculative.interface import SpecWorkerBase + +pytestmark = pytest.mark.cpu_only + + +class _RecordingLogitsProcessor: + def __init__(self, logits): + self.logits = logits + self.calls = [] + + def forward(self, *args): + self.calls.append(args) + return self.logits + + +def test_default_target_logits_are_the_logits_processors(): + logits = torch.randn(3, 16) + processor = _RecordingLogitsProcessor(logits) + hidden, lm_head, attn_metadata = torch.randn(3, 8), object(), object() + + out = SpecWorkerBase.target_logits( + object(), hidden, lm_head, processor, attn_metadata, object(), object() + ) + + assert out is logits + assert len(processor.calls) == 1 + args = processor.calls[0] + assert args[0] is hidden and args[1] is lm_head and args[2] is attn_metadata + assert args[3] is True + + +def _generation_request(post_processors, tokens=(1, 2, 3)): + return SimpleNamespace( + py_request_id=7, + py_logits_post_processors=post_processors, + py_beam_width=1, + get_beam_width_by_iter=lambda for_next_iteration: 1, + get_tokens=lambda beam: list(tokens), + ) + + +def _engine(vocab_size=None, attention_dp=False): + engine = object.__new__(PyTorchModelEngine) + engine.mapping = SimpleNamespace(is_last_pp_rank=lambda: True, enable_attention_dp=attention_dp) + engine.model = SimpleNamespace(lm_head=SimpleNamespace(num_embeddings=vocab_size)) + return engine + + +def _scheduled(*requests): + return SimpleNamespace( + context_requests=[], generation_requests=list(requests), all_requests=lambda: list(requests) + ) + + +def _recording_allgather(monkeypatch, vocab): + calls = [] + + def allgather(tensor, mapping, dim=-1): + calls.append((tensor, mapping, dim)) + return torch.zeros(tensor.shape[0], vocab, dtype=torch.bfloat16) + + monkeypatch.setattr(distributed, "allgather", allgather) + return calls + + +def test_sharded_logits_are_gathered_for_post_processors(monkeypatch): + calls = _recording_allgather(monkeypatch, vocab=32) + seen = [] + request = _generation_request( + [lambda req_id, rows, tokens, s, c: seen.append(tuple(rows.shape))] + ) + shard = torch.ones(1, 8, dtype=torch.bfloat16) + engine = _engine() + + engine._execute_logit_post_processors( + _scheduled(request), {"logits": shard, "logits_vocab_shard": True} + ) + + assert len(calls) == 1 + assert calls[0][0] is shard and calls[0][1] is engine.mapping and calls[0][2] == -1 + assert seen == [(1, 1, 32)] # the post-processor sees the whole vocabulary + + +def test_sharded_logits_without_post_processors_are_not_gathered(monkeypatch): + calls = _recording_allgather(monkeypatch, vocab=32) + + _engine()._execute_logit_post_processors( + _scheduled(_generation_request(None)), + {"logits": torch.ones(1, 8, dtype=torch.bfloat16), "logits_vocab_shard": True}, + ) + + assert calls == [] + + +def test_full_logits_are_never_gathered(monkeypatch): + calls = _recording_allgather(monkeypatch, vocab=32) + seen = [] + request = _generation_request( + [lambda req_id, rows, tokens, s, c: seen.append(tuple(rows.shape))] + ) + + _engine()._execute_logit_post_processors(_scheduled(request), {"logits": torch.ones(1, 32)}) + + assert calls == [] + assert seen == [(1, 1, 32)] + + +def test_sharded_logits_drop_the_vocab_padding(monkeypatch): + # 4 ranks x 8 columns hold a 30-token vocabulary: the gather's last 2 columns are TP padding. + _recording_allgather(monkeypatch, vocab=32) + seen = [] + request = _generation_request( + [lambda req_id, rows, tokens, s, c: seen.append(tuple(rows.shape))] + ) + + _engine(vocab_size=30)._execute_logit_post_processors( + _scheduled(request), + {"logits": torch.ones(1, 8, dtype=torch.bfloat16), "logits_vocab_shard": True}, + ) + + assert seen == [(1, 1, 30)] + + +def test_sharded_logits_under_attention_dp_raise(monkeypatch): + calls = _recording_allgather(monkeypatch, vocab=32) + request = _generation_request([lambda *args: None]) + + with pytest.raises(RuntimeError, match="attention DP"): + _engine(attention_dp=True)._execute_logit_post_processors( + _scheduled(request), + {"logits": torch.ones(1, 8, dtype=torch.bfloat16), "logits_vocab_shard": True}, + ) + assert calls == [] + + +def test_the_spec_model_hands_the_workers_target_logits_to_its_forward(): + hidden = torch.randn(4, 8) + gather_ids = torch.tensor([1, 3]) + calls = {} + + class Worker: + def target_logits(self, *args): + calls["target_logits"] = args + return "worker logits" + + def __call__(self, **kwargs): + calls["forward"] = kwargs + return "worker outputs" + + model = object.__new__(SpecDecOneEngineForCausalLM) + torch.nn.Module.__init__(model) + model.model = lambda **kwargs: hidden + model.layer_idx = -1 + model.spec_worker = Worker() + model.lm_head, model.logits_processor, model.draft_model = "lm_head", "processor", "drafter" + spec_metadata = SimpleNamespace(is_layer_capture=lambda layer: False, gather_ids=gather_ids) + attn_metadata = SimpleNamespace(padded_num_tokens=None, num_tokens=4) + + out = SpecDecOneEngineForCausalLM.forward( + model, attn_metadata, input_ids=torch.arange(4), spec_metadata=spec_metadata + ) + + assert out == "worker outputs" + rows, *rest = calls["target_logits"] + assert torch.equal(rows, hidden[gather_ids]) + assert rest == ["lm_head", "processor", attn_metadata, spec_metadata, "drafter"] + assert calls["forward"]["logits"] == "worker logits" From b12bfad22ce1492b1c60be149a0014d216c06f62 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 04:24:34 -0700 Subject: [PATCH 36/67] DFlash: empty the dummy slot in the drafting step, not in prepare() The dummy slot's context length grows only where the drafting step adds the step's accepted tokens. Emptying it right after that add, with a zero_() on the fixed slot that a captured step replays, keeps it at 0 between steps without host work. prepare() again writes slot lengths only for evicted requests, instead of two pinned tensors, two H2D copies and an index_copy_ on every step. Dummy rows also attend no context in the step itself, rather than that step's accepted tokens. test_dflash_dummy_slot.py drives the drafting-step update (_advance_ctx_len), including a captured CUDA graph replayed over a dirtied dummy slot, and checks that prepare() writes only evicted slots. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/speculative/dflash.py | 20 ++++-- .../hw_agnostic/test_dflash_dummy_slot.py | 70 ++++++++++++++----- 2 files changed, 67 insertions(+), 23 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index d760e15fbc78..15cf8645d8e4 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -551,10 +551,6 @@ def prepare(self): evicted[slot] = 0 worker._req_ctx_pos.pop(rid, None) worker._free_slots.append(slot) - # The dummy slot starts every step empty. Its rows (CUDA-graph padding, warmup dummies) add their accepted - # tokens to its length like any request, and their table rows are the padding request's pages, the rest - # mapped to page 0 (another request's). Left growing, a padding row's context K / V would land there. - evicted[worker._dummy_slot] = 0 worker._write_ctx_len(evicted) # A disagg generation worker receives prompt KV instead of @@ -1345,6 +1341,19 @@ def _apply_kv_rewind_after_draft(self, attn_metadata, spec_metadata): attn_metadata.kv_lens_cuda[nc:bs] -= self._kv_rewind_amount attn_metadata.kv_lens_cuda[nc:bs].clamp_(min=0) + def _advance_ctx_len(self, slots: torch.Tensor, num_accepted: torch.Tensor) -> None: + """Add this step's accepted tokens to the context lengths of ``slots``, on the device. + + Padding, warmup and attention-DP dummy rows share the dummy slot, which no request owns + or frees. Emptying it right after the add keeps it at 0 between steps: those rows attend + no context, and their next context K / V starts at column 0, inside the padding + request's own pages (past them, its page-table row maps page 0, another request's). + The store is to a fixed slot, so a captured step replays it. + """ + self._ctx_len[slots] += num_accepted + self._ctx_len.clamp_(max=self._max_ctx) + self._ctx_len[self._dummy_slot].zero_() + def _write_ctx_len(self, updates: dict[int, int]) -> None: """Apply a slot-to-length mapping to _ctx_len in one async scatter.""" if not updates: @@ -2065,8 +2074,7 @@ def prepare_1st_drafter_inputs( if v_new is not None: self._ctx_v_buf[slot_long, :, col_long] = v_new - self._ctx_len[slots] += gen_num_accepted_long - self._ctx_len.clamp_(max=self._max_ctx) + self._advance_ctx_len(slots, gen_num_accepted_long) num_ctx_per_req_t = self._ctx_len[slots] if self._ctx_block_tables is not None: diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py index 37d837aa7c28..ace488b3fad5 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -12,12 +12,13 @@ # 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. -"""DFlash's dummy slot (CUDA-graph padding and warmup dummies) starts every step at context length 0. +"""DFlash's dummy slot (CUDA-graph padding, warmup and attention-DP dummies) is at context length 0 between steps. -Padding rows add their accepted tokens to the dummy slot's length like any request, while their page-table rows are -the padding request's pages with the rest mapped to page 0, which another request owns. A dummy length that kept -growing would put the padding rows' context K / V on that page. ``DFlashSpecMetadata.prepare`` runs before every step -(eager or replayed graph); after it the dummy slot's length is 0 and every real slot's is untouched. +Dummy rows add their accepted tokens to the dummy slot's length like any request, while a padding row's page-table row +maps past the padding request's own pages to page 0, which another request owns. A dummy length that kept growing +would put the padding rows' context K / V on that page. The drafting step's ``DFlashWorker._advance_ctx_len`` empties +the slot right after the add, inside the captured step, so ``DFlashSpecMetadata.prepare`` writes slot lengths only for +evicted requests. """ from functools import partial @@ -43,11 +44,13 @@ def _worker(lengths): _free_slots=[3], _graph_dummy_id_floor=FLOOR, _dummy_slot=len(lengths) - 1, + _max_ctx=1000, _ctx_len=torch.tensor(lengths, dtype=torch.long, device="cuda"), _ctx_len_host=list(lengths), _batch_to_slot=torch.zeros(8, dtype=torch.long, device="cuda"), ) w._write_ctx_len = partial(DFlashWorker._write_ctx_len, w) + w._advance_ctx_len = partial(DFlashWorker._advance_ctx_len, w) w._assign_slot = lambda rid, *a, **k: None return w @@ -63,21 +66,54 @@ def _prepare(worker, request_ids, num_generations): torch.cuda.synchronize() -def test_padding_steps_keep_the_dummy_slot_empty(): +def _cuda(values): + return torch.tensor(values, dtype=torch.long, device="cuda") + + +def test_padded_steps_keep_the_dummy_slot_empty(): worker = _worker([40, 300, 7, 0, 0]) - padded = [11, 12, 13, FLOOR + 1, FLOOR + 1] # three requests in a graph of five - for step in range(4): - # What the step's acceptance does to the slots of its rows (k3_ctx_kv / the torch path): + accepted. - _prepare(worker, padded, num_generations=5) + # Three requests in a graph of five. + _prepare(worker, [11, 12, 13, FLOOR + 1, FLOOR + 1], num_generations=5) + slots = worker._batch_to_slot[:5].clone() + assert slots.tolist() == [0, 1, 2, 4, 4] + for step in range(1, 5): + worker._advance_ctx_len(slots, _cuda([8] * 5)) + torch.cuda.synchronize() assert worker._ctx_len.tolist() == [40 + 8 * step, 300 + 8 * step, 7 + 8 * step, 0, 0], step - assert worker._batch_to_slot[:5].tolist() == [0, 1, 2, 4, 4] - assert worker._ctx_len_host[4] == 0 - worker._ctx_len[[0, 1, 2]] += 8 - worker._ctx_len[4] += 2 * 8 # both padding rows land on the dummy slot -def test_evicted_request_and_dummy_reset_together(): +def test_dummy_slot_reset_replays_in_a_cuda_graph(): + worker = _worker([40, 300, 7, 0, 0]) + slots = _cuda([0, 1, 4]) + accepted = _cuda([2, 3, 5]) + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + worker._advance_ctx_len(slots, accepted) + torch.cuda.current_stream().wait_stream(side) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + worker._advance_ctx_len(slots, accepted) + worker._ctx_len.copy_(_cuda([40, 300, 7, 0, 0])) + for step in range(1, 4): + worker._ctx_len[worker._dummy_slot] = 99 # whatever the slot holds before the step + graph.replay() + torch.cuda.synchronize() + assert worker._ctx_len.tolist() == [40 + 2 * step, 300 + 3 * step, 7, 0, 0], step + + +def test_prepare_resets_only_evicted_slots(): worker = _worker([40, 300, 7, 0, 96]) - _prepare(worker, [11, 13, FLOOR + 1], num_generations=3) # request 12 left; one padding row - assert worker._ctx_len.tolist() == [40, 0, 7, 0, 0] + written = [] + write = worker._write_ctx_len + + def record(updates): + written.append(dict(updates)) + write(updates) + + worker._write_ctx_len = record + _prepare(worker, [11, 12, 13, FLOOR + 1], num_generations=4) + _prepare(worker, [11, 13, FLOOR + 1], num_generations=3) # request 12 left + assert written == [{}, {1: 0}] + assert worker._ctx_len.tolist() == [40, 0, 7, 0, 96] assert 1 in worker._free_slots and 12 not in worker._req_to_slot From ce1ae633f172a795b85d81c692675b2ba0bda983 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 05:34:48 -0700 Subject: [PATCH 37/67] [None][feat] KDA: commit pending accepted drafts into the live pools The fused KDA verify leaves the state after the last golden token in the SSM pool, the drafts the sampler accepted pending in per-slot replay caches, and the conv window only in those caches. commit_kda_pending_drafts moves the drafts and the window into the live pools of one layer, so that the pools hold exactly what the next verify would start from: - kda_decode_mtp_kernel gains a REPLAY_ONLY compile-time mode that runs only the replay steps of a verify and writes the state after them. The replay is the verify's own code, so the result is bit-identical to the state the next verify reaches before its golden token. The mode reads no token inputs or weights; kda_mtp_commit_pending_drafts launches it. - With per-token verify states (kda_state_tok), the accepted drafts are replayed from the per-draft records the Kimi K3 verify kernels keep there (row innovations, beta * k and the decay), with the kernels' arithmetic (replay_kda_token_records). - store_kda_replay_conv_window stores the W - 1 conv columns starting at the accepted count into the conv pool. They are raw bf16 inputs, so the bf16 pool holds them exactly. The default kernel variant is unchanged: the new code is behind const_expr gates and the new parameter is the last one. The new tests check that committing the drafts, storing the window in a bf16 pool and seeding the replay caches from it before a verify gives that verify's outputs, state, caches and windows bit for bit, and the replay of the per-token records. Signed-off-by: Vasanth Sabavat --- .../cute_dsl_kimi_k3_kda_mtp_ops.py | 104 +++++++- .../blackwell/kimi_k3_kda/kda_mtp_decode.py | 68 ++++-- .../_torch/modules/kimi_kda/_kda_kernels.py | 222 ++++++++++++++++++ .../test_kda_mtp_decode_cute_parity.py | 161 +++++++++++++ 4 files changed, 532 insertions(+), 23 deletions(-) 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 14543b1d0868..98e4a1d434fa 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 @@ -124,6 +124,7 @@ def _run_kda_decode_mtp( FUSE_PRECOMPUTE: cutlass.Constexpr[bool], RUNTIME_PRECOMPUTE_FLAG: cutlass.Constexpr[bool], PROFILE_STAGES: cutlass.Constexpr[bool], + REPLAY_ONLY: cutlass.Constexpr[bool], stream: cuda.CUstream, ): if cutlass.const_expr(USE_ZERO_ACCEPTED): @@ -173,6 +174,7 @@ def _run_kda_decode_mtp( RUNTIME_PRECOMPUTE_FLAG, stage_timing, PROFILE_STAGES, + REPLAY_ONLY, ).launch(grid=(HV, N, 1), block=[NUM_THREADS, 1, 1], stream=stream) @@ -451,6 +453,7 @@ def kda_mtp_decode_impl( out: Optional[torch.Tensor] = None, zero_accepted_hint: bool = False, regular_metadata_hint: bool = False, + replay_only: bool = False, ) -> torch.Tensor: """Launch the fused KDA MTP verify kernel. See the module docstring. @@ -477,6 +480,9 @@ def kda_mtp_decode_impl( regular_metadata_hint: caller asserts ``cu_seqlens`` is the uniform ``arange * (2*num_spec+1)`` pattern and ``ssm_state_indices`` is ``arange(N)`` (benchmark identity layout). + replay_only: run only the replay steps and commit after them (see + ``kda_mtp_commit_pending_drafts``); the token inputs, weights + and ``out`` are not accessed. Returns the output ``[1, T_total, H, V]`` bf16 (rows for replayed positions are zero; only new-token rows are written). @@ -545,7 +551,9 @@ def kda_mtp_decode_impl( V=V_dim, ) pool_size = h0_arg.shape[0] - is_benchmark_static_shape = _is_benchmark_static_shape(N, H, HV, K, V_dim, W, num_spec) + is_benchmark_static_shape = ( + _is_benchmark_static_shape(N, H, HV, K, V_dim, W, num_spec) and not replay_only + ) use_setmaxreg = is_benchmark_static_shape use_reg_q_weights = is_benchmark_static_shape use_regular_metadata = bool(regular_metadata_hint) @@ -555,6 +563,8 @@ def kda_mtp_decode_impl( # num_accepted_tokens is all zeros, so the generic path computes the # same result). use_zero_accepted = bool(zero_accepted_hint) and num_spec == 2 + if replay_only and (use_zero_accepted or use_regular_metadata): + raise ValueError("replay_only takes neither zero_accepted_hint nor regular_metadata_hint") # stage_timing is unused (PROFILE_STAGES=False); pass `out` as the # placeholder tensor argument like the drop's runner does. The alias is # only valid while profiling stays off: with PROFILE_STAGES=True the @@ -605,6 +615,7 @@ def kda_mtp_decode_impl( use_regular_metadata, use_reg_q_weights, use_zero_accepted, + replay_only, ) if key not in _compiled_cache: @@ -612,7 +623,7 @@ def kda_mtp_decode_impl( f"kda_mtp_decode: compiling variant N={N} H={HV} T={T_total} " f"num_spec={num_spec} zero_accepted={use_zero_accepted} " f"regular_metadata={use_regular_metadata} " - f"static_shape={is_benchmark_static_shape}" + f"static_shape={is_benchmark_static_shape} replay_only={replay_only}" ) _compiled_cache[key] = cute.compile( _run_kda_decode_mtp, @@ -657,6 +668,7 @@ def kda_mtp_decode_impl( FUSE_PRECOMPUTE=True, RUNTIME_PRECOMPUTE_FLAG=False, PROFILE_STAGES=profile_stages, + REPLAY_ONLY=replay_only, stream=stream, ) @@ -692,6 +704,94 @@ def kda_mtp_decode_impl( return out +# (device_index, N, H, K, V, W) -> stand-ins for the inputs the replay-only +# variant never reads. Keys are plain values, like _precompute_control_cache. +_replay_only_inputs_cache = {} + + +def _replay_only_inputs(device: torch.device, N: int, H: int, K: int, V: int, W: int) -> dict: + dev = torch.device(device) + key = (dev.index if dev.index is not None else torch.cuda.current_device(), N, H, K, V, W) + if key not in _replay_only_inputs_cache: + + def unused(*shape, dtype=torch.bfloat16): + return torch.empty(shape, dtype=dtype, device=dev) + + _replay_only_inputs_cache[key] = { + "x_q": unused(1, N, H, K), + "x_k": unused(1, N, H, K), + "x_v": unused(1, N, H, V), + "g": unused(1, N, H, K), + "beta": unused(1, N, H), + "w_q": unused(H * K, W, dtype=torch.float32), + "w_k": unused(H * K, W, dtype=torch.float32), + "w_v": unused(H * V, W, dtype=torch.float32), + "A_log": unused(H, dtype=torch.float32), + "dt_bias": unused(H * K, dtype=torch.float32), + "out": unused(1, N, H, V), + # One row per request keeps every request's [bos, eos) non-empty. + "cu_seqlens": torch.arange(N + 1, dtype=torch.int32, device=dev), + } + return _replay_only_inputs_cache[key] + + +def kda_mtp_commit_pending_drafts( + recurrent_state: torch.Tensor, + cs_q: torch.Tensor, + cs_k: torch.Tensor, + cs_v: torch.Tensor, + qkg_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + ssm_state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + num_spec: int, + conv_width: int, +) -> None: + """Commit the pending accepted drafts of the selected slots in place. + + For request ``n`` this runs the first ``num_accepted_tokens[n]`` steps of + a ``kda_mtp_decode`` verify of slot ``ssm_state_indices[n]``, the replay + of its cached drafts, with the same kernel code, and then commits the + recurrent state and the base conv windows. The slot then holds what the + next verify would hold just before its golden token, with nothing left + pending; the caller sets the slot's accepted-draft count to 0. The + tensors are those of ``kda_mtp_decode``; the draft caches are only read. + """ + N = ssm_state_indices.shape[0] + if N == 0: + return + _, H, V, K = recurrent_state.shape + inputs = _replay_only_inputs(recurrent_state.device, N, H, K, V, conv_width) + kda_mtp_decode_impl( + x_q=inputs["x_q"], + x_k=inputs["x_k"], + x_v=inputs["x_v"], + w_q=inputs["w_q"], + w_k=inputs["w_k"], + w_v=inputs["w_v"], + cs_q=cs_q, + cs_k=cs_k, + cs_v=cs_v, + g=inputs["g"], + beta=inputs["beta"], + A_log=inputs["A_log"], + dt_bias=inputs["dt_bias"], + recurrent_state=recurrent_state, + qkg_cache=qkg_cache, + v_cache=v_cache, + beta_cache=beta_cache, + ssm_state_indices=ssm_state_indices, + cu_seqlens=inputs["cu_seqlens"], + num_spec=num_spec, + num_accepted_tokens=num_accepted_tokens, + # Only the new-token steps use the gate bound and the q scale. + lower_bound=0.0, + out=inputs["out"], + replay_only=True, + ) + + if IS_CUTLASS_DSL_AVAILABLE: @torch.library.custom_op( diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py index e00a1ecd670e..822b7f1e5097 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/kimi_k3_kda/kda_mtp_decode.py @@ -108,6 +108,7 @@ def kda_decode_mtp_kernel( RUNTIME_PRECOMPUTE_FLAG: cutlass.Constexpr[bool], stage_timing: cute.Tensor, PROFILE_STAGES: cutlass.Constexpr[bool], + REPLAY_ONLY: cutlass.Constexpr[bool], ): """KDA MTP decode — SMEM pre-compute + register-resident state. @@ -115,6 +116,12 @@ def kda_decode_mtp_kernel( with at least ``HV * N * 4`` elements, indexed as ``(i_hv * grid_n + i_n) * 4``. With profiling off it is never accessed and the host may pass any placeholder tensor. + + With ``REPLAY_ONLY=True`` the kernel runs only the replay steps: it + replays the ``num_accepted_tokens[n]`` pending drafts and commits the + recurrent state after them to ``ht``. It reads only ``h0``, the conv and + replay caches and the metadata, and writes only ``ht``; ``x_*``, ``w_*``, + ``g``, ``beta``, ``A_log``, ``dt_bias`` and ``o`` are not accessed. """ tidx, _, _ = cute.arch.thread_idx() in_warp_tid = tidx % 32 @@ -156,7 +163,10 @@ def kda_decode_mtp_kernel( T_loop = 1 + NUM_SPEC t_max = 1 + NUM_SPEC else: - T_loop = commit_len + 1 + NUM_SPEC + if cutlass.const_expr(REPLAY_ONLY): + T_loop = commit_len + else: + T_loop = commit_len + 1 + NUM_SPEC t_max = 2 * NUM_SPEC + 1 vec_size = TILE_K // 32 num_v_tiles = V // TILE_V @@ -211,35 +221,39 @@ def kda_decode_mtp_kernel( r_state[(KERNEL_WIDTH - 1) * vec_size + w * vec_size + i] = cutlass.Float32( cs_k[slot, hk_off + k_idx, w] ) - for w in range(KERNEL_WIDTH): - if tidx < K: - if cutlass.const_expr(not USE_REG_Q_WEIGHTS): - sConvW[w * K + tidx] = cutlass.Float32(w_q[hk_off + tidx, w]) - sConvW[k_weight_base + w * K + tidx] = cutlass.Float32( - w_k[hk_off + tidx, w] + if cutlass.const_expr(not REPLAY_ONLY): + for w in range(KERNEL_WIDTH): + if tidx < K: + if cutlass.const_expr(not USE_REG_Q_WEIGHTS): + sConvW[w * K + tidx] = cutlass.Float32(w_q[hk_off + tidx, w]) + sConvW[k_weight_base + w * K + tidx] = cutlass.Float32( + w_k[hk_off + tidx, w] + ) + for ld in range(V * KERNEL_WIDTH // NUM_THREADS): + flat = ld * NUM_THREADS + tidx + sConvW[v_weight_base + flat] = cutlass.Float32( + w_v[hv_off + flat % V, flat // V] ) - for ld in range(V * KERNEL_WIDTH // NUM_THREADS): - flat = ld * NUM_THREADS + tidx - sConvW[v_weight_base + flat] = cutlass.Float32( - w_v[hv_off + flat % V, flat // V] - ) - if cutlass.const_expr(USE_REG_Q_WEIGHTS): - if warp_idx == 0: - for _w in range(KERNEL_WIDTH): - for _i in range(vec_size): - r_wq[_w * vec_size + _i] = cutlass.Float32( - w_q[hk_off + _i * 32 + in_warp_tid, _w] - ) + if cutlass.const_expr(USE_REG_Q_WEIGHTS): + if warp_idx == 0: + for _w in range(KERNEL_WIDTH): + for _i in range(vec_size): + r_wq[_w * vec_size + _i] = cutlass.Float32( + w_q[hk_off + _i * 32 + in_warp_tid, _w] + ) cute.arch.barrier() if cutlass.const_expr(USE_SETMAXREG): cute.arch.warpgroup_reg_dealloc(64) if warp_idx < 3: - if warp_idx == 2: - r_exp_A = cute.math.exp(cutlass.Float32(A_log[i_h]), fastmath=True) + if cutlass.const_expr(not REPLAY_ONLY): + if warp_idx == 2: + r_exp_A = cute.math.exp(cutlass.Float32(A_log[i_h]), fastmath=True) i_t = 0 while i_t < T_loop: if cutlass.const_expr(USE_ZERO_ACCEPTED): replay_from_cache = False + elif cutlass.const_expr(REPLAY_ONLY): + replay_from_cache = True else: replay_from_cache = i_t < commit_len if replay_from_cache: @@ -772,6 +786,18 @@ def kda_decode_mtp_kernel( row * vec_size + i ] i_t = i_t + 1 + if cutlass.const_expr(REPLAY_ONLY): + for row in range(NUM_V_ROWS): + v_row = warp_idx * NUM_V_ROWS + row + for i in range(vec_size): + if cutlass.const_expr(USE_FLAT_LAYOUT): + ht[h0_idx, v_base + v_row, i * 32 + in_warp_tid] = r_state[ + row * vec_size + i + ] + else: + ht[slot, i_hv, v_base + v_row, i * 32 + in_warp_tid] = r_state[ + row * vec_size + i + ] if cutlass.const_expr(PROFILE_STAGES): cute.arch.barrier() t_stage2 = read_globaltimer() diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index 90b3e4f1030d..c694b90eba1e 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -264,6 +264,228 @@ def copy_kda_replay_conv_window( ) +@triton.jit +def _store_kda_replay_conv_window_kernel( + conv_ptr, + q_cache_ptr, + k_cache_ptr, + v_cache_ptr, + state_indices_ptr, + num_accepted_ptr, + conv_stride_slot, + conv_stride_dim, + conv_stride_window, + cache_stride_slot, + cache_stride_dim, + cache_stride_window, + PROJECTION_SIZE: tl.constexpr, + COMMITTED: tl.constexpr, + NUM_SPEC: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + request = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < PROJECTION_SIZE * COMMITTED + dim = offsets // COMMITTED + window = offsets % COMMITTED + slot = tl.load(state_indices_ptr + request).to(tl.int64) + # The window after the accepted drafts starts at their count, clamped + # like the verify kernels clamp it. + shift = tl.minimum(tl.maximum(tl.load(num_accepted_ptr + request), 0), NUM_SPEC) + + conv_offset = slot * conv_stride_slot + dim * conv_stride_dim + window * conv_stride_window + cache_offset = ( + slot * cache_stride_slot + dim * cache_stride_dim + (window + shift) * cache_stride_window + ) + section_offset = PROJECTION_SIZE * conv_stride_dim + tl.store(conv_ptr + conv_offset, tl.load(q_cache_ptr + cache_offset, mask=mask), mask=mask) + tl.store( + conv_ptr + conv_offset + section_offset, + tl.load(k_cache_ptr + cache_offset, mask=mask), + mask=mask, + ) + tl.store( + conv_ptr + conv_offset + 2 * section_offset, + tl.load(v_cache_ptr + cache_offset, mask=mask), + mask=mask, + ) + + +def store_kda_replay_conv_window( + conv_pool: torch.Tensor, + q_cache: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, +) -> None: + """Store each selected slot's conv window after its accepted drafts in + the live convolution pool. + + The replay caches hold the window before the golden token in columns + ``[0, W - 1)`` and the raw inputs of the round's drafts after it, so the + ``W - 1`` columns starting at the request's ``num_accepted_tokens`` entry + are the window the next verify would start from. They are raw + convolution inputs, so storing them in a bf16 pool is exact. + """ + projection_size = conv_pool.shape[1] // 3 + committed = conv_pool.shape[2] + caches = (q_cache, k_cache, v_cache) + if any(cache.shape[1] != projection_size or cache.shape[2] < committed for cache in caches): + raise ValueError("KDA replay convolution caches do not match the live pool geometry") + if any(cache.stride() != q_cache.stride() for cache in caches[1:]): + raise ValueError("KDA replay convolution caches must share one layout") + if num_accepted_tokens.shape != state_indices.shape: + raise ValueError("Expected one accepted-draft count per state index") + if state_indices.numel() == 0: + return + + block_size = 256 + grid = ( + state_indices.numel(), + triton.cdiv(projection_size * committed, block_size), + ) + with torch.cuda.device(conv_pool.device.index): + _store_kda_replay_conv_window_kernel[grid]( + conv_pool, + q_cache, + k_cache, + v_cache, + state_indices, + num_accepted_tokens, + conv_pool.stride(0), + conv_pool.stride(1), + conv_pool.stride(2), + q_cache.stride(0), + q_cache.stride(1), + q_cache.stride(2), + PROJECTION_SIZE=projection_size, + COMMITTED=committed, + NUM_SPEC=q_cache.shape[2] - committed, + BLOCK_SIZE=block_size, + ) + + +@triton.jit +def _replay_kda_token_records_kernel( + state_ptr, + records_ptr, + state_indices_ptr, + num_accepted_ptr, + state_stride_slot, + records_stride_slot, + NUM_SPEC: tl.constexpr, + H: tl.constexpr, + V: tl.constexpr, + K: tl.constexpr, + BLOCK_V: tl.constexpr, +): + request = tl.program_id(0) + head = tl.program_id(1) + rows = tl.program_id(2) * BLOCK_V + tl.arange(0, BLOCK_V) + keys = tl.arange(0, K) + slot = tl.load(state_indices_ptr + request).to(tl.int64) + accepted = tl.minimum(tl.maximum(tl.load(num_accepted_ptr + request), 0), NUM_SPEC) + state_ptrs = ( + state_ptr + slot * state_stride_slot + (head * V + rows[:, None]) * K + keys[None, :] + ) + state = tl.load(state_ptrs) + records = records_ptr + slot * records_stride_slot + for t in range(0, accepted): + vn = tl.load(records + (t * H + head) * V + rows) + beta_k = tl.load(records + NUM_SPEC * H * V + (t * H + head) * K + keys) + decay = tl.load(records + NUM_SPEC * H * (V + K) + (t * H + head) * K + keys) + state = tl.fma(decay[None, :], state, vn[:, None] * beta_k[None, :]) + tl.store(state_ptrs, state) + + +def replay_kda_token_records( + recurrent_state: torch.Tensor, + state_tok: torch.Tensor, + state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, +) -> None: + """Replay each selected slot's accepted drafts onto its pool state from + the per-draft records the Kimi K3 verify kernels keep in the slot's + per-token state region. + + A slot's region of ``state_tok`` (``[slots, num_spec, H, V, K]`` fp32) + starts with the records of the last round's drafts: the row innovations + ``vn`` ``[num_spec][H][V]``, then ``beta * k`` and the decay + ``[num_spec][H][K]``. Draft ``t`` updates the state as the kernels do, + ``S = fma(decay, S, vn (beta * k))``, so the result is bit-identical to the + state their next verify starts from. + """ + num_spec, num_heads, v_dim, k_dim = state_tok.shape[1:] + if state_indices.numel() == 0: + return + block_v = 32 + grid = (state_indices.numel(), num_heads, v_dim // block_v) + with torch.cuda.device(recurrent_state.device.index): + _replay_kda_token_records_kernel[grid]( + recurrent_state, + state_tok, + state_indices, + num_accepted_tokens, + recurrent_state.stride(0), + state_tok.stride(0), + NUM_SPEC=num_spec, + H=num_heads, + V=v_dim, + K=k_dim, + BLOCK_V=block_v, + ) + + +def commit_kda_pending_drafts( + recurrent_state: torch.Tensor, + conv_pool: torch.Tensor, + conv_q: torch.Tensor, + conv_k: torch.Tensor, + conv_v: torch.Tensor, + qkg_cache: torch.Tensor, + v_cache: torch.Tensor, + beta_cache: torch.Tensor, + state_indices: torch.Tensor, + num_accepted_tokens: torch.Tensor, + state_tok: Optional[torch.Tensor] = None, +) -> None: + """Commit one layer's pending accepted drafts into its live pools. + + ``num_accepted_tokens[n]`` is the number of drafts of slot + ``state_indices[n]`` the last verify accepted, which the next verify + would take up first. Afterwards ``recurrent_state`` and ``conv_pool`` + hold the state and the conv window that verify would start from, bit for + bit, and the caller clears the slot's count: + + * with per-token states (``state_tok``), the drafts are replayed from the + per-draft records the Kimi K3 verify kernels keep there + (``replay_kda_token_records``); + * otherwise they are replayed from the replay caches by the fused verify + kernel's own replay steps (``kda_mtp_commit_pending_drafts``). + """ + num_spec = conv_q.shape[2] - conv_pool.shape[2] + if state_tok is not None: + replay_kda_token_records(recurrent_state, state_tok, state_indices, num_accepted_tokens) + else: + _load_mtp_module().kda_mtp_commit_pending_drafts( + recurrent_state, + conv_q, + conv_k, + conv_v, + qkg_cache, + v_cache, + beta_cache, + state_indices, + num_accepted_tokens, + num_spec, + conv_width=conv_pool.shape[2] + 1, + ) + store_kda_replay_conv_window( + conv_pool, conv_q, conv_k, conv_v, state_indices, num_accepted_tokens + ) + + # --------------------------------------------------------------------------- # In-tree KDA prefill op (CuTe DSL, trtllm::kda_prefill). # --------------------------------------------------------------------------- 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 64747b380a42..c909afad0c83 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 @@ -661,6 +661,167 @@ def test_beta_cache_sibling_layer_after_aligned_warmup( ) +def _dim_contiguous_clone(cache): + clone = torch.empty( + cache.shape[0], cache.shape[2], cache.shape[1], dtype=cache.dtype, device=cache.device + ).transpose(1, 2) + clone.copy_(cache) + return clone + + +def _verify_after_round_one(B, H, num_spec, seed): + """Inputs of a verify that follows a first verify run by the kernel, with + a different number of the first round's drafts accepted per request. + + The conv windows start at bf16 values, as they do when the cache manager + seeds them from its bf16 conv pool.""" + data = make_conv_data(B, H, M=num_spec, seed=seed) + for name in ("cs_q", "cs_k", "cs_v"): + data[name].copy_(data[name].bfloat16().float()) + first = cute_run(data) + data = make_conv_data(B, H, M=num_spec, seed=seed + 1) + for name in ("qkg_cache", "v_cache", "beta_cache", "cs_q", "cs_k", "cs_v"): + data[name] = first[name] + data["initial_state_cute"] = first["state_v_first"] + data["num_accepted_tokens"] = torch.arange(B, dtype=torch.int32, device="cuda") % (num_spec + 1) + return data + + +@pytest.mark.parametrize("B,H,num_spec", [(32, 12, 2), (32, 6, 2), (4, 6, 7)], ids=lambda v: str(v)) +def test_commit_pending_drafts_matches_the_next_verify(B, H, num_spec): + """Committing the pending drafts first leaves the next verify's results + unchanged, bit for bit. + + This is what the cache manager does across a suspend. + ``kda_mtp_commit_pending_drafts`` folds the accepted drafts into the + state, ``store_kda_replay_conv_window`` puts the conv window after them in + the bf16 conv pool, and the resume seeds the replay caches from that pool + with nothing pending. The verify that follows must give exactly the + outputs, state, caches and windows of the verify that replays the drafts + itself.""" + from tensorrt_llm._torch.custom_ops.cute_dsl_kimi_k3_kda_mtp_ops import ( + kda_mtp_commit_pending_drafts, + ) + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import ( + copy_kda_replay_conv_window, + store_kda_replay_conv_window, + ) + + data = _verify_after_round_one(B, H, num_spec, seed=41) + accepted = data["num_accepted_tokens"] + slots = data["ssm_state_indices"] + replayed = cute_run(data) + + state = data["initial_state_cute"].clone() + cs = [_dim_contiguous_clone(data[name]) for name in ("cs_q", "cs_k", "cs_v")] + kda_mtp_commit_pending_drafts( + state, + *cs, + data["qkg_cache"].clone(), + data["v_cache"].clone(), + data["beta_cache"].clone(), + slots, + accepted, + num_spec, + conv_width=data["W"], + ) + for n in range(B): + unchanged = torch.equal(state[n], data["initial_state_cute"][n]) + assert unchanged == (int(accepted[n]) == 0), n + conv_pool = torch.zeros( + B, 3 * H * data["K"], data["W"] - 1, dtype=torch.bfloat16, device="cuda" + ) + store_kda_replay_conv_window(conv_pool, *cs, slots, accepted) + for cache in cs: + cache.zero_() + copy_kda_replay_conv_window(conv_pool, *cs, slots) + + resumed = dict(data) + resumed.update( + initial_state_cute=state, + cs_q=cs[0], + cs_k=cs[1], + cs_v=cs[2], + qkg_cache=torch.zeros_like(data["qkg_cache"]), + v_cache=torch.zeros_like(data["v_cache"]), + beta_cache=torch.zeros_like(data["beta_cache"]), + num_accepted_tokens=torch.zeros_like(accepted), + # The same token rows: the verify reads request n's new tokens from + # cu_seqlens[n] + num_accepted_tokens[n]. + cu_seqlens=data["cu_seqlens"] + torch.cat([accepted, accepted.new_zeros(1)]), + ) + resumed = cute_run(resumed) + + for name in ( + "out", + "state_v_first", + "qkg_cache", + "v_cache", + "beta_cache", + "cs_q", + "cs_k", + "cs_v", + ): + assert torch.equal(resumed[name], replayed[name]), name + + +def test_commit_pending_drafts_replays_the_per_token_records(): + """With per-token verify states, a slot's accepted drafts are replayed onto + its pool state, in order, from the per-draft records the Kimi K3 verify + kernels keep in its region (vn, beta * k, decay), and the conv window after + them is stored; other slots are left alone.""" + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import commit_kda_pending_drafts + + torch.manual_seed(5) + num_slots, H, K, W, num_spec = 9, 6, 128, 4, 7 + + def small(*shape, scale): + # Few significant bits: every product and sum below is exact in fp32, + # so the reference does not depend on fused or separate rounding. + return torch.randint(-4, 5, shape, device="cuda").float() * scale + + state = small(num_slots, H, K, K, scale=2.0**-4) + vn = small(num_slots, num_spec, H, K, scale=2.0**-2) + beta_k = small(num_slots, num_spec, H, K, scale=2.0**-3) + decay = torch.randint(1, 3, (num_slots, num_spec, H, K), device="cuda").float() * 0.5 + state_tok = torch.zeros(num_slots, num_spec, H, K, K, device="cuda") + records = state_tok.view(num_slots, -1) + size = num_spec * H * K + for i, part in enumerate((vn, beta_k, decay)): + records[:, i * size : (i + 1) * size] = part.reshape(num_slots, -1) + cs = [ + torch.randn(num_slots, W - 1 + num_spec, H * K, device="cuda") + .bfloat16() + .float() + .transpose(1, 2) + for _ in range(3) + ] + conv_pool = torch.zeros(num_slots, 3 * H * K, W - 1, dtype=torch.bfloat16, device="cuda") + slots = torch.tensor([7, 0, 3, 5, 1, 8], dtype=torch.int32, device="cuda") + accepted = torch.tensor([0, 1, 3, 7, 2, 5], dtype=torch.int32, device="cuda") + before = state.clone() + unread = torch.empty(0, device="cuda") + + commit_kda_pending_drafts( + state, conv_pool, *cs, unread, unread, unread, slots, accepted, state_tok=state_tok + ) + + for slot, count in zip(slots.tolist(), accepted.tolist()): + expected = before[slot].clone() + for t in range(count): + expected = ( + decay[slot, t][:, None, :] * expected + + vn[slot, t][:, :, None] * beta_k[slot, t][:, None, :] + ) + assert torch.equal(state[slot], expected), slot + for section, cache in enumerate(cs): + window = conv_pool[slot, section * H * K : (section + 1) * H * K].float() + assert torch.equal(window, cache[slot, :, count : count + W - 1]), (slot, section) + others = [slot for slot in range(num_slots) if slot not in slots.tolist()] + assert torch.equal(state[others], before[others]) + assert not conv_pool[others].any() + + if __name__ == "__main__": import sys From f9d32b4c4d9937f191498d87c0d635847bea2a6f Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 05:34:50 -0700 Subject: [PATCH 38/67] [None][fix] KV cache manager V2: commit a suspended request's KDA replay state into V2 A suspend now leaves the request's whole KDA state in its V2 SSM and conv states instead of a host copy of the replay caches. Before the pages are released, suspend_request commits the slot's pending accepted drafts and its conv window into them (commit_kda_pending_drafts) and clears the draft count; the replay caches hold nothing the request needs any more, and no extra memory is held while it is suspended. The request's next generation step seeds the replay caches of whichever slot it resumes in from the V2 conv state, as a disaggregated generation request's are. Both run on the KV caches' stream: the commit after the current stream's work and before the event the suspend records there, which later page migrations wait for; the seeding after the pages a resume brings back. With a host drafter, update_resources records a batch's acceptance, which the commit needs. Executors record it before they can suspend the request: the host drafter runs without the overlap scheduler, and the PP loop does not evict requests in flight. A suspend between a batch's prepare_resources and its update_resources raises instead of committing a stale count. The host copies and the late acceptance hand-off are removed. The tests resume a request in a new slot and in its old one, check the host drafter guard, and on the V2 runtime with the fused verify kernel compare every verify after a suspend and resume with the run that is never suspended, bit for bit, with the slot taken by another cache meanwhile and without. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 184 ++++++++---- .../kv_cache/test_mamba_cache_manager.py | 269 ++++++++++++++---- 2 files changed, 336 insertions(+), 117 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index bd299fb07164..3be8d84f564a 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -16,6 +16,7 @@ import math import os from abc import ABC, abstractmethod +from contextlib import contextmanager from dataclasses import dataclass, replace from typing import (TYPE_CHECKING, Dict, Iterable, List, Literal, NamedTuple, Optional, Protocol, Sequence, Tuple, Union, cast) @@ -3235,10 +3236,12 @@ def __init__( } self._request_id_to_state_index = {} self._request_id_to_is_dummy = {} - # KDA replay state of suspended requests (suspend_request), and the - # draft counts update_resources accepts for them after their suspend. - self._suspended_kda_replay_state: Dict[int, List[torch.Tensor]] = {} - self._suspended_kda_replay_accepted: Dict[int, int] = {} + # Requests whose suspend committed their KDA replay state into the V2 + # states; their next generation step re-seeds the replay caches. + self._kda_replay_suspended: set[int] = set() + # With a host drafter: requests of a prepared batch whose acceptance + # update_resources has not recorded yet. + self._kda_replay_acceptance_pending: set[int] = set() state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) @@ -4318,38 +4321,88 @@ def _relocate_kda_replay_slots(self, old_slots: List[int], buffer.index_select(dim, source_slots)) def suspend_request(self, req: LlmRequest) -> None: - """Suspend the request's KV cache and keep its KDA replay state. - - The replay state is indexed by SSM slot, outside the V2 pools. While - the request is suspended, its pages can move to a lower tier and its - slot can go to another request; it resumes in whichever slot is free. - The slot's replay state is copied out here, and the request's next - ``_setup_state_indices`` writes it into the slot it resumes in. + """Suspend the request's KV cache with its whole KDA state in V2. + + The KDA replay state is indexed by SSM slot, outside the V2 pools: the + drafts the last verify accepted are pending there, and so is the conv + window, which the fused verify does not write to the V2 conv state. + While the request is suspended, its pages can move to a lower tier + and its slot can go to another request. So the pending drafts and the + conv window are committed into the slot's V2 states here, before the + pages are released, and the request's next generation step re-seeds + the replay caches of whichever slot it resumes in from them. """ request_id = req.py_request_id kv_cache = self.kv_cache_map.get(request_id) slot = self._request_id_to_state_index.get(request_id, -1) if (self.use_kda_replay_update and slot >= 0 and kv_cache is not None and kv_cache.is_active): - saved = [] - for buffer, dim in self._kda_replay_state_buffers(): - state = buffer.select(dim, slot) - host = torch.empty(state.shape, - dtype=state.dtype, - pin_memory=prefer_pinned()) - host.copy_(state, non_blocking=True) - saved.append(host) - self._suspended_kda_replay_state[request_id] = saved + # A host drafter's acceptance is recorded by update_resources. + # Executors record it before they can suspend the request: the + # host drafter runs without the overlap scheduler, and the PP + # loop does not evict requests in flight. + if request_id in self._kda_replay_acceptance_pending: + raise RuntimeError( + f"Request {request_id} is being suspended before " + "update_resources recorded the acceptance of its last " + "verify, so its pending KDA drafts are unknown") + self._commit_kda_pending_drafts(slot) del self._request_id_to_state_index[request_id] + self._kda_replay_suspended.add(request_id) super().suspend_request(req) - def _restore_kda_replay_state(self, request_id: int, slot: int, - saved: List[torch.Tensor]) -> None: - for (buffer, dim), host in zip(self._kda_replay_state_buffers(), saved): - buffer.select(dim, slot).copy_(host, non_blocking=True) - accepted = self._suspended_kda_replay_accepted.pop(request_id, None) - if accepted is not None: - self.prev_num_accepted_tokens[slot] = accepted + @contextmanager + def _on_kv_cache_stream(self): + """Run the enclosed work on the KV caches' stream, after the current + stream's earlier work and before its later work. + + V2 orders page migration against that stream: a suspend records the + event later migrations of the pages wait for, and a resume makes the + stream wait for the pages it brings back. + """ + current = torch.cuda.current_stream() + self._stream.wait_stream(current) + with torch.cuda.stream(self._stream): + yield + current.wait_stream(self._stream) + + def _commit_kda_pending_drafts(self, slot: int) -> None: + """Commit the slot's pending accepted drafts and its conv window into + its V2 SSM and conv states, in place, and clear its draft count. + + Afterwards the V2 states hold exactly what the next verify would have + started from (``commit_kda_pending_drafts``). + """ + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import \ + commit_kda_pending_drafts + + assert self.prev_num_accepted_tokens is not None + # With per-token verify states, the accepted drafts are replayed from + # the records the verify kernels keep there instead of the caches. + state_tok = getattr(self, "kda_state_tok", None) + with self._on_kv_cache_stream(): + slots = torch.full((1, ), + slot, + dtype=torch.int32, + device=self.prev_num_accepted_tokens.device) + num_accepted = self.prev_num_accepted_tokens.index_select(0, slots) + for layer_offset, (ssm_states, conv_states) in enumerate( + zip(self.all_ssm_states, self.all_conv_states)): + commit_kda_pending_drafts( + ssm_states, + conv_states, + self.kda_conv_q[layer_offset], + self.kda_conv_k[layer_offset], + self.kda_conv_v[layer_offset], + self.kda_qkg_cache[layer_offset], + self.kda_v_cache[layer_offset], + self.kda_beta_cache[layer_offset], + slots, + num_accepted, + state_tok=(state_tok[layer_offset] + if state_tok is not None else None), + ) + self.prev_num_accepted_tokens[slot] = 0 def add_dummy_requests( self, @@ -4390,8 +4443,8 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self._request_id_to_state_index.pop(request.py_request_id, None) self._request_id_to_is_dummy.pop(request.py_request_id, None) self._branch_snapshot_points.pop(request.py_request_id, None) - self._suspended_kda_replay_state.pop(request.py_request_id, None) - self._suspended_kda_replay_accepted.pop(request.py_request_id, None) + self._kda_replay_suspended.discard(request.py_request_id) + self._kda_replay_acceptance_pending.discard(request.py_request_id) super().free_resources(request, pin_on_release) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -4403,6 +4456,17 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): num_contexts = len(scheduled_batch.context_requests) self._setup_state_indices(requests, num_contexts=num_contexts) self._reset_context_mamba_slots(num_contexts) + if self._record_kda_replay_in_update_resources: + # Recorded by this batch's update_resources, or dropped by + # revert_allocate_generation when the batch is not run. + self._kda_replay_acceptance_pending.update( + request.py_request_id + for request in scheduled_batch.generation_requests + if not request.is_dummy) + + def revert_allocate_generation(self, req: LlmRequest) -> None: + self._kda_replay_acceptance_pending.discard(req.py_request_id) + super().revert_allocate_generation(req) def _setup_state_indices(self, requests: List[LlmRequest], @@ -4447,32 +4511,33 @@ def _setup_state_indices(self, if self.use_kda_replay_update: # A suspended request that comes back as a context request starts - # over; one that comes back generating takes its replay state. + # over. One that comes back generating has its whole KDA state in + # the V2 states of the slot it resumes in (suspend_request). for request in requests[:num_contexts]: - self._suspended_kda_replay_state.pop(request.py_request_id, - None) - self._suspended_kda_replay_accepted.pop(request.py_request_id, - None) - saved_states = [ - self._suspended_kda_replay_state.pop(request.py_request_id, - None) - for request in replay_requests - ] + self._kda_replay_suspended.discard(request.py_request_id) new_state_values = [ self._request_id_to_state_index[request.py_request_id] for request in replay_requests ] self._relocate_kda_replay_slots(old_state_values, new_state_values) - self._reset_kda_replay_slots([ - new for old, new, saved in zip(old_state_values, - new_state_values, saved_states) - if old < 0 and saved is None - ]) - for request, new, saved in zip(replay_requests, new_state_values, - saved_states): - if saved is not None: - self._restore_kda_replay_state(request.py_request_id, new, - saved) + new_slots = [] + resumed_slots = [] + for request, old, new in zip(replay_requests, old_state_values, + new_state_values): + if request.py_request_id in self._kda_replay_suspended: + self._kda_replay_suspended.remove(request.py_request_id) + resumed_slots.append(new) + elif old < 0: + new_slots.append(new) + self._reset_kda_replay_slots(new_slots) + if resumed_slots: + resumed_indices = torch.tensor(resumed_slots, + dtype=torch.long, + pin_memory=prefer_pinned()) + with self._on_kv_cache_stream(): + self._seed_kda_replay_slots( + resumed_indices.to(self.prev_num_accepted_tokens.device, + non_blocking=True)) def _record_kda_replay_acceptance( self, @@ -4515,6 +4580,12 @@ def seed_kda_replay_caches_for_disagg_gen(self, state_indices = torch.tensor(sorted(set(slots)), dtype=torch.long, device=device) + self._seed_kda_replay_slots(state_indices) + + def _seed_kda_replay_slots(self, state_indices: torch.Tensor) -> None: + """Seed the replay caches of distinct slots from their V2 conv + states, with no drafts pending.""" + assert self.prev_num_accepted_tokens is not None committed_window = self.conv_state_shape[1] section_offsets = [0] for section_dim in self.conv_section_dims: @@ -4585,6 +4656,8 @@ def update_resources( return generation_requests = scheduled_batch.generation_requests + self._kda_replay_acceptance_pending.difference_update( + request.py_request_id for request in generation_requests) drafted_requests = [ request for request in generation_requests if request.py_draft_tokens is not None @@ -4597,15 +4670,6 @@ def update_resources( "Mixed drafted/undrafted generation batch is not supported " "for KDA replay bookkeeping") - for request in generation_requests: - # Suspended after this batch was scheduled: its saved replay state - # predates this acceptance, so the acceptance goes with it. - request_id = request.py_request_id - if (request_id in self._suspended_kda_replay_state and - not self._request_id_to_is_dummy.get(request_id, False)): - self._suspended_kda_replay_accepted[request_id] = max( - 0, request.py_num_accepted_draft_tokens) - active_requests = [ request for request in generation_requests if request.py_request_id in self._request_id_to_state_index @@ -4855,6 +4919,6 @@ def shutdown(self): self.kda_qkg_cache = None self.kda_v_cache = None self.kda_beta_cache = None - self._suspended_kda_replay_state.clear() - self._suspended_kda_replay_accepted.clear() + self._kda_replay_suspended.clear() + self._kda_replay_acceptance_pending.clear() super().shutdown() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 0b379e3a2877..d76aaf4a6fcb 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Regression tests for Python, Cpp, and V2 Mamba cache managers.""" +import contextlib import os from types import SimpleNamespace from unittest.mock import MagicMock @@ -3740,8 +3741,7 @@ def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): mgr._record_kda_replay_in_update_resources = True mgr._request_id_to_state_index = {101: 1, 303: 2} mgr._request_id_to_is_dummy = {101: False, 303: True} - mgr._suspended_kda_replay_state = {} - mgr._suspended_kda_replay_accepted = {} + mgr._kda_replay_acceptance_pending = {101, 202, 404} mgr.prev_num_accepted_tokens = torch.tensor([7, 8, 9], dtype=torch.int32) base_update = MagicMock() monkeypatch.setattr(KVCacheManagerV2, "update_resources", base_update) @@ -3766,6 +3766,8 @@ def test_v2_kda_replay_host_drafter_records_active_requests(monkeypatch): base_update.assert_called_once_with(scheduled_batch, None, None) assert mgr.prev_num_accepted_tokens.tolist() == [7, 2, 9] + # The batch's acceptances are recorded; another batch's are still pending. + assert mgr._kda_replay_acceptance_pending == {404} _KDA_REPLAY_BUFFERS = ( @@ -3793,15 +3795,18 @@ def suspend(self) -> None: def _v2_kda_replay_stub(num_slots=4): - """A V2 hybrid manager stub with CPU KDA replay buffers, for the state-index setup that prepare_resources runs.""" + """A V2 hybrid manager stub with CPU KDA replay buffers and one layer's V2 conv states ([Q | K | V] sections of 2 + channels, a committed window of 3), for the state-index setup that prepare_resources runs. Committing pending drafts + is recorded instead of run (it needs the GPU kernels).""" mgr = object.__new__(MambaHybridCacheManagerV2) mgr.local_num_mamba_layers = 1 mgr._use_kda_replay_update = True mgr._use_replay_state_update = False + mgr._record_kda_replay_in_update_resources = False mgr._request_id_to_state_index = {} mgr._request_id_to_is_dummy = {} - mgr._suspended_kda_replay_state = {} - mgr._suspended_kda_replay_accepted = {} + mgr._kda_replay_suspended = set() + mgr._kda_replay_acceptance_pending = set() mgr._host_state_indices = torch.zeros(num_slots, dtype=torch.int32) mgr.cuda_state_indices = torch.zeros(num_slots, dtype=torch.int32) mgr._dummy_request_mask = None @@ -3811,6 +3816,12 @@ def _v2_kda_replay_stub(num_slots=4): mgr.prev_num_accepted_tokens = torch.zeros(num_slots, dtype=torch.int32) for name in _KDA_REPLAY_BUFFERS: setattr(mgr, name, torch.zeros(1, num_slots, 2, 4)) + mgr.conv_state_shape = [6, 3] + mgr.conv_section_dims = [2, 2, 2] + mgr.all_conv_states = [torch.zeros(num_slots, 6, 3)] + mgr.committed_slots = [] + mgr._commit_kda_pending_drafts = mgr.committed_slots.append + mgr._on_kv_cache_stream = contextlib.nullcontext return mgr @@ -3828,9 +3839,11 @@ def _kda_replay_slot(mgr, slot): @pytest.mark.parametrize("resume_slot", [1, 0], ids=["new_slot", "old_slot"]) -def test_v2_kda_replay_state_follows_a_suspended_request(resume_slot): - """While a request is suspended its SSM page can move to a lower tier and its slot can go to another request. It - resumes in a new slot, or in its old one after that request has finished, with its own replay state either way.""" +def test_v2_kda_replay_suspended_request_resumes_from_its_v2_states(resume_slot): + """A suspend commits the request's pending drafts and conv window into its slot's V2 states and forgets the slot, + which can then go to another request. The request resumes in a new slot, or in its old one after that request has + finished. Either way its replay caches are seeded from the V2 conv state it resumes with, nothing is pending, and + the other request's replay state is left alone.""" mgr = _v2_kda_replay_stub() first = SimpleNamespace(py_request_id=1, is_dummy=False) second = SimpleNamespace(py_request_id=2, is_dummy=False) @@ -3838,6 +3851,8 @@ def test_v2_kda_replay_state_follows_a_suspended_request(resume_slot): mgr._setup_state_indices([first], num_contexts=0) _write_kda_replay_slot(mgr, 0, accepted=2, value=1.0) mgr.suspend_request(first) + assert mgr.committed_slots == [0] + assert 1 not in mgr._request_id_to_state_index assert not mgr.kv_cache_map[1].is_active mgr.kv_cache_map[2] = _SlotCache(0) @@ -3851,39 +3866,57 @@ def test_v2_kda_replay_state_follows_a_suspended_request(resume_slot): del mgr._request_id_to_state_index[2] else: resumed.append(second) + # The V2 conv state the runtime brings back with the request's page. + window = torch.arange(18.0).view(6, 3) + mgr.all_conv_states[0][resume_slot] = window mgr.kv_cache_map[1].slot = resume_slot mgr.kv_cache_map[1].is_active = True mgr._setup_state_indices(resumed, num_contexts=0) - assert _kda_replay_slot(mgr, resume_slot) == (2, [1.0] * len(_KDA_REPLAY_BUFFERS)) + assert int(mgr.prev_num_accepted_tokens[resume_slot]) == 0 + for section, name in enumerate(("kda_conv_q", "kda_conv_k", "kda_conv_v")): + cache = getattr(mgr, name)[0, resume_slot] + assert torch.equal(cache[:, :3], window[2 * section : 2 * section + 2]) + assert not cache[:, 3:].any() + for name in ("kda_qkg_cache", "kda_v_cache", "kda_beta_cache"): + assert not getattr(mgr, name)[0, resume_slot].any() if resume_slot != 0: assert _kda_replay_slot(mgr, 0) == (3, [7.0] * len(_KDA_REPLAY_BUFFERS)) - assert not mgr._suspended_kda_replay_state + assert not mgr._kda_replay_suspended -def test_v2_kda_replay_late_acceptance_follows_a_suspended_request(monkeypatch): - """The host drafter's acceptance is recorded by update_resources after the next batch is scheduled. A request - suspended in between takes that acceptance with its saved replay state instead of writing its old slot.""" +@pytest.mark.parametrize("recorded_by", ["update_resources", "revert_allocate_generation"]) +def test_v2_kda_replay_refuses_a_suspend_before_the_host_drafter_acceptance( + monkeypatch, recorded_by +): + """With a host drafter, update_resources records the acceptance of a batch's verify, and a suspend needs it to + commit the request's pending drafts. A suspend between the batch's prepare_resources and its update_resources + raises instead of committing a stale count. Once update_resources has recorded it, or the batch is reverted without + running, the suspend commits.""" mgr = _v2_kda_replay_stub() mgr._record_kda_replay_in_update_resources = True - monkeypatch.setattr(KVCacheManagerV2, "update_resources", MagicMock()) + for method in ("prepare_resources", "update_resources", "revert_allocate_generation"): + monkeypatch.setattr(KVCacheManagerV2, method, MagicMock()) request = SimpleNamespace( - py_request_id=1, is_dummy=False, py_draft_tokens=[11, 12], py_num_accepted_draft_tokens=0 + py_request_id=1, is_dummy=False, py_draft_tokens=[11, 12], py_num_accepted_draft_tokens=2 ) mgr.kv_cache_map[1] = _SlotCache(0) - mgr._setup_state_indices([request], num_contexts=0) - _write_kda_replay_slot(mgr, 0, accepted=1, value=1.0) - mgr.suspend_request(request) - mgr.kv_cache_map[2] = _SlotCache(0) - _write_kda_replay_slot(mgr, 0, accepted=3, value=7.0) - request.py_num_accepted_draft_tokens = 2 - mgr.update_resources(SimpleNamespace(generation_requests=[request])) + batch = SimpleNamespace(context_requests=[], generation_requests=[request]) + mgr.prepare_resources(batch) - assert int(mgr.prev_num_accepted_tokens[0]) == 3 - mgr.kv_cache_map[1].slot = 2 - mgr.kv_cache_map[1].is_active = True - mgr._setup_state_indices([request], num_contexts=0) - assert _kda_replay_slot(mgr, 2) == (2, [1.0] * len(_KDA_REPLAY_BUFFERS)) + with pytest.raises(RuntimeError, match="acceptance"): + mgr.suspend_request(request) + assert not mgr.committed_slots + assert mgr.kv_cache_map[1].is_active + + if recorded_by == "update_resources": + mgr.update_resources(batch) + assert int(mgr.prev_num_accepted_tokens[0]) == 2 + else: + mgr.revert_allocate_generation(request) + mgr.suspend_request(request) + assert mgr.committed_slots == [0] + assert not mgr.kv_cache_map[1].is_active def _v2_kda_replay_manager_with_large_slots(): @@ -3921,43 +3954,141 @@ def _v2_kda_replay_manager_with_large_slots(): ) -@skip_no_cuda -def test_v2_kda_replay_state_follows_a_suspended_request_on_the_runtime(): - """The same on the V2 runtime: once new caches fill the SSM pool, the suspended request's SSM page moves to the - host tier and one of them takes its slot; the request resumes in another slot with its own replay state.""" +def _take_ssm_slot(mgr, slot, others): + """Create caches until the V2 runtime gives one of them ``slot``, moving the suspended owner's SSM page to the host + tier; keep that one and close the rest.""" + holder = None + for _ in range(mgr.prev_num_accepted_tokens.shape[0] + 1): + cache = mgr.impl.create_kv_cache() + assert cache.resume(mgr._stream.cuda_stream) + others.append(cache) + assert cache.resize(1) + if cache.get_ssm_block_base_index(mgr.ssm_layer_group_id) == slot: + holder = cache + break + assert holder is not None, "no cache took the suspended request's SSM slot" + for cache in others: + if cache is not holder: + cache.close() + others[:] = [holder] + + +def _fill_kda_slot(mgr, slot, generator): + """Random V2 states and replay state at ``slot``: a new owner's prefill and verify rounds.""" + for layer_offset, (ssm_states, conv_states) in enumerate( + zip(mgr.all_ssm_states, mgr.all_conv_states) + ): + ssm_states[slot] = torch.randn(ssm_states.shape[1:], generator=generator, device="cuda") + conv_states[slot] = torch.randn(conv_states.shape[1:], generator=generator, device="cuda") + for name in _KDA_REPLAY_BUFFERS: + buffer = getattr(mgr, name)[layer_offset] + buffer[slot] = torch.randn(buffer.shape[1:], generator=generator, device="cuda") + mgr.prev_num_accepted_tokens[slot] = 1 + + +def _kda_verify_round(mgr, slot, layer_weights, round_idx, num_spec): + """One verify of 1 + num_spec tokens on every KDA layer for the request in ``slot``, as the KDA mixer runs it; the + inputs depend only on the round. Returns the outputs.""" + tokens = 1 + num_spec + pending = mgr.prev_num_accepted_tokens[slot : slot + 1].clone() + slots = torch.tensor([slot], dtype=torch.int32, device="cuda") + cu_seqlens = torch.tensor([0, tokens], dtype=torch.int32, device="cuda") + cu_seqlens[:1].sub_(pending) + outputs = [] + for layer_offset, weights in enumerate(layer_weights): + generator = torch.Generator(device="cuda").manual_seed(1000 * round_idx + layer_offset) + num_heads, head_dim = weights["A_log"].shape[0], 128 + + def inputs(*shape): + return torch.randn(shape, generator=generator, device="cuda").bfloat16() + + outputs.append( + torch.ops.trtllm.kda_mtp_decode( + x_q=inputs(1, tokens, num_heads, head_dim), + x_k=inputs(1, tokens, num_heads, head_dim), + x_v=inputs(1, tokens, num_heads, head_dim), + w_q=weights["w_q"], + w_k=weights["w_k"], + w_v=weights["w_v"], + cs_q=mgr.kda_conv_q[layer_offset], + cs_k=mgr.kda_conv_k[layer_offset], + cs_v=mgr.kda_conv_v[layer_offset], + g=inputs(1, tokens, num_heads, head_dim), + beta=inputs(1, tokens, num_heads), + A_log=weights["A_log"], + dt_bias=weights["dt_bias"], + recurrent_state=mgr.all_ssm_states[layer_offset], + qkg_cache=mgr.kda_qkg_cache[layer_offset], + v_cache=mgr.kda_v_cache[layer_offset], + beta_cache=mgr.kda_beta_cache[layer_offset], + ssm_state_indices=slots, + cu_seqlens=cu_seqlens, + num_spec=num_spec, + num_accepted_tokens=pending, + lower_bound=-5.0, + scale=head_dim**-0.5, + ) + ) + return torch.stack(outputs) + + +def _run_kda_verify_rounds(accepted_per_round, suspend_after=None, slot_taken=False): + """Verify rounds of one request on a V2 manager at Kimi K3's per-rank KDA shapes, each accepting the given number + of drafts. With ``suspend_after``, the request is suspended after that round's acceptance and resumed, its slot + taken and overwritten by another cache meanwhile when ``slot_taken``. Returns every round's outputs and the + request's final SSM states.""" + num_spec = 2 mgr = _v2_kda_replay_manager_with_large_slots() - stream = mgr._stream.cuda_stream others = [] request = None try: (request,) = mgr.add_dummy_requests( - [1], token_nums=[8], is_gen=True, max_num_draft_tokens=2 + [1], token_nums=[8], is_gen=True, max_num_draft_tokens=num_spec ) + generator = torch.Generator(device="cuda").manual_seed(7) + num_heads, head_dim, width = 6, 128, 4 + layer_weights = [ + { + "w_q": torch.randn(num_heads * head_dim, width, generator=generator, device="cuda") + * 0.5, + "w_k": torch.randn(num_heads * head_dim, width, generator=generator, device="cuda") + * 0.5, + "w_v": torch.randn(num_heads * head_dim, width, generator=generator, device="cuda") + * 0.5, + "A_log": torch.randn(num_heads, generator=generator, device="cuda"), + "dt_bias": torch.randn(num_heads * head_dim, generator=generator, device="cuda"), + } + for _ in mgr.all_ssm_states + ] + # The state a prefill leaves: the V2 states, and the replay caches seeded from them with nothing pending. slot = mgr._request_id_to_state_index[1] - _write_kda_replay_slot(mgr, slot, accepted=2, value=1.0) - own = _kda_replay_slot(mgr, slot) - mgr.suspend_request(request) - holder = None - for _ in range(mgr.prev_num_accepted_tokens.shape[0] + 1): - cache = mgr.impl.create_kv_cache() - assert cache.resume(stream) - others.append(cache) - assert cache.resize(1) - if cache.get_ssm_block_base_index(mgr.ssm_layer_group_id) == slot: - holder = cache - break - assert holder is not None, "no cache took the suspended request's SSM slot" - _write_kda_replay_slot(mgr, slot, accepted=3, value=7.0) - for cache in others: - if cache is not holder: - cache.close() - others = [holder] - assert mgr.resume_request(request) - mgr._setup_state_indices([request], num_contexts=0) - new_slot = mgr._request_id_to_state_index[1] - - assert new_slot != slot - assert _kda_replay_slot(mgr, new_slot) == own + for ssm_states, conv_states in zip(mgr.all_ssm_states, mgr.all_conv_states): + ssm_states[slot] = 0.1 * torch.randn( + ssm_states.shape[1:], generator=generator, device="cuda" + ) + conv_states[slot] = torch.randn( + conv_states.shape[1:], generator=generator, device="cuda" + ) + mgr.seed_kda_replay_caches_for_disagg_gen([1]) + + outputs = [] + for round_idx, accepted in enumerate(accepted_per_round): + slot = mgr._request_id_to_state_index[1] + outputs.append(_kda_verify_round(mgr, slot, layer_weights, round_idx, num_spec)) + mgr.prev_num_accepted_tokens[slot] = accepted + if round_idx == suspend_after: + mgr.suspend_request(request) + if slot_taken: + _take_ssm_slot(mgr, slot, others) + _fill_kda_slot(mgr, slot, generator) + assert mgr.resume_request(request) + mgr._setup_state_indices([request], num_contexts=0) + if slot_taken: + assert mgr._request_id_to_state_index[1] != slot + slot = mgr._request_id_to_state_index[1] + states = torch.stack([ssm_states[slot] for ssm_states in mgr.all_ssm_states]) + torch.cuda.synchronize() + return outputs, states finally: for cache in others: cache.close() @@ -3966,6 +4097,30 @@ def test_v2_kda_replay_state_follows_a_suspended_request_on_the_runtime(): mgr.shutdown() +@skip_no_cuda +@pytest.mark.parametrize("slot_taken", [True, False], ids=["slot_taken", "slot_kept"]) +def test_v2_kda_replay_suspend_resume_matches_an_unsuspended_run(slot_taken): + """A request is suspended with two accepted drafts pending and resumes; every later verify gives the outputs and + leaves the states of the same request never suspended, bit for bit, so its tokens cannot change. + + slot_taken: while the request is suspended, new caches fill the SSM pool, so the V2 runtime moves its SSM page to + the host tier and gives its slot to one of them, which overwrites the slot's states; the request resumes in + another slot. slot_kept: nothing else runs while it is suspended.""" + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import is_kda_mtp_verify_available + + if not is_kda_mtp_verify_available(): + pytest.skip("needs the fused KDA verify kernel (sm_100 / sm_103 with the CuTe DSL)") + accepted_per_round = [2, 1, 2, 0, 2, 1] + expected_outputs, expected_states = _run_kda_verify_rounds(accepted_per_round) + outputs, states = _run_kda_verify_rounds( + accepted_per_round, suspend_after=2, slot_taken=slot_taken + ) + + for round_idx, (actual, expected) in enumerate(zip(outputs, expected_outputs)): + assert torch.equal(actual, expected), f"verify round {round_idx}" + assert torch.equal(states, expected_states) + + @skip_no_cuda def test_v2_kda_replay_relocates_live_slot_history(): mgr = _build_v2_hybrid_with_mamba_layer( From 93c0b28bbdd11fbe457af85aa73a979bf07d65e2 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 18:24:33 -0700 Subject: [PATCH 39/67] MNNVL all-reduce: remove LamportFlags' unused cluster-arrival mode Since both RMSNorm-fused kernels count Lamport arrivals per CTA, nothing instantiates LamportFlags<..., true>. Remove the UseCGA template parameter, the cluster branches of ctaArrive and waitAndUpdate, and the cooperative_groups include and alias that only they used. Also delete the commented-out cluster.sync() placeholders that the added cluster.sync() calls replace. No functional change: every kernel already took the per-CTA path. Signed-off-by: Vasanth Sabavat --- cpp/tensorrt_llm/common/lamportUtils.cuh | 66 ++++--------------- .../mnnvlAllreduceKernels.cu | 8 +-- 2 files changed, 16 insertions(+), 58 deletions(-) diff --git a/cpp/tensorrt_llm/common/lamportUtils.cuh b/cpp/tensorrt_llm/common/lamportUtils.cuh index 28a06d1d062a..15ca85c640d6 100644 --- a/cpp/tensorrt_llm/common/lamportUtils.cuh +++ b/cpp/tensorrt_llm/common/lamportUtils.cuh @@ -26,8 +26,6 @@ #include #include -#include - #include "tensorrt_llm/common/cudaTypeUtils.cuh" TRTLLM_NAMESPACE_BEGIN @@ -196,6 +194,7 @@ struct LamportBufferLayout (lamportIndex * numStages + stageIndex) * static_cast(bytesPerBuffer / numStages))); } }; + // Current Index // Dirty Index // bytes_per_buffer @@ -203,10 +202,8 @@ struct LamportBufferLayout // Dirty bytes_to_clear = {stage0, stage1, stage2, stage3} # We fix this to 4 stages // offset_access_ptr -namespace cg = cooperative_groups; - // PackedType is the one used in kernel for Lamport buffer (LDG.128 or LDG.64) -template +template struct __attribute__((aligned(32))) LamportFlags { public: @@ -258,65 +255,28 @@ public: __device__ void ctaArrive() { - if constexpr (UseCGA) +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) + uint32_t const barrierThreads + = ((static_cast(blockDim.x) + kWARP_SIZE - 1U) / kWARP_SIZE) * kWARP_SIZE; + if (threadIdx.x < kWARP_SIZE) { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - // Every thread arrives on the cluster barrier but only rank 0's first warp waits, so a kernel using this - // must not arrive on the cluster barrier again (each thread arrives and waits once per phase). - cg::cluster_group cluster = cg::this_cluster(); - __cluster_barrier_arrive(); - if (cluster.block_rank() == 0 && threadIdx.x < kWARP_SIZE) - { - __cluster_barrier_wait(); - arriveCounter(threadIdx.x); - } -#else - __syncthreads(); + asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory"); arriveCounter(threadIdx.x); -#endif } else { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 700)) - uint32_t const barrierThreads - = ((static_cast(blockDim.x) + kWARP_SIZE - 1U) / kWARP_SIZE) * kWARP_SIZE; - if (threadIdx.x < kWARP_SIZE) - { - asm volatile("barrier.cta.sync 1, %0;" ::"r"(barrierThreads) : "memory"); - arriveCounter(threadIdx.x); - } - else - { - asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory"); - } + asm volatile("barrier.cta.arrive 1, %0;" ::"r"(barrierThreads) : "memory"); + } #else - __syncthreads(); - arriveCounter(threadIdx.x); + __syncthreads(); + arriveCounter(threadIdx.x); #endif - } } __device__ void waitAndUpdate(uint4 bytesToClearPerStage) { - bool isLastCtaT0{false}; - int targetCount{0}; - if constexpr (UseCGA) - { -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cg::grid_group grid = cg::this_grid(); - // Use the first thread instead of the last thread as the last thread may exit early. - isLastCtaT0 = grid.thread_rank() == 0; - targetCount = grid.num_clusters(); -#else - isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; - targetCount = gridDim.x * gridDim.y * gridDim.z; -#endif - } - else - { - isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; - targetCount = gridDim.x * gridDim.y * gridDim.z; - } + bool const isLastCtaT0 = threadIdx.x == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0; + int const targetCount = gridDim.x * gridDim.y * gridDim.z; if (isLastCtaT0) { uint4* flagPtr = reinterpret_cast(mBufferFlagsPtr); diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu index e60a2f233c15..0e25a25eb46c 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu @@ -599,7 +599,7 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed // We only use 1 stage for the oneshot allreduce. Arrivals are counted per CTA: the RMSNorm epilogue's cluster // reduction below must be each thread's only phase of the cluster barrier. - LamportFlags flag(params.bufferFlags, 1); + LamportFlags flag(params.bufferFlags, 1); T* stagePtrMcast = reinterpret_cast(flag.getCurLamportBuf(params.mcastPtr, 0)); T* stagePtrLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], 0)); bool const inBounds = packedIdx * kELTS_PER_THREAD < params.tokenDim; @@ -717,7 +717,6 @@ __global__ void __launch_bounds__(1024) oneshotAllreduceFusionKernel(MnnvlAllRed { cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; } - // cluster.sync(); cluster.barrier_wait(cluster.barrier_arrive()); for (int i = 0; i < numBlocks; ++i) { @@ -893,7 +892,7 @@ __global__ __launch_bounds__(128) void twoshotAllreduceKernel(MnnvlAllReduceKern #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) cudaGridDependencySynchronize(); #endif - LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); T* scatterBufLocal = reinterpret_cast(flag.getCurLamportBuf(params.inputPtrs[params.rank], MNNVLTwoShotStage::SCATTER)); @@ -1021,7 +1020,7 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam // Arrivals are counted per CTA: with UseCGA the RMSNorm cluster reduction below must be each thread's only phase // of the cluster barrier. - LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); + LamportFlags flag(params.bufferFlags, MNNVLTwoShotStage::NUM_STAGES); T* input = reinterpret_cast( flag.getCurLamportBuf(reinterpret_cast(params.bufferInputPtr), MNNVLTwoShotStage::BROADCAST)); @@ -1148,7 +1147,6 @@ __global__ __launch_bounds__(1024) void rmsNormLamport(MnnvlAllReduceKernelParam { cluster.map_shared_rank(&sharedVal[0], threadIdx.x)[blockRank] = blockSum; } - // cluster.sync(); cluster.barrier_wait(cluster.barrier_arrive()); for (int i = 0; i < numBlocks; ++i) { From eda9db209ff4358946550a5bb763ba3f90b4da3a Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 23:27:42 -0700 Subject: [PATCH 40/67] [None][fix] DSpark GQA drafter: check the kept mask row at load and cast it once Follow-ups on the trained mask-token row this branch keeps: - The row is checked at load, before any weight is copied: the shipped embedding must be [rows, hidden_size], and the drafter's mask_token_id must be one of its rows. A ValueError names the problem. Before, an id past the rows kept an empty row, which failed later in the worker's noise block. - The row is cast to the drafter's dtype and placed on the device of the drafter's own fc weight once, at load. The noise block uses it as given, with no per-step cast and no reliance on the current CUDA device. - The docstring says why only the mask row is kept: in the Kimi K3 drafter (RedHatAI/Kimi-K3-speculator.dspark) it is the only one of the 163,840 embedding rows that differs from the target's, and the shipped lm_head is bit-exact with the target's. Tests: test_gqa_dspark_mask_row_takes_the_drafter_dtype_and_device_at_load and test_gqa_dspark_rejects_a_mask_row_the_shipped_embedding_does_not_have (out-of-range, negative and wrong-width cases). Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/models/modeling_dspark.py | 40 +++++++++++---- tensorrt_llm/_torch/speculative/dflash.py | 4 +- .../test_kimi_k3_dspark_semantics.py | 50 +++++++++++++++++-- 3 files changed, 81 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 1855c56b4312..e0fcd0b8e44f 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2673,22 +2673,44 @@ def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): """Take the DSpark head weights, then hand the rest to DFlash.""" weights, _ = self._take_dspark_head_weights(weights) - self._keep_trained_mask_embedding(weights) - return super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + # Checked before the load so a bad row fails early; placed after it, once ``fc`` (its device) exists. + mask_row = self._trained_mask_row(weights) + loaded = super().load_weights(weights, weight_mapper=weight_mapper, **kwargs) + if mask_row is not None: + self.mask_token_embedding = mask_row.to( + device=self.fc.weight.device, dtype=self.model_config.torch_dtype + ) + return loaded - def _keep_trained_mask_embedding(self, weights: Dict) -> None: - """Keep the checkpoint's own embedding row for the mask token. + def _trained_mask_row(self, weights: Dict) -> Optional[torch.Tensor]: + """The checkpoint's own embedding row for the mask token, [hidden_size]; None if it ships no embedding. This drafter takes the target's embedding (``load_weights_from_target_model``). A drafter checkpoint that ships an embedding has trained the mask token's row, which the target's embedding does not have: the block decode - reads ``mask_token_embedding`` for every masked slot instead of the shared lookup. + reads ``mask_token_embedding`` for every masked slot instead of the shared lookup. Only that row is kept: in + the Kimi K3 drafter (RedHatAI/Kimi-K3-speculator.dspark) it is the only row of the shipped embedding that + differs from the target's, and the shipped ``lm_head`` is bit-exact with the target's. + + Raises: + ValueError: the embedding is not ``[rows, hidden_size]``, or ``mask_token_id`` is not one of its rows. """ names = [k for k in ("embed_tokens.weight", "model.embed_tokens.weight") if k in weights] if not names: - return - weight = weights[names[0]] - row = weight[self.mask_token_id : self.mask_token_id + 1] - self.mask_token_embedding = torch.as_tensor(row).reshape(-1).to("cuda") + return None + name = names[0] + weight = weights[name] + hidden_size = self.config.hidden_size + if weight.dim() != 2 or weight.shape[1] != hidden_size: + raise ValueError( + f"{type(self).__name__}: the checkpoint's {name} has shape {list(weight.shape)}, " + f"expected [rows, hidden_size={hidden_size}]." + ) + if not 0 <= self.mask_token_id < weight.shape[0]: + raise ValueError( + f"{type(self).__name__}: the drafter's mask_token_id {self.mask_token_id} is not a row of the " + f"checkpoint's {name} ({weight.shape[0]} rows)." + ) + return weight[self.mask_token_id] class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 02be1e4f6aba..5eec1868d5d1 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -445,6 +445,8 @@ def dflash_noise_block_embedding( The mask row is the drafter's own trained one when it has one (``trained_mask_embedding``); otherwise the shared lookup's. Both rows go through ``embed_tokens.forward`` (NOT ``.weight[...]``) so TP-sharded vocabs mask out ranks that don't own the token id and all-reduce, the mask row included, so every rank makes the same call. + + ``trained_mask_embedding`` ([hidden]) is used as given: the drafter casts it to its own dtype and device at load. """ num_gens = bonus.shape[0] mask_tok = torch.full((1,), int(mask_token_id), dtype=torch.long, device=bonus.device) @@ -452,7 +454,7 @@ def dflash_noise_block_embedding( embed_bonus = combined_embed[:num_gens] embed_mask = combined_embed[num_gens] if trained_mask_embedding is not None: - embed_mask = trained_mask_embedding.to(embed_mask.dtype) + embed_mask = trained_mask_embedding noise_embed_2d = embed_mask.expand(num_gens, block_size, -1).clone() noise_embed_2d[:, 0, :] = embed_bonus return noise_embed_2d diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index 0939bfdca109..4d895340b216 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -235,7 +235,7 @@ def test_swa_window_conventions(): NUM_CAPTURE = 2 -def _tiny_config(dspark: bool, *, published_spelling: bool = False): +def _tiny_config(dspark: bool, *, published_spelling: bool = False, mask_token_id: int = VOCAB - 2): """Tiny drafter config. ``published_spelling`` reproduces the two public K3 DSpark checkpoints @@ -248,7 +248,7 @@ def _tiny_config(dspark: bool, *, published_spelling: bool = False): from transformers import Qwen3Config cfg = dict(TINY) - dflash = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} + dflash = {"mask_token_id": mask_token_id, "target_layer_ids": [0, 1]} if dspark and published_spelling: cfg.update( markov_rank=RANK, @@ -330,11 +330,14 @@ def _build_drafter( *, published_spelling: bool = False, dflash_attention_backend: str = "VANILLA", + mask_token_id: int = VOCAB - 2, ): from tensorrt_llm._torch.model_config import ModelConfig model_config = ModelConfig( - pretrained_config=_tiny_config(dspark, published_spelling=published_spelling), + pretrained_config=_tiny_config( + dspark, published_spelling=published_spelling, mask_token_id=mask_token_id + ), attn_backend="TRTLLM", ) # The DSpark head set lives in the DSpark drafter, not in the DFlash base. @@ -498,6 +501,47 @@ def test_gqa_dspark_keeps_its_trained_mask_row_over_the_target_embedding(): torch.testing.assert_close(noise[:, j].cpu(), trained.expand(2, -1), rtol=0, atol=0) +@needs_gpu +def test_gqa_dspark_mask_row_takes_the_drafter_dtype_and_device_at_load(): + """The kept row is cast once, at load, to the drafter's dtype (bf16 here, from an fp32 checkpoint) and placed on + the device of the drafter's own parameters, so the worker's noise block uses it as is.""" + weights = _tiny_weights() + g = torch.Generator().manual_seed(29) + weights["embed_tokens.weight"] = torch.randn(VOCAB, TINY["hidden_size"], generator=g) * 0.05 + drafter = _build_drafter(True, weights) + + row = drafter.mask_token_embedding + assert row.dtype == torch.bfloat16 + assert row.device == drafter.fc.weight.device + trained = weights["embed_tokens.weight"][drafter.mask_token_id].to(torch.bfloat16) + torch.testing.assert_close(row.cpu(), trained, rtol=0, atol=0) + + +@needs_gpu +@pytest.mark.parametrize( + "width, mask_token_id, message", + [ + (TINY["hidden_size"], VOCAB, f"mask_token_id {VOCAB} is not a row"), + (TINY["hidden_size"], -1, "mask_token_id -1 is not a row"), + ( + TINY["hidden_size"] + 8, + VOCAB - 2, + rf"expected \[rows, hidden_size={TINY['hidden_size']}\]", + ), + ], + ids=["mask_id_past_the_rows", "negative_mask_id", "wrong_width"], +) +def test_gqa_dspark_rejects_a_mask_row_the_shipped_embedding_does_not_have( + width, mask_token_id, message +): + """A mask id that is not a row of the shipped embedding, or rows that are not hidden_size wide, fail at load, + before any weight is copied, with an error that names the problem.""" + weights = _tiny_weights() + weights["embed_tokens.weight"] = torch.zeros(VOCAB, width, dtype=torch.bfloat16) + with pytest.raises(ValueError, match=message): + _build_drafter(True, weights, mask_token_id=mask_token_id) + + @needs_gpu def test_published_drafter_spelling_activates_the_heads(): """Both public K3 DSpark checkpoints load with their heads live. From a0fa306023241e6de394c074e9dc64ff96e9d0b1 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 23:29:44 -0700 Subject: [PATCH 41/67] [None][fix] Mamba2 metadata: upload CPU-tensor state indices through copy_to_device_if_changed Mamba2Metadata.prepare() copied state indices that the cache manager returns as a CPU tensor straight into state_indices, so the values copy_to_device_if_changed keeps for that buffer went stale. Indices given as a list, then as a CPU tensor, then as the first list again were taken as unchanged on the third step, and the device kept the tensor's values. Both host branches now upload through the helper. The new test runs that list -> tensor -> list sequence. Signed-off-by: Vasanth Sabavat --- .../_torch/modules/mamba/mamba2_metadata.py | 4 +- .../modules/mamba/test_mamba2_metadata.py | 63 +++++++++++++------ 2 files changed, 46 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index f4fde357773e..b4380937db8e 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -493,8 +493,8 @@ def prepare(self, attn_metadata: AttentionMetadata): elif isinstance(indices, torch.Tensor): # CPU tensor → bulk H2D self.state_indices_cpu[:batch_size].copy_(indices[:batch_size]) - self.state_indices[:batch_size].copy_( - self.state_indices_cpu[:batch_size], non_blocking=True) + copy_to_device_if_changed(self.state_indices, + self.state_indices_cpu[:batch_size]) else: # indices is a Python sequence (e.g. List[int]); data # already lives on host, CPU staging is fine. One bulk diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 16958c073f84..fd5fcfff902d 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -59,6 +59,30 @@ def get_replay_state_update_metadata(self): ) +class _StateIndexCacheManager: + """Returns ``state_indices`` as set, a list or a CPU tensor, for a batch of decode requests.""" + + def __init__(self, state_indices): + self.state_indices = state_indices + + def get_state_indices(self, request_ids, is_padding): + return self.state_indices[: len(request_ids)] + + +def _decode_attn_metadata(kv_cache_manager): + """Attention metadata of a two-request decode step.""" + seq_lens = torch.tensor([1, 1], dtype=torch.int) + return SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=0, + num_ctx_tokens=0, + kv_cache_manager=kv_cache_manager, + request_ids=[10, 11], + kv_cache_params=None, + ) + + def _torch_reference_work_items(state_indices, prev_num_accepted_tokens, cache_buf_idx): """Run the production ATen path into fresh buffers. @@ -141,26 +165,9 @@ def test_prepare_handles_tensor_cached_tokens(self): def test_prepare_skips_an_unchanged_state_index_upload(self): """A decode step whose state indices (a list from the cache manager) did not change enqueues no upload: the device buffer keeps values the host never sent. A changed step uploads again.""" - - class ListCacheManager: - def __init__(self): - self.state_indices = [3, 1] - - def get_state_indices(self, request_ids, is_padding): - return self.state_indices[: len(request_ids)] - - manager = ListCacheManager() + manager = _StateIndexCacheManager([3, 1]) metadata = Mamba2Metadata(max_batch_size=2, chunk_size=8) - seq_lens = torch.tensor([1, 1], dtype=torch.int) - attn_metadata = SimpleNamespace( - seq_lens=seq_lens, - seq_lens_cuda=seq_lens.cuda(), - num_contexts=0, - num_ctx_tokens=0, - kv_cache_manager=manager, - request_ids=[10, 11], - kv_cache_params=None, - ) + attn_metadata = _decode_attn_metadata(manager) metadata.prepare(attn_metadata) torch.cuda.synchronize() @@ -176,6 +183,24 @@ def get_state_indices(self, request_ids, is_padding): torch.cuda.synchronize() assert metadata.state_indices[:2].tolist() == [3, 2] + def test_prepare_uploads_state_indices_from_a_list_or_a_cpu_tensor(self): + """State indices that come as a list, then a CPU tensor, then the first list again all reach the device: the + tensor step's upload updates the record of what the device holds, so the last step is not taken as + unchanged.""" + manager = _StateIndexCacheManager([3, 1]) + metadata = Mamba2Metadata(max_batch_size=2, chunk_size=8) + attn_metadata = _decode_attn_metadata(manager) + + for state_indices, expected in ( + ([3, 1], [3, 1]), + (torch.tensor([5, 6], dtype=torch.int32), [5, 6]), + ([3, 1], [3, 1]), + ): + manager.state_indices = state_indices + metadata.prepare(attn_metadata) + torch.cuda.synchronize() + assert metadata.state_indices[:2].tolist() == expected + def test_prepare_replay_work_items_write_first(self): class ReplayCacheManager: use_replay_state_update = True From fd37763e93e21de634c1f398ec3c79bbe30a0fc2 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 23:33:50 -0700 Subject: [PATCH 42/67] [None][test] copy_to_device_if_changed: check skipped copies in debug mode and unit tests copy_to_device_if_changed skips a copy whose values equal the ones it last copied into the buffer, which is right only while nothing else writes that buffer. With TLLM_DEBUG_MODE=1, and in every unit test through an autouse fixture in tests/unittest/conftest.py, a skipped copy now synchronizes, reads the buffer back and asserts that it holds the values, so a second writer fails instead of leaving stale values on the device. The check is skipped during CUDA graph capture. Each buffer written through the helper says so where it is allocated: the engine's gather_ids_cuda, Mamba2Metadata.state_indices, the hybrid Mamba cache managers' _dummy_request_mask (C++ and V2) and the V2 manager's cuda_state_indices, DFlash's batch_indices_cuda and _batch_to_slot, and the speculative sampler's _slot_table. The tests that write such a buffer on purpose, to see a copy skipped, turn the check off. New tests: the check passes a buffer only the helper wrote and reports another write, on the host and through Mamba2Metadata.prepare() on the device. Signed-off-by: Vasanth Sabavat --- .../_torch/modules/mamba/mamba2_metadata.py | 3 +++ .../kv_cache/mamba_cache_manager.py | 6 +++++ .../_torch/pyexecutor/model_engine.py | 3 +++ tensorrt_llm/_torch/speculative/dflash.py | 4 +++ .../_torch/speculative/spec_sampler_base.py | 2 ++ tensorrt_llm/_utils.py | 25 ++++++++++++++++++ .../kv_cache/test_mamba_cache_manager.py | 5 +++- .../test_copy_to_device_if_changed.py | 26 ++++++++++++++++++- .../modules/mamba/test_mamba2_metadata.py | 18 ++++++++++++- tests/unittest/conftest.py | 10 +++++++ 10 files changed, 99 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index b4380937db8e..0a2478801bce 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -371,6 +371,9 @@ def __init__( self.state_indices_cpu = torch.zeros(max_batch_size, dtype=torch.int32, pin_memory=prefer_pinned()) + # Only copy_to_device_if_changed writes this buffer: it skips values + # it already copied here, so another write would leave stale values. + # prepare() aliases a cache manager's own device buffer instead. self.state_indices = torch.zeros(max_batch_size, dtype=torch.int32, device="cuda") diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 075bd85fbc13..448282b282af 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -2952,6 +2952,8 @@ def _setup_replay_buffers(self, spec_config) -> None: device): return + # Only copy_to_device_if_changed writes this buffer: it skips values + # it already copied here, so another write would leave stale values. self._dummy_request_mask = torch.zeros(self.max_batch_size, dtype=torch.bool, device=device) @@ -3236,6 +3238,8 @@ def __init__( state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) + # Only copy_to_device_if_changed writes this buffer: it skips values + # it already copied here, so another write would leave stale values. self.cuda_state_indices = torch.zeros([state_index_capacity], dtype=torch.int32, device="cuda") @@ -4153,6 +4157,8 @@ def _setup_replay_buffers(self, spec_config) -> None: return mask_capacity = self._host_state_indices.shape[0] + # Only copy_to_device_if_changed writes this buffer: it skips values + # it already copied here, so another write would leave stale values. self._dummy_request_mask = torch.zeros(mask_capacity, dtype=torch.bool, device=device) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 9aa68f42b830..d4573691c7c6 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -816,6 +816,9 @@ def _allocate_decoder_buffers(self) -> None: self.draft_tokens_cuda = torch.empty((max_num_draft_tokens, ), dtype=torch.int, device='cuda') + # Only copy_to_device_if_changed writes this buffer: it skips + # values it already copied here, so another write would leave + # stale values. self.gather_ids_cuda = torch.empty((self.max_num_tokens, ), dtype=torch.int, device='cuda') diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index e0bf55de16b3..aad46d860ad3 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -490,6 +490,8 @@ class DFlashSpecMetadata(SpecMetadata): captured_hidden_states: Optional[torch.Tensor] = None def __post_init__(self): + # Only copy_to_device_if_changed writes this buffer: it skips values it already copied here, so another write + # would leave stale values. self.batch_indices_cuda = torch.empty( [self.max_num_requests], dtype=torch.int, @@ -1046,6 +1048,8 @@ def _lazy_init_ctx_buffers( self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") self._ctx_len_host = [0] * num_slots + # Only copy_to_device_if_changed writes this buffer: it skips values it already copied here, so another write + # would leave stale values. self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") self._free_slots = deque(range(max_batch)) diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 8cd5ad530b76..595a9ac502d3 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -309,6 +309,8 @@ def __init__( self._slot_table: Optional[torch.Tensor] = None if spec_step_copies.is_supported(): self._store_scatter = spec_step_copies.SlotScatter() + # Only copy_to_device_if_changed writes this buffer: it skips values it already copied here, so another + # write would leave stale values. self._slot_table = torch.zeros((seq_slots,), dtype=torch.int32, device="cuda") # Recorded after the last step's host copies of the stores, which read # them on the D2H side stream. diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index b45bb486b740..d2f76483b441 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -1336,6 +1336,25 @@ def maybe_pin_memory(tensor: torch.Tensor) -> torch.Tensor: return tensor +# Set by TLLM_DEBUG_MODE=1, and for every unit test by a fixture: a copy that +# copy_to_device_if_changed skips first checks that the buffer holds the values. +_verify_skipped_device_copies = os.environ.get("TLLM_DEBUG_MODE", + "")[0:1] == "1" + + +def _assert_device_holds(dst: torch.Tensor, flat: torch.Tensor) -> None: + """Assert that the leading elements of ``dst`` hold ``flat`` once all queued + work has finished.""" + if dst.is_cuda: + if torch.cuda.is_current_stream_capturing(): + return + torch.cuda.synchronize(dst.device) + held = dst.view(-1)[:flat.numel()].cpu() + assert torch.equal(held, flat.to(held.dtype)), ( + "copy_to_device_if_changed skipped a copy, but the buffer holds " + f"{held.tolist()} instead of {flat.tolist()}: something else wrote it") + + def copy_to_device_if_changed(dst: torch.Tensor, host: torch.Tensor) -> None: """Copy ``host`` into the leading elements of the device buffer ``dst`` unless they already hold it. @@ -1347,11 +1366,17 @@ def copy_to_device_if_changed(dst: torch.Tensor, host: torch.Tensor) -> None: objects over the same memory, such as the views a shared buffer pool hands to each metadata object. The copy itself goes through fresh pinned staging, so ``host`` may be reused right away. + + With ``TLLM_DEBUG_MODE=1`` a skipped copy reads ``dst`` back (a device + synchronization) and asserts that it holds ``host``, which catches a write + by anything else. """ flat = host.reshape(-1) n = flat.numel() last = getattr(dst, "_last_host_values", None) if last is not None and last.numel() >= n and torch.equal(last[:n], flat): + if _verify_skipped_device_copies: + _assert_device_holds(dst, flat) return staging = torch.empty_like(flat, device="cpu", pin_memory=prefer_pinned()) staging.copy_(flat) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 5f964814d2d3..39ea1c957ca2 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3842,9 +3842,12 @@ def close(self) -> None: @skip_no_cuda -def test_v2_state_index_setup_skips_unchanged_uploads(): +def test_v2_state_index_setup_skips_unchanged_uploads(monkeypatch): """A step whose state indices and dummy flags did not change enqueues no upload of either: the device buffers keep values the host never sent. A changed step uploads both again.""" + # The test writes the buffers behind _setup_state_indices to see the skip, which the check of skipped copies + # reports. + monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", False) class StateCache: def __init__(self, slot: int) -> None: diff --git a/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py b/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py index c4305439b5a4..e7772ac2d217 100644 --- a/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py +++ b/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py @@ -15,7 +15,8 @@ """``copy_to_device_if_changed``: which calls copy and what the buffer then holds. The skip decision does not depend on the device, so the buffer here is a host tensor: a skipped copy shows as a -sentinel written behind the function's back that survives the call. +sentinel written behind the function's back that survives the call. Such a write is what the debug check of a skipped +copy reports, so these tests turn the check off, except the ones that test it. """ import pytest @@ -28,6 +29,11 @@ SENTINEL = -7 +@pytest.fixture(autouse=True) +def _no_skipped_copy_check(monkeypatch) -> None: + monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", False) + + def _host(values: list[int]) -> torch.Tensor: return torch.tensor(values, dtype=torch.int32) @@ -117,3 +123,21 @@ def test_two_dimensional_values_are_flattened() -> None: _tamper(dst) copy_to_device_if_changed(dst, _host([1, 2, 3, 4])) assert dst.eq(SENTINEL).all() + + +def test_debug_check_passes_a_buffer_only_the_function_wrote(monkeypatch) -> None: + monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", True) + dst = torch.zeros(8, dtype=torch.int64) + copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) + copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) + copy_to_device_if_changed(dst, _host([5, 6])) + assert dst[:4].tolist() == [5, 6, 7, 8] + + +def test_debug_check_reports_a_write_by_anything_else(monkeypatch) -> None: + monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", True) + dst = torch.zeros(8, dtype=torch.int32) + copy_to_device_if_changed(dst, _host([3, 1])) + dst[:2].copy_(_host([5, 6])) + with pytest.raises(AssertionError, match="holds \\[5, 6\\] instead of \\[3, 1\\]"): + copy_to_device_if_changed(dst, _host([3, 1])) diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index fd5fcfff902d..44d2fd4c3055 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -162,9 +162,11 @@ def test_prepare_handles_tensor_cached_tokens(self): assert metadata.chunk_indices is not None assert metadata.chunk_offsets is not None - def test_prepare_skips_an_unchanged_state_index_upload(self): + def test_prepare_skips_an_unchanged_state_index_upload(self, monkeypatch): """A decode step whose state indices (a list from the cache manager) did not change enqueues no upload: the device buffer keeps values the host never sent. A changed step uploads again.""" + # The test writes the buffer behind prepare() to see the skip, which the check of skipped copies reports. + monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", False) manager = _StateIndexCacheManager([3, 1]) metadata = Mamba2Metadata(max_batch_size=2, chunk_size=8) attn_metadata = _decode_attn_metadata(manager) @@ -201,6 +203,20 @@ def test_prepare_uploads_state_indices_from_a_list_or_a_cpu_tensor(self): torch.cuda.synchronize() assert metadata.state_indices[:2].tolist() == expected + def test_check_of_skipped_copies_reports_a_write_behind_prepare(self, monkeypatch): + """With the check of skipped copies on (TLLM_DEBUG_MODE=1, or the unit-test fixture), a state index buffer + that something other than prepare() wrote fails the next prepare() whose indices did not change.""" + monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", True) + manager = _StateIndexCacheManager([3, 1]) + metadata = Mamba2Metadata(max_batch_size=2, chunk_size=8) + attn_metadata = _decode_attn_metadata(manager) + + metadata.prepare(attn_metadata) + metadata.prepare(attn_metadata) + metadata.state_indices[:2].copy_(torch.tensor([5, 6], dtype=torch.int32)) + with pytest.raises(AssertionError, match="holds \\[5, 6\\] instead of \\[3, 1\\]"): + metadata.prepare(attn_metadata) + def test_prepare_replay_work_items_write_first(self): class ReplayCacheManager: use_replay_state_update = True diff --git a/tests/unittest/conftest.py b/tests/unittest/conftest.py index a94b107b461f..364cc6f2d670 100644 --- a/tests/unittest/conftest.py +++ b/tests/unittest/conftest.py @@ -390,6 +390,16 @@ def torch_empty_cache() -> None: torch.cuda.empty_cache() +@pytest.fixture(autouse=True) +def verify_skipped_device_copies(monkeypatch) -> None: + """Check that every buffer copy_to_device_if_changed skips holds the values. + + A test in which anything else writes such a buffer then fails. + """ + monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", + True) + + @pytest.fixture(scope="module") def mpi_pool_executor(request): """ From bea60022cbfbfe9f30b7fc936a1239b14b0f1cd7 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 23:35:05 -0700 Subject: [PATCH 43/67] [None][perf] spec_step_copies: run on the SM 100 family (SM 100, 103, 107) The step-copy kernels (SlotScatter and StepInputStage) use only thread and block indices and global loads and stores, so is_supported() now admits SM 103 and SM 107 as well as SM 100; before, GB300 and Rubin fell back to the torch ops. The kernel tests and the engine's staged-step test skip on the same condition, so they now run there too. Validated on SM 100 (GB200) only. A TODO at the gate says so: no SM 103 or SM 107 machine was available. Signed-off-by: Vasanth Sabavat --- .../_torch/cute_dsl_kernels/spec_step_copies/op.py | 13 +++++++++++-- .../cute_dsl_kernels/test_spec_step_copies.py | 2 +- .../_torch/executor/test_pytorch_model_engine.py | 2 +- 3 files changed, 13 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py index 9bc81fef63bb..4629879fded8 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py @@ -36,8 +36,17 @@ def is_supported() -> bool: - """Whether the kernels run on the current device: SM 100 with the CuTe DSL installed.""" - return IS_CUTLASS_DSL_AVAILABLE and torch.cuda.is_available() and get_sm_version() == 100 + """Whether the kernels run on the current device: the SM 100 family (SM 100, 103, 107) with the CuTe DSL installed. + + The kernels are plain SIMT code (thread and block indices, global loads and stores), with nothing specific to an + architecture. + """ + # TODO: validated on SM 100 (B200 / GB200) only; SM 103 (GB300) and SM 107 (Rubin) are untested. + return ( + IS_CUTLASS_DSL_AVAILABLE + and torch.cuda.is_available() + and get_sm_version() in (100, 103, 107) + ) def _is_i32(t: torch.Tensor) -> bool: diff --git a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py index a2a44221d1da..e497791f2f19 100644 --- a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py +++ b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py @@ -75,7 +75,7 @@ def _supported() -> bool: pytestmark = pytest.mark.skipif( - not _supported(), reason="the step-copy kernels run on SM 100 with the CuTe DSL" + not _supported(), reason="the step-copy kernels run on SM 100 / 103 / 107 with the CuTe DSL" ) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 0a9c19984751..74739d6c37c8 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1741,7 +1741,7 @@ def _check_staged_spec_decode_graph_step( from tensorrt_llm.llmapi.llm_args import \ KvCacheConfig as LlmKvCacheConfig if not spec_step_copies.is_supported(): - self.skipTest("the step-copy kernels run on SM 100") + self.skipTest("the step-copy kernels run on SM 100 / 103 / 107") max_draft_len = 3 tokens_per_step = max_draft_len + 1 model_engine, kv_cache_manager = create_model_engine_and_kvcache( From 0e299645a334e83eaa0ceac2b9723931dfc149a6 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 23:41:12 -0700 Subject: [PATCH 44/67] [None][perf] spec_step_copies: drop the staged host-to-device copies, keep the gathers The decode step no longer stages its host-to-device copies on the step-copy kernel. Staging measured within noise of copying them one by one, and it coupled the engine, TrtllmAttentionMetadata and KVCacheManagerV2 through a 52-argument kernel, a pinned host record guarded by an event, h2d_stage and stage_batch_block_offsets. - StepInputStage becomes StepInputGather: one launch of the overlap scheduler's four gathers (18 arguments), with no begin/copy/ block_copy/commit protocol and no pinned record. The kernel loses its copy and block-copy branches. - TrtllmAttentionMetadata loses h2d_stage, _copy_step_input and _copy_block_offsets; prompt lengths, KV lengths and block offsets are copied as on main. - KVCacheManagerV2.stage_batch_block_offsets is removed. - PyTorchModelEngine copies the positions and runs attn_metadata.prepare() as on main, and launches the gathers where they happen. - SlotScatter is unchanged. Tests: the staged-copy and block-copy cases leave test_spec_step_copies.py; the graph replays capture scatter and gather. The engine's two staged-step tests become one test of the gather kernel against the torch path (l0_b200.yml follows). Signed-off-by: Vasanth Sabavat --- .../_torch/attention/backends/trtllm.py | 57 +-- .../cute_dsl_kernels/spec_step_copies/op.py | 185 +--------- .../spec_step_copies_kernel.py | 185 +--------- .../kv_cache/kv_cache_manager_v2.py | 37 -- .../_torch/pyexecutor/model_engine.py | 47 +-- .../test_lists/test-db/l0_b200.yml | 6 +- .../cute_dsl_kernels/test_spec_step_copies.py | 343 +++--------------- .../executor/test_pytorch_model_engine.py | 68 ++-- 8 files changed, 142 insertions(+), 786 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 3b835cfc3ef0..249aa78ca087 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -25,7 +25,6 @@ if TYPE_CHECKING: from tensorrt_llm.mapping import Mapping - from ...cute_dsl_kernels.spec_step_copies.op import StepInputStage from ...model_config import ModelConfig from ...speculative.interface import SpecMetadata from ...speculative.spec_tree_manager import SpecTreeManager @@ -124,10 +123,6 @@ class TrtllmAttentionMetadata(AttentionMetadata): workspace: Optional[torch.Tensor] = None cuda_graph_workspace: Optional[torch.Tensor] = None workspace_reclaimable: bool = field(default=True, init=False) - # Set by the model engine around prepare() on a CUDA graph decode step: - # the per-step host-to-device copies are staged on it (a StepInputStage) - # and written by its commit, instead of being copied here. - h2d_stage: Optional["StepInputStage"] = field(default=None, init=False) # TrtllmAttention needs to know the beam width to access to the cache indirection buffer, # when beam search is enabled. @@ -878,29 +873,6 @@ def restore_after_draft_forward(self, saved_state: dict | None) -> None: """Restore backend state modified for draft-forward execution.""" return None - def _copy_block_offsets(self, manager, dst: torch.Tensor, - max_blocks: Optional[int]) -> None: - """``manager``'s block-offset copy for this batch, staged on - ``h2d_stage`` when the engine set one and the manager can stage it - (``stage_batch_block_offsets``).""" - stage_copy = getattr(manager, "stage_batch_block_offsets", None) - if self.h2d_stage is not None and stage_copy is not None and stage_copy( - self.h2d_stage, dst, self.request_ids, self.beam_width, - self.num_contexts, self.num_seqs): - return - manager.copy_batch_block_offsets(dst, - self.request_ids, - self.beam_width, - self.num_contexts, - self.num_seqs, - max_blocks=max_blocks) - - def _copy_step_input(self, dst: torch.Tensor, src: torch.Tensor) -> None: - """``dst.copy_(src)`` for a per-step input, staged on ``h2d_stage`` - when the engine set one.""" - if self.h2d_stage is None or not self.h2d_stage.copy(dst, src): - dst.copy_(maybe_pin_memory(src), non_blocking=True) - def prepare(self) -> None: super().prepare() # Recomputed on first use this iteration; see mla_prepare_scheduler_buffers. @@ -929,8 +901,8 @@ def prepare(self) -> None: device='cpu', ) self.prompt_lens_cpu[:self.num_seqs].copy_(prompt_lens) - self._copy_step_input(self.prompt_lens_cuda[:self.num_seqs], - self.prompt_lens_cpu[:self.num_seqs]) + self.prompt_lens_cuda[:self.num_seqs].copy_( + self.prompt_lens_cpu[:self.num_seqs], non_blocking=True) # number of tokens in the kv cache for each sequence in the batch cached_token_lens = torch.tensor( @@ -983,8 +955,9 @@ def prepare(self) -> None: # the sequence length including the cached tokens and the input tokens. self.kv_lens[:self.num_seqs].copy_( kv_lens + self.kv_cache_params.num_extra_kv_tokens) - self._copy_step_input(self.kv_lens_cuda[:self.num_seqs], - kv_lens[:self.num_seqs]) + self.kv_lens_cuda[:self.num_seqs].copy_(maybe_pin_memory( + kv_lens[:self.num_seqs]), + non_blocking=True) # total kv lens for context requests and generation requests, without extra tokens self.host_total_kv_lens[0] = kv_lens[:self.num_contexts].sum().item() self.host_total_kv_lens[1] = kv_lens[self.num_contexts:self. @@ -1022,14 +995,24 @@ def prepare(self) -> None: if not spec_active and self.kv_cache_manager.tokens_per_block: max_blocks = ceil_div(max_kv_len, self.kv_cache_manager.tokens_per_block) - self._copy_block_offsets(self.kv_cache_manager, - self.kv_cache_block_offsets, max_blocks) + self.kv_cache_manager.copy_batch_block_offsets( + self.kv_cache_block_offsets, + self.request_ids, + self.beam_width, + self.num_contexts, + self.num_seqs, + max_blocks=max_blocks) # Also prepare draft KV cache block offsets if draft_kv_cache_manager exists if self.draft_kv_cache_manager is not None: - self._copy_block_offsets(self.draft_kv_cache_manager, - self.draft_kv_cache_block_offsets, - max_blocks) + # Use the wrapper method which works for both V1 and V2 + self.draft_kv_cache_manager.copy_batch_block_offsets( + self.draft_kv_cache_block_offsets, + self.request_ids, + self.beam_width, + self.num_contexts, + self.num_seqs, + max_blocks=max_blocks) # Don't pass self.kv_lens as kv_lens here because it includes extra # tokens. Use the actual KV length (without extra tokens) for diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py index 4629879fded8..88208283e6cb 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py @@ -15,14 +15,13 @@ """Host side of the one-model speculative decoding step's copy kernels (``spec_step_copies_kernel``). ``SlotScatter`` moves a step's per-row outputs into the sampler's slot stores (one launch for the sampler's four -``index_copy_``). ``StepInputStage`` writes a decode step's per-step inputs with one launch: the overlap scheduler's -gathers from those stores, host-to-device copies whose values the kernel reads from a pinned host record, and KV cache -block-offset copies. Each kernel takes every size and address as a launch argument, so it is compiled once per process, -on the first launch, which must happen outside CUDA-graph capture. The kernels are called through TVM-FFI: a launch -costs a few microseconds of host time. +``index_copy_``). ``StepInputGather`` performs the overlap scheduler's gathers from those stores into a decode step's +inputs (one launch for the engine's four). Each kernel takes every size and address as a launch argument, so it is +compiled once per process, on the first launch, which must happen outside CUDA-graph capture. The kernels are called +through TVM-FFI: a launch costs a few microseconds of host time. Callers use them only where ``is_supported()`` holds. A call whose arguments are outside what the kernel covers launches -and stages nothing and returns False; the caller then does that work itself. +nothing and returns False; the caller then does that work itself. """ from __future__ import annotations @@ -31,7 +30,7 @@ import torch -from ...._utils import get_sm_version, prefer_pinned +from ...._utils import get_sm_version from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE @@ -54,11 +53,6 @@ def _is_i32(t: torch.Tensor) -> bool: return t.is_cuda and t.dtype == torch.int32 and t.is_contiguous() -def _is_pinned_i32(t: torch.Tensor) -> bool: - """A contiguous pinned host int32 tensor.""" - return not t.is_cuda and t.dtype == torch.int32 and t.is_contiguous() and t.is_pinned() - - def _compile(name: str, args: tuple): """Compile kernel ``name`` for its scalar ``args`` with TVM-FFI; a call then takes ``args`` and the CUDA stream handle to launch on.""" @@ -200,41 +194,11 @@ def scatter( return True -class StepInputStage: - """One decode step's per-step device inputs, written by one ``stage_kernel`` launch. - - A step calls ``begin()``, then stages the overlap gathers (``gather``), host-to-device copies whose values go into a - pinned host record that the kernel reads in place (``copy``, up to ``max_copies``) and KV cache block-offset copies - (``block_copy``, up to ``max_block_copies``), then ``commit()`` launches one kernel on the current stream that - performs all of them. The record is reused every step: the first ``copy`` after a commit that staged copies waits - for that commit's kernel (in the overlap loop it has run by then: the host is at most one step ahead). - - Args: - capacity: The record's size in int32 values, the most a step's copies stage in total. - """ +class StepInputGather: + """The overlap scheduler's gathers of a decode step's inputs from the sampler's slot stores, as one kernel.""" _kernel: ClassVar[Any] = None - def __init__(self, capacity: int) -> None: - from . import spec_step_copies_kernel as kernel - - self.max_copies = kernel.STAGED_COPIES - self.max_block_copies = kernel.BLOCK_COPIES - self.record = torch.empty((capacity,), dtype=torch.int32, pin_memory=prefer_pinned()) - self._copies: list[tuple[int, int, int]] = [] - self._block_copies: list[tuple[int, ...]] = [] - self._gather: tuple[int, ...] | None = None - self._used = 0 - self._done: torch.cuda.Event | None = None - self._record_free = True - - def begin(self) -> None: - """Start a step: drop anything staged by a step that did not commit (its preparation raised).""" - self._copies = [] - self._block_copies = [] - self._gather = None - self._used = 0 - def gather( self, store_next_new_tokens: torch.Tensor, @@ -254,7 +218,7 @@ def gather( kv_offsets: torch.Tensor, kv_begin: int, ) -> bool: - """Stage the overlap scheduler's gathers of the rows whose request ran in the previous step. + """Launch the gathers of the rows whose request ran in the previous step on the current stream. For ``r < rows`` with ``s = slots[r]`` and ``j < tokens_per_row``: ``input_ids[input_begin + r * tokens_per_row + j] = store_next_new_tokens[j, s, 0]``, @@ -264,8 +228,7 @@ def gather( stores are [width, slots, 1], [slots, width] and [slots]. Returns: - False, staging nothing, when an argument is outside what the kernel covers or a gather is already - staged; True otherwise. + False, launching nothing, when an argument is outside what the kernel covers; True otherwise. """ tensors = ( store_next_new_tokens, @@ -279,8 +242,7 @@ def gather( kv_offsets, ) if not ( - self._gather is None - and all(_is_i32(t) for t in tensors) + all(_is_i32(t) for t in tensors) and _gather_shapes_ok( store_next_new_tokens, store_next_draft, @@ -300,7 +262,7 @@ def gather( return False if rows == 0: return True - self._gather = ( + args = ( store_next_new_tokens.data_ptr(), store_next_draft.data_ptr(), store_lens.data_ptr(), @@ -320,124 +282,7 @@ def gather( pos_begin, kv_begin, ) + if StepInputGather._kernel is None: + StepInputGather._kernel = _compile("gather", args) + StepInputGather._kernel(*args, torch.cuda.current_stream().cuda_stream) return True - - def copy(self, dst: torch.Tensor, values: torch.Tensor) -> bool: - """Stage ``dst[:n] = values`` for the ``n`` int32 host ``values``; ``dst`` is a 1-D int32 device tensor. - - Returns: - False, staging nothing, when ``dst`` is not a contiguous 1-D CUDA int32 tensor of at least ``n`` - values, the step already staged ``max_copies`` copies or the record cannot hold the values; True - otherwise. - """ - host = values.reshape(-1) - n = host.numel() - if not ( - _is_i32(dst) - and dst.dim() == 1 - and n <= dst.numel() - and host.dtype == torch.int32 - and not host.is_cuda - and len(self._copies) < self.max_copies - and self._used + n <= self.record.numel() - and self.record.is_pinned() - ): - return False - if n == 0: - return True - if not self._record_free: - self._done.synchronize() - self._record_free = True - self.record[self._used : self._used + n].copy_(host) - self._copies.append((dst.data_ptr(), self._used, n)) - self._used += n - return True - - def block_copy( - self, - offsets: torch.Tensor, - table: torch.Tensor, - copy_index: torch.Tensor, - index_scales: torch.Tensor, - kv_offset: torch.Tensor, - ) -> bool: - """Stage a KV cache manager's block-offset copy (``copy_batch_block_offsets_to_device``). - - Args: - offsets: The device block offsets, int32 [pools, seqs_cap, 2, blocks]. - table: The pinned host page table, int32 [pools, table_seqs, 2, blocks]. - copy_index: The table row of each of the step's sequences, pinned host int32 [seqs]. - index_scales: The per-pool page index scale, pinned host int32 [pools]. - kv_offset: The per-pool V offset, pinned host int32 [pools]. - - For every pool and sequence the K row is ``index_scale * page`` and the V row ``index_scale * page + - kv_offset`` over the table row ``copy_index[sequence]``; a bad page (-1) gives 0. - - Returns: - False, staging nothing, when an argument is outside what the kernel covers or the step already staged - ``max_block_copies`` block copies; True otherwise. - """ - if not ( - _is_i32(offsets) - and offsets.dim() == 4 - and table.dim() == 4 - and all(_is_pinned_i32(t) for t in (table, copy_index, index_scales, kv_offset)) - and len(self._block_copies) < self.max_block_copies - ): - return False - pools, table_seqs, kv, blocks = table.shape - seqs = copy_index.numel() - if not ( - kv == 2 - and offsets.shape[0] >= pools - and offsets.shape[1] >= seqs - and offsets.shape[2] == 2 - and offsets.shape[3] == blocks - and index_scales.numel() >= pools - and kv_offset.numel() >= pools - ): - return False - if pools * seqs * blocks == 0: - return True - self._block_copies.append( - ( - table.data_ptr(), - offsets.data_ptr(), - copy_index.data_ptr(), - index_scales.data_ptr(), - kv_offset.data_ptr(), - pools, - table_seqs, - offsets.shape[1], - blocks, - seqs, - ) - ) - return True - - def commit(self) -> None: - """Launch the staged gathers and copies as one kernel on the current stream, then start the next staging.""" - if self._gather is None and not self._copies and not self._block_copies: - return - gather = self._gather or (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 0, 0, 0, 0) - copies = self._copies + [(0, 0, 0)] * (self.max_copies - len(self._copies)) - blocks = self._block_copies + [(0, 0, 0, 0, 0, 0, 0, 0, 0, 0)] * ( - self.max_block_copies - len(self._block_copies) - ) - args = ( - *gather, - self.record.data_ptr(), - *(c[0] for c in copies), - *(c[1] for c in copies), - *(c[2] for c in copies), - *(v for b in blocks for v in b), - ) - if StepInputStage._kernel is None: - StepInputStage._kernel = _compile("stage", args) - StepInputStage._kernel(*args, torch.cuda.current_stream().cuda_stream) - if self._copies: - if self._done is None: - self._done = torch.cuda.Event() - self._done.record() - self._record_free = False - self.begin() diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py index 269da57f8d01..90378e25ed00 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py @@ -17,11 +17,9 @@ * ``scatter_kernel``: the sampler moves the forward's per-row outputs into its slot-indexed stores (``index_copy_`` of new tokens, next new tokens, accepted lengths and next draft tokens, each padded or cut to its store width; a row's new tokens past its accepted length are written as zeros). -* ``stage_kernel``: one launch writes a decode step's per-step inputs: the overlap scheduler's gathers from those - stores for the rows whose request ran in the previous step (``index_select`` into input ids and draft tokens by - slot, into the position offsets by the per-token index list, and into the KV-length offsets by slot, minus the - tokens per step), up to four host-to-device copies whose values it reads in place from a pinned host record, and up - to two KV cache block-offset copies. +* ``gather_kernel``: the overlap scheduler's gathers from those stores into a decode step's inputs, for the rows whose + request ran in the previous step (``index_select`` into input ids and draft tokens by slot, into the position + offsets by the per-token index list, and into the KV-length offsets by slot, minus the tokens per step). Every buffer is int32 and passed as a raw address, so a call passes integers only. Row and element offsets are arguments. One thread per element. @@ -37,7 +35,7 @@ def _i32_at(address, index): - """A one-element int32 tensor at ``address + 4 * index`` (trace-time helper; global or mapped host memory).""" + """A one-element int32 tensor at ``address + 4 * index`` in global memory (trace-time helper).""" ptr = cute.make_ptr( cutlass.Int32, address + cutlass.Int64(index) * cutlass.Int64(4), @@ -134,12 +132,8 @@ def scatter( ) # fmt: skip -STAGED_COPIES = 4 -BLOCK_COPIES = 2 - - @cute.kernel -def stage_kernel( +def gather_kernel( store_next_new_tokens: cutlass.Int64, # int32 [next_width, num_slots] (token-major) store_next_draft: cutlass.Int64, # int32 [num_slots, draft_stride] store_lens: cutlass.Int64, # int32 [num_slots] @@ -158,54 +152,15 @@ def stage_kernel( draft_begin: cutlass.Int32, pos_begin: cutlass.Int32, kv_begin: cutlass.Int32, - record: cutlass.Int64, # int32 values in pinned host memory, read in place - dst0: cutlass.Int64, - dst1: cutlass.Int64, - dst2: cutlass.Int64, - dst3: cutlass.Int64, - off0: cutlass.Int32, - off1: cutlass.Int32, - off2: cutlass.Int32, - off3: cutlass.Int32, - n0: cutlass.Int32, - n1: cutlass.Int32, - n2: cutlass.Int32, - n3: cutlass.Int32, - table0: cutlass.Int64, # block copy 0: int32 [pools, table_seqs, 2, blocks] in pinned host memory - offsets0: cutlass.Int64, # int32 [pools, offsets_seqs, 2, blocks] (device) - copy_index0: cutlass.Int64, # int32 [seqs], pinned host - index_scales0: cutlass.Int64, # int32 [pools], pinned host - kv_offset0: cutlass.Int64, # int32 [pools], pinned host - pools0: cutlass.Int32, - table_seqs0: cutlass.Int32, - offsets_seqs0: cutlass.Int32, - blocks0: cutlass.Int32, - seqs0: cutlass.Int32, - table1: cutlass.Int64, - offsets1: cutlass.Int64, - copy_index1: cutlass.Int64, - index_scales1: cutlass.Int64, - kv_offset1: cutlass.Int64, - pools1: cutlass.Int32, - table_seqs1: cutlass.Int32, - offsets_seqs1: cutlass.Int32, - blocks1: cutlass.Int32, - seqs1: cutlass.Int32, ): - """The overlap gathers for threads [0, rows * tokens_per_row) (for row r, slot s = slots[r] and token j: - ``input_ids[input_begin + r * tokens_per_row + j] = store_next_new_tokens[j, s]``, ``pos_offsets[pos_begin + r * - tokens_per_row + j] = store_lens[pos_indices[r * tokens_per_row + j]]``, ``draft_tokens[draft_begin + r * - draft_width + j] = store_next_draft[s, j]`` for j < draft_width, ``kv_offsets[kv_begin + r] = store_lens[s] - - tokens_per_row``); then up to four staged copies ``dst_i[j] = record[off_i + j]`` for j < n_i; then up to two KV - block-offset copies, each the - ``copyBatchBlockOffsetsToDeviceKernel`` of kvCacheManagerV2Utils.cu (the K and V rows of every (pool, sequence) - from the pinned host table's row ``copy_index[sequence]``: ``index_scale * page``, and ``+ kv_offset`` for V; a - bad page (-1) gives 0). One thread per int32 (per K / V pair for the block copies).""" + """For row r < rows, slot s = slots[r] and token j < tokens_per_row: ``input_ids[input_begin + r * tokens_per_row + + j] = store_next_new_tokens[j, s]``, ``pos_offsets[pos_begin + r * tokens_per_row + j] = store_lens[pos_indices[r + * tokens_per_row + j]]``, ``draft_tokens[draft_begin + r * draft_width + j] = store_next_draft[s, j]`` for j < + draft_width and ``kv_offsets[kv_begin + r] = store_lens[s] - tokens_per_row``. One thread per (row, token).""" tidx, _, _ = cute.arch.thread_idx() bidx, _, _ = cute.arch.block_idx() t = bidx * cutlass.Int32(THREADS) + tidx - gathered = rows * tokens_per_row - if t < gathered: + if t < rows * tokens_per_row: row = t // tokens_per_row col = t - row * tokens_per_row slot = _i32_at(slots, row)[0] @@ -219,78 +174,10 @@ def stage_kernel( )[0] if col == cutlass.Int32(0): _i32_at(kv_offsets, kv_begin + row)[0] = _i32_at(store_lens, slot)[0] - tokens_per_row - else: - u = t - gathered - dst = cutlass.Int64(0) - src = cutlass.Int32(0) - j = cutlass.Int32(0) - active = cutlass.Int32(0) - if u < n0: - dst = dst0 - src = off0 + u - j = u - active = cutlass.Int32(1) - elif u < n0 + n1: - dst = dst1 - src = off1 + u - n0 - j = u - n0 - active = cutlass.Int32(1) - elif u < n0 + n1 + n2: - dst = dst2 - src = off2 + u - n0 - n1 - j = u - n0 - n1 - active = cutlass.Int32(1) - elif u < n0 + n1 + n2 + n3: - dst = dst3 - src = off3 + u - n0 - n1 - n2 - j = u - n0 - n1 - n2 - active = cutlass.Int32(1) - if active == cutlass.Int32(1): - _i32_at(dst, j)[0] = _i32_at(record, src)[0] - v = u - n0 - n1 - n2 - n3 - per0 = pools0 * seqs0 * blocks0 - per1 = pools1 * seqs1 * blocks1 - if v >= cutlass.Int32(0): - if v < per0 + per1: - table = table0 - offsets = offsets0 - copy_index = copy_index0 - index_scales = index_scales0 - kv_offset = kv_offset0 - table_seqs = table_seqs0 - offsets_seqs = offsets_seqs0 - blocks = blocks0 - seqs = seqs0 - w = v - if v >= per0: - table = table1 - offsets = offsets1 - copy_index = copy_index1 - index_scales = index_scales1 - kv_offset = kv_offset1 - table_seqs = table_seqs1 - offsets_seqs = offsets_seqs1 - blocks = blocks1 - seqs = seqs1 - w = v - per0 - pool = w // (seqs * blocks) - rest = w - pool * seqs * blocks - seq = rest // blocks - block = rest - seq * blocks - row = pool * table_seqs + _i32_at(copy_index, seq)[0] - page = _i32_at(table, row * cutlass.Int32(2) * blocks + block)[0] - key = cutlass.Int32(0) - value = cutlass.Int32(0) - if page != cutlass.Int32(-1): - key = _i32_at(index_scales, pool)[0] * page - value = key + _i32_at(kv_offset, pool)[0] - out = (pool * offsets_seqs + seq) * cutlass.Int32(2) * blocks + block - _i32_at(offsets, out)[0] = key - _i32_at(offsets, out + blocks)[0] = value @cute.jit -def stage( +def gather( store_next_new_tokens: cutlass.Int64, store_next_draft: cutlass.Int64, store_lens: cutlass.Int64, @@ -309,58 +196,14 @@ def stage( draft_begin: cutlass.Int32, pos_begin: cutlass.Int32, kv_begin: cutlass.Int32, - record: cutlass.Int64, - dst0: cutlass.Int64, - dst1: cutlass.Int64, - dst2: cutlass.Int64, - dst3: cutlass.Int64, - off0: cutlass.Int32, - off1: cutlass.Int32, - off2: cutlass.Int32, - off3: cutlass.Int32, - n0: cutlass.Int32, - n1: cutlass.Int32, - n2: cutlass.Int32, - n3: cutlass.Int32, - table0: cutlass.Int64, - offsets0: cutlass.Int64, - copy_index0: cutlass.Int64, - index_scales0: cutlass.Int64, - kv_offset0: cutlass.Int64, - pools0: cutlass.Int32, - table_seqs0: cutlass.Int32, - offsets_seqs0: cutlass.Int32, - blocks0: cutlass.Int32, - seqs0: cutlass.Int32, - table1: cutlass.Int64, - offsets1: cutlass.Int64, - copy_index1: cutlass.Int64, - index_scales1: cutlass.Int64, - kv_offset1: cutlass.Int64, - pools1: cutlass.Int32, - table_seqs1: cutlass.Int32, - offsets_seqs1: cutlass.Int32, - blocks1: cutlass.Int32, - seqs1: cutlass.Int32, stream: cuda_driver.CUstream, ) -> None: - total = ( - rows * tokens_per_row - + n0 - + n1 - + n2 - + n3 - + pools0 * seqs0 * blocks0 - + pools1 * seqs1 * blocks1 - ) - stage_kernel( + gather_kernel( store_next_new_tokens, store_next_draft, store_lens, slots, pos_indices, input_ids, draft_tokens, pos_offsets, kv_offsets, rows, tokens_per_row, draft_width, num_slots, draft_stride, input_begin, draft_begin, pos_begin, - kv_begin, record, dst0, dst1, dst2, dst3, off0, off1, off2, off3, n0, n1, n2, n3, table0, offsets0, - copy_index0, index_scales0, kv_offset0, pools0, table_seqs0, offsets_seqs0, blocks0, seqs0, table1, offsets1, - copy_index1, index_scales1, kv_offset1, pools1, table_seqs1, offsets_seqs1, blocks1, seqs1, + kv_begin, ).launch( - grid=[(total + THREADS - 1) // THREADS, 1, 1], + grid=[(rows * tokens_per_row + THREADS - 1) // THREADS, 1, 1], block=[THREADS, 1, 1], stream=stream, ) # fmt: skip diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 878d236ed231..737d42c05936 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -6109,43 +6109,6 @@ def update_resources( ) self._allocated_draft_lens.pop(req.py_request_id, None) - def stage_batch_block_offsets( - self, - stage, - dst_tensor: torch.Tensor, - request_ids: List[int], - beam_width: int, - num_contexts: int, - num_seqs: int, - ) -> bool: - """Stage ``copy_batch_block_offsets`` on a StepInputStage instead of launching its kernel. - - Returns False, staging nothing, where that copy is not this class's single-table copy on the current - stream: the per-layer page-table layout, a subclass that overrides ``copy_batch_block_offsets``, a manager - whose stream is not the current one, or a copy the stage does not cover. The caller then copies as usual. - """ - if ( - self._use_per_layer_page_tables - or type(self).copy_batch_block_offsets is not KVCacheManagerV2.copy_batch_block_offsets - or self._stream != torch.cuda.current_stream() - ): - return False - assert beam_width <= self.max_copy_beam_width - copy_idx = self.index_mapper.get_copy_index( - request_ids, - num_contexts, - beam_width, - self.kv_cache_type == CacheTypeCpp.CROSS, - ) - assert copy_idx.shape[0] == num_seqs - return stage.block_copy( - dst_tensor, - self.host_kv_cache_block_offsets, - copy_idx, - self.index_scales, - self.kv_offset, - ) - def copy_batch_block_offsets( self, dst_tensor: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index d4573691c7c6..7d7f29ed13db 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -887,12 +887,11 @@ def _init_decoder_state(self) -> None: attn_backend=self.attn_backend, cuda_graph_manager=None) self._prepare_inputs_event: Optional[torch.cuda.Event] = None - # Where the step-copy kernels run: writes a speculative decode step's - # per-step inputs with one launch. Positions (up to max_num_tokens) and - # two per-sequence lengths fit its record. - self._step_input_stage: Optional[spec_step_copies.StepInputStage] = ( - spec_step_copies.StepInputStage(3 * self.max_num_tokens) if - self.is_spec_decode and spec_step_copies.is_supported() else None) + # Where the step-copy kernels run: performs a speculative decode step's + # overlap gathers with one launch. + self._step_input_gather: Optional[spec_step_copies.StepInputGather] = ( + spec_step_copies.StepInputGather() if self.is_spec_decode + and spec_step_copies.is_supported() else None) # Let the first CUDA graph capture create its private pool. Piecewise # CUDA graphs use a separate pool owned by their runners, so sharing a # pre-created pool handle with the outer graph runner is unnecessary. @@ -5578,18 +5577,8 @@ def previous_seq_slots_device(): num_draft_tokens = len(draft_tokens) total_num_tokens = len(position_ids) # Where the step-copy kernels run, the overlap gathers below are one - # kernel. On a CUDA graph step whose attention metadata is a plain - # TrtllmAttentionMetadata (its prepare() reads none of these inputs on - # the device), the gathers, the positions and the metadata's prompt - # lengths, KV lengths and block offsets are staged instead and written - # by one launch right after attn_metadata.prepare(). - step_copies = (self._step_input_stage if enable_spec_decode else None) - if step_copies is not None: - step_copies.begin() - stage = (step_copies if step_copies is not None - and attn_metadata.is_cuda_graph and not self.use_mrope - and type(attn_metadata) is TrtllmAttentionMetadata - and attn_metadata.fp4_mla_state is None else None) + # kernel. + step_gather = (self._step_input_gather if enable_spec_decode else None) assert total_num_tokens <= self.max_num_tokens, ( f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self.max_num_tokens})" ) @@ -5654,7 +5643,7 @@ def previous_seq_slots_device(): # Zeroed above. previous_begin = (num_extend_reqeust_wo_dummy - previous_batch_len) - if step_copies is not None and step_copies.gather( + if step_gather is None or not step_gather.gather( new_tokens_device, next_draft_tokens_device, new_tokens_lens_device, self.previous_batch_indices_cuda, @@ -5665,9 +5654,6 @@ def previous_seq_slots_device(): self.previous_pos_id_offsets_cuda, previous_begin * runtime_tokens_per_gen_step, self.previous_kv_lens_offsets_cuda, previous_begin): - if stage is None: - step_copies.commit() - else: # previous input ids new_tokens = new_tokens_device.transpose(0, 1)[ previous_slots, :runtime_tokens_per_gen_step].flatten() @@ -5765,11 +5751,8 @@ def previous_seq_slots_device(): final_position_ids = self.mrope_position_ids_cuda[:, :, : total_num_tokens] else: - if stage is None or not stage.copy( - self.position_ids_cuda[:total_num_tokens], - host_position_ids): - self.position_ids_cuda[:total_num_tokens].copy_( - host_position_ids, non_blocking=True) + self.position_ids_cuda[:total_num_tokens].copy_(host_position_ids, + non_blocking=True) final_position_ids = self.position_ids_cuda[: total_num_tokens].unsqueeze( 0) @@ -5852,15 +5835,7 @@ def previous_seq_slots_device(): # pre-prepare counts so the steady-gen recording below stores values # that the per-step prepare() can re-clamp from scratch. num_cached_tokens_snapshot = list(num_cached_tokens_per_seq) - if stage is None: - attn_metadata.prepare() - else: - attn_metadata.h2d_stage = stage - try: - attn_metadata.prepare() - finally: - attn_metadata.h2d_stage = None - stage.commit() + attn_metadata.prepare() cross_attention_inputs = (self._prepare_enc_dec_cross_attn_inputs( cross_encoder_hidden_states, cross_encoder_seq_lens, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index b71146eb7f05..d7ddb7ff9e89 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -243,10 +243,10 @@ l0_b200: # accuracy deviation found so far is Blackwell-only), so it is not # hw-agnostic and needs its own Blackwell run. - unittest/_torch/speculative/test_fused_sampling_op.py - # One-model speculative decoding step-copy kernels and the engine's staged - # decode-step inputs (SM 100 only; they skip on every other list). + # One-model speculative decoding step-copy kernels and the engine's decode + # step that runs them (SM 100 family; they skip on other GPUs). - unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py - - unittest/_torch/executor/test_pytorch_model_engine.py -k "test_staged_spec_decode_graph_step" + - unittest/_torch/executor/test_pytorch_model_engine.py -k "test_spec_decode_graph_step_gather_kernel" - unittest/_torch/thop/parallel TIMEOUT (90) - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/thop/serial diff --git a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py index e497791f2f19..58f9ba03102e 100644 --- a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py +++ b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py @@ -23,20 +23,15 @@ > 0; shuffled distinct slot tables; accepted lengths from 0 to past the output width. A mixed context / generation batch laid out as the one-model worker writes it (a context row's tokens past column 0 never written) stores zeros there. -* ``StepInputStage.gather`` committed on its own vs PyTorchModelEngine._prepare_tp_inputs's torch overlap gathers (the - stores by slot into the input ids and draft tokens, the lengths by the per-token index list into the position - offsets, ``new_tokens_lens - T`` by slot into the KV-length offsets): the engine's offsets (the requests without a - previous batch first), zero and arbitrary offsets, a draft width below T - 1, the engine's and arbitrary index lists. -* ``StepInputStage`` vs those gathers plus what it stages instead of copying: ``dst.copy_(pinned, non_blocking=True)`` - per staged copy (positions, prompt and KV lengths, a fourth) and the KV cache manager's block-offset copy - (``copy_batch_block_offsets_to_device``, target and draft managers, bad pages included): the engine's step, all four - copy slots at offset destinations, the gathers alone, the copies alone. +* ``StepInputGather.gather`` vs PyTorchModelEngine._prepare_tp_inputs's torch overlap gathers (the stores by slot into + the input ids and draft tokens, the lengths by the per-token index list into the position offsets, + ``new_tokens_lens - T`` by slot into the KV-length offsets): the engine's offsets (the requests without a previous + batch first), zero and arbitrary offsets, a draft width below T - 1, the engine's and arbitrary index lists. * The calls the kernels do not cover launch nothing and return False, so the caller keeps its torch path. Every buffer starts random and is compared whole (what an op must not touch included), every case runs twice (reruns -bit-identical), and one CUDA graph per split family (R = 1 .. 8 at one T) captures each step's scatter, gather and -staged commit; it is replayed with every input (and the staged values in the stages' pinned records) rewritten in -place, bit-identical to the eager ops and to the torch ops on every replay. +bit-identical), and one CUDA graph per split family (R = 1 .. 8 at one T) captures each step's scatter and gather; it +is replayed with every input rewritten in place, bit-identical to the eager ops and to the torch ops on every replay. Table: ``python3 test_spec_step_copies.py report``. Timing: ``python3 test_spec_step_copies.py time`` (CUDA graphs of back-to-back calls, median us per call over 15 replays, each kernel vs the torch ops it replaces, at every split). @@ -53,11 +48,6 @@ STORE_WIDTH = MAX_DRAFT + 1 # the sampler's new_tokens / next_new_tokens stores MAX_NUM_TOKENS = 8192 # the engine's per-token buffers (input ids, positions, previous_*) DRAFT_BUFFER = MAX_DRAFT * MAX_BATCH # draft_tokens_cuda: max_draft_loop_tokens * batch_size -STAGE_CAPACITY = 3 * MAX_NUM_TOKENS # PyTorchModelEngine._get_step_input_stage -BLOCKS = 40 # KV blocks per sequence (copy_batch_block_offsets_to_device needs a multiple of 4) -TABLE_SEQS = 24 # rows of a KV cache manager's pinned host block-offset table -POOLS = (2, 1) # pools of the target and the draft KV cache managers -COPY_NAMES = ("position_ids", "prompt_lens", "kv_lens", "extra") # staged copies, in order ROWS = list(range(1, MAX_BATCH + 1)) TOKENS = [1, 2, 4, 8] @@ -79,33 +69,14 @@ def _supported() -> bool: ) -def _pinned_ok() -> bool: - """StepInputStage reads a pinned host record in place; it refuses to run where pinned memory is not preferred.""" - from tensorrt_llm._utils import prefer_pinned - - return prefer_pinned() - - -_eager = {} - - -def eager_stage(): - """The StepInputStage of the eager runs (never captured: a captured commit's stage cannot stage again).""" - if "stage" not in _eager: - _eager["stage"] = _op().StepInputStage(STAGE_CAPACITY) - return _eager["stage"] - - def kernel_scatter(args) -> None: """``SlotScatter.scatter`` on ``args``, which it must cover.""" assert _op().SlotScatter().scatter(*args), "scatter declined a covered call" -def kernel_gather(stage, args) -> None: - """The overlap gathers on ``args`` (which they must cover) committed on their own through ``stage``.""" - stage.begin() - assert stage.gather(*args), "gather declined a covered call" - stage.commit() +def kernel_gather(args) -> None: + """``StepInputGather.gather`` on ``args``, which it must cover.""" + assert _op().StepInputGather().gather(*args), "gather declined a covered call" # ---------------------------------------------------------------------------------------------------------------- @@ -166,7 +137,7 @@ def torch_gather( previous_kv_lens_offsets_cuda, kv_begin, ): - """PyTorchModelEngine._prepare_tp_inputs's torch overlap gathers (StepInputStage.gather's arguments).""" + """PyTorchModelEngine._prepare_tp_inputs's torch overlap gathers (StepInputGather.gather's arguments).""" tokens = runtime_tokens_per_gen_step width = runtime_draft_token_buffer_width previous_slots = previous_batch_indices_cuda[:previous_batch_len] @@ -190,17 +161,6 @@ def torch_gather( ) -def torch_block_copy(offsets, table, copy_index, index_scales, kv_offset): - """KVCacheManagerV2.copy_batch_block_offsets's copy (StepInputStage.block_copy's arguments), current stream.""" - from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import ( - copy_batch_block_offsets_to_device, - ) - - copy_batch_block_offsets_to_device( - table, offsets, copy_index, index_scales, kv_offset, torch.cuda.current_stream().cuda_stream - ) - - # ---------------------------------------------------------------------------------------------------------------- # One decode step's buffers # ---------------------------------------------------------------------------------------------------------------- @@ -228,19 +188,15 @@ def engine_begins(rows, tokens): class Step: - """One decode step of ``rows`` requests x ``tokens`` tokens on engine-sized int32 buffers, all random so that a - stray or a missing write shows (index buffers hold valid indices past their used part): + """One decode step of ``rows`` requests x ``tokens`` tokens on engine-sized int32 device buffers, all random so + that a stray or a missing write shows (index buffers hold valid indices past their used part): * the forward's outputs: new_tokens [N, a], next_new_tokens [N, b], new_tokens_lens [N], next_draft_tokens [N, c] (default a = b = T, c = T - 1); * the sampler's slot table [slots] and stores new_tokens / next_new_tokens [8, slots, 1], new_tokens_lens [slots], next_draft_tokens [slots, 7]; - * the engine's previous_batch_indices / previous_pos_indices, and two output sets ("a": the gathers committed on - their own, "b": the gathers in a full staged step) of input_ids, draft_tokens [56], previous_pos_id_offsets and - previous_kv_lens_offsets [8] (grown only when an offset needs it); - * the staged copies' destinations (positions, prompt and KV lengths [slots], a fourth buffer) and host values; - * the target and draft KV cache managers' device block offsets [pools, slots, 2, 40] and pinned host tables - [pools, 24, 2, 40] (a fifth of the pages bad, -1), copy indices, index scales and KV offsets. + * the engine's previous_batch_indices / previous_pos_indices and the gathers' outputs input_ids, draft_tokens [56], + previous_pos_id_offsets and previous_kv_lens_offsets [8] (grown only when an offset needs it). ``chain``: the gather reads the scatter's slots in another order (the next step's inputs from this step's stores). """ @@ -259,20 +215,19 @@ def __init__( draft_width=None, pos_list="engine", chain=False, - copy_at=(0, 0, 0, 0), ): self.rows, self.tokens, self.num_slots = rows, tokens, num_slots self.widths = widths or (tokens, tokens, tokens - 1) self.row_begin, self.begins = row_begin, begins self.draft_width = tokens - 1 if draft_width is None else draft_width - self.pos_list, self.chain, self.copy_at = pos_list, chain, copy_at + self.pos_list, self.chain = pos_list, chain self.g = torch.Generator().manual_seed(seed) n_out = max(MAX_BATCH, row_begin + rows) if out_rows is None else out_rows assert n_out >= row_begin + rows n_draft = max(DRAFT_BUFFER, begins[1] + rows * self.draft_width) n_kv = max(MAX_BATCH, begins[3] + rows) a, b, c = self.widths - device = { + shapes = { "out_new": (n_out, a), "out_next": (n_out, b), "out_lens": (n_out,), @@ -284,36 +239,20 @@ def __init__( "st_draft": (num_slots, MAX_DRAFT), "prev_slots": (MAX_NUM_TOKENS,), "prev_pos": (MAX_NUM_TOKENS,), - "position_ids": (MAX_NUM_TOKENS,), - "prompt_lens": (num_slots,), - "kv_lens": (num_slots,), - "extra": (MAX_NUM_TOKENS,), + "input_ids": (MAX_NUM_TOKENS,), + "draft": (n_draft,), + "pos_off": (MAX_NUM_TOKENS,), + "kv_off": (n_kv,), } - host = {} - for s in "ab": - device[f"input_ids_{s}"] = (MAX_NUM_TOKENS,) - device[f"draft_{s}"] = (n_draft,) - device[f"pos_off_{s}"] = (MAX_NUM_TOKENS,) - device[f"kv_off_{s}"] = (n_kv,) - for m, pools in zip("td", POOLS): - device[f"blk_{m}"] = (pools, num_slots, 2, BLOCKS) - host[f"tab_{m}"] = (pools, TABLE_SEQS, 2, BLOCKS) - host[f"cidx_{m}"] = (rows,) - host[f"scale_{m}"] = (pools,) - host[f"kvoff_{m}"] = (pools,) - self.bufs = {k: torch.empty(v, dtype=torch.int32, device="cuda") for k, v in device.items()} - for k, v in host.items(): - self.bufs[k] = torch.empty(v, dtype=torch.int32, pin_memory=True) + self.bufs = {k: torch.empty(v, dtype=torch.int32, device="cuda") for k, v in shapes.items()} self.rewrite() def rewrite(self) -> None: - """New random contents for every buffer and host value, written in place (a captured graph keeps the - addresses).""" + """New random contents for every buffer, written in place (a captured graph keeps the addresses).""" g, rows, tokens, num_slots = self.g, self.rows, self.tokens, self.num_slots b = self.bufs for t in b.values(): - if t.is_cuda: - t.copy_(rand_i32(t.shape, g)) + t.copy_(rand_i32(t.shape, g)) # Accepted lengths from none to past the new-token width, so every row cuts its new tokens somewhere. b["out_lens"].copy_(rand_i32(b["out_lens"].shape, g, 0, self.widths[0] + 2)) slots = shuffled_distinct(rows, num_slots, g) @@ -329,20 +268,6 @@ def rewrite(self) -> None: if self.pos_list == "engine": # each row's slot, once per token per_token[: rows * tokens] = previous.repeat_interleave(tokens) b["prev_pos"].copy_(per_token) - self.values = { - "position_ids": rand_i32((rows * tokens,), g, 0, 1 << 20), - "prompt_lens": rand_i32((rows,), g, 1, 1 << 20), - "kv_lens": rand_i32((rows,), g, 1, 1 << 20), - "extra": rand_i32((rows,), g), - } - self.pinned = {k: v.pin_memory() for k, v in self.values.items()} - for m in "td": - pages = rand_i32(b[f"tab_{m}"].shape, g, 0, 1 << 20) - pages[torch.rand(pages.shape, generator=g) < 0.2] = -1 - b[f"tab_{m}"].copy_(pages) - b[f"cidx_{m}"].copy_(torch.tensor(shuffled_distinct(rows, TABLE_SEQS, g))) - b[f"scale_{m}"].copy_(rand_i32(b[f"scale_{m}"].shape, g, 1, 8)) - b[f"kvoff_{m}"].copy_(rand_i32(b[f"kvoff_{m}"].shape, g, 0, 1 << 16)) def outputs(self, b): return { @@ -357,60 +282,16 @@ def scatter_args(self, b): return (self.outputs(b), self.row_begin, self.rows, b["slot_table"], b["st_new"], b["st_next"], b["st_lens"], b["st_draft"]) # fmt: skip - def gather_args(self, b, out): - """StepInputStage.gather's arguments on the buffers ``b``, into output set ``out``.""" + def gather_args(self, b): + """StepInputGather.gather's arguments on the buffers ``b``.""" ib, db, pb, kb = self.begins return (b["st_next"], b["st_draft"], b["st_lens"], b["prev_slots"], b["prev_pos"], self.rows, self.tokens, - self.draft_width, b[f"input_ids_{out}"], ib, b[f"draft_{out}"], db, b[f"pos_off_{out}"], pb, - b[f"kv_off_{out}"], kb) # fmt: skip - - def copies(self, b, count): - """The first ``count`` staged copies: (destination view, host value name).""" - return [ - (b[name][at : at + self.values[name].numel()], name) - for name, at in zip(COPY_NAMES[:count], self.copy_at) - ] - - def block_copies(self, b): - """StepInputStage.block_copy's arguments for the target and the draft KV cache managers.""" - return [ - (b[f"blk_{m}"], b[f"tab_{m}"], b[f"cidx_{m}"], b[f"scale_{m}"], b[f"kvoff_{m}"]) - for m in "td" - ] - - def run_stage(self, b, stage, gather=True, copies=3, blocks=True): - """The step's inputs through ``stage`` (a StepInputStage): one commit.""" - stage.begin() - if gather: - assert stage.gather(*self.gather_args(b, "b")), "gather declined" - for dst, name in self.copies(b, copies): - assert stage.copy(dst, self.values[name]), "copy declined" - if blocks: - for args in self.block_copies(b): - assert stage.block_copy(*args), "block copy declined" - stage.commit() - - def torch_stage(self, b, gather=True, copies=3, blocks=True): - """What the stage replaces: the gathers, a non-blocking copy from pinned memory per staged copy, and the KV - cache managers' block-offset copies.""" - if gather: - torch_gather(*self.gather_args(b, "b")) - for dst, name in self.copies(b, copies): - dst.copy_(self.pinned[name], non_blocking=True) - if blocks: - for args in self.block_copies(b): - torch_block_copy(*args) + self.draft_width, b["input_ids"], ib, b["draft"], db, b["pos_off"], pb, b["kv_off"], kb) # fmt: skip def clone_bufs(bufs): - """A copy of every buffer (pinned host buffers stay pinned).""" - out = {} - for k, t in bufs.items(): - if t.is_cuda: - out[k] = t.clone() - else: - out[k] = torch.empty(t.shape, dtype=t.dtype, pin_memory=t.is_pinned()).copy_(t) - return out + """A copy of every buffer.""" + return {k: t.clone() for k, t in bufs.items()} def mismatches(x, y): @@ -476,21 +357,6 @@ def gather_cases(rows, tokens): ] # fmt: skip -def stage_cases(rows, tokens): - """(label, Step arguments, run_stage arguments).""" - begins = engine_begins(rows, tokens) - return [ - (f"engine step: gathers at {begins}, positions / prompt / KV lengths, target + draft block offsets", - dict(num_slots=MAX_BATCH, begins=begins), dict()), - ("4 copies to offset destinations, gathers at (37, 5, 11, 3), block offsets, 16 slots", - dict(num_slots=2 * MAX_BATCH, begins=(37, 5, 11, 3), pos_list="any", copy_at=(7, 1, 2, 3)), - dict(copies=4)), - ("gathers only, 16 slots", dict(num_slots=2 * MAX_BATCH), dict(copies=0, blocks=False)), - ("copies + block offsets only (no previous batch), 16 slots", - dict(num_slots=2 * MAX_BATCH, copy_at=(5, 3, 0, 0)), dict(gather=False)), - ] # fmt: skip - - def measure_scatter(rows, tokens): out = [] for i, (case, kw) in enumerate(scatter_cases(rows, tokens)): @@ -505,91 +371,53 @@ def measure_scatter(rows, tokens): def measure_gather(rows, tokens): - stage = eager_stage() out = [] for i, (case, kw) in enumerate(gather_cases(rows, tokens)): st = Step(rows, tokens, _seed(2, rows, tokens, i), **kw) bad = compare( st.bufs, - lambda b: torch_gather(*st.gather_args(b, "a")), - lambda b: kernel_gather(stage, st.gather_args(b, "a")), - ) - out.append(result(case, **bad)) - return out - - -def measure_stage(rows, tokens): - stage = eager_stage() - out = [] - for i, (case, kw, run_kw) in enumerate(stage_cases(rows, tokens)): - st = Step(rows, tokens, _seed(3, rows, tokens, i), **kw) - bad = compare( - st.bufs, - lambda b: st.torch_stage(b, **run_kw), - lambda b: st.run_stage(b, stage, **run_kw), + lambda b: torch_gather(*st.gather_args(b)), + lambda b: kernel_gather(st.gather_args(b)), ) out.append(result(case, **bad)) return out def measure_graph(tokens, replays=6): - """One CUDA graph for the split family R = 1 .. 8 at ``tokens``: each step's scatter, its gather (from the stores - the scatter wrote, as the next step reads them) and its staged commit, captured once and replayed with every - input rewritten in place; each replay against the eager ops and the torch ops on the same inputs, then a replay - of the last inputs against the first.""" - op = _op() - with_stage = _pinned_ok() - gather_stage = op.StepInputStage(STAGE_CAPACITY) # gathers only: its record is never read + """One CUDA graph for the split family R = 1 .. 8 at ``tokens``: each step's scatter and its gather (from the + stores the scatter wrote, as the next step reads them), captured once and replayed with every input rewritten in + place; each replay against the eager ops and the torch ops on the same inputs, then a replay of the last inputs + against the first.""" steps = [ Step(r, tokens, _seed(4, r, tokens, 0), num_slots=2 * MAX_BATCH, row_begin=MAX_BATCH - r, out_rows=MAX_BATCH, begins=engine_begins(r, tokens), chain=True) for r in ROWS ] # fmt: skip - stage = eager_stage() if with_stage else None - # A captured commit reads its stage's record at every replay, so each step keeps its own stage. - captured = [op.StepInputStage(STAGE_CAPACITY) if with_stage else None for _ in steps] - def run_kernels(st, b, step_stage): + def run_kernels(st, b): kernel_scatter(st.scatter_args(b)) - kernel_gather(gather_stage, st.gather_args(b, "a")) - if step_stage is not None: - st.run_stage(b, step_stage) + kernel_gather(st.gather_args(b)) def run_torch(st, b): torch_scatter(*st.scatter_args(b)) - torch_gather(*st.gather_args(b, "a")) - if with_stage: - st.torch_stage(b) + torch_gather(*st.gather_args(b)) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): for st in steps: # the kernels compile on their first call, outside capture - run_kernels(st, clone_bufs(st.bufs), stage) + run_kernels(st, clone_bufs(st.bufs)) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=stream): - for st, step_stage in zip(steps, captured): - run_kernels(st, st.bufs, step_stage) + for st in steps: + run_kernels(st, st.bufs) torch.cuda.synchronize() - # The staged values sit in each captured stage's pinned record in staging order; a replay reads them in place. - spans = [] - for st, step_stage in zip(steps, captured): - at, step_spans = 0, [] - for _, name in st.copies(st.bufs, 3) if with_stage else []: - n = st.values[name].numel() - assert torch.equal(step_stage.record[at : at + n], st.values[name]), "record layout" - step_spans.append((at, n, name)) - at += n - spans.append(step_spans) - what = "scatter, gather" + (", staged commit" if with_stage else "") out = [] before = None for rep in range(replays): - for st, step_stage, step_spans in zip(steps, captured, spans): + for st in steps: st.rewrite() - for at, n, name in step_spans: - step_stage.record[at : at + n].copy_(st.values[name]) before = [clone_bufs(st.bufs) for st in steps] graph.replay() torch.cuda.synchronize() @@ -597,11 +425,11 @@ def run_torch(st, b): for st, start in zip(steps, before): want, got = clone_bufs(start), clone_bufs(start) run_torch(st, want) - run_kernels(st, got, stage) + run_kernels(st, got) torch.cuda.synchronize() bad["torch"] += [f"{st.rows}x{tokens} {k}" for k in mismatches(want, st.bufs)] bad["eager"] += [f"{st.rows}x{tokens} {k}" for k in mismatches(got, st.bufs)] - out.append(result(f"replay {rep}: {what}, every input rewritten", **bad)) + out.append(result(f"replay {rep}: scatter, gather, every input rewritten", **bad)) after = [clone_bufs(st.bufs) for st in steps] for st, start in zip(steps, before): for k, t in st.bufs.items(): @@ -632,14 +460,6 @@ def test_gather(rows, tokens): assert not bad, bad -@pytest.mark.parametrize("rows,tokens", SPLITS, ids=[f"{r}x{t}" for r, t in SPLITS]) -def test_stage(rows, tokens): - if not _pinned_ok(): - pytest.skip("StepInputStage needs pinned host memory (not preferred here)") - bad = [r for r in measure_stage(rows, tokens) if not r["ok"]] - assert not bad, bad - - @pytest.mark.parametrize("tokens", TOKENS, ids=[f"Rx{t}" for t in TOKENS]) def test_graph_replay(tokens): bad = [r for r in measure_graph(tokens) if not r["ok"]] @@ -701,20 +521,16 @@ def test_scatter_zeros_past_accepted_length(poison): def test_contract(): """rows == 0 launches nothing; every call the kernels do not cover returns False and launches nothing, so the - caller keeps its torch path; StepInputStage's per-step limits (max_copies = 4 copies, max_block_copies = 2 block - copies, one gather) decline the call past them.""" + caller keeps its torch path.""" op = _op() assert op.is_supported() st = Step(MAX_BATCH, STORE_WIDTH, _seed(6, 0, 0, 0)) b = clone_bufs(st.bufs) before = clone_bufs(b) scatter = list(st.scatter_args(b)) - gather = list(st.gather_args(b, "a")) - stage = op.StepInputStage(STAGE_CAPACITY) + gather = list(st.gather_args(b)) assert op.SlotScatter().scatter(*scatter[:2], 0, *scatter[3:]) - stage.begin() - assert stage.gather(*gather[:5], 0, *gather[6:]) - stage.commit() + assert op.StepInputGather().gather(*gather[:5], 0, *gather[6:]) torch.cuda.synchronize() assert not mismatches(before, b), "rows == 0 wrote" # Argument 7 (draft width) == tokens, argument 6 (tokens) wider than the 8-wide store, argument 15 (KV begin) @@ -727,8 +543,9 @@ def test_contract(): ): args = list(gather) args[index] = value - stage.begin() - assert not stage.gather(*args), f"gather argument {index} = {value!r} was staged" + assert not op.StepInputGather().gather(*args), ( + f"gather argument {index} = {value!r} was launched" + ) # Outputs of 8 rows with rows 1 .. 8 asked, an int64 slot table, a non-contiguous output. assert not op.SlotScatter().scatter(scatter[0], 1, *scatter[2:]) assert not op.SlotScatter().scatter(*scatter[:3], b["slot_table"].long(), *scatter[4:]) @@ -736,27 +553,6 @@ def test_contract(): scatter[0], next_draft_tokens=torch.empty_like(b["out_draft"]).t().contiguous().t() ) assert not op.SlotScatter().scatter(strided, *scatter[1:]) - if _pinned_ok(): - stage.begin() - for _ in range(stage.max_copies): - assert stage.copy(b["extra"][:2], torch.tensor([1, 2], dtype=torch.int32)) - assert not stage.copy(b["extra"][:2], torch.tensor([1, 2], dtype=torch.int32)) - stage.begin() - assert not stage.copy(b["extra"][:2], torch.tensor([1, 2])) # int64 values - assert not stage.copy( - b["out_new"], torch.tensor([1, 2], dtype=torch.int32) - ) # 2-D destination - assert stage.gather(*gather) - assert not stage.gather(*gather) - blocks = st.block_copies(b) - for args in blocks: - assert stage.block_copy(*args) - assert not stage.block_copy(*blocks[0]) - stage.begin() # drops the staged step: nothing was launched - unpinned = blocks[0][1].clone() # a pageable host table - assert not unpinned.is_pinned() - assert not stage.block_copy(blocks[0][0], unpinned, *blocks[0][2:]) - stage.begin() torch.cuda.synchronize() assert not mismatches(before, b), "a declined call wrote" @@ -769,11 +565,10 @@ def test_contract(): def report() -> int: print(f"{torch.cuda.get_device_name()}") ok_all = True - sections = [("SlotScatter.scatter", measure_scatter), ("StepInputStage.gather", measure_gather)] - if _pinned_ok(): - sections.append(("StepInputStage", measure_stage)) - else: - print("StepInputStage: skipped (pinned host memory not preferred here)") + sections = [ + ("SlotScatter.scatter", measure_scatter), + ("StepInputGather.gather", measure_gather), + ] for name, measure in sections: print(f"\n## {name}\n") print("| split | case | identical | result |") @@ -822,45 +617,22 @@ def time_graph(body, calls, replays=15): return statistics.median(per_call), min(per_call), max(per_call) -class StageArm: - """Back-to-back staged commits: fresh stages for every capture (a captured commit's stage cannot stage again), - and one of its own for the warm-up call.""" - - def __init__(self, step, bufs, calls): - self.step, self.bufs, self.calls = step, bufs, calls - self.warm = _op().StepInputStage(STAGE_CAPACITY) - self.fresh = [] - - def __call__(self, i): - if i < 0: - self.fresh = [_op().StepInputStage(STAGE_CAPACITY) for _ in range(self.calls)] - self.step.run_stage(self.bufs, self.warm) - else: - self.step.run_stage(self.bufs, self.fresh[i]) - - def timing() -> None: - op = _op() calls = 32 - with_stage = _pinned_ok() print(f"{torch.cuda.get_device_name()}; graphs of {calls} back-to-back calls on one step's buffers (the engine's " "layout), 15 replays: median (min-max) us per call") # fmt: skip - print("| split | scatter: torch | scatter: kernel | gather: torch | gather: kernel | stage: torch " - "| stage: kernel |") # fmt: skip - print("| :-- | --: | --: | --: | --: | --: | --: |") + print("| split | scatter: torch | scatter: kernel | gather: torch | gather: kernel |") + print("| :-- | --: | --: | --: | --: |") for rows, tokens in SPLITS: st = Step(rows, tokens, _seed(5, rows, tokens, 0), row_begin=MAX_BATCH - rows, out_rows=MAX_BATCH, begins=engine_begins(rows, tokens)) # fmt: skip b = st.bufs - gather_stage = op.StepInputStage(STAGE_CAPACITY) arms = [ lambda i: torch_scatter(*st.scatter_args(b)), lambda i: kernel_scatter(st.scatter_args(b)), - lambda i: torch_gather(*st.gather_args(b, "a")), - lambda i: kernel_gather(gather_stage, st.gather_args(b, "a")), + lambda i: torch_gather(*st.gather_args(b)), + lambda i: kernel_gather(st.gather_args(b)), ] - if with_stage: - arms += [lambda i: st.torch_stage(b), StageArm(st, b, calls)] res = [[] for _ in arms] for rep in range(3): # alternating order order = range(len(arms)) if rep % 2 == 0 else reversed(range(len(arms))) @@ -871,7 +643,6 @@ def timing() -> None: meds = sorted(x[0] for x in timings) lo, hi = min(x[1] for x in timings), max(x[2] for x in timings) cells.append(f"{meds[1]:.2f} ({lo:.2f}-{hi:.2f})") - cells += ["n/a"] * (6 - len(cells)) print(f"| {rows}x{tokens} | " + " | ".join(cells) + " |", flush=True) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 74739d6c37c8..601365578604 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1727,40 +1727,23 @@ def test_promoted_context_precedes_speculative_overlap_generation( [generation.py_seq_slot], 0) kv_cache_manager.shutdown() - def _check_staged_spec_decode_graph_step( - self, use_kv_cache_manager_v2: bool) -> None: - """A CUDA graph decode step of a speculative engine whose step inputs - are staged on the StepInputStage (one launch after the attention - metadata's prepare) writes what the torch path writes.""" + def test_spec_decode_graph_step_gather_kernel_matches_torch_path( + self) -> None: + """A CUDA graph decode step of a speculative engine whose overlap + gathers run as one StepInputGather launch writes what the torch path + writes.""" from tensorrt_llm._torch.attention.backends.trtllm import \ TrtllmAttentionMetadata from tensorrt_llm._torch.cute_dsl_kernels.spec_step_copies import \ op as spec_step_copies - from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import \ - KVCacheManagerV2 - from tensorrt_llm.llmapi.llm_args import \ - KvCacheConfig as LlmKvCacheConfig if not spec_step_copies.is_supported(): self.skipTest("the step-copy kernels run on SM 100 / 103 / 107") max_draft_len = 3 tokens_per_step = max_draft_len + 1 model_engine, kv_cache_manager = create_model_engine_and_kvcache( spec_config=SADecodingConfig(max_draft_len=max_draft_len)) - if use_kv_cache_manager_v2: - kv_cache_manager.shutdown() - kv_cache_manager = KVCacheManagerV2( - LlmKvCacheConfig(max_tokens=512, enable_block_reuse=False), - tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, - num_layers=1, - num_kv_heads=model_engine.model.config.num_key_value_heads, - head_dim=model_engine.model.config.head_dim, - tokens_per_block=4, - max_seq_len=256, - max_batch_size=8, - mapping=Mapping(world_size=1, tp_size=1, rank=0), - dtype=tensorrt_llm.bindings.DataType.HALF) - stage = model_engine._step_input_stage - self.assertIsNotNone(stage) + step_gather = model_engine._step_input_gather + self.assertIsNotNone(step_gather) resource_manager = ResourceManager( {ResourceManagerType.KV_CACHE_MANAGER: kv_cache_manager}) @@ -1806,19 +1789,16 @@ def _check_staged_spec_decode_graph_step( device="cuda"), ) - def step_inputs(step_input_stage): + def step_inputs(step_input_gather): """The buffers one step writes, from a filler they all start at.""" - model_engine._step_input_stage = step_input_stage + model_engine._step_input_gather = step_input_gather for request, slot in zip(requests, slots): request.py_batch_idx = slot written = (model_engine.input_ids_cuda, model_engine.position_ids_cuda, model_engine.draft_tokens_cuda, model_engine.previous_pos_id_offsets_cuda, - model_engine.previous_kv_lens_offsets_cuda, - graph_metadata.prompt_lens_cuda, - graph_metadata.kv_lens_cuda, - graph_metadata.kv_cache_block_offsets) + model_engine.previous_kv_lens_offsets_cuda) for buffer in written: buffer.fill_(-3) model_engine._prepare_tp_inputs(scheduled_requests=batch, @@ -1834,29 +1814,25 @@ def step_inputs(step_input_stage): return [buffer.clone() for buffer in written] torch_path = step_inputs(None) - with patch.object(stage, "commit", wraps=stage.commit) as commit, \ - patch.object(stage, "block_copy", - wraps=stage.block_copy) as block_copy: - staged = step_inputs(stage) - commit.assert_called_once() - self.assertEqual(block_copy.call_count, - 1 if use_kv_cache_manager_v2 else 0) - for expected, actual in zip(torch_path, staged): + launched = [] + launch = step_gather.gather + + def recorded_launch(*args, **kwargs): + launched.append(launch(*args, **kwargs)) + return launched[-1] + + with patch.object(step_gather, "gather", side_effect=recorded_launch): + kernel_path = step_inputs(step_gather) + self.assertEqual(launched, [True]) + for expected, actual in zip(torch_path, kernel_path): torch.testing.assert_close(actual, expected, rtol=0, atol=0) # The previous rows' input ids came from the stores, by slot. num_tokens = len(requests) * tokens_per_step expected_input_ids = previous.new_tokens[:, slots, 0].t().reshape(-1) - self.assertEqual(staged[0][:num_tokens].tolist(), + self.assertEqual(kernel_path[0][:num_tokens].tolist(), expected_input_ids.tolist()) kv_cache_manager.shutdown() - def test_staged_spec_decode_graph_step_matches_torch_path(self) -> None: - self._check_staged_spec_decode_graph_step(use_kv_cache_manager_v2=False) - - def test_staged_spec_decode_graph_step_matches_torch_path_kv_v2( - self) -> None: - self._check_staged_spec_decode_graph_step(use_kv_cache_manager_v2=True) - def test_multimodal_encoder_max_seq_len(self) -> None: class CapturingEncoder(torch.nn.Module, MultimodalEncoderMixin): From 0cbb559b2528f2c2a9104289d0a624eb07612580 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 00:01:41 -0700 Subject: [PATCH 45/67] [None][fix] DSpark GQA drafter: keep no mask row for DFlash's default mask id eda9db209f raised for any mask id outside the shipped embedding, including DFlash's default (vocab_size). The drafter takes that default when its config names no mask token in a key it reads. The speculative config can still name the mask token then (a user value, or a key only it reads, such as dspark_noise_token_id), and such a drafter loaded before. It now keeps no row, as a checkpoint without an embedding does. Any other id outside the shipped rows still raises. Test: test_gqa_dspark_keeps_no_mask_row_when_its_config_names_no_mask_token. The out-of-range case of the rejection test now uses vocab_size + 1. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/models/modeling_dspark.py | 19 +++++++++------ .../test_kimi_k3_dspark_semantics.py | 24 +++++++++++++++---- 2 files changed, 32 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index e0fcd0b8e44f..639c40293b46 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2692,7 +2692,8 @@ def _trained_mask_row(self, weights: Dict) -> Optional[torch.Tensor]: differs from the target's, and the shipped ``lm_head`` is bit-exact with the target's. Raises: - ValueError: the embedding is not ``[rows, hidden_size]``, or ``mask_token_id`` is not one of its rows. + ValueError: the embedding is not ``[rows, hidden_size]``, or the mask id the drafter config names is not + one of its rows. """ names = [k for k in ("embed_tokens.weight", "model.embed_tokens.weight") if k in weights] if not names: @@ -2705,12 +2706,16 @@ def _trained_mask_row(self, weights: Dict) -> Optional[torch.Tensor]: f"{type(self).__name__}: the checkpoint's {name} has shape {list(weight.shape)}, " f"expected [rows, hidden_size={hidden_size}]." ) - if not 0 <= self.mask_token_id < weight.shape[0]: - raise ValueError( - f"{type(self).__name__}: the drafter's mask_token_id {self.mask_token_id} is not a row of the " - f"checkpoint's {name} ({weight.shape[0]} rows)." - ) - return weight[self.mask_token_id] + if 0 <= self.mask_token_id < weight.shape[0]: + return weight[self.mask_token_id] + if self.mask_token_id == self.config.vocab_size: + # DFlash's default when the drafter config names no mask token in a key the drafter reads. The speculative + # config may still name one (a user value, or a key only it reads), so there is no known row to keep. + return None + raise ValueError( + f"{type(self).__name__}: the drafter's mask_token_id {self.mask_token_id} is not a row of the " + f"checkpoint's {name} ({weight.shape[0]} rows)." + ) class MLADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index 4d895340b216..0dd80eb191d7 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -22,6 +22,7 @@ import math import re from types import SimpleNamespace +from typing import Optional import pytest import torch @@ -235,7 +236,9 @@ def test_swa_window_conventions(): NUM_CAPTURE = 2 -def _tiny_config(dspark: bool, *, published_spelling: bool = False, mask_token_id: int = VOCAB - 2): +def _tiny_config( + dspark: bool, *, published_spelling: bool = False, mask_token_id: Optional[int] = VOCAB - 2 +): """Tiny drafter config. ``published_spelling`` reproduces the two public K3 DSpark checkpoints @@ -248,7 +251,9 @@ def _tiny_config(dspark: bool, *, published_spelling: bool = False, mask_token_i from transformers import Qwen3Config cfg = dict(TINY) - dflash = {"mask_token_id": mask_token_id, "target_layer_ids": [0, 1]} + dflash = {"target_layer_ids": [0, 1]} + if mask_token_id is not None: + dflash["mask_token_id"] = mask_token_id if dspark and published_spelling: cfg.update( markov_rank=RANK, @@ -330,7 +335,7 @@ def _build_drafter( *, published_spelling: bool = False, dflash_attention_backend: str = "VANILLA", - mask_token_id: int = VOCAB - 2, + mask_token_id: Optional[int] = VOCAB - 2, ): from tensorrt_llm._torch.model_config import ModelConfig @@ -521,7 +526,7 @@ def test_gqa_dspark_mask_row_takes_the_drafter_dtype_and_device_at_load(): @pytest.mark.parametrize( "width, mask_token_id, message", [ - (TINY["hidden_size"], VOCAB, f"mask_token_id {VOCAB} is not a row"), + (TINY["hidden_size"], VOCAB + 1, f"mask_token_id {VOCAB + 1} is not a row"), (TINY["hidden_size"], -1, "mask_token_id -1 is not a row"), ( TINY["hidden_size"] + 8, @@ -542,6 +547,17 @@ def test_gqa_dspark_rejects_a_mask_row_the_shipped_embedding_does_not_have( _build_drafter(True, weights, mask_token_id=mask_token_id) +@needs_gpu +def test_gqa_dspark_keeps_no_mask_row_when_its_config_names_no_mask_token(): + """A drafter config that names no mask token leaves DFlash's default id (vocab_size), one past the shipped rows. + The drafter keeps no row and loads; the speculative config can still name the mask token.""" + weights = _tiny_weights() + weights["embed_tokens.weight"] = torch.zeros(VOCAB, TINY["hidden_size"], dtype=torch.bfloat16) + drafter = _build_drafter(True, weights, mask_token_id=None) + assert drafter.mask_token_id == VOCAB + assert getattr(drafter, "mask_token_embedding", None) is None + + @needs_gpu def test_published_drafter_spelling_activates_the_heads(): """Both public K3 DSpark checkpoints load with their heads live. From b83f960a1d81cf16bfa022e59cac3b96893a95fc Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 02:57:32 -0700 Subject: [PATCH 46/67] [None][fix] DSpark GQA drafter: keep the trained row of the worker's mask id The worker embeds the speculative config's mask id when it has one, and that config also reads a checkpoint key the drafter's own config resolution does not: dspark_noise_token_id. For a checkpoint that names its mask token only that way, the drafter fell back to DFlash's default id (vocab_size) and kept no trained row. The worker's id check then failed, so every masked slot silently got the target's row again. The drafter's own ModelConfig carries no spec_config, so _build_dspark_draft now hands the speculative config's mask id to GQADSparkForCausalLM, which adopts it. The drafter keeps the row of the id the worker uses, and the worker's check matches. MLA DSpark keeps its own embedding and is unchanged. When a shipped embedding still keeps no row (no config names a mask token), the drafter logs a warning. Tests: test_gqa_dspark_keeps_the_trained_row_of_the_speculative_configs_mask_id and test_gqa_dspark_warns_when_its_shipped_embedding_keeps_no_mask_row. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/models/modeling_dspark.py | 44 +++++++++++---- .../test_kimi_k3_dspark_semantics.py | 53 ++++++++++++++++++- 2 files changed, 84 insertions(+), 13 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_dspark.py b/tensorrt_llm/_torch/models/modeling_dspark.py index 639c40293b46..57611e48d00d 100644 --- a/tensorrt_llm/_torch/models/modeling_dspark.py +++ b/tensorrt_llm/_torch/models/modeling_dspark.py @@ -2666,8 +2666,23 @@ class GQADSparkForCausalLM(_DSparkHeadMixin, DFlashForCausalLM): # _paged_ctx_cache is False here, so VANILLA is the only route to the arena. _default_attention_backend = "TRTLLM" - def __init__(self, draft_config, *, dflash_attention_backend: str = "AUTO"): + def __init__( + self, + draft_config, + *, + dflash_attention_backend: str = "AUTO", + spec_mask_token_id: Optional[int] = None, + ): super().__init__(draft_config, dflash_attention_backend=dflash_attention_backend) + # The worker embeds the speculative config's mask id when it has one, and that config also reads keys the + # drafter config resolution does not (``dspark_noise_token_id``). Adopting it keeps the trained row of the id + # the worker uses. + if spec_mask_token_id is not None and int(spec_mask_token_id) != self.mask_token_id: + logger.info( + f"{type(self).__name__}: mask_token_id {int(spec_mask_token_id)} from the speculative config " + f"(the drafter config resolves {self.mask_token_id})." + ) + self.mask_token_id = int(spec_mask_token_id) self._init_dspark_heads(draft_config.pretrained_config) def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): @@ -2692,8 +2707,7 @@ def _trained_mask_row(self, weights: Dict) -> Optional[torch.Tensor]: differs from the target's, and the shipped ``lm_head`` is bit-exact with the target's. Raises: - ValueError: the embedding is not ``[rows, hidden_size]``, or the mask id the drafter config names is not - one of its rows. + ValueError: the embedding is not ``[rows, hidden_size]``, or the mask id in use is not one of its rows. """ names = [k for k in ("embed_tokens.weight", "model.embed_tokens.weight") if k in weights] if not names: @@ -2709,12 +2723,16 @@ def _trained_mask_row(self, weights: Dict) -> Optional[torch.Tensor]: if 0 <= self.mask_token_id < weight.shape[0]: return weight[self.mask_token_id] if self.mask_token_id == self.config.vocab_size: - # DFlash's default when the drafter config names no mask token in a key the drafter reads. The speculative - # config may still name one (a user value, or a key only it reads), so there is no known row to keep. + # DFlash's default id, used when no config names a mask token. + logger.warning( + f"{type(self).__name__}: the checkpoint ships {name}, but mask_token_id is DFlash's default " + f"({self.mask_token_id} = vocab_size, outside its {weight.shape[0]} rows): no trained mask row is " + "kept, and the masked slots use the shared embedding's row." + ) return None raise ValueError( - f"{type(self).__name__}: the drafter's mask_token_id {self.mask_token_id} is not a row of the " - f"checkpoint's {name} ({weight.shape[0]} rows)." + f"{type(self).__name__}: mask_token_id {self.mask_token_id} is not a row of the checkpoint's {name} " + f"({weight.shape[0]} rows)." ) @@ -3458,12 +3476,16 @@ def _build_dspark_draft(model_config, draft_config, lm_head, model): # (``DFlashForCausalLM.__init__``) already resolves the backbone through the # model registry, so keying on model_type a second time would force a new # entry for every GQA family that already works. - drafter_cls = ( - MLADSparkForCausalLM if is_mla(draft_config.pretrained_config) else GQADSparkForCausalLM - ) - return drafter_cls( + if is_mla(draft_config.pretrained_config): + return MLADSparkForCausalLM( + draft_config, + dflash_attention_backend=model_config.spec_config.attention_backend, + ) + # The drafter's own ModelConfig has no spec_config, so the worker's mask id is handed over here. + return GQADSparkForCausalLM( draft_config, dflash_attention_backend=model_config.spec_config.attention_backend, + spec_mask_token_id=model_config.spec_config.mask_token_id, ) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py index 0dd80eb191d7..ff4d75f354dc 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -549,8 +549,8 @@ def test_gqa_dspark_rejects_a_mask_row_the_shipped_embedding_does_not_have( @needs_gpu def test_gqa_dspark_keeps_no_mask_row_when_its_config_names_no_mask_token(): - """A drafter config that names no mask token leaves DFlash's default id (vocab_size), one past the shipped rows. - The drafter keeps no row and loads; the speculative config can still name the mask token.""" + """With no mask token named by the drafter config or a speculative config, the drafter keeps DFlash's default id + (vocab_size), one past the shipped rows: it keeps no row and loads.""" weights = _tiny_weights() weights["embed_tokens.weight"] = torch.zeros(VOCAB, TINY["hidden_size"], dtype=torch.bfloat16) drafter = _build_drafter(True, weights, mask_token_id=None) @@ -558,6 +558,55 @@ def test_gqa_dspark_keeps_no_mask_row_when_its_config_names_no_mask_token(): assert getattr(drafter, "mask_token_embedding", None) is None +@needs_gpu +def test_gqa_dspark_warns_when_its_shipped_embedding_keeps_no_mask_row(): + """Keeping no row for DFlash's default id is logged as a warning, since the masked slots then use the shared + embedding's row instead of a trained one.""" + from unittest.mock import patch + + from tensorrt_llm._torch.models import modeling_dspark + + weights = _tiny_weights() + weights["embed_tokens.weight"] = torch.zeros(VOCAB, TINY["hidden_size"], dtype=torch.bfloat16) + with patch.object(modeling_dspark.logger, "warning") as warning: + drafter = _build_drafter(True, weights, mask_token_id=None) + assert getattr(drafter, "mask_token_embedding", None) is None + messages = [" ".join(map(str, call.args)) for call in warning.call_args_list] + assert any("no trained mask row is kept" in m for m in messages), messages + + +@needs_gpu +def test_gqa_dspark_keeps_the_trained_row_of_the_speculative_configs_mask_id(): + """A checkpoint that names its mask token only as ``dspark_noise_token_id``: the speculative config resolves that + key and the worker embeds its id, while the drafter config resolution falls back to DFlash's default. The builder + hands the drafter the speculative config's id, so the drafter keeps that id's trained row and its mask id is the + worker's.""" + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.models.modeling_dspark import _build_dspark_draft + + mask_id = VOCAB - 2 + pretrained = _tiny_config(True, mask_token_id=None) + pretrained.dspark_noise_token_id = mask_id + spec = SimpleNamespace( + mask_token_id=mask_id, attention_backend="VANILLA", draft_is_embedded_in_target=False + ) + drafter = _build_dspark_draft( + SimpleNamespace(spec_config=spec), + ModelConfig(pretrained_config=pretrained, attn_backend="TRTLLM"), + None, + None, + ).to("cuda") + weights = _tiny_weights() + g = torch.Generator().manual_seed(31) + trained = (torch.randn(VOCAB, TINY["hidden_size"], generator=g) * 0.05).to(torch.bfloat16) + weights["embed_tokens.weight"] = trained + drafter.load_weights(dict(weights)) + + assert isinstance(drafter, GQADSparkForCausalLM) + assert drafter.mask_token_id == spec.mask_token_id + torch.testing.assert_close(drafter.mask_token_embedding.cpu(), trained[mask_id], rtol=0, atol=0) + + @needs_gpu def test_published_drafter_spelling_activates_the_heads(): """Both public K3 DSpark checkpoints load with their heads live. From c95eac12a69081f1ba62db657ab6cc6564fd341e Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 09:30:56 -0700 Subject: [PATCH 47/67] [None][fix] KDA suspend: read the drafts' records as [slots, 3, num_spec, H, K] The Kimi K3 verify kernels' kda_state_tok now holds only the drafts' records, [slots, 3, num_spec, H, K] fp32 per layer: the row innovations vn, then beta * k and the decay, in the order they have always been laid out. It no longer has room for a full state per draft. replay_kda_token_records takes num_spec, H and K from that shape and V from the state pool. It raises ValueError for any other layout instead of misreading it. The records' offsets within a slot are unchanged. The suspend commit test builds its records in the new shape. Signed-off-by: Vasanth Sabavat --- .../_torch/modules/kimi_kda/_kda_kernels.py | 32 +++++++++++-------- .../kv_cache/mamba_cache_manager.py | 4 +-- .../test_kda_mtp_decode_cute_parity.py | 10 +++--- 3 files changed, 26 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index c694b90eba1e..48189c103d47 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -406,17 +406,24 @@ def replay_kda_token_records( num_accepted_tokens: torch.Tensor, ) -> None: """Replay each selected slot's accepted drafts onto its pool state from - the per-draft records the Kimi K3 verify kernels keep in the slot's - per-token state region. - - A slot's region of ``state_tok`` (``[slots, num_spec, H, V, K]`` fp32) - starts with the records of the last round's drafts: the row innovations - ``vn`` ``[num_spec][H][V]``, then ``beta * k`` and the decay - ``[num_spec][H][K]``. Draft ``t`` updates the state as the kernels do, - ``S = fma(decay, S, vn (beta * k))``, so the result is bit-identical to the - state their next verify starts from. + the per-draft records the Kimi K3 verify kernels keep in ``state_tok``. + + A slot of ``state_tok`` (``[slots, 3, num_spec, H, K]`` fp32) holds the + records of the last round's drafts: the row innovations ``vn`` + ``[num_spec][H][V]``, then ``beta * k`` and the decay + ``[num_spec][H][K]`` (V = K). Draft ``t`` updates the state as the + kernels do, ``S = fma(decay, S, vn (beta * k))``, so the result is + bit-identical to the state their next verify starts from. """ - num_spec, num_heads, v_dim, k_dim = state_tok.shape[1:] + _, num_spec, num_heads, k_dim = state_tok.shape[1:] + v_dim = recurrent_state.shape[-2] + if state_tok.shape[1] != 3 or v_dim != k_dim: + raise ValueError( + "replay_kda_token_records: expected the drafts' records " + "[slots, 3, num_spec, H, K] with V = K, got " + f"{tuple(state_tok.shape)} for a state of " + f"{tuple(recurrent_state.shape)}" + ) if state_indices.numel() == 0: return block_v = 32 @@ -458,9 +465,8 @@ def commit_kda_pending_drafts( hold the state and the conv window that verify would start from, bit for bit, and the caller clears the slot's count: - * with per-token states (``state_tok``), the drafts are replayed from the - per-draft records the Kimi K3 verify kernels keep there - (``replay_kda_token_records``); + * with the Kimi K3 verify kernels' draft records (``state_tok``), the + drafts are replayed from those (``replay_kda_token_records``); * otherwise they are replayed from the replay caches by the fused verify kernel's own replay steps (``kda_mtp_commit_pending_drafts``). """ diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 3be8d84f564a..c7cafe41f02b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -4377,8 +4377,8 @@ def _commit_kda_pending_drafts(self, slot: int) -> None: commit_kda_pending_drafts assert self.prev_num_accepted_tokens is not None - # With per-token verify states, the accepted drafts are replayed from - # the records the verify kernels keep there instead of the caches. + # With the verify kernels' draft records (kda_state_tok), the accepted + # drafts are replayed from those instead of the caches. state_tok = getattr(self, "kda_state_tok", None) with self._on_kv_cache_stream(): slots = torch.full((1, ), 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 c909afad0c83..b2672bb5d20c 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 @@ -766,10 +766,10 @@ def test_commit_pending_drafts_matches_the_next_verify(B, H, num_spec): def test_commit_pending_drafts_replays_the_per_token_records(): - """With per-token verify states, a slot's accepted drafts are replayed onto - its pool state, in order, from the per-draft records the Kimi K3 verify - kernels keep in its region (vn, beta * k, decay), and the conv window after - them is stored; other slots are left alone.""" + """With the Kimi K3 verify kernels' draft records, a slot's accepted drafts + are replayed onto its pool state, in order, from its records (vn, + beta * k, decay), and the conv window after them is stored; other slots + are left alone.""" from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import commit_kda_pending_drafts torch.manual_seed(5) @@ -784,7 +784,7 @@ def small(*shape, scale): vn = small(num_slots, num_spec, H, K, scale=2.0**-2) beta_k = small(num_slots, num_spec, H, K, scale=2.0**-3) decay = torch.randint(1, 3, (num_slots, num_spec, H, K), device="cuda").float() * 0.5 - state_tok = torch.zeros(num_slots, num_spec, H, K, K, device="cuda") + state_tok = torch.zeros(num_slots, 3, num_spec, H, K, device="cuda") records = state_tok.view(num_slots, -1) size = num_spec * H * K for i, part in enumerate((vn, beta_k, decay)): From dc903ae9d6168c9013c550a05d8cc5e39a3f4960 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 09:24:34 -0700 Subject: [PATCH 48/67] [None][fix] KDA verify: size kda_state_tok for the drafts' records The Kimi K3 verify kernels keep, per draft of the last verify round, the row innovations vn [H, V], beta * k and the decay [H, K] in kda_state_tok, and replay the accepted drafts from them. They never write a draft's full state. The buffer was sized and documented as the full state after every draft, [layers, slots, num_spec, H, V, K] fp32, so ~98 % of it was never used; at the SSM live floor those bytes came out of the KV cache quota. kda_state_tok is now [layers, slots, 3, num_spec, H, K] fp32, in the kernels' record order (vn, then beta * k, then the decay). The per-slot budget and the live-floor quota count those bytes. V = K holds wherever the replay caches exist, since their [Q | K | V] sections are equal. For Kimi K3 at TP16 with 7 drafts, a slot's records take 4.45 MB over the 69 KDA layers instead of 190 MB. The token-state tests now expect the records' shape and bytes. The runtime rounds the pools' quota up to whole 2 MiB grains, and the live-floor test's records are now less than one grain, so that test allows the pools and the records up to one grain over the quota. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/pyexecutor/_util.py | 6 +- .../kv_cache/mamba_cache_manager.py | 78 +++++++++++-------- .../kv_cache/test_mamba_cache_manager.py | 47 ++++++----- 3 files changed, 73 insertions(+), 58 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 2028050aa034..98ef12781bac 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3048,9 +3048,9 @@ def _create_kv_cache_manager( if is_kda_mtp_verify_available(): kda_extra_kwargs["kda_replay_num_spec"] = ( spec_config.tokens_per_gen_step - 1) - # A model whose KDA verify starts from the state after the - # accepted drafts (instead of replaying them) asks for the - # state after every draft with `kda_token_states`. + # A model whose KDA verify replays the accepted drafts from + # its own per-draft records (instead of the replay caches) + # asks for a buffer of them with `kda_token_states`. backbone = getattr(getattr(model_engine, "model", None), "model", None) if (issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 6fb06cb63aac..2ee88b7c871c 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -335,10 +335,9 @@ def use_kda_replay_update(self) -> bool: @property def keeps_kda_token_states(self) -> bool: - """Whether the KDA replay caches come with the state after every - draft of the last verify round (``kda_state_tok``), so that a verify - kernel can start the next round from the state after the accepted - drafts instead of replaying them.""" + """Whether the KDA replay caches come with the records of every draft + of the last verify round (``kda_state_tok``), from which a verify + kernel replays the accepted drafts onto the slot's state.""" return False @abstractmethod @@ -449,13 +448,15 @@ class SpeculativeState(State): kda_qkg_cache: torch.Tensor | None = None kda_v_cache: torch.Tensor | None = None kda_beta_cache: torch.Tensor | None = None - # Optional: the state after every draft of the last verify round, - # [slots, num_spec, H, V, K] fp32 (the SSM pool holds the state - # after the round's first, non-draft token). A verify kernel starts - # the next round from entry n - 1 when the round accepted n > 0 - # drafts (prev_num_accepted_tokens) instead of replaying them from - # the caches above; a slot reset for a new request has n = 0, so a - # previous owner's entries are never read. + # Optional: the records of every draft of the last verify round, + # [slots, 3, num_spec, H, K] fp32: the row innovations vn [num_spec, + # H, V], then beta * k and the decay [num_spec, H, K] (V = K). The + # SSM pool holds the state after the round's first, non-draft token; + # when the round accepted n > 0 drafts (prev_num_accepted_tokens), a + # verify kernel starts the next round by replaying the first n + # records onto it, S = decay * S + vn (beta * k), instead of + # replaying them from the caches above. A slot reset for a new + # request has n = 0, so a previous owner's records are never read. kda_state_tok: torch.Tensor | None = None # Replay path: compact double-buffered cache @@ -3060,8 +3061,8 @@ def __init__( self._kda_replay_num_spec = kda_replay_num_spec self._use_kda_replay_update = kda_replay_num_spec is not None - # Only with the KDA replay caches: also keep the state after every - # verify token (kda_state_tok). + # Only with the KDA replay caches: also keep the records of every + # draft of the last verify round (kda_state_tok). self._kda_token_states = kda_token_states and self._use_kda_replay_update if self._use_kda_replay_update: if use_replay_state_update: @@ -3536,12 +3537,12 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: # Beside the V2 pools. With the SSM pool at its live floor, # _build_cache_config took these bytes out of the GPU quota. self.kda_state_tok = torch.zeros( - (self.local_num_mamba_layers, cache_size, num_spec, - *self.ssm_state_shape), + (self.local_num_mamba_layers, cache_size, + *self._kda_draft_records_shape()), dtype=torch.float32, device=device, ) - per_token = (", with per-token states" + per_token = (", with the drafts' records" if self.kda_state_tok is not None else "") logger.info("Mamba Cache (kda-replay) is allocated for " f"{cache_size} state slots{per_token}") @@ -3837,13 +3838,21 @@ def _get_pool_roles(self, def _max_resident_sequences(self) -> int: return self.max_batch_size * self.mapping.pp_size + def _kda_draft_records_shape(self) -> Tuple[int, int, int, int]: + """One SSM slot's shape in a layer of ``kda_state_tok``: per draft, + the row innovations vn, then beta * k and the decay, each [H, K] (the + replay caches' equal [Q | K | V] sections make V = K).""" + return (3, self._kda_replay_num_spec, self.ssm_state_shape[0], + self.ssm_state_shape[-1]) + def _kda_token_state_bytes_per_slot(self) -> int: - """Bytes of ``kda_state_tok`` per SSM slot: an fp32 SSM state per - draft of every local Mamba layer, or 0 without per-token states.""" + """Bytes of ``kda_state_tok`` per SSM slot: the fp32 records of every + draft of every local Mamba layer, or 0 without them.""" if not getattr(self, "_kda_token_states", False): return 0 - return (self.local_num_mamba_layers * self._kda_replay_num_spec * - math.prod(self.ssm_state_shape) * torch.float32.itemsize) + return (self.local_num_mamba_layers * + math.prod(self._kda_draft_records_shape()) * + torch.float32.itemsize) def _mamba_state_bytes_per_slot(self) -> int: base_bytes = (self.local_num_mamba_layers * @@ -3945,8 +3954,8 @@ def _get_typical_request_capacity( return fallback_capacity def _state_quota(self, state_slots: int) -> int: - """Quota for ``state_slots`` recurrent-state slots, plus the token - states of the further slots the live floor keeps + """Quota for ``state_slots`` recurrent-state slots, plus the drafts' + records of the further slots the live floor keeps (``_live_floor_token_state_slots``).""" extra_slots = self._live_floor_token_state_slots( self.kv_cache_config.max_util_for_resume) - state_slots @@ -4007,13 +4016,14 @@ def _ssm_pool_at_live_floor(self, kv_cache_config: KvCacheConfig) -> bool: per resident sequence plus the reserved dummies) instead of its share of the typical step's ratio. - With per-token KDA states, ``kda_state_tok`` takes num_spec states per - SSM slot beside the V2 pools (``_allocate_pool_replay_buffers``), - addressed by the request's SSM slot. Without block reuse no SSM slot - holds anything but a live request or a reserved dummy, so slots past - the floor are never used and would only add that memory. At the floor - the pool's slot count is known before sizing, so - ``_build_cache_config`` takes those states' bytes out of the GPU quota. + With per-token KDA states, ``kda_state_tok`` takes num_spec drafts' + records per SSM slot beside the V2 pools + (``_allocate_pool_replay_buffers``), addressed by the request's SSM + slot. Without block reuse no SSM slot holds anything but a live + request or a reserved dummy, so slots past the floor are never used + and would only add that memory. At the floor the pool's slot count is + known before sizing, so ``_build_cache_config`` takes those records' + bytes out of the GPU quota. """ return (getattr(self, "_kda_token_states", False) and self.local_num_mamba_layers > 0 @@ -4027,8 +4037,8 @@ def _live_floor_token_state_slots(self, max_util_for_resume: float) -> int: every slot the pool keeps: the floor's min-slots constraint as the runtime scales it by 1 / max_util_for_resume. (Rounding each pool up to whole allocation grains can add slots when they are smaller than a - grain.) ``_build_cache_config`` takes their token states out of the - GPU quota before the pools are sized. Otherwise 0. + grain.) ``_build_cache_config`` takes their records out of the GPU + quota before the pools are sized. Otherwise 0. """ kv_cache_config = self.kv_cache_config if (kv_cache_config.pool_ratio is not None @@ -4078,7 +4088,7 @@ def _build_cache_config( self._ssm_pool_at_live_floor(kv_cache_config)) if ssm_pool_at_live_floor: # kda_state_tok is allocated beside the pools, one entry per SSM - # slot: take those states' bytes out of the GPU quota before the + # slot: take those records' bytes out of the GPU quota before the # pools are sized. token_state_slots = self._live_floor_token_state_slots( config.max_util_for_resume) @@ -4089,7 +4099,7 @@ def _build_cache_config( GpuCacheTierConfig(quota=gpu_quota), *cache_tiers[1:] ] logger.info( - "KV cache manager v2: the per-token KDA states of " + "KV cache manager v2: the KDA draft records of " f"{token_state_slots} SSM slots take " f"{token_state_quota / (1 << 30)}GiB of the device quota; the " f"pools get {gpu_quota / (1 << 30)}GiB") @@ -4437,7 +4447,7 @@ def _relocate_kda_replay_slots(self, old_slots: List[int], ] assert all(replay_buffer is not None for replay_buffer in replay_buffers) - # The per-token verify states belong to the slot as well. + # The drafts' records belong to the slot as well. if self.kda_state_tok is not None: replay_buffers.append(self.kda_state_tok) for replay_buffer in replay_buffers: diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index b4805d3ed744..26ed075837ec 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -2,7 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 """Regression tests for Python, Cpp, and V2 Mamba cache managers.""" -import math import os from types import SimpleNamespace from unittest.mock import MagicMock @@ -3673,8 +3672,8 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): @skip_no_cuda @pytest.mark.parametrize("kda_replay_num_spec", [2, None]) def test_v2_kda_token_states_allocated_with_the_replay_caches(kda_replay_num_spec): - """With the replay caches, ``kda_token_states`` adds the fp32 state after every - draft of each slot, per layer; without them it allocates nothing.""" + """With the replay caches, ``kda_token_states`` adds the fp32 records of every draft + of each slot (vn, beta * k, decay), per layer; without them it allocates nothing.""" mgr = _build_v2_hybrid_with_mamba_layer( max_batch_size=4, num_mamba_layers=2, @@ -3697,7 +3696,8 @@ def test_v2_kda_token_states_allocated_with_the_replay_caches(kda_replay_num_spe layer_cache = mgr.mamba_layer_cache(layer_idx) cache_size = layer_cache.temporal.shape[0] states = layer_cache.kda_state_tok - assert states.shape == (cache_size, 2, *layer_cache.temporal.shape[1:]) + num_heads, _, head_dim = layer_cache.temporal.shape[1:] + assert states.shape == (cache_size, 3, 2, num_heads, head_dim) assert states.dtype is torch.float32 assert states.data_ptr() == mgr.kda_state_tok[layer_idx].data_ptr() assert not states.any() @@ -3779,8 +3779,8 @@ def _v2_kda_token_state_manager(kda_token_states, enable_block_reuse): mgr._attention_cache_bytes_per_token = lambda: 8 mgr._use_kda_replay_update = True mgr._kda_token_states = kda_token_states - # The per-token states count in each slot's state bytes (_mamba_state_bytes_per_slot): num_spec fp32 states of - # the SSM state shape, 2 MiB each here. + # The drafts' records count in each slot's state bytes (_mamba_state_bytes_per_slot): per draft, vn, beta * k + # and the decay, fp32 [H, K] each (24 KiB here). mgr._kda_replay_num_spec = 2 mgr.ssm_state_shape = [8, 256, 256] mgr.kv_cache_config = KvCacheConfig( @@ -3798,9 +3798,9 @@ def _v2_kda_token_state_manager(kda_token_states, enable_block_reuse): def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( kda_token_states, enable_block_reuse, at_floor ): - """The per-token KDA states take memory per SSM slot beside the V2 pools. Without block reuse the SSM pool keeps - only its live floor, the states' bytes come out of the quota, and attention gets the rest; without the states - (replay caches only, with or without block reuse) the typical step's ratio sizes the pool.""" + """The KDA drafts' records take memory per SSM slot beside the V2 pools. Without block reuse the SSM pool keeps + only its live floor, the records' bytes come out of the quota, and attention gets the rest; without them (replay + caches only, with or without block reuse) the typical step's ratio sizes the pool.""" mgr = _v2_kda_token_state_manager(kda_token_states, enable_block_reuse) base_config = KVCacheManagerConfig( tokens_per_block=32, @@ -3822,16 +3822,19 @@ def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( if at_floor: # The floor's min-slots constraint is scaled by 1 / max_util_for_resume (0.97). assert floor <= slots["ssm"] <= floor + 1 - # kda_state_tok, allocated beside the pools: (layers, SSM slots, num_spec, *ssm_state_shape) fp32. + # kda_state_tok, allocated beside the pools: (layers, SSM slots, 3, num_spec, H, K) fp32. token_state_bytes = ( mgr.local_num_mamba_layers * slots["ssm"] + * 3 * mgr._kda_replay_num_spec - * math.prod(mgr.ssm_state_shape) + * mgr.ssm_state_shape[0] + * mgr.ssm_state_shape[-1] * 4 ) - # The pools and the states fit in the quota, and attention takes the rest of it, to within one 2 MiB grain. - assert (128 << 20) - (2 << 20) < pool_bytes + token_state_bytes <= 128 << 20 + # The records' bytes come out of the quota and attention takes the rest of it. The runtime rounds the pools' + # quota up to whole 2 MiB grains, so the pools can take back less than one grain of the records' bytes. + assert 128 << 20 <= pool_bytes + token_state_bytes < (128 << 20) + (2 << 20) else: assert slots["ssm"] > 10 * floor @@ -3861,14 +3864,14 @@ def test_v2_kda_token_states_refuse_a_pool_off_the_live_floor(off_the_floor): def test_v2_kda_token_states_reserve_quota_before_sizing(): - """At the live floor the token states of every slot the SSM pool keeps come out of the GPU quota before the pools - are sized, and the minimum quota and the quota for max_tokens count them all.""" + """At the live floor the drafts' records of every slot the SSM pool keeps come out of the GPU quota before the + pools are sized, and the minimum quota and the quota for max_tokens count them all.""" mgr = _v2_kda_token_state_manager(kda_token_states=True, enable_block_reuse=False) plain = _v2_kda_token_state_manager(kda_token_states=False, enable_block_reuse=False) assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 3 - # The pool keeps ceil(3 / max_util_for_resume) = 4 slots, each with 2 fp32 token states of 2 MiB; the quotas - # without token states differ by exactly those. - reserved = 4 * 2 * (2 << 20) + # The pool keeps ceil(3 / max_util_for_resume) = 4 slots, each with the records of 2 drafts (3 x 8 x 256 fp32 + # each); the quotas without them differ by exactly those. + reserved = 4 * 2 * 3 * 8 * 256 * 4 minimum = mgr._minimum_live_gpu_quota() assert minimum - plain._minimum_live_gpu_quota() == reserved for max_tokens in (0, 4096): @@ -3907,7 +3910,7 @@ def test_v2_kda_token_states_reserve_the_slots_the_runtime_keeps(): max_util_for_resume=0.95, ) assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 19 - reserved = 21 * 2 * (2 << 20) + reserved = 21 * 2 * 3 * 8 * 256 * 4 minimum = mgr._minimum_live_gpu_quota() assert minimum - plain._minimum_live_gpu_quota() == reserved quota = minimum + (64 << 20) @@ -4061,7 +4064,8 @@ def test_v2_kda_token_states_relocate_with_their_slot(): @skip_no_cuda def test_v2_kda_token_states_count_in_the_per_slot_budget(): - """The capacity math sees the per-token states: a slot costs one more fp32 SSM state per draft and layer.""" + """The capacity math sees the drafts' records: per draft and layer, a slot costs vn, beta * k and the decay, fp32 + [H, K] each.""" kwargs = dict( num_mamba_layers=2, spec_config=MTPDecodingConfig(max_draft_len=2), @@ -4079,7 +4083,8 @@ def test_v2_kda_token_states_count_in_the_per_slot_budget(): plain.shutdown() with_states = _build_v2_hybrid_with_mamba_layer(kda_token_states=True, **kwargs) try: - per_token = 2 * 2 * math.prod(with_states.ssm_state_shape) * 4 + num_heads, _, head_dim = with_states.ssm_state_shape + per_token = 2 * 2 * 3 * num_heads * head_dim * 4 assert with_states._mamba_state_bytes_per_slot() == plain_bytes + per_token finally: with_states.shutdown() From c05a0c3a452055dfe04f0b902f02c96ab4c36b4a Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 11:03:49 -0700 Subject: [PATCH 49/67] [None][fix] KDA verify: drop the replay caches the drafts' records replace With kda_token_states, the verify kernels replay the accepted drafts from the drafts' records (kda_state_tok). Nothing then reads the per-draft replay caches kda_qkg_cache, kda_v_cache and kda_beta_cache: their only reader is the built-in fused verify (kda_mtp_decode), which a model that asks for the records replaces. They are no longer allocated in that mode. The extended conv caches stay, since the records' verify reads them too. has_kda_replay_caches now follows the conv caches, so the replay path stays on with either the replay caches or the records. Slot relocation, the disaggregated-generation slot seeding and the per-layer view handle the absent caches. Without kda_token_states nothing changes. For Kimi K3 at TP16 with 7 drafts this frees 86,240 B per slot and KDA layer, 59.5 MB per rank at 10 SSM slots. The tests check that the caches are not allocated with the records, that the view, relocation and seeding work without them, and that the replay path holds in both modes. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 110 ++++++++++-------- .../kv_cache/test_mamba_cache_manager.py | 105 +++++++++++++++-- 2 files changed, 153 insertions(+), 62 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 2ee88b7c871c..35c599d5b8fa 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -417,7 +417,9 @@ class SpeculativeState(State): - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) - KDA replay: per-slot draft-token caches consumed by the fused ``trtllm::kda_mtp_decode`` verify kernel, which replays accepted - drafts and commits states in place (kda_conv_*, kda_*_cache) + drafts and commits states in place (kda_conv_*, kda_*_cache); or, + for a verify kernel that keeps the drafts' records, kda_conv_* and + kda_state_tok in place of the kda_*_cache """ _SHARED_FIELDS = frozenset({ "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed" @@ -444,7 +446,7 @@ class SpeculativeState(State): kda_conv_v: torch.Tensor | None = None # Post-processed per-draft quantities for replay: # [layers, slots, num_spec, 3, H*K] (q/k/gate), [.., num_spec, H*V], - # [.., num_spec, H] — all fp32. + # [.., num_spec, H] — all fp32. None with kda_state_tok. kda_qkg_cache: torch.Tensor | None = None kda_v_cache: torch.Tensor | None = None kda_beta_cache: torch.Tensor | None = None @@ -455,8 +457,9 @@ class SpeculativeState(State): # when the round accepted n > 0 drafts (prev_num_accepted_tokens), a # verify kernel starts the next round by replaying the first n # records onto it, S = decay * S + vn (beta * k), instead of - # replaying them from the caches above. A slot reset for a new - # request has n = 0, so a previous owner's records are never read. + # replaying them from the caches above, which are then not + # allocated. A slot reset for a new request has n = 0, so a previous + # owner's records are never read. kda_state_tok: torch.Tensor | None = None # Replay path: compact double-buffered cache @@ -476,9 +479,11 @@ class SpeculativeState(State): @property def has_kda_replay_caches(self) -> bool: - """True when the KDA fused-verify replay caches were allocated - (the fused ``trtllm::kda_mtp_decode`` verify path).""" - return self.kda_qkg_cache is not None + """True on the KDA fused-verify replay path: the extended conv + caches are allocated, with the per-draft replay caches of the + fused ``trtllm::kda_mtp_decode`` verify or, in their place, the + drafts' records (``kda_state_tok``).""" + return self.kda_conv_q is not None def commit_conv_window(self, slot_indices: torch.Tensor, conv_pool: torch.Tensor) -> None: @@ -3511,39 +3516,42 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: self.kda_conv_q = allocate_dim_contiguous_conv_cache() self.kda_conv_k = allocate_dim_contiguous_conv_cache() self.kda_conv_v = allocate_dim_contiguous_conv_cache() - self.kda_qkg_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - 3, - section_dim, - dtype=torch.float32, - device=device, - ) - self.kda_v_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - section_dim, - dtype=torch.float32, - device=device, - ) - self.kda_beta_cache = _allocate_kda_beta_cache( - (self.local_num_mamba_layers, cache_size, num_spec, - self.ssm_state_shape[0]), - device, - ) if self._kda_token_states: # Beside the V2 pools. With the SSM pool at its live floor, - # _build_cache_config took these bytes out of the GPU quota. + # _build_cache_config took these bytes out of the GPU quota. The + # verify replays the accepted drafts from these records, so the + # per-draft replay caches are not allocated. self.kda_state_tok = torch.zeros( (self.local_num_mamba_layers, cache_size, *self._kda_draft_records_shape()), dtype=torch.float32, device=device, ) - per_token = (", with the drafts' records" - if self.kda_state_tok is not None else "") + else: + self.kda_qkg_cache = torch.zeros( + self.local_num_mamba_layers, + cache_size, + num_spec, + 3, + section_dim, + dtype=torch.float32, + device=device, + ) + self.kda_v_cache = torch.zeros( + self.local_num_mamba_layers, + cache_size, + num_spec, + section_dim, + dtype=torch.float32, + device=device, + ) + self.kda_beta_cache = _allocate_kda_beta_cache( + (self.local_num_mamba_layers, cache_size, num_spec, + self.ssm_state_shape[0]), + device, + ) + per_token = (", with the drafts' records in place of the replay " + "caches" if self.kda_state_tok is not None else "") logger.info("Mamba Cache (kda-replay) is allocated for " f"{cache_size} state slots{per_token}") return True @@ -4437,19 +4445,16 @@ def _relocate_kda_replay_slots(self, old_slots: List[int], destination_slots = torch.tensor([new for _, new in moves], dtype=torch.long, device=device) - replay_buffers = [ - self.kda_conv_q, - self.kda_conv_k, - self.kda_conv_v, - self.kda_qkg_cache, - self.kda_v_cache, - self.kda_beta_cache, - ] + replay_buffers = [self.kda_conv_q, self.kda_conv_k, self.kda_conv_v] assert all(replay_buffer is not None for replay_buffer in replay_buffers) - # The drafts' records belong to the slot as well. - if self.kda_state_tok is not None: - replay_buffers.append(self.kda_state_tok) + # The per-draft replay caches, or the drafts' records in their place. + replay_buffers += [ + replay_buffer + for replay_buffer in (self.kda_qkg_cache, self.kda_v_cache, + self.kda_beta_cache, self.kda_state_tok) + if replay_buffer is not None + ] for replay_buffer in replay_buffers: replay_buffer.index_copy_( 1, @@ -4624,10 +4629,12 @@ def seed_kda_replay_caches_for_disagg_gen(self, replay_layer[:, :, :committed_window].index_copy_( 0, state_indices, section.to(replay_layer.dtype)) + # The drafts' records need no reset: with no pending drafts, none of + # them is read. for replay_buffer in (self.kda_qkg_cache, self.kda_v_cache, self.kda_beta_cache): - assert replay_buffer is not None - replay_buffer.index_fill_(1, state_indices, 0) + if replay_buffer is not None: + replay_buffer.index_fill_(1, state_indices, 0) self.prev_num_accepted_tokens[state_indices] = 0 def mamba_layer_cache( @@ -4642,20 +4649,21 @@ def mamba_layer_cache( assert self.kda_conv_q is not None assert self.kda_conv_k is not None assert self.kda_conv_v is not None - assert self.kda_qkg_cache is not None - assert self.kda_v_cache is not None - assert self.kda_beta_cache is not None spec_kwargs = { "prev_num_accepted_tokens": self.prev_num_accepted_tokens, "kda_conv_q": self.kda_conv_q[layer_offset], "kda_conv_k": self.kda_conv_k[layer_offset], "kda_conv_v": self.kda_conv_v[layer_offset], - "kda_qkg_cache": self.kda_qkg_cache[layer_offset], - "kda_v_cache": self.kda_v_cache[layer_offset], - "kda_beta_cache": self.kda_beta_cache[layer_offset], } if self.kda_state_tok is not None: spec_kwargs["kda_state_tok"] = self.kda_state_tok[layer_offset] + else: + assert self.kda_qkg_cache is not None + assert self.kda_v_cache is not None + assert self.kda_beta_cache is not None + spec_kwargs["kda_qkg_cache"] = self.kda_qkg_cache[layer_offset] + spec_kwargs["kda_v_cache"] = self.kda_v_cache[layer_offset] + spec_kwargs["kda_beta_cache"] = self.kda_beta_cache[layer_offset] if self.mamba_ssm_rand_seed is not None: spec_kwargs["mamba_ssm_rand_seed"] = self.mamba_ssm_rand_seed return PythonMambaCacheManager.SpeculativeState( diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 26ed075837ec..799373ba639f 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3665,15 +3665,17 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): assert layer_cache.intermediate_conv_window is None assert layer_cache.kda_state_tok is None assert not mgr.keeps_kda_token_states + assert layer_cache.has_kda_replay_caches finally: mgr.shutdown() @skip_no_cuda @pytest.mark.parametrize("kda_replay_num_spec", [2, None]) -def test_v2_kda_token_states_allocated_with_the_replay_caches(kda_replay_num_spec): - """With the replay caches, ``kda_token_states`` adds the fp32 records of every draft - of each slot (vn, beta * k, decay), per layer; without them it allocates nothing.""" +def test_v2_kda_token_states_replace_the_replay_caches(kda_replay_num_spec): + """On the KDA replay path, ``kda_token_states`` allocates the fp32 records of every draft of each slot (vn, + beta * k, decay), per layer, in place of the per-draft replay caches; the conv caches and the replay path stay. + Without the replay path it allocates nothing.""" mgr = _build_v2_hybrid_with_mamba_layer( max_batch_size=4, num_mamba_layers=2, @@ -3692,9 +3694,17 @@ def test_v2_kda_token_states_allocated_with_the_replay_caches(kda_replay_num_spe assert not mgr.keeps_kda_token_states return assert mgr.keeps_kda_token_states + assert mgr.kda_qkg_cache is None + assert mgr.kda_v_cache is None + assert mgr.kda_beta_cache is None for layer_idx in range(2): layer_cache = mgr.mamba_layer_cache(layer_idx) cache_size = layer_cache.temporal.shape[0] + assert layer_cache.has_kda_replay_caches + assert layer_cache.kda_conv_q.shape == (cache_size, 48, 6) + assert layer_cache.kda_qkg_cache is None + assert layer_cache.kda_v_cache is None + assert layer_cache.kda_beta_cache is None states = layer_cache.kda_state_tok num_heads, _, head_dim = layer_cache.temporal.shape[1:] assert states.shape == (cache_size, 3, 2, num_heads, head_dim) @@ -4036,6 +4046,7 @@ def test_v2_kda_replay_relocates_live_slot_history(): @skip_no_cuda def test_v2_kda_token_states_relocate_with_their_slot(): + """Without the per-draft replay caches, a slot's conv caches, drafts' records and draft count move with it.""" mgr = _build_v2_hybrid_with_mamba_layer( spec_config=MTPDecodingConfig(max_draft_len=2), conv_state_layout="q_k_v", @@ -4045,17 +4056,27 @@ def test_v2_kda_token_states_relocate_with_their_slot(): kda_token_states=True, ) try: - states = mgr.kda_state_tok - assert states is not None - states.zero_() - states[:, 0].fill_(1.0) - states[:, 1].fill_(2.0) - source_zero, source_one = states[:, 0].clone(), states[:, 1].clone() + assert mgr.kda_qkg_cache is None + assert mgr.kda_v_cache is None + assert mgr.kda_beta_cache is None + mgr.prev_num_accepted_tokens[:3].copy_( + torch.tensor([10, 20, 30], dtype=torch.int32, device="cuda") + ) + slot_buffers = (mgr.kda_conv_q, mgr.kda_conv_k, mgr.kda_conv_v, mgr.kda_state_tok) + source_rows = [] + for offset, slot_buffer in enumerate(slot_buffers): + assert slot_buffer is not None + slot_buffer.zero_() + slot_buffer[:, 0].fill_(offset * 10 + 1) + slot_buffer[:, 1].fill_(offset * 10 + 2) + source_rows.append((slot_buffer[:, 0].clone(), slot_buffer[:, 1].clone())) mgr._relocate_kda_replay_slots([0, 1], [1, 2]) - torch.testing.assert_close(states[:, 1], source_zero, rtol=0, atol=0) - torch.testing.assert_close(states[:, 2], source_one, rtol=0, atol=0) + assert mgr.prev_num_accepted_tokens[:3].tolist() == [10, 10, 20] + for slot_buffer, (source_zero, source_one) in zip(slot_buffers, source_rows): + torch.testing.assert_close(slot_buffer[:, 1], source_zero, rtol=0, atol=0) + torch.testing.assert_close(slot_buffer[:, 2], source_one, rtol=0, atol=0) finally: mgr.shutdown() # Released with the other replay buffers. @@ -4190,6 +4211,68 @@ def dim_contiguous_conv_cache() -> torch.Tensor: assert mgr.prev_num_accepted_tokens[2] == 5 +def test_v2_kda_token_states_seed_disaggregated_generation_slots(): + """Without the per-draft replay caches, seeding a transferred request's slot sets its conv caches and clears its + draft count; its drafts' records stay as they are, since with no pending drafts none of them is read.""" + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._use_kda_replay_update = True + mgr.local_num_mamba_layers = 2 + mgr.conv_state_shape = [6, 4] + mgr.conv_section_dims = [2, 2, 2] + mgr._request_id_to_state_index = {101: 1, 202: 3} + mgr.prev_num_accepted_tokens = torch.full((4,), 5, dtype=torch.int32) + mgr.all_conv_states = [ + torch.arange(4 * 6 * 4, dtype=torch.float32).reshape(4, 6, 4), + torch.arange(4 * 6 * 4, dtype=torch.float32).reshape(4, 6, 4) + 1000, + ] + + def dim_contiguous_conv_cache() -> torch.Tensor: + return torch.full((2, 4, 6, 2), 7.0).transpose(-1, -2) + + mgr.kda_conv_q = dim_contiguous_conv_cache() + mgr.kda_conv_k = dim_contiguous_conv_cache() + mgr.kda_conv_v = dim_contiguous_conv_cache() + mgr.kda_qkg_cache = None + mgr.kda_v_cache = None + mgr.kda_beta_cache = None + mgr.kda_state_tok = torch.full((2, 4, 3, 2, 1, 2), 7.0) + + mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) + + for layer_offset, conv_state in enumerate(mgr.all_conv_states): + for slot in (1, 3): + for replay_buffer, start in ( + (mgr.kda_conv_q, 0), + (mgr.kda_conv_k, 2), + (mgr.kda_conv_v, 4), + ): + torch.testing.assert_close( + replay_buffer[layer_offset, slot, :, :4], + conv_state[slot, start : start + 2, :], + ) + assert torch.count_nonzero(replay_buffer[layer_offset, slot, :, 4:]) == 0 + assert mgr.prev_num_accepted_tokens[slot] == 0 + assert mgr.prev_num_accepted_tokens.tolist() == [5, 0, 5, 0] + assert (mgr.kda_state_tok == 7.0).all() + + +def test_kda_replay_path_stays_on_with_the_drafts_records(): + """``has_kda_replay_caches`` follows the replay conv caches: it holds with the per-draft replay caches and with + the drafts' records in their place, and not without the replay path.""" + t = torch.zeros(1) + conv_caches = dict(kda_conv_q=t, kda_conv_k=t, kda_conv_v=t) + with_caches = PythonMambaCacheManager.SpeculativeState( + conv=t, temporal=t, kda_qkg_cache=t, kda_v_cache=t, kda_beta_cache=t, **conv_caches + ) + with_records = PythonMambaCacheManager.SpeculativeState( + conv=t, temporal=t, kda_state_tok=t, **conv_caches + ) + assert with_caches.has_kda_replay_caches + assert with_records.has_kda_replay_caches + assert not PythonMambaCacheManager.SpeculativeState(conv=t, temporal=t).has_kda_replay_caches + assert not PythonMambaCacheManager.State(conv=t, temporal=t).has_kda_replay_caches + + def test_v2_kda_replay_seeds_bf16_conv_state_from_disagg_transfer(): """The conv half of the KDA state crosses ctx->gen in bf16. From 003f8d088d93e333a0ca5f5fd6e6dc716997607f Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 13:09:19 -0700 Subject: [PATCH 50/67] [None][fix] KDA suspend: commit and relocate without the replay caches With the Kimi K3 verify kernels' draft records (kda_state_tok, #19817), the per-draft replay caches kda_qkg_cache, kda_v_cache and kda_beta_cache are not allocated. The suspend-time commit indexed them per layer and the slot relocation asserted them, so both would fail on such a manager. The commit now passes None for absent caches (the records' replay never reads them), and the relocation moves the buffers that exist, the drafts' records included. With the replay caches nothing changes. The new test builds a manager with the records and without the caches. It compares the suspend-time commit bit for bit with the records' replay and the conv store run on their own, and checks that relocation moves a slot's conv caches, records and draft count. Signed-off-by: Vasanth Sabavat --- .../_torch/modules/kimi_kda/_kda_kernels.py | 10 ++- .../kv_cache/mamba_cache_manager.py | 25 +++++-- .../kv_cache/test_mamba_cache_manager.py | 73 +++++++++++++++++++ 3 files changed, 97 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index 48189c103d47..4eabac524b5f 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -450,9 +450,9 @@ def commit_kda_pending_drafts( conv_q: torch.Tensor, conv_k: torch.Tensor, conv_v: torch.Tensor, - qkg_cache: torch.Tensor, - v_cache: torch.Tensor, - beta_cache: torch.Tensor, + qkg_cache: Optional[torch.Tensor], + v_cache: Optional[torch.Tensor], + beta_cache: Optional[torch.Tensor], state_indices: torch.Tensor, num_accepted_tokens: torch.Tensor, state_tok: Optional[torch.Tensor] = None, @@ -466,7 +466,8 @@ def commit_kda_pending_drafts( bit, and the caller clears the slot's count: * with the Kimi K3 verify kernels' draft records (``state_tok``), the - drafts are replayed from those (``replay_kda_token_records``); + drafts are replayed from those (``replay_kda_token_records``), and the + replay caches, which are then not allocated, may be None; * otherwise they are replayed from the replay caches by the fused verify kernel's own replay steps (``kda_mtp_commit_pending_drafts``). """ @@ -474,6 +475,7 @@ def commit_kda_pending_drafts( if state_tok is not None: replay_kda_token_records(recurrent_state, state_tok, state_indices, num_accepted_tokens) else: + assert qkg_cache is not None and v_cache is not None and beta_cache is not None _load_mtp_module().kda_mtp_commit_pending_drafts( recurrent_state, conv_q, diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index c7cafe41f02b..523a9fa42471 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -4292,11 +4292,17 @@ def _kda_replay_state_buffers(self) -> List[Tuple[torch.Tensor, int]]: (buffer, slot dimension) pairs.""" assert self.prev_num_accepted_tokens is not None buffers = [(self.prev_num_accepted_tokens, 0)] - for replay_buffer in (self.kda_conv_q, self.kda_conv_k, self.kda_conv_v, - self.kda_qkg_cache, self.kda_v_cache, - self.kda_beta_cache): + for replay_buffer in (self.kda_conv_q, self.kda_conv_k, + self.kda_conv_v): assert replay_buffer is not None buffers.append((replay_buffer, 1)) + # The per-draft replay caches, or the verify kernels' draft records + # (kda_state_tok) in their place. + for replay_buffer in (self.kda_qkg_cache, self.kda_v_cache, + self.kda_beta_cache, + getattr(self, "kda_state_tok", None)): + if replay_buffer is not None: + buffers.append((replay_buffer, 1)) return buffers def _relocate_kda_replay_slots(self, old_slots: List[int], @@ -4378,7 +4384,8 @@ def _commit_kda_pending_drafts(self, slot: int) -> None: assert self.prev_num_accepted_tokens is not None # With the verify kernels' draft records (kda_state_tok), the accepted - # drafts are replayed from those instead of the caches. + # drafts are replayed from those, and the replay caches are not + # allocated. state_tok = getattr(self, "kda_state_tok", None) with self._on_kv_cache_stream(): slots = torch.full((1, ), @@ -4388,15 +4395,19 @@ def _commit_kda_pending_drafts(self, slot: int) -> None: num_accepted = self.prev_num_accepted_tokens.index_select(0, slots) for layer_offset, (ssm_states, conv_states) in enumerate( zip(self.all_ssm_states, self.all_conv_states)): + qkg_cache, v_cache, beta_cache = ( + None if replay_cache is None else replay_cache[layer_offset] + for replay_cache in (self.kda_qkg_cache, self.kda_v_cache, + self.kda_beta_cache)) commit_kda_pending_drafts( ssm_states, conv_states, self.kda_conv_q[layer_offset], self.kda_conv_k[layer_offset], self.kda_conv_v[layer_offset], - self.kda_qkg_cache[layer_offset], - self.kda_v_cache[layer_offset], - self.kda_beta_cache[layer_offset], + qkg_cache, + v_cache, + beta_cache, slots, num_accepted, state_tok=(state_tok[layer_offset] diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index d76aaf4a6fcb..96a5d57824f1 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -4121,6 +4121,79 @@ def test_v2_kda_replay_suspend_resume_matches_an_unsuspended_run(slot_taken): assert torch.equal(states, expected_states) +@skip_no_cuda +def test_v2_kda_replay_commit_and_relocation_without_the_replay_caches(): + """With the verify kernels' draft records (``kda_state_tok``) in place of the per-draft replay caches, the + suspend-time commit replays a slot's records onto its V2 SSM state and stores its conv window, bit for bit as the + records' replay and the conv store do on their own, and relocation moves a slot's conv caches, records and draft + count.""" + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import ( + replay_kda_token_records, + store_kda_replay_conv_window, + ) + + num_layers, num_slots, num_spec, num_heads, head_dim, width = 2, 4, 2, 6, 128, 4 + section_dim = num_heads * head_dim + generator = torch.Generator(device="cuda").manual_seed(13) + + def randn(*shape): + return 0.1 * torch.randn(*shape, generator=generator, device="cuda") + + mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._use_kda_replay_update = True + mgr._stream = torch.cuda.Stream() + mgr.all_ssm_states = [ + randn(num_slots, num_heads, head_dim, head_dim) for _ in range(num_layers) + ] + mgr.all_conv_states = [randn(num_slots, 3 * section_dim, width - 1) for _ in range(num_layers)] + mgr.kda_conv_q, mgr.kda_conv_k, mgr.kda_conv_v = ( + randn(num_layers, num_slots, width - 1 + num_spec, section_dim).transpose(-1, -2) + for _ in range(3) + ) + mgr.kda_qkg_cache = None + mgr.kda_v_cache = None + mgr.kda_beta_cache = None + mgr.kda_state_tok = randn(num_layers, num_slots, 3, num_spec, num_heads, head_dim) + mgr.prev_num_accepted_tokens = torch.tensor([0, 2, 1, 0], dtype=torch.int32, device="cuda") + + slot = 1 + slots = torch.tensor([slot], dtype=torch.int32, device="cuda") + num_accepted = mgr.prev_num_accepted_tokens[slots.long()].clone() + expected_ssm, expected_conv = [], [] + for layer_offset in range(num_layers): + ssm = mgr.all_ssm_states[layer_offset].clone() + conv = mgr.all_conv_states[layer_offset].clone() + replay_kda_token_records(ssm, mgr.kda_state_tok[layer_offset], slots, num_accepted) + store_kda_replay_conv_window( + conv, + mgr.kda_conv_q[layer_offset], + mgr.kda_conv_k[layer_offset], + mgr.kda_conv_v[layer_offset], + slots, + num_accepted, + ) + expected_ssm.append(ssm) + expected_conv.append(conv) + + mgr._commit_kda_pending_drafts(slot) + torch.cuda.synchronize() + + assert mgr.prev_num_accepted_tokens.tolist() == [0, 0, 1, 0] + for actual, expected in zip(mgr.all_ssm_states, expected_ssm): + assert torch.equal(actual, expected) + for actual, expected in zip(mgr.all_conv_states, expected_conv): + assert torch.equal(actual, expected) + + slot_buffers = (mgr.kda_conv_q, mgr.kda_conv_k, mgr.kda_conv_v, mgr.kda_state_tok) + source_rows = [slot_buffer[:, 2].clone() for slot_buffer in slot_buffers] + + mgr._relocate_kda_replay_slots([2], [3]) + + assert mgr.prev_num_accepted_tokens.tolist() == [0, 0, 1, 1] + for slot_buffer, source_row in zip(slot_buffers, source_rows): + assert torch.equal(slot_buffer[:, 3], source_row) + + @skip_no_cuda def test_v2_kda_replay_relocates_live_slot_history(): mgr = _build_v2_hybrid_with_mamba_layer( From de2f7e12efd66140b4184030da9832b25168fe5c Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 13:13:32 -0700 Subject: [PATCH 51/67] [None][fix] KV cache manager V2: count every live-floor SSM slot in the KDA quotas At the live floor the runtime keeps ceil(floor / max_util_for_resume) SSM slots, and kda_state_tok has an entry for each. The max_tokens and minimum quotas counted the states of the floor's slots only and the drafts' records of all of them, so the slots past the floor were left without their states. The full-state token buffers' share of the 1 / max_util_for_resume headroom used to cover that; the records do not. For Kimi K3 at batch 8 (9 slots, 10 kept), the max_tokens 16,896 estimation pool could resume about 16,469 attention tokens. The quotas now count every slot the floor keeps with its states and its records. For that K3 configuration the estimation pool resumes about 17,485 tokens (arithmetic on the logged sizes). The pools that serving sizes from the memory fraction, and models without the records, do not change. The quota tests now expect the extra slots' states. Signed-off-by: Vasanth Sabavat --- .../pyexecutor/kv_cache/mamba_cache_manager.py | 15 ++++++++------- .../executor/kv_cache/test_mamba_cache_manager.py | 14 ++++++++++---- 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 35c599d5b8fa..8f4767ca0efb 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3962,13 +3962,14 @@ def _get_typical_request_capacity( return fallback_capacity def _state_quota(self, state_slots: int) -> int: - """Quota for ``state_slots`` recurrent-state slots, plus the drafts' - records of the further slots the live floor keeps - (``_live_floor_token_state_slots``).""" - extra_slots = self._live_floor_token_state_slots( - self.kv_cache_config.max_util_for_resume) - state_slots - return (state_slots * self._mamba_state_bytes_per_slot() + - max(0, extra_slots) * self._kda_token_state_bytes_per_slot()) + """Quota for ``state_slots`` recurrent-state slots, or for every slot + the live floor keeps (``_live_floor_token_state_slots``) when that is + more: each holds its states and its drafts' records.""" + slots = max( + state_slots, + self._live_floor_token_state_slots( + self.kv_cache_config.max_util_for_resume)) + return slots * self._mamba_state_bytes_per_slot() def _get_quota_from_max_tokens(self, max_tokens: int) -> int: attention_quota = super()._get_quota_from_max_tokens(max_tokens) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 799373ba639f..074c5c19ad16 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3880,15 +3880,17 @@ def test_v2_kda_token_states_reserve_quota_before_sizing(): plain = _v2_kda_token_state_manager(kda_token_states=False, enable_block_reuse=False) assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 3 # The pool keeps ceil(3 / max_util_for_resume) = 4 slots, each with the records of 2 drafts (3 x 8 x 256 fp32 - # each); the quotas without them differ by exactly those. + # each). The quotas count the 4 slots' records and the 4th slot's states, which the quotas without the records + # leave out; the build takes the records out of the quota. reserved = 4 * 2 * 3 * 8 * 256 * 4 + extra_slot_states = plain._mamba_state_bytes_per_slot() minimum = mgr._minimum_live_gpu_quota() - assert minimum - plain._minimum_live_gpu_quota() == reserved + assert minimum - plain._minimum_live_gpu_quota() == reserved + extra_slot_states for max_tokens in (0, 4096): assert ( mgr._get_quota_from_max_tokens(max_tokens) - plain._get_quota_from_max_tokens(max_tokens) - == reserved + == reserved + extra_slot_states ) def build(quota): @@ -3922,7 +3924,11 @@ def test_v2_kda_token_states_reserve_the_slots_the_runtime_keeps(): assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 19 reserved = 21 * 2 * 3 * 8 * 256 * 4 minimum = mgr._minimum_live_gpu_quota() - assert minimum - plain._minimum_live_gpu_quota() == reserved + # The 2 slots past the 19 resident ones count with their states as well. + assert ( + minimum - plain._minimum_live_gpu_quota() + == reserved + 2 * plain._mamba_state_bytes_per_slot() + ) quota = minimum + (64 << 20) built = mgr._build_cache_config( KVCacheManagerConfig( From 0d83ef50b6fd4f7eff5aca9e679fc0a82dad03d5 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 20:33:36 -0700 Subject: [PATCH 52/67] Runtime branch: the blank line ruff-format wants between two merged test functions Signed-off-by: Vasanth Sabavat --- .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 6d8f63b48196..7d966ae51f81 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3963,6 +3963,7 @@ def test_v2_kda_token_states_reserve_the_slots_the_runtime_keeps(): ) assert quota - built.cache_tiers[0].quota == reserved + def test_mamba_cache_manager_delegates_kda_replay_capability() -> None: mgr = object.__new__(MambaCacheManager) mgr._impl = SimpleNamespace(use_kda_replay_update=True) From 35433fb616cd757b4726771bd447f3ea642d2fa1 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 7 Oct 2026 21:34:04 -0700 Subject: [PATCH 53/67] [None][fix] thop_attention catalog test: do not probe a context length past the call's tokens test_bf16_paged_context_reads_the_pool_and_batch_state checks that a context row's context_lengths must be the call's new-token count by running two wrong values against the right answer. One of them, the sequence's full KV length (84 on a 20-token call), makes the paged context FMHA write 84 rows into the op's 20-row output: compute-sanitizer memcheck reports 441 invalid 16-byte global writes in fmhaSm100aKernel_...PagedKvCausal...StaticContext from that call, in a run of modeling_v2/activation and this file on main. Natively the writes land in whatever the allocator placed after the output, and a later test in the same process can fail far from its cause. The out-of-bounds value is not a contract violation the op reports, so it cannot be probed safely here; the test keeps the 0 case. Signed-off-by: Vasanth Sabavat --- .../test_modeling_v2_thop_attention.py | 38 ++++++++++--------- 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py index cb0a1f090d4c..dd5ca7590e7b 100644 --- a/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py +++ b/tests/unittest/_torch/modeling_v2/attention/test_modeling_v2_thop_attention.py @@ -2300,8 +2300,11 @@ def test_bf16_paged_context_reads_the_pool_and_batch_state() -> None: point and is read straight back from there). 3. use_paged_context_fmha=False on a context call with a cached prefix is silently wrong — no exception, an answer far off the reference. - 4. context_lengths on a context row must be this call's new-token count. - The sequence's full KV length and 0 are both far off. + 4. context_lengths on a context row must be this call's new-token count; + 0 is far off. A larger value is not probed: the context FMHA then + writes past the output buffer, which holds this call's new tokens + only (compute-sanitizer memcheck: invalid global writes), and corrupts + whatever lies beyond it. The wrong batch states of (3) and (4) also corrupt the pool, so every rival runs against its own freshly prepared copy of one fixed state: @@ -2387,22 +2390,21 @@ def prepared() -> Tuple[_PagedAttnEnv, torch.Tensor]: base, "cached-prefix context at use_paged_context_fmha=False", ) - for ctx_len in (cached + new, 0): - cenv, cq = prepared() - _assert_far_outside( - cenv.call_op( - cq, - [new], - 1, - [0], - MASK_CAUSAL, - record=False, - ctx_lens_override=[ctx_len], - use_paged_context_fmha=True, - ), - base, - f"context_lengths={ctx_len}", - ) + cenv, cq = prepared() + _assert_far_outside( + cenv.call_op( + cq, + [new], + 1, + [0], + MASK_CAUSAL, + record=False, + ctx_lens_override=[0], + use_paged_context_fmha=True, + ), + base, + "context_lengths=0", + ) # (2) page by page. Pages 0-1 hold the cached prefix and must be read; # page 2 holds only tokens 64..83, which this call appends itself. From bdadd409bd0b94f3f658e665267f86cf108b89f3 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 8 Oct 2026 23:22:48 -0700 Subject: [PATCH 54/67] [None][test] List the MNNVL capture-guard test in l0_dgx_b200.yml Review r4227132758: test_mnnvl_workspace_creation_refuses_graph_capture lives under unittest/_torch/multi_gpu, so it belongs with the multi-GPU B200 lists. Move it from l0_b200.yml to the 2-GPU pre_merge block of l0_dgx_b200.yml, next to the other unittest/_torch/multi_gpu entry. Signed-off-by: Vasanth Sabavat --- tests/integration/test_lists/test-db/l0_b200.yml | 3 --- tests/integration/test_lists/test-db/l0_dgx_b200.yml | 2 ++ 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 052dfcc1d22c..2999194426d3 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -153,9 +153,6 @@ l0_b200: - unittest/_torch/modeling/test_kimi_linear_checkpoint.py # CPU self-test of the Kimi K3 disagg parity harness comparison logic. - test_kimi_k3_specdec.py::test_kimi_k3_disagg_parity_selftest - # ------------- distributed --------------- - # One process: the MNNVL workspace's capture guard raises before any MNNVL call. - - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_creation_refuses_graph_capture # ------------- modules (non-MoE) --------------- - unittest/_torch/modules/test_fused_add_rms_norm_quant.py - unittest/_torch/modules/test_fused_activation_quant.py diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 1f3b81d69be7..066b51ae50c4 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -17,6 +17,8 @@ l0_dgx_b200: tests: - unittest/_torch/misc/test_autotuner.py::test_autotuner_distributed_strategy - unittest/_torch/multi_gpu/test_minimax_m3_piecewise.py::test_minimax_m3_piecewise_empty_adp_rank_preserves_caches + # One process: the MNNVL workspace's capture guard raises before any MNNVL call. + - unittest/_torch/multi_gpu/test_mnnvl_allreduce.py::test_mnnvl_workspace_creation_refuses_graph_capture - accuracy/test_llm_api_pytorch.py::TestQwen3_5_35B_A3B::test_bf16[tp2-TRTLLM] - disaggregated/test_disaggregated_beam_search.py::test_disaggregated_beam_search_v2[beam2-reuse] - disaggregated/test_disaggregated_beam_search.py::test_disaggregated_beam_search_v2[beam4-no-reuse] From cd21be6a568ab4188b5cbdfb543b6b17591f639f Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 9 Oct 2026 01:45:33 -0700 Subject: [PATCH 55/67] [None][fix] Enable the spec step-copy kernels on SM 100 only Review r4228224944: is_supported() also admitted SM 103 and SM 107, where the kernels are untested. Admit SM 100 only; other GPUs keep the torch path in the engine and the sampler. Signed-off-by: Vasanth Sabavat --- .../_torch/cute_dsl_kernels/spec_step_copies/op.py | 11 +++-------- tests/integration/test_lists/test-db/l0_b200.yml | 2 +- .../_torch/cute_dsl_kernels/test_spec_step_copies.py | 2 +- .../_torch/executor/test_pytorch_model_engine.py | 2 +- 4 files changed, 6 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py index 88208283e6cb..feba742bdefb 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py @@ -35,17 +35,12 @@ def is_supported() -> bool: - """Whether the kernels run on the current device: the SM 100 family (SM 100, 103, 107) with the CuTe DSL installed. + """Whether the kernels run on the current device: SM 100 (B200 / GB200) with the CuTe DSL installed. The kernels are plain SIMT code (thread and block indices, global loads and stores), with nothing specific to an - architecture. + architecture, but they are validated on SM 100 only; SM 103 and SM 107 keep the torch path. """ - # TODO: validated on SM 100 (B200 / GB200) only; SM 103 (GB300) and SM 107 (Rubin) are untested. - return ( - IS_CUTLASS_DSL_AVAILABLE - and torch.cuda.is_available() - and get_sm_version() in (100, 103, 107) - ) + return IS_CUTLASS_DSL_AVAILABLE and torch.cuda.is_available() and get_sm_version() == 100 def _is_i32(t: torch.Tensor) -> bool: diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index f1201e9caecb..6ed1802473e1 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -246,7 +246,7 @@ l0_b200: # hw-agnostic and needs its own Blackwell run. - unittest/_torch/speculative/test_fused_sampling_op.py # One-model speculative decoding step-copy kernels and the engine's decode - # step that runs them (SM 100 family; they skip on other GPUs). + # step that runs them (SM 100; they skip on other GPUs). - unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py - unittest/_torch/executor/test_pytorch_model_engine.py -k "test_spec_decode_graph_step_gather_kernel" - unittest/_torch/thop/parallel TIMEOUT (90) diff --git a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py index 58f9ba03102e..3edb55c42aef 100644 --- a/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py +++ b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py @@ -65,7 +65,7 @@ def _supported() -> bool: pytestmark = pytest.mark.skipif( - not _supported(), reason="the step-copy kernels run on SM 100 / 103 / 107 with the CuTe DSL" + not _supported(), reason="the step-copy kernels run on SM 100 with the CuTe DSL" ) diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index b8703e54316a..774c1a8b496e 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -1772,7 +1772,7 @@ def test_spec_decode_graph_step_gather_kernel_matches_torch_path( from tensorrt_llm._torch.cute_dsl_kernels.spec_step_copies import \ op as spec_step_copies if not spec_step_copies.is_supported(): - self.skipTest("the step-copy kernels run on SM 100 / 103 / 107") + self.skipTest("the step-copy kernels run on SM 100") max_draft_len = 3 tokens_per_step = max_draft_len + 1 model_engine, kv_cache_manager = create_model_engine_and_kvcache( From 4552c8d489fe17faddcaff53197c81e733c7c748 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 6 Oct 2026 13:45:56 -0700 Subject: [PATCH 56/67] [None][feat] DSpark: load speculators-format drafter checkpoints A DSpark drafter saved in the vLLM speculators format, such as RedHatAI/Kimi-K3-speculator.dspark as downloaded, now loads as speculative_model without a converted copy. Its config.json has no top-level model_type, so it could not be read before. The config is detected by speculators_model_type == "dspark" with a nested transformer_layer_config, and translated the way the hand conversion (make_redhat_dspark_adapter.py) did: - the qwen3 backbone is flattened from transformer_layer_config; - target_layer_ids = aux_hidden_state_layer_ids - 1, since vLLM-style aux ids name the layer whose input is captured; - the sliding window is dropped. The drafter attention backends reject non-causal windows, and this is exact while the context fits in the trained window (2048 tokens for the Kimi K3 drafter); - mask_token_id and the markov / confidence-head fields are carried over; block_size is not. Every reader of the drafter's config.json uses the same translation: the model config (load_pretrained_config), the DFlash / DSpark spec-config resolution, and the context-buffer budget. Other drafter configs are unchanged. Signed-off-by: Vasanth Sabavat (cherry picked from commit e2536950a6bf9b110cb3120dfb71d059cec3c9f6) --- .../_torch/pyexecutor/config_utils.py | 16 +- tensorrt_llm/llmapi/llm_args.py | 116 ++++++++- .../test_dspark_speculators_checkpoint.py | 236 ++++++++++++++++++ 3 files changed, 364 insertions(+), 4 deletions(-) create mode 100644 tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index ce1d56453ea2..8c49ec9783a7 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -9,7 +9,9 @@ import transformers from tensorrt_llm._utils import str_dtype_to_torch -from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig +from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, + is_speculators_dspark_config, + translate_speculators_dspark_config) from tensorrt_llm.logger import logger if TYPE_CHECKING: @@ -1018,6 +1020,18 @@ def load_pretrained_config(model_name_or_path: str, MistralConfigLoader model_config = MistralConfigLoader().load( model_name_or_path).pretrained_config + elif is_speculators_dspark_config(config_dict): + # A DSpark drafter in the vLLM speculators format: no top-level + # model_type, and its auto_map needs the speculators package. Build + # the backbone's config from the translation the spec config uses. + from transformers.models.auto.configuration_auto import CONFIG_MAPPING + translated = translate_speculators_dspark_config(config_dict) + model_config = CONFIG_MAPPING["qwen3"].from_dict(translated) + logger.info( + f"{model_name_or_path}: speculators-format DSpark drafter, read as " + f"qwen3 with target_layer_ids " + f"{translated['dflash_config']['target_layer_ids']} " + f"(aux_hidden_state_layer_ids - 1) and no sliding window") elif is_qwen_image_bench_config(config_dict): from tensorrt_llm._torch.models.modeling_qwen3_5 import ( Qwen35ConfigCompat, _normalize_qwen35_quantization_config) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f241528a2b31..2fb8ae95f71e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3144,6 +3144,113 @@ def spec_dec_mode(self): return TorchSpeculativeDecodingMode.PARD +# Backbone fields a speculators drafter nests under transformer_layer_config. +# The rest of that dict (flex_attention_backend, use_cache, ...) is training +# configuration. +_SPECULATORS_BACKBONE_KEYS = ( + "hidden_size", + "intermediate_size", + "num_hidden_layers", + "num_attention_heads", + "num_key_value_heads", + "head_dim", + "hidden_act", + "rms_norm_eps", + "max_position_embeddings", + "vocab_size", + "attention_bias", + "attention_dropout", + "bos_token_id", + "eos_token_id", + "pad_token_id", + "initializer_range", + "rope_parameters", +) +_SPECULATORS_DSPARK_HEAD_KEYS = ( + "markov_rank", + "markov_head_type", + "enable_confidence_head", + "confidence_head_with_markov", + "draft_vocab_size", +) + + +def is_speculators_dspark_config(config_dict: dict) -> bool: + """Whether a drafter ``config.json`` is a speculators-format DSpark drafter. + + The vLLM speculators format (e.g. RedHatAI/Kimi-K3-speculator.dspark) has + no top-level ``model_type``: it names the algorithm in + ``speculators_model_type`` and nests the backbone under + ``transformer_layer_config``. + """ + return (config_dict.get("speculators_model_type") == "dspark" + and isinstance(config_dict.get("transformer_layer_config"), dict)) + + +def translate_speculators_dspark_config(config_dict: dict) -> dict: + """Translate a speculators-format DSpark ``config.json`` for TRT-LLM. + + The result is the config TRT-LLM's GQA DSpark drafter reads. Any other + config is returned unchanged. + + - The qwen3 backbone is flattened from ``transformer_layer_config``. + - ``aux_hidden_state_layer_ids`` name the layer whose input is captured; + ``target_layer_ids`` name the 0-indexed layer whose output is, so each id + is one less. + - The sliding window is dropped, because the drafter attention backends + reject non-causal windows. That is exact while the context fits in the + trained window (2048 tokens for the Kimi K3 drafter). + - The checkpoint's ``block_size`` is not carried over: the runtime block + follows ``max_draft_len``. + + The weights need no renaming. Every config.json reader of the drafter + (model config, spec config resolution, the context-buffer budget) goes + through this function, so they all see the same drafter. + """ + if not is_speculators_dspark_config(config_dict): + return config_dict + layer = config_dict["transformer_layer_config"] + if layer.get("model_type") != "qwen3": + raise ValueError( + "speculators DSpark drafters are supported with a qwen3 backbone; " + f"this one has model_type {layer.get('model_type')!r}") + aux_layer_ids = config_dict.get("aux_hidden_state_layer_ids") + if not aux_layer_ids or min(aux_layer_ids) < 1: + raise ValueError( + "a speculators DSpark drafter needs aux_hidden_state_layer_ids >= 1 " + "(layer 0's input is the embedding output, which is not captured); " + f"got {aux_layer_ids}") + num_layers = layer["num_hidden_layers"] + dflash_config = { + "target_layer_ids": [int(layer_id) - 1 for layer_id in aux_layer_ids] + } + if config_dict.get("mask_token_id") is not None: + dflash_config["mask_token_id"] = config_dict["mask_token_id"] + translated = { + key: layer[key] + for key in _SPECULATORS_BACKBONE_KEYS if key in layer + } + translated.update( + architectures=["Qwen3ForCausalLM"], + model_type="qwen3", + tie_word_embeddings=bool(config_dict.get("tie_word_embeddings", False)), + torch_dtype=(config_dict.get("dtype") or config_dict.get("torch_dtype") + or "bfloat16"), + use_sliding_window=False, + sliding_window=None, + max_window_layers=num_layers, + layer_types=["full_attention"] * num_layers, + dflash_config=dflash_config, + ) + rope_theta = (layer.get("rope_parameters") or {}).get("rope_theta") + if rope_theta is not None: + translated["rope_theta"] = rope_theta + for key in _SPECULATORS_DSPARK_HEAD_KEYS: + if config_dict.get(key) is not None: + translated[key] = config_dict[key] + return translated + + class DFlashDecodingConfig(DecodingBaseConfig): """Configuration for DFlash speculative decoding. @@ -3212,7 +3319,8 @@ def resolve_from_checkpoint(self) -> None: if not os.path.exists(draft_config_path): return with open(draft_config_path) as f: - dflash_cfg = json.load(f).get("dflash_config", {}) + dflash_cfg = translate_speculators_dspark_config(json.load(f)).get( + "dflash_config", {}) if self.target_layer_ids is None: layer_ids = dflash_cfg.get("target_layer_ids") @@ -6639,7 +6747,8 @@ def _validate_dflash_ctx_budget(self, "config.json") if os.path.exists(draft_config_path): with open(draft_config_path) as f: - draft_config = json.load(f) + draft_config = translate_speculators_dspark_config( + json.load(f)) arena_before_pool = self._kv_cache_estimation_runs() if spec_cfg.skip_ctx_buffer_budget_check: @@ -6896,7 +7005,8 @@ def validate_speculative_config(self): "config.json") if os.path.exists(draft_config_path): with open(draft_config_path) as f: - draft_cfg = json.load(f) + draft_cfg = translate_speculators_dspark_config( + json.load(f)) dspark_cfg = draft_cfg.get("dspark_config") or {} dflash_cfg = draft_cfg.get("dflash_config") or {} diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py new file mode 100644 index 000000000000..496db57dc970 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py @@ -0,0 +1,236 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""A DSpark drafter saved in the vLLM speculators format (CPU-only, no weights). + +The reference is the config that a hand conversion +(make_redhat_dspark_adapter.py) wrote for RedHatAI/Kimi-K3-speculator.dspark, +which every Kimi K3 DSpark measurement so far used. The checkpoint as downloaded +must give exactly the drafter that the converted one gave. +""" + +import copy +import json + +import pytest + +from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config +from tensorrt_llm.llmapi.llm_args import ( + DSparkDecodingConfig, + TorchLlmArgs, + is_speculators_dspark_config, + translate_speculators_dspark_config, +) + +pytestmark = pytest.mark.cpu_only + +# RedHatAI/Kimi-K3-speculator.dspark config.json, as published. +SPECULATORS_CONFIG = { + "architectures": ["DSparkDraftModel"], + "auto_map": {"": "config.DSparkSpeculatorConfig"}, + "aux_hidden_state_layer_ids": [24, 48, 72, 88, 92], + "block_size": 8, + "confidence_head_with_markov": True, + "draft_vocab_size": 163840, + "dtype": "bfloat16", + "enable_confidence_head": True, + "markov_head_type": "vanilla", + "markov_rank": 256, + "mask_token_id": 163837, + "sample_from_anchor": True, + "sliding_window_non_causal": False, + "speculators_config": { + "algorithm": "dspark", + "default_proposal_method": "greedy", + "proposal_methods": [ + { + "accept_tolerance": 0.0, + "proposal_type": "greedy", + "speculative_tokens": 8, + "verifier_accept_k": 1, + } + ], + "verifier": { + "architectures": ["KimiK3ForConditionalGeneration"], + "name_or_path": "moonshotai/Kimi-K3", + }, + }, + "speculators_model_type": "dspark", + "speculators_version": "0.7.0.dev141", + "target_hidden_size": None, + "tie_word_embeddings": False, + "transformer_layer_config": { + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 163584, + "eos_token_id": 163586, + "flex_attention_backend": "FLASH", + "head_dim": 64, + "hidden_act": "silu", + "hidden_size": 7168, + "initializer_range": 0.02, + "intermediate_size": 14336, + "layer_types": ["sliding_attention"] * 5, + "max_position_embeddings": 1048576, + "max_window_layers": 28, + "model_type": "qwen3", + "num_attention_heads": 96, + "num_hidden_layers": 5, + "num_key_value_heads": 16, + "pad_token_id": 163839, + "rms_norm_eps": 1e-05, + "rope_parameters": {"rope_theta": 10000.0, "rope_type": "default"}, + "sliding_window": 2048, + "tie_word_embeddings": False, + "use_cache": True, + "use_sliding_window": True, + "vocab_size": 163840, + }, + "transformers_version": "5.14.1", +} + +# What the hand conversion wrote for it: the backbone flattened, the capture +# layers as 0-indexed layer outputs, the sliding window off. +CONVERTED_CONFIG = { + "hidden_size": 7168, + "intermediate_size": 14336, + "num_hidden_layers": 5, + "num_attention_heads": 96, + "num_key_value_heads": 16, + "head_dim": 64, + "hidden_act": "silu", + "rms_norm_eps": 1e-05, + "max_position_embeddings": 1048576, + "vocab_size": 163840, + "attention_bias": False, + "attention_dropout": 0.0, + "bos_token_id": 163584, + "eos_token_id": 163586, + "pad_token_id": 163839, + "initializer_range": 0.02, + "rope_parameters": {"rope_theta": 10000.0, "rope_type": "default"}, + "architectures": ["Qwen3ForCausalLM"], + "model_type": "qwen3", + "rope_theta": 10000.0, + "tie_word_embeddings": False, + "torch_dtype": "bfloat16", + "use_sliding_window": False, + "sliding_window": None, + "max_window_layers": 5, + "layer_types": ["full_attention"] * 5, + "dflash_config": {"mask_token_id": 163837, "target_layer_ids": [23, 47, 71, 87, 91]}, + "markov_rank": 256, + "markov_head_type": "vanilla", + "enable_confidence_head": True, + "confidence_head_with_markov": True, + "draft_vocab_size": 163840, +} + +# Bookkeeping that differs between a config read from a directory and one built +# from a dict, without changing the model. +_LOAD_METADATA = ("_name_or_path", "_commit_hash", "transformers_version") + + +def _checkpoint(tmp_path, name, config): + path = tmp_path / name + path.mkdir() + (path / "config.json").write_text(json.dumps(config)) + return str(path) + + +def _spec_config(draft_dir, max_draft_len=7): + args = TorchLlmArgs( + model="/tmp/dummy_model", + skip_tokenizer_init=True, + speculative_config=DSparkDecodingConfig( + max_draft_len=max_draft_len, speculative_model=draft_dir + ), + ) + return args.speculative_config + + +def _comparable(config): + values = config.to_dict() + for key in _LOAD_METADATA: + values.pop(key, None) + return values + + +def test_translation_is_the_converted_config(): + # Every config.json reader of the drafter goes through this translation. + assert translate_speculators_dspark_config(SPECULATORS_CONFIG) == CONVERTED_CONFIG + + +def test_drafter_config_matches_the_converted_checkpoint(tmp_path): + native = load_pretrained_config(_checkpoint(tmp_path, "native", SPECULATORS_CONFIG)) + converted = load_pretrained_config(_checkpoint(tmp_path, "converted", CONVERTED_CONFIG)) + + assert type(native) is type(converted) + assert native.architectures == ["Qwen3ForCausalLM"] + assert _comparable(native) == _comparable(converted) + + +def test_spec_config_matches_the_converted_checkpoint(tmp_path): + native = _spec_config(_checkpoint(tmp_path, "native", SPECULATORS_CONFIG)) + converted = _spec_config(_checkpoint(tmp_path, "converted", CONVERTED_CONFIG)) + + # vLLM-style aux ids name the layer whose input is captured. + assert native.target_layer_ids == [23, 47, 71, 87, 91] + assert native.mask_token_id == 163837 + assert native.markov_rank == 256 + # The checkpoint's block_size (8 under sample_from_anchor) does not pin + # max_draft_len: the converted checkpoint ran at 7. + assert native.block_size == 7 + for field in ( + "target_layer_ids", + "mask_token_id", + "markov_rank", + "markov_head_type", + "block_size", + ): + assert getattr(native, field) == getattr(converted, field), field + + +def test_trained_draft_length_is_accepted(tmp_path): + native = _spec_config(_checkpoint(tmp_path, "native", SPECULATORS_CONFIG), max_draft_len=8) + + assert native.block_size == 8 + assert native.target_layer_ids == [23, 47, 71, 87, 91] + + +@pytest.mark.parametrize( + "field, value, message", + [ + ("model_type", "llama", "supported with a qwen3 backbone"), + ("aux_hidden_state_layer_ids", [0, 48, 72, 88, 92], "aux_hidden_state_layer_ids"), + ], +) +def test_untranslatable_drafters_are_rejected(tmp_path, field, value, message): + config = copy.deepcopy(SPECULATORS_CONFIG) + if field == "model_type": + config["transformer_layer_config"]["model_type"] = value + else: + config[field] = value + + with pytest.raises(ValueError, match=message): + load_pretrained_config(_checkpoint(tmp_path, "native", config)) + + +def test_other_drafter_configs_are_left_alone(): + assert is_speculators_dspark_config(SPECULATORS_CONFIG) + assert not is_speculators_dspark_config(CONVERTED_CONFIG) + eagle3 = dict(SPECULATORS_CONFIG, speculators_model_type="eagle3") + assert not is_speculators_dspark_config(eagle3) + assert translate_speculators_dspark_config(CONVERTED_CONFIG) is CONVERTED_CONFIG + assert translate_speculators_dspark_config(eagle3) is eagle3 From 0fd5a3a30a877cb5e225b09f485b19c0f5ebef2d Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Fri, 9 Oct 2026 13:37:03 -0700 Subject: [PATCH 57/67] [None][fix] Ragged spec sampler tests: give the stub sampler B's store-copy state Main's test_ragged_spec_sampler builds SpecSampler without __init__, so it lacked the store scatter, slot table and pending store-copy event this branch adds; the stub now sets them and takes the main-stream host-copy path its transfer checks expect. The async-worker path copies verify_lens after the stores, as main does. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/speculative/spec_sampler_base.py | 8 ++++++-- .../speculative/hw_agnostic/test_ragged_spec_sampler.py | 4 ++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index fd7f17181c16..2dff89858628 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -567,13 +567,14 @@ def sample_async( next_draft_tokens=self.store.next_draft_tokens, ) - host_verify_lens = self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None if self._async_worker_active(): host_tensors = SampleStateTensorsSpec( new_tokens=self._copy_to_host(self.store.new_tokens), new_tokens_lens=self._copy_to_host(self.store.new_tokens_lens), next_draft_tokens=self._copy_to_host(self.store.next_draft_tokens), - verify_lens=host_verify_lens, + verify_lens=( + self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None + ), ) sampler_event = self._record_sampler_event() else: @@ -581,6 +582,9 @@ def sample_async( # needs the stores, so the host copies run on the side stream. # update_requests syncs their event before reading them, and the # next step's store update waits for it on this stream. + host_verify_lens = ( + self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None + ) with self._make_side_stream_copier() as copier: host_tensors = SampleStateTensorsSpec( new_tokens=copier.stage_copy_to_host(self.store.new_tokens), diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py index 7afbea6046a3..899f4620d830 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -608,6 +608,10 @@ def record_event(): sampler._copy_to_host = copy_to_host sampler._record_sampler_event = record_event + sampler._async_worker_active = lambda: True + sampler._store_scatter = None + sampler._slot_table = None + sampler._store_copies_done = None return sampler, transfers From 6a95a9e3ee939ea3bcfa3f00a509bf0af8e1c0c8 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Mon, 12 Oct 2026 00:13:12 -0700 Subject: [PATCH 58/67] [None][fix] KV cache manager V2: do not commit an unseeded disagg generation request's KDA replay caches On a disaggregated generation server a request's SSM slot is allocated at gen-init, before its conv and SSM states are transferred, and its replay caches are seeded from the transferred conv state only when its first generation step is prepared. A request whose transfer has completed is evictable before that step. Its suspend committed the slot's replay caches, which still held another request's window, over the transferred conv state, and the resume seeded the replay caches from it. The manager now tracks the requests whose slot is allocated at gen-init and whose replay caches are not seeded yet. A suspend of such a request keeps its transferred V2 states as they are; its resume seeds the replay caches from them. Signed-off-by: Vasanth Sabavat --- .../kv_cache/mamba_cache_manager.py | 14 +++++++- .../kv_cache/test_mamba_cache_manager.py | 36 +++++++++++++++++++ 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 9c50e331bf4d..0108de49cb14 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -3257,6 +3257,7 @@ def __init__( # With a host drafter: requests of a prepared batch whose acceptance # update_resources has not recorded yet. self._kda_replay_acceptance_pending: set[int] = set() + self._kda_replay_unseeded: set[int] = set() state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) @@ -4361,7 +4362,10 @@ def suspend_request(self, req: LlmRequest) -> None: f"Request {request_id} is being suspended before " "update_resources recorded the acceptance of its last " "verify, so its pending KDA drafts are unknown") - self._commit_kda_pending_drafts(slot) + if request_id in self._kda_replay_unseeded: + self._kda_replay_unseeded.remove(request_id) + else: + self._commit_kda_pending_drafts(slot) del self._request_id_to_state_index[request_id] self._kda_replay_suspended.add(request_id) super().suspend_request(req) @@ -4460,6 +4464,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self._branch_snapshot_points.pop(request.py_request_id, None) self._kda_replay_suspended.discard(request.py_request_id) self._kda_replay_acceptance_pending.discard(request.py_request_id) + self._kda_replay_unseeded.discard(request.py_request_id) super().free_resources(request, pin_on_release) def prepare_resources(self, scheduled_batch: ScheduledRequests): @@ -4471,6 +4476,11 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): num_contexts = len(scheduled_batch.context_requests) self._setup_state_indices(requests, num_contexts=num_contexts) self._reset_context_mamba_slots(num_contexts) + if self.use_kda_replay_update: + self._kda_replay_unseeded.update( + request.py_request_id + for request in scheduled_batch.context_requests + if request.is_disagg_generation_init_state) if self._record_kda_replay_in_update_resources: # Recorded by this batch's update_resources, or dropped by # revert_allocate_generation when the batch is not run. @@ -4587,6 +4597,7 @@ def seed_kda_replay_caches_for_disagg_gen(self, for request_id in request_ids if request_id in self._request_id_to_state_index ] + self._kda_replay_unseeded.difference_update(request_ids) if not slots: return @@ -4936,4 +4947,5 @@ def shutdown(self): self.kda_beta_cache = None self._kda_replay_suspended.clear() self._kda_replay_acceptance_pending.clear() + self._kda_replay_unseeded.clear() super().shutdown() diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index cc41b02f320d..b5fe79e5d740 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -3808,6 +3808,7 @@ def _v2_kda_replay_stub(num_slots=4): mgr._request_id_to_is_dummy = {} mgr._kda_replay_suspended = set() mgr._kda_replay_acceptance_pending = set() + mgr._kda_replay_unseeded = set() mgr._host_state_indices = torch.zeros(num_slots, dtype=torch.int32) mgr.cuda_state_indices = torch.zeros(num_slots, dtype=torch.int32) mgr._dummy_request_mask = None @@ -3920,6 +3921,39 @@ def test_v2_kda_replay_refuses_a_suspend_before_the_host_drafter_acceptance( assert not mgr.kv_cache_map[1].is_active +@pytest.mark.parametrize("seeded", [False, True], ids=["before_seeding", "after_seeding"]) +def test_v2_kda_replay_disagg_generation_request_suspends_its_transferred_state( + monkeypatch, seeded +): + mgr = _v2_kda_replay_stub() + monkeypatch.setattr(KVCacheManagerV2, "prepare_resources", MagicMock()) + request = SimpleNamespace(py_request_id=1, is_dummy=False, is_disagg_generation_init_state=True) + mgr.kv_cache_map[1] = _SlotCache(0) + _write_kda_replay_slot(mgr, 0, accepted=0, value=5.0) + mgr.prepare_resources(SimpleNamespace(context_requests=[request], generation_requests=[])) + request.is_disagg_generation_init_state = False + window = torch.arange(18.0).view(6, 3) + mgr.all_conv_states[0][0] = window + if seeded: + mgr.seed_kda_replay_caches_for_disagg_gen([1]) + + mgr.suspend_request(request) + assert mgr.committed_slots == ([0] if seeded else []) + assert 1 not in mgr._request_id_to_state_index + assert not mgr.kv_cache_map[1].is_active + + mgr.kv_cache_map[1].slot = 2 + mgr.kv_cache_map[1].is_active = True + mgr.all_conv_states[0][2] = window + mgr._setup_state_indices([request], num_contexts=0) + for section, name in enumerate(("kda_conv_q", "kda_conv_k", "kda_conv_v")): + assert torch.equal(getattr(mgr, name)[0, 2][:, :3], window[2 * section : 2 * section + 2]) + assert not mgr._kda_replay_unseeded + + mgr.suspend_request(request) + assert mgr.committed_slots == ([0, 2] if seeded else [2]) + + def _v2_kda_replay_manager_with_large_slots(): """A real V2 hybrid manager at Kimi K3's per-rank KDA shapes (6 heads, K = V = 128, conv width 4) on 6 layers. An SSM slot (~2.25 MiB) is larger than the 2 MiB allocation grain, so the SSM pool holds a handful of slots and runs @@ -4220,6 +4254,7 @@ def close(self) -> None: def test_v2_kda_replay_seeds_disaggregated_generation_slots(): mgr = object.__new__(MambaHybridCacheManagerV2) mgr._use_kda_replay_update = True + mgr._kda_replay_unseeded = set() mgr.local_num_mamba_layers = 2 mgr.conv_state_shape = [6, 4] mgr.conv_section_dims = [2, 2, 2] @@ -4277,6 +4312,7 @@ def test_v2_kda_replay_seeds_bf16_conv_state_from_disagg_transfer(): torch.manual_seed(0) mgr = object.__new__(MambaHybridCacheManagerV2) mgr._use_kda_replay_update = True + mgr._kda_replay_unseeded = set() mgr.local_num_mamba_layers = 2 mgr.conv_state_shape = [6, 4] mgr.conv_section_dims = [2, 2, 2] From 63aa7a223d548d572c1b0c8554970ddd6a7fee07 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Mon, 12 Oct 2026 07:15:38 -0700 Subject: [PATCH 59/67] [None][chore] KV cache manager V2: separate KDA draft-record accounting from the managed states Review follow-ups on the V2 manager's KDA draft records: - The constructor flag kda_token_states is now kda_draft_records (with keeps_kda_draft_records, _kda_draft_record_bytes_per_slot and _live_floor_draft_record_slots), so it is not confused with the kda_state_tok tensor that holds the records. The model contract that asks for them follows (`kda_draft_records` on the backbone). - _managed_state_bytes_per_slot gives the bytes of the states the V2 pools manage (SSM, conv and PLE); _mamba_state_bytes_per_slot is that plus the draft records, as before. - _build_cache_config keeps the device quota as total_gpu_quota and names the draft-record reservation and the remaining pool quota separately. - _quota_filling_request_capacity uses ceil_div. No behavior change. Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_torch/pyexecutor/_util.py | 6 +- .../kv_cache/mamba_cache_manager.py | 66 ++++++++++-------- .../kv_cache/test_mamba_cache_manager.py | 68 ++++++++++--------- 3 files changed, 74 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index a1c8466391c0..9a40f91c5c56 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3307,12 +3307,12 @@ def _create_kv_cache_manager( spec_config.tokens_per_gen_step - 1) # A model whose KDA verify replays the accepted drafts from # its own per-draft records (instead of the replay caches) - # asks for a buffer of them with `kda_token_states`. + # asks for a buffer of them with `kda_draft_records`. backbone = getattr(getattr(model_engine, "model", None), "model", None) if (issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2) - and getattr(backbone, "kda_token_states", False)): - kda_extra_kwargs["kda_token_states"] = True + and getattr(backbone, "kda_draft_records", False)): + kda_extra_kwargs["kda_draft_records"] = True if is_glm5_next: # The manager places an indexer buffer on every attention layer. from ..attention.backends.sparse.glm_kpool import \ diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 0cdf04c7fee0..5a91e35e52c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -52,6 +52,7 @@ from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from tensorrt_llm.math_utils import ceil_div from tensorrt_llm.runtime.kv_cache_manager_v2 import (DEFAULT_BEAM_INDEX, BatchDesc, BufferConfig, DataRole, @@ -335,7 +336,7 @@ def use_kda_replay_update(self) -> bool: return getattr(self, "_use_kda_replay_update", False) @property - def keeps_kda_token_states(self) -> bool: + def keeps_kda_draft_records(self) -> bool: """Whether the KDA replay caches come with the records of every draft of the last verify round (``kda_state_tok``), from which a verify kernel replays the accepted drafts onto the slot's state.""" @@ -3065,7 +3066,7 @@ def __init__( conv_state_layout: Literal["x_b_c", "q_k_v"] = "x_b_c", kda_replay_num_spec: Optional[int] = None, qwen4_exp_ple_cache_params: Optional["Qwen4ExpPLECacheParams"] = None, - kda_token_states: bool = False, + kda_draft_records: bool = False, **kwargs, ) -> None: if conv_state_layout not in ("x_b_c", "q_k_v"): @@ -3084,7 +3085,7 @@ def __init__( self._use_kda_replay_update = kda_replay_num_spec is not None # Only with the KDA replay caches: also keep the records of every # draft of the last verify round (kda_state_tok). - self._kda_token_states = kda_token_states and self._use_kda_replay_update + self._kda_draft_records = kda_draft_records and self._use_kda_replay_update if self._use_kda_replay_update: if use_replay_state_update: raise ValueError( @@ -3475,7 +3476,7 @@ def use_gdn_cached_replay_all_layer_commit(self) -> bool: return getattr(self, "_use_gdn_cached_replay_all_layer_commit", False) @property - def keeps_kda_token_states(self) -> bool: + def keeps_kda_draft_records(self) -> bool: return getattr(self, "kda_state_tok", None) is not None def _setup_mtp_intermediate_states(self, spec_config, @@ -3541,7 +3542,7 @@ def allocate_dim_contiguous_conv_cache() -> torch.Tensor: self.kda_conv_q = allocate_dim_contiguous_conv_cache() self.kda_conv_k = allocate_dim_contiguous_conv_cache() self.kda_conv_v = allocate_dim_contiguous_conv_cache() - if self._kda_token_states: + if self._kda_draft_records: # Beside the V2 pools. With the SSM pool at its live floor, # _build_cache_config took these bytes out of the GPU quota. The # verify replays the accepted drafts from these records, so the @@ -3878,19 +3879,18 @@ def _kda_draft_records_shape(self) -> Tuple[int, int, int, int]: return (3, self._kda_replay_num_spec, self.ssm_state_shape[0], self.ssm_state_shape[-1]) - def _kda_token_state_bytes_per_slot(self) -> int: + def _kda_draft_record_bytes_per_slot(self) -> int: """Bytes of ``kda_state_tok`` per SSM slot: the fp32 records of every draft of every local Mamba layer, or 0 without them.""" - if not getattr(self, "_kda_token_states", False): + if not getattr(self, "_kda_draft_records", False): return 0 return (self.local_num_mamba_layers * math.prod(self._kda_draft_records_shape()) * torch.float32.itemsize) - def _mamba_state_bytes_per_slot(self) -> int: + def _managed_state_bytes_per_slot(self) -> int: base_bytes = (self.local_num_mamba_layers * - (self.ssm_bytes + self.conv_bytes) + - self._kda_token_state_bytes_per_slot()) + (self.ssm_bytes + self.conv_bytes)) local_ple_layers = sum(layer_id in self.pp_layers for layer_id in self._ple_layer_ids) if local_ple_layers == 0: @@ -3901,6 +3901,10 @@ def _mamba_state_bytes_per_slot(self) -> int: math.prod(self._ple_ngram_context_shape) * torch.int64.itemsize) return base_bytes + local_ple_layers * ple_bytes + def _mamba_state_bytes_per_slot(self) -> int: + return (self._managed_state_bytes_per_slot() + + self._kda_draft_record_bytes_per_slot()) + def _ple_buffer_configs(self) -> list[BufferConfig]: if not self._ple_layer_ids: return [] @@ -3988,11 +3992,11 @@ def _get_typical_request_capacity( def _state_quota(self, state_slots: int) -> int: """Quota for ``state_slots`` recurrent-state slots, or for every slot - the live floor keeps (``_live_floor_token_state_slots``) when that is + the live floor keeps (``_live_floor_draft_record_slots``) when that is more: each holds its states and its drafts' records.""" slots = max( state_slots, - self._live_floor_token_state_slots( + self._live_floor_draft_record_slots( self.kv_cache_config.max_util_for_resume)) return slots * self._mamba_state_bytes_per_slot() @@ -4059,12 +4063,12 @@ def _ssm_pool_at_live_floor(self, kv_cache_config: KvCacheConfig) -> bool: known before sizing, so ``_build_cache_config`` takes those records' bytes out of the GPU quota. """ - return (getattr(self, "_kda_token_states", False) + return (getattr(self, "_kda_draft_records", False) and self.local_num_mamba_layers > 0 and not kv_cache_config.enable_block_reuse and self._attention_cache_bytes_per_token() > 0) - def _live_floor_token_state_slots(self, max_util_for_resume: float) -> int: + def _live_floor_draft_record_slots(self, max_util_for_resume: float) -> int: """SSM slots whose ``kda_state_tok`` entries come out of the GPU quota. With the SSM pool at its live floor and no ``pool_ratio``, that is @@ -4093,7 +4097,7 @@ def _quota_filling_request_capacity(self, gpu_quota: int) -> int: bytes_per_token = (self._max_resident_sequences() * self._attention_cache_bytes_per_token()) # KVCacheDesc.capacity is a 32-bit int. - return min(-(-gpu_quota // bytes_per_token), 1 << 30) + return min(ceil_div(gpu_quota, bytes_per_token), 1 << 30) def _build_cache_config( self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: @@ -4110,7 +4114,7 @@ def _build_cache_config( "are sized once from the SSM pool's slot count, and a KV pool " "rebalance can grow that pool past them: set " "kv_cache_config.enable_kv_pool_rebalance=False.") - if (getattr(self, "_kda_token_states", False) + if (getattr(self, "_kda_draft_records", False) and self.local_num_mamba_layers > 0 and (config.initial_pool_ratio is not None or not self._ssm_pool_at_live_floor(kv_cache_config))): @@ -4118,37 +4122,39 @@ def _build_cache_config( # only at the live floor is that count known before sizing, so # that its bytes can come out of the quota. raise ValueError( - "Per-token KDA states (kda_token_states) need the SSM pool at " + "KDA draft records (kda_draft_records) need the SSM pool at " "its live floor: kv_cache_config.enable_block_reuse off, " "kv_cache_config.pool_ratio unset, and attention layers on " "this rank.") cache_tiers = config.cache_tiers - gpu_quota = cache_tiers[0].quota + total_gpu_quota = cache_tiers[0].quota minimum_live_quota = self._minimum_live_gpu_quota() - if minimum_live_quota > gpu_quota: + if minimum_live_quota > total_gpu_quota: raise ValueError( "The V2 Mamba GPU cache quota is too small for live recurrent " - f"states and attention pages: got {gpu_quota} bytes, need at " - f"least {minimum_live_quota} bytes.") + f"states and attention pages: got {total_gpu_quota} bytes, " + f"need at least {minimum_live_quota} bytes.") + pool_gpu_quota = total_gpu_quota ssm_pool_at_live_floor = (config.initial_pool_ratio is None and self._ssm_pool_at_live_floor(kv_cache_config)) if ssm_pool_at_live_floor: # kda_state_tok is allocated beside the pools, one entry per SSM # slot: take those records' bytes out of the GPU quota before the # pools are sized. - token_state_slots = self._live_floor_token_state_slots( + draft_record_slots = self._live_floor_draft_record_slots( config.max_util_for_resume) - token_state_quota = (token_state_slots * - self._kda_token_state_bytes_per_slot()) - gpu_quota -= token_state_quota + draft_record_quota = (draft_record_slots * + self._kda_draft_record_bytes_per_slot()) + pool_gpu_quota = total_gpu_quota - draft_record_quota cache_tiers = [ - GpuCacheTierConfig(quota=gpu_quota), *cache_tiers[1:] + GpuCacheTierConfig(quota=pool_gpu_quota), *cache_tiers[1:] ] logger.info( "KV cache manager v2: the KDA draft records of " - f"{token_state_slots} SSM slots take " - f"{token_state_quota / (1 << 30)}GiB of the device quota; the " - f"pools get {gpu_quota / (1 << 30)}GiB") + f"{draft_record_slots} SSM slots take " + f"{draft_record_quota / (1 << 30)}GiB of the " + f"{total_gpu_quota / (1 << 30)}GiB device quota; the pools " + f"get {pool_gpu_quota / (1 << 30)}GiB") # _build_base_config already constructed every attention layer, # including dtype-specific scale and subclass-provided side buffers. # Preserve those configs and replace only the local Mamba layers. @@ -4186,7 +4192,7 @@ def _build_cache_config( if ssm_pool_at_live_floor: typical_capacity = max( typical_capacity, - self._quota_filling_request_capacity(gpu_quota)) + self._quota_filling_request_capacity(pool_gpu_quota)) request_descs = self._typical_request_descs(typical_capacity, kv_cache_config) typical_step = BatchDesc(request_descs * diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 99ef49600f4a..884653c0e6ff 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -386,21 +386,21 @@ def test_kimi_explicit_v2_manager_enables_kda_replay( assert kwargs["kda_replay_num_spec"] == spec_config.tokens_per_gen_step - 1 assert kwargs["conv_state_layout"] == "q_k_v" - assert "kda_token_states" not in kwargs + assert "kda_draft_records" not in kwargs @pytest.mark.parametrize("backbone_asks", [True, False]) -def test_kimi_explicit_v2_manager_kda_token_states_follow_the_model( +def test_kimi_explicit_v2_manager_kda_draft_records_follow_the_model( monkeypatch: pytest.MonkeyPatch, backbone_asks: bool, ) -> None: """The factory asks the V2 manager for per-token KDA verify states only when - the model's backbone sets ``kda_token_states``.""" + the model's backbone sets ``kda_draft_records``.""" monkeypatch.setattr( "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", lambda: True, ) - backbone = SimpleNamespace(kda_token_states=True) if backbone_asks else SimpleNamespace() + backbone = SimpleNamespace(kda_draft_records=True) if backbone_asks else SimpleNamespace() model_engine = SimpleNamespace( model=SimpleNamespace(model=backbone), _max_cuda_graph_batch_size=4, @@ -412,7 +412,7 @@ def test_kimi_explicit_v2_manager_kda_token_states_follow_the_model( ) assert kwargs["kda_replay_num_spec"] == 3 - assert kwargs.get("kda_token_states", False) is backbone_asks + assert kwargs.get("kda_draft_records", False) is backbone_asks def test_kimi_explicit_v2_manager_uses_qkv_convolution_layout( @@ -2388,7 +2388,7 @@ def _build_v2_hybrid_with_mamba_layer( mamba_ssm_cache_dtype=torch.float16, kda_replay_num_spec=None, enable_kv_pool_rebalance=False, - kda_token_states=False, + kda_draft_records=False, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2444,7 +2444,7 @@ def _build_v2_hybrid_with_mamba_layer( dtype=dtype, conv_state_layout=conv_state_layout, kda_replay_num_spec=kda_replay_num_spec, - kda_token_states=kda_token_states, + kda_draft_records=kda_draft_records, ) @@ -3668,7 +3668,7 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): assert layer_cache.intermediate_ssm is None assert layer_cache.intermediate_conv_window is None assert layer_cache.kda_state_tok is None - assert not mgr.keeps_kda_token_states + assert not mgr.keeps_kda_draft_records assert layer_cache.has_kda_replay_caches finally: mgr.shutdown() @@ -3676,8 +3676,8 @@ def test_v2_kda_replay_allocates_logical_slot_caches(): @skip_no_cuda @pytest.mark.parametrize("kda_replay_num_spec", [2, None]) -def test_v2_kda_token_states_replace_the_replay_caches(kda_replay_num_spec): - """On the KDA replay path, ``kda_token_states`` allocates the fp32 records of every draft of each slot (vn, +def test_v2_kda_draft_records_replace_the_replay_caches(kda_replay_num_spec): + """On the KDA replay path, ``kda_draft_records`` allocates the fp32 records of every draft of each slot (vn, beta * k, decay), per layer, in place of the per-draft replay caches; the conv caches and the replay path stay. Without the replay path it allocates nothing.""" mgr = _build_v2_hybrid_with_mamba_layer( @@ -3690,14 +3690,14 @@ def test_v2_kda_token_states_replace_the_replay_caches(kda_replay_num_spec): mamba_n_groups=6, mamba_ssm_cache_dtype=torch.float32, kda_replay_num_spec=kda_replay_num_spec, - kda_token_states=True, + kda_draft_records=True, ) try: if kda_replay_num_spec is None: assert getattr(mgr, "kda_state_tok", None) is None - assert not mgr.keeps_kda_token_states + assert not mgr.keeps_kda_draft_records return - assert mgr.keeps_kda_token_states + assert mgr.keeps_kda_draft_records assert mgr.kda_qkg_cache is None assert mgr.kda_v_cache is None assert mgr.kda_beta_cache is None @@ -3781,7 +3781,7 @@ def test_v2_kda_replay_refuses_kv_pool_rebalance(kda_replay): ) -def _v2_kda_token_state_manager(kda_token_states, enable_block_reuse): +def _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse): """A V2 hybrid manager stub for _build_cache_config: one KDA layer with the replay caches and one attention layer.""" mgr = object.__new__(MambaHybridCacheManagerV2) @@ -3811,7 +3811,7 @@ def _v2_kda_token_state_manager(kda_token_states, enable_block_reuse): # Layer 1's 256-byte pages of 32 tokens (_base_attention_layer_configs; layer 0 becomes the SSM layer). mgr._attention_cache_bytes_per_token = lambda: 8 mgr._use_kda_replay_update = True - mgr._kda_token_states = kda_token_states + mgr._kda_draft_records = kda_draft_records # The drafts' records count in each slot's state bytes (_mamba_state_bytes_per_slot): per draft, vn, beta * k # and the decay, fp32 [H, K] each (24 KiB here). mgr._kda_replay_num_spec = 2 @@ -3825,16 +3825,16 @@ def _v2_kda_token_state_manager(kda_token_states, enable_block_reuse): @pytest.mark.parametrize( - ("kda_token_states", "enable_block_reuse", "at_floor"), + ("kda_draft_records", "enable_block_reuse", "at_floor"), [(True, False, True), (False, False, False), (False, True, False)], ) -def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( - kda_token_states, enable_block_reuse, at_floor +def test_v2_kda_draft_records_keep_ssm_pool_at_live_floor( + kda_draft_records, enable_block_reuse, at_floor ): """The KDA drafts' records take memory per SSM slot beside the V2 pools. Without block reuse the SSM pool keeps only its live floor, the records' bytes come out of the quota, and attention gets the rest; without them (replay caches only, with or without block reuse) the typical step's ratio sizes the pool.""" - mgr = _v2_kda_token_state_manager(kda_token_states, enable_block_reuse) + mgr = _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse) base_config = KVCacheManagerConfig( tokens_per_block=32, cache_tiers=[GpuCacheTierConfig(quota=128 << 20)], @@ -3873,12 +3873,12 @@ def test_v2_kda_token_states_keep_ssm_pool_at_live_floor( @pytest.mark.parametrize("off_the_floor", ["block_reuse", "pool_ratio", "no_attention"]) -def test_v2_kda_token_states_refuse_a_pool_off_the_live_floor(off_the_floor): +def test_v2_kda_draft_records_refuse_a_pool_off_the_live_floor(off_the_floor): """kda_state_tok is allocated for every slot the SSM pool keeps, so with block reuse, a pool_ratio, or no attention layers on the rank (the SSM pool sized by its ratio, its slot count unknown before sizing) the build refuses the per-token states instead of allocating them outside the quota.""" - mgr = _v2_kda_token_state_manager( - kda_token_states=True, enable_block_reuse=off_the_floor == "block_reuse" + mgr = _v2_kda_draft_record_manager( + kda_draft_records=True, enable_block_reuse=off_the_floor == "block_reuse" ) pool_ratio = None if off_the_floor == "pool_ratio": @@ -3896,11 +3896,11 @@ def test_v2_kda_token_states_refuse_a_pool_off_the_live_floor(off_the_floor): ) -def test_v2_kda_token_states_reserve_quota_before_sizing(): +def test_v2_kda_draft_records_reserve_quota_before_sizing(): """At the live floor the drafts' records of every slot the SSM pool keeps come out of the GPU quota before the pools are sized, and the minimum quota and the quota for max_tokens count them all.""" - mgr = _v2_kda_token_state_manager(kda_token_states=True, enable_block_reuse=False) - plain = _v2_kda_token_state_manager(kda_token_states=False, enable_block_reuse=False) + mgr = _v2_kda_draft_record_manager(kda_draft_records=True, enable_block_reuse=False) + plain = _v2_kda_draft_record_manager(kda_draft_records=False, enable_block_reuse=False) assert mgr._max_resident_sequences() + mgr._num_reserved_dummy_slots == 3 # The pool keeps ceil(3 / max_util_for_resume) = 4 slots, each with the records of 2 drafts (3 x 8 x 256 fp32 # each). The quotas count the 4 slots' records and the 4th slot's states, which the quotas without the records @@ -3931,11 +3931,11 @@ def build(quota): assert build(128 << 20).cache_tiers[0].quota == (128 << 20) - reserved -def test_v2_kda_token_states_reserve_the_slots_the_runtime_keeps(): +def test_v2_kda_draft_records_reserve_the_slots_the_runtime_keeps(): """The quota checks and the build count the same slots when floor / max_util_for_resume is a whole number in exact arithmetic: the runtime holds the ratio as a C++ float, so 19 / 0.95 keeps 21 slots, not 20.""" - mgr = _v2_kda_token_state_manager(kda_token_states=True, enable_block_reuse=False) - plain = _v2_kda_token_state_manager(kda_token_states=False, enable_block_reuse=False) + mgr = _v2_kda_draft_record_manager(kda_draft_records=True, enable_block_reuse=False) + plain = _v2_kda_draft_record_manager(kda_draft_records=False, enable_block_reuse=False) for m in (mgr, plain): m.max_batch_size = 18 m.kv_cache_config = KvCacheConfig( @@ -4535,7 +4535,7 @@ def test_v2_kda_replay_relocates_live_slot_history(): @skip_no_cuda -def test_v2_kda_token_states_relocate_with_their_slot(): +def test_v2_kda_draft_records_relocate_with_their_slot(): """Without the per-draft replay caches, a slot's conv caches, drafts' records and draft count move with it.""" mgr = _build_v2_hybrid_with_mamba_layer( spec_config=MTPDecodingConfig(max_draft_len=2), @@ -4543,7 +4543,7 @@ def test_v2_kda_token_states_relocate_with_their_slot(): mamba_n_groups=4, mamba_ssm_cache_dtype=torch.float32, kda_replay_num_spec=2, - kda_token_states=True, + kda_draft_records=True, ) try: assert mgr.kda_qkg_cache is None @@ -4574,7 +4574,7 @@ def test_v2_kda_token_states_relocate_with_their_slot(): @skip_no_cuda -def test_v2_kda_token_states_count_in_the_per_slot_budget(): +def test_v2_kda_draft_records_count_in_the_per_slot_budget(): """The capacity math sees the drafts' records: per draft and layer, a slot costs vn, beta * k and the decay, fp32 [H, K] each.""" kwargs = dict( @@ -4592,10 +4592,12 @@ def test_v2_kda_token_states_count_in_the_per_slot_budget(): plain_bytes = plain._mamba_state_bytes_per_slot() finally: plain.shutdown() - with_states = _build_v2_hybrid_with_mamba_layer(kda_token_states=True, **kwargs) + with_states = _build_v2_hybrid_with_mamba_layer(kda_draft_records=True, **kwargs) try: num_heads, _, head_dim = with_states.ssm_state_shape per_token = 2 * 2 * 3 * num_heads * head_dim * 4 + assert with_states._managed_state_bytes_per_slot() == plain_bytes + assert with_states._kda_draft_record_bytes_per_slot() == per_token assert with_states._mamba_state_bytes_per_slot() == plain_bytes + per_token finally: with_states.shutdown() @@ -4753,7 +4755,7 @@ def dim_contiguous_conv_cache() -> torch.Tensor: assert mgr.prev_num_accepted_tokens[2] == 5 -def test_v2_kda_token_states_seed_disaggregated_generation_slots(): +def test_v2_kda_draft_records_seed_disaggregated_generation_slots(): """Without the per-draft replay caches, seeding a transferred request's slot sets its conv caches and clears its draft count; its drafts' records stay as they are, since with no pending drafts none of them is read.""" mgr = object.__new__(MambaHybridCacheManagerV2) From cdbba03f749f6012d8288d71c915c7af7f863ea6 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 13 Oct 2026 01:16:06 -0700 Subject: [PATCH 60/67] test: isolate MNNVL workspace lifecycle from CUDA capture state Signed-off-by: Vasanth Sabavat --- .../_torch/distributed/test_mnnvl_allreduce_lifecycle.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unittest/_torch/distributed/test_mnnvl_allreduce_lifecycle.py b/tests/unittest/_torch/distributed/test_mnnvl_allreduce_lifecycle.py index 72bd82ac0e7f..d78e1a5d6f73 100644 --- a/tests/unittest/_torch/distributed/test_mnnvl_allreduce_lifecycle.py +++ b/tests/unittest/_torch/distributed/test_mnnvl_allreduce_lifecycle.py @@ -124,6 +124,7 @@ def _patch_workspace_dependencies(monkeypatch: pytest.MonkeyPatch) -> _FakeWorld monkeypatch.setattr(MNNVLAllReduce, "_allreduce_mnnvl_workspace_locks", {}) monkeypatch.setattr(ops_module, "mpi_comm", lambda: world) monkeypatch.setattr(ops_module.torch, "device", lambda *args, **kwargs: cpu_device) + monkeypatch.setattr(ops_module.torch.cuda, "is_current_stream_capturing", lambda: False) monkeypatch.setattr(ops_module, "_initialize_allreduce_mnnvl_protocol", lambda workspace: None) return world From 02025194cbecac897ffea7d5739167f1995ffffa Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 13 Oct 2026 01:16:07 -0700 Subject: [PATCH 61/67] test: align hybrid fixtures and CUDA requirements with main Signed-off-by: Vasanth Sabavat --- .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index b31dca4820b1..e0dd25b97a3f 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -153,6 +153,8 @@ def _hybrid_model_config(): return SimpleNamespace( pretrained_config=config, sparse_attention_config=None, + mapping=None, + spec_config=None, get_num_mamba_layers=lambda: 1, ) @@ -3771,6 +3773,7 @@ def test_v2_kda_replay_validates_configuration( ) +@skip_no_cuda @pytest.mark.parametrize("kda_replay", [True, False], ids=["kda_replay", "no_replay"]) def test_v2_kda_replay_refuses_kv_pool_rebalance(kda_replay): """The KDA replay caches and prev_num_accepted_tokens hold one entry per SSM slot the pool keeps at construction, @@ -3833,6 +3836,7 @@ def _v2_kda_draft_record_manager(kda_draft_records, enable_block_reuse): return mgr +@skip_no_cuda @pytest.mark.parametrize( ("kda_draft_records", "enable_block_reuse", "at_floor"), [(True, False, True), (False, False, False), (False, True, False)], From 72bb8ca0b8ad4f5086f6fe4e1162afc6de533f61 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Tue, 13 Oct 2026 01:31:11 -0700 Subject: [PATCH 62/67] test: remove obsolete MNNVL workspace dictionary assertion Signed-off-by: Vasanth Sabavat --- .../_torch/distributed/test_mnnvl_allreduce_lifecycle.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unittest/_torch/distributed/test_mnnvl_allreduce_lifecycle.py b/tests/unittest/_torch/distributed/test_mnnvl_allreduce_lifecycle.py index d78e1a5d6f73..bfb9833d33ce 100644 --- a/tests/unittest/_torch/distributed/test_mnnvl_allreduce_lifecycle.py +++ b/tests/unittest/_torch/distributed/test_mnnvl_allreduce_lifecycle.py @@ -149,12 +149,11 @@ def fail_once(*args: object, **kwargs: object) -> _FakeMcastBuffer: assert world.split_count == 1 assert MNNVLAllReduce.allreduce_mnnvl_pending_comms[mapping] is world.comm - workspace = get_or_scale_allreduce_mnnvl_workspace(mapping, torch.float32) + get_or_scale_allreduce_mnnvl_workspace(mapping, torch.float32) assert construction_count == 2 assert world.split_count == 1 assert MNNVLAllReduce.allreduce_mnnvl_pending_comms == {} - assert workspace["mpi_comm"] is world.comm def test_concurrent_workspace_construction_is_serialized(monkeypatch: pytest.MonkeyPatch) -> None: From 1f94ed5ca2e232b1c2545f6eb7f71a9c7742661f Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 14 Oct 2026 09:32:23 -0700 Subject: [PATCH 63/67] Own cached host-to-device copies and share metadata buffer snapshots Signed-off-by: Vasanth Sabavat --- .../torch_compile_and_piecewise_cuda_graph.md | 12 + tensorrt_llm/_device_copy.py | 45 ++++ .../_torch/modules/mamba/mamba2_metadata.py | 17 +- .../engine/runners/decoder/runner.py | 27 +-- .../kv_cache/mamba_cache_manager.py | 65 +++--- tensorrt_llm/_torch/speculative/dflash.py | 36 ++- .../_torch/speculative/spec_sampler_base.py | 18 +- tensorrt_llm/_utils.py | 53 ----- .../kv_cache/test_mamba_cache_manager.py | 21 +- .../test_copy_to_device_if_changed.py | 143 ------------ .../executor/test_device_copy_buffer.py | 205 ++++++++++++++++++ .../kimi_kda/test_kimi_k3_mamba_metadata.py | 7 +- .../kimi_kda/test_kimi_kda_verify_parity.py | 7 +- .../modules/mamba/test_mamba2_metadata.py | 51 ++++- .../hw_agnostic/test_dflash_copy_skip.py | 28 +-- .../hw_agnostic/test_dflash_dummy_slot.py | 6 +- .../hw_agnostic/test_ragged_spec_sampler.py | 2 +- tests/unittest/conftest.py | 8 +- 18 files changed, 427 insertions(+), 324 deletions(-) create mode 100644 tensorrt_llm/_device_copy.py delete mode 100644 tests/unittest/_torch/executor/test_copy_to_device_if_changed.py create mode 100644 tests/unittest/_torch/executor/test_device_copy_buffer.py diff --git a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md index 32ebadc25bea..fbca75885083 100644 --- a/docs/source/features/torch_compile_and_piecewise_cuda_graph.md +++ b/docs/source/features/torch_compile_and_piecewise_cuda_graph.md @@ -139,6 +139,18 @@ the Multimodal Model Family. ## Development Guide +Executor host-upload buffers must keep stable device addresses across capture and replay. Their internal +`DeviceCopyBuffer` owns the tensor and last-uploaded host snapshot; every cached writer must use the same +owner's `update(host)` method. A Mamba metadata CUDA alias shares its cache manager's owner, rather than +creating another cache over the same storage. Returned tensors are borrowed views, not immutable tensors; +writing them directly bypasses the cached-upload contract. + +Changed updates use independent host staging, pinned according to the owner's configuration, and write only +the supplied prefix. Shorter updates preserve the device tail and its cached values; unchanged prefixes +enqueue no upload. The C++ Mamba manager uses `force=True` to retain its unconditional-upload behavior while +keeping aliases' shared snapshot coherent. `TLLM_DEBUG_MODE=1` verifies skipped values outside CUDA graph +capture; capture does not add a device synchronization for this check. + ### Background Knowledge Currently, TRT-LLM mainly relies on torch.compile **fullgraph** mode to enable Piecewise CUDA Graph feature, which means all the operations in the model must be recognized by torch.compile. diff --git a/tensorrt_llm/_device_copy.py b/tensorrt_llm/_device_copy.py new file mode 100644 index 000000000000..416911bf493e --- /dev/null +++ b/tensorrt_llm/_device_copy.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import os + +import torch + +_verify_skipped_device_copies = os.environ.get("TLLM_DEBUG_MODE", "")[0:1] == "1" + + +class DeviceCopyBuffer: + def __init__(self, tensor: torch.Tensor, *, pin_memory: bool = False): + self._tensor = tensor + self._pin_memory = pin_memory + self._last_host_values = None + + @property + def tensor(self) -> torch.Tensor: + return self._tensor + + def update(self, host: torch.Tensor, *, force: bool = False) -> None: + flat = host.reshape(-1) + n = flat.numel() + last = self._last_host_values + if not force and last is not None and last.numel() >= n and torch.equal(last[:n], flat): + if _verify_skipped_device_copies: + if self._tensor.is_cuda: + if torch.cuda.is_current_stream_capturing(): + return + torch.cuda.synchronize(self._tensor.device) + held = self._tensor.view(-1)[:n].cpu() + assert torch.equal(held, flat.to(held.dtype)), ( + f"DeviceCopyBuffer skipped a copy, but the buffer holds {held.tolist()} " + f"instead of {flat.tolist()}: something else wrote it" + ) + return + staging = torch.empty_like(flat, device="cpu", pin_memory=self._pin_memory) + staging.copy_(flat) + self._tensor.view(-1)[:n].copy_(staging, non_blocking=True) + if last is None or last.numel() <= n: + self._last_host_values = flat.clone() + else: + last = last.clone() + last[:n] = flat + self._last_host_values = last diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 0a2478801bce..0fdf51bdafcb 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -21,10 +21,11 @@ import triton import triton.language as tl +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import \ CUDA_GRAPH_DUMMY_REQUEST_ID -from tensorrt_llm._utils import copy_to_device_if_changed, prefer_pinned +from tensorrt_llm._utils import prefer_pinned REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0 REPLAY_WORK_CACHE_SLOT = 1 @@ -371,12 +372,11 @@ def __init__( self.state_indices_cpu = torch.zeros(max_batch_size, dtype=torch.int32, pin_memory=prefer_pinned()) - # Only copy_to_device_if_changed writes this buffer: it skips values - # it already copied here, so another write would leave stale values. - # prepare() aliases a cache manager's own device buffer instead. self.state_indices = torch.zeros(max_batch_size, dtype=torch.int32, device="cuda") + self._state_indices_buffer = DeviceCopyBuffer( + self.state_indices, pin_memory=prefer_pinned()) # Stable data_ptr() of the CUDA tensor we alias (if any) — used to # detect cache-manager buffer reallocation that would silently break # CUDA graph replays. @@ -493,11 +493,12 @@ def prepare(self, attn_metadata: AttentionMetadata): "are used; got a different address than the first " "call.") self.state_indices = indices + self._state_indices_buffer = kv_cache_manager._state_indices_buffer elif isinstance(indices, torch.Tensor): # CPU tensor → bulk H2D self.state_indices_cpu[:batch_size].copy_(indices[:batch_size]) - copy_to_device_if_changed(self.state_indices, - self.state_indices_cpu[:batch_size]) + self._state_indices_buffer.update( + self.state_indices_cpu[:batch_size]) else: # indices is a Python sequence (e.g. List[int]); data # already lives on host, CPU staging is fine. One bulk @@ -508,8 +509,8 @@ def prepare(self, attn_metadata: AttentionMetadata): self.state_indices_cpu[:batch_size].copy_( torch.as_tensor(indices, dtype=self.state_indices_cpu.dtype)) - copy_to_device_if_changed(self.state_indices, - self.state_indices_cpu[:batch_size]) + self._state_indices_buffer.update( + self.state_indices_cpu[:batch_size]) self._prepare_replay_work_items(kv_cache_manager, batch_size, num_contexts) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 15ac46c72ebd..616b8dd33e0e 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -15,6 +15,7 @@ import torch +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.autotuner import AutoTuner, autotune @@ -79,13 +80,7 @@ set_per_request_prefill_cuda_graph_flag, with_model_extra_attrs, ) -from tensorrt_llm._utils import ( - copy_to_device_if_changed, - global_mpi_rank, - maybe_pin_memory, - nvtx_range, - prefer_pinned, -) +from tensorrt_llm._utils import global_mpi_rank, maybe_pin_memory, nvtx_range, prefer_pinned from tensorrt_llm.inputs.multimodal import ( MultimodalParams, MultimodalRuntimeData, @@ -455,11 +450,9 @@ def _allocate_decoder_buffers(self) -> None: self.draft_tokens_cuda = torch.empty( (max_num_draft_tokens,), dtype=torch.int, device="cuda" ) - # Only copy_to_device_if_changed writes this buffer: it skips - # values it already copied here, so another write would leave - # stale values. - self.gather_ids_cuda = torch.empty( - (self._config.max_num_tokens,), dtype=torch.int, device="cuda" + self._gather_ids_buffer = DeviceCopyBuffer( + torch.empty((self._config.max_num_tokens,), dtype=torch.int, device="cuda"), + pin_memory=prefer_pinned(), ) self.previous_pos_indices_cuda = torch.empty( (self._config.max_num_tokens,), dtype=torch.int, device="cuda" @@ -4221,9 +4214,7 @@ def previous_seq_slots_device(): final_position_ids = self.position_ids_cuda[:total_num_tokens].unsqueeze(0) if enable_spec_decode: - copy_to_device_if_changed( - self.gather_ids_cuda, torch.tensor(gather_ids, dtype=torch.int) - ) + self._gather_ids_buffer.update(torch.tensor(gather_ids, dtype=torch.int)) if self.mapping.has_cp_helix(): # A non-None owned-count list is what arms @@ -4461,7 +4452,7 @@ def previous_seq_slots_device(): total_draft_lens = sum(draft_lens) spec_metadata.draft_tokens = self.draft_tokens_cuda[:total_draft_lens] spec_metadata.request_ids = request_ids - spec_metadata.gather_ids = self.gather_ids_cuda[: len(gather_ids)] + spec_metadata.gather_ids = self._gather_ids_buffer.tensor[: len(gather_ids)] # num_generations / num_tokens / seq_lens are set above, before the # attention-DP allgather that must agree with prepare(). spec_metadata.host_position_ids = host_position_ids @@ -4561,7 +4552,9 @@ def previous_seq_slots_device(): "use_mrope": _use_mrope, } - gather_ids_device = self.gather_ids_cuda[: len(gather_ids)] if enable_spec_decode else None + gather_ids_device = ( + self._gather_ids_buffer.tensor[: len(gather_ids)] if enable_spec_decode else None + ) return inputs, gather_ids_device, runtime_draft_len @nvtx_range("_prepare_inputs") diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py index 5a91e35e52c1..8ecddb4b29d6 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -32,6 +32,7 @@ from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig from tensorrt_llm.sampling_params import SamplingParams +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.disaggregation.resource.page import (MapperKind, RoleLayout) from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import ( @@ -45,8 +46,8 @@ PoolConfiguration, get_pp_layers) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm._utils import (TensorWrapper, convert_to_torch_tensor, - copy_to_device_if_changed, nvtx_range, - prefer_pinned, torch_dtype_to_binding) + nvtx_range, prefer_pinned, + torch_dtype_to_binding) from tensorrt_llm.bindings.internal.batch_manager import ( LinearAttentionMetadata, LinearCacheType) from tensorrt_llm.llmapi.llm_args import KvCacheConfig @@ -1396,6 +1397,16 @@ class MambaHybridCacheManager(BaseResourceManager, BaseMambaCacheManager): _supports_additional_snapshot_offsets = False + @property + def cuda_state_indices(self) -> torch.Tensor: + return self._state_indices_buffer.tensor + + @property + def _dummy_request_mask(self) -> Optional[torch.Tensor]: + if self._dummy_request_mask_buffer is None: + return None + return self._dummy_request_mask_buffer.tensor + def _setup_mtp_intermediate_states(self, spec_config, max_batch_size: int) -> None: self.spec_config = spec_config @@ -1436,7 +1447,7 @@ def _allocate_pool_replay_buffers( self.prev_num_accepted_tokens = None self.cache_buf_idx = None self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None + self._dummy_request_mask_buffer = None self._dummy_request_mask_host = None self.old_x = None self.old_B = None @@ -1497,8 +1508,7 @@ def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: if n > 0: self._dummy_request_mask_host[:n].copy_( torch.tensor(is_dummy, dtype=torch.bool)) - copy_to_device_if_changed(self._dummy_request_mask, - self._dummy_request_mask_host) + self._dummy_request_mask_buffer.update(self._dummy_request_mask_host) def _reset_context_mamba_slots(self, num_contexts: int) -> None: if num_contexts == 0: @@ -2528,9 +2538,10 @@ def __init__( self.recurrent_states_pool_index = self.kv_cache_pool_mapping[ self.layer_offsets[self.mamba_pp_layers[0]]][0] - self.cuda_state_indices = torch.zeros([self.max_batch_size], - dtype=torch.int32, - device="cuda") + self._state_indices_buffer = DeviceCopyBuffer( + torch.zeros([self.max_batch_size], dtype=torch.int32, + device="cuda"), + pin_memory=prefer_pinned()) self._host_state_indices = torch.zeros([self.max_batch_size], dtype=torch.int32, pin_memory=prefer_pinned()) @@ -2639,7 +2650,7 @@ def shutdown(self): self.prev_num_accepted_tokens = None self.cache_buf_idx = None self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None + self._dummy_request_mask_buffer = None self._dummy_request_mask_host = None self.old_x = None self.old_B = None @@ -2921,10 +2932,7 @@ def _setup_state_indices(self, requests=None) -> None: ) self._host_state_indices[:n] = values - idx_staging = torch.empty_like(self._host_state_indices, - pin_memory=prefer_pinned()) - idx_staging.copy_(self._host_state_indices) - self.cuda_state_indices.copy_(idx_staging, non_blocking=True) + self._state_indices_buffer.update(self._host_state_indices, force=True) is_dummy = [req.is_dummy for req in requests] self._refresh_dummy_request_mask(is_dummy) @@ -2989,11 +2997,9 @@ def _setup_replay_buffers(self, spec_config) -> None: device): return - # Only copy_to_device_if_changed writes this buffer: it skips values - # it already copied here, so another write would leave stale values. - self._dummy_request_mask = torch.zeros(self.max_batch_size, - dtype=torch.bool, - device=device) + self._dummy_request_mask_buffer = DeviceCopyBuffer( + torch.zeros(self.max_batch_size, dtype=torch.bool, device=device), + pin_memory=prefer_pinned()) self._dummy_request_mask_host = torch.zeros( self.max_batch_size, dtype=torch.bool, @@ -3291,11 +3297,11 @@ def __init__( state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) - # Only copy_to_device_if_changed writes this buffer: it skips values - # it already copied here, so another write would leave stale values. - self.cuda_state_indices = torch.zeros([state_index_capacity], - dtype=torch.int32, - device="cuda") + self._state_indices_buffer = DeviceCopyBuffer( + torch.zeros([state_index_capacity], + dtype=torch.int32, + device="cuda"), + pin_memory=prefer_pinned()) self._host_state_indices = torch.zeros([state_index_capacity], dtype=torch.int32, pin_memory=prefer_pinned()) @@ -4359,11 +4365,9 @@ def _setup_replay_buffers(self, spec_config) -> None: return mask_capacity = self._host_state_indices.shape[0] - # Only copy_to_device_if_changed writes this buffer: it skips values - # it already copied here, so another write would leave stale values. - self._dummy_request_mask = torch.zeros(mask_capacity, - dtype=torch.bool, - device=device) + self._dummy_request_mask_buffer = DeviceCopyBuffer( + torch.zeros(mask_capacity, dtype=torch.bool, device=device), + pin_memory=prefer_pinned()) self._dummy_request_mask_host = torch.zeros( mask_capacity, dtype=torch.bool, @@ -4705,8 +4709,7 @@ def _setup_state_indices(self, f"request {req.py_request_id}") self._host_state_indices[i] = base_index - copy_to_device_if_changed(self.cuda_state_indices, - self._host_state_indices) + self._state_indices_buffer.update(self._host_state_indices) is_dummy = [req.is_dummy for req in requests] self._refresh_dummy_request_mask(is_dummy) state_values = self._host_state_indices[:n].tolist() @@ -5117,7 +5120,7 @@ def shutdown(self): self.prev_num_accepted_tokens = None self.cache_buf_idx = None self.mamba_ssm_rand_seed = None - self._dummy_request_mask = None + self._dummy_request_mask_buffer = None self._dummy_request_mask_host = None self.old_x = None self.old_B = None diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 6e02d7530a54..d23cf26a6a46 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -21,8 +21,9 @@ import torch from torch import nn +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.custom_ops import inplace_slice_copy -from tensorrt_llm._utils import copy_to_device_if_changed, prefer_pinned +from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -518,13 +519,11 @@ class DFlashSpecMetadata(SpecMetadata): captured_hidden_states: Optional[torch.Tensor] = None def __post_init__(self): - # Only copy_to_device_if_changed writes this buffer: it skips values it already copied here, so another write - # would leave stale values. - self.batch_indices_cuda = torch.empty( - [self.max_num_requests], - dtype=torch.int, - device="cuda", + self._batch_indices_buffer = DeviceCopyBuffer( + torch.empty([self.max_num_requests], dtype=torch.int, device="cuda"), + pin_memory=prefer_pinned(), ) + self.batch_indices_cuda = self._batch_indices_buffer.tensor self.is_spec_dec_tree = False self.is_spec_dec_dynamic_tree = False @@ -565,7 +564,7 @@ def prepare(self): assert self.request_ids is not None num_seqs = len(self.request_ids) - copy_to_device_if_changed(self.batch_indices_cuda, torch.arange(num_seqs, dtype=torch.int)) + self._batch_indices_buffer.update(torch.arange(num_seqs, dtype=torch.int)) # Update slot mapping for DFlash context buffers worker = getattr(self, "_dflash_worker", None) @@ -604,7 +603,7 @@ def prepare(self): [worker._req_to_slot.get(rid, worker._dummy_slot) for rid in self.request_ids], dtype=torch.long, ) - copy_to_device_if_changed(worker._batch_to_slot, mapping) + worker._batch_to_slot_buffer.update(mapping) def is_layer_capture(self, layer_id: int) -> bool: return layer_id in self._capture_layer_set @@ -706,7 +705,7 @@ def __init__( self._kv_rewind_amount = None self._kv_rewind_nc = None self._kv_rewind_bs = None - self._batch_to_slot = None + self._batch_to_slot_buffer = None self._max_ctx = 0 self._ctx_k_buf = None # [max_batch+1, L, max_ctx+block, nkv, hd] self._ctx_v_buf = None @@ -1076,9 +1075,10 @@ def _lazy_init_ctx_buffers( self._ctx_len = torch.zeros(num_slots, dtype=torch.long, device="cuda") self._ctx_len_host = [0] * num_slots - # Only copy_to_device_if_changed writes this buffer: it skips values it already copied here, so another write - # would leave stale values. - self._batch_to_slot = torch.zeros(max_batch, dtype=torch.long, device="cuda") + self._batch_to_slot_buffer = DeviceCopyBuffer( + torch.zeros(max_batch, dtype=torch.long, device="cuda"), + pin_memory=prefer_pinned(), + ) self._free_slots = deque(range(max_batch)) self._req_to_slot = {} @@ -1382,14 +1382,8 @@ def _advance_ctx_len(self, slots: torch.Tensor, num_accepted: torch.Tensor) -> N self._ctx_len[self._dummy_slot].zero_() def _rebuild_batch_to_slot(self, request_ids: List[int]) -> None: - """Point _batch_to_slot at each request's slot after prefill assigned new ones. - - Like DFlashSpecMetadata.prepare, this writes through copy_to_device_if_changed, - which keeps the values it last copied on the buffer: every write of the buffer - must go through it, or a later unchanged batch skips its copy over other values. - """ mapping = [self._req_to_slot.get(rid, self._dummy_slot) for rid in request_ids] - copy_to_device_if_changed(self._batch_to_slot, torch.tensor(mapping, dtype=torch.long)) + self._batch_to_slot_buffer.update(torch.tensor(mapping, dtype=torch.long)) def _write_ctx_len(self, updates: dict[int, int]) -> None: """Apply a slot-to-length mapping to _ctx_len in one async scatter.""" @@ -2012,7 +2006,7 @@ def prepare_1st_drafter_inputs( query_tokens_per_req = block_size # Get slots for gen requests from pre-computed mapping - slots = self._batch_to_slot[num_contexts : num_contexts + num_gens] + slots = self._batch_to_slot_buffer.tensor[num_contexts : num_contexts + num_gens] gen_rows_out = torch.arange( num_contexts, num_contexts + num_gens, dtype=torch.long, device="cuda" ) diff --git a/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 2dff89858628..dc0c69f7d909 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -27,7 +27,8 @@ import torch -from ..._utils import copy_to_device_if_changed +from ..._device_copy import DeviceCopyBuffer +from ..._utils import prefer_pinned from ...sampling_params import SamplingParams from ..cute_dsl_kernels.spec_step_copies import op as spec_step_copies from ..pyexecutor.llm_request import LlmRequest, LlmRequestState, get_draft_token_length @@ -306,15 +307,14 @@ def __init__( # tree (K=6, T=60), MTP dynamic tree, PARD (T=2K-1) and the linear # modes -- none exceed it. self.max_accepted_path_len = args.max_draft_len + 1 - # Where the step-copy kernels run, one kernel moves a step's outputs - # into the stores below, reading each row's slot from _slot_table. self._store_scatter: Optional[spec_step_copies.SlotScatter] = None - self._slot_table: Optional[torch.Tensor] = None + self._slot_table_buffer: Optional[DeviceCopyBuffer] = None if spec_step_copies.is_supported(): self._store_scatter = spec_step_copies.SlotScatter() - # Only copy_to_device_if_changed writes this buffer: it skips values it already copied here, so another - # write would leave stale values. - self._slot_table = torch.zeros((seq_slots,), dtype=torch.int32, device="cuda") + self._slot_table_buffer = DeviceCopyBuffer( + torch.zeros((seq_slots,), dtype=torch.int32, device="cuda"), + pin_memory=prefer_pinned(), + ) # Recorded after the last step's host copies of the stores, which read # them on the D2H side stream. self._store_copies_done: Optional[torch.cuda.Event] = None @@ -546,12 +546,12 @@ def sample_async( scattered = False if self._store_scatter is not None: # A step whose batch is unchanged leaves the slot table as it is. - copy_to_device_if_changed(self._slot_table, torch.tensor(slots, dtype=torch.int32)) + self._slot_table_buffer.update(torch.tensor(slots, dtype=torch.int32)) scattered = self._store_scatter.scatter( outputs, num_skip, len(slots), - self._slot_table, + self._slot_table_buffer.tensor, self.store.new_tokens, self.store.next_new_tokens, self.store.new_tokens_lens, diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index d2f76483b441..c875715c4501 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -1336,59 +1336,6 @@ def maybe_pin_memory(tensor: torch.Tensor) -> torch.Tensor: return tensor -# Set by TLLM_DEBUG_MODE=1, and for every unit test by a fixture: a copy that -# copy_to_device_if_changed skips first checks that the buffer holds the values. -_verify_skipped_device_copies = os.environ.get("TLLM_DEBUG_MODE", - "")[0:1] == "1" - - -def _assert_device_holds(dst: torch.Tensor, flat: torch.Tensor) -> None: - """Assert that the leading elements of ``dst`` hold ``flat`` once all queued - work has finished.""" - if dst.is_cuda: - if torch.cuda.is_current_stream_capturing(): - return - torch.cuda.synchronize(dst.device) - held = dst.view(-1)[:flat.numel()].cpu() - assert torch.equal(held, flat.to(held.dtype)), ( - "copy_to_device_if_changed skipped a copy, but the buffer holds " - f"{held.tolist()} instead of {flat.tolist()}: something else wrote it") - - -def copy_to_device_if_changed(dst: torch.Tensor, host: torch.Tensor) -> None: - """Copy ``host`` into the leading elements of the device buffer ``dst`` - unless they already hold it. - - ``dst`` is the whole buffer (not a slice), and the values it last received - are kept on it, so a reallocated buffer starts over. Only for per-step - metadata buffers that nothing else writes: a decode step whose batch is - unchanged then enqueues no copy. "Nothing else" includes other tensor - objects over the same memory, such as the views a shared buffer pool hands - to each metadata object. The copy itself goes through fresh pinned staging, - so ``host`` may be reused right away. - - With ``TLLM_DEBUG_MODE=1`` a skipped copy reads ``dst`` back (a device - synchronization) and asserts that it holds ``host``, which catches a write - by anything else. - """ - flat = host.reshape(-1) - n = flat.numel() - last = getattr(dst, "_last_host_values", None) - if last is not None and last.numel() >= n and torch.equal(last[:n], flat): - if _verify_skipped_device_copies: - _assert_device_holds(dst, flat) - return - staging = torch.empty_like(flat, device="cpu", pin_memory=prefer_pinned()) - staging.copy_(flat) - dst.view(-1)[:n].copy_(staging, non_blocking=True) - if last is None or last.numel() <= n: - dst._last_host_values = flat.clone() - else: - last = last.clone() - last[:n] = flat - dst._last_host_values = last - - def async_tensor_h2d(data, dtype: torch.dtype, device: Union[str, torch.device]) -> torch.Tensor: """Build a CPU tensor from `data` and ship it to `device` with a diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index e0dd25b97a3f..324070ce6fed 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -10,6 +10,7 @@ import pytest import torch +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.disaggregation.resource.kv_extractor import build_page_table_from_manager from tensorrt_llm._torch.disaggregation.resource.page import ( MAMBA_CONV_ROLE, @@ -1442,7 +1443,7 @@ def test_cpp_hybrid_replay_bookkeeping_is_fused_into_conv_promotion( mgr = object.__new__(CppMambaHybridCacheManager) mgr.local_num_mamba_layers = 1 mgr._use_replay_state_update = True - mgr._dummy_request_mask = torch.tensor([False, True]) + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.tensor([False, True])) mgr.prev_num_accepted_tokens = torch.tensor([13, 7], dtype=torch.int32) mgr.cache_buf_idx = torch.tensor([1, 0], dtype=torch.int32) mgr.replay_step_width = 4 @@ -2104,11 +2105,11 @@ def test_cpp_hybrid_state_indices_skip_context_placeholders( mgr.recurrent_states_pool_index = 0 mgr.blocks_per_window = {LinearCacheType.RECURRENT_STATES.value: (1264, 0)} mgr._host_state_indices = torch.zeros(1, dtype=torch.int32) - mgr.cuda_state_indices = torch.zeros(1, dtype=torch.int32) + mgr._state_indices_buffer = DeviceCopyBuffer(torch.zeros(1, dtype=torch.int32)) mgr._row_indices = torch.arange(1, dtype=torch.long) mgr._request_id_to_state_index = {} mgr._request_id_to_is_dummy = {} - mgr._dummy_request_mask = None + mgr._dummy_request_mask_buffer = None mgr._setup_state_indices() @@ -3990,7 +3991,7 @@ def test_v2_kda_replay_records_acceptance_and_skips_dummy_rows(): mgr._use_kda_replay_update = True mgr._use_replay_state_update = False mgr.prev_num_accepted_tokens = torch.tensor([7, 8, 9], dtype=torch.int32) - mgr._dummy_request_mask = torch.tensor([False, False, True]) + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.tensor([False, False, True])) mgr.update_mamba_states( SimpleNamespace(num_seqs=3, num_contexts=1), @@ -4005,7 +4006,7 @@ def test_v2_kda_replay_resets_context_slots(): mgr = object.__new__(MambaHybridCacheManagerV2) mgr._use_kda_replay_update = True mgr._use_replay_state_update = False - mgr.cuda_state_indices = torch.tensor([2, 1], dtype=torch.int32) + mgr._state_indices_buffer = DeviceCopyBuffer(torch.tensor([2, 1], dtype=torch.int32)) mgr.prev_num_accepted_tokens = torch.tensor([5, 6, 7], dtype=torch.int32) mgr.mamba_ssm_rand_seed = None @@ -4088,8 +4089,8 @@ def _v2_kda_replay_stub(num_slots=4): mgr._kda_replay_acceptance_pending = set() mgr._kda_replay_unseeded = set() mgr._host_state_indices = torch.zeros(num_slots, dtype=torch.int32) - mgr.cuda_state_indices = torch.zeros(num_slots, dtype=torch.int32) - mgr._dummy_request_mask = None + mgr._state_indices_buffer = DeviceCopyBuffer(torch.zeros(num_slots, dtype=torch.int32)) + mgr._dummy_request_mask_buffer = None mgr.mamba_ssm_rand_seed = None mgr.kv_cache_map = {} mgr.ssm_layer_group_id = 0 @@ -4677,7 +4678,7 @@ def test_v2_state_index_setup_skips_unchanged_uploads(monkeypatch): keep values the host never sent. A changed step uploads both again.""" # The test writes the buffers behind _setup_state_indices to see the skip, which the check of skipped copies # reports. - monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", False) + monkeypatch.setattr("tensorrt_llm._device_copy._verify_skipped_device_copies", False) class StateCache: def __init__(self, slot: int) -> None: @@ -5054,7 +5055,7 @@ def test_v2_gdn_replay_commits_before_advancing_bookkeeping(monkeypatch): mgr.prev_num_accepted_tokens = torch.zeros(batch_size, dtype=torch.int32) mgr.cache_buf_idx = torch.zeros(batch_size, dtype=torch.int32) mgr.intermediate_state_indices = torch.arange(batch_size, dtype=torch.int32) - mgr._dummy_request_mask = torch.zeros(batch_size, dtype=torch.bool) + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(batch_size, dtype=torch.bool)) mgr.all_conv_states = [torch.empty(0)] mgr._stacked_conv_states = None mgr.intermediate_conv_states = torch.empty(0) @@ -5094,7 +5095,7 @@ def test_v2_hybrid_replay_update_skips_dummy_and_padding_rows(monkeypatch): 102: True, 103: False, } - mgr._dummy_request_mask = torch.zeros(4, dtype=torch.bool) + mgr._dummy_request_mask_buffer = DeviceCopyBuffer(torch.zeros(4, dtype=torch.bool)) mgr._dummy_request_mask_host = torch.zeros(4, dtype=torch.bool) mgr._use_replay_state_update = True mgr.replay_step_width = 5 diff --git a/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py b/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py deleted file mode 100644 index e7772ac2d217..000000000000 --- a/tests/unittest/_torch/executor/test_copy_to_device_if_changed.py +++ /dev/null @@ -1,143 +0,0 @@ -# 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. -"""``copy_to_device_if_changed``: which calls copy and what the buffer then holds. - -The skip decision does not depend on the device, so the buffer here is a host tensor: a skipped copy shows as a -sentinel written behind the function's back that survives the call. Such a write is what the debug check of a skipped -copy reports, so these tests turn the check off, except the ones that test it. -""" - -import pytest -import torch - -from tensorrt_llm._utils import copy_to_device_if_changed - -pytestmark = pytest.mark.cpu_only - -SENTINEL = -7 - - -@pytest.fixture(autouse=True) -def _no_skipped_copy_check(monkeypatch) -> None: - monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", False) - - -def _host(values: list[int]) -> torch.Tensor: - return torch.tensor(values, dtype=torch.int32) - - -def _tamper(dst: torch.Tensor) -> None: - """Overwrite the buffer without the function knowing, so only a real copy restores it.""" - dst.fill_(SENTINEL) - - -def test_first_call_copies() -> None: - dst = torch.zeros(8, dtype=torch.int32) - copy_to_device_if_changed(dst, _host([3, 1, 2])) - assert dst.tolist() == [3, 1, 2, 0, 0, 0, 0, 0] - - -def test_unchanged_values_skip_the_copy() -> None: - dst = torch.zeros(8, dtype=torch.int32) - copy_to_device_if_changed(dst, _host([3, 1, 2])) - _tamper(dst) - copy_to_device_if_changed(dst, _host([3, 1, 2])) - assert dst.eq(SENTINEL).all() - - -def test_changed_values_copy() -> None: - dst = torch.zeros(8, dtype=torch.int32) - copy_to_device_if_changed(dst, _host([3, 1, 2])) - _tamper(dst) - copy_to_device_if_changed(dst, _host([3, 1, 4])) - assert dst[:3].tolist() == [3, 1, 4] - assert dst[3:].eq(SENTINEL).all(), "only the leading elements are written" - - -def test_prefix_of_the_last_values_skips_the_copy() -> None: - dst = torch.zeros(8, dtype=torch.int32) - copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) - _tamper(dst) - copy_to_device_if_changed(dst, _host([5, 6])) - assert dst.eq(SENTINEL).all() - - -def test_longer_values_copy_and_extend_what_is_kept() -> None: - dst = torch.zeros(8, dtype=torch.int32) - copy_to_device_if_changed(dst, _host([5, 6])) - copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) - assert dst[:4].tolist() == [5, 6, 7, 8] - _tamper(dst) - copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) - assert dst.eq(SENTINEL).all() - - -def test_shorter_changed_values_keep_the_tail() -> None: - # A shorter copy rewrites only its leading elements; the buffer's tail still holds the earlier values, and what - # the function keeps says so. - dst = torch.zeros(8, dtype=torch.int32) - copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) - copy_to_device_if_changed(dst, _host([9, 9])) - assert dst[:4].tolist() == [9, 9, 7, 8] - _tamper(dst) - copy_to_device_if_changed(dst, _host([9, 9, 7, 8])) - assert dst.eq(SENTINEL).all() - copy_to_device_if_changed(dst, _host([9, 9, 7, 1])) - assert dst[:4].tolist() == [9, 9, 7, 1] - - -def test_host_values_may_be_reused_right_away() -> None: - dst = torch.zeros(4, dtype=torch.int32) - host = _host([1, 2, 3]) - copy_to_device_if_changed(dst, host) - host.fill_(0) - assert dst[:3].tolist() == [1, 2, 3] - copy_to_device_if_changed(dst, _host([1, 2, 3])) - assert dst[:3].tolist() == [1, 2, 3] - - -def test_a_new_buffer_starts_over() -> None: - first = torch.zeros(4, dtype=torch.int32) - copy_to_device_if_changed(first, _host([1, 2])) - second = torch.full((4,), SENTINEL, dtype=torch.int32) - copy_to_device_if_changed(second, _host([1, 2])) - assert second[:2].tolist() == [1, 2] - - -def test_two_dimensional_values_are_flattened() -> None: - dst = torch.zeros(6, dtype=torch.int32) - copy_to_device_if_changed(dst, torch.tensor([[1, 2], [3, 4]], dtype=torch.int32)) - assert dst[:4].tolist() == [1, 2, 3, 4] - _tamper(dst) - copy_to_device_if_changed(dst, _host([1, 2, 3, 4])) - assert dst.eq(SENTINEL).all() - - -def test_debug_check_passes_a_buffer_only_the_function_wrote(monkeypatch) -> None: - monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", True) - dst = torch.zeros(8, dtype=torch.int64) - copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) - copy_to_device_if_changed(dst, _host([5, 6, 7, 8])) - copy_to_device_if_changed(dst, _host([5, 6])) - assert dst[:4].tolist() == [5, 6, 7, 8] - - -def test_debug_check_reports_a_write_by_anything_else(monkeypatch) -> None: - monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", True) - dst = torch.zeros(8, dtype=torch.int32) - copy_to_device_if_changed(dst, _host([3, 1])) - dst[:2].copy_(_host([5, 6])) - with pytest.raises(AssertionError, match="holds \\[5, 6\\] instead of \\[3, 1\\]"): - copy_to_device_if_changed(dst, _host([3, 1])) diff --git a/tests/unittest/_torch/executor/test_device_copy_buffer.py b/tests/unittest/_torch/executor/test_device_copy_buffer.py new file mode 100644 index 000000000000..6aa126d67b6c --- /dev/null +++ b/tests/unittest/_torch/executor/test_device_copy_buffer.py @@ -0,0 +1,205 @@ +# 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. + +import importlib.util +from pathlib import Path + +import pytest +import torch + +_MODULE_PATH = Path(__file__).resolve().parents[4] / "tensorrt_llm" / "_device_copy.py" +_SPEC = importlib.util.spec_from_file_location("device_copy", _MODULE_PATH) +device_copy = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(device_copy) +DeviceCopyBuffer = device_copy.DeviceCopyBuffer + +pytestmark = pytest.mark.cpu_only + +SENTINEL = -7 + + +@pytest.fixture(autouse=True) +def _no_skipped_copy_check(monkeypatch) -> None: + monkeypatch.setattr(device_copy, "_verify_skipped_device_copies", False) + + +def _host(values: list[int]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.int32) + + +def _tamper(dst: torch.Tensor) -> None: + dst.fill_(SENTINEL) + + +def test_first_call_copies() -> None: + dst = torch.zeros(8, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + buffer.update(_host([3, 1, 2])) + assert dst.tolist() == [3, 1, 2, 0, 0, 0, 0, 0] + + +def test_unchanged_values_skip_the_copy() -> None: + dst = torch.zeros(8, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + buffer.update(_host([3, 1, 2])) + _tamper(dst) + buffer.update(_host([3, 1, 2])) + assert dst.eq(SENTINEL).all() + + +def test_changed_values_copy() -> None: + dst = torch.zeros(8, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + buffer.update(_host([3, 1, 2])) + _tamper(dst) + buffer.update(_host([3, 1, 4])) + assert dst[:3].tolist() == [3, 1, 4] + assert dst[3:].eq(SENTINEL).all(), "only the leading elements are written" + + +def test_prefix_of_the_last_values_skips_the_copy() -> None: + dst = torch.zeros(8, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + buffer.update(_host([5, 6, 7, 8])) + _tamper(dst) + buffer.update(_host([5, 6])) + assert dst.eq(SENTINEL).all() + + +def test_longer_values_copy_and_extend_what_is_kept() -> None: + dst = torch.zeros(8, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + buffer.update(_host([5, 6])) + buffer.update(_host([5, 6, 7, 8])) + assert dst[:4].tolist() == [5, 6, 7, 8] + _tamper(dst) + buffer.update(_host([5, 6, 7, 8])) + assert dst.eq(SENTINEL).all() + + +def test_shorter_changed_values_keep_the_tail() -> None: + dst = torch.zeros(8, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + buffer.update(_host([5, 6, 7, 8])) + buffer.update(_host([9, 9])) + assert dst[:4].tolist() == [9, 9, 7, 8] + _tamper(dst) + buffer.update(_host([9, 9, 7, 8])) + assert dst.eq(SENTINEL).all() + buffer.update(_host([9, 9, 7, 1])) + assert dst[:4].tolist() == [9, 9, 7, 1] + + +def test_host_values_may_be_reused_right_away() -> None: + dst = torch.zeros(4, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + host = _host([1, 2, 3]) + buffer.update(host) + host.fill_(0) + assert dst[:3].tolist() == [1, 2, 3] + buffer.update(_host([1, 2, 3])) + assert dst[:3].tolist() == [1, 2, 3] + + +def test_a_new_buffer_starts_over() -> None: + first = torch.zeros(4, dtype=torch.int32) + first_buffer = DeviceCopyBuffer(first) + first_buffer.update(_host([1, 2])) + second = torch.full((4,), SENTINEL, dtype=torch.int32) + second_buffer = DeviceCopyBuffer(second) + second_buffer.update(_host([1, 2])) + assert second[:2].tolist() == [1, 2] + + +def test_two_dimensional_values_are_flattened() -> None: + dst = torch.zeros(6, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + buffer.update(torch.tensor([[1, 2], [3, 4]], dtype=torch.int32)) + assert dst[:4].tolist() == [1, 2, 3, 4] + _tamper(dst) + buffer.update(_host([1, 2, 3, 4])) + assert dst.eq(SENTINEL).all() + + +def test_debug_check_passes_a_buffer_only_the_function_wrote(monkeypatch) -> None: + monkeypatch.setattr(device_copy, "_verify_skipped_device_copies", True) + dst = torch.zeros(8, dtype=torch.int64) + buffer = DeviceCopyBuffer(dst) + buffer.update(_host([5, 6, 7, 8])) + buffer.update(_host([5, 6, 7, 8])) + buffer.update(_host([5, 6])) + assert dst[:4].tolist() == [5, 6, 7, 8] + + +def test_debug_check_reports_a_write_by_anything_else(monkeypatch) -> None: + monkeypatch.setattr(device_copy, "_verify_skipped_device_copies", True) + dst = torch.zeros(8, dtype=torch.int32) + buffer = DeviceCopyBuffer(dst) + buffer.update(_host([3, 1])) + dst[:2].copy_(_host([5, 6])) + with pytest.raises(AssertionError, match="holds \\[5, 6\\] instead of \\[3, 1\\]"): + buffer.update(_host([3, 1])) + + +def test_two_consumers_share_changed_then_restored_copy_state(monkeypatch) -> None: + storage = torch.full((4,), SENTINEL, dtype=torch.int32) + manager = DeviceCopyBuffer(storage) + metadata = manager + manager_view = manager.tensor.view_as(storage) + metadata_view = metadata.tensor[:2] + address = storage.data_ptr() + writes = [] + original = torch.Tensor.copy_ + + def counted(dst, src, *args, **kwargs): + if dst.data_ptr() == address: + writes.append(src.tolist()) + return original(dst, src, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "copy_", counted) + for consumer, values in ( + (manager, [1, 2]), + (metadata, [9, 2]), + (manager, [1, 2]), + ): + consumer.update(_host(values)) + assert manager_view.tolist() == values + [SENTINEL, SENTINEL] + assert metadata_view.tolist() == values + assert storage.data_ptr() == address + metadata.update(_host([1, 2])) + manager.update(_host([1, 2])) + assert writes == [[1, 2], [9, 2], [1, 2]] + + +def test_forced_writer_keeps_a_shared_alias_cache_coherent(monkeypatch) -> None: + storage = torch.zeros(2, dtype=torch.int32) + manager = DeviceCopyBuffer(storage) + metadata = manager + writes = [] + original = torch.Tensor.copy_ + + def counted(dst, src, *args, **kwargs): + if dst.data_ptr() == storage.data_ptr(): + writes.append(src.tolist()) + return original(dst, src, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "copy_", counted) + manager.update(_host([1, 2]), force=True) + manager.update(_host([1, 2]), force=True) + metadata.update(_host([9, 2])) + manager.update(_host([1, 2]), force=True) + metadata.update(_host([1, 2])) + assert storage.tolist() == [1, 2] + assert writes == [[1, 2], [1, 2], [9, 2], [1, 2]] diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py index da5ab48c7370..2553fe8f23ce 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_k3_mamba_metadata.py @@ -7,6 +7,7 @@ import pytest import torch +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.modules.fla.index import prepare_chunk_indices from tensorrt_llm._torch.modules.kimi_kda.kimi_k3_mamba_metadata import KimiK3MambaMetadata @@ -61,10 +62,12 @@ class KdaCacheManager: use_kda_replay_update = True def __init__(self) -> None: - self.state_indices = torch.tensor([9, 4, 7], dtype=torch.int32, device="cuda") + self._state_indices_buffer = DeviceCopyBuffer( + torch.tensor([9, 4, 7], dtype=torch.int32, device="cuda") + ) def get_state_indices(self, request_ids: list[int], is_padding: list[bool]) -> torch.Tensor: - return self.state_indices[: len(request_ids)] + return self._state_indices_buffer.tensor[: len(request_ids)] manager = KdaCacheManager() metadata = KimiK3MambaMetadata(max_batch_size=3, chunk_size=8, max_num_tokens=8) diff --git a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py index 167200bf9d37..2695a57e6f3a 100644 --- a/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py +++ b/tests/unittest/_torch/modules/kimi_kda/test_kimi_kda_verify_parity.py @@ -25,6 +25,7 @@ from kimi_kda_test_utils import get_production_decode_kernel_path +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.modules.kimi_kda import KimiKDALinearAttention from tensorrt_llm._torch.modules.kimi_kda.kimi_k3_mamba_metadata import KimiK3MambaMetadata from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream @@ -128,10 +129,12 @@ class KdaCacheManager: use_kda_replay_update = True def __init__(self) -> None: - self.state_indices = torch.tensor([9, 4], dtype=torch.int32, device="cuda") + self._state_indices_buffer = DeviceCopyBuffer( + torch.tensor([9, 4], dtype=torch.int32, device="cuda") + ) def get_state_indices(self, request_ids: list[int], is_padding: list[bool]) -> torch.Tensor: - return self.state_indices[: len(request_ids)] + return self._state_indices_buffer.tensor[: len(request_ids)] def mamba_layer_cache(self, layer_idx: int) -> SimpleNamespace: return layer_cache diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index 44d2fd4c3055..04f06439febd 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -19,6 +19,7 @@ import pytest import torch +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.modules.mamba import mamba2_metadata from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( REPLAY_WORK_CACHE_BUF_IDX, @@ -166,7 +167,7 @@ def test_prepare_skips_an_unchanged_state_index_upload(self, monkeypatch): """A decode step whose state indices (a list from the cache manager) did not change enqueues no upload: the device buffer keeps values the host never sent. A changed step uploads again.""" # The test writes the buffer behind prepare() to see the skip, which the check of skipped copies reports. - monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", False) + monkeypatch.setattr("tensorrt_llm._device_copy._verify_skipped_device_copies", False) manager = _StateIndexCacheManager([3, 1]) metadata = Mamba2Metadata(max_batch_size=2, chunk_size=8) attn_metadata = _decode_attn_metadata(manager) @@ -206,7 +207,7 @@ def test_prepare_uploads_state_indices_from_a_list_or_a_cpu_tensor(self): def test_check_of_skipped_copies_reports_a_write_behind_prepare(self, monkeypatch): """With the check of skipped copies on (TLLM_DEBUG_MODE=1, or the unit-test fixture), a state index buffer that something other than prepare() wrote fails the next prepare() whose indices did not change.""" - monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", True) + monkeypatch.setattr("tensorrt_llm._device_copy._verify_skipped_device_copies", True) manager = _StateIndexCacheManager([3, 1]) metadata = Mamba2Metadata(max_batch_size=2, chunk_size=8) attn_metadata = _decode_attn_metadata(manager) @@ -217,6 +218,52 @@ def test_check_of_skipped_copies_reports_a_write_behind_prepare(self, monkeypatc with pytest.raises(AssertionError, match="holds \\[5, 6\\] instead of \\[3, 1\\]"): metadata.prepare(attn_metadata) + def test_prepare_alias_keeps_one_owner_across_manager_and_metadata_writes(self, monkeypatch): + class CacheManager: + def __init__(self): + self._state_indices_buffer = DeviceCopyBuffer( + torch.zeros(2, dtype=torch.int32, device="cuda") + ) + self.alias = True + self.values = [9, 2] + + def get_state_indices(self, request_ids, is_padding): + if self.alias: + return self._state_indices_buffer.tensor + return self.values + + manager = CacheManager() + metadata = Mamba2Metadata(max_batch_size=2, chunk_size=8) + attn_metadata = _decode_attn_metadata(manager) + address = manager._state_indices_buffer.tensor.data_ptr() + writes = [] + original = torch.Tensor.copy_ + + def counted(dst, src, *args, **kwargs): + if dst.is_cuda and dst.data_ptr() == address: + writes.append(src.tolist()) + return original(dst, src, *args, **kwargs) + + monkeypatch.setattr(torch.Tensor, "copy_", counted) + manager._state_indices_buffer.update(torch.tensor([1, 2], dtype=torch.int32)) + metadata.prepare(attn_metadata) + torch.cuda.synchronize() + assert metadata.state_indices.tolist() == [1, 2] + manager.alias = False + metadata.prepare(attn_metadata) + torch.cuda.synchronize() + assert manager._state_indices_buffer.tensor.tolist() == [9, 2] + manager._state_indices_buffer.update(torch.tensor([1, 2], dtype=torch.int32)) + manager.alias = True + metadata.prepare(attn_metadata) + manager._state_indices_buffer.update(torch.tensor([1, 2], dtype=torch.int32)) + metadata.prepare(attn_metadata) + torch.cuda.synchronize() + assert metadata.state_indices.tolist() == [1, 2] + assert manager._state_indices_buffer.tensor.tolist() == [1, 2] + assert metadata.state_indices.data_ptr() == address + assert writes == [[1, 2], [9, 2], [1, 2]] + def test_prepare_replay_work_items_write_first(self): class ReplayCacheManager: use_replay_state_update = True diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py index e58a89378a6c..f88ba9687763 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py @@ -1,20 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""CPU unit tests for the DFlash worker's per-step slot buffers and copy_to_device_if_changed. - -DFlashSpecMetadata.prepare writes batch_indices_cuda and the worker's _batch_to_slot, and the -worker's post-prefill rebuild writes _batch_to_slot again. Both buffers go through -copy_to_device_if_changed, which skips a copy whose values equal the last ones it copied. The -tests check that each buffer always holds the latest mapping, including a step whose mapping -returns to an earlier one after the rebuild wrote another, and that an unchanged step enqueues -no copy. -""" from types import SimpleNamespace import pytest import torch +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker pytestmark = pytest.mark.cpu_only @@ -30,7 +22,7 @@ def _worker(req_to_slot): _free_slots=[], _graph_dummy_id_floor=1 << 30, _dummy_slot=DUMMY_SLOT, - _batch_to_slot=torch.zeros(8, dtype=torch.long), + _batch_to_slot_buffer=DeviceCopyBuffer(torch.zeros(8, dtype=torch.long)), _write_ctx_len=lambda updates: None, ) # Generation requests are already assigned in these tests; prepare assigns nothing. @@ -39,12 +31,14 @@ def _worker(req_to_slot): def _metadata(worker, request_ids, num_generations): - return SimpleNamespace( + metadata = SimpleNamespace( request_ids=list(request_ids), num_generations=num_generations, batch_indices_cuda=torch.full((8,), -1, dtype=torch.int), _dflash_worker=worker, ) + metadata._batch_indices_buffer = DeviceCopyBuffer(metadata.batch_indices_cuda) + return metadata def _prepare(metadata, request_ids, num_generations): @@ -63,15 +57,15 @@ def test_mapping_that_returns_after_a_rebuild_is_written_again(): worker = _worker({5: 1}) metadata = _metadata(worker, [8, 5], num_generations=1) _prepare(metadata, [8, 5], num_generations=1) - assert worker._batch_to_slot[:2].tolist() == [DUMMY_SLOT, 1] + assert worker._batch_to_slot_buffer.tensor[:2].tolist() == [DUMMY_SLOT, 1] worker._req_to_slot[8] = 4 _rebuild(worker, [8, 5]) - assert worker._batch_to_slot[:2].tolist() == [4, 1] + assert worker._batch_to_slot_buffer.tensor[:2].tolist() == [4, 1] # Request 8 is gone and context request 9 takes its row: prepare maps [dummy, 1] again, the # values it copied two writes ago. The buffer must hold them, not the rebuild's [4, 1]. del worker._req_to_slot[8] _prepare(metadata, [9, 5], num_generations=1) - assert worker._batch_to_slot[:2].tolist() == [DUMMY_SLOT, 1] + assert worker._batch_to_slot_buffer.tensor[:2].tolist() == [DUMMY_SLOT, 1] def test_rebuild_then_prepare_back_to_back_keep_the_latest_mapping(): @@ -79,9 +73,9 @@ def test_rebuild_then_prepare_back_to_back_keep_the_latest_mapping(): metadata = _metadata(worker, [5, 6], num_generations=2) for request_ids, slots in (([5, 6], [1, 2]), ([6, 5], [2, 1]), ([5, 6], [1, 2])): _rebuild(worker, request_ids) - assert worker._batch_to_slot[:2].tolist() == slots + assert worker._batch_to_slot_buffer.tensor[:2].tolist() == slots _prepare(metadata, request_ids, num_generations=2) - assert worker._batch_to_slot[:2].tolist() == slots + assert worker._batch_to_slot_buffer.tensor[:2].tolist() == slots def test_batch_indices_hold_the_batch_rows(): @@ -108,4 +102,4 @@ def counted(dst, src, *args, **kwargs): _prepare(metadata, [5, 6], num_generations=2) monkeypatch.undo() assert copies == [] - assert worker._batch_to_slot[:2].tolist() == [1, 2] + assert worker._batch_to_slot_buffer.tensor[:2].tolist() == [1, 2] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py index ace488b3fad5..269e1de42197 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_dummy_slot.py @@ -27,6 +27,7 @@ import pytest import torch +from tensorrt_llm._device_copy import DeviceCopyBuffer from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata, DFlashWorker pytestmark = pytest.mark.skipif( @@ -47,7 +48,7 @@ def _worker(lengths): _max_ctx=1000, _ctx_len=torch.tensor(lengths, dtype=torch.long, device="cuda"), _ctx_len_host=list(lengths), - _batch_to_slot=torch.zeros(8, dtype=torch.long, device="cuda"), + _batch_to_slot_buffer=DeviceCopyBuffer(torch.zeros(8, dtype=torch.long, device="cuda")), ) w._write_ctx_len = partial(DFlashWorker._write_ctx_len, w) w._advance_ctx_len = partial(DFlashWorker._advance_ctx_len, w) @@ -62,6 +63,7 @@ def _prepare(worker, request_ids, num_generations): batch_indices_cuda=torch.zeros(8, dtype=torch.int, device="cuda"), _dflash_worker=worker, ) + meta._batch_indices_buffer = DeviceCopyBuffer(meta.batch_indices_cuda) DFlashSpecMetadata.prepare(meta) torch.cuda.synchronize() @@ -74,7 +76,7 @@ def test_padded_steps_keep_the_dummy_slot_empty(): worker = _worker([40, 300, 7, 0, 0]) # Three requests in a graph of five. _prepare(worker, [11, 12, 13, FLOOR + 1, FLOOR + 1], num_generations=5) - slots = worker._batch_to_slot[:5].clone() + slots = worker._batch_to_slot_buffer.tensor[:5].clone() assert slots.tolist() == [0, 1, 2, 4, 4] for step in range(1, 5): worker._advance_ctx_len(slots, _cuda([8] * 5)) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py index b1bcff6d8c73..947ee792d725 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_ragged_spec_sampler.py @@ -612,7 +612,7 @@ def record_event(): sampler._record_sampler_event = record_event sampler._async_worker_active = lambda: True sampler._store_scatter = None - sampler._slot_table = None + sampler._slot_table_buffer = None sampler._store_copies_done = None return sampler, transfers diff --git a/tests/unittest/conftest.py b/tests/unittest/conftest.py index 364cc6f2d670..fe4538843d23 100644 --- a/tests/unittest/conftest.py +++ b/tests/unittest/conftest.py @@ -392,12 +392,8 @@ def torch_empty_cache() -> None: @pytest.fixture(autouse=True) def verify_skipped_device_copies(monkeypatch) -> None: - """Check that every buffer copy_to_device_if_changed skips holds the values. - - A test in which anything else writes such a buffer then fails. - """ - monkeypatch.setattr("tensorrt_llm._utils._verify_skipped_device_copies", - True) + monkeypatch.setattr( + "tensorrt_llm._device_copy._verify_skipped_device_copies", True) @pytest.fixture(scope="module") From 597e23a47721aaa8fdaf50b93ffaa4738bf9e3d2 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Wed, 14 Oct 2026 11:53:02 -0700 Subject: [PATCH 64/67] Reject malformed DSpark auxiliary layer IDs before translation Signed-off-by: Vasanth Sabavat --- docs/source/models/supported-models.md | 2 +- tensorrt_llm/llmapi/llm_args.py | 6 ++- .../test_dspark_speculators_checkpoint.py | 48 +++++++++++++++++++ 3 files changed, 53 insertions(+), 3 deletions(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index fd9721a4acfe..118e43653600 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -96,7 +96,7 @@ statuses for the same architecture in the two matrices. [^14]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`). [^15]: Kimi K3 is only supported on NVIDIA Blackwell GPUs (`SM100` family); see the [Kimi K3 deployment guide](../deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md). Which recipes fit is set by the attention layout rather than by the GPU count: DEP16 (`enable_attention_dp: true`) replicates the BF16 non-expert weights on every rank (114 GB) on top of the MXFP4 routed experts at 16-way expert parallelism (90 GB), needing 210 GB per rank, so it requires GB300-class per-GPU memory. TEP16 (`enable_attention_dp: false`) shards those non-expert weights instead and needs 115 GB per rank; it is validated end-to-end on GB200 (`SM100`) at 16 GPUs. On B200 (`SM100`), kernel and module support is functional and covered by CI. Note that the FP8 weight-read path (TRTLLM-14765) does not relax the DEP16 requirement: the conversion runs after the weights are already resident in BF16, so it lowers the steady-state footprint but not the load-time peak. [^16]: Guided decoding for `Glm4MoeForCausalLM` is currently supported only with the `xgrammar` backend; `llguidance` is not working. -[^17]: Kimi K3 has no MTP or EAGLE-3 head, and its DSpark checkpoints are not compatible with plain `DFlash`. +[^17]: Kimi K3 has no MTP or EAGLE-3 head, and its DSpark checkpoints are not compatible with plain `DFlash`. Speculators-format DSpark checkpoints require `aux_hidden_state_layer_ids` to be a nonempty list of positive integers; booleans, floats and strings are rejected rather than coerced. [^18]: NGram and standalone Suffix Automaton (SA) use model-free drafting on the PyTorch backend, so they are not listed in individual entries. This does not imply universal end-to-end support: compatibility depends on each model's multi-token verification and cache-management paths and may be untested or explicitly restricted. [^19]: KV cache reuse for hybrid recurrent-attention models requires an explicit recurrent-state snapshot policy, such as `kv_cache_config.mamba_state_config.periodic_snapshot_interval`; the model default disables reuse when no snapshot policy is configured. [^20]: Supports text, image, and video inputs, MTP (including with attention data parallelism), and FP8 KV cache. Image and video inputs require `transformers==5.17.0`. Beam search is not supported. See the [GLM-5.3-Flash deployment guide](../deployment-guide/deployment-guide-for-glm-5.3-flash-on-trtllm.md) for setup and feature-specific requirements. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 01f76b2886f9..f3962c6e93a6 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3383,14 +3383,16 @@ def translate_speculators_dspark_config(config_dict: dict) -> dict: "speculators DSpark drafters are supported with a qwen3 backbone; " f"this one has model_type {layer.get('model_type')!r}") aux_layer_ids = config_dict.get("aux_hidden_state_layer_ids") - if not aux_layer_ids or min(aux_layer_ids) < 1: + if (not isinstance(aux_layer_ids, list) or not aux_layer_ids or any( + type(layer_id) is not int or layer_id < 1 + for layer_id in aux_layer_ids)): raise ValueError( "a speculators DSpark drafter needs aux_hidden_state_layer_ids >= 1 " "(layer 0's input is the embedding output, which is not captured); " f"got {aux_layer_ids}") num_layers = layer["num_hidden_layers"] dflash_config = { - "target_layer_ids": [int(layer_id) - 1 for layer_id in aux_layer_ids] + "target_layer_ids": [layer_id - 1 for layer_id in aux_layer_ids] } if config_dict.get("mask_token_id") is not None: dflash_config["mask_token_id"] = config_dict["mask_token_id"] diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py index 496db57dc970..07db5bbe1a9e 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dspark_speculators_checkpoint.py @@ -172,6 +172,54 @@ def test_translation_is_the_converted_config(): assert translate_speculators_dspark_config(SPECULATORS_CONFIG) == CONVERTED_CONFIG +@pytest.mark.parametrize( + "aux_layer_ids", + [ + None, + [], + 1, + 1.9, + True, + "1", + {"1": 2}, + (1, 2), + [True], + [False], + [1.9], + [1.0], + ["1"], + [1, "2"], + [1, 2.0], + [1, True], + [[1]], + [None], + [0], + [-1], + [1, 0], + [1, -1], + ], +) +def test_translation_rejects_invalid_aux_layer_ids(aux_layer_ids): + config = copy.deepcopy(SPECULATORS_CONFIG) + config["aux_hidden_state_layer_ids"] = aux_layer_ids + + with pytest.raises(ValueError): + translate_speculators_dspark_config(config) + + +@pytest.mark.parametrize( + "aux_layer_ids, target_layer_ids", + [([1], [0]), ([3, 1, 3], [2, 0, 2])], +) +def test_translation_preserves_positive_integer_layer_ids(aux_layer_ids, target_layer_ids): + config = copy.deepcopy(SPECULATORS_CONFIG) + config["aux_hidden_state_layer_ids"] = aux_layer_ids + expected = copy.deepcopy(CONVERTED_CONFIG) + expected["dflash_config"]["target_layer_ids"] = target_layer_ids + + assert translate_speculators_dspark_config(config) == expected + + def test_drafter_config_matches_the_converted_checkpoint(tmp_path): native = load_pretrained_config(_checkpoint(tmp_path, "native", SPECULATORS_CONFIG)) converted = load_pretrained_config(_checkpoint(tmp_path, "converted", CONVERTED_CONFIG)) From 3261a6fe280620e6b3c93571351811f04ef8d427 Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 15 Oct 2026 08:12:09 -0700 Subject: [PATCH 65/67] Address speculative runtime review nits Signed-off-by: Vasanth Sabavat --- tensorrt_llm/_device_copy.py | 2 +- tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py | 4 ---- tensorrt_llm/llmapi/llm_args.py | 3 ++- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_device_copy.py b/tensorrt_llm/_device_copy.py index 416911bf493e..0aef62d3fa05 100644 --- a/tensorrt_llm/_device_copy.py +++ b/tensorrt_llm/_device_copy.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import os diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 8b174c8d1785..1247cd5cc53c 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -796,10 +796,6 @@ def capture(self, # Do not keep the eager result live from this runner across graph # setup/capture; release its reference before entering. output = None - # A captured forward does not run, so there is no in-place input - # change to undo (postprocess_fn is for the warmup forwards above). - # Undoing one here would shift the static inputs, which are the - # engine's own buffers when static_input_ids is set. with torch.cuda.graph(graph, pool=self.memory_pool): output = forward_fn(capture_inputs) # The addresses the graph baked in, for the strict buffer check. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index f3962c6e93a6..6ee4a2e20242 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3387,7 +3387,8 @@ def translate_speculators_dspark_config(config_dict: dict) -> dict: type(layer_id) is not int or layer_id < 1 for layer_id in aux_layer_ids)): raise ValueError( - "a speculators DSpark drafter needs aux_hidden_state_layer_ids >= 1 " + "a speculators DSpark drafter needs aux_hidden_state_layer_ids to be " + "a non-empty list of positive integers (>= 1, excluding bool) " "(layer 0's input is the embedding output, which is not captured); " f"got {aux_layer_ids}") num_layers = layer["num_hidden_layers"] From 6b3d878d0783014746d09a8841bf029e4ddcdbba Mon Sep 17 00:00:00 2001 From: Vasanth Sabavat Date: Thu, 15 Oct 2026 22:54:58 -0700 Subject: [PATCH 66/67] [None][fix] Align Kimi cache test fixture with upstream mapping fields Signed-off-by: Vasanth Sabavat --- .../_torch/executor/kv_cache/test_mamba_cache_manager.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 324070ce6fed..7e5f568191de 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py @@ -302,12 +302,11 @@ def _kimi_model_config() -> SimpleNamespace: ) return SimpleNamespace( pretrained_config=config, + mapping=Mapping(), quant_config=None, sparse_attention_config=None, - get_num_mamba_layers=lambda: 2, - # get_kv_cache_manager_cls reads the mapping and spec config (helix check). - mapping=None, spec_config=None, + get_num_mamba_layers=lambda: 2, ) From 14a76b0fd116f2dee5d1a2cebcc67829c212436b Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Sat, 17 Oct 2026 00:56:50 +0800 Subject: [PATCH 67/67] [None][fix] Refresh serving configuration schemas (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../schemas/trtllm-serve-config.schema.json | 66 ++++++++++++++++++- .../trtllm-serve-disagg-config.schema.json | 66 ++++++++++++++++++- ...trtllm-serve-visual-gen-config.schema.json | 18 +++++ 3 files changed, 146 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 1211e43a6390..bdbb70d826af 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -2798,7 +2798,7 @@ }, "enable_kv_pool_rebalance": { "default": false, - "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by V2's internal 2000-sample / 120s cooldown. When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", + "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by the ``kv_pool_rebalance_config`` thresholds (2000 samples / 120s by default). When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", "status": "prototype", "title": "Enable Kv Pool Rebalance", "type": "boolean" @@ -2897,6 +2897,11 @@ "description": "Streaming (push-based) KV cache event publishing (KV cache manager V2 only). When set, each rank publishes its own events directly (e.g. over ZeroMQ) instead of the buffered event_buffer_max_size gather/poll path.", "status": "prototype" }, + "kv_pool_rebalance_config": { + "$ref": "#/$defs/KvPoolRebalanceConfig", + "description": "Tuning for the KVCacheManagerV2 pool rebalance auto-tuner. Only used when ``enable_kv_pool_rebalance`` is True.", + "status": "prototype" + }, "mamba_ssm_cache_dtype": { "default": "auto", "description": "The data type to use for the Mamba SSM cache. If set to 'auto', the data type will be inferred from the model config.", @@ -3159,6 +3164,63 @@ "title": "KvCacheConnectorConfig", "type": "object" }, + "KvPoolRebalanceConfig": { + "additionalProperties": false, + "description": "Tuning for the KV cache manager v2 pool rebalance auto-tuner.\n\nOnly used when ``kv_cache_config.enable_kv_pool_rebalance`` is True.", + "properties": { + "check_interval": { + "default": 10, + "description": "Number of PyExecutor iterations between checks of whether a pool rebalance should be attempted.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Check Interval", + "type": "integer" + }, + "cooldown_secs": { + "default": 120.0, + "description": "Minimum time, in seconds, between successive rebalance adjustments.", + "minimum": 0, + "status": "prototype", + "title": "Cooldown Secs", + "type": "number" + }, + "min_sampled_kv_caches": { + "default": 2000, + "description": "Minimum number of sampled (closed) KvCaches observed before the auto-tuner will consider making an adjustment.", + "minimum": 0, + "status": "prototype", + "title": "Min Sampled Kv Caches", + "type": "integer" + }, + "moving_average_decay": { + "default": 0.9999, + "description": "Decay factor for the exponential moving averages (reused length, capacity, history length) used to compute target pool ratios. Higher values react more slowly to change.", + "exclusiveMaximum": 1, + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Moving Average Decay", + "type": "number" + }, + "ratio_tolerance": { + "default": 0.25, + "description": "How far a pool group's current ratio may drift from its target before a rebalance is triggered, as a fraction of the smaller of the two. The comparison is symmetric, so 0.25 means a rebalance is triggered once one of the two exceeds the other by more than 25%.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Ratio Tolerance", + "type": "number" + }, + "target_ratio_update_interval": { + "default": 100, + "description": "Number of newly sampled KvCaches between recomputations of the target pool-group ratios.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Target Ratio Update Interval", + "type": "integer" + } + }, + "title": "KvPoolRebalanceConfig", + "type": "object" + }, "LayerwiseBenchmarksConfig": { "additionalProperties": false, "description": "Configuration for layer-wise benchmarks calibration.", @@ -5653,7 +5715,7 @@ } ], "default": null, - "description": "The dynamic batch config to use. This only applies for the TensorRT backend and cannot currently be used with the PyTorch backend." + "description": "The dynamic batch configuration used by the scheduler to adjust batch size and token limits at runtime. This option currently has no effect since the scheduler does not support dynamic batch adjustments in the pytorch backend." }, "enable_prefix_aware_scheduling": { "default": true, diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 15599ad03a1f..0c7fd1773e10 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -4348,7 +4348,7 @@ }, "enable_kv_pool_rebalance": { "default": false, - "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by V2's internal 2000-sample / 120s cooldown. When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", + "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by the ``kv_pool_rebalance_config`` thresholds (2000 samples / 120s by default). When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", "status": "prototype", "title": "Enable Kv Pool Rebalance", "type": "boolean" @@ -4447,6 +4447,11 @@ "description": "Streaming (push-based) KV cache event publishing (KV cache manager V2 only). When set, each rank publishes its own events directly (e.g. over ZeroMQ) instead of the buffered event_buffer_max_size gather/poll path.", "status": "prototype" }, + "kv_pool_rebalance_config": { + "$ref": "#/$defs/KvPoolRebalanceConfig", + "description": "Tuning for the KVCacheManagerV2 pool rebalance auto-tuner. Only used when ``enable_kv_pool_rebalance`` is True.", + "status": "prototype" + }, "mamba_ssm_cache_dtype": { "default": "auto", "description": "The data type to use for the Mamba SSM cache. If set to 'auto', the data type will be inferred from the model config.", @@ -4709,6 +4714,63 @@ "title": "KvCacheConnectorConfig", "type": "object" }, + "KvPoolRebalanceConfig": { + "additionalProperties": false, + "description": "Tuning for the KV cache manager v2 pool rebalance auto-tuner.\n\nOnly used when ``kv_cache_config.enable_kv_pool_rebalance`` is True.", + "properties": { + "check_interval": { + "default": 10, + "description": "Number of PyExecutor iterations between checks of whether a pool rebalance should be attempted.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Check Interval", + "type": "integer" + }, + "cooldown_secs": { + "default": 120.0, + "description": "Minimum time, in seconds, between successive rebalance adjustments.", + "minimum": 0, + "status": "prototype", + "title": "Cooldown Secs", + "type": "number" + }, + "min_sampled_kv_caches": { + "default": 2000, + "description": "Minimum number of sampled (closed) KvCaches observed before the auto-tuner will consider making an adjustment.", + "minimum": 0, + "status": "prototype", + "title": "Min Sampled Kv Caches", + "type": "integer" + }, + "moving_average_decay": { + "default": 0.9999, + "description": "Decay factor for the exponential moving averages (reused length, capacity, history length) used to compute target pool ratios. Higher values react more slowly to change.", + "exclusiveMaximum": 1, + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Moving Average Decay", + "type": "number" + }, + "ratio_tolerance": { + "default": 0.25, + "description": "How far a pool group's current ratio may drift from its target before a rebalance is triggered, as a fraction of the smaller of the two. The comparison is symmetric, so 0.25 means a rebalance is triggered once one of the two exceeds the other by more than 25%.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Ratio Tolerance", + "type": "number" + }, + "target_ratio_update_interval": { + "default": 100, + "description": "Number of newly sampled KvCaches between recomputations of the target pool-group ratios.", + "exclusiveMinimum": 0, + "status": "prototype", + "title": "Target Ratio Update Interval", + "type": "integer" + } + }, + "title": "KvPoolRebalanceConfig", + "type": "object" + }, "LayerwiseBenchmarksConfig": { "additionalProperties": false, "description": "Configuration for layer-wise benchmarks calibration.", @@ -7203,7 +7265,7 @@ } ], "default": null, - "description": "The dynamic batch config to use. This only applies for the TensorRT backend and cannot currently be used with the PyTorch backend." + "description": "The dynamic batch configuration used by the scheduler to adjust batch size and token limits at runtime. This option currently has no effect since the scheduler does not support dynamic batch adjustments in the pytorch backend." }, "enable_prefix_aware_scheduling": { "default": true, diff --git a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json index 6cab6b77e810..cffa1974bcdc 100644 --- a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json @@ -397,6 +397,24 @@ "title": "Ring Size", "type": "integer" }, + "tp_layout": { + "anyOf": [ + { + "enum": [ + "replicated", + "token_sharded" + ], + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Activation layout between TP projections: 'replicated' (all-reduce TP; the default when unset) or 'token_sharded' (each TP rank holds 1/tp_size of the tokens; each all-reduce becomes a reduce-scatter plus an all-gather). 'token_sharded' requires tp_size > 1 and no ulysses/ring/attn2d.", + "status": "prototype", + "title": "Tp Layout" + }, "tp_size": { "default": 1, "description": "Tensor parallel group size. Heads are sharded across tp_size GPUs.",