#3381 [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

Descrição

## 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 e88ee39d (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 e88ee39d 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](https://claude.com/claude-code)