TensorRT-LLM 源码完全拆解（十）：投机解码——one-model 统一 worker、KV 回滚与验收流水线

1
def get_spec_drafter(model_engine, draft_model_engine, sampler,
2
                     spec_resource_manager, guided_decoder=None):
3
    if spec_config.spec_dec_mode.is_user_provided():
4
        return spec_config.drafter
5
    if spec_config.spec_dec_mode.is_ngram():
6
        return NGramDrafter(spec_config, spec_resource_manager)
7
    return None

1
class Drafter(ABC):
2
    @abstractmethod
3
    def prepare_draft_tokens(self, scheduled_requests,
4
                             resource_manager=None) -> None: ...
5

6
    @final
7
    def pad_draft_tokens_for_cuda_graph(self, scheduled_requests,
8
                                        resource_manager=None) -> None:
9
        pad_to = self._static_max_total_draft_tokens
10
        for req in scheduled_requests.generation_requests:
11
            num_draft_tokens = get_draft_token_length(req)
12
            req.py_draft_tokens_effective_len = num_draft_tokens
13
            req.py_draft_tokens.extend(0 for _ in range(pad_to - num_draft_tokens))
14
        ...

1
def get_spec_decoder(sampler_args: TorchSampler.Args, spec_config):
2
    spec_dec_mode = spec_config.spec_dec_mode
3
    if spec_dec_mode.use_one_engine():
4
        penalty_supported = not _is_effective_dynamic_tree(spec_config)
5
        fused_sampling = (spec_config.advanced_sampling_mode.is_fused)
6
        return SpecSampler(sampler_args,
7
                           enable_penalty=spec_config.enable_penalty,
8
                           penalty_supported=penalty_supported,
9
                           fused_sampling=fused_sampling)
10
    raise ValueError(...)

1
{
2
    'logits': raw_logits,
3
    'new_tokens': accepted_tokens,        # [batch, draft_len+1]
4
    'new_tokens_lens': num_accepted_tokens,  # [batch]
5
    'next_draft_tokens': next_draft_tokens,  # [batch, draft_len]
6
    'next_new_tokens': next_new_tokens,      # [batch, draft_len+1]
7
}

1
slots = torch.as_tensor([r.py_seq_slot for r in sampling_requests], dtype=torch.long)
2
self.store.new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
3
self.store.next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
4
self.store.new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
5
self.store.next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)

1
if o_next_draft_tokens.shape[1] < draft_tokens_width:
2
    o_next_draft_tokens = torch.nn.functional.pad(
3
        o_next_draft_tokens, (0, draft_tokens_width - o_next_draft_tokens.shape[1]))
4
elif o_next_draft_tokens.shape[1] > draft_tokens_width:
5
    o_next_draft_tokens = o_next_draft_tokens[:, :draft_tokens_width]

1
draft_lens = [
2
    min(r.py_draft_tokens_effective_len, get_draft_token_length(r))
3
    if r.py_draft_tokens_effective_len is not None
4
    else get_draft_token_length(r)
5
    for r in sampling_requests
6
]

1
req.py_num_accepted_draft_tokens = num_new_tokens - 1
2
verified_len = (
3
    token_windows[req_idx] - 1 if token_windows is not None else runtime_draft_len)
4
drafted_len = state.draft_lens[req_idx] if state.draft_lens is not None else 0
5
req.py_num_draft_tokens_verified = min(drafted_len, verified_len)
6
req.py_rewind_len = verified_len - req.py_num_accepted_draft_tokens
7
self._request_common_handling(req, next_draft_tokens_list, runtime_draft_len)

1
request.py_draft_tokens = next_draft_tokens[request.py_seq_slot][:runtime_draft_len]
2
request.py_decoding_iter += 1

1
stats.specdec_stats.acceptance_length = float(
2
    total_accepted_tokens + num_requests_with_draft) / float(num_requests_with_draft)

1
def rejection_sampling_one_model(draft_probs, draft_token_ids, target_probs,
2
                                 seed, offset):
3
    """Verify each request's linear draft against the target distribution."""
4
    return torch.ops.trtllm.chain_speculative_sampling(
5
        draft_probs, draft_token_ids, target_probs, seed, offset)

1
def _can_use_rejection_sampling(self, spec_metadata) -> bool:
2
    return (spec_metadata.use_rejection_sampling
3
            and not spec_metadata.is_all_greedy_sample)

1
if not self._penalty_supported:
2
    raise ValueError(
3
        "repetition_penalty / presence_penalty / frequency_penalty are not "
4
        "supported with tree speculative decoding (dynamic tree) yet. ...")

1
for req in scheduled_batch.generation_requests:
2
    ...
3
    draft_len = get_draft_token_length(req)
4
    self.impl.add_token(req.py_request_id)
5
    for _ in range(max(draft_len, self._kv_reserve_draft_tokens)):
6
        self.impl.add_token(req.py_request_id)

1
self._kv_reserve_draft_tokens = self.max_total_draft_tokens
2
if (self.is_draft and spec_config is not None
3
        and getattr(spec_config, 'use_dynamic_tree', False)
4
        and getattr(spec_config, 'dynamic_tree_max_topK', 0) > 0):
