diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index dfa4daac7549..af3d60c9e0e2 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/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 diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index 800631bbd5d5..00349c6d6d43 100644 --- a/examples/kimi_k3/README.md +++ b/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. diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py index 4ac52386bfcc..8be8e9e83db9 100644 --- a/tensorrt_llm/_torch/attention_backend/trtllm.py +++ b/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) diff --git a/tensorrt_llm/_torch/attention_backend/utils.py b/tensorrt_llm/_torch/attention_backend/utils.py index ef83c99159ed..7bf3c6c305e7 100644 --- a/tensorrt_llm/_torch/attention_backend/utils.py +++ b/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, diff --git a/tensorrt_llm/_torch/models/modeling_speculative.py b/tensorrt_llm/_torch/models/modeling_speculative.py index a991e266965d..12edaf21eb1d 100755 --- a/tensorrt_llm/_torch/models/modeling_speculative.py +++ b/tensorrt_llm/_torch/models/modeling_speculative.py @@ -842,6 +842,91 @@ def forward( return hidden_states_out, hidden_states_out +def dspark_layer_window_size(use_swa: bool, swa_window: int, layer_types, + layer_idx: int) -> tuple[int, int]: + """flash-attn ``window_size`` for one draft layer of the block decode. + + DSpark drafters (deepseek-ai/DeepSpec) run the draft block through HF + attention with ``sliding_window`` set on 'sliding_attention' layers and + is_causal=False. HF's flash path + (transformers/modeling_flash_attention_utils.py) translates that to + ``window_size = (sliding_window - 1, sliding_window - 1)``, i.e. each + query attends keys within ``swa_window - 1`` KV-index distance on both + sides. In the DFlash pool layout KV index == token position, so this + limits draft queries to the most recent ``swa_window`` context tokens + plus the (nearby) draft block. Full-attention layers and non-dspark + drafters keep flash-attn's default ``(-1, -1)`` (no window). + """ + if not use_swa: + return (-1, -1) + if layer_types is not None and layer_idx < len(layer_types) and \ + layer_types[layer_idx] != 'sliding_attention': + return (-1, -1) + return (swa_window - 1, swa_window - 1) + + +def dspark_markov_step_bias(prev_tokens: torch.Tensor, markov_w1: torch.Tensor, + markov_w2: torch.Tensor) -> torch.Tensor: + """Vanilla Markov head logit bias for one intra-block draft step. + + Reference: DeepSpec ``VanillaMarkov`` (deepspec/modeling/dspark/ + markov_head.py): ``bias = markov_w2(markov_w1(prev_token))`` where + markov_w1 is nn.Embedding(vocab, rank) and markov_w2 is + nn.Linear(rank, vocab, bias=False). With both weights stored + [vocab, rank] this is ``markov_w1[prev] @ markov_w2.T``. + + Args: + prev_tokens: [B] long, previous token per request (draft vocab). + markov_w1: [vocab, rank]. + markov_w2: [vocab_or_shard, rank] (rows may be a TP vocab shard). + Returns: + [B, vocab_or_shard] bias in the markov weights' dtype. + """ + return F.linear(F.embedding(prev_tokens, markov_w1), markov_w2) + + +def dspark_markov_chain_logits( + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + argmax_fn=None, +) -> torch.Tensor: + """Apply the vanilla Markov intra-block bias across a drafted block. + + Reference: DeepSpec ``VanillaMarkov.sample_block_tokens`` at + temperature 0: for step i, ``logits_i += bias(prev_i)`` with + ``prev_0`` = the anchor token (last accepted token, block slot 0) and + ``prev_{i>0}`` = the greedy token from step i-1's *biased* logits. + llama.cpp PR #25173 implements the same greedy chain. + + Args: + base_logits: [B, K, vocab_or_shard] shared-lm_head logits. + first_prev_tokens: [B] long, anchor token ids (draft vocab). + markov_w1 / markov_w2: see :func:`dspark_markov_step_bias`. + argmax_fn: callable([B, vocab_or_shard]) -> [B] token ids in the + full draft vocab; defaults to plain argmax. Workers pass a + TP-aware argmax when the draft logits are vocab-sharded. + Returns: + [B, K, vocab_or_shard] biased logits. Greedy per-position argmax of + the result reproduces the reference sampled chain exactly. + """ + K = base_logits.shape[1] + if K == 0: + return base_logits + prev = first_prev_tokens.long() + steps = [] + for i in range(K): + bias = dspark_markov_step_bias(prev, markov_w1, markov_w2) + step_logits = base_logits[:, i] + bias.to(base_logits.dtype) + steps.append(step_logits) + if argmax_fn is not None: + prev = argmax_fn(step_logits).long() + else: + prev = torch.argmax(step_logits, dim=-1) + return torch.stack(steps, dim=1) + + class DFlashForCausalLM(nn.Module): """Draft model wrapper for DFlash speculative decoding. @@ -904,6 +989,75 @@ def __init__(self, draft_config): f"target_layer_ids: {self.target_layer_ids}, block_size: {self.block_size}" ) + # DSpark drafters (DFlash + low-rank Markov head + confidence head, + # arXiv 2607.05147; reference: deepseek-ai/DeepSpec). The weights- + # independent drafter-forward semantics ARE implemented here: + # - vanilla Markov intra-block logit bias (applied by DFlashWorker + # through apply_markov_chain_logits), + # - sliding-window attention on 'sliding_attention' draft layers + # during the block decode (use_swa / swa_window_size), + # - the shift_label output convention (hidden state at block slot j + # predicts draft token j+1; slot 0 holds the anchor token). + # Confidence-scheduled verification is NOT implemented yet: the + # confidence_proj weights are loaded (for the follow-up MR) but never + # used, and drafting always proposes the full K tokens. + self._dspark_shift_label = bool(dflash_config.get('shift_label', False)) + self._dspark_use_swa = bool(dflash_config.get('use_swa', False)) + self._dspark_swa_window = int( + dflash_config.get('swa_window_size', 0) or 0) + self._dspark_markov_rank = int(dflash_config.get('markov_rank', 0) or 0) + self._dspark_markov_head_type = str( + dflash_config.get('markov_head_type', 'vanilla') + or 'vanilla').lower() + self._dspark_use_confidence_head = bool( + dflash_config.get('use_confidence_head', False)) + # Plain None placeholders rather than nn.Parameter/buffer: most + # DFlash checkpoints don't ship these heads, and their shapes + # ([vocab, rank]) are checkpoint-dependent, so nothing is + # pre-allocated. load_weights() fills them in only when the + # checkpoint ships them; consumers treat None as "head absent". + self.markov_w1 = None # [vocab, rank] (nn.Embedding weight layout) + self.markov_w2 = None # [vocab, rank] (nn.Linear(rank->vocab) weight) + self.confidence_proj_weight = None # loaded, unused (follow-up MR) + self.confidence_proj_bias = None + + if self._dspark_markov_rank > 0 and \ + self._dspark_markov_head_type != 'vanilla': + raise ValueError( + f"DFlash dspark drafter declares markov_head_type=" + f"'{self._dspark_markov_head_type}'; only 'vanilla' is " + "supported (gated/rnn heads need per-step hidden features).") + if self._dspark_use_swa and self._dspark_swa_window < 1: + raise ValueError( + "DFlash dspark drafter sets use_swa but swa_window_size=" + f"{dflash_config.get('swa_window_size')} is invalid.") + # causal=true is only invalid under the dspark convention. Legacy + # DFlash drafter configs (e.g. Laguna) also carry a causal field; + # their causality is handled by the legacy decode path + # (_sliding_layers_causal), so don't reject them here. + is_dspark = (str(dflash_config.get('projector_type', '') + or '').lower() == 'dspark' or self._dspark_shift_label + or self._dspark_use_swa or self._dspark_markov_rank > 0 + or self._dspark_use_confidence_head) + if is_dspark and dflash_config.get('causal'): + raise ValueError( + "DFlash dspark drafter sets causal=true; the block decode " + "only supports the non-causal dspark convention.") + # Per-layer flash-attn window for the block decode, resolved once. + num_draft_layers = getattr(pretrained_config, 'num_hidden_layers', 0) + layer_types = getattr(pretrained_config, 'layer_types', None) + self._dspark_layer_windows = [ + dspark_layer_window_size(self._dspark_use_swa, + self._dspark_swa_window, layer_types, i) + for i in range(num_draft_layers) + ] + if self._dspark_use_confidence_head: + logger.warning( + "DFlash dspark drafter declares use_confidence_head; " + "confidence-scheduled verification is not implemented yet " + "(confidence_proj weights are loaded but unused, drafting " + "always proposes the full K tokens).") + self.logits_processor = None # Set by caller after construction # RoPE - lazily initialized from draft model's attention module @@ -989,6 +1143,35 @@ def project_target_hidden(self, hidden_states = hidden_states.to(self.fc.weight.dtype) return self.hidden_norm(self.fc(hidden_states)) + @property + def has_markov_head(self) -> bool: + return self._dspark_markov_rank > 0 and self.markov_w1 is not None + + def apply_markov_chain_logits( + self, + base_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + argmax_fn=None, + vocab_slice: slice | None = None) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + No-op (returns ``base_logits`` unchanged) for non-dspark drafters. + See :func:`dspark_markov_chain_logits` for the semantics; when + ``base_logits`` is a TP vocab shard, the caller must pass this + rank's ``vocab_slice`` (to shard the markov_w2 rows identically) + and an ``argmax_fn`` returning full-vocab token ids — DFlashWorker + handles both. + """ + if not self.has_markov_head: + return base_logits + markov_w2 = self.markov_w2 if vocab_slice is None else \ + self.markov_w2[vocab_slice] + return dspark_markov_chain_logits(base_logits, + first_prev_tokens, + self.markov_w1, + markov_w2, + argmax_fn=argmax_fn) + def _post_attention_gate(self, attn_output, gate_input, attn_mod, num_heads, head_dim): """Hook applied to the block-attention output before o_proj. @@ -1035,6 +1218,42 @@ def load_weights(self, weights: Dict, weight_mapper=None, **kwargs): split[k] = v weights = split + # DSpark head weights: keep them out of the backbone remap (they'd + # get a 'model.' prefix and be dropped by allow_partial_loading). + # markov_w1/markov_w2 drive the intra-block logit bias; the + # confidence_proj weights are loaded for the confidence-scheduling + # follow-up MR but are not used yet. + dspark_keys = ('markov_w1.weight', 'markov_w2.weight', + 'confidence_proj.weight', 'confidence_proj.bias') + dspark_weights = {k: weights[k] for k in dspark_keys if k in weights} + if dspark_weights: + weights = { + k: v + for k, v in weights.items() if k not in dspark_weights + } + if self._dspark_markov_rank > 0: + vocab = self.config.vocab_size + rank = self._dspark_markov_rank + for k in ('markov_w1.weight', 'markov_w2.weight'): + if k not in dspark_weights: + raise ValueError( + f"DFlash dspark drafter declares markov_rank=" + f"{self._dspark_markov_rank} but the checkpoint is " + f"missing {k}.") + if tuple(dspark_weights[k].shape) != (vocab, rank): + raise ValueError( + f"DFlash dspark {k} has shape " + f"{tuple(dspark_weights[k].shape)}, expected " + f"[vocab, markov_rank] = ({vocab}, {rank}).") + self.markov_w1 = dspark_weights['markov_w1.weight'].to('cuda') + self.markov_w2 = dspark_weights['markov_w2.weight'].to('cuda') + if 'confidence_proj.weight' in dspark_weights: + self.confidence_proj_weight = dspark_weights[ + 'confidence_proj.weight'].to('cuda') + if 'confidence_proj.bias' in dspark_weights: + self.confidence_proj_bias = dspark_weights[ + 'confidence_proj.bias'].to('cuda') + # Remap: add 'model.' prefix where needed, and extract DFlash-specific weights remapped = {} for key, value in weights.items(): @@ -1493,6 +1712,12 @@ def dflash_forward( layer_types = getattr(self.config, 'layer_types', None) causal = (self._sliding_layers_causal and bool(layer_types) and layer_types[layer_idx] == 'sliding_attention') + # DSpark SWA: non-causal sliding window on 'sliding_attention' + # layers ((-1, -1) == flash-attn default == no window otherwise). + # KV index == token position in the pool, so this restricts draft + # queries to the last swa_window context tokens + the block. + window_size = (self._dspark_layer_windows[layer_idx] if layer_idx + < len(self._dspark_layer_windows) else (-1, -1)) out = flash_attn_with_kvcache( q=Q_bshd, k_cache=layer_k_cache, @@ -1502,6 +1727,7 @@ def dflash_forward( cache_seqlens=cache_seqlens_i32, cache_batch_idx=cache_batch_idx_i32, causal=causal, + window_size=window_size, ) attn_output = out.reshape(B * block_size, q_size) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 9c5ee02f6ac3..96d60a4be5c6 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6025,6 +6025,24 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): ResourceManagerType.SEQ_SLOT_MANAGER].prepare_resources( requests) self._setup_sampler_step(requests) + # KDA fused-verify replay caches: the ctx->gen state transfer + # populates only the base mamba conv/ssm pools; with one-model + # spec decoding (e.g. SA) the first forward for these requests + # takes the fused verify path, which reads the per-slot + # kda_conv_* replay caches instead of the conv pool. Seed them + # from the transferred conv states before the first step + # (otherwise the recurrent state is permanently contaminated by + # a zero/stale conv window; K3 SA-in-disagg GSM8K -0.6pp). + if self.model_engine.enable_spec_decode: + kv_mgr = self.resource_manager.resource_managers.get( + ResourceManagerType.KV_CACHE_MANAGER) + seed = getattr(kv_mgr, 'seed_kda_replay_caches_for_disagg_gen', + None) + if seed is not None: + seed([ + req.py_request_id + for req in cache_trans_complete_requests + ]) for req in scheduled_batch.generation_requests: if req.is_disagg_generation_transmission_complete: diff --git a/tensorrt_llm/_torch/speculative/accept_stats.py b/tensorrt_llm/_torch/speculative/accept_stats.py new file mode 100644 index 000000000000..a9a1958b9113 --- /dev/null +++ b/tensorrt_llm/_torch/speculative/accept_stats.py @@ -0,0 +1,390 @@ +# 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. +"""Opt-in DFlash/DSpark acceptance-statistics recorder. + +Collects, per rank, from the DFlash one-engine accept/reject site: + +- a histogram of accepted-draft-token counts per target verify step + (position-k acceptance rates and AL derive from it, see + :func:`summarize_hist`), +- per-request step / accepted-draft totals, +- optional per-position (confidence, accepted) calibration counts for + DSpark confidence-scheduled verification (binned; used to calibrate + ``confidence_threshold`` and Sequential Temperature Scaling). + +Default OFF and zero-overhead: :func:`maybe_create_recorder` returns +``None`` unless ``TLLM_DFLASH_ACCEPT_STATS_DIR`` is set, and the DFlash +worker only calls into the recorder when it exists. When enabled, the +recorder performs a small device->host copy per step, so it is intended +for measurement runs (eager mode), never perf-reference runs. + +Environment variables: + TLLM_DFLASH_ACCEPT_STATS_DIR: directory to write per-rank JSON stats + (``dflash_accept_stats_rank{rank}.json``). Enables recording. + TLLM_DFLASH_ACCEPT_STATS_FLUSH_EVERY: flush period in verify steps + (default 50; also flushed at interpreter exit). + +Calibration caveat: with ``confidence_threshold`` > 0 the runtime trims +low-confidence draft positions to the mask sentinel, which forces their +rejection — pairs recorded for trimmed positions are biased. Collect +calibration data with the confidence head present but +``confidence_threshold`` unset/0 so every drafted position is genuinely +verified. + +Confidence provider: the recorder itself never computes confidences; it +owns an optional ``confidence_provider`` callable + + provider(draft_model, gen_hidden, first_prev_tokens, gen_draft_tokens) + -> Optional[Sequence[Sequence[float]]] # [num_gens][K] host rows + +which the DFlash worker invokes at draft time (via +:meth:`DFlashAcceptStatsRecorder.record_draft_confidence`). With the +provider ``None`` (this MR), no calibration rows are collected and the +calibration table stays empty. The DSpark confidence-scheduled +verification MR ships ``tensorrt_llm._torch.speculative.dspark_confidence`` +with the real provider (built on its ``compute_draft_confidence`` / +``build_confidence_prev_tokens`` helpers); when that module is importable, +:func:`maybe_create_recorder` picks it up automatically. + +This module is import-light on purpose (no torch import) so the +aggregation math is unit-testable on CPU-only hosts. +""" + +import atexit +import json +import os +from typing import Any, Callable, Dict, List, Optional, Sequence + +from tensorrt_llm.logger import logger + +ENV_STATS_DIR = "TLLM_DFLASH_ACCEPT_STATS_DIR" +ENV_FLUSH_EVERY = "TLLM_DFLASH_ACCEPT_STATS_FLUSH_EVERY" +DEFAULT_NUM_CONF_BINS = 20 +STATS_FILE_PREFIX = "dflash_accept_stats_rank" + +# provider(draft_model, gen_hidden, first_prev_tokens, gen_draft_tokens) +# -> per-request per-position confidence rows (host floats), or None to +# skip this step (e.g. the drafter has no confidence head). +ConfidenceProvider = Callable[..., Optional[Sequence[Sequence[float]]]] + + +def _resolve_default_confidence_provider() -> Optional[ConfidenceProvider]: + """Return the DSpark confidence provider if its module is in the tree. + + The module ships with the DSpark confidence-scheduled verification MR; + without it (this MR alone) calibration rows are simply not collected. + """ + try: + from .dspark_confidence import dspark_confidence_provider + except ImportError: + return None + return dspark_confidence_provider + + +def maybe_create_recorder(max_draft_len: int, rank: int) -> Optional["DFlashAcceptStatsRecorder"]: + """Create a recorder iff TLLM_DFLASH_ACCEPT_STATS_DIR is set (else None).""" + stats_dir = os.environ.get(ENV_STATS_DIR) + if not stats_dir: + return None + logger.warning( + "DFlash acceptance-stats recording is EXPERIMENTAL and specific to " + "the DFlash/DSpark drafter: enabling %s with other speculative " + "decoding methods records nothing and is unsupported.", + ENV_STATS_DIR, + ) + flush_every = int(os.environ.get(ENV_FLUSH_EVERY, "50")) + return DFlashAcceptStatsRecorder( + stats_dir, + max_draft_len, + rank, + flush_every=flush_every, + confidence_provider=_resolve_default_confidence_provider(), + ) + + +class DFlashAcceptStatsRecorder: + """Accumulates DFlash acceptance statistics and dumps them to JSON. + + All inputs arrive as host lists (the caller does the single small + ``.tolist()`` device->host copy); the accumulation itself is pure + Python so it is unit-testable without a GPU. + """ + + # Request id 0 is never a real request in the py executor (ids are + # assigned from 1): it is the id carried by executor-warmup dummies and + # by the padding requests idle attention-DP ranks step in lock-step + # with active ranks. Without this filter, idle DEP ranks + # log hundreds of spurious all-K "acceptances" of padding blocks. + DEFAULT_EXCLUDE_REQUEST_IDS = frozenset({0}) + + def __init__( + self, + stats_dir: str, + max_draft_len: int, + rank: int, + flush_every: int = 50, + num_conf_bins: int = DEFAULT_NUM_CONF_BINS, + exclude_request_ids=DEFAULT_EXCLUDE_REQUEST_IDS, + confidence_provider: Optional[ConfidenceProvider] = None, + ): + os.makedirs(stats_dir, exist_ok=True) + self.path = os.path.join(stats_dir, f"{STATS_FILE_PREFIX}{rank}.json") + self.max_draft_len = max_draft_len + self.rank = rank + self.flush_every = max(1, flush_every) + self.num_conf_bins = num_conf_bins + self.exclude_request_ids = frozenset(int(r) for r in exclude_request_ids) + # Optional confidence source (see module docstring); None means + # no calibration rows are collected. + self.confidence_provider = confidence_provider + + # hist[a] = number of verify steps that accepted exactly `a` draft + # tokens (0..K); the bonus token is not counted here. + self.hist: List[int] = [0] * (max_draft_len + 1) + # request_id -> [verify_steps, accepted_draft_total] + self.per_request: Dict[int, List[int]] = {} + self.num_steps = 0 + # Confidence calibration, binned: position k (0-based), bin b covers + # confidence in [b/nbins, (b+1)/nbins). + self.calib_attempts = [[0] * num_conf_bins for _ in range(max_draft_len)] + self.calib_accepted = [[0] * num_conf_bins for _ in range(max_draft_len)] + # request_id -> per-position confidences of the in-flight draft block + # (recorded at draft time, joined with the accept outcome next step). + self._pending_conf: Dict[int, List[float]] = {} + + atexit.register(self.flush) + + # ---- recording ------------------------------------------------------- + + def on_accept(self, request_ids: Sequence[int], num_accepted_tokens: Sequence[int]) -> None: + """Record one verify step for the gen requests. + + Args: + request_ids: gen-request ids, batch order. + num_accepted_tokens: per gen request, bonus token + accepted + draft tokens (the runtime's ``new_tokens_lens`` semantics). + """ + K = self.max_draft_len + for rid, n in zip(request_ids, num_accepted_tokens): + if int(rid) in self.exclude_request_ids: + continue # executor-warmup / DEP-padding dummy + accepted_draft = min(max(int(n) - 1, 0), K) + self.hist[accepted_draft] += 1 + entry = self.per_request.setdefault(int(rid), [0, 0]) + entry[0] += 1 + entry[1] += accepted_draft + conf = self._pending_conf.pop(int(rid), None) + if conf is not None: + for k, c in enumerate(conf[:K]): + b = _conf_bin(c, self.num_conf_bins) + self.calib_attempts[k][b] += 1 + if accepted_draft > k: + self.calib_accepted[k][b] += 1 + self.num_steps += 1 + if self.num_steps % self.flush_every == 0: + self.flush() + + def on_draft_confidence( + self, request_ids: Sequence[int], confidence_rows: Sequence[Sequence[float]] + ) -> None: + """Stash per-position confidences of the block drafted this step. + + They are joined with the accept outcome when the same request is + verified on the next step (dropped if the request finishes first). + """ + for rid, row in zip(request_ids, confidence_rows): + if int(rid) in self.exclude_request_ids: + continue # executor-warmup / DEP-padding dummy + self._pending_conf[int(rid)] = [float(c) for c in row] + + def record_draft_confidence( + self, + request_ids: Sequence[int], + draft_model, + gen_hidden, + first_prev_tokens, + gen_draft_tokens, + ) -> None: + """Ask the confidence provider for this step's rows and stash them. + + No-op when no provider is configured or the provider declines + (returns None). The provider arguments are opaque to the recorder; + they are forwarded verbatim from the DFlash draft site. + """ + if self.confidence_provider is None: + return + rows = self.confidence_provider( + draft_model, gen_hidden, first_prev_tokens, gen_draft_tokens + ) + if rows is None: + return + self.on_draft_confidence(request_ids, rows) + + # ---- output ---------------------------------------------------------- + + def snapshot(self) -> Dict[str, Any]: + return { + "rank": self.rank, + "max_draft_len": self.max_draft_len, + "num_steps": self.num_steps, + "accepted_draft_hist": list(self.hist), + "per_request": { + str(rid): {"steps": v[0], "accepted_draft": v[1]} + for rid, v in self.per_request.items() + }, + "confidence_calibration": { + "num_bins": self.num_conf_bins, + "attempts": [list(r) for r in self.calib_attempts], + "accepted": [list(r) for r in self.calib_accepted], + }, + } + + def flush(self) -> None: + tmp = self.path + ".tmp" + with open(tmp, "w") as f: + json.dump(self.snapshot(), f) + os.replace(tmp, self.path) + + +def _conf_bin(c: float, num_bins: int) -> int: + return min(max(int(float(c) * num_bins), 0), num_bins - 1) + + +# ---- aggregation math (pure functions, CPU unit-tested) ------------------- + + +def summarize_hist(hist: Sequence[int]) -> Dict[str, Any]: + """Summary statistics from an accepted-draft-count histogram. + + Args: + hist: hist[a] = verify steps that accepted exactly ``a`` draft + tokens, a in 0..K. + + Returns dict with: + num_steps, mean_accepted_draft, + al: mean accepted tokens per target step INCLUDING the bonus token + (mean_accepted_draft + 1) — the AL figure of merit, + ar_per_position: [K] where entry k-1 = P(draft position k accepted) + = P(accepted_draft >= k). Prefix acceptance makes this exact. + """ + K = len(hist) - 1 + num_steps = sum(hist) + if num_steps == 0: + return { + "num_steps": 0, + "mean_accepted_draft": 0.0, + "al": 0.0, + "ar_per_position": [0.0] * K, + } + total_accepted = sum(a * n for a, n in enumerate(hist)) + ar = [] + tail = num_steps + for k in range(1, K + 1): + tail -= hist[k - 1] # steps with accepted_draft < k drop out + ar.append(tail / num_steps) + return { + "num_steps": num_steps, + "mean_accepted_draft": total_accepted / num_steps, + "al": total_accepted / num_steps + 1.0, + "ar_per_position": ar, + } + + +def merge_snapshots(snapshots: Sequence[Dict[str, Any]]) -> Dict[str, Any]: + """Merge per-rank recorder snapshots (attention-DP: ranks hold disjoint + request sets; TP-replicated ranks should pass a single rank's file).""" + if not snapshots: + raise ValueError("no snapshots to merge") + K = snapshots[0]["max_draft_len"] + nbins = snapshots[0]["confidence_calibration"]["num_bins"] + merged = { + "max_draft_len": K, + "num_steps": 0, + "accepted_draft_hist": [0] * (K + 1), + "per_request": {}, + "confidence_calibration": { + "num_bins": nbins, + "attempts": [[0] * nbins for _ in range(K)], + "accepted": [[0] * nbins for _ in range(K)], + }, + } + for s in snapshots: + if s["max_draft_len"] != K: + raise ValueError("mismatched max_draft_len across snapshots") + if s["confidence_calibration"]["num_bins"] != nbins: + raise ValueError("mismatched num_bins across snapshots") + merged["num_steps"] += s["num_steps"] + for a, n in enumerate(s["accepted_draft_hist"]): + merged["accepted_draft_hist"][a] += n + for rid, v in s["per_request"].items(): + entry = merged["per_request"].setdefault(rid, {"steps": 0, "accepted_draft": 0}) + entry["steps"] += v["steps"] + entry["accepted_draft"] += v["accepted_draft"] + cc = s["confidence_calibration"] + for k in range(K): + for b in range(nbins): + merged["confidence_calibration"]["attempts"][k][b] += cc["attempts"][k][b] + merged["confidence_calibration"]["accepted"][k][b] += cc["accepted"][k][b] + return merged + + +def calibration_table( + attempts: Sequence[Sequence[int]], accepted: Sequence[Sequence[int]] +) -> Dict[str, Any]: + """Empirical acceptance rate per (position, confidence bin). + + Returns bin centers plus, per position, the empirical acceptance rate + in each bin (None where the bin has no samples) and per-position + expected calibration error (ECE, attempts-weighted |confidence - + empirical acceptance|). A well-calibrated confidence head has the + empirical rate tracking the bin center (ECE ~ 0); a monotone but + shifted/scaled curve is what Sequential Temperature Scaling + corrects. + """ + K = len(attempts) + nbins = len(attempts[0]) if K else 0 + centers = [(b + 0.5) / nbins for b in range(nbins)] + per_position = [] + for k in range(K): + rates: List[Optional[float]] = [] + ece_num = 0.0 + n_total = 0 + for b in range(nbins): + n = attempts[k][b] + if n == 0: + rates.append(None) + continue + rate = accepted[k][b] / n + rates.append(rate) + ece_num += n * abs(centers[b] - rate) + n_total += n + per_position.append( + { + "empirical_acceptance": rates, + "num_samples": list(attempts[k]), + "ece": (ece_num / n_total) if n_total else None, + } + ) + return {"bin_centers": centers, "per_position": per_position} + + +def load_rank_snapshots(stats_dir: str) -> List[Dict[str, Any]]: + """Load every per-rank stats JSON found in ``stats_dir``.""" + snaps = [] + for name in sorted(os.listdir(stats_dir)): + if name.startswith(STATS_FILE_PREFIX) and name.endswith(".json"): + with open(os.path.join(stats_dir, name)) as f: + snaps.append(json.load(f)) + return snaps diff --git a/tensorrt_llm/_torch/speculative/dflash.py b/tensorrt_llm/_torch/speculative/dflash.py index 585b796a6dfb..84db6eb9abf8 100644 --- a/tensorrt_llm/_torch/speculative/dflash.py +++ b/tensorrt_llm/_torch/speculative/dflash.py @@ -27,12 +27,39 @@ from ..attention_backend import AttentionMetadata from ..pyexecutor.mamba_cache_manager import MambaHybridCacheManager from ..pyexecutor.resource_manager import BaseResourceManager +from .accept_stats import maybe_create_recorder from .interface import SpecMetadata, SpecWorkerBase if TYPE_CHECKING: from ...llmapi.llm_args import DFlashDecodingConfig +def dflash_draft_slot_ids( + num_gens: int, + block_size: int, + num_draft_tokens: int, + shift_label: bool, + device="cuda", +) -> torch.Tensor: + """Flat indices into the [num_gens * block_size] drafter block outputs + whose hidden states produce the K draft-token logits per request. + + Plain DFlash (K2.7 convention, shift_label off): the hidden state at + mask slot j (j = 1..K) predicts draft token j; slot 0 (which holds the + anchor/bonus token) is unused. + + DSpark shift_label convention (DeepSpec: labels for a block anchored at + position p are input_ids[p+1 .. p+block_size], so the hidden state at + block slot j predicts the token at position p+1+j): slots 0..K-1 are + used, and slot 0 — the anchor token slot — predicts the first draft + token. + """ + request_bases = torch.arange(num_gens, dtype=torch.long, device=device) * block_size + offsets = torch.arange(num_draft_tokens, dtype=torch.long, device=device) + first_slot = 0 if shift_label else 1 + return (request_bases.unsqueeze(1) + first_slot + offsets.unsqueeze(0)).flatten() + + @dataclass class DFlashSpecMetadata(SpecMetadata): """Metadata for DFlash speculative decoding. @@ -192,6 +219,17 @@ def __init__( self._free_slots = deque() # available slot indices self._dummy_slot = None # for cudagraph padding or warmup dummy requests + # Opt-in acceptance-statistics recorder (None unless + # TLLM_DFLASH_ACCEPT_STATS_DIR is set; see accept_stats.py). Only + # consulted behind `is not None` checks — zero overhead when off. + self._accept_stats = maybe_create_recorder( + spec_config.max_draft_len, getattr(mapping, "rank", 0) or 0 + ) + if self._accept_stats is not None: + logger.info( + f"DFlash: acceptance-statistics recording enabled -> {self._accept_stats.path}" + ) + logger.info( f"DFlashWorker initialized with use_separate_draft_kv_cache={use_separate_draft_kv_cache}" ) @@ -460,6 +498,21 @@ def _forward_impl( logits, attn_metadata, spec_metadata ) + # Opt-in acceptance recording (env-gated; eager-mode measurement + # runs only). Skipped for CUDA-graph batches (capture/replay/warmup + # use synthetic requests and forbid the host sync). + if ( + self._accept_stats is not None + and num_gens > 0 + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + ): + self._accept_stats.on_accept( + spec_metadata.request_ids[num_contexts:batch_size], + num_accepted_tokens[num_contexts:batch_size].tolist(), + ) + # Update GDN/Mamba recurrent states to the accepted token's state. if num_gens > 0 and isinstance(attn_metadata.kv_cache_manager, MambaHybridCacheManager): attn_metadata.kv_cache_manager.update_mamba_states( @@ -525,13 +578,15 @@ def _forward_impl( ctx_cache_batch_idx=inputs["ctx_cache_batch_idx"], ) - # Gather K logits per gen request from mask positions (1..K). - # hidden_states_out is flat: [num_gens * block_size, hidden_dim] + # Gather K logits per gen request from the block outputs. + # hidden_states_out is flat: [num_gens * block_size, hidden_dim]. + # Plain DFlash reads mask slots 1..K; dspark shift_label reads + # slots 0..K-1 (see dflash_draft_slot_ids). block_size = self._resolved_block_size - request_bases = torch.arange(num_gens, dtype=torch.long, device="cuda") * block_size - offsets = torch.arange(K, dtype=torch.long, device="cuda") - # Masks are at positions 1..K in each request's block_size output - gen_gather_ids = (request_bases.unsqueeze(1) + 1 + offsets.unsqueeze(0)).flatten() + shift_label = getattr(draft_model, "_dspark_shift_label", False) + gen_gather_ids = dflash_draft_slot_ids( + num_gens, block_size, K, shift_label, device="cuda" + ) gen_gather_ids = gen_gather_ids.clamp(max=hidden_states_out.shape[0] - 1) gen_logits = draft_model.logits_processor( @@ -541,6 +596,13 @@ def _forward_impl( vocab_size = gen_logits.shape[-1] gen_logits = gen_logits.reshape(num_gens, K, vocab_size) + # DSpark Markov head: add the greedy-chained intra-block + # logit bias before sampling (no-op for plain DFlash). + if getattr(draft_model, "has_markov_head", False): + gen_logits = self._apply_dspark_markov_bias( + draft_model, gen_logits, inputs["first_prev_tokens"], spec_metadata + ) + gen_draft_tokens = self.sample_draft_tokens( gen_logits, spec_metadata, @@ -548,6 +610,28 @@ def _forward_impl( num_contexts=num_contexts, ) + # Opt-in confidence-calibration recording (env-gated). The + # confidence values come from an optional provider on the + # recorder (None here -> no calibration rows collected; the + # DSpark confidence-scheduled verification MR supplies the + # real provider, see accept_stats.py). Guarded off for + # CUDA-graph batches and d2t vocab-mapped drafters. + if ( + self._accept_stats is not None + and self._accept_stats.confidence_provider is not None + and not spec_metadata.is_cuda_graph + and not torch.cuda.is_current_stream_capturing() + and spec_metadata.request_ids is not None + and self._d2t is None + ): + self._accept_stats.record_draft_confidence( + spec_metadata.request_ids[num_contexts:batch_size], + draft_model, + hidden_states_out[gen_gather_ids].reshape(num_gens, K, -1), + inputs["first_prev_tokens"], + gen_draft_tokens, + ) + else: gen_draft_tokens = torch.empty((0, K), dtype=torch.int32, device="cuda") @@ -589,6 +673,61 @@ def _forward_impl( "next_new_tokens": next_new_tokens, } + def _apply_dspark_markov_bias( + self, + draft_model, + gen_logits: torch.Tensor, + first_prev_tokens: torch.Tensor, + spec_metadata, + ) -> torch.Tensor: + """Apply the dspark vanilla-Markov intra-block bias to block logits. + + Reference (DeepSpec VanillaMarkov.sample_block_tokens, temperature 0): + step i adds bias = markov_w2 @ markov_w1[prev_i] to the shared-lm_head + logits, where prev_0 is the anchor (last accepted) token and prev_{i>0} + is the greedy token from step i-1's biased logits. Greedy per-position + argmax of the returned logits therefore reproduces the reference + sampled chain; the rejection-sampling path samples from the same + biased distributions (proposal conditioned on the greedy chain). + + Handles a TP vocab-sharded draft lm_head by slicing markov_w2's rows + to this rank's contiguous shard and chaining through the TP-aware + global argmax. + """ + if self._d2t is not None: + raise NotImplementedError( + "DSpark Markov head requires a shared draft/target vocab " + "(d2t vocab mapping is not supported)." + ) + full_vocab = draft_model.markov_w2.shape[0] + shard = gen_logits.shape[-1] + vocab_slice = None + if shard != full_vocab: + mapping = self.mapping + if ( + mapping is None + or getattr(mapping, "enable_attention_dp", False) + or shard * mapping.tp_size != full_vocab + ): + raise NotImplementedError( + f"DSpark Markov head: draft logits width {shard} does not " + f"match the drafter vocab {full_vocab} and is not a plain " + "TP column shard of it." + ) + vocab_slice = slice(mapping.tp_rank * shard, (mapping.tp_rank + 1) * shard) + + def argmax_fn(step_logits): + # Full-vocab token ids (TP-aware when sharded); tokens stay in + # draft-vocab space, which is what markov_w1 indexes. + return self.greedy_sample_draft_with_tp_gather(step_logits, spec_metadata).long() + + return draft_model.apply_markov_chain_logits( + gen_logits, + first_prev_tokens, + argmax_fn=argmax_fn, + vocab_slice=vocab_slice, + ) + def prepare_1st_drafter_inputs( self, input_ids: torch.LongTensor, @@ -739,6 +878,7 @@ def prepare_1st_drafter_inputs( query_positions = torch.empty(0, 0, dtype=torch.long, device="cuda") num_ctx_per_req_t = torch.empty(0, dtype=torch.long, device="cuda") slots = torch.empty(0, dtype=torch.long, device="cuda") + bonus = torch.empty(0, dtype=torch.long, device="cuda") return { "noise_embedding": noise_embedding, @@ -747,4 +887,7 @@ def prepare_1st_drafter_inputs( "ctx_k_cache": self._ctx_k_buf, "ctx_v_cache": self._ctx_v_buf, "ctx_cache_batch_idx": slots, + # Anchor token per gen request (block slot 0): last accepted + # token. The dspark Markov chain conditions its first step on it. + "first_prev_tokens": bonus, } diff --git a/tensorrt_llm/_torch/speculative/suffix_automaton.py b/tensorrt_llm/_torch/speculative/suffix_automaton.py index 12bb77208fe4..4eba0781d46f 100644 --- a/tensorrt_llm/_torch/speculative/suffix_automaton.py +++ b/tensorrt_llm/_torch/speculative/suffix_automaton.py @@ -26,6 +26,7 @@ from tensorrt_llm._utils import prefer_pinned from tensorrt_llm.bindings.internal import suffix_automaton as _sa_native +from tensorrt_llm.logger import logger as trtllm_logger from ..pyexecutor.llm_request import LlmRequest from ..pyexecutor.resource_manager import BaseResourceManager @@ -598,14 +599,56 @@ def extend_global( # --- BaseResourceManager interface --- def prepare_resources(self, scheduled_batch: ScheduledRequests): - """Prepare SA states for new context requests.""" + """Prepare SA states for new context and disagg-generation requests.""" for req in scheduled_batch.context_requests: + # Disaggregated serving: the executor's _prepare_disagg_gen_init + # also routes DISAGG_GENERATION_INIT requests through + # prepare_resources as context_requests_last_chunk, BEFORE the + # context server's first generated token has arrived (it is only + # appended later, by _prepare_disagg_gen_transmission_complete). + # Initializing here would (a) freeze the automaton at prompt-only + # -- permanently missing the ctx-produced first token relative to + # an aggregated run, since _initialized_requests makes init + # one-shot -- and (b) pin an SA slot for the whole KV-transfer + # duration. Defer to the generation-request loop below, which + # runs once the request is scheduled with its full history. + if req.is_generation_only_request(): + continue if req.is_first_context_chunk: if req.request_id not in self._initialized_requests: context_tokens = req.get_tokens(0) self.add_request(req.request_id, context_tokens) self._initialized_requests.add(req.request_id) + # Disaggregated serving: on a generation server, requests skip the + # local context phase -- they arrive as DISAGG_GENERATION_INIT + # and are scheduled as generation requests (the executor's + # _prepare_disagg_gen_transmission_complete appends the context + # server's first generated token before resources are prepared). + # Build the automaton here from the full token history (prompt + + # first generated token), matching the state an aggregated run would + # have after its context step. _initialized_requests keeps this a + # one-shot init; free_resources clears it on completion. + for req in scheduled_batch.generation_requests: + if ( + req.is_generation_only_request() + and not req.is_dummy + and req.request_id not in self._initialized_requests + ): + history = req.get_tokens(0) + self.add_request(req.request_id, history) + self._initialized_requests.add(req.request_id) + # Observability marker for disagg SA: + # confirms the automaton is built in the GENERATION schedule + # with the full history (prompt + ctx first token, i.e. + # len(history) == prompt_len + 1), not in the premature + # context-phase pass at _prepare_disagg_gen_init. + trtllm_logger.debug( + f"[SA] disagg gen-init: request {req.request_id} automaton " + f"built in generation schedule from {len(history)} history " + f"tokens (prompt_len={getattr(req, 'py_prompt_len', None)})" + ) + def update_resources(self, scheduled_batch: ScheduledRequests): """Update resources after forward pass (no-op for SA).""" pass diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index 70ddbd8f055d..0ba4b9915d8c 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -15,7 +15,7 @@ import torch from PIL import Image from torchvision.transforms import ToTensor -from transformers import AutoProcessor, ProcessorMixin +from transformers import AutoProcessor, PreTrainedTokenizerBase, ProcessorMixin from transformers.utils import logging from tensorrt_llm.inputs.content_format import (ContentFormat, @@ -586,6 +586,15 @@ def interleave_mm_placeholders( return separator.join(parts) +def _has_python_chat_template(tokenizer: PreTrainedTokenizerBase) -> bool: + """Return whether a tokenizer overrides Hugging Face's Jinja renderer.""" + apply_chat_template_method = getattr(type(tokenizer), "apply_chat_template", + None) + return (apply_chat_template_method is not None + and apply_chat_template_method + is not PreTrainedTokenizerBase.apply_chat_template) + + def resolve_hf_chat_template( tokenizer: TokenizerBase, processor: ProcessorMixin, @@ -691,6 +700,7 @@ def apply_chat_template( Uses content-format-driven dispatch: - PASSTHROUGH: skip template rendering, just concatenate content strings + - PYTHON: use a tokenizer-native renderer when no Jinja template is declared - OPENAI: reconstructs content as list of dicts for the template to handle - STRING: keeps flattened text with pre-inserted placeholders """ @@ -716,6 +726,21 @@ def apply_chat_template( if isinstance(tokenizer, TransformersTokenizer): tokenizer = tokenizer.tokenizer # we need the TokenizerBase for apply_chat_template + if (chat_template is None + and getattr(processor, "chat_template", None) is None + and getattr(tokenizer, "chat_template", None) is None + and _has_python_chat_template(tokenizer)): + native_kwargs = dict(chat_template_kwargs or {}) + if documents is not None: + native_kwargs["documents"] = documents + return tokenizer.apply_chat_template( + conversation, + tools=tools, + tokenize=enable_tokenize, + add_generation_prompt=add_generation_prompt, + **native_kwargs, + ) + hf_chat_template = resolve_hf_chat_template(tokenizer, processor, chat_template, tools) if hf_chat_template is None: diff --git a/tensorrt_llm/llmapi/reasoning_parser.py b/tensorrt_llm/llmapi/reasoning_parser.py index c04456bdf6d3..d6217e8e287d 100644 --- a/tensorrt_llm/llmapi/reasoning_parser.py +++ b/tensorrt_llm/llmapi/reasoning_parser.py @@ -72,9 +72,26 @@ def create_reasoning_parser( def keys(cls): return cls._parsers.keys() + @classmethod + def needs_raw_special_tokens(cls, reasoning_parser: str) -> bool: + """Whether the registered parser must see special tokens. + + See ``BaseReasoningParser.needs_raw_special_tokens``. + """ + entry = cls._parsers.get(reasoning_parser.lower()) + return bool(entry + and getattr(entry[0], "needs_raw_special_tokens", False)) + class BaseReasoningParser(ABC): + # Parsers whose delimiters are registered special tokens must see the + # raw decoded stream; the serving layer checks this flag and disables + # ``skip_special_tokens`` for the request (mirrors + # ``BaseToolParser.needs_raw_special_tokens``, which only takes effect + # when the request carries tools). + needs_raw_special_tokens: bool = False + def __init__(self, *, chat_template_kwargs: Optional[dict[str, Any]] = None) -> None: @@ -334,6 +351,7 @@ def parse(self, text: str) -> ReasoningParserResult: "gemma4": "gemma4", "kimi_k2": "kimi_k2", "kimi_k25": "kimi_k25", + "kimi_k3": "kimi_k3", "minimax_m3": "minimax_m3", "minimax_m3_vl": "minimax_m3", } @@ -832,3 +850,189 @@ def parse_delta(self, delta_text: str) -> ReasoningParserResult: raise RuntimeError( "Unreachable code reached in `KimiK2ReasoningParser.parse_delta`") + + +@register_reasoning_parser("kimi_k3") +class KimiK3ReasoningParser(BaseReasoningParser): + """Reasoning parser for Kimi-K3 XTML output. + + K3 renders assistant messages as an XTML tag stream built from the + special tokens ``<|open|>`` / ``<|close|>`` / ``<|sep|>`` / + ``<|end_of_msg|>`` with plain-text tag headers (see the checkpoint's + ``encoding_k3.py``):: + + <|open|>think<|sep|>REASONING<|close|>think<|sep|> + <|open|>response<|sep|>CONTENT<|close|>response<|sep|> + [<|open|>tools<|sep|>...calls...<|close|>tools<|sep|>] + <|close|>message<|sep|><|end_of_msg|> + + The generation prompt already ends inside ``<|open|>think<|sep|>`` + (thinking mode, the default) or ``<|open|>response<|sep|>`` + (``chat_template_kwargs={"thinking": False}``), so the model output + begins mid-channel with no opening tag. + + This parser emits the think body as ``reasoning_content`` and the + response body as ``content``. A ``tools`` section is passed through + into ``content`` verbatim (special tokens included) so the ``kimi_k3`` + tool parser can consume it; all other structural markup is dropped. + """ + + needs_raw_special_tokens = True + + OPEN = "<|open|>" + CLOSE = "<|close|>" + SEP = "<|sep|>" + EOM = "<|end_of_msg|>" + TOOLS_END = CLOSE + "tools" + SEP + + def __init__(self, + *, + chat_template_kwargs: Optional[dict[str, Any]] = None) -> None: + super().__init__(chat_template_kwargs=chat_template_kwargs) + thinking = True + if chat_template_kwargs is not None: + thinking = chat_template_kwargs.get("thinking", True) is not False + # Channel the model starts generating in (its opening tag is part + # of the prompt). + self._initial_channel = "think" if thinking else "response" + self._reset() + + def _reset(self) -> None: + self._buffer = "" + # 'body' | 'open_header' | 'close_header' | 'tools_pass' | 'done' + self._state = "body" + self._channel: Optional[str] = self._initial_channel + + @staticmethod + def _partial_suffix_len(text: str, markers: tuple[str, ...]) -> int: + """Length of the longest text suffix that is a proper prefix of any marker. + + Markers may contain internal ``<`` (e.g. ``<|close|>tools<|sep|>``), + so every suffix length up to ``len(marker) - 1`` must be checked, not + just the one starting at the last ``<``. + """ + best = 0 + for marker in markers: + for length in range(min(len(text), len(marker) - 1), best, -1): + if marker.startswith(text[-length:]): + best = length + break + return best + + def _emit(self, text: str, content: list, reasoning: list) -> None: + if not text: + return + if self._channel == "think": + reasoning.append(text) + elif self._channel == "response": + content.append(text) + # channel None: structural gap – dropped + + def _step(self, content: list, reasoning: list) -> bool: + """Consume as much of the buffer as possible. + + Returns False when more input is needed. + """ + buf = self._buffer + if self._state == "done": + self._buffer = "" + return False + if self._state == "body": + markers = (self.OPEN, self.CLOSE, self.EOM) + indices = [(buf.find(m), m) for m in markers] + indices = [(i, m) for i, m in indices if i != -1] + if not indices: + hold = self._partial_suffix_len(buf, markers) + emit_len = len(buf) - hold + self._emit(buf[:emit_len], content, reasoning) + self._buffer = buf[emit_len:] + return False + idx, marker = min(indices) + self._emit(buf[:idx], content, reasoning) + self._buffer = buf[idx + len(marker):] + if marker == self.OPEN: + self._state = "open_header" + elif marker == self.CLOSE: + self._state = "close_header" + else: + self._state = "done" + self._buffer = "" + return True + if self._state in ("open_header", "close_header"): + idx = buf.find(self.SEP) + if idx == -1: + return False + header = buf[:idx] + self._buffer = buf[idx + len(self.SEP):] + tag = header.split(None, 1)[0] if header.split() else "" + if self._state == "open_header": + if tag == "think": + self._channel = "think" + elif tag == "response": + self._channel = "response" + elif tag == "tools": + # Replay the section opener for the tool parser. + self._channel = "response" + self._emit(self.OPEN + header + self.SEP, content, + reasoning) + self._state = "tools_pass" + return True + else: + self._channel = None + else: + if tag == "message": + self._state = "done" + self._buffer = "" + return False + self._channel = None + self._state = "body" + return True + if self._state == "tools_pass": + idx = buf.find(self.TOOLS_END) + if idx == -1: + hold = self._partial_suffix_len(buf, (self.TOOLS_END, )) + emit_len = len(buf) - hold + self._emit(buf[:emit_len], content, reasoning) + self._buffer = buf[emit_len:] + return False + end = idx + len(self.TOOLS_END) + self._emit(buf[:end], content, reasoning) + self._buffer = buf[end:] + self._channel = None + self._state = "body" + return True + raise RuntimeError("Unreachable state in `KimiK3ReasoningParser._step`") + + def _feed(self, text: str) -> ReasoningParserResult: + self._buffer += text + content: list = [] + reasoning: list = [] + while self._step(content, reasoning): + pass + return ReasoningParserResult(content="".join(content), + reasoning_content="".join(reasoning)) + + def parse(self, text: str) -> ReasoningParserResult: + self._reset() + result = self._feed(text) + tail = self.finish() + return ReasoningParserResult( + content=result.content + tail.content, + reasoning_content=result.reasoning_content + tail.reasoning_content, + ) + + def parse_delta(self, delta_text: str) -> ReasoningParserResult: + return self._feed(delta_text) + + def finish(self) -> ReasoningParserResult: + remaining = self._buffer + self._buffer = "" + if not remaining or self._state == "done": + return ReasoningParserResult() + if self._state in ("body", "tools_pass"): + if self._channel == "think": + return ReasoningParserResult(reasoning_content=remaining) + if self._channel == "response": + return ReasoningParserResult(content=remaining) + # Header fragments and structural-gap text are dropped. + return ReasoningParserResult() diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 367caab10ec1..aa56a508e856 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -48,6 +48,7 @@ from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole) from tensorrt_llm.llmapi.llm import LLM, RequestOutput +from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory from tensorrt_llm.llmapi.thinking_budget import \ add_thinking_budget_logits_processor from tensorrt_llm.logger import logger @@ -162,6 +163,36 @@ async def route_handler(request: Request): TIMEOUT_KEEP_ALIVE = 5 # seconds. +def _configure_parser_special_token_decoding( + sampling_params: SamplingParams, reasoning_parser_name: Optional[str], + tool_parser_name: Optional[str], has_tools: bool) -> None: + """Configure detokenization for parsers that consume special tokens.""" + needs_compact_special_tokens = False + if reasoning_parser_name and ReasoningParserFactory.needs_raw_special_tokens( + reasoning_parser_name): + # Unlike the tool-parser flag below, this applies to every chat + # request: reasoning delimiters are special tokens regardless of + # whether tools are attached. + sampling_params.skip_special_tokens = False + needs_compact_special_tokens = reasoning_parser_name.lower( + ) == "kimi_k3" + + if tool_parser_name and has_tools: + tool_parser_cls = ToolParserFactory.parsers.get( + tool_parser_name.lower()) + if tool_parser_cls and getattr(tool_parser_cls, + 'needs_raw_special_tokens', False): + sampling_params.skip_special_tokens = False + needs_compact_special_tokens |= tool_parser_name.lower( + ) == "kimi_k3" + + if needs_compact_special_tokens: + # K3 XTML places ordinary-text tag names directly between special + # tokens, for example ``<|open|>tools<|sep|>``. Inserting spaces here + # changes the protocol and prevents the K3 parsers from matching it. + sampling_params.spaces_between_special_tokens = False + + def _build_tool_strict_guided_decoding_params(tools, tool_parser_name): """Build GuidedDecodingParams with structural tags for tools with strict=True. @@ -1484,12 +1515,13 @@ async def chat_stream_generator( tokenizer=self.tokenizer, chat_template_kwargs=request.chat_template_kwargs, ) + reasoning_parser_name = self.generator.args.reasoning_parser + _configure_parser_special_token_decoding( + sampling_params, + reasoning_parser_name=reasoning_parser_name, + tool_parser_name=self.tool_parser, + has_tools=bool(request.tools)) if self.tool_parser and request.tools: - tool_parser_cls = ToolParserFactory.parsers.get( - self.tool_parser.lower()) - if tool_parser_cls and getattr( - tool_parser_cls, 'needs_raw_special_tokens', False): - sampling_params.skip_special_tokens = False # When strict=True on any tool, apply constrained decoding # via structural tags (only if response_format doesn't already # set guided decoding). diff --git a/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py new file mode 100644 index 000000000000..b434e019be61 --- /dev/null +++ b/tensorrt_llm/serve/tool_parser/kimi_k3_tool_parser.py @@ -0,0 +1,201 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tool-call parser for the Kimi K3 XTML output format. + +K3 emits tool calls as an XTML tag stream built from the special tokens +``<|open|>`` / ``<|close|>`` / ``<|sep|>`` with plain-text tag headers +(authoritative rendering: the checkpoint's ``encoding_k3.py``):: + + <|open|>tools<|sep|> + <|open|>call tool="NAME" index="1"<|sep|> + <|open|>argument key="K" type="string|number|boolean|null|object|array"<|sep|> + VALUE + <|close|>argument<|sep|> + ... + <|close|>call<|sep|> + ... + <|close|>tools<|sep|> + +Alternatively a call body may carry one raw JSON block:: + + <|open|>json type="object"<|sep|>{...}<|close|>json<|sep|> + +Attribute values are escaped (``&`` -> ``&``, ``"`` -> ``"``). +``argument`` bodies are raw text for ``type="string"`` and JSON text for +every other type. The section arrives after the response body, so streaming +buffers the section and emits complete calls once ``<|close|>tools<|sep|>`` +is seen. +""" + +import json +import re +from typing import Any, Dict, List + +from tensorrt_llm.logger import logger + +from ..openai_protocol import ChatCompletionToolsParam as Tool +from .base_tool_parser import BaseToolParser +from .core_types import StreamingParseResult, ToolCallItem, _GetInfoFunc + + +def _unescape_attr(value: str) -> str: + return value.replace(""", '"').replace("&", "&") + + +def _parse_attrs(header: str) -> Dict[str, str]: + return {key: _unescape_attr(value) for key, value in re.findall(r'(\w+)="([^"]*)"', header)} + + +class KimiK3ToolParser(BaseToolParser): + """Detector for the Kimi K3 XTML function-call format.""" + + needs_raw_special_tokens = True + + def __init__(self): + super().__init__() + self.bot_token = "<|open|>tools<|sep|>" # nosec B105 + self.eot_token = "<|close|>tools<|sep|>" # nosec B105 + # Structural leftovers that may trail the tools section when the + # reasoning parser is not in front of this parser. + self._trailing_structural = re.compile( + r"(?:<\|close\|>message<\|sep\|>|<\|end_of_msg\|>)+\s*$" + ) + + self._call_regex = re.compile( + r"<\|open\|>call(?P[^<]*?)<\|sep\|>" + r"(?P.*?)<\|close\|>call<\|sep\|>", + re.DOTALL, + ) + self._argument_regex = re.compile( + r"<\|open\|>argument(?P[^<]*?)<\|sep\|>" + r"(?P.*?)<\|close\|>argument<\|sep\|>", + re.DOTALL, + ) + self._json_regex = re.compile( + r"<\|open\|>json(?P[^<]*?)<\|sep\|>" + r"(?P.*?)<\|close\|>json<\|sep\|>", + re.DOTALL, + ) + + def has_tool_call(self, text: str) -> bool: + return self.bot_token in text + + def supports_structural_tag(self) -> bool: + # XTML argument bodies are tag-structured text, not JSON — the + # JSON-schema-driven structural-tag constrained decoding used for + # strict tools does not apply. + return False + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError( + "kimi_k3 XTML tool calls do not support structural-tag constrained decoding" + ) + + @staticmethod + def _coerce_value(value: str, value_type: str) -> Any: + if value_type == "string": + return value + try: + return json.loads(value) + except json.JSONDecodeError: + logger.warning( + "kimi_k3 tool parser: argument declared type=%s but body is " + "not valid JSON; keeping raw text", + value_type, + ) + return value + + def _parse_call_arguments(self, body: str) -> str: + """Reconstruct the OpenAI ``function.arguments`` JSON string from a call body.""" + json_match = self._json_regex.search(body) + if json_match is not None: + raw = json_match.group("value").strip() + try: + return json.dumps(json.loads(raw), ensure_ascii=False) + except json.JSONDecodeError: + logger.warning( + "kimi_k3 tool parser: json block is not valid JSON; passing raw text through" + ) + return raw + arguments: Dict[str, Any] = {} + for match in self._argument_regex.finditer(body): + attrs = _parse_attrs(match.group("attrs")) + key = attrs.get("key") + if key is None: + continue + arguments[key] = self._coerce_value(match.group("value"), attrs.get("type", "string")) + return json.dumps(arguments, ensure_ascii=False) + + def _parse_tools_section(self, section: str, tools: List[Tool]) -> List[ToolCallItem]: + tool_indices = self._get_tool_indices(tools) + calls: List[ToolCallItem] = [] + for position, match in enumerate(self._call_regex.finditer(section)): + attrs = _parse_attrs(match.group("attrs")) + name = attrs.get("tool") + if not name: + logger.warning( + "kimi_k3 tool parser: call without tool attribute: %s", match.group("attrs") + ) + continue + if name not in tool_indices: + logger.warning("Model attempted to call undefined function: %s", name) + calls.append( + ToolCallItem( + tool_index=position, + name=name, + parameters=self._parse_call_arguments(match.group("body")), + ) + ) + return calls + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + bot_idx = text.find(self.bot_token) + if bot_idx == -1: + return StreamingParseResult(normal_text=self._trailing_structural.sub("", text)) + normal_text = text[:bot_idx] + section = text[bot_idx + len(self.bot_token) :] + eot_idx = section.find(self.eot_token) + if eot_idx != -1: + section = section[:eot_idx] + calls = self._parse_tools_section(section, tools) + self.prev_tool_call_arr = [] + for call in calls: + try: + arguments = json.loads(call.parameters) + except json.JSONDecodeError: + arguments = call.parameters + self.prev_tool_call_arr.append( + { + "name": call.name, + "arguments": arguments, + } + ) + return StreamingParseResult(normal_text=normal_text, calls=calls) + + def parse_streaming_increment(self, new_text: str, tools: List[Tool]) -> StreamingParseResult: + self._buffer += new_text + bot_idx = self._buffer.find(self.bot_token) + if bot_idx == -1: + hold = self._ends_with_partial_token(self._buffer, self.bot_token) + emit_len = len(self._buffer) - hold + normal_text = self._buffer[:emit_len] + self._buffer = self._buffer[emit_len:] + return StreamingParseResult(normal_text=normal_text) + + # Flush any response text preceding the section, then buffer the + # whole section until it completes: K3 tool calls terminate the + # message, so latency cost is negligible and complete calls avoid + # partial-argument reconstruction entirely. + normal_text = self._buffer[:bot_idx] + self._buffer = self._buffer[bot_idx:] + eot_idx = self._buffer.find(self.eot_token) + if eot_idx == -1: + return StreamingParseResult(normal_text=normal_text) + section_end = eot_idx + len(self.eot_token) + result = self.detect_and_parse(self._buffer[:section_end], tools) + # Anything after the section (normally empty) is re-examined on the + # next increment rather than dropped. + self._buffer = self._buffer[section_end:] + return StreamingParseResult( + normal_text=normal_text + result.normal_text, calls=result.calls + ) diff --git a/tensorrt_llm/serve/tool_parser/tool_parser_factory.py b/tensorrt_llm/serve/tool_parser/tool_parser_factory.py index 87bc6d090b48..2087b50530a3 100644 --- a/tensorrt_llm/serve/tool_parser/tool_parser_factory.py +++ b/tensorrt_llm/serve/tool_parser/tool_parser_factory.py @@ -11,6 +11,7 @@ from .glm4_parser import Glm4ToolParser from .glm47_parser import Glm47ToolParser from .kimi_k2_tool_parser import KimiK2ToolParser +from .kimi_k3_tool_parser import KimiK3ToolParser from .minimax_m2_parser import MiniMaxM2ToolParser from .minimax_m3_parser import MiniMaxM3ToolParser from .poolside_v1_parser import PoolsideV1ToolParser @@ -29,6 +30,7 @@ "deepseek_v4": "deepseek_v4", "kimi_k2": "kimi_k2", "kimi_k25": "kimi_k2", + "kimi_k3": "kimi_k3", "glm4": "glm4", "glm4_moe": "glm47", "glm4_moe_lite": "glm47", @@ -58,6 +60,7 @@ class ToolParserFactory: "qwen3": Qwen3ToolParser, "qwen3_coder": Qwen3CoderToolParser, "kimi_k2": KimiK2ToolParser, + "kimi_k3": KimiK3ToolParser, "deepseek_v3": DeepSeekV3Parser, "deepseek_v31": DeepSeekV31Parser, "deepseek_v32": DeepSeekV32Parser, diff --git a/tests/integration/defs/kimi_k3_sa_harness.py b/tests/integration/defs/kimi_k3_sa_harness.py new file mode 100644 index 000000000000..10ace5b4ae55 --- /dev/null +++ b/tests/integration/defs/kimi_k3_sa_harness.py @@ -0,0 +1,558 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 SA spec-dec / sanity test harness via the LLM API. + +Runs greedy prompts through the standard TRT-LLM PyTorch backend and +asserts the outputs contain expected content, optionally with SA +(suffix-automaton) speculative decoding and parity checking against a +baseline run. Two modes: + +* full model (default): needs 16 GPUs (4x GB300 trays), asserts output + quality (coherent completions, correct GSM8K-style answer). +* truncated (``KIMI_K3_NUM_LAYERS=``, e.g. 4 on a single 4-GPU tray): + the harness builds a temporary truncated copy of the checkpoint config + (weight shards are symlinked; extra layers are ignored at load time) and + only asserts the pipeline runs e2e (load -> prefill -> decode -> + shutdown); output text is NOT checked (a 4/93-layer model produces + gibberish by construction). + +Driven by tests/integration/defs/test_kimi_k3_specdec.py; can also be run +standalone: + + python tests/integration/defs/kimi_k3_sa_harness.py + +Environment: + KIMI_K3_CKPT model dir (required) + KIMI_K3_TP tensor_parallel_size == EP width (default 16; + 896 % TP must be 0) + KIMI_K3_ADP 1 (default) = DEP deployment (attention-DP + + MoE EP dispatch/combine, EP width == TP); + 0 = plain EP (replicated attention + + allreduce) — both modes are test-coverable + KIMI_K3_NUM_LAYERS truncate to first N layers via a temporary + checkpoint copy (skips output-quality + assertions) + KIMI_K3_MOE_BACKEND moe_config.backend passed to the LLM + (default AUTO). Currently a no-op for + Kimi K3: KimiK3MoERuntime pins the routed + MoE backend to TRTLLM regardless; parity + holds because the baseline and spec runs + share the same backend + KIMI_K3_SPEC_MODE 'off' (default) or 'sa' (suffix + automaton, one-engine) + KIMI_K3_SPEC_DRAFT_LEN SA max_draft_len (default 2) + KIMI_K3_SA_NGRAM_SIZE SA matching mode (default -1 = longest + match; >=1 = fixed ngram size) + KIMI_K3_SPEC_PARITY 0 (off) | 1/text (bit-identical outputs; + full model) | logits (tolerance logprob + comparison along shared prefix + near-tie + check at divergence; truncated models) + KIMI_K3_OUTPUT_JSON dump this run's outputs (text/token_ids/ + logprobs) to a JSON file + KIMI_K3_BASELINE_JSON load the parity baseline from a prior run's + JSON instead of a second in-process model + (required for full-model parity: the model + does not fit twice in one process) + KIMI_K3_SPEC_LP_TOL logits-parity drift tolerance (default 1.0; + catastrophic-corruption net on truncated + noise models — see _compare_logits_parity) + KIMI_K3_SPEC_TIE_TOL logits-parity near-tie gap bound (default 0.3) +Exit code 0 = PASS, 1 = FAIL. +""" + +import json +import os +import sys +import tempfile +import time + +# tensorrt_llm imports live inside _build_llm/_generate so the comparison +# machinery (_compare_logits_parity, _parity_prompts, _chosen_logprob) stays +# importable from client-side tools (kimi_k3_disagg_parity.py) that run +# without a GPU-side tensorrt_llm install. + +# Keeps the truncated-checkpoint temp dir alive for the process lifetime. +_TRUNCATED_CKPT_DIR = None + + +def _truncated_checkpoint(ckpt: str, num_layers: int) -> str: + """Build a temporary checkpoint dir truncated to the first N layers. + + The doctored ``config.json`` sets ``num_hidden_layers`` and filters the + 1-indexed ``linear_attn_config`` layer schedules consistently; every + other file (tokenizer, index, weight shards) is symlinked. Extra + checkpoint layers are ignored at load time, so the full shard set can + stay in place. + + The copy lives in node-local tmp, so truncated runs are single-node + only (remote ranks could not resolve the doctored config); the 4-GPU + integration test this serves always fits one node. + """ + global _TRUNCATED_CKPT_DIR + _TRUNCATED_CKPT_DIR = tempfile.TemporaryDirectory(prefix="kimi-k3-truncated-") + out = _TRUNCATED_CKPT_DIR.name + with open(os.path.join(ckpt, "config.json")) as f: + config = json.load(f) + text = config.get("text_config", config) + assert 0 < num_layers <= text["num_hidden_layers"] + text["num_hidden_layers"] = num_layers + lin = dict(text["linear_attn_config"]) + lin["kda_layers"] = [n for n in lin["kda_layers"] if n <= num_layers] + lin["full_attn_layers"] = [n for n in lin["full_attn_layers"] if n <= num_layers] + text["linear_attn_config"] = lin + with open(os.path.join(out, "config.json"), "w") as f: + json.dump(config, f) + for entry in os.scandir(ckpt): + if entry.name != "config.json": + os.symlink(entry.path, os.path.join(out, entry.name)) + return out + + +PROMPTS_AND_CHECKS = [ + ("The capital of France is", "Paris"), + ("1 + 1 = 2, 2 + 2 = 4, 4 + 4 =", "8"), + ("Water is made of hydrogen and", "oxygen"), + ( + "Question: Natalia sold clips to 48 of her friends in April, and " + "then she sold half as many clips in May. How many clips did " + "Natalia sell altogether in April and May?\nAnswer:", + "#### 72", + ), +] + + +def _graphs_enabled() -> bool: + return os.environ.get("KIMI_K3_CUDA_GRAPHS", "1") == "1" + + +def _build_llm(ckpt: str, tp: int, spec_mode: str, adp: bool): + from tensorrt_llm import LLM + from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, MoeConfig + + speculative_config = None + if spec_mode == "sa": + from tensorrt_llm.llmapi import SADecodingConfig + + speculative_config = SADecodingConfig( + max_draft_len=int(os.environ.get("KIMI_K3_SPEC_DRAFT_LEN", "2")), + # -1 = longest match via the suffix automaton (the SA + # differentiator); >=1 pins a fixed ngram size. + max_matching_ngram_size=int(os.environ.get("KIMI_K3_SA_NGRAM_SIZE", "-1")), + # Cross-request pattern reuse (SA's differentiator on + # homogeneous workloads); pool must be >= max_batch_size. + enable_global_pool=os.environ.get("KIMI_K3_SA_GLOBAL_POOL", "0") == "1", + ) + elif spec_mode != "off": + raise ValueError(f"KIMI_K3_SPEC_MODE={spec_mode!r} (expected 'off' or 'sa')") + return LLM( + model=ckpt, + tensor_parallel_size=tp, + enable_attention_dp=adp, + moe_expert_parallel_size=tp if adp else None, + trust_remote_code=True, # tiktoken tokenizer ships with the ckpt + # Currently a no-op for Kimi K3: KimiK3MoERuntime pins the routed + # MoE backend to TRTLLM regardless of moe_config.backend. Kept as + # a passthrough for checkpoints that honor it; parity is + # unaffected because the baseline and spec runs share the backend. + moe_config=MoeConfig(backend=os.environ.get("KIMI_K3_MOE_BACKEND", "AUTO")), + max_batch_size=int(os.environ.get("KIMI_K3_MAX_BATCH_SIZE", "8")), + max_seq_len=int(os.environ.get("KIMI_K3_MAX_SEQ_LEN", "4096")), + max_num_tokens=int(os.environ.get("KIMI_K3_MAX_NUM_TOKENS", "4096")), + enable_chunked_prefill=False, + # Non-spec standalone runs use the upstream defaults (overlap + + # CUDA graphs). Spec-dec runs keep the validated conservative + # settings (SA is graph-safe/overlap-capable in principle, but the + # K3 verify/promote path is only certified eager without overlap). + # Parity runs force BOTH instances into the conservative regime via + # KIMI_K3_CUDA_GRAPHS=0 (set in main) — a baseline compared against + # an eager spec run must not execute under a different regime. + # KIMI_K3_SPEC_CUDA_GRAPHS=1 opts a spec-dec run into CUDA graphs + # (EXPERIMENTAL: SA is graph-safe by design — draft padding keeps + # shapes static — but the K3 verify/promote path is not yet + # certified under capture). Overlap stays off for spec runs. + disable_overlap_scheduler=(speculative_config is not None or not _graphs_enabled()), + cuda_graph_config=( + CudaGraphConfig(enable_padding=True, max_batch_size=8) + if _graphs_enabled() + and ( + speculative_config is None or os.environ.get("KIMI_K3_SPEC_CUDA_GRAPHS", "0") == "1" + ) + else None + ), + speculative_config=speculative_config, + kv_cache_config=KvCacheConfig( + enable_block_reuse=False, + free_gpu_memory_fraction=float(os.environ.get("KIMI_K3_FREE_GPU_FRACTION", "0.25")), + # tokens_per_block=64 keeps the MLA (576, 512) generation path + # on the flashinfer trtllm-gen kernel (32 falls back to a C++ + # path requiring num_heads % 64 == 0; K3 has 96 query heads). + tokens_per_block=64, + ), + ) + + +def _parity_prompts(extra: int): + """Deterministic extra prompt pool for spec-dec parity statistics. + + More prompts = more independent trajectories: shared-prefix drift is + sampled widely and divergence events (near-tie vs not) become + meaningful in aggregate — a state bug corrupts trajectories + systematically, rounding flips are scattered. Templates mix factual + stubs, arithmetic, and repetitive structure (the latter gives the SA + real acceptances, exercising the state-promotion path). Model loading + dominates runtime, so extra prompts are nearly free. + """ + subjects = [ + "France", + "Japan", + "Brazil", + "Canada", + "Egypt", + "Kenya", + "Norway", + "Peru", + "Thailand", + "Greece", + "Chile", + "Poland", + ] + templates = [ + "The capital of {s} is", + "Q: Name three facts about {s}.\nA: 1.", + "{s}, {s}, {s}. The word repeated above is", + "Count by twos: 2, 4, 6, 8, 10,", + ] + return [ + ( + templates[i % len(templates)].format(s=subjects[(i // len(templates)) % len(subjects)]), + None, + ) + for i in range(extra) + ] + + +def _generate(llm, prompts, max_tokens: int, want_logprobs: bool = False): + from tensorrt_llm import SamplingParams + + want_stats = os.environ.get("KIMI_K3_SPEC_STATS", "0") == "1" + sampling = SamplingParams( + max_tokens=max_tokens, + temperature=0.0, + logprobs=5 if want_logprobs else None, + return_perf_metrics=want_stats, + ) + t0 = time.monotonic() + outputs = llm.generate(prompts, sampling) + wall = time.monotonic() - t0 + print( + f"[sanity] generate wall: {wall:.1f}s for {len(prompts)} prompts x {max_tokens} max_tokens" + ) + if want_stats: + _print_spec_stats(outputs) + return [out.outputs[0] for out in outputs] + + +def _print_spec_stats(outputs): + """Aggregate speculative-decoding stats (KIMI_K3_SPEC_STATS=1). + + tokens/step (avg_decoded_tokens_per_iter) is the headline spec-dec + win; acceptance rate = accepted draft tokens / drafted tokens. On a + non-spec baseline run both simply report the trivial values. + """ + tokens_per_iter = [ + out.avg_decoded_tokens_per_iter + for out in outputs + if getattr(out, "avg_decoded_tokens_per_iter", None) is not None + ] + accepted = drafted = 0 + for out in outputs: + pm = getattr(out.outputs[0], "request_perf_metrics", None) + sd = getattr(pm, "speculative_decoding", None) if pm else None + if sd is not None: + accepted += sd.total_accepted_draft_tokens + drafted += sd.total_draft_tokens + if tokens_per_iter: + mean_tpi = sum(tokens_per_iter) / len(tokens_per_iter) + print( + f"[sanity] spec stats: tokens/step mean {mean_tpi:.3f} " + f"(min {min(tokens_per_iter):.3f}, " + f"max {max(tokens_per_iter):.3f}, n={len(tokens_per_iter)})" + ) + if drafted > 0: + print(f"[sanity] spec stats: acceptance {accepted}/{drafted} = {accepted / drafted:.1%}") + + +def _chosen_logprob(logprob_dict, token_id): + entry = logprob_dict.get(token_id) + if entry is None: + return None + return getattr(entry, "logprob", entry) + + +def _dump_completions(path, completions, want_logprobs): + import json + + payload = [] + for comp in completions: + entry = {"text": comp.text, "token_ids": list(comp.token_ids)} + if want_logprobs and comp.logprobs is not None: + entry["logprobs"] = [ + {str(tid): float(getattr(lp, "logprob", lp)) for tid, lp in pos.items()} + for pos in comp.logprobs + ] + payload.append(entry) + with open(path, "w") as f: + json.dump(payload, f) + print(f"[sanity] outputs dumped to {path}") + + +def _load_completions(path): + import json + from types import SimpleNamespace + + with open(path) as f: + payload = json.load(f) + loaded = [] + for entry in payload: + logprobs = None + if "logprobs" in entry: + logprobs = [{int(t): lp for t, lp in pos.items()} for pos in entry["logprobs"]] + loaded.append( + SimpleNamespace(text=entry["text"], token_ids=entry["token_ids"], logprobs=logprobs) + ) + return loaded + + +def _compare_logits_parity(base, spec, prompt, failures, tol=None, tie_tol=None): + """Tolerance-based spec-dec parity along the shared output prefix. + + Logits at position i depend only on tokens 0..i-1, so positions up to + and INCLUDING the first divergence were computed from identical + histories and must agree to rounding tolerance. At the divergence + itself the flip must be a near-tie (top-2 logprob gap < tie_tol) — + that positively identifies benign reduction-order rounding, whereas a + confident token flipping or a drifting prefix indicates a real + verification/state bug. Positions after the divergence are + uncomparable (legitimately different histories) and ignored. + + Tolerances: KIMI_K3_SPEC_LP_TOL (default 1.0) bounds shared-prefix + drift; KIMI_K3_SPEC_TIE_TOL (default 0.3) bounds the divergence gap. + On truncated NOISE models the reduction-order wobble is amplified by + the recurrent state layer-over-layer AND position-over-position + (~0.2 observed at position 1-5, ~0.4 at position 11, with + kernel-exact verify math), so the drift tolerance is calibrated as a + catastrophic-corruption net: real state bugs (wrong slot / wrong + step / missed promotion) produce O(1-10) logprob errors. Precision + correctness is owned by the exact KDA verify unit test; full-model + strict text parity is the e2e proof. + """ + if tol is None: + tol = float(os.environ.get("KIMI_K3_SPEC_LP_TOL", "1.0")) + if tie_tol is None: + tie_tol = float(os.environ.get("KIMI_K3_SPEC_TIE_TOL", "0.3")) + b_ids, s_ids = list(base.token_ids), list(spec.token_ids) + b_lp, s_lp = base.logprobs, spec.logprobs + if not b_lp or len(b_lp) != len(b_ids): + failures.append( + f"baseline for {prompt!r} has " + f"{len(b_lp) if b_lp else 0} logprob entries for " + f"{len(b_ids)} tokens; logits parity requires aligned " + "baseline logprobs (dump the baseline with " + "KIMI_K3_DUMP_LOGPROBS=1)" + ) + return "drift" + shared = 0 + while shared < min(len(b_ids), len(s_ids)) and b_ids[shared] == s_ids[shared]: + shared += 1 + # One-engine spec samplers (SA/MTP/Eagle3) do not emit per-token + # logprobs (SpecSamplerBase stores tokens only); the spec run's + # logprobs list is then shorter than its token_ids and NOT + # position-aligned. Fall back to one-sided certification: the + # divergence near-tie classification below (baseline-side logprobs) + # still applies; shared-prefix drift is only checkable when the spec + # sampler returned aligned per-token logprobs (host-drafter modes). + s_lp_aligned = bool(s_lp) and len(s_lp) == len(s_ids) + if not s_lp_aligned: + print( + f"[sanity] NOTE: spec run returned " + f"{len(s_lp) if s_lp else 0} logprob entries for " + f"{len(s_ids)} tokens (one-engine sampler); one-sided " + "parity — drift check skipped, near-tie classification " + "(baseline-side) active" + ) + for i in range(shared if s_lp_aligned else 0): + lb = _chosen_logprob(b_lp[i], b_ids[i]) + ls = _chosen_logprob(s_lp[i], s_ids[i]) + if lb is None or ls is None: + continue + if abs(lb - ls) > tol: + failures.append( + f"logit drift for {prompt!r} at shared position {i}: " + f"baseline lp={lb:.4f} vs specdec lp={ls:.4f} " + f"(|diff| > {tol})" + ) + return "drift" + if shared < min(len(b_ids), len(s_ids)): + # First divergence: expected to be a near-tie in the baseline + # distribution. The drift check above is the hard regression + # signal (prompt-agnostic, stable); a single non-tie flip is only + # a warning (borderline gaps occur legitimately on noise models), + # but the caller hard-fails when non-ties dominate in aggregate — + # a real state bug corrupts trajectories systematically. + top = sorted((getattr(v, "logprob", v) for v in b_lp[shared].values()), reverse=True) + gap = top[0] - top[1] if len(top) > 1 else float("inf") + if gap > tie_tol: + message = ( + f"non-tie divergence for {prompt!r} at position {shared}: " + f"baseline top-2 logprob gap {gap:.4f} > {tie_tol} " + f"(confident token flipped — investigate if drift also seen)" + ) + if os.environ.get("KIMI_K3_SPEC_TIE_STRICT", "0") == "1": + failures.append(message) + else: + print(f"[sanity] WARNING: {message}") + return "non_tie" + print( + f"[sanity] {prompt!r}: benign divergence at position " + f"{shared} (top-2 gap {gap:.4f}, shared prefix verified)" + ) + return "benign" + print(f"[sanity] {prompt!r}: full output identical ({shared} tokens)") + return "identical" + + +def main() -> int: + ckpt = os.environ["KIMI_K3_CKPT"] + tp = int(os.environ.get("KIMI_K3_TP", "16")) + num_layers = os.environ.get("KIMI_K3_NUM_LAYERS") + truncated = num_layers is not None + if truncated: + ckpt = _truncated_checkpoint(ckpt, int(num_layers)) + max_tokens = int(os.environ.get("KIMI_K3_MAX_TOKENS", "64")) + spec_mode = os.environ.get("KIMI_K3_SPEC_MODE", "off") + # Spec-dec parity modes (baseline greedy runs first, then spec): + # 1 / text : outputs must be BIT-IDENTICAL (use on the full model, + # where confident logits absorb kernel-rounding noise) + # logits : tolerance-based logprob comparison along the shared + # prefix; divergences must be near-ties (use on truncated + # models, where noise logits flip argmax on rounding) + spec_parity = os.environ.get("KIMI_K3_SPEC_PARITY", "0") + if spec_parity == "1": + spec_parity = "text" + if spec_parity not in ("0", "text", "logits"): + raise ValueError( + f"KIMI_K3_SPEC_PARITY={spec_parity!r} (expected '0', '1'/'text', or 'logits')" + ) + + # DEP deployment (attention data-parallel + MoE expert-parallel + # dispatch/combine; EP width == tp) is the default. KIMI_K3_ADP=0 + # selects the plain EP mode (replicated attention + latent allreduce). + adp = os.environ.get("KIMI_K3_ADP", "1") == "1" + print( + f"[sanity] ckpt={ckpt} tp(EP)={tp} adp={adp} truncated={truncated} " + f"moe_backend={os.environ.get('KIMI_K3_MOE_BACKEND', 'AUTO')} " + f"spec_mode={spec_mode} spec_parity={spec_parity}" + ) + + want_logprobs = spec_parity == "logits" or os.environ.get("KIMI_K3_DUMP_LOGPROBS", "0") == "1" + # Cross-process parity: a prior baseline run dumps its outputs via + # KIMI_K3_OUTPUT_JSON; this run loads them via KIMI_K3_BASELINE_JSON + # instead of loading a second in-process model (the full model does + # not fit twice — shutdown does not release everything). + baseline_json = os.environ.get("KIMI_K3_BASELINE_JSON") + output_json = os.environ.get("KIMI_K3_OUTPUT_JSON") + + # Extra parity-only prompts (default 0). More trajectories make the + # aggregate divergence statistics meaningful; quality checks below + # apply only to the base PROMPTS_AND_CHECKS. + # Honored whenever explicitly requested (default 0 = no change): + # baseline-dump runs (KIMI_K3_OUTPUT_JSON, parity off) need the same + # extra prompts as the parity run that will consume the dump. + extra = int(os.environ.get("KIMI_K3_SPEC_NUM_PROMPTS", "0")) + prompt_set = PROMPTS_AND_CHECKS + _parity_prompts(extra) + prompt_texts = [p for p, _ in prompt_set] + + baseline = None + if spec_parity != "0": + # Regime-match the parity baseline to the (eager) spec run. + os.environ.setdefault("KIMI_K3_CUDA_GRAPHS", "0") + # A cross-process baseline permits no-spec regression checks + # (e.g. verifying a model-class change is behavior-neutral by + # comparing two spec-off runs across code states). + assert spec_mode != "off" or baseline_json, ( + "KIMI_K3_SPEC_PARITY requires a spec mode, or a cross-process " + "baseline via KIMI_K3_BASELINE_JSON for no-spec regression " + "checks" + ) + if baseline_json: + baseline = _load_completions(baseline_json) + if len(baseline) != len(prompt_set): + raise ValueError( + f"baseline {baseline_json} holds {len(baseline)} " + f"completions but this run uses {len(prompt_set)} " + "prompts; set KIMI_K3_SPEC_NUM_PROMPTS to the value " + "used for the baseline run" + ) + else: + llm = _build_llm(ckpt, tp, "off", adp) + baseline = _generate(llm, prompt_texts, max_tokens, want_logprobs) + llm.shutdown() + + llm = _build_llm(ckpt, tp, spec_mode, adp) + completions = _generate(llm, prompt_texts, max_tokens, want_logprobs) + llm.shutdown() + if output_json: + _dump_completions(output_json, completions, want_logprobs) + + failures = [] + for comp, (prompt, expected) in zip(completions, prompt_set): + text = comp.text + print("=" * 80) + print(f"PROMPT: {prompt!r}") + print(f"OUTPUT: {text!r}") + if not truncated and expected is not None and expected not in text: + failures.append(f"expected {expected!r} in output of {prompt!r}") + if spec_parity == "text": + for base, spec, (prompt, _) in zip(baseline, completions, prompt_set): + if base.text != spec.text: + failures.append( + f"spec-dec parity mismatch for {prompt!r}:\n" + f" baseline: {base.text!r}\n" + f" specdec: {spec.text!r}" + ) + elif spec_parity == "logits": + outcomes = [ + _compare_logits_parity(base, spec, prompt, failures) + for base, spec, (prompt, _) in zip(baseline, completions, prompt_set) + ] + divergences = [o for o in outcomes if o in ("benign", "non_tie")] + non_ties = outcomes.count("non_tie") + print( + f"[sanity] logits-parity summary: {len(outcomes)} prompts, " + f"{outcomes.count('identical')} identical, " + f"{len(divergences)} divergences ({non_ties} non-tie), " + f"{outcomes.count('drift')} drift" + ) + # Aggregate systemic check: scattered non-tie flips are rounding; + # a real state bug makes them dominate. + if non_ties >= 2 and non_ties > 0.25 * max(len(divergences), 1): + failures.append( + f"non-tie divergences dominate: {non_ties}/" + f"{len(divergences)} divergences exceeded the tie bound " + "(systemic — suspect state bug)" + ) + + if failures: + print("[sanity] FAIL") + for f in failures: + print(f" - {f}") + return 1 + suffix = " (pipeline only, truncated model)" if truncated else "" + if spec_parity != "0": + suffix += f" (spec-dec {spec_parity} parity verified)" + print("[sanity] PASS" + suffix) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/integration/defs/test_kimi_k3_specdec.py b/tests/integration/defs/test_kimi_k3_specdec.py new file mode 100644 index 000000000000..df8eb2fa2151 --- /dev/null +++ b/tests/integration/defs/test_kimi_k3_specdec.py @@ -0,0 +1,95 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 SA (suffix automaton) spec-dec integration test (truncated model). + +Runs the kimi_k3_sa_harness (same directory) on the first 4 layers (KDA + +the first MLA layer) with SA spec dec and logits-parity checking: baseline +and spec logprobs must agree along the shared output prefix (hard failure +on drift — the state-bug signature), while non-tie divergences only warn +(benign reduction-order rounding flips argmax on truncated-model noise +logits; see the harness docstring). + +Requirements: 4 GPUs and the Kimi K3 checkpoint (env KIMI_K3_CKPT or +/Kimi-K3). Skips cleanly when the +checkpoint is absent. The MoE backend defaults to VANILLA (the reference +dequant path — the bit-parity oracle; slow but fine at 4 layers) so the +test has no fused-kernel dependency and runs on any arch. +""" + +import os +import subprocess +import sys + +import pytest + +from defs.conftest import llm_models_root +from defs.trt_test_alternative import print_info + + +def _find_checkpoint(): + ckpt = os.environ.get("KIMI_K3_CKPT") + if ckpt and os.path.isdir(ckpt): + return ckpt + models_root = llm_models_root() + if models_root: + candidate = os.path.join(models_root, "Kimi-K3") + if os.path.isdir(candidate): + return candidate + return None + + +@pytest.mark.skip_less_device(4) +def test_kimi_k3_sa_specdec_logits_parity(): + ckpt = _find_checkpoint() + if ckpt is None: + pytest.skip( + "Kimi K3 checkpoint not available (set KIMI_K3_CKPT or stage under LLM_MODELS_ROOT)" + ) + + env = os.environ.copy() + env.update( + { + "KIMI_K3_CKPT": ckpt, + "KIMI_K3_TP": "4", + "KIMI_K3_NUM_LAYERS": "4", + "KIMI_K3_SPEC_MODE": "sa", + "KIMI_K3_SPEC_DRAFT_LEN": "2", + "KIMI_K3_SPEC_PARITY": "logits", + # ~50 trajectories make the aggregate divergence statistics + # meaningful; loading dominates runtime so this is nearly free. + "KIMI_K3_SPEC_NUM_PROMPTS": "48", + # No KIMI_K3_MOE_BACKEND override: KimiK3MoERuntime pins the + # routed MoE backend to TRTLLM regardless of moe_config.backend, + # so a VANILLA default here would not take effect. Parity holds + # because the baseline and spec runs share the same backend. + } + ) + script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kimi_k3_sa_harness.py") + print_info(f"running {script} with ckpt={ckpt}") + result = subprocess.run( + [sys.executable, script], env=env, capture_output=True, text=True, timeout=1800 + ) + sys.stdout.write(result.stdout[-8000:]) + sys.stderr.write(result.stderr[-4000:]) + assert result.returncode == 0, "sanity harness reported FAIL" + assert "[sanity] PASS" in result.stdout + + +def test_kimi_k3_disagg_parity_selftest(): + """Self-test of the two-endpoint (aggregated vs disagg) parity harness. + + Comparison logic only: canned responses, no servers or GPUs. + """ + script = os.path.join(os.path.dirname(os.path.abspath(__file__)), "kimi_k3_disagg_parity.py") + if not os.path.exists(script): + pytest.skip( + "kimi_k3_disagg_parity.py harness not present on this branch " + "(ships with the disagg parity PR)" + ) + result = subprocess.run( + [sys.executable, script, "--self-test"], capture_output=True, text=True, timeout=120 + ) + sys.stdout.write(result.stdout[-4000:]) + sys.stderr.write(result.stderr[-2000:]) + assert result.returncode == 0, "parity harness self-test FAILED" + assert "[self-test] PASS" in result.stdout diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_accept_stats.py b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_accept_stats.py new file mode 100644 index 000000000000..33f37709d6a1 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_dflash_accept_stats.py @@ -0,0 +1,271 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the DFlash/DSpark acceptance-statistics aggregation +math (tensorrt_llm/_torch/speculative/accept_stats.py). + +accept_stats.py is import-light (no torch), so these tests also run on +hosts without a GPU stack; the fallback loader below keeps them runnable +even where the tensorrt_llm package itself cannot be imported. +""" + +import json +import os + +import pytest + +try: + from tensorrt_llm._torch.speculative import accept_stats +except ImportError: # pragma: no cover - torch-less host fallback + import importlib.util + + _path = os.path.join( + os.path.dirname(__file__), + *([os.pardir] * 5), + "tensorrt_llm", + "_torch", + "speculative", + "accept_stats.py", + ) + _spec = importlib.util.spec_from_file_location("accept_stats", os.path.abspath(_path)) + accept_stats = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(accept_stats) + + +def test_summarize_hist_basic(): + # K = 3; 10 steps: 4 steps accepted 0 drafts, 3 accepted 1, + # 2 accepted 2, 1 accepted 3. + hist = [4, 3, 2, 1] + s = accept_stats.summarize_hist(hist) + assert s["num_steps"] == 10 + assert s["mean_accepted_draft"] == pytest.approx(1.0) + assert s["al"] == pytest.approx(2.0) + # AR_k = P(accepted_draft >= k): [6/10, 3/10, 1/10] + assert s["ar_per_position"] == pytest.approx([0.6, 0.3, 0.1]) + # Monotone non-increasing by construction (prefix acceptance). + ar = s["ar_per_position"] + assert all(a >= b for a, b in zip(ar, ar[1:])) + + +def test_summarize_hist_empty_and_dummy_regimes(): + assert accept_stats.summarize_hist([0, 0, 0])["al"] == 0.0 + # Dummy-drafter regime: every step accepts 0 drafts -> AL exactly 1.0 + # (bonus token only), AR 0 at every position. + s = accept_stats.summarize_hist([100, 0, 0, 0, 0, 0, 0, 0]) + assert s["al"] == pytest.approx(1.0) + assert s["ar_per_position"] == pytest.approx([0.0] * 7) + # Perfect drafter: every step accepts all K=7 drafts -> AL = 8. + s = accept_stats.summarize_hist([0, 0, 0, 0, 0, 0, 0, 50]) + assert s["al"] == pytest.approx(8.0) + assert s["ar_per_position"] == pytest.approx([1.0] * 7) + + +def test_recorder_accumulation_and_flush(tmp_path): + rec = accept_stats.DFlashAcceptStatsRecorder( + str(tmp_path), max_draft_len=3, rank=0, flush_every=1000, num_conf_bins=4 + ) + # Step 1: req 7 accepts bonus+2 drafts, req 8 bonus only. + rec.on_draft_confidence([7, 8], [[0.9, 0.6, 0.1], [0.2, 0.2, 0.2]]) + rec.on_accept([7, 8], [3, 1]) + # Step 2: req 7 accepts everything (bonus + 3 drafts, clamped input 5). + rec.on_accept([7], [5]) + + snap = rec.snapshot() + assert snap["num_steps"] == 2 + assert snap["accepted_draft_hist"] == [1, 0, 1, 1] + assert snap["per_request"]["7"] == {"steps": 2, "accepted_draft": 5} + assert snap["per_request"]["8"] == {"steps": 1, "accepted_draft": 0} + + cc = snap["confidence_calibration"] + # req 7: conf (0.9, 0.6, 0.1) -> bins (3, 2, 0); accepted at pos 1,2 only. + assert cc["attempts"][0][3] == 1 and cc["accepted"][0][3] == 1 + assert cc["attempts"][1][2] == 1 and cc["accepted"][1][2] == 1 + assert cc["attempts"][2][0] == 2 and cc["accepted"][2][0] == 0 + # req 8: conf 0.2 -> bin 0 at each position, nothing accepted. + assert cc["attempts"][0][0] == 1 and cc["accepted"][0][0] == 0 + # Step 2 had no pending confidences for req 7 (consumed in step 1). + assert sum(sum(r) for r in cc["attempts"]) == 6 + + rec.flush() + with open(rec.path) as f: + assert json.load(f) == snap + + +def test_merge_snapshots_and_load(tmp_path): + r0 = accept_stats.DFlashAcceptStatsRecorder(str(tmp_path), 2, rank=0, num_conf_bins=4) + r1 = accept_stats.DFlashAcceptStatsRecorder(str(tmp_path), 2, rank=1, num_conf_bins=4) + r0.on_draft_confidence([1], [[0.9, 0.9]]) + r0.on_accept([1], [3]) # 2 drafts accepted + r1.on_accept([2], [1]) # 0 drafts accepted + r0.flush() + r1.flush() + + snaps = accept_stats.load_rank_snapshots(str(tmp_path)) + assert len(snaps) == 2 + merged = accept_stats.merge_snapshots(snaps) + assert merged["num_steps"] == 2 + assert merged["accepted_draft_hist"] == [1, 0, 1] + assert merged["per_request"]["1"]["accepted_draft"] == 2 + assert merged["per_request"]["2"]["steps"] == 1 + s = accept_stats.summarize_hist(merged["accepted_draft_hist"]) + assert s["al"] == pytest.approx(2.0) + assert s["ar_per_position"] == pytest.approx([0.5, 0.5]) + + with pytest.raises(ValueError): + accept_stats.merge_snapshots([]) + + +def test_merge_snapshots_shape_mismatch(tmp_path): + a = accept_stats.DFlashAcceptStatsRecorder(str(tmp_path), 2, rank=0).snapshot() + b = accept_stats.DFlashAcceptStatsRecorder(str(tmp_path), 3, rank=1).snapshot() + with pytest.raises(ValueError): + accept_stats.merge_snapshots([a, b]) + + +def test_calibration_table(): + # One position, 4 bins; perfectly calibrated in bins 1 and 3. + attempts = [[0, 10, 0, 8]] + accepted = [[0, 4, 0, 7]] # empirical 0.4 vs center 0.375; 0.875 = center + table = accept_stats.calibration_table(attempts, accepted) + assert table["bin_centers"] == pytest.approx([0.125, 0.375, 0.625, 0.875]) + pos = table["per_position"][0] + assert pos["empirical_acceptance"][0] is None + assert pos["empirical_acceptance"][1] == pytest.approx(0.4) + assert pos["empirical_acceptance"][3] == pytest.approx(0.875) + # ECE = (10*|0.375-0.4| + 8*|0.875-0.875|) / 18 + assert pos["ece"] == pytest.approx(10 * 0.025 / 18) + + +def test_recorder_excludes_dummy_request_id_zero(tmp_path): + # Executor-warmup dummies and idle attention-DP ranks' padding requests + # carry request id 0 and must not pollute any counter. + rec = accept_stats.DFlashAcceptStatsRecorder( + str(tmp_path), max_draft_len=3, rank=0, num_conf_bins=4 + ) + rec.on_draft_confidence([0, 5], [[0.9, 0.9, 0.9], [0.4, 0.4, 0.4]]) + rec.on_accept([0, 5], [4, 2]) # id 0: all-K padding "accept" -> dropped + snap = rec.snapshot() + assert snap["accepted_draft_hist"] == [0, 1, 0, 0] + assert list(snap["per_request"]) == ["5"] + assert sum(sum(r) for r in snap["confidence_calibration"]["attempts"]) == 3 + + +def test_confidence_provider_none_collects_nothing(tmp_path): + # Default recorder (no provider): record_draft_confidence is a no-op, + # the calibration table stays empty, AL/AR still accumulate. + rec = accept_stats.DFlashAcceptStatsRecorder( + str(tmp_path), max_draft_len=2, rank=0, num_conf_bins=4 + ) + assert rec.confidence_provider is None + rec.record_draft_confidence([1], "model", "hidden", "prev", "drafts") + rec.on_accept([1], [2]) + snap = rec.snapshot() + assert snap["accepted_draft_hist"] == [0, 1, 0] + assert sum(sum(r) for r in snap["confidence_calibration"]["attempts"]) == 0 + + +def test_confidence_provider_decline_and_forwarding(tmp_path): + # A provider may decline a step by returning None; the recorder must + # forward the draft-site arguments verbatim (they are opaque to it). + calls = [] + + def declining_provider(*args): + calls.append(args) + return None + + rec = accept_stats.DFlashAcceptStatsRecorder( + str(tmp_path), + max_draft_len=2, + rank=0, + num_conf_bins=4, + confidence_provider=declining_provider, + ) + sentinel = ("model", "hidden", "prev", "drafts") + rec.record_draft_confidence([1], *sentinel) + assert calls == [sentinel] + assert not rec._pending_conf + + +def test_mocked_calibration_end_to_end(tmp_path): + """Integration: scripted mock provider + scripted accept outcomes, + driven through the recorder exactly as the DFlash worker drives it + (verify previous block via on_accept, then draft a new block via + record_draft_confidence), then flushed / loaded / merged / tabulated. + Asserts exact bin counts and calibration-table contents.""" + K, NBINS = 3, 4 # bins: [0,.25) [.25,.5) [.5,.75) [.75,1] + + scripted_rows = [ + # step 1 block: rows for requests (1, 2) + [[0.9, 0.6, 0.1], [0.3, 0.3, 0.3]], + # step 2 block + [[0.8, 0.2, 0.55], [0.7, 0.1, 0.9]], + ] + provider_calls = [] + + def mock_provider(draft_model, gen_hidden, first_prev_tokens, gen_draft_tokens): + provider_calls.append(draft_model) + return scripted_rows[len(provider_calls) - 1] + + rec = accept_stats.DFlashAcceptStatsRecorder( + str(tmp_path), + max_draft_len=K, + rank=0, + flush_every=1000, + num_conf_bins=NBINS, + confidence_provider=mock_provider, + ) + + rids = [1, 2] + # Step 1: nothing pending yet; both requests accept 0 drafts (n=1). + rec.on_accept(rids, [1, 1]) + rec.record_draft_confidence(rids, "m", "h", "p", "d") + # Step 2: req 1 accepts 2 drafts (n=3), req 2 accepts 0 (n=1). + rec.on_accept(rids, [3, 1]) + rec.record_draft_confidence(rids, "m", "h", "p", "d") + # Step 3: req 1 accepts 1 draft (n=2), req 2 accepts all 3 (n=4). + rec.on_accept(rids, [2, 4]) + + assert len(provider_calls) == 2 + rec.flush() + + snaps = accept_stats.load_rank_snapshots(str(tmp_path)) + merged = accept_stats.merge_snapshots(snaps) + assert merged["num_steps"] == 3 + assert merged["accepted_draft_hist"] == [3, 1, 1, 1] + + cc = merged["confidence_calibration"] + # Exact per-(position, bin) attempt/accept counts (bin = int(4c)): + # step-1 rows joined with step-2 outcomes, step-2 rows with step-3. + assert cc["attempts"] == [ + [0, 1, 1, 2], # pos 1: 0.3 | 0.7 | 0.9, 0.8 + [2, 1, 1, 0], # pos 2: 0.2, 0.1 | 0.3 | 0.6 + [1, 1, 1, 1], # pos 3: 0.1 | 0.3 | 0.55 | 0.9 + ] + assert cc["accepted"] == [ + [0, 0, 1, 2], + [1, 0, 1, 0], + [0, 0, 0, 1], + ] + + table = accept_stats.calibration_table(cc["attempts"], cc["accepted"]) + assert table["bin_centers"] == pytest.approx([0.125, 0.375, 0.625, 0.875]) + p1, p2, p3 = table["per_position"] + assert p1["empirical_acceptance"] == [None, 0.0, 1.0, 1.0] + assert p2["empirical_acceptance"] == [0.5, 0.0, 1.0, None] + assert p3["empirical_acceptance"] == [0.0, 0.0, 0.0, 1.0] + assert p1["num_samples"] == [0, 1, 1, 2] + # pos-1 ECE: (1*|.375-0| + 1*|.625-1| + 2*|.875-1|) / 4 + assert p1["ece"] == pytest.approx((0.375 + 0.375 + 2 * 0.125) / 4) + + +def test_default_provider_absent_in_core_tree(): + # The optional dspark_confidence module ships in a separate MR; on this + # tree the default provider must resolve to None (not raise). + assert accept_stats._resolve_default_confidence_provider() is None + + +def test_conf_bin_edges(): + assert accept_stats._conf_bin(-0.1, 20) == 0 + assert accept_stats._conf_bin(0.0, 20) == 0 + assert accept_stats._conf_bin(0.999, 20) == 19 + assert accept_stats._conf_bin(1.0, 20) == 19 + assert accept_stats._conf_bin(1.5, 20) == 19 diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py new file mode 100644 index 000000000000..ef778488ad8a --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dflash_scaffold.py @@ -0,0 +1,344 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 DFlash structural scaffold: config schema + capture plumbing. + +Mostly CPU-only (no checkpoint): validates that the synthetic drafter +generator emits the exact K2.7-Code-DFlash checkpoint schema the generic +DFlashForCausalLM loader expects, that DFlashDecodingConfig resolves to the +predicates the K3 one-engine wrapper relies on, and that KimiLinearModel +threads spec_metadata for per-layer hidden-state capture. The capture-buffer +round-trip needs a CUDA device (DFlashSpecMetadata allocates its buffer on +cuda) and is skip-guarded. +""" + +import importlib.util +import inspect +import json +import os + +import pytest +import torch + +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode + + +def _load_generator(): + path = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "..", + "..", + "examples", + "kimi_k3", + "make_synthetic_dflash_drafter.py", + ) + ) + if not os.path.exists(path): + # The synthetic-drafter generator ships with the examples/kimi_k3 + # PR; on branches without it, the schema tests below skip. + return None + spec = importlib.util.spec_from_file_location("make_synthetic_dflash_drafter", path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +GEN = _load_generator() + +requires_generator = pytest.mark.skipif( + GEN is None, + reason="examples/kimi_k3/make_synthetic_dflash_drafter.py not present on this branch", +) + +# Reference: nvidia/Kimi-K2.7-Code-DFlash (69 tensors). The generator must +# reproduce this key/shape schema exactly at K2.7 dims — the generic +# DFlashForCausalLM.load_weights() key remapping is written against it. +K27_HIDDEN = 7168 +K27_FC_SHAPE = (7168, 43008) # hidden x (hidden * 6 capture layers) +K27_NUM_TENSORS = 69 +K27_TARGET_LAYER_IDS = [1, 12, 24, 35, 47, 58] # over 61 target layers + +DFLASH_EXTRA_KEYS = {"fc.weight", "hidden_norm.weight"} +PER_LAYER_SUFFIXES = { + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + "mlp.gate_proj.weight", + "mlp.up_proj.weight", + "mlp.down_proj.weight", + "input_layernorm.weight", + "post_attention_layernorm.weight", +} + + +@requires_generator +def test_even_spacing_matches_k27_reference(): + assert GEN.even_target_layer_ids(61, 6) == K27_TARGET_LAYER_IDS + + +@requires_generator +def test_tensor_plan_matches_k27_schema(): + plan = GEN.drafter_tensor_plan(K27_HIDDEN, GEN.K27_DRAFTER, 6) + assert len(plan) == K27_NUM_TENSORS + assert plan["fc.weight"] == K27_FC_SHAPE + keys = set(plan) + assert DFLASH_EXTRA_KEYS | {"norm.weight"} <= keys + for i in range(GEN.K27_DRAFTER["num_hidden_layers"]): + for suffix in PER_LAYER_SUFFIXES: + assert f"layers.{i}.{suffix}" in keys + # No embeddings or lm_head: shared with the target model. + assert not any("embed" in k or "lm_head" in k for k in keys) + + +# Reference: the K3 DSpark drafter dummy checkpoint from the training team +# (dummy-dspark0724, 2026-07-24; 73 tensors = K2.7's 69 + markov_w1/w2 + +# confidence_proj.{weight,bias}). Dims from its config.json. +K3_HIDDEN = 7168 +K3_VOCAB = 163840 +K3_NUM_TARGET_LAYERS = 93 +K3_TARGET_LAYER_IDS = [1, 19, 37, 54, 72, 90] +K3_MASK_TOKEN_ID = 163606 # NOT vocab-2 +K3_MARKOV_RANK = 256 +K3_NUM_TENSORS = 73 + + +@requires_generator +def test_even_spacing_matches_real_k3_drafter(): + """The real config's target_layer_ids follow the even-spacing convention.""" + assert GEN.even_target_layer_ids(K3_NUM_TARGET_LAYERS, 6) == K3_TARGET_LAYER_IDS + + +@requires_generator +def test_tensor_plan_matches_real_dspark_schema(): + """Plan must match the dummy-dspark0724 safetensors header exactly.""" + plan = GEN.drafter_tensor_plan( + K3_HIDDEN, + GEN.K3_DRAFTER, + len(K3_TARGET_LAYER_IDS), + vocab=K3_VOCAB, + markov_rank=K3_MARKOV_RANK, + use_confidence_head=True, + ) + assert len(plan) == K3_NUM_TENSORS + assert plan["fc.weight"] == (7168, 43008) + assert plan["markov_w1.weight"] == (K3_VOCAB, K3_MARKOV_RANK) + assert plan["markov_w2.weight"] == (K3_VOCAB, K3_MARKOV_RANK) + # Confidence head reads [hidden, markov_features] concat -> in-dim 7424. + assert plan["confidence_proj.weight"] == (1, K3_HIDDEN + K3_MARKOV_RANK) + assert plan["confidence_proj.bias"] == (1,) + # Per-layer backbone: K3 drafter has 32 Q heads / 8 KV heads / hd 128. + assert plan["layers.0.self_attn.q_proj.weight"] == (4096, 7168) + assert plan["layers.0.self_attn.k_proj.weight"] == (1024, 7168) + assert plan["layers.0.mlp.gate_proj.weight"] == (12288, 7168) + assert not any("embed" in k or "lm_head" in k for k in plan) + + +@requires_generator +def test_k3_drafter_config_is_dspark(): + cfg = GEN.drafter_config( + K3_HIDDEN, + K3_VOCAB, + K3_NUM_TARGET_LAYERS, + K3_TARGET_LAYER_IDS, + K3_MASK_TOKEN_ID, + GEN.K3_DRAFTER, + ) + assert cfg["architectures"] == ["DFlashDraftModel"] + assert cfg["model_type"] == "qwen3" + dflash_cfg = cfg["dflash_config"] + assert dflash_cfg["mask_token_id"] == K3_MASK_TOKEN_ID + assert dflash_cfg["target_layer_ids"] == K3_TARGET_LAYER_IDS + # DSpark extras, exactly as the real config declares them. + assert dflash_cfg["projector_type"] == "dspark" + assert dflash_cfg["causal"] is False + assert dflash_cfg["use_swa"] is True and dflash_cfg["swa_window_size"] == 1024 + assert dflash_cfg["shift_label"] is True + assert dflash_cfg["markov_rank"] == K3_MARKOV_RANK + assert dflash_cfg["markov_head_type"] == "vanilla" + assert dflash_cfg["use_confidence_head"] is True + assert cfg["layer_types"] == ["sliding_attention"] * 6 + assert cfg["sliding_window"] == 1024 + assert cfg["rope_theta"] == 10000.0 and cfg["rope_scaling"] is None + + +@requires_generator +def test_drafter_config_drives_generic_dflash_path(): + cfg = GEN.drafter_config(K27_HIDDEN, 163840, 61, K27_TARGET_LAYER_IDS, 163838, GEN.K27_DRAFTER) + # Unknown architecture label + model_type=qwen3 selects the generic + # (non-Laguna) DFlashForCausalLM with a Qwen3ForCausalLM backbone. + assert cfg["architectures"] == ["DFlashDraftModel"] + assert not any("Laguna" in a for a in cfg["architectures"]) + assert cfg["model_type"] == "qwen3" + assert cfg["dflash_config"] == { + "mask_token_id": 163838, + "target_layer_ids": K27_TARGET_LAYER_IDS, + } + assert cfg["synthetic_random_weights"] is True + + +@requires_generator +def test_generator_tiny_roundtrip(tmp_path): + import subprocess + import sys + + from safetensors import safe_open + + out = tmp_path / "tiny_dflash" + subprocess.run( + [sys.executable, GEN.__file__, "--tiny", "--out", str(out)], + check=True, + ) + with open(out / "config.json") as f: + cfg = json.load(f) + dflash_cfg = cfg["dflash_config"] + plan = GEN.drafter_tensor_plan( + cfg["hidden_size"], + GEN.TINY, + len(dflash_cfg["target_layer_ids"]), + vocab=cfg["vocab_size"], + markov_rank=GEN.TINY["markov_rank"], + use_confidence_head=True, + ) + with safe_open(out / "model.safetensors", "pt") as f: + keys = set(f.keys()) + assert keys == set(plan) + for k in keys: + assert tuple(f.get_tensor(k).shape) == tuple(plan[k]) + assert dflash_cfg["mask_token_id"] == cfg["vocab_size"] - 2 + assert all(0 <= t < cfg["num_target_layers"] for t in dflash_cfg["target_layer_ids"]) + + +@requires_generator +def test_generator_real_config_mode(tmp_path): + """--config adopts a real drafter config.json verbatim (random weights, + exact real module structure).""" + import subprocess + import sys + + from safetensors import safe_open + + real_cfg = { + "architectures": ["DFlashDraftModel"], + "model_type": "qwen3", + "block_size": 4, + "dflash_config": { + "mask_token_id": 510, + "target_layer_ids": [2, 5], + "projector_type": "dspark", + "markov_rank": 16, + "use_confidence_head": True, + }, + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "intermediate_size": 128, + "vocab_size": 512, + "rope_theta": 12345.0, # pass-through field the generator must keep + "num_target_layers": 8, + } + cfg_path = tmp_path / "real_config.json" + cfg_path.write_text(json.dumps(real_cfg)) + out = tmp_path / "synth" + subprocess.run( + [sys.executable, GEN.__file__, "--config", str(cfg_path), "--out", str(out)], + check=True, + ) + with open(out / "config.json") as f: + emitted = json.load(f) + # Verbatim adoption plus the synthetic marker. + assert emitted["dflash_config"] == real_cfg["dflash_config"] + assert emitted["rope_theta"] == 12345.0 + assert emitted["synthetic_random_weights"] is True + plan = GEN.drafter_tensor_plan( + real_cfg["hidden_size"], + dict( + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + intermediate_size=128, + block_size=4, + ), + len(real_cfg["dflash_config"]["target_layer_ids"]), + vocab=real_cfg["vocab_size"], + markov_rank=16, + use_confidence_head=True, + ) + assert { + "markov_w1.weight", + "markov_w2.weight", + "confidence_proj.weight", + "confidence_proj.bias", + } <= set(plan) + with safe_open(out / "model.safetensors", "pt") as f: + assert set(f.keys()) == set(plan) + + +def test_dflash_decoding_config_predicates(): + from tensorrt_llm.llmapi import DFlashDecodingConfig + + cfg = DFlashDecodingConfig(max_draft_len=7) + mode = cfg.spec_dec_mode + assert mode == SpeculativeDecodingMode.DFLASH + assert mode.is_dflash() + # The K3 one-engine wrapper (SpecDecOneEngineForCausalLM) relies on + # these to attach the external drafter and its capture metadata. + assert mode.use_one_engine() + assert mode.is_external_drafter() + # KimiLinearForCausalLM admission: SA or DFlash only. + assert mode.is_sa() or mode.is_dflash() + assert not SpeculativeDecodingMode.NGRAM.is_sa() + assert not SpeculativeDecodingMode.NGRAM.is_dflash() + + +def test_kimi_linear_model_threads_spec_metadata(): + pytest.importorskip("fla") + from tensorrt_llm._torch.models.modeling_kimi_linear import KimiLinearModel + + # The capture hook only fires if spec_metadata is an explicit parameter + # (previously it was silently swallowed by **kwargs). + params = inspect.signature(KimiLinearModel.forward).parameters + assert "spec_metadata" in params + src = inspect.getsource(KimiLinearModel.forward) + assert "maybe_capture_hidden_states" in src + + +@pytest.mark.skipif( + not torch.cuda.is_available(), reason="DFlashSpecMetadata allocates capture buffer on cuda" +) +def test_dflash_spec_metadata_capture_prefix_sum_convention(): + """K3 passes the full post-layer prefix sum with residual=None; the + buffer must then hold the hidden state verbatim (no residual add).""" + from tensorrt_llm._torch.speculative.dflash import DFlashSpecMetadata + + hidden_size, max_tokens = 16, 8 + md = DFlashSpecMetadata( + max_draft_len=4, + max_total_draft_tokens=4, + spec_dec_mode=SpeculativeDecodingMode.DFLASH, + max_num_requests=2, + max_num_tokens=max_tokens, + hidden_size=hidden_size, + layers_to_capture=[1, 3], + dtype=torch.bfloat16, + ) + assert md.is_layer_capture(1) and md.is_layer_capture(3) + assert not md.is_layer_capture(2) + + h1 = torch.randn(max_tokens, hidden_size, dtype=torch.bfloat16, device="cuda") + h3 = torch.randn_like(h1) + md.maybe_capture_hidden_states(1, h1, None) + md.maybe_capture_hidden_states(2, torch.randn_like(h1), None) # no-op + md.maybe_capture_hidden_states(3, h3, None) + captured = md.get_hidden_states(max_tokens) + assert captured.shape == (max_tokens, 2 * hidden_size) + torch.testing.assert_close(captured[:, :hidden_size], h1) + torch.testing.assert_close(captured[:, hidden_size:], h3) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py new file mode 100644 index 000000000000..a2cbae410d27 --- /dev/null +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_kimi_k3_dspark_semantics.py @@ -0,0 +1,494 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Kimi K3 DSpark drafter-forward semantics (MR A). + +Validates the weights-independent dspark math against a port of the +DeepSpec reference (github.com/deepseek-ai/DeepSpec): + - vanilla Markov intra-block logit bias (deepspec/modeling/dspark/ + markov_head.py VanillaMarkov, greedy chain of sample_block_tokens), + - the shift_label output convention (block slot j predicts draft token + j+1; deepspec/eval/dspark/draft_ops.py build_dspark_proposal), + - sliding-window attention on 'sliding_attention' draft layers + (HF flash translation: window_size = (w-1, w-1), non-causal), +and the no-regression property: a dflash_config WITHOUT dspark fields +resolves to the exact pre-dspark behavior (slots 1..K, no window, no +Markov bias). CPU-only where possible; the tiny end-to-end block-decode +parity test needs CUDA + flash_attn and is skip-guarded. + +Confidence-scheduled verification is MR B: here we only check that +confidence_proj weights load without being used. +""" + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.models.modeling_speculative import ( + DFlashForCausalLM, + dspark_layer_window_size, + dspark_markov_chain_logits, + dspark_markov_step_bias, +) +from tensorrt_llm._torch.speculative.dflash import dflash_draft_slot_ids + +# --------------------------------------------------------------------------- +# Reference oracle: line-for-line port of DeepSpec VanillaMarkov +# (deepspec/modeling/dspark/markov_head.py) at temperature 0. +# --------------------------------------------------------------------------- + + +class _RefVanillaMarkov: + def __init__(self, markov_w1: torch.Tensor, markov_w2: torch.Tensor): + # markov_w1: nn.Embedding(vocab, rank).weight -> [vocab, rank] + # markov_w2: nn.Linear(rank, vocab, bias=False).weight -> [vocab, rank] + self.w1 = markov_w1 + self.w2 = markov_w2 + + def compute_step_bias(self, token_ids: torch.Tensor) -> torch.Tensor: + # markov_w2(markov_w1(ids)) + return F.linear(F.embedding(token_ids.long(), self.w1), self.w2) + + def sample_block_tokens(self, base_logits: torch.Tensor, first_prev_token_ids: torch.Tensor): + """Greedy (temperature 0) reference chain.""" + sampled, corrected = [], [] + prev = first_prev_token_ids.long() + for step in range(base_logits.shape[1]): + step_logits = base_logits[:, step, :] + self.compute_step_bias(prev) + corrected.append(step_logits.unsqueeze(1)) + nxt = torch.argmax(step_logits, dim=-1) + sampled.append(nxt) + prev = nxt + return torch.stack(sampled, dim=1), torch.cat(corrected, dim=1) + + +VOCAB, RANK, B, K = 512, 16, 3, 7 + + +def _random_markov(seed=1234, dtype=torch.float32): + g = torch.Generator().manual_seed(seed) + w1 = torch.randn(VOCAB, RANK, generator=g, dtype=dtype) + w2 = torch.randn(VOCAB, RANK, generator=g, dtype=dtype) + base = torch.randn(B, K, VOCAB, generator=g, dtype=dtype) + anchor = torch.randint(0, VOCAB, (B,), generator=g) + return w1, w2, base, anchor + + +def test_markov_step_bias_formula(): + """bias over vocab = markov_w1[prev] @ markov_w2.T (both [vocab, rank]).""" + w1, w2, _, anchor = _random_markov() + bias = dspark_markov_step_bias(anchor, w1, w2) + expected = w1[anchor] @ w2.T + torch.testing.assert_close(bias, expected) + + +def test_markov_chain_matches_deepspec_reference(): + """Corrected block logits and the greedy token chain match the ported + DeepSpec VanillaMarkov.sample_block_tokens bitwise (same dtype/ops).""" + w1, w2, base, anchor = _random_markov() + ref_tokens, ref_logits = _RefVanillaMarkov(w1, w2).sample_block_tokens(base, anchor) + + out = dspark_markov_chain_logits(base, anchor, w1, w2) + assert torch.equal(out, ref_logits) + # Greedy per-position argmax of the corrected logits reproduces the + # reference sequentially-sampled chain (what sample_draft_tokens does). + assert torch.equal(torch.argmax(out, dim=-1), ref_tokens) + + +def test_markov_chain_empty_block_is_noop(): + w1, w2, base, anchor = _random_markov() + empty = base[:, :0, :] + assert dspark_markov_chain_logits(empty, anchor, w1, w2) is empty + + +def test_markov_chain_sharded_matches_full_vocab(): + """The DFlashWorker TP path — every rank runs the chain on its + contiguous markov_w2/logits vocab shard, chained through a global + argmax over all shards — reassembles to the full-vocab chain.""" + w1, w2, base, anchor = _random_markov() + full = dspark_markov_chain_logits(base, anchor, w1, w2) + + tp = 4 + shard_w = VOCAB // tp + shards = [slice(r * shard_w, (r + 1) * shard_w) for r in range(tp)] + + # Lockstep emulation: per-rank shard bias, "TP gather" = global argmax + # across the concatenated shards (what greedy_sample_draft_with_tp_gather + # computes), returning full-vocab ids for the next markov_w1 lookup. + prev = anchor.long() + rank_outputs = [[] for _ in range(tp)] + for i in range(K): + step_shards = [] + for r, sl in enumerate(shards): + bias = dspark_markov_step_bias(prev, w1, w2[sl]) + step = base[:, i, sl] + bias + rank_outputs[r].append(step) + step_shards.append(step) + prev = torch.argmax(torch.cat(step_shards, dim=-1), dim=-1) + + reassembled = torch.cat([torch.stack(rank_outputs[r], dim=1) for r in range(tp)], dim=-1) + torch.testing.assert_close(reassembled, full) + + +# --------------------------------------------------------------------------- +# shift_label slot convention +# --------------------------------------------------------------------------- + + +def test_slot_ids_plain_dflash_matches_old_formula(): + """No-regression: shift_label off reproduces the previous inline + formula (mask slots 1..K).""" + num_gens, block, k = 3, 8, 7 + ids = dflash_draft_slot_ids(num_gens, block, k, False, device="cpu") + bases = torch.arange(num_gens, dtype=torch.long) * block + offs = torch.arange(k, dtype=torch.long) + old = (bases.unsqueeze(1) + 1 + offs.unsqueeze(0)).flatten() + assert torch.equal(ids, old) + + +def test_slot_ids_shift_label_uses_anchor_slot(): + """DSpark shift_label: slots 0..K-1; slot 0 (anchor token slot) + predicts the first draft token (DeepSpec build_dspark_proposal reads + block_hidden[:, :block_size]).""" + ids = dflash_draft_slot_ids(2, 8, 8, True, device="cpu") + assert ids.tolist() == list(range(8)) + [8 + j for j in range(8)] + # With shift_label, K == block_size stays in range (plain would not). + assert ids.max().item() == 2 * 8 - 1 + + +# --------------------------------------------------------------------------- +# SWA window convention +# --------------------------------------------------------------------------- + + +def test_swa_window_conventions(): + sliding = ["sliding_attention", "full_attention"] + # HF flash translation: window_size = (w-1, w-1) on sliding layers. + assert dspark_layer_window_size(True, 1024, sliding, 0) == (1023, 1023) + assert dspark_layer_window_size(True, 1024, sliding, 1) == (-1, -1) + # use_swa off -> flash-attn default regardless of layer_types. + assert dspark_layer_window_size(False, 1024, sliding, 0) == (-1, -1) + # No layer_types declared + use_swa -> window on every layer. + assert dspark_layer_window_size(True, 8, None, 0) == (7, 7) + + +# --------------------------------------------------------------------------- +# Tiny end-to-end drafter: config parsing, weight loading, block-decode +# parity vs an fp32 eager oracle (needs CUDA + flash_attn). +# --------------------------------------------------------------------------- + +TINY = dict( + architectures=["DFlashDraftModel"], + model_type="qwen3", + block_size=4, + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + # 128 = the real K3 drafter head_dim; small head dims are rejected by + # the fusedQKNormRope kernel the bf16 block decode uses. + head_dim=128, + intermediate_size=128, + hidden_act="silu", + rms_norm_eps=1e-6, + vocab_size=VOCAB, + max_position_embeddings=2048, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + torch_dtype="bfloat16", + num_target_layers=4, + tie_word_embeddings=False, +) + +SWA_WINDOW = 8 +CTX_LEN = 24 # > SWA_WINDOW so the window binds +NUM_CAPTURE = 2 + + +def _tiny_config(dspark: bool): + from transformers import Qwen3Config + + cfg = dict(TINY) + dflash = {"mask_token_id": VOCAB - 2, "target_layer_ids": [0, 1]} + if dspark: + dflash.update( + projector_type="dspark", + causal=False, + use_swa=True, + swa_window_size=SWA_WINDOW, + shift_label=True, + markov_rank=RANK, + markov_head_type="vanilla", + use_confidence_head=True, + ) + cfg["layer_types"] = ["sliding_attention"] * cfg["num_hidden_layers"] + cfg["sliding_window"] = SWA_WINDOW + cfg["dflash_config"] = dflash + return Qwen3Config.from_dict(cfg) + + +def _tiny_weights(seed=7): + g = torch.Generator().manual_seed(seed) + + def rnd(*shape): + return (torch.randn(*shape, generator=g) * 0.05).to(torch.bfloat16) + + h, inter = TINY["hidden_size"], TINY["intermediate_size"] + nh, nkv, hd = (TINY["num_attention_heads"], TINY["num_key_value_heads"], TINY["head_dim"]) + w = { + "fc.weight": rnd(h, h * NUM_CAPTURE), + "hidden_norm.weight": rnd(h) + 1.0, + "norm.weight": rnd(h) + 1.0, + "markov_w1.weight": rnd(VOCAB, RANK), + "markov_w2.weight": rnd(VOCAB, RANK), + "confidence_proj.weight": rnd(1, h + RANK), + "confidence_proj.bias": rnd(1), + } + for i in range(TINY["num_hidden_layers"]): + p = f"layers.{i}." + w[p + "self_attn.q_proj.weight"] = rnd(nh * hd, h) + w[p + "self_attn.k_proj.weight"] = rnd(nkv * hd, h) + w[p + "self_attn.v_proj.weight"] = rnd(nkv * hd, h) + w[p + "self_attn.o_proj.weight"] = rnd(h, nh * hd) + w[p + "self_attn.q_norm.weight"] = rnd(hd) + 1.0 + w[p + "self_attn.k_norm.weight"] = rnd(hd) + 1.0 + w[p + "input_layernorm.weight"] = rnd(h) + 1.0 + w[p + "post_attention_layernorm.weight"] = rnd(h) + 1.0 + w[p + "mlp.gate_proj.weight"] = rnd(inter, h) + w[p + "mlp.up_proj.weight"] = rnd(inter, h) + w[p + "mlp.down_proj.weight"] = rnd(h, inter) + return w + + +def _build_drafter(dspark: bool, weights): + from tensorrt_llm._torch.model_config import ModelConfig + + model_config = ModelConfig(pretrained_config=_tiny_config(dspark), attn_backend="TRTLLM") + drafter = DFlashForCausalLM(model_config).to("cuda") + # Drop dspark head tensors for the plain drafter (schema without them). + if not dspark: + weights = {k: v for k, v in weights.items() if not k.startswith(("markov_", "confidence_"))} + drafter.load_weights(dict(weights)) + return drafter + + +def _rms(x, w, eps=1e-6): + xf = x.float() + return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)) * w.float() + + +def _rope(x, positions, theta=10000.0): + # NeoX half-split convention, matching RotaryEmbedding(is_neox=True). + # x: [T, heads, hd]; positions: [T] + hd = x.shape[-1] + inv = 1.0 / theta ** (torch.arange(0, hd, 2, dtype=torch.float64) / hd) + ang = positions.double().unsqueeze(-1) * inv # [T, hd/2] + cos = ang.cos().float().unsqueeze(1) + sin = ang.sin().float().unsqueeze(1) + x1, x2 = x[..., : hd // 2], x[..., hd // 2 :] + return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1) + + +def _oracle_block_decode(weights, captured, noise_embed, use_swa): + """fp32 eager port of the DeepSpec dspark block decode + (Qwen3DSparkDecoderLayer stack over [context ; draft block]).""" + w = {k: v.float() for k, v in weights.items()} + nh, nkv, hd = (TINY["num_attention_heads"], TINY["num_key_value_heads"], TINY["head_dim"]) + ctx = captured.shape[0] + blk = noise_embed.shape[0] + ctx_pos = torch.arange(ctx, dtype=torch.long) + q_pos = torch.arange(ctx, ctx + blk, dtype=torch.long) + all_pos = torch.cat([ctx_pos, q_pos]) + + # Target feature projection: hidden_norm(fc(captured)); constant across + # layers, no input_layernorm on the context path (generic DFlash). + ctx_feat = _rms(captured.float() @ w["fc.weight"].T, w["hidden_norm.weight"]) + + hs = noise_embed.float() + for i in range(TINY["num_hidden_layers"]): + p = f"layers.{i}." + h = _rms(hs, w[p + "input_layernorm.weight"]) + q = (h @ w[p + "self_attn.q_proj.weight"].T).view(blk, nh, hd) + k_ctx = (ctx_feat @ w[p + "self_attn.k_proj.weight"].T).view(ctx, nkv, hd) + k_noise = (h @ w[p + "self_attn.k_proj.weight"].T).view(blk, nkv, hd) + v_ctx = (ctx_feat @ w[p + "self_attn.v_proj.weight"].T).view(ctx, nkv, hd) + v_noise = (h @ w[p + "self_attn.v_proj.weight"].T).view(blk, nkv, hd) + k = torch.cat([k_ctx, k_noise], dim=0) + v = torch.cat([v_ctx, v_noise], dim=0) + q = _rms(q, w[p + "self_attn.q_norm.weight"]) + k = _rms(k, w[p + "self_attn.k_norm.weight"]) + q = _rope(q, q_pos) + k = _rope(k, all_pos) + # GQA expand + rep = nh // nkv + k = k.repeat_interleave(rep, dim=1) + v = v.repeat_interleave(rep, dim=1) + scores = torch.einsum("qhd,khd->hqk", q, k) / hd**0.5 + if use_swa: + dist = (q_pos.unsqueeze(1) - all_pos.unsqueeze(0)).abs() + scores = scores.masked_fill(dist.unsqueeze(0) > SWA_WINDOW - 1, float("-inf")) + attn = torch.softmax(scores, dim=-1) + o = torch.einsum("hqk,khd->qhd", attn, v).reshape(blk, nh * hd) + hs = hs + o @ w[p + "self_attn.o_proj.weight"].T + h2 = _rms(hs, w[p + "post_attention_layernorm.weight"]) + gate = h2 @ w[p + "mlp.gate_proj.weight"].T + up = h2 @ w[p + "mlp.up_proj.weight"].T + hs = hs + (F.silu(gate) * up) @ w[p + "mlp.down_proj.weight"].T + return _rms(hs, w["norm.weight"]) + + +def _has_flash_attn(): + try: + import flash_attn # noqa: F401 + + return True + except ImportError: + return False + + +needs_gpu = pytest.mark.skipif( + not torch.cuda.is_available() or not _has_flash_attn(), + reason="tiny block-decode parity needs CUDA + flash_attn", +) + + +@needs_gpu +def test_dspark_drafter_loads_head_weights_and_parses_config(): + weights = _tiny_weights() + drafter = _build_drafter(True, weights) + assert drafter._dspark_shift_label and drafter._dspark_use_swa + assert drafter._dspark_swa_window == SWA_WINDOW + assert drafter._dspark_layer_windows == [(SWA_WINDOW - 1, SWA_WINDOW - 1)] * 2 + assert drafter.has_markov_head + torch.testing.assert_close(drafter.markov_w1.cpu(), weights["markov_w1.weight"]) + torch.testing.assert_close(drafter.markov_w2.cpu(), weights["markov_w2.weight"]) + # Confidence weights loaded for MR B, but never consumed here. + torch.testing.assert_close( + drafter.confidence_proj_weight.cpu(), weights["confidence_proj.weight"] + ) + assert drafter.confidence_proj_bias is not None + + +@needs_gpu +def test_plain_dflash_drafter_keeps_old_gates(): + """No-regression: a config WITHOUT dspark fields resolves to the exact + old code path (no window, no markov, mask slots 1..K).""" + drafter = _build_drafter(False, _tiny_weights()) + assert not drafter._dspark_shift_label + assert not drafter._dspark_use_swa + assert drafter._dspark_layer_windows == [(-1, -1)] * 2 + assert not drafter.has_markov_head + assert drafter.markov_w1 is None and drafter.confidence_proj_weight is None + x = torch.randn(2, 3, VOCAB) + t = torch.zeros(2, dtype=torch.long) + assert drafter.apply_markov_chain_logits(x, t) is x + + +@needs_gpu +def test_legacy_causal_dflash_config_constructs(): + """No-regression: legacy DFlash drafter configs (e.g. Laguna) declare + causal=true without any dspark fields; their causality is handled by + the legacy decode path, so construction must not raise.""" + from tensorrt_llm._torch.model_config import ModelConfig + + cfg = _tiny_config(False) + cfg.dflash_config = dict(cfg.dflash_config, causal=True) + drafter = DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + assert drafter._dspark_layer_windows == [(-1, -1)] * 2 + assert not drafter.has_markov_head + + +@needs_gpu +def test_dspark_causal_config_rejected(): + """The dspark block decode only supports the non-causal convention.""" + from tensorrt_llm._torch.model_config import ModelConfig + + cfg = _tiny_config(True) + cfg.dflash_config = dict(cfg.dflash_config, causal=True) + with pytest.raises(ValueError, match="non-causal dspark convention"): + DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + + +@needs_gpu +def test_dspark_projector_type_alone_rejects_causal(): + """projector_type='dspark' marks the dspark convention even when no + dspark feature flag is enabled; causal=true must still be rejected.""" + from tensorrt_llm._torch.model_config import ModelConfig + + cfg = _tiny_config(False) + cfg.dflash_config = dict(cfg.dflash_config, projector_type="dspark", causal=True) + with pytest.raises(ValueError, match="non-causal dspark convention"): + DFlashForCausalLM(ModelConfig(pretrained_config=cfg, attn_backend="TRTLLM")) + + +def _run_block_decode(drafter, weights, captured, noise_embed): + dev = "cuda" + blk = TINY["block_size"] + proj = drafter.project_target_hidden(captured.to(dev, torch.bfloat16)) + ctx_pos = torch.arange(CTX_LEN, device=dev) + k, v = drafter.precompute_context_kv(proj, ctx_pos) + L = drafter._num_attn_layers + nkv, hd = drafter._num_kv_heads, drafter._head_dim + pool_k = torch.zeros(1, L, CTX_LEN + blk, nkv, hd, dtype=torch.bfloat16, device=dev) + pool_v = torch.zeros_like(pool_k) + pool_k[0, :, :CTX_LEN] = k.permute(1, 0, 2, 3) + pool_v[0, :, :CTX_LEN] = v.permute(1, 0, 2, 3) + q_pos = torch.arange(CTX_LEN, CTX_LEN + blk, device=dev).unsqueeze(0) + out = drafter.dflash_forward( + noise_embedding=noise_embed.to(dev, torch.bfloat16).unsqueeze(0), + query_positions=q_pos, + num_ctx_per_req=torch.tensor([CTX_LEN], device=dev), + ctx_k_cache=pool_k, + ctx_v_cache=pool_v, + ctx_cache_batch_idx=torch.tensor([0], device=dev), + ) + return out.float().cpu() + + +@needs_gpu +def test_dspark_block_decode_matches_reference_oracle(): + """The full drafter block decode (fc/hidden_norm projection, per-layer + QKV + q/k-norm + RoPE, non-causal SWA flash attention over + [context ; block], MLP, final norm) matches the fp32 eager oracle; the + no-window oracle does NOT match (the window demonstrably binds).""" + torch.manual_seed(0) + weights = _tiny_weights() + drafter = _build_drafter(True, weights) + + g = torch.Generator().manual_seed(42) + captured = torch.randn(CTX_LEN, TINY["hidden_size"] * NUM_CAPTURE, generator=g) * 0.5 + noise_embed = torch.randn(TINY["block_size"], TINY["hidden_size"], generator=g) * 0.5 + + out = _run_block_decode(drafter, weights, captured, noise_embed) + + # Oracle consumes the same bf16-quantized inputs the drafter sees. + captured_q = captured.to(torch.bfloat16) + noise_q = noise_embed.to(torch.bfloat16) + oracle_swa = _oracle_block_decode(weights, captured_q, noise_q, True) + oracle_full = _oracle_block_decode(weights, captured_q, noise_q, False) + + diff_swa = (out - oracle_swa).abs().max().item() + diff_full = (out - oracle_full).abs().max().item() + # bf16 forward vs fp32 oracle: tolerance well below the SWA-vs-full gap. + assert diff_swa < 0.02, f"SWA parity failed: max abs diff {diff_swa}" + assert diff_full > 4 * max(diff_swa, 1e-4), ( + f"negative control failed: no-window oracle too close " + f"({diff_full} vs swa {diff_swa}) — window may not be applied" + ) + + +@needs_gpu +def test_plain_dflash_block_decode_matches_full_attention_oracle(): + """No-regression numeric check: the plain-DFlash drafter (no dspark + fields) still runs full non-causal attention over the whole context.""" + weights = _tiny_weights() + drafter = _build_drafter(False, weights) + g = torch.Generator().manual_seed(43) + captured = torch.randn(CTX_LEN, TINY["hidden_size"] * NUM_CAPTURE, generator=g) * 0.5 + noise_embed = torch.randn(TINY["block_size"], TINY["hidden_size"], generator=g) * 0.5 + out = _run_block_decode(drafter, weights, captured, noise_embed) + oracle = _oracle_block_decode( + weights, captured.to(torch.bfloat16), noise_embed.to(torch.bfloat16), False + ) + diff = (out - oracle).abs().max().item() + assert diff < 0.02, f"plain DFlash parity failed: max abs diff {diff}" diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py index bd3f54683d5c..99c233a58ce0 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_sa.py @@ -6,6 +6,8 @@ import torch from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests +from tensorrt_llm._torch.speculative.suffix_automaton import SAConfig, SuffixAutomatonManager from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig, SADecodingConfig sys.path.append(os.path.join(os.path.dirname(__file__), "..")) @@ -236,5 +238,280 @@ def test_sa_config_invalid_zero(): ) +class _FakeSARequest: + """Minimal request stand-in for SuffixAutomatonManager.prepare_resources. + + Provides only the attributes/methods that prepare_resources and + free_resources touch. Mirrors an LlmRequest on a disaggregated + generation server, where the request skips the context phase + (DISAGG_GENERATION_INIT) and is scheduled directly as a generation + request with LLMREQUEST_TYPE_GENERATION_ONLY type. + """ + + def __init__( + self, + request_id: int, + tokens: list, + *, + generation_only: bool = False, + is_dummy: bool = False, + is_first_context_chunk: bool = True, + ): + self.request_id = request_id + self._tokens = list(tokens) + self._generation_only = generation_only + self.is_dummy = is_dummy + self.is_first_context_chunk = is_first_context_chunk + + def get_tokens(self, beam: int) -> list: + assert beam == 0 + return list(self._tokens) + + def is_generation_only_request(self) -> bool: + return self._generation_only + + +class TestSADisaggGenInit: + """CPU-only tests for the disagg-generation SA init path. + + On a disagg generation server, requests never appear as context + requests, so prepare_resources must build the automaton from the + generation request's token history (prompt + first generated token). + These tests exercise only host-side automaton construction and slot + bookkeeping; no GPU workspace is allocated. + """ + + @staticmethod + def _make_manager(max_slots: int = 8) -> SuffixAutomatonManager: + config = SAConfig(max_seq_len=1024, max_slots=max_slots) + return SuffixAutomatonManager(config, max_num_requests=max_slots) + + def test_disagg_gen_request_initializes_automaton(self): + """A generation-only request must be initialized from its tokens.""" + manager = self._make_manager() + try: + # Prompt tokens + first generated token (appended by + # _prepare_disagg_gen_transmission_complete on the gen server). + tokens = [1, 2, 3, 4, 5, 1, 2, 3] + [6] + req = _FakeSARequest(7, tokens, generation_only=True) + + batch = ScheduledRequests() + batch.generation_requests = [req] + + free_slots_before = len(manager._free_slots) + manager.prepare_resources(batch) + + assert req.request_id in manager._initialized_requests + assert req.request_id in manager._request_to_slot + assert req.request_id in manager._host_states_native + assert req.request_id in manager._pending_copies + assert len(manager._free_slots) == free_slots_before - 1 + + # A second prepare_resources with the same request (every + # subsequent decode iteration) must be a no-op: same slot, no + # extra slot consumed, host state not rebuilt. + slot = manager._request_to_slot[req.request_id] + state = manager._host_states_native[req.request_id] + manager.prepare_resources(batch) + assert manager._request_to_slot[req.request_id] == slot + assert manager._host_states_native[req.request_id] is state + assert len(manager._free_slots) == free_slots_before - 1 + finally: + manager.shutdown() + + def test_regular_and_dummy_generation_requests_are_skipped(self): + """Only disagg (generation-only) non-dummy requests are initialized.""" + manager = self._make_manager() + try: + regular_gen = _FakeSARequest(1, [1, 2, 3], generation_only=False) + dummy_gen = _FakeSARequest(2, [1, 2, 3], generation_only=True, is_dummy=True) + + batch = ScheduledRequests() + batch.generation_requests = [regular_gen, dummy_gen] + + free_slots_before = len(manager._free_slots) + manager.prepare_resources(batch) + + assert regular_gen.request_id not in manager._initialized_requests + assert dummy_gen.request_id not in manager._initialized_requests + assert len(manager._free_slots) == free_slots_before + finally: + manager.shutdown() + + def test_context_initialized_request_not_reinitialized_in_generation(self): + """Aggregated flow: context-phase init must not be redone in generation.""" + manager = self._make_manager() + try: + req = _FakeSARequest(3, [10, 20, 30], generation_only=False) + + ctx_batch = ScheduledRequests() + ctx_batch.context_requests_last_chunk = [req] + manager.prepare_resources(ctx_batch) + slot = manager._request_to_slot[req.request_id] + state = manager._host_states_native[req.request_id] + + # The same request later shows up as a generation request. + gen_batch = ScheduledRequests() + gen_batch.generation_requests = [req] + manager.prepare_resources(gen_batch) + + assert manager._request_to_slot[req.request_id] == slot + assert manager._host_states_native[req.request_id] is state + finally: + manager.shutdown() + + def test_free_resources_after_disagg_gen_init(self): + """free_resources must fully release slots/bookkeeping for the gen path.""" + manager = self._make_manager() + try: + req = _FakeSARequest(4, [1, 2, 3, 4], generation_only=True) + + batch = ScheduledRequests() + batch.generation_requests = [req] + + free_slots_before = len(manager._free_slots) + manager.prepare_resources(batch) + manager.free_resources(req) + + assert req.request_id not in manager._initialized_requests + assert req.request_id not in manager._request_to_slot + assert req.request_id not in manager._host_states_native + assert req.request_id not in manager._pending_copies + assert len(manager._free_slots) == free_slots_before + + # The request can be initialized again from scratch. + manager.prepare_resources(batch) + assert req.request_id in manager._initialized_requests + finally: + manager.shutdown() + + def test_disagg_gen_init_defers_to_generation_schedule(self): + """Ctx/gen spec split: init must NOT happen at _prepare_disagg_gen_init. + + The executor's _prepare_disagg_gen_init routes DISAGG_GENERATION_INIT + requests through prepare_resources as context_requests_last_chunk + BEFORE the ctx server's first generated token has been appended + (that happens later, in _prepare_disagg_gen_transmission_complete). + Initializing there would freeze the automaton at prompt-only — + permanently missing the ctx first token relative to an aggregated + run — and pin an SA slot for the whole KV-transfer duration. The + context loop must skip generation-only requests and let the + generation loop initialize them with the full token history. + """ + manager = self._make_manager() + try: + prompt = [1, 2, 3, 4, 5] + req = _FakeSARequest(11, prompt, generation_only=True) + + # Phase 1: _prepare_disagg_gen_init — request arrives as a + # context_requests_last_chunk entry with prompt-only tokens. + init_batch = ScheduledRequests() + init_batch.context_requests_last_chunk = [req] + + free_slots_before = len(manager._free_slots) + manager.prepare_resources(init_batch) + + assert req.request_id not in manager._initialized_requests + assert req.request_id not in manager._request_to_slot + # No SA slot pinned during the KV transfer. + assert len(manager._free_slots) == free_slots_before + + # Phase 2: transmission complete — the executor appends the ctx + # first token and the request is scheduled as a generation + # request. Init must now use the FULL history. + req._tokens.append(6) + + seen_tokens = {} + orig_add_request = manager.add_request + + def spy_add_request(request_id, context_tokens): + seen_tokens[request_id] = list(context_tokens) + return orig_add_request(request_id, context_tokens) + + manager.add_request = spy_add_request + gen_batch = ScheduledRequests() + gen_batch.generation_requests = [req] + manager.prepare_resources(gen_batch) + manager.add_request = orig_add_request + + assert req.request_id in manager._initialized_requests + assert seen_tokens[req.request_id] == prompt + [6] + finally: + manager.shutdown() + + +class TestKdaReplaySeedOnDisaggTransfer(unittest.TestCase): + """seed_kda_replay_caches_for_disagg_gen must mirror + _sync_kda_replay_conv_window for transferred requests: committed conv + window seeded from the (transferred) conv pool, draft tail columns and + pending-draft scratch cleared, other slots untouched.""" + + L, SLOTS, D, W, M, NHEADS = 2, 4, 6, 4, 2, 3 # committed = W - 1 = 3 + + def _make_manager(self, use_kda_replay=True): + from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import PythonMambaCacheManager + + L, SLOTS, D, W, M, NH = (self.L, self.SLOTS, self.D, self.W, self.M, self.NHEADS) + committed = W - 1 + + class _FakeSpecState: + pass + + cache = _FakeSpecState() + torch.manual_seed(0) + cache.conv = torch.randn(L, SLOTS, 3 * D, W) + cache.kda_conv_q = torch.full((L, SLOTS, D, committed + M), 7.0) + cache.kda_conv_k = torch.full((L, SLOTS, D, committed + M), 7.0) + cache.kda_conv_v = torch.full((L, SLOTS, D, committed + M), 7.0) + cache.kda_qkg_cache = torch.full((L, SLOTS, M, 3, D), 7.0) + cache.kda_v_cache = torch.full((L, SLOTS, M, D), 7.0) + cache.kda_beta_cache = torch.full((L, SLOTS, M, NH), 7.0) + cache.prev_num_accepted_tokens = torch.full((SLOTS,), 5, dtype=torch.int32) + + mgr = PythonMambaCacheManager.__new__(PythonMambaCacheManager) + mgr._use_kda_replay_update = use_kda_replay + mgr.SpeculativeState = _FakeSpecState + mgr.mamba_cache = cache + mgr.mamba_cache_index = {101: 1, 202: 3} + return mgr, cache + + def test_seeds_committed_window_and_clears_scratch(self): + mgr, cache = self._make_manager() + committed = self.W - 1 + conv_before = cache.conv.clone() + mgr.seed_kda_replay_caches_for_disagg_gen([101, 202]) + d = self.D + for slot in (1, 3): + for kda, lo, hi in ( + (cache.kda_conv_q, 0, d), + (cache.kda_conv_k, d, 2 * d), + (cache.kda_conv_v, 2 * d, 3 * d), + ): + torch.testing.assert_close( + kda[:, slot, :, :committed], conv_before[:, slot, lo:hi, 1:].to(kda.dtype) + ) + assert (kda[:, slot, :, committed:] == 0).all() + assert (cache.kda_qkg_cache[:, slot] == 0).all() + assert (cache.kda_v_cache[:, slot] == 0).all() + assert (cache.kda_beta_cache[:, slot] == 0).all() + assert cache.prev_num_accepted_tokens[slot] == 0 + # Conv pool itself must not be modified. + torch.testing.assert_close(cache.conv, conv_before) + # Untouched slots keep their contents. + for slot in (0, 2): + assert (cache.kda_conv_q[:, slot] == 7.0).all() + assert (cache.kda_qkg_cache[:, slot] == 7.0).all() + assert cache.prev_num_accepted_tokens[slot] == 5 + + def test_noop_without_kda_replay_or_unknown_ids(self): + mgr, cache = self._make_manager(use_kda_replay=False) + mgr.seed_kda_replay_caches_for_disagg_gen([101]) + assert (cache.kda_conv_q == 7.0).all() + + mgr2, cache2 = self._make_manager() + mgr2.seed_kda_replay_caches_for_disagg_gen([999]) # not in index + assert (cache2.kda_conv_q == 7.0).all() + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 8b2a8774c292..ee3a2dd7cb40 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -377,7 +377,7 @@ commands: flags: - "--post_processor_hook" reasoning_parser: - type: Choice(['auto', 'deepseek-r1', 'deepseek_v4', 'gemma4', 'kimi_k2', 'kimi_k25', 'laguna', 'minimax_m2', 'minimax_m2_append_think', 'minimax_m3', 'nano-v3', 'nemotron-v3', 'qwen3', 'qwen3_5']) + type: Choice(['auto', 'deepseek-r1', 'deepseek_v4', 'gemma4', 'kimi_k2', 'kimi_k25', 'kimi_k3', 'laguna', 'minimax_m2', 'minimax_m2_append_think', 'minimax_m3', 'nano-v3', 'nemotron-v3', 'qwen3', 'qwen3_5']) default: null status: prototype required: false @@ -443,7 +443,7 @@ commands: flags: - "--tokenizer" tool_parser: - type: Choice(['auto', 'deepseek_v3', 'deepseek_v31', 'deepseek_v32', 'deepseek_v4', 'gemma4', 'glm4', 'glm47', 'kimi_k2', 'minimax_m2', 'minimax_m3', 'poolside_v1', 'qwen3', 'qwen3_coder']) + type: Choice(['auto', 'deepseek_v3', 'deepseek_v31', 'deepseek_v32', 'deepseek_v4', 'gemma4', 'glm4', 'glm47', 'kimi_k2', 'kimi_k3', 'minimax_m2', 'minimax_m3', 'poolside_v1', 'qwen3', 'qwen3_coder']) default: null status: prototype required: false diff --git a/tests/unittest/inputs/test_chat_template_dispatch.py b/tests/unittest/inputs/test_chat_template_dispatch.py index 9bb730b45655..8f7a823a14fd 100644 --- a/tests/unittest/inputs/test_chat_template_dispatch.py +++ b/tests/unittest/inputs/test_chat_template_dispatch.py @@ -5,6 +5,7 @@ import threading import pytest +from transformers import PreTrainedTokenizerBase from tensorrt_llm.inputs.content_format import ContentFormat from tensorrt_llm.inputs.registry import ( @@ -17,9 +18,11 @@ _build_openai_content, _resolve_content_format, add_multimodal_placeholders, + apply_chat_template, async_apply_chat_template, interleave_mm_placeholders, ) +from tensorrt_llm.tokenizer import TransformersTokenizer pytestmark = pytest.mark.cpu_only @@ -361,6 +364,102 @@ def apply_chat_template(self, **_): assert tokenizer.worker_thread_id != event_loop_thread_id +class TestPythonChatTemplate: + @pytest.mark.parametrize( + ("enable_tokenize", "expected"), + [(False, "native-rendered"), (True, [1, 2, 3])], + ) + def test_uses_native_renderer_without_jinja_template(self, enable_tokenize, expected): + class NativeTokenizer: + chat_template = None + + def __init__(self): + self.captured = None + + def get_chat_template(self, *_args, **_kwargs): + raise AssertionError("Jinja resolution should be bypassed") + + def apply_chat_template(self, conversation, **kwargs): + self.captured = (conversation, kwargs) + return [1, 2, 3] if kwargs["tokenize"] else "native-rendered" + + class Processor: + chat_template = None + + conversation = [ConversationMessage(role="user", content="hello", media=[])] + tools = [{"type": "function", "function": {"name": "echo"}}] + inner_tokenizer = NativeTokenizer() + + result = apply_chat_template( + model_type="kimi_linear", + tokenizer=TransformersTokenizer(inner_tokenizer), + processor=Processor(), + conversation=conversation, + add_generation_prompt=True, + mm_placeholder_counts=[{}], + tools=tools, + chat_template_kwargs={"thinking_effort": "high"}, + enable_tokenize=enable_tokenize, + ) + + assert result == expected + captured_conversation, captured_kwargs = inner_tokenizer.captured + assert captured_conversation == conversation + assert captured_kwargs == { + "tools": tools, + "tokenize": enable_tokenize, + "add_generation_prompt": True, + "thinking_effort": "high", + } + + def test_explicit_jinja_template_takes_precedence(self): + class NativeTokenizer: + chat_template = None + + def __init__(self): + self.captured = None + + def apply_chat_template(self, conversation, **kwargs): + self.captured = (conversation, kwargs) + return "jinja-rendered" + + conversation = [ConversationMessage(role="user", content="hello", media=[])] + inner_tokenizer = NativeTokenizer() + template = "{{ messages }}" + + result = apply_chat_template( + model_type="kimi_linear", + tokenizer=TransformersTokenizer(inner_tokenizer), + processor=None, + conversation=conversation, + add_generation_prompt=True, + mm_placeholder_counts=[{}], + chat_template=template, + ) + + assert result == "jinja-rendered" + _, captured_kwargs = inner_tokenizer.captured + assert captured_kwargs["chat_template"] == template + + def test_missing_template_still_errors_without_native_override(self): + class DefaultTokenizer: + chat_template = None + apply_chat_template = PreTrainedTokenizerBase.apply_chat_template + + def get_chat_template(self, *_args, **_kwargs): + return None + + with pytest.raises(ValueError, match="No chat template found"): + apply_chat_template( + model_type="test_string_model", + tokenizer=DefaultTokenizer(), + processor=None, + conversation=[ConversationMessage(role="user", content="hello", media=[])], + add_generation_prompt=True, + mm_placeholder_counts=[{}], + ) + + class TestServingChatTemplateGather: """Cover the asyncio.gather integration in the serving chat-template paths.""" diff --git a/tests/unittest/llmapi/apps/test_tool_parsers.py b/tests/unittest/llmapi/apps/test_tool_parsers.py index f551ed534d93..7b24d4554337 100644 --- a/tests/unittest/llmapi/apps/test_tool_parsers.py +++ b/tests/unittest/llmapi/apps/test_tool_parsers.py @@ -19,6 +19,7 @@ import pytest +from tensorrt_llm.sampling_params import SamplingParams from tensorrt_llm.serve.openai_protocol import (ChatCompletionToolsParam, FunctionDefinition) from tensorrt_llm.serve.tool_parser.base_tool_parser import BaseToolParser @@ -32,6 +33,7 @@ from tensorrt_llm.serve.tool_parser.glm4_parser import Glm4ToolParser from tensorrt_llm.serve.tool_parser.glm47_parser import Glm47ToolParser from tensorrt_llm.serve.tool_parser.kimi_k2_tool_parser import KimiK2ToolParser +from tensorrt_llm.serve.tool_parser.kimi_k3_tool_parser import KimiK3ToolParser from tensorrt_llm.serve.tool_parser.minimax_m2_parser import MiniMaxM2ToolParser from tensorrt_llm.serve.tool_parser.poolside_v1_parser import \ PoolsideV1ToolParser @@ -4068,3 +4070,288 @@ def test_strict_tool_with_deepseek_parser(self): if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +class TestKimiK3ToolParser(BaseToolParserTestClass): + """Test suite for KimiK3ToolParser (XTML tool-call format). + + Fixture strings follow the checkpoint's `encoding_k3.py` rendering: + `<|open|>tag key="value"<|sep|>` / `<|close|>tag<|sep|>`, attributes + space-prefixed and `&`/`"`-escaped, call indices 1-based, string + argument bodies raw and non-string bodies JSON. + """ + + BOT = "<|open|>tools<|sep|>" + EOT = "<|close|>tools<|sep|>" + + @staticmethod + def _call(name: str, index: int, body: str) -> str: + return (f'<|open|>call tool="{name}" index="{index}"<|sep|>' + f'{body}<|close|>call<|sep|>') + + @staticmethod + def _argument(key: str, type_: str, value: str) -> str: + return (f'<|open|>argument key="{key}" type="{type_}"<|sep|>' + f'{value}<|close|>argument<|sep|>') + + def _section(self, *calls: str) -> str: + return self.BOT + "".join(calls) + self.EOT + + def make_parser(self): + return KimiK3ToolParser() + + def make_tool_parser_test_cases(self): + single_call = self._section( + self._call("get_weather", 1, + self._argument("location", "string", "NYC"))) + return ToolParserTestCases( + has_tool_call_true="Some text " + single_call, + detect_and_parse_single_tool=( + "Normal text" + single_call, + "Normal text", + "get_weather", + { + "location": "NYC" + }, + ), + detect_and_parse_multiple_tools=( + self._section( + self._call("get_weather", 1, + self._argument("location", "string", "LA")), + self._call("search_web", 2, + self._argument("query", "string", "AI")), + ), + ("get_weather", "search_web"), + ), + # A call without the mandatory tool="..." attribute is skipped. + detect_and_parse_malformed_tool=self._section( + '<|open|>call index="1"<|sep|>' + '<|open|>argument key="location" type="string"<|sep|>NYC' + '<|close|>argument<|sep|><|close|>call<|sep|>'), + # K3 has no JSON "parameters" key wrapper; the closest analogue + # is the raw-JSON call body variant. + detect_and_parse_with_parameters_key=( + self._section( + self._call( + "search_web", 1, + '<|open|>json type="object"<|sep|>{"query": "test"}' + '<|close|>json<|sep|>')), + "search_web", + { + "query": "test" + }, + ), + parse_streaming_increment_partial_bot_token="<|open|>too", + undefined_tool=self._section( + self._call("undefined_func", 1, + self._argument("arg", "string", "any value"))), + ) + + def test_initialization(self, parser): + assert parser.bot_token == self.BOT + assert parser.eot_token == self.EOT + + def test_undefined_tool(self, sample_tools, parser, tool_parser_test_cases): + """Keep undefined-tool calls at their positional index. + + K3 warns about undefined tools rather than remapping ``tool_index`` to + ``-1``. + """ + text = tool_parser_test_cases.undefined_tool + + result = parser.detect_and_parse(text, sample_tools) + + assert len(result.calls) == 1 + assert result.calls[0].name == "undefined_func" + assert result.calls[0].tool_index == 0 + + def test_supports_structural_tag(self): + """Reject JSON-schema structural tagging for XTML bodies. + + XTML bodies already use tag-structured text, so JSON-schema + structural-tag constrained decoding does not apply. + """ + parser = KimiK3ToolParser() + assert parser.supports_structural_tag() is False + with pytest.raises(NotImplementedError): + parser.structure_info() + + def test_argument_type_coercion(self, sample_tools, parser): + """Non-string argument bodies are JSON; string bodies stay raw.""" + text = self._section( + self._call( + "get_weather", 1, + self._argument("location", "string", '"quoted" & raw') + + self._argument("count", "number", "3") + + self._argument("celsius", "boolean", "true") + + self._argument("extra", "null", "null") + + self._argument("nested", "object", '{"a": [1, 2]}') + + self._argument("tags", "array", '["x", "y"]'))) + + result = parser.detect_and_parse(text, sample_tools) + + assert len(result.calls) == 1 + assert json.loads(result.calls[0].parameters) == { + "location": '"quoted" & raw', + "count": 3, + "celsius": True, + "extra": None, + "nested": { + "a": [1, 2] + }, + "tags": ["x", "y"], + } + + def test_argument_invalid_json_falls_back_to_raw(self, sample_tools, + parser): + """Keep a non-string argument's invalid JSON body as raw text. + + Invalid JSON should fall back to its original text instead of raising. + """ + text = self._section( + self._call("get_weather", 1, + self._argument("count", "number", "not-a-number"))) + + result = parser.detect_and_parse(text, sample_tools) + + assert json.loads(result.calls[0].parameters) == { + "count": "not-a-number" + } + + def test_attribute_unescaping(self, sample_tools, parser): + """Unescape encoded XTML attribute values. + + K3 attributes arrive escaped by ``encoding_k3._escape_attr_value``. + """ + text = self._section( + self._call( + "get_weather", 1, + self._argument("say "hi" & bye", "string", "v"))) + + result = parser.detect_and_parse(text, sample_tools) + + assert json.loads(result.calls[0].parameters) == {'say "hi" & bye': "v"} + + def test_empty_arguments(self, sample_tools, parser): + """A call with no argument tags yields an empty JSON object.""" + text = self._section(self._call("search_web", 1, "")) + + result = parser.detect_and_parse(text, sample_tools) + + assert len(result.calls) == 1 + assert result.calls[0].parameters == "{}" + + def test_trailing_structural_markup_stripped(self, sample_tools, parser): + """Strip trailing XTML terminators from standalone normal text. + + This covers parsing without a tools section or reasoning parser. + """ + text = "The answer is 4.<|close|>message<|sep|><|end_of_msg|>" + + result = parser.detect_and_parse(text, sample_tools) + + assert result.normal_text == "The answer is 4." + assert len(result.calls) == 0 + + def test_streaming_buffers_section_until_complete(self, sample_tools, + parser): + """Buffer an incomplete tools section while streaming response text. + + Response text is emitted immediately, but calls wait for the closing + ``<|close|>tools<|sep|>`` marker. + """ + result = parser.parse_streaming_increment("Checking. ", sample_tools) + assert result.normal_text == "Checking. " + assert result.calls == [] + + # Section opener + call header: everything buffered. + result = parser.parse_streaming_increment( + self.BOT + '<|open|>call tool="get_weather" index="1"<|sep|>', + sample_tools) + assert result.normal_text == "" + assert result.calls == [] + + # Arguments still buffered. + result = parser.parse_streaming_increment( + self._argument("location", "string", "NYC"), sample_tools) + assert result.normal_text == "" + assert result.calls == [] + + # Section close: the complete call is emitted. + result = parser.parse_streaming_increment( + "<|close|>call<|sep|>" + self.EOT, sample_tools) + assert result.normal_text == "" + assert len(result.calls) == 1 + assert result.calls[0].name == "get_weather" + assert json.loads(result.calls[0].parameters) == {"location": "NYC"} + + def test_composes_with_kimi_k3_reasoning_parser(self, sample_tools, parser): + """Parse tools passed through the Kimi-K3 reasoning parser. + + The reasoning parser preserves the tools section verbatim for this + parser. + """ + from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory + + completion = ( + "Need the weather.<|close|>think<|sep|>" + "<|open|>response<|sep|>Checking." + "<|close|>response<|sep|>" + self._section( + self._call("get_weather", 1, + self._argument("location", "string", "NYC"))) + + "<|close|>message<|sep|><|end_of_msg|>") + + reasoning = ReasoningParserFactory.create_reasoning_parser("kimi_k3") + stage1 = reasoning.parse(completion) + assert stage1.reasoning_content == "Need the weather." + + stage2 = parser.detect_and_parse(stage1.content, sample_tools) + assert stage2.normal_text == "Checking." + assert len(stage2.calls) == 1 + assert stage2.calls[0].name == "get_weather" + assert json.loads(stage2.calls[0].parameters) == {"location": "NYC"} + + +class TestConfigureParserSpecialTokenDecoding: + """Test parser-specific detokenization settings in the OpenAI server.""" + + @staticmethod + def _configure(reasoning_parser_name: str | None = None, + tool_parser_name: str | None = None, + has_tools: bool = False) -> SamplingParams: + from tensorrt_llm.serve.openai_server import \ + _configure_parser_special_token_decoding + + sampling_params = SamplingParams() + _configure_parser_special_token_decoding( + sampling_params, + reasoning_parser_name=reasoning_parser_name, + tool_parser_name=tool_parser_name, + has_tools=has_tools) + return sampling_params + + def test_kimi_k3_reasoning_parser_preserves_compact_xtml(self) -> None: + sampling_params = self._configure(reasoning_parser_name="kimi_k3") + + assert sampling_params.skip_special_tokens is False + assert sampling_params.spaces_between_special_tokens is False + + def test_kimi_k3_tool_parser_preserves_compact_xtml(self) -> None: + sampling_params = self._configure(tool_parser_name="KIMI_K3", + has_tools=True) + + assert sampling_params.skip_special_tokens is False + assert sampling_params.spaces_between_special_tokens is False + + def test_tool_parser_does_not_apply_without_tools(self) -> None: + sampling_params = self._configure(tool_parser_name="kimi_k3") + + assert sampling_params.skip_special_tokens is True + assert sampling_params.spaces_between_special_tokens is True + + def test_other_raw_token_parser_keeps_spacing_contract(self) -> None: + sampling_params = self._configure(tool_parser_name="deepseek_v32", + has_tools=True) + + assert sampling_params.skip_special_tokens is False + assert sampling_params.spaces_between_special_tokens is True diff --git a/tests/unittest/llmapi/test_reasoning_parser.py b/tests/unittest/llmapi/test_reasoning_parser.py index b2a8f78457a5..cf9da1d1eb58 100644 --- a/tests/unittest/llmapi/test_reasoning_parser.py +++ b/tests/unittest/llmapi/test_reasoning_parser.py @@ -873,3 +873,192 @@ def test_gemma4_reasoning_parser_finish_unterminated_reasoning(): result = parser.finish() assert result.content == "" assert result.reasoning_content == "tag key="value"<|sep|>` / `<|close|>tag<|sep|>` with no +# whitespace between segments, and the generation prompt ends inside +# `<|open|>think<|sep|>` (or `<|open|>response<|sep|>` when thinking=False), +# so completions start mid-channel and end with +# `<|close|>message<|sep|><|end_of_msg|>`. +# --------------------------------------------------------------------------- + +K3_OPEN, K3_CLOSE, K3_SEP, K3_EOM = ("<|open|>", "<|close|>", "<|sep|>", + "<|end_of_msg|>") + +# One get_weather call, exactly as encoding_k3._render_assistant_segments +# renders it (attributes space-prefixed, index 1-based, string args raw). +K3_TOOLS_SECTION = (f'{K3_OPEN}tools{K3_SEP}' + f'{K3_OPEN}call tool="get_weather" index="1"{K3_SEP}' + f'{K3_OPEN}argument key="location" type="string"{K3_SEP}' + f'NYC' + f'{K3_CLOSE}argument{K3_SEP}' + f'{K3_CLOSE}call{K3_SEP}' + f'{K3_CLOSE}tools{K3_SEP}') + +K3_MSG_END = f"{K3_CLOSE}message{K3_SEP}{K3_EOM}" + + +def _k3_completion(reasoning: str, + content: str, + tools_section: str = "", + terminated: bool = True) -> str: + """A thinking-mode completion (prompt already opened the think channel).""" + text = (f"{reasoning}{K3_CLOSE}think{K3_SEP}" + f"{K3_OPEN}response{K3_SEP}{content}") + if terminated: + text += f"{K3_CLOSE}response{K3_SEP}{tools_section}{K3_MSG_END}" + return text + + +@pytest.mark.parametrize( + ("text", "content", "reasoning_content"), + [ + # Fully terminated think + response message. + (_k3_completion("step by step", "The answer is 4."), "The answer is 4.", + "step by step"), + # Tool-calling message: the tools section passes through into content + # verbatim so the kimi_k3 tool parser can consume it downstream. + (_k3_completion("pick a tool", "Checking.", K3_TOOLS_SECTION), + "Checking." + K3_TOOLS_SECTION, "pick a tool"), + # Length-capped mid-think: everything is reasoning. + ("unterminated reasoning", "", "unterminated reasoning"), + # Length-capped mid-response: finish() flushes the response tail. + (_k3_completion("r", "partial resp", + terminated=False), "partial resp", "r"), + # Empty think body. + (_k3_completion("", "only content"), "only content", ""), + ]) +def test_kimi_k3_reasoning_parser(text: str, content: str, + reasoning_content: str): + parser = ReasoningParserFactory.create_reasoning_parser("kimi_k3") + result = parser.parse(text) + assert result.content == content + assert result.reasoning_content == reasoning_content + + +@pytest.mark.parametrize( + ("text", "content", "reasoning_content"), + [ + # thinking=False: prompt opened the response channel directly. + (f"plain answer{K3_CLOSE}response{K3_SEP}{K3_MSG_END}", "plain answer", + ""), + (f"answer{K3_CLOSE}response{K3_SEP}{K3_TOOLS_SECTION}{K3_MSG_END}", + "answer" + K3_TOOLS_SECTION, ""), + ]) +def test_kimi_k3_reasoning_parser_non_thinking(text: str, content: str, + reasoning_content: str): + parser = ReasoningParserFactory.create_reasoning_parser( + "kimi_k3", {"thinking": False}) + result = parser.parse(text) + assert result.content == content + assert result.reasoning_content == reasoning_content + + +@pytest.mark.parametrize( + ("delta_texts", "content", "reasoning_content"), + [ + # Plain reasoning streams straight through. + (["a", "b"], ["", ""], ["a", "b"]), + # Channel switch split across deltas mid-marker. + ( + [ + "rea", + f"son{K3_CLOSE}thi", # codespell:ignore thi + f"nk{K3_SEP}{K3_OPEN}response{K3_SEP}c", + "d" + ], + ["", "", "c", "d"], + ["rea", "son", "", ""]), + # A partial marker at the end of a delta is held back, then released + # as reasoning once it turns out not to be a marker. + (["a<|clo", "x"], ["", ""], ["a", "<|clox"]), + # Structural close/open pair arriving as one delta. + ([f"r{K3_CLOSE}think{K3_SEP}{K3_OPEN}response{K3_SEP}c"], ["c"], ["r"]), + # tools_pass: `<|close|>tools<|sep|>` has an internal `<`, so the + # suffix hold must consider mid-marker splits like `...<|close|>to`. + ([ + f"r{K3_CLOSE}think{K3_SEP}{K3_OPEN}response{K3_SEP}", + f"{K3_OPEN}tools{K3_SEP}CALL{K3_CLOSE}to", + f"ols{K3_SEP}{K3_MSG_END}", + ], [ + "", + f"{K3_OPEN}tools{K3_SEP}CALL", + f"{K3_CLOSE}tools{K3_SEP}", + ], ["r", "", ""]), + # Message terminator split across deltas produces no output. + ([ + f"r{K3_CLOSE}think{K3_SEP}{K3_OPEN}response{K3_SEP}c", + "<|close|>mes", f"sage{K3_SEP}{K3_EOM}" + ], ["c", "", ""], ["r", "", ""]), + ]) +def test_kimi_k3_reasoning_parser_stream(delta_texts: list, content: list, + reasoning_content: list): + parser = ReasoningParserFactory.create_reasoning_parser("kimi_k3") + for i, delta_text in enumerate(delta_texts): + result = parser.parse_delta(delta_text) + assert result.content == content[i], \ + f"Step {i}: delta={delta_text!r}, expected content={content[i]!r}, got {result.content!r}" + assert result.reasoning_content == reasoning_content[i], \ + f"Step {i}: delta={delta_text!r}, expected reasoning={reasoning_content[i]!r}, got {result.reasoning_content!r}" + + +@pytest.mark.parametrize("chunk_size", [1, 2, 3, 7]) +@pytest.mark.parametrize("thinking", [True, False]) +def test_kimi_k3_reasoning_parser_stream_matches_parse(chunk_size: int, + thinking: bool): + """Streaming in arbitrary chunkings must reproduce the one-shot parse. + + This sweeps every marker-split position, which is the riskiest logic in + the parser (`_partial_suffix_len` suffix holds). + """ + if thinking: + text = _k3_completion("Let me think.", "Answer: 4.", K3_TOOLS_SECTION) + else: + text = (f"Answer: 4.{K3_CLOSE}response{K3_SEP}{K3_TOOLS_SECTION}" + f"{K3_MSG_END}") + kwargs = None if thinking else {"thinking": False} + + oneshot = ReasoningParserFactory.create_reasoning_parser("kimi_k3", + kwargs).parse(text) + + streamer = ReasoningParserFactory.create_reasoning_parser("kimi_k3", kwargs) + content, reasoning = [], [] + for start in range(0, len(text), chunk_size): + result = streamer.parse_delta(text[start:start + chunk_size]) + content.append(result.content) + reasoning.append(result.reasoning_content) + tail = streamer.finish() + content.append(tail.content) + reasoning.append(tail.reasoning_content) + + assert "".join(content) == oneshot.content + assert "".join(reasoning) == oneshot.reasoning_content + + +def test_kimi_k3_needs_raw_special_tokens(): + """The K3 delimiters are special tokens. + + The serving layer keys off this flag to disable skip_special_tokens for + the request. + """ + assert ReasoningParserFactory.needs_raw_special_tokens("kimi_k3") is True + assert ReasoningParserFactory.needs_raw_special_tokens("KIMI_K3") is True + assert ReasoningParserFactory.needs_raw_special_tokens( + "deepseek-r1") is False + assert ReasoningParserFactory.needs_raw_special_tokens( + "no_such_parser") is False + + +def test_auto_detect_kimi_k3(tmp_path): + """Kimi K3 model → 'kimi_k3' parser.""" + model_dir = str(tmp_path / "Kimi-K3") + os.makedirs(model_dir) + _write_config(model_dir, "kimi_k3") + + result = resolve_auto_reasoning_parser(model_dir) + assert result == "kimi_k3"