Speculative decoding

Take one position. The draft model proposes token x with probability p(x). The target model wants probabilities q(x). The rule:

1. Accept x with probability min(1, q(x)/p(x)).
2. On a reject, draw from r(x) = max(q(x) − p(x), 0) / Z, with Z = Σ max(q − p, 0).

Proof. The chance to output x by acceptance is p(x) · min(1, q(x)/p(x)) = min(p(x), q(x)). The total chance of a reject is 1 − Σ min(p, q). Because Σ p = Σ q = 1, that is equal to Σ max(q − p, 0) = Z. So the residual adds Z · max(q(x) − p(x), 0) / Z, which is max(q(x) − p(x), 0). The sum of both parts is min(p, q) + max(q − p, 0) = q(x).

Worked example (panel 2 of the figure). Tokens A, B, C, D have p = [0.6, 0.3, 0.1, 0]. The target gives q = [0.4, 0.4, 0.1, 0.1].

Z = 0.2, so the residual is [0, 0.5, 0, 0.5]. If the draft is A and U = 0.8, then 0.667 > 0.8 is false: reject. The sampler draws B or D, each with chance 0.5 (about 0.5 in DeepInfer, because of the penalty in the next deep dive).

Many drafts. Each row applies the same rule, given that all drafts before it were accepted. So the accepted prefix plus the final token has the distribution of the target model, token by token. Take a general case: each draft passes with rate a, and there are γ drafts. Then a step gives (1 − a^(γ+1)) / (1 − a) tokens on average. For a = 0.8 and γ = 5, that is about 3.69.

Measured numbers. The deployment Q&A gives a "fixed-5 accepted length" of 3.25 to 3.32 (doc/inference_deployment_qa.md:249-256). The Q&A does not say if this number includes the token that the last row adds.

DeepInfer does not use max(q − p, 0) for the residual. It uses max(q − α·p, 0) with α = exp(−DRAFT_PROB_PENALTY) and DRAFT_PROB_PENALTY = 5e-5:

# sample
DRAFT_POS_HASH_MUL = 10000019
DRAFT_PROB_PENALTY = 5e-5  # 防止 reject sampling 采不出东西

Why. If q and p are equal up to rounding, max(q − p, 0) can be all zeros. Then a reject has nothing to sample. With α < 1, a token with q = p > 0 keeps a weight of at least (1 − α)·p. A test checks such a case: all residual weights are below 2^-31, and the kernel still returns a valid token (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:620-651).

The cost. The module docstring proves a bound. Translated from deepinfer/model_executor/layers/tilelang_ops/common.py:3-19:

• Let p be the draft distribution and q the target distribution. Sample from p. If the sample j has p[j] ≤ q[j], accept. Otherwise sample from the normalized weights r[i] = max(q[i] − α·p[i], 0).
• Let ε = 1 − α, and let A[j] be the real output probability of token j. Then q[j] − ε/(1 + ε) ≤ A[j] ≤ (1 + ε)·q[j].

With DRAFT_PROB_PENALTY = 5e-5, ε ≈ 5.0e-5. So no token gets more than 1.00005 times its target probability. No token loses more than about 0.00005 (an absolute amount, not a factor). The distribution test checks the upper bound: each count must stay below (1 + ε)·q times the number of samples, plus a statistical margin (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:566-573, deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:598-600).

