Skip to content

perf(qwen3.5): prefill the empty think block for single-turn models - #779

Open
manueldomke wants to merge 2 commits into
ROCm:mainfrom
manueldomke:perf-qwen35-prefill-think-block
Open

manueldomke wants to merge 2 commits into
ROCm:mainfrom
manueldomke:perf-qwen35-prefill-think-block

Conversation

@manueldomke

Copy link
Copy Markdown

Motivation

Without thinking, the Qwen3.5 chat template ends with an empty think block,
<think>\n\n</think>\n\n. Qwen3_5VL::insert() cuts those four tokens off
before the prefill, so the checkpoint lands where the next turn of a
conversation diverges. generate() then feeds them in with one forward()
each. Only after those four decode steps comes the draw for the answer's first
token.

A single-turn model never returns to that checkpoint, so the four decode steps
buy nothing. On Qwen3.5-2B they are ~118 ms of a ~600 ms one-token request.
The decision models from #776 and #777 are single-turn by nature.

Technical Details

  • The change. For single_turn && !enable_think, insert() prefills the
    block together with the prompt. The token _shared_insert() draws from the
    prefill logits is the answer's first token, and generate() takes
    last_token instead of running the forced block. Multi-turn models keep the
    old path, because there the checkpoint pays off. So does think: true.
  • Opt-in per catalog entry. The existing "single-turn" label now also
    sets AutoModel::single_turn (in _shared_initialize_model_state()). Until
    now only the hunyuan, Qwen3VL_Flash and gemma4e classes set the flag
    themselves, and their catalog entries already carry the label, so nothing
    changes for them. The flag is only ever switched on, never off. Models
    without the label run exactly the old path.
    The label keeps its existing
    effects: /v1/chat/completions rejects requests with assistant messages,
    and a context-length override is ignored.
  • Second commit: prefill again when a request matches the whole history.
    With the block prefilled, a single-turn request puts exactly its own tokens
    into token_history. If the same request comes again and nothing is left to
    prefill, _chunked_insert() returns an empty buffer and Sampler::sample()
    reads it, which is a segfault at address 0. This PR alone cannot hit it,
    because the decode loop still appends one more token. Together with
    #<max_tokens PR> it can. A complete match is now treated like a divergence.
    Stepping back by one token is not possible, because the linear attention
    layers carry recurrent state.

Numerics. The template's four tokens are exactly the ids generate() used
to feed: [248068, 271, 248069, 271], checked with the model's tokenizer. They
now go through the prefill kernels instead of four decode steps, so the result
is not bit-identical; see Test Result. Neither path is a reference, since the
prompt itself always runs through the prefill kernels.

With #777, the first answer token is now drawn in _shared_insert(), where
#777 does not capture. That is handled in #.

Test Plan

  • Latency. Same harness as #<max_tokens PR>: 20 independent one-token
    prompts, temperature: 0, three alternating rounds of main against this
    branch, with the label set in a copy of the catalog.
  • Answers. This branch against main + fix: don't count a discarded sample as generated text #776. main's own answers are
    distorted by the discarded-sample penalty that fix: don't count a discarded sample as generated text #776 removes.
  • Tokenization of the template tail.
  • Crash. The same request twice, with this branch and the max_tokens fix
    combined, under gdb. Then eight identical requests in a row with the guard.

Test Result

                      main      this PR
median wall           604 ms    486 ms
prefill_duration      430 ms    431 ms
decoding_duration     117 ms     30 ms   (what is left is the max_tokens off-by-one)

Answers.

Crash. With this PR and #<max_tokens PR> combined, the second identical
request segfaulted in Sampler::sample() called from _shared_insert(). Its
log shows no Prefill chunk line. With the guard, eight identical requests in
a row pass, answers and confidences are unchanged, and the extra prefill is
not measurable.

Submission Checklist

🤖 Generated with Claude Code

manueldomke and others added 2 commits October 12, 2026 14:56
Without thinking, the chat template ends with an empty think block.
Qwen3_5VL::insert() cut those four tokens off before the prefill, so the
checkpoint lands where the next turn of a conversation diverges, and
generate() then fed them in with one forward() each. Only after those four
decode steps came the draw for the answer's first token.

A single-turn model never returns to that checkpoint. For such a model the
block is now prefilled with the prompt, and the token _shared_insert() draws
from the prefill logits is the answer's first token. Multi-turn models keep
the old path, because there the checkpoint pays off. So does thinking, where
the model writes the block itself.

The catalog's "single-turn" label now also sets AutoModel::single_turn. Until
now only three classes set the flag themselves, all with the label, so
nothing changes for them. It is only ever switched on, never off.

Measured on Qwen3.5-2B (NPU2) with the label set, 20 independent one-token
prompts, three alternating rounds against main: median wall clock 604 -> 486
ms per request. Prefill time is unchanged (430 vs 431 ms); the four tokens
it gains cost nothing measurable.

The numbers are not bit-identical. The template's tokens are exactly the
four ids generate() used to feed ([248068, 271, 248069, 271]), but they now
go through the prefill kernels instead of four decode steps. Against main
with the discarded-sample fix (ROCm#776) 1 of 20 answers differs, and that was an
exact tie (A and C both at logprob -1.939). One of two 60-token generations
diverges after 90 characters. With logprobs applied on top, the log-odds
between options move by 0.25 nats median, 0.5 at most.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
_shared_insert() skips the prefix a request shares with token_history. If
the request matches the history completely, nothing is left to prefill,
_chunked_insert() returns an empty buffer, and Sampler::sample() reads it:
a segfault at address 0.

With the empty think block prefilled, a single-turn Qwen3.5 request puts
exactly its own tokens into token_history. The first generated token is
never added there. On its own this change cannot hit the bug: with
max_tokens=1 the decode loop still draws one more token and appends it.
Together with the max_tokens fix from the companion PR the loop no longer
runs, and the same request sent twice crashed the server on the second one.
Reproduced under gdb: the crashing request logs no "Prefill chunk" line.

A complete match is now treated like a divergence: clear and prefill
everything. Stepping back by one token is not an option, because the linear
attention layers carry recurrent state that cannot be rewound. With both
changes applied, eight identical requests in a row pass, and answers and
confidences are unchanged. The extra prefill is not measurable, because a
single-turn model cleared its context on these requests before anyway.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant