[None][perf] One-model speculative decoding: host work between steps, worker-produced target logits, KDA drafts' records

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.

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.

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

tensorrt_llm/_torch/attention/backends/trtllm.py

@@ -1415,6 +1415,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,

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.
+"""

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