Skip to content

[Engine] STRATA_PROMPT_ATTN_RDNA2=1 (opt-in, gfx103x, INT8 KV): prompt QSA attention without the decode kernel's split and merge, -3.5% prompt time at 32K and 119K - #3381

Open
Vop3p wants to merge 1 commit into
Niko1221:mainfrom
Vop3p:rdna2-prompt-attn

Conversation

@Vop3p

@Vop3p Vop3p commented Oct 20, 2026

Copy link
Copy Markdown
Contributor

Problem and result

On gfx103x the prompt's QSA attention has no kernel of its own. qsa_prompt_attn_batch returns false (the matrix-core kernels are RDNA4/CUDA), so prefill.cpp falls back to the decode path: attn_chunk_kernel_pre75 32 queries at a time, each query's ~2,051 selected cells split into 64-cell chunks over many blocks, then attn_merge_kernel. The split is there to fill the card with one decode query; a prompt chunk already has thousands of queries.

With STRATA_PROMPT_ATTN_RDNA2=1, one block per (query, KV head) walks all of the query's chunks with an online softmax: no per-chunk scratch and no merge pass. The score phase is attn_chunk_kernel_pre75's, unchanged. The value phase takes 4 dimensions per thread (one 4-byte load of int8 codes feeds 48 FMAs instead of 12), and each chunk is summed from zero and folded into the running sum with one fma, so the summation is blocked like the split path's.

Result on 2x RX 6900 XT (below): prompt time -3.49% at 31K and at 119K tokens, 10 of 10 interleaved cold-start pairs each; the kernel itself 1.36-1.68x faster in the synthetic parity test; error against FP64 no larger than the split path's, on synthetic and on real data. Without the variable the default path is byte-identical (checked, below).

Related discussion

None. It is opt-in and limited to gfx103x with INT8 K/V; nothing changes for other cards or without the variable.

Approach

  • qsa_prompt_attn_rdna2() (qsa_decode_attn.cu) runs only when the variable is 1, pre75_attn() holds (gfx103x, or STRATA_ATTN_PRE75=1 on another wave32 card), K/V are INT8 (not Q4_0 / FP16), head_dim 256 with 12 query heads per KV head, and n_q <= 65535; otherwise it returns false and prefill.cpp keeps its current path. STRATA_PROMPT_ATTN_OLD still wins.
  • HIP only: the kernel is compiled under #if defined(__HIPCC__) (inside the existing STRATA_ATTN_PRE75_BUILT block, which CUDA builds also define), the entry point is a stub returning false elsewhere, and prefill.cpp calls it only under STRATA_USE_HIP - so the CUDA, SYCL and Metal engines never reference it (Metal links root src/prefill/prefill.cpp against generated stubs; SYCL builds its own migrated copies).
  • Not bit-identical to the split path: the chunks fold into one running maximum instead of the merge's two passes (FP32-level differences, quantified below).
  • Test: hip_prompt_attn_rdna2 is qsa_prompt_attn_parity.cpp built with STRATA_PARITY_RDNA2 - the same checks the RDNA4 kernel has (hip_prompt_attn_wmma): against an FP64 host reference and against the decode kernel; it skips (77) off gfx103x.

Validation

Environment: 2x RX 6900 XT (gfx1030, 16 GB, PCIe 4.0 x8 each), Ryzen 5 5600X, 128 GB DDR4-3200, Ubuntu 26.04 (kernel 7.0.0), ROCm 10.0.0. Revision: main e88ee39 (v0.2.0-beta + 12 commits) and this branch on it. Built with setup.py's HIP options: -DSTRATA_ENABLE_HIP=ON -DSTRATA_ENABLE_CUDA=OFF -DSTRATA_PREFILL_MMQ=ON -DSTRATA_Q6K_EXPERTS=ON -DSTRATA_MMQ_KQUANTS=ON -DCMAKE_HIP_ARCHITECTURES=gfx1030, the pinned ggml, Release. Model: Qwen3.8-Flash-Next GSQ-RCO IQ3_S, layer split over both cards, --kv int8.