5
    draft_loop_tokens = spec_config.dynamic_tree_max_topK * spec_config.max_draft_len
6
    self._kv_reserve_draft_tokens = max(self.max_total_draft_tokens, draft_loop_tokens)

1
def get_num_extra_kv_tokens(spec_config):
2
    if spec_config is None:
3
        return 0
4
    if getattr(spec_config, "_use_shared_kv_cache", False):
5
        return 0
6
    if spec_config.spec_dec_mode.use_one_engine():
7
        return spec_config.max_draft_len - 1
8
    return 0

1
def rewind_kv_cache(self, request: LlmRequest, rewind_len: int):
2
    self.impl.rewind_kv_cache(request.py_request_id, rewind_len)

1
def rewind_cache(self, request: LlmRequest, rewind_len: int):
2
    if rewind_len == 0:
3
        return
4
    request_id = request.py_request_id
5
    self.num_sequences[request_id] -= rewind_len
6
    updated_token_num = max(self.num_sequences[request_id], 0)
7
    block_count_needed = self.compute_block_count(updated_token_num,
8
                                                  self.tokens_per_block)
9
    num_rewind_pages = len(self.block_ids[request_id]) - block_count_needed
10
    if num_rewind_pages > 0:
11
        self._free_blocks(self.block_ids[request_id][-num_rewind_pages:])
12
        self.block_ids[request_id] = self.block_ids[request_id][:-num_rewind_pages]

1
def _update_kv_cache_draft_token_location(cache_manager, scheduled_batch,
2
                                          attn_metadata, kv_cache_dtype_byte_size):
3
    run_kv_cache_relocation = False
4
    for request in scheduled_batch.generation_requests:
5
        if request.state != LlmRequestState.GENERATION_COMPLETE:
6
            if (request.py_num_accepted_draft_tokens > 0
7
                    and len(request.py_num_accepted_draft_tokens_indices) > 0):
8
                run_kv_cache_relocation = True
9
    if not run_kv_cache_relocation:
10
        return
11
    (accepted_draft_token_offsets, packed_accepted_draft_tokens_indices,
12
     rewind_draft_token_separate_adjustments) = _locate_accepted_draft_tokens(requests)
13
    ...

1
assert use_paged_kv_cache, "Only paged kv cache is supported"
2
assert len(set(cache_manager.max_attention_window_vec)) == 1, (
3
    "Currently, only one max attention window size is supported.")
4
assert len(set(cache_manager.num_kv_heads_per_layer)) == 1, (
5
    "update_kv_cache_draft_token_location requires uniform num_kv_heads across all layers, ...")

1
if is_kv_cache_manager_v2 and self.max_draft_len > 0:
2
    raise ValueError(
3
        "KV Cache Connector is not supported with speculative "
4
        "decoding. Rejected draft tokens ..."
5
        "Disable speculative decoding to run a connector.")

1
def _maybe_rewind_reused_context(self, batch_llm_requests):
2
    """Recompute a tail of reused prompt tokens for drafter capture.
3

4
    Hidden-state-conditioned drafters (DFlash/DSpark) capture target
5
    hidden states only for tokens that physically pass through a target
6
    forward. A prefix-cache hit skips the reused tokens, so the drafter's
7
    cross-attention context never sees them and acceptance length drops
8
    (measured -15% AL at a 100% prefix-hit rate). Rewinding the context
9
    position after ``add_sequence_batch`` keeps the blocks reused ... while
10
    restoring the drafter's inputs.
11
    """

1
if any(window is not None and 0 < window < self.max_seq_len
2
       for window in self.max_attention_window_vec):
3
    logger.warning("context_recompute_tail is unsupported with sliding-window "
4
                   "attention layers on this KV cache manager: ...")
5
    self._spec_recompute_tail = 0
6
    return
7
if not self.enable_chunked_prefill:
8
    logger.warning("context_recompute_tail requires chunked prefill: ...")
9
    self._spec_recompute_tail = 0
10
    return

1
Generation phase 1: accept partial draft tokens
2
Target model:
3
    - input tokens: E + FGH
4
    - sampling tokens: FGXY
5
    - accepted tokens: FGX
