TensorRT-LLM 源码完全拆解（十七）：Sampler——从 logits 到 token 的采样管线、分档执行与输出闭环

1
setup_sampler_step(scheduled_requests)          # 批被调度之后、前向之前：新请求入槽
2
sample_async(scheduled_requests, model_outputs, ...) -> SampleState
3
                                                # 前向之后：真正采样，同步返回结果
4
update_requests(state, resource_manager)        # 把结果写回请求对象、判定结束

1
class SlotSamplingParams:
2
    """TorchSampler's resolved sampling params, indexed by sequence slot."""
3
    def __init__(self, num_slots: int, device):
4
        self.modes = [SamplingMode.GREEDY] * num_slots          # host: 分组依据
5
        self.need_processed_logprobs = [False] * num_slots      # host
6
        self.need_raw_logprobs = [False] * num_slots            # host
7
        self.request_ids = [None] * num_slots                   # 用途：assert_admitted
8
        # device: (temperature, top_p, min_p) 叠在一个张量里，index_copy_ 一次搬完
9
        self.floats = torch.tensor([[1.0], [1.0], [0.0]], dtype=torch.float32,
10
                                   device=device).repeat(1, num_slots)
11
        self.top_k = torch.zeros(num_slots, dtype=torch.int32, device=device)

1
class SampleType(Enum):
2
    """Sampling tiers, ordered by capability.
3
    - FAST: only temperature / top_p / top_k / min_p. Sampled inside the graph.
4
    - FULL: every sampling feature. Sampled eagerly after the forward.
5
    """
6
    FAST = auto()
7
    FULL = auto()

1
def _can_use_fast_path(self, requests) -> bool:
2
    if self._use_beam_search:
3
        return False
4
    for req in requests:
5
        slot = req.py_seq_slot
6
        if modes[slot] is not SamplingMode.SAMPLING:   # 贪心刻意排除
7
            return False
8
        if self._top_p_decay.uses_decay(slot):        # 每步变化的 top_p，不是固定行参数
9
            return False
10
        if req.py_return_log_probs or self._request_rewrites_logits(req):
11
            return False
12
    return True

1
def _request_bypasses_in_graph_sampling(request) -> bool:
2
    """Whether sampling this request in-graph would skip work it depends on.
3

4
    The hook runs at the tail of the forward, before anything the executor
5
    does to the logits afterwards: the guided-decoding bitmask applied by
6
    ``guided_decoder.execute()``, and the user's own logits processors run
7
    by ``_execute_logit_post_processors``...
8
    """
9
    return (request.py_guided_decoding_params is not None
10
            or bool(getattr(request, "py_logits_post_processors", None))
11
            or bool(getattr(request, "py_is_draft", False)))

1
return (request.py_embedding_bias is not None
2
        or bool(getattr(request, "py_bad_words", None))
3
        or bool(getattr(request, "py_no_repeat_ngram_size", None))
4
        or has_occurrence_penalty(request)      # repetition / presence / frequency
5
        or bool(request.py_min_length))         # min_length 靠「封掉 EOS」实现

1
self.NEW_TOKENS_SHAPE = (self.max_tokens, self.max_num_sequences + 1, self.max_beam_width)
2
# 「每个 0..max_num_sequences-1 的索引都属于某个请求」，
3
# 而重放会为每个 padding 行散布一个 token，所以它们需要一个不属于任何人的槽位。

1
"""Per-row fused-kernel params the in-graph fast tier samples from.
2

3
A captured graph holds these addresses, so they are allocated once for the
4
widest batch and only their leading rows are written. ``floats`` stacks
5
temperature, top_p and min_p, in the layout of ``SlotSamplingParams.floats``.
6
"""
7
# 初始值中性：temperature 1, top_k 0（全保留）, top_p 1, min_p 0
8
# 这样即使重放读到未写入的行，也不会过滤任何东西。

