Skip to content

[Engine] sycl: prompt attention over 8-query selection unions on XMX (opt-in STRATA_PROMPT_ATTN_UNION=1) - #2808

Open
hotschmoe wants to merge 1 commit into
Niko1221:mainfrom
hotschmoe:pr/sycl-prefill-attn-xmx
Open

hotschmoe wants to merge 1 commit into
Niko1221:mainfrom
hotschmoe:pr/sycl-prefill-attn-xmx

Conversation

@hotschmoe

@hotschmoe hotschmoe commented Oct 16, 2026 •

Copy link
Copy Markdown
Contributor

Problem and result

On the SYCL port the QSA prompt attention runs the FP32 decode kernel, 32 positions per launch, and every position gathers its own ~2,051 selected K/V cells. The XMX port (STRATA_PROMPT_ATTN_XMX=1) is slower than that, and docs/INTEL.md lists a kernel over the union of neighbouring positions' selections as "not built".

This adds that kernel, opt-in (STRATA_PROMPT_ATTN_UNION=1). For every 8 neighbouring positions a first kernel builds the sorted union of their selected cells, each cell marked with the positions that chose it. The attention kernel dequantizes the INT8 K/V of 64-cell chunks of that union once to FP16 and runs S^T = K Q^T and O += P V on joint_matrix (8x16x16, FP32 sums) for all 96 rows (8 positions x 12 heads). A cell is -inf in the online softmax for the positions that did not select it, so every position attends to exactly the cells it attends to with the FP32 kernel.

Result on 0.2.0-beta, two B70s: prompt reading 1,220 -> 1,414 tok/s at 4K (+16%) and 2,084 -> 2,387 tok/s at 32K (+15%). Decode is not touched (verify window 39.4 / 40.0 -> 39.4 / 39.2 ms). Unset, outputs and draft counts are identical to main.

Related discussion

Feature issue #2807 (closed in the 0.2.0 tracker reset). No maintainer agreement on scope yet. Opt-in, INT8 KV on Battlemage only; any other KV format or card, or a failed scratch allocation, keeps the existing path (one stderr line says why).

Approach

  • sycl/src/kernels/cuda/qsa_prompt_attn_union.dp.cpp (new): the union kernel (one work-group per 8 positions, 16K-cell local bitmaps), the attention kernel (one work-group per 8 positions and KV head, 12 sub-groups of 16, 256 GRF), the launcher and its per-queue scratch.
  • sycl/include/strata/kernels/qsa_prompt_attn_union.hpp (new): the entry point (the contract of qsa_prompt_attn_batch) and the switch.
  • sycl/src/kernels/cuda/qsa_prompt_attn.dp.cpp: the include and the guarded call before the XMX v2 call (4 lines).
  • sycl/src/kernels/qsa_prompt_attn_union_parity.cpp (new): the parity test (--selftest in ctest, skip code 77 off Battlemage).
  • sycl/CMakeLists.txt, docs/INTEL.md.

Left alone: the selection (block scores, top-k), decode, the FP32 and XMX v2 prompt kernels, other KV formats, every file outside sycl/ and docs/.

Validation