Passed:

  1. hip_prompt_attn_rdna2 32768 2048 3 (synthetic, GPU 1):

    case max abs error vs FP64: split / new new vs split (of output scale) ms per chunk: split -> new
    ctx 32768, 2048 queries 2.13e-6 / 2.00e-6 6.3e-7 19.30 -> 14.22 (1.36x)
    ctx 1500, 1500 queries (every cell) 1.78e-6 / 1.90e-6 5.1e-7 6.48 -> 3.85 (1.68x)
    ctx 2100, 256 queries (identity-to-sparse edge) 2.03e-6 / 1.54e-6 5.9e-7 2.43 -> 1.68 (1.45x)

    RMS error vs FP64 (an extra print, not in the test): new / split = 0.94, 0.93, 0.94. An earlier version that rescaled one running sum every chunk passed the test's bounds too but had 1.16-1.41x the split path's RMS error; the blocked fold costs ~1.6% kernel time.

  2. ctest (STRATA_BUILD_TESTS=ON): 138 tests, 132 passed (including hip_prompt_attn_rdna2, 3.8 s). The 6 failures are the same 6 on main e88ee39 built the same way (137 tests): ple_parity (needs a Q2_0 GGUF this machine does not have), hip_prefill_mmq_parity (SEGFAULT on main too, with PREFILL_MMQ=ON), native_hadamard_test, tq_cpu_parity, tq_cpu_parity_avx2 ("this ggml has no trellis types"), isa_guard_check (tq_rows_avx512 not listed).

  3. Default path byte-identical: deterministic config (parallel 1, no pipeline windows, --short-read 640 --adapt-every 100000, STRATA_PREFILL_CPU_SHARE=0), STRATA_LOGPOS with top-256 log-probabilities at every scored position: 8 chats of real text (first message 6K-105K tokens read through the prompt path, then ~580 tokens of continuation scored), 4,067 positions. main and this branch without the variable: cmp identical.

  4. Quality, same harness, against a local build of the same kernel in double (scores, softmax and sums in FP64; not part of this PR). Mean over the 8 chats, criteria fixed before the run (KL <= 1.5x main's, argmax agreement >= main's - 2 points):

    mean KL to the FP64-attention reference argmax agreement
    main 0.0334 94.81%
    this branch, =1 0.0325 95.88%

    Per chat the two differ in both directions; each has one chat well above its mean (main: the 14K chat, 0.066; this branch: the 65K chat, 0.131, also high in a second build). For that 65K chat a local diagnostic build re-ran the last 512 queries of every one of the 96 attention calls through the split path and the FP64 kernel on the real inputs: new/split RMS error 0.829-0.924 (median 0.883), the same distribution as the 14K control chat (median 0.893), no call above 1.2x. So the chat's higher KL is downstream amplification (QSA's top-k selection is discrete), not a larger attention error; I am reporting it rather than averaging it away.

  5. Speed: production-like config (parallel 2, --pipeline-windows 2 --batch-mtp, --kv-resident 32768, STRATA_PREFILL_CPU_SHARE=auto, STRATA_HC_PERSIST=1, STRATA_HIP_PROMPT_F16=1, STRATA_SH_STREAM=1, STRATA_PREFILL_STREAM_MIN=2048, a rocBLAS table via STRATA_ROCBLAS_TUNING), the same binary with and without the variable, 20 cold starts alternating A B A B, each reading two real-text prompts with max_tokens 1, prompt_ms from /metrics (0 tokens reused):

    prompt without (median, range) =1 (median, range) per pair median faster
    31,039 tokens 19,126 ms (19,112-19,151) 18,453 ms (18,448-18,481) -3.43% to -3.59% -3.49% 10/10
    119,290 tokens 62,639 ms (62,617-62,666) 60,456 ms (60,419-60,501) -3.44% to -3.55% -3.49% 10/10

    The =1 arms all printed the confirmation line, the others none; 0 stalls.

  6. CUDA compile only (no NVIDIA GPU here): CUDA 13.2.2 (nvcc 13.2.86) with GCC 15.2, -DSTRATA_ENABLE_CUDA=ON -DSTRATA_BUILD_TESTS=ON -DSTRATA_Q6K_EXPERTS=ON -DSTRATA_MMQ_KQUANTS=ON -DCMAKE_CUDA_ARCHITECTURES=86: the engine and every test target build, including qsa_prompt_attn_parity with the changed source; no warning from the changed files (the build's 261 warnings are all in files this PR does not touch). The HIP engine binary is byte-identical before and after adding the STRATA_USE_HIP guard (sha256 unchanged), so items 1-5 hold for the submitted commit.

Failed: none observed.

Not run:

  • CUDA execution (compiled only, item 6), other CUDA architectures than sm_86, and the SYCL and Metal builds (no toolkits here). SYCL compiles its own copies of the changed files; Metal compiles the changed prefill.cpp, where the new call is under STRATA_USE_HIP, so it needs no new stub - not verified by a Mac build.
  • Other gfx103x cards (RX 6800 / 6700, W6800, V620) and Windows: not tested.
  • gfx11 / gfx12: the dispatch keeps their path (pre75_attn() is false there unless STRATA_ATTN_PRE75=1); not executed.

Limits

  • Opt-in, gfx103x with INT8 K/V and the Flash-Next attention geometry (head_dim 256, 12 query heads per KV head) only; Q4_0 or FP16 K/V keep the split path.
  • Only the prompt path; decode is unchanged.
  • Quality is "not worse" by the measurements above, not "more accurate": FP32-level differences against the split path, which is itself not bit-identical to FP64.

Developed with an AI coding assistant; every number above was measured on the machine described here.

🤖 Generated with Claude Code

… a prompt's attention one block per (query, KV head)

On gfx103x the prompt's QSA attention has no kernel of its own: it runs the decode kernel (attn_chunk_kernel_pre75)
32 queries at a time, each query's selected cells split into 64-cell chunks over many blocks, then a merge pass.
That split is there to fill the card with one decode query; a prompt chunk already has thousands of queries.

With the variable set, one block per (query, KV head) walks all of the query's chunks with an online softmax (the
score phase is attn_chunk_kernel_pre75's, unchanged), so there is no per-chunk scratch and no merge, and the value
phase takes 4 dimensions per thread (one 4-byte load of int8 codes for 48 FMAs).  The kernel is built for HIP only;
elsewhere the entry point returns false and the caller keeps its path.  Without the variable nothing changes.

hip_prompt_attn_rdna2 (ctest, skips off gfx103x): qsa_prompt_attn_parity's checks - against an FP64 host reference
and against the decode kernel - for this kernel.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
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