1
staged_matches_batch = (
2
    self._in_graph_staged
3
    and self._in_graph_dest_indices is not None
4
    # 草稿请求复用目标请求的 py_request_id，所以 id 列表本身分辨不了
5
    # 「这个批是 drafter 的批」还是「暂存它的目标批」
6
    and not any(request.py_is_draft for request in live_requests)
7
    and [request.py_request_id for request in live_requests] == self._in_graph_request_ids)

1
"""Consumption is one-shot. The id list is not a unique key for a batch -- a draft
2
request reuses its target's py_request_id and the drafter shares this sampler --
3
so leaving the state staged would let a later drafter step with the same ids be
4
served these tokens as draft tokens. The next forward restages."""
5
self._consume_in_graph_staging()

1
single_step_greedy = not any(
2
    self._request_needs_finish_reason_kernel(request) for request in live_requests)

1
new_tokens:       (max_tokens, max_num_sequences + 1, max_beam_width)   # 最新采样的 token
2
sampled_log_probs:(max_num_sequences, max_beam_width, max_tokens)       # 只给 logprobs / beam 用
3
log_probs_store:  若干同形状缓冲（indices / ranks / topk_indices / topk_vals）

1
assert (new_tokens_shape[0] == finish_reasons_shape[0]
2
        and new_tokens_shape[1] == finish_reasons_shape[1] + 1
3
        and new_tokens_shape[2] == finish_reasons_shape[2]), \
4
    f"{new_tokens_shape} is not {finish_reasons_shape} plus a scratch slot row"

1
@dataclass(kw_only=True)
2
class SampleStateTensors:
3
    new_tokens: torch.Tensor
4
    log_probs: torch.Tensor | None = None
5

6
@dataclass(kw_only=True)
7
class SampleState(Generic[...]):
8
    requests: list[LlmRequest]
9
    device: Optional[...] = None
10
    host: Optional[...] = None
11
    sampler_event: Optional[SamplerEvent] = None
12
    runtime_draft_len: Optional[int] = None

1
class _StepIndexTranslator(ABC):
2
    """...selecting the memory locations of tokens associated with given sets of requests."""
3
    @staticmethod
4
    def calculate_request_offsets(req_num_steps, pin_memory=False):
5
        # 前缀和右移一位：req_offsets[i] = sum(num_steps[:i])
6
        req_offsets = torch.cumsum(req_num_steps, 0)
7
        req_offsets_rolled[1:] = req_offsets[:-1]
8
        req_offsets_rolled[0] = 0

1
class LogitsProcessor(nn.Module):
2
    @nvtx_range("LogitsProcessor")
3
    def forward(self, hidden_states, lm_head, attn_metadata, return_context_logits=False):
4
        if not return_context_logits:
5
            if attn_metadata is not None:
6
                last_tokens = torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1
7
                hidden_states = hidden_states[last_tokens]
8
            else:
9
                hidden_states = hidden_states[-1]
10
        logits = lm_head(hidden_states)
11
        logits = logits.float()
12
        return logits

1
for batch_index, llm_req in enumerate(context_requests):
2
    logits_begin = num_context_logits_prefix_sum[batch_index]
3
    logits_end   = num_context_logits_prefix_sum[batch_index + 1]
4
    if llm_req.py_return_context_logits:
5
        if llm_req.prepopulated_prompt_len > 0:
6
            logger.warning(f"Because of KV cache reuse, not all context logits could be produced ...")
7
        llm_req.py_result.append_context_logits(logits[logits_begin:logits_end])
8
    ...
9
total_context_logits = num_context_logits_prefix_sum[-1]
10
for batch_index, llm_req in enumerate(generation_requests):
11
    logits_begin = total_context_logits + batch_index * beam_width
12
    ...

1
prefix_sum += context_chunk_size if request.py_return_context_logits else 1

1
@dataclass(kw_only=True)
2
class SamplingRequestsMetadata:
3
    req_num_generated_tokens: torch.Tensor      # 每请求本拍产出的 token 数（入向 beam 宽度）
4
    req_num_generated_tokens_output: torch.Tensor  # 出向 beam 宽度
5
    req_num_beams: torch.Tensor                  # 每请求的 beam 数
6
    req_num_steps: torch.Tensor                  # 每请求的生成步数 = 1 + draft_len
