Repository navigation
[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
Conversation
… 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>
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
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.
Problem and result
On gfx103x the prompt's QSA attention has no kernel of its own.
qsa_prompt_attn_batchreturns false (the matrix-core kernels are RDNA4/CUDA), so prefill.cpp falls back to the decode path:attn_chunk_kernel_pre7532 queries at a time, each query's ~2,051 selected cells split into 64-cell chunks over many blocks, thenattn_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 isattn_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, orSTRATA_ATTN_PRE75=1on 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_OLDstill wins.#if defined(__HIPCC__)(inside the existingSTRATA_ATTN_PRE75_BUILTblock, which CUDA builds also define), the entry point is a stub returning false elsewhere, and prefill.cpp calls it only underSTRATA_USE_HIP- so the CUDA, SYCL and Metal engines never reference it (Metal links rootsrc/prefill/prefill.cppagainst generated stubs; SYCL builds its own migrated copies).hip_prompt_attn_rdna2isqsa_prompt_attn_parity.cppbuilt withSTRATA_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:
hip_prompt_attn_rdna2 32768 2048 3(synthetic, GPU 1):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.
ctest(STRATA_BUILD_TESTS=ON): 138 tests, 132 passed (includinghip_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_avx512not listed).Default path byte-identical: deterministic config (
parallel1, no pipeline windows,--short-read 640 --adapt-every 100000,STRATA_PREFILL_CPU_SHARE=0),STRATA_LOGPOSwith 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:cmpidentical.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):
=1Per 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.
Speed: production-like config (
parallel2,--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 viaSTRATA_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_msfrom /metrics (0 tokens reused):=1(median, range)The
=1arms all printed the confirmation line, the others none; 0 stalls.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, includingqsa_prompt_attn_paritywith 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 theSTRATA_USE_HIPguard (sha256 unchanged), so items 1-5 hold for the submitted commit.Failed: none observed.
Not run:
STRATA_USE_HIP, so it needs no new stub - not verified by a Mac build.pre75_attn()is false there unlessSTRATA_ATTN_PRE75=1); not executed.Limits
Developed with an AI coding assistant; every number above was measured on the machine described here.
🤖 Generated with Claude Code