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/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index d6943819348c..28075bae9997 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -97,7 +97,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/_device_copy.py b/tensorrt_llm/_device_copy.py new file mode 100644 index 000000000000..0aef62d3fa05 --- /dev/null +++ b/tensorrt_llm/_device_copy.py @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 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/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 909e208129b0..e60f2cbbcfc2 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1424,6 +1424,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..feba742bdefb --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/op.py @@ -0,0 +1,283 @@ +# 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_``). ``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 +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 +from ...cute_dsl_utils import IS_CUTLASS_DSL_AVAILABLE + + +def is_supported() -> bool: + """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, but they are validated on SM 100 only; SM 103 and SM 107 keep the torch path. + """ + 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 _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. 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 + + 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 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 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: + """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]``, + ``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, launching nothing, when an argument is outside what the kernel covers; 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 ( + 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 + args = ( + 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, + ) + if StepInputGather._kernel is None: + StepInputGather._kernel = _compile("gather", args) + StepInputGather._kernel(*args, torch.cuda.current_stream().cuda_stream) + return True 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..90378e25ed00 --- /dev/null +++ b/tensorrt_llm/_torch/cute_dsl_kernels/spec_step_copies/spec_step_copies_kernel.py @@ -0,0 +1,209 @@ +# 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; + a row's new tokens past its accepted length are written as zeros). +* ``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. +""" + +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`` in global memory (trace-time helper).""" + 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] + 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: + 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 + 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] = accepted + + +@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 + + +@cute.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] + 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, +): + """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 + if t < rows * tokens_per_row: + 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 + + +@cute.jit +def gather( + 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, + stream: cuda_driver.CUstream, +) -> None: + 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, + ).launch( + grid=[(rows * tokens_per_row + THREADS - 1) // THREADS, 1, 1], + block=[THREADS, 1, 1], + stream=stream, + ) # fmt: skip diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index 32109f3f637c..40ca00aa3ec4 100644 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -1894,12 +1894,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/modules/kimi_kda/_kda_kernels.py b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py index 151e6b4a799e..a88956f19d84 100644 --- a/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py +++ b/tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py @@ -434,17 +434,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 @@ -471,10 +478,10 @@ def commit_kda_pending_drafts( conv_q: torch.Tensor, conv_k: torch.Tensor, conv_v: torch.Tensor, - k_cache: torch.Tensor, - g_cache: torch.Tensor, - v_cache: torch.Tensor, - beta_cache: torch.Tensor, + k_cache: Optional[torch.Tensor], + g_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, @@ -487,19 +494,25 @@ def commit_kda_pending_drafts( hold the state and the conv window that verify would start from (the window 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 (``k_cache`` / - ``g_cache`` / ``v_cache`` / ``beta_cache``) by the fused verify kernel's - own replay steps (``kda_mtp_commit_pending_drafts``). The tensor-core - verify computes a chunk in one pass, so this commit rounds the state - differently from a verify that replays the drafts itself. + * with the Kimi K3 verify kernels' draft records (``state_tok``), the + 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``). The + tensor-core verify computes a chunk in one pass, so this commit + rounds the state differently from a verify that replays the drafts + itself. """ 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: + assert ( + k_cache is not None + and g_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/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index c1aa57f19e99..efead8e0d799 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -21,6 +21,7 @@ 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 @@ -374,6 +375,8 @@ def __init__( 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. @@ -490,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]) - self.state_indices[:batch_size].copy_( - self.state_indices_cpu[:batch_size], non_blocking=True) + 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 @@ -505,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)) - self.state_indices[:batch_size].copy_( - self.state_indices_cpu[:batch_size], non_blocking=True) + 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/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9ad518e3db5f..d92cecdaa033 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -3511,6 +3511,14 @@ def _create_kv_cache_manager( spec_config.tokens_per_gen_step - 1) if kda_replay_num_spec is not None: kda_extra_kwargs["kda_replay_num_spec"] = kda_replay_num_spec + # 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_draft_records`. + backbone = getattr(getattr(model_engine, "model", None), + "model", None) + if (issubclass(kv_cache_manager_cls, MambaHybridCacheManagerV2) + 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/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 1100ced91274..d2b2aa28a6f1 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -150,6 +150,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: @@ -225,6 +232,25 @@ def _create_shared_static_tensors(self): max_total_tokens = min(max_total_tokens, self.config.max_num_tokens) self._pre_replay_position_ids = None + 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), @@ -765,11 +791,8 @@ def capture(self, output = None with torch.cuda.graph(graph, pool=self.memory_pool): output = forward_fn(capture_inputs) - # Snapshot before postprocessing so the check compares against - # the addresses the graph actually baked in. + # The addresses the graph baked in, for the strict buffer check. self._record_strict_buffer_snapshot(key) - if postprocess_fn is not None: - postprocess_fn(capture_inputs) _restore_spec_decode_capture_state(attn_metadata, saved_kv_lens_cuda) @@ -809,7 +832,11 @@ def replay( 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: @@ -856,7 +883,9 @@ def replay( 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/engine/runners/decoder/inputs.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py index e6b5f2df85c1..6baf6e5d9bc9 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/inputs.py @@ -11,8 +11,10 @@ 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.cute_dsl_kernels.spec_step_copies import op as spec_step_copies from tensorrt_llm._torch.distributed import Distributed from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input @@ -135,7 +137,6 @@ class DecoderInputBuffers(InputBuffers): previous_batch_indices_cuda: torch.Tensor # Host counter of the steady-state generation prepare cache. steady_gen_positions_pinned: torch.Tensor - num_accepted_draft_tokens_cuda: torch.Tensor | None = None previous_pos_indices_cuda: torch.Tensor | None = None previous_pos_id_offsets_cuda: torch.Tensor | None = None previous_kv_lens_offsets_cuda: torch.Tensor | None = None @@ -167,11 +168,6 @@ def for_config( steady_gen_positions_pinned=torch.empty( (config.max_num_tokens,), dtype=torch.int, pin_memory=prefer_pinned() ), - num_accepted_draft_tokens_cuda=( - torch.empty((config.max_batch_size,), dtype=torch.int, device="cuda") - if spec - else None - ), previous_pos_indices_cuda=( torch.empty((config.max_num_tokens,), dtype=torch.int, device="cuda") if spec @@ -241,7 +237,6 @@ class HostLists: mrope_position_ids: list[tuple[int, int, torch.Tensor]] = field(default_factory=list) mrope_delta_write_seq_slots: list[int] = field(default_factory=list) mrope_delta_read_seq_slots: list[int] = field(default_factory=list) - num_accepted_draft_tokens: list[int] = field(default_factory=list) # per request # will contain previous batch indices of generation requests previous_batch_indices: list[int] = field(default_factory=list) previous_pos_indices: list[int] = field(default_factory=list) @@ -339,6 +334,19 @@ def __init__( self._steady_gen_cache_supported = steady_gen_cache_supported self.guided_decoder: CapturableGuidedDecoder | None = None self._prepare_inputs_event: torch.cuda.Event | None = None + self._gather_ids_buffer = ( + DeviceCopyBuffer(buffers.gather_ids_cuda, pin_memory=prefer_pinned()) + if config.is_spec_decode + else None + ) + # Reuse device indices while the scheduled batch stays unchanged. + self._staged_previous_batch_indices: list[int] | None = None + self._staged_previous_pos_indices: list[int] | None = None + self._step_input_gather: spec_step_copies.StepInputGather | None = ( + spec_step_copies.StepInputGather() + if config.is_spec_decode and spec_step_copies.is_supported() + else None + ) # Steady-state generation-only prepare cache (non-speculative overlap # decode). Holds the per-request lists that are invariant while the # scheduled generation batch keeps the same composition, plus a pinned @@ -832,14 +840,13 @@ def prepare_inputs( new_tokens_lens_device, next_draft_tokens_device, runtime_draft_len, + attn_metadata=attn_metadata, + enable_spec_decode=enable_spec_decode, ) self._copy_position_ids_to_device(host) if enable_spec_decode: - self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)].copy_( - torch.tensor(host.lists.gather_ids, dtype=torch.int, pin_memory=prefer_pinned()), - non_blocking=True, - ) + self._gather_ids_buffer.update(torch.tensor(host.lists.gather_ids, dtype=torch.int)) self._prepare_attn_metadata( host, @@ -891,7 +898,7 @@ def prepare_inputs( ) gather_ids_device = ( - self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)] + self._gather_ids_buffer.tensor[: len(host.lists.gather_ids)] if enable_spec_decode else None ) @@ -965,7 +972,6 @@ def _pack_context_requests( position_ids = host.lists.position_ids num_cached_tokens_per_seq = host.lists.num_cached_tokens_per_seq draft_lens = host.lists.draft_lens - num_accepted_draft_tokens = host.lists.num_accepted_draft_tokens multimodal_params_list = host.lists.multimodal_params_list extra_inputs = host.extra_inputs context_prompt_lookahead = host.context_prompt_lookahead @@ -1000,7 +1006,6 @@ def _pack_context_requests( 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 - request.py_num_compressed_tokens) @@ -1261,7 +1266,6 @@ def _pack_extend_requests( draft_lens = host.lists.draft_lens prompt_lengths = host.lists.prompt_lengths sequence_lengths = host.lists.sequence_lengths - num_accepted_draft_tokens = host.lists.num_accepted_draft_tokens gather_ids = host.lists.gather_ids position_ids = host.lists.position_ids num_cached_tokens_per_seq = host.lists.num_cached_tokens_per_seq @@ -1321,7 +1325,6 @@ def _pack_extend_requests( 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), len(position_ids) + 1 + num_draft_tokens)) ) @@ -1361,7 +1364,6 @@ def _pack_extend_requests( 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) @@ -1418,7 +1420,6 @@ def _pack_generation_requests( generation_requests = host.lists.generation_requests draft_lens = host.lists.draft_lens sequence_lengths = host.lists.sequence_lengths - num_accepted_draft_tokens = host.lists.num_accepted_draft_tokens request_ids = host.lists.request_ids input_ids = host.lists.input_ids previous_batch_indices = host.lists.previous_batch_indices @@ -1498,7 +1499,6 @@ def _pack_generation_requests( # 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) @@ -1674,10 +1674,12 @@ def _copy_token_ids_to_device( new_tokens_lens_device: torch.Tensor | None, next_draft_tokens_device: torch.Tensor | None, runtime_draft_len: int, + *, + attn_metadata: AttentionMetadata, + enable_spec_decode: bool, ) -> None: input_ids = host.lists.input_ids draft_tokens = host.lists.draft_tokens - num_accepted_draft_tokens = host.lists.num_accepted_draft_tokens position_ids = host.lists.position_ids draft_lens = host.lists.draft_lens previous_batch_indices = host.lists.previous_batch_indices @@ -1688,11 +1690,13 @@ def _copy_token_ids_to_device( previous_batch_len = len(previous_batch_indices) def previous_seq_slots_device() -> torch.Tensor: - previous_batch_indices_host = torch.tensor( - previous_batch_indices, dtype=torch.int, pin_memory=prefer_pinned() - ) previous_slots = self._buffers.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 if scheduled_requests.is_attention_dp_phase_idle: @@ -1702,6 +1706,7 @@ def previous_seq_slots_device() -> torch.Tensor: num_tokens = len(input_ids) num_draft_tokens = len(draft_tokens) total_num_tokens = len(position_ids) + step_gather = self._step_input_gather if enable_spec_decode else None assert total_num_tokens <= self._config.max_num_tokens, ( f"total_num_tokens ({total_num_tokens}) should be less than or equal to max_num_tokens ({self._config.max_num_tokens})" # noqa: E501 ) @@ -1715,50 +1720,36 @@ def previous_seq_slots_device() -> torch.Tensor: self._buffers.draft_tokens_cuda[: len(draft_tokens)].copy_( draft_tokens, non_blocking=True ) - if self._config.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._buffers.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._buffers.previous_pos_id_offsets_cuda.zero_() - self._buffers.previous_kv_lens_offsets_cuda.zero_() runtime_tokens_per_gen_step = self.get_runtime_tokens_per_gen_step(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 host.lists.generation_requests + and scheduled_requests.num_context_requests == 0 + ): + self._buffers.previous_pos_id_offsets_cuda.zero_() + self._buffers.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._buffers.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._buffers.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, + 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() ) - # 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._buffers.previous_pos_indices_cuda[0:previous_batch_tokens].copy_( - previous_pos_indices_host, non_blocking=True - ) + self._buffers.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 @@ -1772,29 +1763,67 @@ def previous_seq_slots_device() -> torch.Tensor: # 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. + # '0' to skip the value changes in _preprocess_inputs. + # 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. + # 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._buffers.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._buffers.previous_pos_indices_cuda[0:previous_batch_tokens] - ], - non_blocking=True, - ) - - self._buffers.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_gather is None or not step_gather.gather( + new_tokens_device, + next_draft_tokens_device, + new_tokens_lens_device, + self._buffers.previous_batch_indices_cuda, + self._buffers.previous_pos_indices_cuda, + previous_batch_len, + runtime_tokens_per_gen_step, + runtime_draft_token_buffer_width, + self._buffers.input_ids_cuda, + num_tokens, + self._buffers.draft_tokens_cuda, + num_draft_tokens, + self._buffers.previous_pos_id_offsets_cuda, + previous_begin * runtime_tokens_per_gen_step, + self._buffers.previous_kv_lens_offsets_cuda, + previous_begin, + ): + # previous input ids + new_tokens = new_tokens_device.transpose(0, 1)[ + previous_slots, :runtime_tokens_per_gen_step + ].flatten() + self._buffers.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._buffers.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._buffers.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._buffers.previous_pos_indices_cuda[0:previous_batch_tokens] + ], + non_blocking=True, + ) + self._buffers.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() @@ -2111,13 +2140,10 @@ def _build_inputs( total_draft_lens = sum(host.lists.draft_lens) spec_metadata.draft_tokens = self._buffers.draft_tokens_cuda[:total_draft_lens] spec_metadata.request_ids = host.lists.request_ids - spec_metadata.gather_ids = self._buffers.gather_ids_cuda[: len(host.lists.gather_ids)] + spec_metadata.gather_ids = self._gather_ids_buffer.tensor[: len(host.lists.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.host_position_ids - spec_metadata.num_accepted_draft_tokens = self._buffers.num_accepted_draft_tokens_cuda[ - : len(host.lists.num_accepted_draft_tokens) - ] if host.context_prompt_lookahead is not None: spec_metadata.populate_context_prompt_lookahead(host.context_prompt_lookahead) # No-op for non 1-model diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 2a16877dd7ee..cf7471f1a2ec 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -432,6 +432,18 @@ def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: kv_cache_manager_key=self.kv_cache_manager_key, sparse_attention_config=self.sparse_attention_config, enable_in_graph_sampling=self._config.enable_in_graph_sampling, + # Speculative steps rewrite every input token and position, so + # captured graphs can read the persistent buffers directly. + static_input_ids=( + self._buffers.input_ids_cuda + if self._config.is_spec_decode and not self._config.use_mrope + else None + ), + static_position_ids=( + self._buffers.position_ids_cuda + if self._config.is_spec_decode and not self._config.use_mrope + else None + ), ) def _initialize_breakable_cuda_graph_runner( @@ -827,6 +839,31 @@ def _execute_logit_post_processors( 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 tensorrt_llm._torch.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/pyexecutor/engine/runners/encoder_decoder.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py index 4caa001e58de..8f15762ff3c7 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/encoder_decoder.py @@ -829,6 +829,7 @@ def _prepare_fast( non_blocking=True, ) self._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( 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 f612de0660c9..434d711b88cf 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/mamba_cache_manager.py @@ -21,6 +21,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 @@ -31,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 ( @@ -51,9 +53,12 @@ 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, 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, @@ -344,6 +349,13 @@ 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_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.""" + return False + @abstractmethod def get_conv_states(self, layer_idx: int) -> torch.Tensor: """Return conv states for specific layer. @@ -421,7 +433,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" @@ -454,6 +468,17 @@ class SpeculativeState(State): kda_g_cache: torch.Tensor | None = None kda_v_cache: torch.Tensor | None = None kda_beta_cache: torch.Tensor | None = None + # 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, 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 # prev_num_accepted_tokens: # accepted tokens (always >= 1 if drafting). @@ -472,9 +497,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_k_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: @@ -1394,6 +1421,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 @@ -1434,7 +1471,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 @@ -1495,10 +1532,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)) - 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) + 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 +2562,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 +2674,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 +2956,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,9 +3021,9 @@ def _setup_replay_buffers(self, spec_config) -> None: device): return - 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, @@ -3064,6 +3096,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_draft_records: bool = False, **kwargs, ) -> None: if conv_state_layout not in ("x_b_c", "q_k_v"): @@ -3080,6 +3113,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 records of every + # draft of the last verify round (kda_state_tok). + 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( @@ -3285,9 +3321,11 @@ def __init__( state_index_capacity = (self.max_batch_size + self._num_reserved_dummy_slots) - 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()) @@ -3467,6 +3505,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_draft_records(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: @@ -3498,6 +3540,7 @@ def _allocate_pool_replay_buffers( self.kda_g_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 @@ -3532,40 +3575,54 @@ 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() - # k / v / beta replay entries in bf16 (halves the per-draft replay - # traffic of the verify kernel); the gate cache stays fp32 because it - # feeds the cumulative-gate exponents of the replayed chunk. - self.kda_k_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - section_dim, - dtype=torch.bfloat16, - device=device, - ) - self.kda_g_cache = torch.zeros( - self.local_num_mamba_layers, - cache_size, - num_spec, - 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.bfloat16, - 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_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 + # 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, + ) + else: + # k / v / beta replay entries in bf16 (halves the per-draft replay + # traffic of the verify kernel); the gate cache stays fp32 because it + # feeds the cumulative-gate exponents of the replayed chunk. + self.kda_k_cache = torch.zeros( + self.local_num_mamba_layers, + cache_size, + num_spec, + section_dim, + dtype=torch.bfloat16, + device=device, + ) + self.kda_g_cache = torch.zeros( + self.local_num_mamba_layers, + cache_size, + num_spec, + 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.bfloat16, + 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") + f"{cache_size} state slots{per_token}") return True def _commit_gdn_cached_replay_history_layers( @@ -3858,9 +3915,25 @@ def _get_pool_roles(self, def _max_resident_sequences(self) -> int: return self.max_batch_size * self.mapping.pp_size - def _mamba_state_bytes_per_slot(self) -> int: - base_bytes = self.local_num_mamba_layers * (self.ssm_bytes + - self.conv_bytes) + 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_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_draft_records", False): + return 0 + return (self.local_num_mamba_layers * + math.prod(self._kda_draft_records_shape()) * + torch.float32.itemsize) + + def _managed_state_bytes_per_slot(self) -> int: + base_bytes = (self.local_num_mamba_layers * + (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: @@ -3871,6 +3944,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 [] @@ -3956,6 +4033,16 @@ 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, or for every slot + 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_draft_record_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) num_request_lineages = self._max_resident_sequences() @@ -3963,7 +4050,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. @@ -3999,12 +4086,62 @@ 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, ) + 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 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_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_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 + 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 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 + or not self._ssm_pool_at_live_floor(kv_cache_config)): + return 0 + floor_slots = (self._max_resident_sequences() + + self._num_reserved_dummy_slots) + # 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 + 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(ceil_div(gpu_quota, bytes_per_token), 1 << 30) + def _build_cache_config( self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: kv_cache_config = self.kv_cache_config @@ -4020,14 +4157,47 @@ 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_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))): + # 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( + "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. + draft_record_slots = self._live_floor_draft_record_slots( + config.max_util_for_resume) + 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=pool_gpu_quota), *cache_tiers[1:] + ] + logger.info( + "KV cache manager v2: the KDA draft records of " + 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. @@ -4062,6 +4232,10 @@ def _build_cache_config( if config.initial_pool_ratio is None: typical_capacity = self._get_typical_request_capacity( kv_cache_config) + if ssm_pool_at_live_floor: + typical_capacity = max( + typical_capacity, + self._quota_filling_request_capacity(pool_gpu_quota)) request_descs = self._typical_request_descs(typical_capacity, kv_cache_config) typical_step = BatchDesc(request_descs * @@ -4088,6 +4262,7 @@ def _build_cache_config( ] return replace( config, + cache_tiers=cache_tiers, layers=layers, typical_step=typical_step, constraints=constraints, @@ -4227,9 +4402,9 @@ def _setup_replay_buffers(self, spec_config) -> None: return mask_capacity = self._host_state_indices.shape[0] - 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, @@ -4353,11 +4528,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_k_cache, self.kda_g_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_k_cache, self.kda_g_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], @@ -4439,8 +4620,9 @@ 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, 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, ), @@ -4450,16 +4632,20 @@ 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)): + k_cache, g_cache, v_cache, beta_cache = ( + None if replay_cache is None else replay_cache[layer_offset] + for replay_cache in (self.kda_k_cache, self.kda_g_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_k_cache[layer_offset], - self.kda_g_cache[layer_offset], - self.kda_v_cache[layer_offset], - self.kda_beta_cache[layer_offset], + k_cache, + g_cache, + v_cache, + beta_cache, slots, num_accepted, state_tok=(state_tok[layer_offset] @@ -4567,10 +4753,7 @@ 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) + 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() @@ -4672,10 +4855,12 @@ def _seed_kda_replay_slots(self, state_indices: torch.Tensor) -> None: 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_k_cache, self.kda_g_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.index_fill_(0, state_indices, 0) def mamba_layer_cache( @@ -4690,20 +4875,23 @@ 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_k_cache is not None - assert self.kda_g_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_k_cache": self.kda_k_cache[layer_offset], - "kda_g_cache": self.kda_g_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_k_cache is not None + assert self.kda_g_cache is not None + assert self.kda_v_cache is not None + assert self.kda_beta_cache is not None + spec_kwargs["kda_k_cache"] = self.kda_k_cache[layer_offset] + spec_kwargs["kda_g_cache"] = self.kda_g_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( @@ -4976,7 +5164,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 @@ -4992,5 +5180,6 @@ def shutdown(self): self.kda_beta_cache = None self._kda_replay_suspended.clear() self._kda_replay_acceptance_pending.clear() + self.kda_state_tok = None self._kda_replay_unseeded.clear() super().shutdown() diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a582900ec098..56297bd52a41 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -73,7 +73,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 @@ -5950,7 +5950,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/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 88da4911cca0..659ef04c34da 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -21,6 +21,7 @@ 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 prefer_pinned from tensorrt_llm.logger import logger @@ -539,11 +540,11 @@ class DFlashSpecMetadata(SpecMetadata): captured_hidden_states: Optional[torch.Tensor] = None def __post_init__(self): - 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 @@ -584,10 +585,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) + 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) @@ -625,10 +623,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) + worker._batch_to_slot_buffer.update(mapping) def is_layer_capture(self, layer_id: int) -> bool: return layer_id in self._capture_layer_set @@ -730,7 +726,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 @@ -1114,7 +1110,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 - 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 = {} @@ -1417,6 +1416,10 @@ def _advance_ctx_len(self, slots: torch.Tensor, num_accepted: torch.Tensor) -> N self._ctx_len.clamp_(max=self._max_ctx) self._ctx_len[self._dummy_slot].zero_() + def _rebuild_batch_to_slot(self, request_ids: List[int]) -> None: + mapping = [self._req_to_slot.get(rid, self._dummy_slot) for rid in request_ids] + 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.""" if not updates: @@ -1712,17 +1715,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, @@ -2090,7 +2083,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/interface.py b/tensorrt_llm/_torch/speculative/interface.py index 8edccc3e3f8a..a97d28355c5c 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -489,8 +489,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 @@ -886,6 +884,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) @@ -1703,6 +1706,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/tensorrt_llm/_torch/speculative/spec_sampler_base.py b/tensorrt_llm/_torch/speculative/spec_sampler_base.py index 5fb5b185dc0b..dc0c69f7d909 100644 --- a/tensorrt_llm/_torch/speculative/spec_sampler_base.py +++ b/tensorrt_llm/_torch/speculative/spec_sampler_base.py @@ -27,7 +27,10 @@ import torch +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 from ..pyexecutor.resource_manager import BaseResourceManager from ..pyexecutor.sampler import ( @@ -304,6 +307,17 @@ 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 + self._store_scatter: Optional[spec_step_copies.SlotScatter] = None + self._slot_table_buffer: Optional[DeviceCopyBuffer] = None + if spec_step_copies.is_supported(): + self._store_scatter = spec_step_copies.SlotScatter() + 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 self.store = self.Store( new_tokens=int_tensor((self.max_accepted_path_len, seq_slots, self.max_beam_width)), next_new_tokens=int_tensor( @@ -396,6 +410,52 @@ def update_requests( req.py_rewind_len = verified_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, @@ -425,7 +485,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 @@ -442,17 +501,7 @@ 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] - + num_sampling_requests = len(sampling_requests) o_verify_lens = outputs.get("verify_lens") native_uniform = outputs.get("native_uniform_verify", False) host_snapshot = outputs.get("host_policy_windows_snapshot", False) @@ -489,36 +538,27 @@ def sample_async( raise ValueError("verify_lens must contain one int32 token window per output row") o_verify_lens = o_verify_lens[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]) + 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. + self._slot_table_buffer.update(torch.tensor(slots, dtype=torch.int32)) + scattered = self._store_scatter.scatter( + outputs, + num_skip, + len(slots), + self._slot_table_buffer.tensor, + self.store.new_tokens, + self.store.next_new_tokens, + self.store.new_tokens_lens, + self.store.next_draft_tokens, ) - 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( @@ -527,13 +567,33 @@ 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), - verify_lens=(self._copy_to_host(o_verify_lens) if o_verify_lens is not None else None), - ) - 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), + 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: + # 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. + 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), + 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), + verify_lens=host_verify_lens, + ) + 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/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 4d3b3b48de4d..c5b7ebb35770 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -3408,14 +3408,17 @@ 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 " + "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"] 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/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index ce332dfae73a..9e012038cc84 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -5715,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 e34b7f2adaa7..8aa5fe63300c 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -7265,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/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 0061bf2389ad..ce8b28893d4f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -266,6 +266,10 @@ l0_b200: - unittest/_torch/speculative/test_chain_speculative_sampling_op.py - unittest/_torch/speculative/test_dynamic_tree_rejection_op.py - unittest/_torch/speculative/test_dynamic_tree_sampling_dispatch.py + # One-model speculative decoding step-copy kernels and the engine's decode + # 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) - 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..3edb55c42aef --- /dev/null +++ b/tests/unittest/_torch/cute_dsl_kernels/test_spec_step_copies.py @@ -0,0 +1,655 @@ +# 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, 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. +* ``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 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). +""" + +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 + +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 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(args) -> None: + """``StepInputGather.gather`` on ``args``, which it must cover.""" + assert _op().StepInputGather().gather(*args), "gather declined a covered call" + + +# ---------------------------------------------------------------------------------------------------------------- +# 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]) + 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) + 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 (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] + 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 + ) + + +# ---------------------------------------------------------------------------------------------------------------- +# 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 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 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). + """ + + 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, + ): + 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 = 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 + shapes = { + "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,), + "input_ids": (MAX_NUM_TOKENS,), + "draft": (n_draft,), + "pos_off": (MAX_NUM_TOKENS,), + "kv_off": (n_kv,), + } + 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, 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(): + 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) + 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) + + 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): + """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["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.""" + return {k: t.clone() for k, t in bufs.items()} + + +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 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): + 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)), + 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 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 + + def run_kernels(st, b): + kernel_scatter(st.scatter_args(b)) + kernel_gather(st.gather_args(b)) + + def run_torch(st, b): + torch_scatter(*st.scatter_args(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)) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + for st in steps: + run_kernels(st, st.bufs) + torch.cuda.synchronize() + out = [] + before = None + for rep in range(replays): + for st in steps: + st.rewrite() + 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) + 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}: 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(): + 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("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 + + +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.""" + 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)) + assert op.SlotScatter().scatter(*scatter[:2], 0, *scatter[3:]) + 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) + # 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 + 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:]) + strided = dict( + scatter[0], next_draft_tokens=torch.empty_like(b["out_draft"]).t().contiguous().t() + ) + assert not op.SlotScatter().scatter(strided, *scatter[1:]) + 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), + ("StepInputGather.gather", measure_gather), + ] + 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() + stream.wait_stream(torch.cuda.current_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) + + +def timing() -> None: + calls = 32 + 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 |") + 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 + 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)), + lambda i: kernel_gather(st.gather_args(b)), + ] + 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})") + 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/kv_cache/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager.py index 0387894bf82c..b0cd2508a088 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, @@ -78,6 +79,7 @@ ) from tensorrt_llm.mapping import Mapping from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + GPU_LEVEL, AttentionLayerConfig, BatchDesc, BufferConfig, @@ -152,6 +154,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, ) @@ -309,6 +313,7 @@ def _kimi_model_config() -> SimpleNamespace: 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.""" @@ -327,7 +332,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, @@ -383,6 +388,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_draft_records" not in kwargs + + +@pytest.mark.parametrize("backbone_asks", [True, False]) +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_draft_records``.""" + monkeypatch.setattr( + "tensorrt_llm._torch.modules.kimi_kda._kda_kernels.is_kda_mtp_verify_available", + lambda: True, + ) + backbone = SimpleNamespace(kda_draft_records=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_draft_records", False) is backbone_asks @pytest.mark.parametrize("max_draft_len, replay_enabled", [(7, True), (8, False)]) @@ -1468,7 +1500,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 @@ -2130,11 +2162,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() @@ -2425,6 +2457,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_draft_records=False, ): """Construct a real MambaHybridCacheManagerV2.""" mamba_mask = [True] * num_mamba_layers + [False] * num_attention_layers @@ -2480,6 +2513,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_draft_records=kda_draft_records, ) @@ -3706,6 +3740,56 @@ 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_draft_records + 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_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( + 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_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_draft_records + return + assert mgr.keeps_kda_draft_records + assert mgr.kda_k_cache is None + assert mgr.kda_g_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_k_cache is None + assert layer_cache.kda_g_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) + 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() @@ -3753,6 +3837,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, @@ -3772,6 +3857,190 @@ def test_v2_kda_replay_refuses_kv_pool_rebalance(kda_replay): ) +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) + 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_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 + 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, + ) + 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)], +) +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_draft_record_manager(kda_draft_records, enable_block_reuse) + 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) + pool_bytes = runtime_manager.get_quota(GPU_LEVEL) + 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 + # 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 + * mgr.ssm_state_shape[0] + * mgr.ssm_state_shape[-1] + * 4 + ) + # 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 + + +@pytest.mark.parametrize("off_the_floor", ["block_reuse", "pool_ratio", "no_attention"]) +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_draft_record_manager( + kda_draft_records=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_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_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 + # 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 + 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 + extra_slot_states + ) + + 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_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_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( + 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 * 3 * 8 * 256 * 4 + minimum = mgr._minimum_live_gpu_quota() + # 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( + 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) @@ -3785,7 +4054,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), @@ -3800,7 +4069,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._host_state_indices = 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 @@ -3885,8 +4154,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 @@ -4240,6 +4509,80 @@ def test_v2_kda_replay_suspend_resume_matches_an_unsuspended_run(slot_taken): torch.testing.assert_close(states, expected_states, atol=1e-3, rtol=1e-3) +@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_k_cache = None + mgr.kda_g_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( @@ -4281,6 +4624,76 @@ def test_v2_kda_replay_relocates_live_slot_history(): mgr.shutdown() +@skip_no_cuda +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), + conv_state_layout="q_k_v", + mamba_n_groups=4, + mamba_ssm_cache_dtype=torch.float32, + kda_replay_num_spec=2, + kda_draft_records=True, + ) + try: + assert mgr.kda_k_cache is None + assert mgr.kda_g_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]) + + 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. + assert mgr.kda_state_tok is None + + +@skip_no_cuda +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( + 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_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() + + @skip_no_cuda def test_v2_kda_state_index_setup_relocates_generation_history(): class StateCache: @@ -4337,6 +4750,57 @@ def close(self) -> None: mgr.shutdown() +@skip_no_cuda +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._device_copy._verify_skipped_device_copies", False) + + 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 @@ -4385,6 +4849,76 @@ def dim_contiguous_conv_cache() -> torch.Tensor: assert mgr.prev_num_accepted_tokens[2] == 5 +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) + 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] + 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_k_cache = None + mgr.kda_g_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_k_cache=t, + kda_g_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. @@ -4614,7 +5148,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) @@ -4654,7 +5188,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 @@ -5135,7 +5669,9 @@ def test_v2_kda_admission_slot_updates_without_sync(device): mgr.local_num_mamba_layers = num_layers mgr.prev_num_accepted_tokens = torch.arange(1, num_slots + 1, dtype=torch.int32, device=device) # Stale device indices: the reset must take its slots from _host_state_indices. - mgr.cuda_state_indices = torch.full((num_slots,), 7, dtype=torch.int32, device=device) + mgr._state_indices_buffer = DeviceCopyBuffer( + torch.full((num_slots,), 7, dtype=torch.int32, device=device) + ) mgr._host_state_indices = torch.tensor([5, 1, 0, 0, 0, 0, 0, 0], dtype=torch.int32) mgr._request_id_to_state_index = {10: 4, 11: 2} mgr._kda_replay_unseeded = set() 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 b9f34f6bc2b4..8f4e4010db1b 100644 --- a/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py +++ b/tests/unittest/_torch/executor/test_cuda_graph_capture_replay.py @@ -332,6 +332,85 @@ def forward_fn(inputs): 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) + 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), + ) + + class _MetadataStub: """Metadata stand-in with one graph-visible CUDA tensor; the strict check walks vars().""" @@ -395,38 +474,45 @@ def test_replay_accepts_in_place_update_to_graph_tensor(self, monkeypatch): assert output.item() == 99 @pytest.mark.parametrize( - "owner, to_none, in_postprocess", + "owner, to_none", [ - ("attn_metadata", False, False), - ("attn_metadata", True, False), - ("spec_metadata", False, False), - ("attn_metadata", False, True), + ("attn_metadata", False), + ("attn_metadata", True), + ("spec_metadata", False), ], - ids=["rebound_tensor", "non_tensor", "spec_metadata", "postprocess_fn"], + ids=["rebound_tensor", "non_tensor", "spec_metadata"], ) - def test_replay_rejects_rebound_graph_attr(self, monkeypatch, owner, to_none, in_postprocess): + def test_replay_rejects_rebound_graph_attr(self, monkeypatch, owner, to_none): """Rebinding a graph-visible tensor attribute raises, naming the attribute.""" runner, key, _, inputs = self._make_runner_and_inputs(monkeypatch, value=10) if owner == "spec_metadata": inputs["spec_metadata"] = _MetadataStub(value=20) - def rebind(fn_inputs): - fn_inputs[owner].some_buf = ( - None if to_none else torch.full((1,), 99, device="cuda", dtype=torch.int32) - ) - - runner.capture( - key, - self._forward_reading_some_buf, - inputs, - postprocess_fn=rebind if in_postprocess else None, + runner.capture(key, self._forward_reading_some_buf, inputs) + inputs[owner].some_buf = ( + None if to_none else torch.full((1,), 99, device="cuda", dtype=torch.int32) ) - if not in_postprocess: - rebind(inputs) with pytest.raises(RuntimeError, match="some_buf"): runner.replay(key, inputs) + def test_postprocess_fn_follows_only_the_warmup_forwards(self, monkeypatch): + """capture() calls postprocess_fn after each warmup forward, not after the captured one, so a rebind in it is + what the graph bakes in and the strict check accepts it.""" + runner, key, _, inputs = self._make_runner_and_inputs(monkeypatch, value=10) + calls = [] + + def rebind(fn_inputs): + calls.append(1) + fn_inputs["attn_metadata"].some_buf = torch.full( + (1,), 99, device="cuda", dtype=torch.int32 + ) + + runner.capture(key, self._forward_reading_some_buf, inputs, postprocess_fn=rebind) + + assert len(calls) == runner.WARMUP_STEPS + assert runner.replay(key, inputs).item() == 99 + def test_replay_accepts_swapped_attr_with_unchanged_saved_target(self, monkeypatch): """A swapped attr is validated via its saved target; clearing the swap flags it.""" runner, key, attn_metadata, inputs = self._make_runner_and_inputs(monkeypatch, value=10) 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/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index 3a7293bd7e9d..d6ab99225820 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -78,7 +78,7 @@ def _input_preparer() -> DecoderInputPreparer: return DecoderInputPreparer(buffers=None, - config=None, + config=SimpleNamespace(is_spec_decode=False), mapping=None, dist=None, model=None, @@ -2312,6 +2312,114 @@ def test_promoted_context_precedes_speculative_overlap_generation( [generation.py_seq_slot], 0) kv_cache_manager.shutdown() + 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 + 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)) + inputs = model_engine._runner._inputs + buffers = model_engine._runner._buffers + step_gather = inputs._step_input_gather + self.assertIsNotNone(step_gather) + 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._runner._metadata.setup_attention( + 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_gather): + """The buffers one step writes, from a filler they all start at.""" + inputs._step_input_gather = step_input_gather + for request, slot in zip(requests, slots): + request.py_batch_idx = slot + written = (buffers.input_ids_cuda, buffers.position_ids_cuda, + buffers.draft_tokens_cuda, + buffers.previous_pos_id_offsets_cuda, + buffers.previous_kv_lens_offsets_cuda) + for buffer in written: + buffer.fill_(-3) + inputs.prepare_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, + enable_spec_decode=True, + runtime_draft_len=max_draft_len, + is_dummy=False) + torch.cuda.synchronize() + return [buffer.clone() for buffer in written] + + torch_path = step_inputs(None) + 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(kernel_path[0][:num_tokens].tolist(), + expected_input_ids.tolist()) + kv_cache_manager.shutdown() + 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/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 9f2f2351a736..aa936b063399 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 @@ -1255,10 +1255,10 @@ def test_commit_pending_drafts_matches_the_replay_reference(B, H, M): 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) @@ -1273,7 +1273,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)): 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 48e4c29b8ca5..5a06177217bb 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 @@ -129,10 +130,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 56fb548b8793..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, @@ -59,6 +60,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. @@ -138,6 +163,107 @@ 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, 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._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) + + 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_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_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._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) + + 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_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 new file mode 100644 index 000000000000..f88ba9687763 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_copy_skip.py @@ -0,0 +1,105 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +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 + +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_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. + worker._assign_slot = lambda rid: worker._req_to_slot.get(rid) + return worker + + +def _metadata(worker, request_ids, num_generations): + 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): + 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_buffer.tensor[:2].tolist() == [DUMMY_SLOT, 1] + worker._req_to_slot[8] = 4 + _rebuild(worker, [8, 5]) + 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_buffer.tensor[: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_buffer.tensor[:2].tolist() == slots + _prepare(metadata, request_ids, num_generations=2) + assert worker._batch_to_slot_buffer.tensor[: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_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_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)) 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 0249d62b4efa..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 @@ -610,6 +610,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_buffer = None + sampler._store_copies_done = None return sampler, transfers 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 93b3f4ce223a..2e39163ef844 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 @@ -247,3 +247,41 @@ def test_block_rows_take_consecutive_slots() -> None: assert offsets.tolist() == [ base + (4 + row) * _RNG_SLOT_SPAN for base in (1000, BASE) for row in range(3) ] + + +# --- all-greedy batches ------------------------------------------------------ + +ROWS = MAX_DRAFT_LEN + 1 + + +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, ROWS) 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 * ROWS) + 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 * ROWS,), sentinel, dtype=torch.int64) + meta.offsets = torch.full((8 * ROWS,), 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, BASE + WINDOW] + assert meta.request_seeds[:2].tolist() == [7, DEFAULT_SAMPLING_SEED] + assert meta.offsets[: 2 * ROWS].tolist() == [ + base + row * _RNG_SLOT_SPAN for base in (WINDOW, BASE + WINDOW) for row in range(ROWS) + ] 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..e5a1e091e937 --- /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.engine.runners.decoder import DecoderRunner +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__(DecoderRunner) + 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" diff --git a/tests/unittest/conftest.py b/tests/unittest/conftest.py index a94b107b461f..fe4538843d23 100644 --- a/tests/unittest/conftest.py +++ b/tests/unittest/conftest.py @@ -390,6 +390,12 @@ def torch_empty_cache() -> None: torch.cuda.empty_cache() +@pytest.fixture(autouse=True) +def verify_skipped_device_copies(monkeypatch) -> None: + monkeypatch.setattr( + "tensorrt_llm._device_copy._verify_skipped_device_copies", True) + + @pytest.fixture(scope="module") def mpi_pool_executor(request): """