7
    req_offsets: torch.Tensor                    # 每请求在原始 logits 里的起始偏移

1
# 1. embedding bias（logit bias）：把用户指定的 token 偏置加进去
2
apply_embedding_bias(logits_cuda, sampling_requests, sampling_requests_metadata.req_num_steps)
3

4
# 2. 罚项：repetition/presence/frequency，原地改写
5
"""Apply repetition/presence/frequency penalties in place, before the greedy fast
6
path, so both greedy and grouped-sampling logits are penalized. With beam search
7
this also re-parents the per-beam counts, which reads the predecessor map the
8
step's sampling is about to overwrite -- so it has to stay ahead of sampling."""
9
self._penalty_handler.apply(logits_cuda, sampling_requests, new_tokens=new_tokens_cuda, ...)
10

11
# 3. token ban：bad words / no-repeat-ngram / min_length，把某些 token 打成 -inf
12
self._token_ban_handler.apply_ban_list(logits_cuda, bans, new_tokens_cuda=new_tokens_cuda)

1
if self._can_use_fast_greedy_path(sampling_requests):
2
    batch_destination_indexer = _UnpackedStepIndexer(...)   # 目标散布位置
3
    fast_greedy_sample_kernel(logits_cuda, new_tokens_cuda,
4
                              batch_dest_indices_cuda, self.max_beam_width, d2t)

1
next_tokens = torch.argmax(logits_cuda, dim=-1).to(dtype=new_tokens_cuda.dtype)
2
if d2t is not None:
3
    next_tokens += d2t[next_tokens]        # Eagle3 的草稿词表 → 目标词表
4
scatter_new_tokens(next_tokens, new_tokens_cuda, batch_dest_indices, max_beam_width)

1
for req_idx, (request, slot) in enumerate(zip(requests, seq_slots)):
2
    slot_params.assert_admitted(request, slot)
3
    mode = modes[slot]
4
    needs_probs = need_processed_logprobs[slot]
5
    if mode is SamplingMode.BEAM:
6
        beam_params = BeamSearchParams.from_request(request)
7
        key = (mode, needs_probs, beam_params.group_key)
8
    else:
9
        key = (mode, needs_probs)
10
    groups[key].req_indices.append(req_idx)

1
def _sample_fused(self, logits_cuda, slots_per_row, *, return_probs):
2
    seeds = self._seed_manager.take_row_seeds(slots_per_row, device=device)
3
    slots_cuda = torch.tensor(slots_per_row, dtype=torch.int64, pin_memory=prefer_pinned()).to(device, non_blocking=True)
4
    temperature, top_k, top_p, min_p = self._slot_params.gather(slots_cuda)
5
    top_p = self._top_p_decay.override_top_p(top_p, slots_cuda)
6
    # 设备端断言，避免主机侧检查带来的每步同步
7
    torch._assert_async(~torch.any(torch.isnan(logits_cuda)))
8
    params = (logits_cuda, temperature, top_k, top_p, min_p)
9
    if return_probs:
10
        tokens, probs = fused_sample_from_logits_with_probs(*params, seed=seeds.seed, offset=seeds.offset)
11
        return tokens, probs, temperature
12
    return fused_sample_from_logits(*params, seed=seeds.seed, offset=seeds.offset), None, temperature

1
def _dims_canonically_ordered(t): return len(t.dim_order(ambiguity_check=[torch.contiguous_format])) == t.ndim
2
assert _dims_canonically_ordered(new_tokens_cuda)

1
fused_sample_from_logits             tokens  -- 目标/草稿采样器
2
fused_compute_probs_from_logits      probs   -- 目标分布的概率（供拒绝采样用）
3
fused_sample_from_logits_with_probs  both    -- 草稿采样器 + 它的草稿概率

1
def fused_sample_from_logits(logits, temperatures, top_ks, top_ps, min_ps, *,
2
                             seed=None, offset=None) -> torch.Tensor:
3
    """Sample one token per row. Returns ``[num_rows]`` int32."""
