Follow-up to #114-B (PR #117). The batched-query SDPA (llm_full_seq_attention / _bf16) uses the shared-scores fast path only, so CudaHybridGdnForwardPass.AttentionBatched is gated to startPos + N <= 4096 (one query block per (head, query), scores held in 16 KB __shared__). Past that window the prefill falls back to the per-position KV-append + Attention loop (bit-identical, just one launch/token).
This means the original #114-B motivation — the trunk ballooning to ~16–20 s at ~25–34 K cumulative tokens, dominated by the per-token attention launches at long context — is only partly addressed: the constant launch-overhead base and the ≤4K regime are batched, but the long-context attention growth past 4096 still runs per-position. There is now a one-shot stderr log when this fallback engages.
Goal
A batched-query SDPA that works past 4096 without needing O(N × heads × maxSeqLen) score scratch (infeasible on a 12 GB card: N=512 × 32 heads × 32 K ctx ≈ 2 GB).
Approach (sketch)
- Wave-based scratch: launch the batched attention in waves of W queries (W chosen so
W × heads × seqLen × 4 B fits a bounded scratch, e.g. ≤512 MB → W≈100 at 25 K ctx), each wave one launch over (heads, W) blocks with per-(wave-query, head) global score slices. ~5 launches/attn-layer/chunk instead of N — still a large launch + GPU-utilisation win, and bit-identical (each block clones llm_attention's global-scratch path).
- Must stay bit-identical to the per-token
llm_attention (same strided score write, same 256-thread tree softmax, same V-weighted sum) — the MtpDecoder_GreedyParity / BatchedTrunk oracles must hold (see feedback_qkv_matvecdual_breaks_mtp_parity).
Tests to add
- Per-kernel bit-exactness above 4096 (e.g.
(5000, 64), (24000, 40)) vs the per-token global-scratch Attention path.
- A model-level oracle prefilling a >4096-token prompt and asserting bit-identical logits to the per-position fallback (today's
BatchedPrefill_* oracles cap ctx at 4096).
- A SnapKV-active prefill oracle with
BatchedAttnEnabled on (today SnapKV-active prefill also takes the per-position fallback — the batched path must either support Q-capture or stay correctly gated).
Quantify with bench-prefill.ps1 at long context (the 16 s→ regime), watching the trunk= field.
Follow-up to #114-B (PR #117). The batched-query SDPA (
llm_full_seq_attention/_bf16) uses the shared-scores fast path only, soCudaHybridGdnForwardPass.AttentionBatchedis gated tostartPos + N <= 4096(one query block per(head, query), scores held in 16 KB__shared__). Past that window the prefill falls back to the per-position KV-append +Attentionloop (bit-identical, just one launch/token).This means the original #114-B motivation — the trunk ballooning to ~16–20 s at ~25–34 K cumulative tokens, dominated by the per-token attention launches at long context — is only partly addressed: the constant launch-overhead base and the ≤4K regime are batched, but the long-context attention growth past 4096 still runs per-position. There is now a one-shot stderr log when this fallback engages.
Goal
A batched-query SDPA that works past 4096 without needing
O(N × heads × maxSeqLen)score scratch (infeasible on a 12 GB card: N=512 × 32 heads × 32 K ctx ≈ 2 GB).Approach (sketch)
W × heads × seqLen × 4 Bfits a bounded scratch, e.g. ≤512 MB → W≈100 at 25 K ctx), each wave one launch over(heads, W)blocks with per-(wave-query, head) global score slices. ~5 launches/attn-layer/chunk instead of N — still a large launch + GPU-utilisation win, and bit-identical (each block clonesllm_attention's global-scratch path).llm_attention(same strided score write, same 256-thread tree softmax, same V-weighted sum) — theMtpDecoder_GreedyParity/BatchedTrunkoracles must hold (seefeedback_qkv_matvecdual_breaks_mtp_parity).Tests to add
(5000, 64),(24000, 40)) vs the per-token global-scratchAttentionpath.BatchedPrefill_*oracles cap ctx at 4096).BatchedAttnEnabledon (today SnapKV-active prefill also takes the per-position fallback — the batched path must either support Q-capture or stay correctly gated).Quantify with
bench-prefill.ps1at long context (the 16 s→ regime), watching thetrunk=field.