6
    - KV cache: (ABCD) + EFGH   (H's KV cache is invalid)
7
    - hidden states: H_E, H_F, H_G, H_H (H_H is invalid)

1
def _commit_target_mamba_states(self, attn_metadata, num_accepted_tokens, batch_size):
2
    """Promote the target state at each generation request's accepted depth.
3

4
    Target verification writes intermediate GDN/Mamba states for every
5
    candidate position. The next draft pass must start from the accepted
6
    candidate, rather than the last (possibly rejected) position.
7
    """

1
def _get_eagle3_default_capture_layers(num_layers: int):
2
    return (1, num_layers // 2 - 1, num_layers - 4)

1
if self.num_layers <= 5:
2
    raise ValueError("Not enough hidden layers for default EAGLE3 capture")

1
def maybe_capture_hidden_states(self, layer_id, hidden_states, residual=None):
2
    for i, captured_layer_id in enumerate(self.layers_to_capture):
3
        if captured_layer_id == layer_id:
4
            num_tokens = self.num_capture_tokens
5
            to_save = hidden_states[:num_tokens]
6
            if residual is not None:
7
                to_save = to_save + residual[:num_tokens]
8
            inplace_slice_copy(self.hidden_states, to_save,
9
                               i * self.hidden_size, (i + 1) * self.hidden_size)
10
            break

1
if self.spec_dec_mode.is_mtp_eagle_one_model():
2
    # MTP Eagle one-model feeds the target model's hidden_states
3
    # directly to the MTP layer ... It never reads spec_metadata.hidden_states,
4
    # so leave layers_to_capture empty: this makes is_layer_capture()
5
    # return False everywhere and avoids the post-MLP/MoE fusion
6
    # disable side effect in modeling_deepseekv3 / glm / etc.
7
    self.layers_to_capture = ()

1
if i == 0:
2
    start_ids_gen = (spec_metadata.batch_indices_cuda[:num_gens] *
3
                     (runtime_draft_len + 1)).long()
4
    gather_ids_gen = (start_ids_gen + num_accepted_tokens[num_contexts:] - 1 +
5
                      attn_metadata.num_ctx_tokens)
6
    gather_ids = torch.concat([last_tokens_idx[:num_contexts], gather_ids_gen], dim=0)
7
else:
8
    gather_ids = slice(0, batch_size)

1
if spec_config is not None and spec_config.spec_dec_mode.is_pard():
2
    # PARD carries K proposals followed by K-1 placeholders.
3
    buffer_width = spec_config.get_runtime_tokens_per_gen_step(draft_len) - 1
4
    current_draft_len = ((current_num_draft_tokens + 1) // 2
5
                         if current_num_draft_tokens > 0 else 0)
6
    draft_tokens = request.py_draft_tokens[:min(current_draft_len, draft_len)]
7
    draft_tokens.extend([draft_buffer_pad] * (draft_len - len(draft_tokens)))
8
    request.py_draft_tokens = draft_tokens + [draft_buffer_pad] * (buffer_width - len(draft_tokens))

1
def get_draft_len_for_batch_size(self, batch_size: int) -> int:
2
    thresholds = list(self.draft_len_schedule.keys())
3
    idx = bisect_right(thresholds, batch_size)
4
    if idx == 0:
5
        logger.warning(f"get_draft_len_for_batch_size called with batch_size={batch_size} < 1. ...")
6
        return 0
7
    threshold = thresholds[idx - 1]
8
    return self.draft_len_schedule[threshold]

1
def update_max_total_draft_tokens(self, new_max_total_draft_tokens):
2
    self.max_total_draft_tokens = new_max_total_draft_tokens
3
    self.max_draft_len = new_max_total_draft_tokens

1
def record_acceptance_rate(self, acceptance_rate, sample_id=None):
2
    if self.disabled or self.window is None or self.window <= 0 or self.threshold is None:
3
        return False, None
4
    if acceptance_rate is None:
5
        return False, None
6
    acceptance_rate = float(acceptance_rate)
7
    if acceptance_rate < 0.0 or acceptance_rate > 1.0:
8
        raise ValueError("acceptance_rate must be in the range [0.0, 1.0], "
9
                         f"got {acceptance_rate}")

1
self.speculation_permanently_disabled = False
2
self.speculation_gate = None
3
spec_config = getattr(self.model_engine, 'spec_config', None)
4
if spec_config is not None:
5
    window = getattr(spec_config, 'acceptance_rate_window_size', None)
6
    threshold = getattr(spec_config, 'acceptance_rate_threshold', None)
7
    if ...:
8
        self.speculation_gate = SpeculationGate(window, threshold)

1
def resolve_draft_len(spec_config, scheduled_batch, *, max_draft_len, static_draft_len,
2
                      draft_len=None, speculation_permanently_disabled=False):
3
    if speculation_permanently_disabled:
4
        for request in scheduled_batch.generation_requests:
5
            request.py_draft_tokens = []
6
        return 0

1
class SaveHiddenStatesSpecMetadata(SpecMetadata):
2
    ...
3
class SaveHiddenStatesResourceManager(ResourceManager):
4
    ...

1
def register_auxiliary_state_handler(self, handler):
2
    """Register one model-owned state handler without re-parenting the module."""
3
    if not any(registered is handler for registered in self._auxiliary_state_handlers):
4
        self._auxiliary_state_handlers.append(handler)
5

6
def commit_auxiliary_speculative_states(self, num_accepted_tokens, state_indices, num_contexts):
7
    """Promote registered model-side states using accepted token counts."""
8
    for handler in self._auxiliary_state_handlers:
9
        handler.commit_speculative_states(num_accepted_tokens, state_indices, num_contexts)

1
if not IS_FLASHINFER_AVAILABLE or Version(flashinfer.__version__) < Version("0.6.4"):
2
    raise ImportError("Speculative decoding requires flashinfer>=0.6.4, "
3
                      "please install the version pinned in requirements.txt.")