4
    return torch.ops.trtllm.fused_sample_from_logits(
5
        logits, temperatures, top_ks, top_ps, min_ps, seed, offset)

1
def resolve_advanced_sampling_filters(advanced_sampling_mode, top_k, top_p):
2
    """Resolve advanced_sampling_mode to effective (top_k, top_p) tensors.
3
    A filter the mode disables (via ``skips_top_k`` / ``skips_top_p``), or one already
4
    None, becomes None so the downstream op skips that kernel..."""
5
    eff_top_k = None if advanced_sampling_mode.skips_top_k or top_k is None else top_k
6
    eff_top_p = None if advanced_sampling_mode.skips_top_p or top_p is None else top_p
7
    return eff_top_k, eff_top_p

1
@torch.compile(options={"max-autotune": True})
2
def sample_from_logits_op(logits, temperatures, top_k=None, top_p=None, seed=None, offset=None):
3
    if top_k is not None:
4
        top_k = sanitize_top_k(top_k, logits.shape[-1])
5
        logits = top_k_mask_logits_op(logits, top_k)
6
    probs = softmax_op(logits, temperatures)
7
    if top_p is not None:
8
        return top_p_sampling_from_probs_op(probs, top_p, seed=seed, offset=offset)
9
    return sampling_from_probs_op(probs, seed=seed, offset=offset)

1
"""Per ``SamplingParams``, ``top_k == 0`` means "all logits" (top-k disabled),
2
but the flashinfer top-k kernels (``top_k_mask_logits``) break on a literal
3
0 — they mask the entire row (all-zero probs). Map any non-positive value
4
(and any oversized disable sentinel such as ``INT32_MAX``) to ``vocab_size``
5
(== keep all tokens), leaving genuine top_k values untouched."""
6
return top_k.clamp(max=vocab_size).masked_fill_(top_k <= 0, vocab_size)

1
def warmup_sampling_module() -> None:
2
    """Build flashinfer's sampling kernels now instead of on first use.
3
    flashinfer ships these kernels as source, so the first sampling call runs
4
    nvcc inline (~90s with a cold cache). Doing that inside the executor loop
5
    stalls the rank past the hang detector's threshold. Cheap no-op once built."""

1
def min_p_renorm_probs(probs, min_p):
2
    """Keep tokens with prob >= ``min_p`` times the per-row max, then renormalize."""
3
    max_probs = probs.max(dim=-1, keepdim=True).values
4
    thresholds = min_p * max_probs
5
    probs = probs.masked_fill(probs < thresholds, 0.0)
6
    return probs / probs.sum(dim=-1, keepdim=True)
7

8
def greedy_search_sampling_batch(logits, *, return_probs=True):
9
    next_tokens = torch.argmax(logits, dim=-1)
10
    ...

1
# Map each logits row to its owning request with a broadcasted range comparison.
2
# This is O(T * R), but T and R are both small (rows per step x requests) and the
3
# whole thing fuses into the surrounding elementwise graph, so it measures faster
4
# than either a searchsorted lookup or a repeat_interleave expansion. Notably
5
# torch.repeat_interleave must NOT be used here: its output length is
6
# sum(num_steps), which -- without an explicit host-provided output_size -- torch
7
# reads back from the device, and that per-step D2H sync destroys the overlap
8
# between the sampler's host work and the model forward (measured ~20x slower).

1
self._seeds   = torch.full((max_num_sequences,), global_seed, dtype=torch.int64)  # 按槽位
2
self._offsets = torch.zeros((max_num_sequences,), dtype=torch.int64)              # 按槽位
3
self._slot_owner: list[Optional[int]] = [None] * max_num_sequences                # 槽位当前属主
4
self._unseeded_offset = UNSEEDED_OFFSET_BASE     # 1 << 62

1
# First offset of the rows that carry no user seed. They all share one seed, which a user
2
# seed may equal, so they take offsets far beyond any a seeded stream reaches counting up
3
# from 0.

1
def observe(self, requests) -> None:
2
    """Seed any slot whose occupant changed.
3
    ...Resetting the offset on ownership change is what makes a seeded request start at
4
    the beginning of its stream instead of wherever the slot's previous occupant left off."""
