Conversation
`unified_attention` disables the 3D (split-KV) path whenever
`max_seqlen_q > 1`. The comment attributes this to prefill, but the real
constraint is scratch capacity: the kernel and `reduce_segments` index
`softmax_segm_*` per *query row* --
query_offset_0 = cur_batch_in_all_start_index + query_pos # token index
segm_output[query_offset_0, head, segm_idx, :]
-- while the builder allocates `seq_threshold_3D` rows, i.e. one per
*sequence*. The two coincide only at `query_len == 1`, which is why the
mismatch has been invisible.
Speculative decoding makes `max_seqlen_q = 1 + num_speculative_tokens`, so
it is swept onto the 2D path even though split-KV is exactly what a
long-KV, few-query-rows batch wants. Widening the gate alone is not safe:
with per-sequence scratch the kernel then writes past the end of all three
buffers. A guard-region probe on this tree, forcing the 3D path at
`query_len = 2`:
batch=32 total_q=64 rows=64 | out-of-bounds writes: 0
batch=33 total_q=66 rows=64 | out-of-bounds writes: 32768
batch=64 total_q=128 rows=64 | out-of-bounds writes: 1048576
Size the scratch per query row instead, and gate on that width rather than
on `max_seqlen_q > 1`, with an exact capacity check so an oversized batch
degrades to 2D instead of corrupting memory. True prefill still takes the
2D path: its query length far exceeds the decode width, and split-KV would
not help it.
The width is derived from the speculative config directly rather than via
`_init_reorder_batch_threshold`, so this does not change the backend's
`reorder_batch_threshold` and therefore does not alter batch reordering.
Without speculation `max_q_len_3D == 1`, every gating decision is
unchanged, and the scratch keeps its current size.
Signed-off-by: Mingfei Guo <1800012773@pku.edu.cn>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
Fixes #48076. The 3D split-KV path is disabled whenever
max_seqlen_q > 1, but the real constraint is thatsoftmax_segm_*is allocated per sequence and indexed per query row, so this sizes the scratch per query row and gates on that width instead — letting speculative decode use split-KV rather than falling back to 2D.Test Plan
Test Result
1584 passed with no new failures (the 290
use_tdfailures are pre-existing and reproduce identically onmain). A guard-region probe forcing the 3D path atquery_len=2shows the current per-sequence sizing writing out of bounds — 32768 floats past the end at batch 33, 1048576 at batch 64 — and zero once the scratch is sized per query row. Kernel time improves 2.6x–8.9x, scaling with KV length (8.9x at KV=32768, batch 1), and without speculationmax_q_len_3D == 1, so every gating decision and the scratch size are unchanged.