The docstring shortens the accept step. The code accepts with the unchanged ratio min(1, q/p), also when p[j] > q[j] (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:238). Only the residual uses α.

        logits = logits.move()
        total_q = logits.shape[0]
        device = logits.device

        # Map rows to requests and steps, and mark draft rows for verification.
        ftp = decode_inputs.mtp_first_token_positions
        seqlen_q = decode_inputs.mtp_seqlens_q
        row_to_req, local_step, is_verify = prepare_verification_metadata(ftp, seqlen_q, total_q, self.max_MTP_value)
        speculative_inputs = driver_inputs.speculative

        # Apply structured-output masks before sampling and draft verification.
        if so_indices is not None:
            for step, indices in enumerate(so_indices):
                if len(indices):
                    apply_token_bitmask_by_step_inplace(logits, so_bitmasks[step], indices, ftp, step)

        (
            logprobs,
            token_index_cur,
            candidate_token_ids,
            accept_result_verify,
            draft_selected_logprobs,
            accept_probs,
        ) = process_sampling(
            logits,
            driver_inputs.kernel,
            self.force_topk,
            row_to_req=row_to_req,
            local_step=local_step,
            is_verify=is_verify,
            last_nextn_token_logprobs=mtp_info_pool.logprobs if speculative_inputs is not None else None,
            last_nextn_token_logprobs_index=mtp_info_pool.logprobs_index if speculative_inputs is not None else None,
            draft_token_ids_pool=mtp_info_pool.token_ids if mtp_info_pool is not None else None,
            mtp_info_read_rows=speculative_inputs.mtp_info_read_rows if speculative_inputs is not None else None,
        )
        del logits

        candidate_vocab_ids = candidate_token_ids
        if self.force_topk:
            candidate_vocab_ids = map_lookup(
                None,
                token_index_cur,
                candidate_token_ids.unsqueeze(-1),
                mode="by_index",
            )

        # Apply golden overrides and find each request's accepted prefix and final row.
        if speculative_inputs is not None:
            sample_results_cur, accept_result, is_final = get_accept_result(
                accept_result_verify,
                driver_inputs.golden_output,
                mtp_info_pool.token_ids,
                speculative_inputs.mtp_info_read_rows,
                candidate_vocab_ids,
                is_verify,
                ftp,
                row_to_req,
            )
        else:
            sample_results_cur = (
                torch.where(
                    driver_inputs.golden_output == -1,
                    candidate_vocab_ids,
                    driver_inputs.golden_output,
                )
                if driver_inputs.golden_output is not None
                else candidate_vocab_ids
            )
            is_final = None

1. Lines 423-425. logits is float32 [total_q, V]. Movable lets the function free the tensor early (line 458).
2. Lines 428-430. ftp (first token positions) and seqlen_q are int32 [req_cnt]. The result tensors have shape [total_q].
3. Lines 434-437. One structured-output mask per step. Step k applies to row ftp + k of each listed request.
4. Lines 446-457. process_sampling with the draft data of the pool: logprobs (and logprobs_index with force top-k), token_ids and the read rows. Without speculative_inputs (nextn = 0), the draft arguments are None and no row is verified.
5. Lines 460-467. With force top-k, candidate_token_ids are column numbers. map_lookup turns them into vocabulary ids.
6. Lines 470-480. get_accept_result gives the token per row, the accepted prefix and is_final.
7. Lines 481-491. Without drafts, every row is a fresh sample. Only the golden override applies.

The rest of the function (lines 493-595) looks up the token logprobs, computes the optional sum_smaller_probs and the top logprobs, and calls finalize_verification.

Request A has seqlen_q = 6 and accepts drafts 1 and 2. Request B has seqlen_q = 4 (3 drafts, for example after DSpark shortened it) and accepts all 3. nextn = 5.

ftp          = [0, 6]                        # int32 [2]
row          =  0  1  2  3  4  5 |  6  7  8  9
row_to_req   =  0  0  0  0  0  0 |  1  1  1  1
local_step   =  0  1  2  3  4  5 |  0  1  2  3
is_verify    =  T  T  T  T  T  F |  T  T  T  F
raw accept   =  T  T  F  T  F  T |  T  T  T  T   # kernel decision per row (fresh rows: T)
blocked      =  F  F  T  F  T  T |  F  F  F  T
alive        =  T  T  T  F  F  F |  T  T  T  T
eff. accept  =  T  T  F  F  F  F |  T  T  T  F
is_final     =  F  F  T  F  F  F |  F  F  F  T

last_valid_offset   = [2, 3]                 # int32 [2]
sampled_token_ids   = [r_A, bonus_B]         # residual sample of row 2; fresh sample of row 9
verified_token_ids  = [d1, d2, -1, -1, -1, -1, e1, e2, e3, -1]   # int32 [10]
num_generated_tokens = [2, 2, 1, 0, 0]       # accepted drafts at step 0..4
num_verified_tokens  = [2, 2, 2, 1, 1]       # verify rows at step 0..4

Row 3 of A has a raw accept, but row 2 before it is blocked, so row 3 does not count. A gets 3 tokens (d1, d2, r_A). B gets 4 tokens (e1, e2, e3, bonus_B). The counters come from a scatter-add by local_step; steps above nextn − 1 go into the last slot (deepinfer/model_executor/layers/tilelang_ops/sampler.py:347-357). _decode_post_func later puts the request count in front: [2, 2, 2, 1, 0, 0] (deepinfer/model_executor/layers/sampler.py:293-296).

1. The sampler returns a SamplerOutput with sampled_token (one per request), verified_tokens (one per row) and first_token_positions (CPU int32 [req_cnt]) (deepinfer/model_executor/layers/sampler.py:277-291).
2. LLMEngine._process_model_outputs passes the flat [total_q] array of accepted ids to C++ without a copy (deepinfer/engine/llm_engine.py:697-703).
3. Scheduler::process_outputs appends the accepted drafts and the sampled token, and checks the stop conditions after each one (Step 7 above).
4. build_request_outputs makes one RequestOutput per new token. The last one carries draft_tok = drafts + 1 and accept_tok = accepted + 1 (csrc/engine/outputs.hpp:214-216, csrc/engine/outputs.hpp:287-292).

For request A of the previous example, the user gets three objects. The third has draft_tok = 6 and accept_tok = 3. With ZOS, these objects arrive one step() call later (Zero-overhead scheduling (ZOS)).

Device copies for the next step. With ZOS, the next step is already scheduled before the CPU sees these results. So the sampler also keeps device copies: sampled_token_ids_gpu [num_seqs], predict_token_ids_gpu [num_seqs, 5] and num_accepted_tokens_gpu (= last_valid_offset) (deepinfer/model_executor/layers/sampler.py:258-291, deepinfer/sequence.py:259-269). A kernel uses them to fix seqlens_k of the next step on the device (deepinfer/worker/worker.py:220-227).

• On the CPU. _check_stop runs after each appended token. A stop at accepted draft 1 drops draft 2 and the sampled token (csrc/engine/scheduler.hpp:982-990, csrc/engine/scheduler.hpp:1035-1052).
• On the device. finalize_verification sets emitted_stop_token for each request whose emitted tokens contain one of its stop_token_ids (deepinfer/model_executor/layers/tilelang_ops/sampler.py:376-382). The sampler then sets the confidences of that request's next drafts to 0:

                if emitted_stop_token is not None and predict_tokens.confidence_rates is not None:
                    # 请求遇到了 eos，它下一轮 draft 的 token 的置信度都设为 0
                    predict_tokens.confidence_rates.masked_fill_(emitted_stop_token.unsqueeze(1), 0.0)

With DSpark, a confidence of 0 gives the request no drafts in the next step. This matters because the CPU schedules that step before it sees the stop.

A decode step writes cache entries for all 6 input tokens before the sampler knows which drafts survive. So every cache that keeps a fixed window must have room for these extra writes. DeepInfer calls this backtracking.

Why 134 is enough. Let the cache hold tokens up to position P − 1. A step writes positions P … P + 5 into ring slots pos % 134, before its attention reads the ring. These slots last held positions P − 134 … P − 129. The query at P is the oldest query of the step. With a 128-token window, it reads keys P − 127 … P (csrc/engine/input_metadata/prepare_tensor_input_v4.cpp:631-649). So the writes never remove a key that a query of this step or a later step needs.

Why the rejected keys do no harm. A rejected key at position y is also never read. A query reads only keys up to its own position. The first later step with a query at y also writes the real key of y before its attention. The docs give the same number: "128-token window + 6 tokens of speculative backtracking" (doc/inference_deployment_qa.md:191).

Nothing is copied back after a reject. The next step simply writes the same positions again. The tour Tour: One speculative decoding round shows this for every cache.