5
    self._batch_is_draft = bool(requests) and requests[0].py_is_draft
6
    if self._batch_is_draft:
7
        return
8
    for request in requests:
9
        if self._slot_owner[seq_slot] != request_id:
10
            self._slot_owner[seq_slot] = request_id
11
            seed = request_random_seed(request)
12
            self._seeds[seq_slot] = self._global_seed if seed is None else seed
13
            self._offsets[seq_slot] = 0

1
# Only the live rows consume RNG stream. Padding rows borrow a slot that
2
# may belong to an active request this batch simply did not schedule;
3
# advancing it would move that request's Philox offset by an amount that
4
# depends on concurrent load, breaking the guarantee _SeedManager exists
5
# for -- that a seeded request's stream depends only on its own draws.
6
live_slots_per_row = slots_per_row[:live_rows]
7
row_seeds = self._seed_manager.take_row_seeds(live_slots_per_row, device=params.seed.device)

1
"""The C++ ``SamplingConfig.randomSeed`` is ``uint64``, so a user seed in
2
``[2**63, 2**64)`` (e.g. from ``random.getrandbits(64)``) does not fit the
3
int64 tensors the torch sampler stores it in. Reinterpret the bit pattern as
4
signed rather than rejecting or clamping: the RNG consumes the seed as an
5
opaque bit pattern, so this keeps distinct seeds distinct, and it avoids an
6
overflow that would abort the whole sampling step rather than one request."""
7
return seed - (1 << 64) if seed >= (1 << 63) else seed

1
counts_cuda: int32[num_slots, vocab_size]，原地累加
2
presence_prefix_cuda: bool[num_slots, vocab_size]   # prompt_ignore_length 用的前缀存在掩码

1
# accumulate=True sums repeated (slot, token) pairs -> occurrence count.
2
counts_cuda.index_put_((counted_slots, counted_tokens), ones, accumulate=True)

1
def _update_beam_occurrence_counts_impl(counts_cuda, active_cuda, has_previous_token_cuda,
2
                                        beam_slot_cuda, new_tokens, predecessor_beams,
3
                                        seq_slots, request_num_beams, max_beam_width):
4
    """Re-parent every beam onto the beam it continues, then fold in its token.
5
    Beam ``b`` of slot ``s`` takes over the counts of ``predecessor_beams[s, b]``
6
    before this step's token is folded in."""

1
"""``armed`` gates per *slot* (``.unsqueeze(1)``), so beams past the current width are
2
re-parented too, from a clamped and possibly stale parent -- safe because nothing
3
reads such a row until the width grows to cover it... The gate is device-side, so this
4
runs unconditionally rather than pay a D2H sync."""
5
parent = predecessor_beams[seq_slots].to(torch.int64).clamp(0, max_beam_width - 1)

1
bans = self._token_ban_handler.generate_ban_list(
2
    sampling_requests, req_num_steps.tolist(), req_num_beams.tolist(), ngram_sizes,
3
    stale_by_one=stale_by_one, pending_steps=pending_steps)
4
self._token_ban_handler.apply_ban_list(logits_cuda, bans, new_tokens_cuda=new_tokens_cuda)

1
"""Overlap-scheduler stale flags (per request): True when the host token history lags
2
the device by one token. Only the overlap handler consumes them; computed here as it
3
needs sampler state. Only the suffix-matching bans (bad words, no-repeat ngram) care:
4
min_length bans EOS from a length count, never from token values, so a lagging history
5
cannot mismatch it."""

1
self._token_ban_handler = (OverlappedTokenBanHandler() if self._track_pending_steps
2
                           else SynchronousTokenBanHandler())

1
runtime_top_p = initial_top_p                                if token == reset_id
2
              = max(runtime_top_p * top_p_decay, top_p_min)  otherwise

1
1. Admission (setup_for_new_requests): 清掉新入槽位的陈旧成员标记，再重新登记；
2
   同时维护主机侧 _slots 集合（O(1) 早退）与设备侧 is_top_p_decay_slot_cuda 门控镜像。