Setup: 0.2.0-beta main e88ee39 against this branch rebased on it. 2x Intel Arc Pro B70 32 GB, kernel 7.1, Compute Runtime 26.22.38646.4 at run time; built with the repository dev image (sycl/tools/Dockerfile on the current llama-sycl-b70:stable, IGC 2.41.5; an older cached base with IGC 2.36.3 cannot build 0.2.0, see #3219), AOT bmg-g31. Model Qwen3.8-Flash-Next ISTA GSQ-RCO IQ3_XXS, --layer-split 30 --prefill 8192 --max-context 65536 --kv int8 --spec 4 --mtp <BF16 draft layer> --mtp-max-t 4 --spec-min-p 0 --stream-experts --expert-cache auto, one slot, greedy, two fixed coding prompts with 512-token replies, a fresh engine per arm, then fresh 4K and 32K prompts. Coherence check and per-card health after every stop.

Passed:

  • qsa_prompt_attn_union_parity --selftest (0.2.0-beta + this branch, head 2c97fc0, one B70): 5 of 5 PASS, FAILURES: 0 (random INT8 pools behind a shuffled page table, 1.5K-40K contexts, ragged last groups, 2% non-resident pages, a bitwise rerun). Error against FP64 2.6e-3 - 5.0e-3 at an output scale of 3.3-3.6 (FP32 kernel 1.7e-6 - 2.6e-6); one chunk 0.88x (fresh selections, union 5.3x) to 5.45x (union 1.0x) of the FP32 kernel.
  • Engine A/B (decode tok/s and verify window ms per prompt; fresh-prompt read tok/s):
prompt 1 prompt 2 4K read 32K read
main e88ee39 70.6 (39.4 ms) 73.9 (40.0 ms) 1,219 2,078
this branch, switch unset 70.6 (39.4 ms) 73.9 (40.0 ms) 1,220 2,084
STRATA_PROMPT_ATTN_UNION=1 75.2 (39.4 ms) 77.3 (39.2 ms) 1,414 2,387
  • Switch unset: output SHAs and draft counts identical to main. Switch set: coherent; the text differs (FP16 products), so decode tok/s moves with draft acceptance while the window time is unchanged. Coherence true and per-card health passed after every stop.

Failed: none observed.

Not run: the teacher-forced NLL comparison on 0.2.0 (on 0.1.x main, 8 chats, 12,737 tokens: NLL 2.6234 unset vs 2.5506 set, inside the spread of a rounding-only control, so no quality change resolved either way); --max-context 200000 (not every expert fits on a plain two-card layer split); prompts above 32K on 0.2.0; CUDA/HIP/Metal (not touched); ctest as a whole.

Limits

  • Not bitwise with the switch set: Q, K, V and the probabilities enter the products as FP16.
  • It pays only when neighbouring selections overlap: with independent selections (union 5.3x one position's cells) it is no faster.
  • About 80 MB of device memory per card that reads a prompt, allocated at the first prompt, out of the --vram-reserve-mib headroom.
  • The docs/INTEL.md entry quotes the 0.1.x measurement (main 296b59d); I can swap in the 0.2.0 numbers.
  • Prepared with AI assistance (Claude Code): the code, the measurements and this text were produced by Claude on the submitter's two Arc Pro B70s, at the submitter's request.

🤖 Generated with Claude Code

…TRATA_PROMPT_ATTN_UNION=1)

qsa_prompt_attn_union: 8 neighbouring prompt positions share one pass over the union of their selected
cells. A first kernel builds the union (each cell marked with the positions that selected it); the
attention kernel dequantizes the INT8 K/V of 64-cell chunks once to FP16 and runs S^T = K Q^T and
O += P V on joint_matrix (8x16x16) for the group's 96 rows, with an online softmax that masks the cells
a row's position did not select. INT8 KV on Battlemage only; anything else keeps the FP32 attention.
Not bitwise (FP16 inputs to the products). Off by default: the hook in qsa_prompt_attn_batch reads the
switch first. qsa_prompt_attn_union_parity checks it against FP64 and the FP32 kernel and times both.

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

Copy link
Copy Markdown
Owner

Thanks for the pull request.

We're resetting the tracker for 0.2.0-beta. There are more than 1,300 open issues and pull requests, too many for us to review properly, so we're closing everything that isn't being worked on right now. Closed here does not mean rejected.

To reopen it takes about two minutes:

  1. Rebase on main if it no longer applies, and re-run your checks on 0.2.0-beta.
  2. Start the title with its area: [Engine], [GUI], [TUI], [Benchmark], [Docs] or [Other].
  3. Fill in the pull request template: the problem and result, and the validation (what you ran, on which revision, what you saw).
  4. Comment /reopen. A check runs and either reopens the pull request or lists what is still missing.

The rules are in CONTRIBUTING.md.

@Niko1221 Niko1221 closed this Oct 19, 2026
@hotschmoe hotschmoe changed the title sycl: prompt attention over 8-query selection unions on XMX (opt-in STRATA_PROMPT_ATTN_UNION=1) [Engine] sycl: prompt attention over 8-query selection unions on XMX (opt-in STRATA_PROMPT_ATTN_UNION=1) Oct 19, 2026
@hotschmoe

Copy link
Copy Markdown
Contributor Author

/reopen

Rebased on 0.2.0-beta (main e88ee39) and re-validated on two Arc Pro B70s; the description has the template and the new numbers. The rebased branch goes up right after the reopen (GitHub cannot reopen a pull request whose branch was force-pushed while it was closed).

@github-actions github-actions Bot reopened this Oct 19, 2026
@github-actions

Copy link
Copy Markdown

Reopened, thank you. A maintainer will pick it up from here. If it conflicts with main, rebase it first; that speeds up review.

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.

2 participants