vllm.v1.worker.gpu.spec_decode.suffix.speculator

GPU suffix decoding speculator for the V2 model runner.

Extends the n-gram speculator with the two parts of SuffixDecoding (https://arxiv.org/abs/2411.04975) that carry most of its acceptance gain:

• Variable-length matching: drafts continue the occurrence whose preceding tokens match the longest suffix of the request's history (up to suffix_decoding_max_tree_depth), instead of a fixed n-gram window.
• Cross-request memory: responses of finished requests are appended to a device-resident token corpus. A hash index keyed on the last 2 and last 4 tokens gives candidate positions, and each candidate is extended to the longest suffix match, so the index is only an entry point for suffix matching.

The propose path never synchronizes with the host. Ingestion reads token history and lengths directly from RequestState and processes finished requests in sorted id order; index updates are order-independent, so every TP rank holds an identical corpus without any collective.

Classes:

• SuffixSpeculator –

  V2-compatible GPU suffix decoding speculator.

SuffixSpeculator

Bases: BaseSpeculator

V2-compatible GPU suffix decoding speculator.

Methods:

• on_requests_finished –

  Add finished responses to the corpus before their slots are freed.

• reset_state –

  Forget everything learned so far (used after the warmup requests).

Source code in vllm/v1/worker/gpu/spec_decode/suffix/speculator.py

class SuffixSpeculator(BaseSpeculator):
    """V2-compatible GPU suffix decoding speculator."""

    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
        req_states: RequestState,
    ):
        if not HAS_TRITON:
            raise RuntimeError("suffix speculative decoding requires Triton.")
        spec = vllm_config.speculative_config
        assert spec is not None
        self.vllm_config = vllm_config
        self.device = device
        self.req_states = req_states
        self.num_speculative_steps: int = spec.num_speculative_tokens
        self.max_depth: int = spec.suffix_decoding_max_tree_depth
        self.max_num_reqs: int = vllm_config.scheduler_config.max_num_seqs
        self.max_model_len: int = vllm_config.model_config.max_model_len

        L = self.max_model_len
        if L >= 1024:
            self.block_l = 256
        elif L >= 256:
            self.block_l = 128
        else:
            self.block_l = max(16, triton.next_power_of_2(max(L, 1)))
        self.n_blocks = triton.cdiv(L, self.block_l)
        self.scratch = torch.zeros(
            (self.max_num_reqs, self.n_blocks), dtype=torch.int64, device=device
        )

        # Corpus of finished responses separated by -1, plus two hash tables
        # keyed on the last 4 and last 2 tokens before each position. A
        # capacity of 0 disables it.
        self.corpus_len: int = (
            spec.suffix_decoding_corpus_tokens
            if spec.suffix_decoding_max_cached_requests != 0
            else 0
        )
        if self.corpus_len > 0 and self.max_depth < 2:
            # Corpus matches need the 2-token hash key.
            logger.warning(
                "suffix_decoding_max_tree_depth=%d is too short for "
                "cross-request matching (needs >= 2); drafting from each "
                "request's own history only.",
                self.max_depth,
            )
            self.corpus_len = 0
        self.has_corpus = self.corpus_len > 0
        self.num_buckets = triton.next_power_of_2(max(16, self.corpus_len // 4))
        self.corpus = torch.full(
            (max(self.corpus_len, 1),), -1, dtype=torch.int32, device=device
        )
        self.table = torch.full(
            (2, self.num_buckets if self.has_corpus else 1, _BUCKET_SLOTS),
            -1,
            dtype=torch.int64,
            device=device,
        )
        self.corpus_head = torch.zeros(1, dtype=torch.int64, device=device)
        self.doc_seq = torch.zeros(1, dtype=torch.int64, device=device)
        # Host mirror of doc_seq: index entries keep it in 31 bits.
        self.num_docs = 0
        self.ingest_out = torch.zeros(2, dtype=torch.int64, device=device)

        # Batch-ordered draft output, scattered into RequestState.draft_tokens
        # by the model runner (same contract as the model-based speculators).
        self.drafts = torch.zeros(
            (self.max_num_reqs, self.num_speculative_steps),
            dtype=torch.int64,
            device=device,
        )
        # Number of real (non-filler) draft tokens per batch row.
        self.num_valid = torch.zeros(
            self.max_num_reqs, dtype=torch.int32, device=device
        )

        # Adaptive verification (optional): per-slot confidences from
        # acceptance statistics learned online.
        self.enable_adaptive_verification = spec.enable_adaptive_verification
        self.draft_token_confidence_probs = torch.zeros(
            (self.max_num_reqs, self.num_speculative_steps),
            dtype=torch.float32,
            device=device,
        )
        self.accept_stats = torch.zeros(
            (2 * _NUM_LEN_BUCKETS, self.num_speculative_steps, 2),
            dtype=torch.int64,
            device=device,
        )
        self.last_bucket = torch.full(
            (self.max_num_reqs,), -1, dtype=torch.int32, device=device
        )
        self.last_valid = torch.zeros(
            self.max_num_reqs, dtype=torch.int32, device=device
        )
        self.num_steps = 0
        # Slots reassigned since the last propose; their stats bookkeeping
        # still belongs to the previous request.
        self.new_slots: list[int] = []

    def add_request(self, req_index: int) -> None:
        if self.enable_adaptive_verification:
            self.new_slots.append(req_index)

    def reset_state(self) -> None:
        """Forget everything learned so far (used after the warmup requests)."""
        self.corpus.fill_(-1)
        self.table.fill_(-1)
        self.corpus_head.zero_()
        self.doc_seq.zero_()
        self.num_docs = 0
        self.accept_stats.zero_()
        self.last_bucket.fill_(-1)
        self.last_valid.zero_()
        self.num_steps = 0
        self.new_slots.clear()

    def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None:
        del cudagraph_mode

    def capture(self) -> None:
        return None

    def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None:
        """Add finished responses to the corpus before their slots are freed."""
        if not self.has_corpus:
            return
        req_states = self.req_states
        # Sorted so every TP rank ingests documents in the same order. Starts
        # come from the host copy of prompt_len: the device-visible copy is
        # republished every step and may be reused for a new request before
        # this kernel runs.
        docs = []
        for req_id in sorted(finished_req_ids):
            slot = req_states.req_id_to_index.get(req_id)
            if slot is None:
                continue
            # Keep the prompt's last few tokens (typically the chat template's
            # assistant header) so the first response tokens are indexed by
            # the tokens that precede them.
            start = max(int(req_states.prompt_len.np[slot]) - _PROMPT_CONTEXT, 0)
            docs += [slot, start]
        if not docs:
            return
        num_docs = len(docs) // 2
        if self.num_docs + num_docs >= 1 << 31:
            # Index entries pack doc_seq into 31 bits; start a new index (the
            # corpus tokens stay, only their entries are dropped).
            self.table.fill_(-1)
            self.doc_seq.zero_()
            self.num_docs = 0
        self.num_docs += num_docs
        docs_gpu = async_tensor_h2d(docs, self.device, torch.int32)
        token_ids = self.req_states.all_token_ids.gpu
        _suffix_ingest_kernel[
            (num_docs, triton.cdiv(self.max_model_len, _INGEST_BLOCK))
        ](
            token_ids,
            token_ids.stride(0),
            docs_gpu,
            num_docs,
            req_states.total_len.gpu,
            self.corpus,
            self.corpus_len,
            self.corpus_head,
            self.doc_seq,
            self.ingest_out[0:1],
            self.ingest_out[1:2],
            self.table,
            HASH_MUL=_HASH_MUL,
            NUM_BUCKETS=self.num_buckets,
            SLOTS=_BUCKET_SLOTS,
            BLOCK=_INGEST_BLOCK,
            num_warps=4,
        )
        self.corpus_head.copy_(self.ingest_out[0:1])
        self.doc_seq.copy_(self.ingest_out[1:2])

    @torch.inference_mode()
    def propose(
        self,
        input_batch: InputBatch,
        attn_metadata: Any,
        slot_mappings: Any,
        last_hidden_states: torch.Tensor,
        aux_hidden_states: list[torch.Tensor] | None,
        num_sampled: torch.Tensor,
        num_rejected: torch.Tensor,
        last_sampled: torch.Tensor,
        next_prefill_tokens: torch.Tensor,
        temperature: torch.Tensor,
        seeds: torch.Tensor,
        dp_sync_state: DPSyncState | None = None,
        dummy_run: bool = False,
        skip_attn_for_dummy_run: bool = False,
        mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
        is_profile: bool = False,
        num_speculative_tokens: int | None = None,
    ) -> torch.Tensor:
        num_reqs = input_batch.num_reqs
        if dummy_run:
            return self.drafts[:num_reqs]

        token_ids = self.req_states.all_token_ids.gpu
        total_len = self.req_states.total_len.gpu
        idx_mapping = input_batch.idx_mapping
        k_po2 = max(1, triton.next_power_of_2(self.num_speculative_steps))

        if self.enable_adaptive_verification:
            if self.new_slots:
                slots = async_tensor_h2d(self.new_slots, self.device, torch.int64)
                self.last_bucket.index_fill_(0, slots, -1)
                self.new_slots.clear()
            _suffix_observe_kernel[(num_reqs,)](
                idx_mapping,
                num_sampled,
                num_rejected,
                self.accept_stats,
                self.last_bucket,
                self.last_valid,
                K=self.num_speculative_steps,
                K_PO2=k_po2,
                num_warps=1,
            )
            self.num_steps += 1
            if self.num_steps % _STATS_HALF_LIFE == 0:
                self.accept_stats.bitwise_right_shift_(1)

        _suffix_local_scan_kernel[(num_reqs, self.n_blocks)](
            token_ids,
            token_ids.stride(0),
            idx_mapping,
            total_len,
            num_sampled,
            self.scratch,
            self.scratch.stride(0),
            self.max_model_len + 1,
            MAX_DEPTH=self.max_depth,
            BLOCK=self.block_l,
            num_warps=4,
        )
        _suffix_finalize_kernel[(num_reqs,)](
            token_ids,
            token_ids.stride(0),
            idx_mapping,
            total_len,
            num_sampled,
            last_sampled.view(-1),
            self.scratch,
            self.scratch.stride(0),
            self.n_blocks,
            self.max_model_len + 1,
            self.corpus,
            self.corpus.shape[0],
            self.table,
            self.drafts,
            self.num_valid,
            self.accept_stats,
            self.last_bucket,
            self.last_valid,
            self.draft_token_confidence_probs,
            HAS_CORPUS=self.has_corpus,
            ADAPTIVE=self.enable_adaptive_verification,
            HASH_MUL=_HASH_MUL,
            NUM_BUCKETS=self.num_buckets,
            SLOTS=_BUCKET_SLOTS,
            MAX_DEPTH=self.max_depth,
            DEPTH_PO2=triton.next_power_of_2(self.max_depth),
            K=self.num_speculative_steps,
            K_PO2=k_po2,
            NB_PO2=max(1, triton.next_power_of_2(self.n_blocks)),
            num_warps=4,
        )
        return self.drafts[:num_reqs]

on_requests_finished(finished_req_ids)

Add finished responses to the corpus before their slots are freed.

Source code in vllm/v1/worker/gpu/spec_decode/suffix/speculator.py

def on_requests_finished(self, finished_req_ids: Iterable[str]) -> None:
    """Add finished responses to the corpus before their slots are freed."""
    if not self.has_corpus:
        return
    req_states = self.req_states
    # Sorted so every TP rank ingests documents in the same order. Starts
    # come from the host copy of prompt_len: the device-visible copy is
    # republished every step and may be reused for a new request before
    # this kernel runs.
    docs = []
    for req_id in sorted(finished_req_ids):
        slot = req_states.req_id_to_index.get(req_id)
        if slot is None:
            continue
        # Keep the prompt's last few tokens (typically the chat template's
        # assistant header) so the first response tokens are indexed by
        # the tokens that precede them.
        start = max(int(req_states.prompt_len.np[slot]) - _PROMPT_CONTEXT, 0)
        docs += [slot, start]
    if not docs:
        return
    num_docs = len(docs) // 2
    if self.num_docs + num_docs >= 1 << 31:
        # Index entries pack doc_seq into 31 bits; start a new index (the
        # corpus tokens stay, only their entries are dropped).
        self.table.fill_(-1)
        self.doc_seq.zero_()
        self.num_docs = 0
    self.num_docs += num_docs
    docs_gpu = async_tensor_h2d(docs, self.device, torch.int32)
    token_ids = self.req_states.all_token_ids.gpu
    _suffix_ingest_kernel[
        (num_docs, triton.cdiv(self.max_model_len, _INGEST_BLOCK))
    ](
        token_ids,
        token_ids.stride(0),
        docs_gpu,
        num_docs,
        req_states.total_len.gpu,
        self.corpus,
        self.corpus_len,
        self.corpus_head,
        self.doc_seq,
        self.ingest_out[0:1],
        self.ingest_out[1:2],
        self.table,
        HASH_MUL=_HASH_MUL,
        NUM_BUCKETS=self.num_buckets,
        SLOTS=_BUCKET_SLOTS,
        BLOCK=_INGEST_BLOCK,
        num_warps=4,
    )
    self.corpus_head.copy_(self.ingest_out[0:1])
    self.doc_seq.copy_(self.ingest_out[1:2])

reset_state()

Forget everything learned so far (used after the warmup requests).

Source code in vllm/v1/worker/gpu/spec_decode/suffix/speculator.py

def reset_state(self) -> None:
    """Forget everything learned so far (used after the warmup requests)."""
    self.corpus.fill_(-1)
    self.table.fill_(-1)
    self.corpus_head.zero_()
    self.doc_seq.zero_()
    self.num_docs = 0
    self.accept_stats.zero_()
    self.last_bucket.fill_(-1)
    self.last_valid.zero_()
    self.num_steps = 0
    self.new_slots.clear()

_corpus_draft(row, seq_len, table_ptr, corpus_ptr, corpus_len, HASH_MUL, NUM_BUCKETS, SLOTS, MAX_DEPTH, DEPTH_PO2, K, K_PO2)

(match_len, draft, num_tokens) from the corpus for one request.

Candidates are the slots of the request's 4-token and 2-token buckets. Each is verified and extended against the corpus, so hash collisions and entries for overwritten documents only cost a slot. The draft is then chosen by frequency, as in SuffixDecoding: for each candidate match length L, the candidates matching at least L tokens vote token by token (prefix-consistent; ties go to the newest candidate), and the chain with the highest sum of per-token probabilities wins (ties: longer L).

Source code in vllm/v1/worker/gpu/spec_decode/suffix/speculator.py

@triton.jit
def _corpus_draft(
    row,
    seq_len,
    table_ptr,
    corpus_ptr,
    corpus_len,
    HASH_MUL: tl.constexpr,
    NUM_BUCKETS: tl.constexpr,
    SLOTS: tl.constexpr,
    MAX_DEPTH: tl.constexpr,
    DEPTH_PO2: tl.constexpr,
    K: tl.constexpr,
    K_PO2: tl.constexpr,
):
    """(match_len, draft, num_tokens) from the corpus for one request.

    Candidates are the slots of the request's 4-token and 2-token buckets.
    Each is verified and extended against the corpus, so hash collisions and
    entries for overwritten documents only cost a slot. The draft is then
    chosen by frequency, as in SuffixDecoding: for each candidate match
    length L, the candidates matching at least L tokens vote token by token
    (prefix-consistent; ties go to the newest candidate), and the chain with
    the highest sum of per-token probabilities wins (ties: longer L).
    """
    C: tl.constexpr = 2 * SLOTS
    c_iota = tl.arange(0, C).to(tl.int64)
    d_iota = tl.arange(0, DEPTH_PO2).to(tl.int64)
    k_iota = tl.arange(0, K_PO2).to(tl.int64)

    tail_pos = seq_len - 1 - d_iota
    tail = tl.load(
        row + tail_pos, mask=(d_iota < MAX_DEPTH) & (tail_pos >= 0), other=-2
    )
    # Hash the trailing tokens newest-first, as ingestion does.
    hv = tl.zeros((), tl.int64)
    for d in tl.static_range(4):
        tok = tl.load(row + seq_len - 1 - d, mask=seq_len - 1 - d >= 0, other=0)
        hv = (hv ^ tok.to(tl.int64)) * HASH_MUL
        if d == 1:
            bucket2 = _bucket(hv, NUM_BUCKETS)
    bucket4 = _bucket(hv, NUM_BUCKETS)
    is4 = c_iota < SLOTS
    key_len = tl.where(is4, 4, 2)
    slot = tl.where(
        is4, bucket4 * SLOTS + c_iota, (NUM_BUCKETS + bucket2) * SLOTS + c_iota - SLOTS
    )
    vals = tl.load(table_ptr + slot)
    cpos = vals & 0xFFFFFFFF

    src_pos = cpos[:, None] - 1 - d_iota[None, :]
    src = tl.load(
        corpus_ptr + src_pos,
        mask=(vals[:, None] >= 0) & (src_pos >= 0) & (d_iota[None, :] < MAX_DEPTH),
        other=-1,
    )
    miss = (src != tail[None, :]).to(tl.int32)
    c_len = tl.sum((tl.cumsum(miss, axis=1) == 0).to(tl.int64), axis=1)
    nxt = tl.load(corpus_ptr + cpos, mask=(vals >= 0) & (cpos < corpus_len), other=-1)
    ok = (vals >= 0) & (nxt >= 0) & (c_len >= key_len) & (seq_len >= key_len)
    # Both buckets often hold the same position; count it once.
    dup = (
        (cpos[:, None] == cpos[None, :])
        & ok[:, None]
        & (c_iota[:, None] < c_iota[None, :])
    )
    ok = ok & (tl.max(dup.to(tl.int32), axis=0) == 0)

    # Continuations, cut at the first document separator.
    cont_pos = cpos[:, None] + k_iota[None, :]
    cont = tl.load(
        corpus_ptr + cont_pos,
        mask=ok[:, None] & (k_iota[None, :] < K) & (cont_pos < corpus_len),
        other=-1,
    ).to(tl.int64)
    cont = tl.where(tl.cumsum((cont < 0).to(tl.int32), axis=1) == 0, cont, -1)

    # Candidates sharing a match length draft identical chains; score each
    # distinct length once.
    same_len = (
        (c_len[:, None] == c_len[None, :])
        & ok[:, None]
        & (c_iota[:, None] < c_iota[None, :])
    )
    first = ok & (tl.max(same_len.to(tl.int32), axis=0) == 0)

    # A token's probability is its vote count over the candidates still
    # alive, which is the previous token's count, so a chain's score
    # telescopes to (sum of its vote counts) / (candidates at the start).
    # Scores are compared as these integer fractions, so ties are exact.
    best_num = tl.full((), -1, tl.int64)
    best_den = tl.full((), 1, tl.int64)
    best_len = tl.zeros((), tl.int64)
    best_n = tl.zeros((), tl.int64)
    best_chain = tl.zeros([K_PO2], tl.int64)
    for i in range(C):
        sel = c_iota == i
        if tl.sum(tl.where(sel, first.to(tl.int32), 0), axis=0) > 0:
            L = tl.sum(tl.where(sel, c_len, 0), axis=0)
            alive = ok & (c_len >= L)
            den = tl.maximum(tl.sum(alive.to(tl.int64), axis=0), 1)
            go = L > 0
            num = tl.zeros((), tl.int64)
            n_emit = tl.zeros((), tl.int64)
            chain = tl.zeros([K_PO2], tl.int64)
            for d in tl.static_range(K):
                col = tl.sum(tl.where(k_iota[None, :] == d, cont, 0), axis=1)
                has = alive & (col >= 0)
                cnt = tl.sum(
                    ((col[:, None] == col[None, :]) & has[None, :]).to(tl.int32), axis=1
                )
                cnt = tl.where(has, cnt, 0)
                cmax = tl.max(cnt, axis=0)
                # Most frequent next token; ties go to the newest candidate.
                newest = tl.max(tl.where(has & (cnt == cmax), vals, -1), axis=0)
                next_tok = tl.sum(tl.where(has & (vals == newest), col, 0), axis=0)
                go = go & (cmax > 0)
                num = tl.where(go, num + cmax.to(tl.int64), num)
                chain = tl.where((k_iota == d) & go, next_tok, chain)
                n_emit += go.to(tl.int64)
                alive = alive & has & (col == next_tok)
            lhs = num * best_den
            rhs = best_num * den
            take = (lhs > rhs) | ((lhs == rhs) & (best_len < L))
            best_num = tl.where(take, num, best_num)
            best_den = tl.where(take, den, best_den)
            best_len = tl.where(take, L, best_len)
            best_n = tl.where(take, n_emit, best_n)
            best_chain = tl.where(take, chain, best_chain)
    return best_len, best_chain, best_n

_suffix_ingest_kernel(token_ids_ptr, token_ids_stride, docs_ptr, num_docs, total_len_ptr, corpus_ptr, corpus_len, head_ptr, doc_seq_ptr, head_out_ptr, doc_seq_out_ptr, table_ptr, HASH_MUL, NUM_BUCKETS, SLOTS, BLOCK)

Append finished responses to the corpus ring and index them, with one program per (document, block).

Documents are placed one after another. A document that does not fit before the end of the ring wraps to offset 0, so no document straddles the boundary. Older documents are overwritten in FIFO order (including earlier documents of the same batch, if the batch exceeds the ring); their index entries become stale and are rejected by verification at lookup. Each program replays the batch's placement to find its offset and skips positions a later document of the batch overwrites, so every position is written by exactly one program and the result does not depend on execution order.

Source code in vllm/v1/worker/gpu/spec_decode/suffix/speculator.py

@triton.jit(do_not_specialize=["num_docs"])
def _suffix_ingest_kernel(
    token_ids_ptr,  # *int32  [max_num_reqs, token_ids_stride]
    token_ids_stride,
    docs_ptr,  # *int32  [F, 2]  (slot, start) of finished requests, sorted
    num_docs,
    total_len_ptr,  # *int32  [max_num_reqs]
    corpus_ptr,  # *int32  [corpus_len]
    corpus_len,
    head_ptr,  # *int64  [1]  next write offset
    doc_seq_ptr,  # *int64  [1]  documents ingested so far
    head_out_ptr,  # *int64  [1]
    doc_seq_out_ptr,  # *int64  [1]
    table_ptr,  # *int64  [2, NUM_BUCKETS, SLOTS]
    HASH_MUL: tl.constexpr,
    NUM_BUCKETS: tl.constexpr,
    SLOTS: tl.constexpr,
    BLOCK: tl.constexpr,
):
    """Append finished responses to the corpus ring and index them, with one
    program per (document, block).

    Documents are placed one after another. A document that does not fit
    before the end of the ring wraps to offset 0, so no document straddles
    the boundary. Older documents are overwritten in FIFO order (including
    earlier documents of the same batch, if the batch exceeds the ring);
    their index entries become stale and are rejected by verification at
    lookup. Each program replays the batch's placement to find its offset
    and skips positions a later document of the batch overwrites, so every
    position is written by exactly one program and the result does not
    depend on execution order.
    """
    f = tl.program_id(0).to(tl.int64)
    blk = tl.program_id(1).to(tl.int64)
    cap = tl.cast(corpus_len, tl.int64)
    req = tl.load(docs_ptr + 2 * f).to(tl.int64)
    end = tl.load(total_len_ptr + req).to(tl.int64)
    # Keep the newest cap - 1 tokens of an over-long response.
    start = tl.maximum(tl.load(docs_ptr + 2 * f + 1).to(tl.int64), end - (cap - 1))
    n = tl.maximum(end - start, 0)
    # Block 0 also writes the separator.
    if (blk > 0) & (blk * BLOCK >= n):
        return
    head = tl.load(head_ptr)
    for i in range(f):
        req_i = tl.load(docs_ptr + 2 * i).to(tl.int64)
        end_i = tl.load(total_len_ptr + req_i).to(tl.int64)
        start_i = tl.maximum(
            tl.load(docs_ptr + 2 * i + 1).to(tl.int64), end_i - (cap - 1)
        )
        n_i = tl.maximum(end_i - start_i, 0)
        head = tl.where(head + n_i + 1 > cap, 0, head) + n_i + 1
    head = tl.where(head + n + 1 > cap, 0, head)
    doc_seq = tl.load(doc_seq_ptr) + f

    j = blk * BLOCK + tl.arange(0, BLOCK).to(tl.int64)
    m = j < n
    pos = head + j
    sep = head + n
    keep = m
    keep_sep = sep >= 0
    later = head + n + 1
    for g in range(f + 1, num_docs):
        req_g = tl.load(docs_ptr + 2 * g).to(tl.int64)
        end_g = tl.load(total_len_ptr + req_g).to(tl.int64)
        start_g = tl.maximum(
            tl.load(docs_ptr + 2 * g + 1).to(tl.int64), end_g - (cap - 1)
        )
        n_g = tl.maximum(end_g - start_g, 0)
        later = tl.where(later + n_g + 1 > cap, 0, later)
        # Document g and its separator occupy [later, later + n_g].
        keep = keep & ((pos < later) | (pos > later + n_g))
        keep_sep = keep_sep & ((sep < later) | (sep > later + n_g))
        later += n_g + 1
    if blk == 0:
        tl.store(corpus_ptr + sep, -1, mask=keep_sep)
        if f == num_docs - 1:
            tl.store(head_out_ptr, sep + 1)
            tl.store(doc_seq_out_ptr, doc_seq + 1)

    row = token_ids_ptr + req * token_ids_stride + start
    tl.store(corpus_ptr + pos, tl.load(row + j, mask=m), mask=keep)
    # Index each position j under its preceding 2 and 4 tokens, when they lie
    # in the document. Entries order by (document, position), so atomic_max
    # keeps the newest occurrence regardless of execution order.
    slot = doc_seq % SLOTS
    val = (doc_seq << 32) | pos
    hv = tl.zeros([BLOCK], tl.int64)
    for d in tl.static_range(4):
        t = tl.load(row + j - 1 - d, mask=m & (j - 1 - d >= 0), other=0)
        hv = (hv ^ t.to(tl.int64)) * HASH_MUL
        if d == 1:
            tl.atomic_max(
                table_ptr + (NUM_BUCKETS + _bucket(hv, NUM_BUCKETS)) * SLOTS + slot,
                val,
                mask=m & (j >= 2),
            )
        if d == 3:
            tl.atomic_max(
                table_ptr + _bucket(hv, NUM_BUCKETS) * SLOTS + slot,
                val,
                mask=m & (j >= 4),
            )

_suffix_local_scan_kernel(token_ids_ptr, token_ids_stride, idx_mapping_ptr, total_len_ptr, num_sampled_ptr, scratch_ptr, scratch_stride, Lp1, MAX_DEPTH, BLOCK)

Per (request, block): best (match_len, end_pos) packed into one int64.

end_pos is the exclusive end of an earlier occurrence of the request's suffix in its own history; the draft continues from end_pos. Longer matches win, then later positions.

Source code in vllm/v1/worker/gpu/spec_decode/suffix/speculator.py

@triton.jit
def _suffix_local_scan_kernel(
    token_ids_ptr,  # *int32  [max_num_reqs, token_ids_stride]
    token_ids_stride,
    idx_mapping_ptr,  # *int32  [B]  batch_idx -> req_state_idx
    total_len_ptr,  # *int32  [max_num_reqs]
    num_sampled_ptr,  # *int32  [B]
    scratch_ptr,  # *int64  [B, scratch_stride]  (output)
    scratch_stride,
    Lp1,  # int64 packing radix (> any position)
    MAX_DEPTH: tl.constexpr,
    BLOCK: tl.constexpr,
):
    """Per (request, block): best (match_len, end_pos) packed into one int64.

    end_pos is the exclusive end of an earlier occurrence of the request's
    suffix in its own history; the draft continues from end_pos. Longer
    matches win, then later positions.
    """
    b = tl.program_id(0).to(tl.int64)
    blk = tl.program_id(1).to(tl.int64)
    req = tl.load(idx_mapping_ptr + b).to(tl.int64)
    seq_len = tl.load(total_len_ptr + req).to(tl.int64)
    num_sampled = tl.load(num_sampled_ptr + b)
    row = token_ids_ptr + req * token_ids_stride

    scratch_off = b * scratch_stride + blk
    # Blocks past the last candidate (an occurrence must leave a token).
    if blk * BLOCK >= seq_len:
        tl.store(scratch_ptr + scratch_off, tl.zeros((), tl.int64))
        return

    pos = blk * BLOCK + tl.arange(0, BLOCK).to(tl.int64)
    active = (pos >= 1) & (pos < seq_len) & (num_sampled > 0)
    match_len = tl.zeros([BLOCK], dtype=tl.int64)
    for t in range(MAX_DEPTH):
        # Request token t back from its end (newest first).
        tail_pos = seq_len - 1 - t
        tail_tok = tl.load(row + tail_pos, mask=tail_pos >= 0, other=-2)
        src_pos = pos - 1 - t
        ok = active & (src_pos >= 0) & (tail_pos >= 0)
        tok = tl.load(row + src_pos, mask=ok, other=-1)
        active = ok & (tok == tail_tok)
        match_len += active.to(tl.int64)

    score = tl.where(match_len > 0, match_len * Lp1 + pos, 0)
    tl.store(scratch_ptr + scratch_off, tl.max(score, axis=0))

_suffix_observe_kernel(idx_mapping_ptr, num_sampled_ptr, num_rejected_ptr, stats_ptr, last_bucket_ptr, last_valid_ptr, K, K_PO2)

Fold the target's verdict on each request's last draft into the stats.

Integer atomics are order-independent, so every TP rank keeps identical stats (and therefore identical confidences).

Source code in vllm/v1/worker/gpu/spec_decode/suffix/speculator.py

@triton.jit
def _suffix_observe_kernel(
    idx_mapping_ptr,  # *int32  [B]
    num_sampled_ptr,  # *int32  [B]
    num_rejected_ptr,  # *int32  [B]
    stats_ptr,  # *int64  [2 * _NUM_LEN_BUCKETS, K, 2]
    last_bucket_ptr,  # *int32  [max_num_reqs]
    last_valid_ptr,  # *int32  [max_num_reqs]
    K: tl.constexpr,
    K_PO2: tl.constexpr,
):
    """Fold the target's verdict on each request's last draft into the stats.

    Integer atomics are order-independent, so every TP rank keeps identical
    stats (and therefore identical confidences).
    """
    b = tl.program_id(0).to(tl.int64)
    req = tl.load(idx_mapping_ptr + b).to(tl.int64)
    bucket = tl.load(last_bucket_ptr + req).to(tl.int64)
    num_sampled = tl.load(num_sampled_ptr + b).to(tl.int64)
    if (bucket < 0) | (num_sampled == 0):
        return
    num_accepted = num_sampled - 1
    num_verified = num_accepted + tl.load(num_rejected_ptr + b).to(tl.int64)
    num_valid = tl.load(last_valid_ptr + req).to(tl.int64)
    k_iota = tl.arange(0, K_PO2).to(tl.int64)
    # Accepted positions plus the first rejected one were observed.
    observed = (
        (k_iota < K)
        & (k_iota <= num_accepted)
        & (k_iota < num_verified)
        & (k_iota < num_valid)
    )
    stat = stats_ptr + (bucket * K + k_iota) * 2
    tl.atomic_add(stat, (k_iota < num_accepted).to(tl.int64), mask=observed)
    tl.atomic_add(stat + 1, tl.full([K_PO2], 1, tl.int64), mask=observed)
    tl.store(last_bucket_ptr + req, -1)