3
2. Pre-sample (override_top_p): 用融合 gather 把 decay-active 行的 top_p 换成运行时值。
4
3. Post-sample (update_after_sample): 用融合 update 就地跑递推式。
5
4. Finish (retire_slot): 该槽位退出成员集合，早退重新武装。

1
"""Per-slot bool gate (device mirror of ``TopPDecayHandler._slots``). Lets the
2
fused post-sample update op filter decay-active slots on the GPU,
3
avoiding a host-side ``.tolist()`` / set intersection each step."""

1
"""Reject unsupported combinations for a top-p-decay-active request.
2
Top-p decay is supported only for single-token decode steps without beam search.
3
Called from validate_request (request admission), so a violating request is failed
4
individually instead of aborting the whole batch."""

1
# cudagraphs is unsafe here (the update mutates persistent per-slot state in place
2
# and the gather's output is consumed outside the compiled region; cudagraph static
3
# output buffers get overwritten by subsequent replays).

1
# NB: The sampler is either called directly by PyExecutor, for the target model,
2
#     or by a drafter's prepare_draft_tokens(), for the draft model. In the former
3
#     case there are 1 + get_draft_token_length(request) tokens per request. In the
4
#     latter case, there is always only 1 token per request because draft tokens
5
#     are sampled one-by-one.

1
"""Batches are homogeneous by construction: a drafter builds all-draft batches for its
2
sample_async/update_requests calls on this shared sampler, and PyExecutor's batches are
3
all-target. The pending-steps accounting relies on this to skip draft batches wholesale;
4
assert it so a mixed batch fails loudly instead of silently corrupting the counters."""

1
def apply_d2t(tokens, model_outputs) -> None:
2
    """Applies draft-to-target token translation table. Modifies tokens in-place."""
3
    if "d2t" in model_outputs:
4
        d2t = model_outputs["d2t"][tokens]
5
        tokens += d2t

1
needs_d2t = "d2t" in model_outputs
2
if needs_d2t and any(group.mode is not SamplingMode.GREEDY for group in groups):
3
    raise ValueError("d2t does not yet support non-greedy sampling")

1
proc_lp_renorm_offset_cuda = (proc_lp_softmax_cuda.gather(dim=1, index=proc_lp_max_logit_indices.unsqueeze(-1))
2
                              .log().squeeze(-1) - proc_lp_max_logits_cuda)

1
batch_raw_logprob_indices_cuda = logits_cuda_indexer[torch.tensor(reqs_indices_needing_raw_logprobs, ...)]

1
"""Sample a token AND return the filtered, renormalized distribution.
2
The rejection path needs both, and needs them to come from the *same* filtering, or
3
acceptance is computed against a distribution neither side sampled from."""

1
num_accepted = self._two_model_spec_dec.process_draft_tokens(
2
    req, new_tokens_list=new_tokens_list, finish_reasons=finish_reasons)
3
if (actual_draft_len := get_draft_token_length(req)) > 0:
4
    req.py_num_accepted_draft_tokens = num_accepted
5
    # Pair the acceptance count with the real proposal count of the same step:
6
    # py_draft_tokens is padded to the static max for CUDA graphs, so use the
7
    # pre-padding length the drafter recorded. py_rewind_len keeps the padded
8
    # length (padding occupies KV cache).
9
    effective_len = req.py_draft_tokens_effective_len
10
    req.py_num_draft_tokens_verified = (min(effective_len, actual_draft_len)
11
                                        if effective_len is not None else actual_draft_len)
12
    req.py_rewind_len = actual_draft_len - num_accepted

1
# req.state can become GENERATION_COMPLETE within this loop iteration
2
# (e.g. process_draft_tokens -> _handle_stop_criteria -> finish_by),
3
# so the comparison is valid at runtime; mypy narrowed it away via the
4
# `continue` at the top of the loop and cannot see the mutating calls.

1
if actual_draft_len > 0:
2
    num_new_tokens = req.get_num_tokens(DEFAULT_BEAM_IDX) - num_tokens_before
3
    if num_new_tokens > 0:
