[TRTLLM-14814][feat] Kimi K3 serving parsers, chat template, and speculative decoding (suffix automaton + DFlash scaffold)

docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md

@@ -224,8 +224,6 @@ When the `Status: 200` code is returned, the server is ready for queries. Note t
 
 ### Basic Test
 
-> **Note:** The `/v1/chat/completions` endpoint requires the Kimi K3 chat template and serving parsers, which are being added in [TRTLLM-14814](https://github.com/NVIDIA/TensorRT-LLM/pull/17327). Until that change lands, use the `/v1/completions` endpoint with a plain `prompt` string instead.
-
 After the TensorRT LLM server is set up and shows `Application startup complete`, you can send requests to the server:
 
 ```shell

examples/kimi_k3/README.md

@@ -192,5 +192,5 @@ default cache manager.
   and the TEP16/TEP8 latency recipes are unaffected. Tracked as
   TRTLLM-14904.
 - FP8 KV cache (`kv_cache_config.dtype: fp8`) is not yet supported.
-- Speculative decoding is not yet supported.
+- Speculative decoding: suffix-automaton speculation is supported for aggregated serving (`speculative_config: {decoding_type: SA}` in the extra LLM API options). Combining speculation with disaggregated serving is not yet supported.
 - Disaggregated serving is not yet supported.

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1597,17 +1597,51 @@ def forward(
             else:
                 forward_args.fmha_scheduler_counter.zero_()
             assert forward_args.latent_cache is not None
-            from .utils import append_mla_latent_cache
-            append_mla_latent_cache(
-                metadata.kv_cache_manager,
-                self.get_local_layer_idx(metadata),
-                metadata.request_ids,
-                metadata.seq_lens.tolist(),
-                metadata.kv_cache_params.num_cached_tokens_per_seq,
-                forward_args.latent_cache,
-                kv_layout=metadata.kv_layout,
-                seq_start=num_ctx,
-            )
+            from ..pyexecutor.mamba_cache_manager import BaseMambaCacheManager
+
+            # Hybrid (mamba/masked-layer) KV managers take the graph-safe
+            # append; the same predicate interface.py uses to detect hybrid
+            # managers for mamba metadata. Dense MLA models keep the
+            # host-side path unchanged.
+            if isinstance(metadata.kv_cache_manager, BaseMambaCacheManager):
+                from .utils import \
+                    append_mla_latent_cache_generation_cuda_graph_safe
+
+                # The write positions must come from device tensors: the host
+                # lists (request ids, seq lens, cached-token counts) are
+                # frozen into the graph at capture time and corrupt the cache
+                # on replay. The helper falls back to the host-side loop for
+                # eager forwards only; under CUDA graphs it scatters
+                # device-side for any uniform q_len (1 for plain decode,
+                # 1 + draft_len for padded spec-dec verification batches).
+                append_mla_latent_cache_generation_cuda_graph_safe(
+                    metadata,
+                    # NOTE: get_buffers / layer_offsets take the GLOBAL layer
+                    # index (they map through layer_offsets internally).
+                    # Passing the local offset double-maps and breaks hybrid
+                    # models whose KV manager covers a masked layer subset
+                    # (e.g. Kimi K3).
+                    self.layer_idx,
+                    forward_args.latent_cache,
+                )
+            else:
+                # TODO(TRTLLM-15193): this host-side path freezes write
+                # positions at CUDA graph capture time and may corrupt the
+                # cache on replay IF a non-hybrid MLA model reaches this path
+                # under graph capture. Investigate whether that is reachable;
+                # if so, the graph-safe branch above is a correctness fix to
+                # generalize, not an optimization.
+                from .utils import append_mla_latent_cache
+                append_mla_latent_cache(
+                    metadata.kv_cache_manager,
+                    self.get_local_layer_idx(metadata),
+                    metadata.request_ids,
+                    metadata.seq_lens.tolist(),
+                    metadata.kv_cache_params.num_cached_tokens_per_seq,
+                    forward_args.latent_cache,
+                    kv_layout=metadata.kv_layout,
+                    seq_start=num_ctx,
+                )
 
         forward_args.sparse_runtime_params = prepare_sparse_runtime_params(
             self, q, k, metadata, forward_args)

tensorrt_llm/_torch/attention_backend/utils.py

@@ -109,6 +109,112 @@ def create_attention(
     )
 
 
+def append_mla_latent_cache_generation_cuda_graph_safe(
+    metadata,
+    layer_idx: int,
+    latent_cache: torch.Tensor,
+) -> None:
+    """Append generation-phase MLA latent tokens, safe under CUDA graphs.
+
+    :func:`append_mla_latent_cache` computes every write location on the host
+    (request ids, per-request block lists, cached-token counts), so a CUDA
+    graph captures its copy kernels with those positions frozen and replays
+    them against stale slots, silently corrupting the cache. This variant
+    derives the destination purely from device tensors living in graph-stable
+    buffers that ``metadata.prepare()`` refreshes every step:
+
+    - ``kv_lens_cuda_runtime`` holds each request's total KV length (cached +
+      new), so the ``q_len`` new tokens sit at positions
+      ``kv_len - q_len .. kv_len - 1`` (clamped to 0 so graph-warmup passes
+      with zeroed lengths stay in bounds).
+    - ``kv_cache_block_offsets[pool_idx, slot, 0]`` is the C++
+      ``setOffsets``-encoded block table: entries hold
+      ``pool_block_index * num_pool_layers * kv_factor`` (+ the K/V field
+      index, always 0 for the single-plane MLA cache), so the raw block index
+      for the per-layer ``get_buffers`` view is recovered by integer division.
+
+    Handles any uniform ``q_len >= 1`` per generation request: plain decode
+    (``q_len == 1``) and speculative-verification batches (``q_len ==
+    1 + draft_len``; the spec workers pad drafts to the static max, so the
+    per-request token count is uniform). This matters for spec-dec under
+    CUDA graphs: the previous ``q_len == 1``-only version silently fell back
+    to the host-side loop for verification batches, whose capture-time write
+    positions (dummy-request block tables) were frozen into the graph — real
+    requests' generation-token latents were never appended on replay,
+    corrupting decode accuracy (Kimi K3 SA GSM8K 88.2 with graphs vs 96.2
+    eager).
+
+    Falls back to the host-side loop for eager forwards (numerics identical
+    to the non-graph baseline) and for ragged generation batches, which
+    cannot occur under CUDA graphs.
+    """
+    kv_cache_manager = metadata.kv_cache_manager
+    num_ctx = metadata.num_contexts
+    n_gen = metadata.num_generations
+    # Tensor shapes are static under CUDA graphs, so this host-side check is
+    # stable across replays: generation-only graph batches carry a uniform
+    # per-request token count (1 for plain decode, 1 + draft_len for
+    # padded speculative verification).
+    q_len_is_uniform = n_gen > 0 and latent_cache.shape[0] % n_gen == 0
+    if not metadata.is_cuda_graph or not q_len_is_uniform:
+        append_mla_latent_cache(
+            kv_cache_manager,
+            layer_idx,
+            metadata.request_ids,
+            metadata.seq_lens.tolist(),
+            metadata.kv_cache_params.num_cached_tokens_per_seq,
+            latent_cache,
+            kv_layout=metadata.kv_layout,
+            seq_start=num_ctx,
+        )
+        return
+
+    kv_layout = metadata.kv_layout
+    kv_cache = kv_cache_manager.get_buffers(layer_idx, kv_layout=kv_layout)
+
+    # Static per-layer facts: plain ints baked into the kernel launches, and
+    # they never change between replays. kv_cache_pool_mapping exists on both
+    # V1 and V2 managers, including hybrid subclasses whose KV manager covers
+    # a masked layer subset (layer_offsets maps the global layer index).
+    layer_offset = kv_cache_manager.layer_offsets[layer_idx]
+    pool_mapping = kv_cache_manager.kv_cache_pool_mapping
+    pool_idx = int(pool_mapping[layer_offset, 0])
+    num_pool_layers = int((pool_mapping[:, 0] == pool_idx).sum())
+    kv_factor = kv_cache_manager.kv_factor
+    tokens_per_block = kv_cache_manager.tokens_per_block
+
+    # Everything below only reads graph-stable device buffers. ``q_len`` is
+    # derived from static tensor shapes, so it is a stable host constant per
+    # captured graph (1 for plain decode, 1 + draft_len for spec verify).
+    q_len = latent_cache.shape[0] // n_gen
+    kv_lens = metadata.kv_lens_cuda_runtime[num_ctx:num_ctx + n_gen]
+    # ``kv_len`` includes the new tokens, so they occupy positions
+    # ``kv_len - q_len .. kv_len - 1``. pos: [n_gen, q_len].
+    pos = ((kv_lens.to(torch.int64) - q_len).clamp_(min=0).unsqueeze(1) +
+           torch.arange(q_len, dtype=torch.int64, device=kv_lens.device))
+    block_slot = pos // tokens_per_block
+    block_offset = pos % tokens_per_block
+    # [num_pools, max_num_sequences, 2, max_blocks_per_seq]; the two K/V
+    # entries are identical for the kv_factor=1 MLA cache, take field 0.
+    block_table = metadata.kv_cache_block_offsets[pool_idx,
+                                                  num_ctx:num_ctx + n_gen, 0]
+    encoded = block_table.gather(1, block_slot)  # [n_gen, q_len]
+    # Placeholder entries are negative; clamp so warmup rows stay in bounds.
+    # TODO(TRTLLM-15199): clamping to block 0 means a padded/warmup row
+    # scatters into a real request's block 0. Exclude invalid rows (or
+    # reserve a scratch block) instead of clamping.
+    dest_block = encoded.to(
+        torch.int64).clamp_(min=0) // (num_pool_layers * kv_factor)
+    src = latent_cache.to(kv_cache.dtype).reshape(n_gen, q_len,
+                                                  latent_cache.shape[-1])
+    if kv_layout == "NHD":
+        kv_cache[dest_block, 0, block_offset, 0, :] = src
+    elif kv_layout == "HND":
+        kv_cache[dest_block, 0, 0, block_offset, :] = src
+    else:
+        raise ValueError(f"Unsupported kv_layout: {kv_layout}")
+
+
 def append_mla_latent_cache(
     kv_cache_manager,
     layer_idx: int,