4
        confirmed_tokens = req.get_tokens(DEFAULT_BEAM_IDX)[-num_new_tokens:]
5
        finalized_token_updates.append((req.py_seq_slot, confirmed_tokens))
6
...
7
self._penalty_handler.update_token_counts(finalized_token_updates)

1
# Reuse sample_async's qualification instead of rechecking every request after
2
# the asynchronous sample completes.
3
if state.host.single_step_greedy:
4
    self._update_requests_single_beam_single_step(state)
5
    return

1
add_new_tokens_to_requests(requests, new_tokens, DEFAULT_BEAM_IDX)
2
assert state.host.finish_reasons is None   # 这条路径刻意不算设备 finish reason
3
for request, new_token in zip(requests, new_tokens):
4
    # Order matches handle_stop_criteria: END_ID, then LENGTH, then STOP_WORDS,
5
    # so a terminal EOS at the token limit is reported as END_ID.
6
    if new_token == request.py_end_id:
7
        request.finish_by(FinishReason.END_ID, DEFAULT_BEAM_IDX)
8
    elif (request.max_beam_num_tokens - request.py_orig_prompt_len >= request.py_max_new_tokens
9
          or request.max_beam_num_tokens >= self.max_seq_len):
10
        request.finish_by(FinishReason.LENGTH, DEFAULT_BEAM_IDX)
11
    elif meet_stop_token_criteria(request, new_token, DEFAULT_BEAM_IDX):
12
        request.finish_by(FinishReason.STOP_WORDS, DEFAULT_BEAM_IDX)

1
# Fast-path (batched pybind): collapse per-request pybind chatter into one batched
2
# add_new_tokens_to_requests call. Single-pass eligibility check with early-break;
3
# falls through when any invariant breaks.
4
if self._batch_fastpath_eligible and logprobs_state_list is None:
5
    for req in state.requests:
6
        if get_draft_token_length(req) != 0 or req.py_stop_words_list:
7
            fastpath_ok = False
8
            break

1
# Context-only (disaggregated prefill) requests reach the sampler already flagged
2
# as finished -- the context server is done with them -- but they are migrating, not
3
# completing: the generation server decodes the rest. Finalizing here would rewrite
4
# the beam rows from the beam history, and after a single step only beam 0 has
5
# history; every other beam is all BEAM_SEARCH_PAD_TOKEN ... The handoff reads the
6
# per-beam first token back out via getTokens().back() (llmRequest.cpp), which would
7
# then hand the generation server the prompt tail instead of that beam's token.
8
if (not req.is_context_only_request and request_finishing
9
    and (beam_history := _maybe_build_beam_history(req_idx)) is not None):
10
    finalize_beam(req, beam_history)
11
else:
12
    for beam_idx in range(_get_beam_width_out(req)):
13
        add_token(req, new_tokens_list, beam_idx=beam_idx)

1
_valid_finish_reasons = {FinishReason.END_ID, FinishReason.LENGTH, FinishReason.STOP_WORDS}

1
if not single_step_greedy:
2
    # write_finish_reasons shifts stop-word history by the accepted-draft count,
3
    # which update_requests only refreshes after this step. Nothing was speculated
4
    # here -- the tier requires zero draft tokens -- so a count left over from an
5
    # earlier speculative step would shift the comparison and could miss the stop word.
6
    accepted = self._finish_reasons_handler.store.num_accepted_draft_tokens_host
7
    for request in live_requests:
8
        if request.py_stop_words_list:
9
            accepted[request.py_seq_slot] = 0

1
"""A request is complete once all of the beams it actually uses have a finish
2
reason, i.e. once its leading ``py_beam_width`` entries are all set. Rather than
3
reducing each request's row separately, reduce the whole tensor once and return,
4
per slot, how many leading beams have finished."""
5
unfinished = finish_reasons == FinishReason.NOT_FINISHED.value
6
return (unfinished.cumsum(dim=1) == 0).sum(dim=1).tolist()

1
self._penalty_handler.update_token_counts(finalized_token_updates)

1
req.py_decoding_iter += 1