feat(rocm): f32-query decode GQA kernel with bf16 and fp8 KV cache support - #2168
Conversation
|
The split from #1936 is clear and the fp8 read-side dispatch is appropriately opt-in, but this head now conflicts with current main. Please rebase it and rerun the ROCm fp8 KV-cache/backend gates plus the branch record checks before merge. |
74510c5 to
a05d760
Compare
|
Done — rebased onto current Integration. The branch is now a linear chain on
Branch record checks. The ratchet activated by the integration found two regressions this PR owns, both repaired:
One gate the branch newly owed. ROCm gates, gfx1100 / RX 7900 XTX, ROCm 7.14.0, built and run at
Checks. |
7ce4a3e to
c24c2b7
Compare
|
Rebased onto current The retired Landing order: this PR (#2168) lands first. #2361 (fp8 prefill fast path) builds on this branch and shares the decode kernel history. It has been slimmed to only its 2 prefill-specific commits rebased onto this head, so the decode kernel history lands exactly once. Record gates at
ROCm gates, gfx1100 (RX 7900 XTX), ROCm 7.15.26333, built and run at
|
c24c2b7 to
934ff06
Compare
localai-org-maint-bot
left a comment
localai-org-maint-bot
left a comment
There was a problem hiding this comment.
Reviewed exact head 934ff06e278178ae056ac94fbef282e2ef25a937. Changes required before merge:
src/vt/rocm/rocm_paged_attn.hip: add an exact-geometry ROCm test that reaches the GQA4 FP8 kernel. Existing FP8 coverage useshq=2,d=16, and never enables this route, so deleting the dispatch remains green.- Gate the new C ABI and
vllm-clipath end to end, not only direct kernel calls. - Add coverage that
ROCM_ATTNaccepts the advertised FP8 cache configuration. - Update public C-ABI/FP8 capability documentation for the new exposed controls.
- Rebuild history so the committed spec precedes implementation.
- Reconcile the issue, spec, and PR body onto one owning row.
The frozen head lacks production-seam and mutation evidence for the behavior it adds.
934ff06 to
db3e689
Compare
Four gates for PR mudler#2168 review 5070827757: G6 (tests/vt/test_rocm_fp8_kv_cache.cpp): exact-geometry reach test for PagedAttnDecodeGqaF32Q with fp8 KV. Uses hq=16, kv=4, d=128, f32 q/out, fp8-e4m3 K/V — the exact guard geometry. G4/G4b (hq=2, H=1, D=16) never reaches the GQA4 dispatch, so deleting it left them green. A file-scope setenv initializer activates VT_ATTN_DECODE_GQA4=1 before any PagedAttentionKernelRocm call; G4/G4b's geometry does not match the guard, so they are not rerouted. C ABI v24 (tests/capi/test_capi.cpp): kv_cache_dtype defaults to NULL, and a valid string reaches FromModelDir (fails with MODEL_LOAD, not INVALID_ARGUMENT). The ABI version pin moves to >= 24. vllm-cli (tests/vllm/entrypoints/test_cli_kv_cache_dtype.cpp): --kv-cache-dtype is accepted and reaches the loader (the load fails on a missing checkpoint, not on an unrecognised argument). Registered unconditionally; exits 77 without VLLM_CPP_BUILD_EXAMPLES. ROCM_ATTN (tests/vllm/v1/attention/test_attn_backend_registry.cpp): supported_kv_cache_dtypes includes fp8 and fp8_e4m3, and supports_kv_cache_dtype binds the declaration to selection. e5m2 is NOT advertised. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-2 [OMP]
|
All six requirements from review 5070827757 are resolved. The branch was force-pushed ( 1. Exact-geometry ROCm test reaching GQA4 FP8 kernel
8/8 cases, 32 assertions, 0 skipped — G6 ran on the RX 7900 XTX. 2. C ABI + vllm-cli end-to-end gate
New 68/68 C ABI cases, 676 assertions; 2/2 CLI cases, 12 assertions. 3. ROCM_ATTN FP8 cache config acceptance test
21/21 cases, 135 assertions. 4. Public C-ABI/FP8 capability documentation
5. Spec before implementationCommit 6. One owning rowIssue #7, spec, and this PR body all carry Gates
|
|
Record anchor fix pushed ( What was wrongThe FixAdvanced the citation from CI status at
|
…sue #7) Scope, upstream chain, design, risks, and gates for widening PagedAttnDecodeGqaF32Q to accept fp8-e4m3 KV cache pages. The spec precedes the implementation per AGENTS.md. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-2 [OMP]
…pport Add PagedAttnDecodeGqaF32Q — a warp-strided decode attention kernel for the f32-query x bf16/fp8-KV x f32-output dtype combination that the GGUF dense path uses. Without this arm, f32-query decode falls through to PagedAttnOnline, which walks each context token with a per-key __syncthreads() reduction (~593us/call, 8 calls/token on Qwen3.5-4B). The kernel fuses QG=4 query heads per KV group, walks the sequence warp-strided, and does register online softmax with no per-key block sync. EPL = d/32 (8 at d=256, 4 at d=128). Dispatch is gated behind VT_ATTN_DECODE_GQA4=1 for the exact geometry (f32 q/out, bf16 or fp8 k/v, d=128 or 256, hq=16, kv=4). fp8 KV cache support: template PagedAttnDecodeGqaF32Q on TKV (default __hip_bfloat16) with k_scale/v_scale params. Add LoadRowEplFp8<EPL> (vectorized uint8_t loads + F8E4M3ToF32Dev dequant x scale) and LoadRowEplKv<EPL, TKV> (if constexpr dispatch between bf16 and fp8). Widen the dispatch guard: kBF16 || (kI8 && kv_cache_dtype != kAuto). The dequant arithmetic is identical to LoadKv(uint8_t*, ...) at line 176 — the fast kernel calls the same dequant the fallback already used. Benchmark: Qwen3.5-4B Q4_K_M, RX 7900 XTX, ROCm 7.14.0, 128-token greedy decode, 4 reps: ctx=256: 99.94 -> 140.19 tok/s (1.40x), bf16=143.0 (0.98x) ctx=1024: 56.28 -> 122.31 tok/s (2.17x), bf16=128.7 (0.95x) ctx=4096: 20.53 -> 79.53 tok/s (3.87x), bf16=92.7 (0.86x) ctx=8192: 11.08 -> 54.55 tok/s (4.93x), bf16=66.4 (0.82x) ctx=16384: 5.78 -> 33.47 tok/s (5.79x), bf16=43.16 (0.78x) Token-exact: 3/3 simple prompts match the PagedAttnOnline fallback output byte-for-byte. Divergence at thinking-tag near-ties is the reduction-order risk the bf16 GQA4 path already carries — same policy class as VT_ATTN_DECODE_D128. The arm ships opt-in (VT_ATTN_DECODE_GQA4=1); near-tie adjudication before any default-ON flip remains owed. Spec: .agents/specs/rocm-fp8-kv-decode-attn.md Evidence: docs/bench-evidence/gfx1100-tg200-t3a-20260823.md Fork issue: #7 FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-2 [OMP]
W6 store/read is on main. This branch's PagedAttnDecodeGqaF32Q arm reads fp8 KV. Without advertising the dtypes, --kv-cache-dtype fp8 is refused at the backend registry before the kernel can run. Fork issue: #7 FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-3-FLASH [OMP]
`check-env-doc` reads `std::getenv("VT_ATTN_DECODE_GQA4")` in
`src/vt/rocm/rocm_paged_attn.hip` and the knob was on neither surface:
not documented in `docs/ENVIRONMENT.md` and not on the kernel-internal
allowlist. It is a behavior-changing dispatch gate, not a tuning
switch, so the row documents what setting it routes where, the exact
geometry it guards, and the near-tie caveat the read site's own comment
carries.
FOLLOWING_AGENTS_PROTOCOL
Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:GLM-5-3-FLASH [OMP]
…ffer The EPL=16 path copied the bf16 LoadRowEplBf16 pattern of loading two uint4 values (32 bytes), but fp8 elements are 1 byte each so 16 elements fit in one uint4 (16 bytes). The second load wrote past r[15] into r[16..31] on a float r[16] array — a stack buffer overflow. The bug was latent: the dispatch guard only instantiates EPL=4 (d=128) and EPL=8 (d=256). EPL=16 would need d=512, which the guard excludes. The static_assert admits EPL=16, so the trap was set for a future d=512 extension. Fixed to a single uint4 load writing r[0..15], matching the spec design section. Verified on RX 7900 XTX (gfx1100) in a ROCm 10.0.0 container: test_rocm_fp8_kv_cache 7/7, test_ops_fp8_kv_cache 8/8 (511 assertions), test_attn_backend_registry 20/20, test_attn_validate_configuration 21/21, test_kv_cache_fp8_wiring 31/31 (487 assertions), test_ops_attention 11/11, test_rocm_backend 9/9 (1065 assertions), test_rocm_arch 9/9. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:glm-5-2 [OMP]
The fp8 KV cache capability landed in W6 but was unreachable from the CLI because the C ABI (vllm_model_params) did not carry a kv_cache_dtype field. The server had --kv-cache-dtype; the CLI did not. Add it as ABI v24: a const char* kv_cache_dtype field appended to vllm_model_params, mapped to EngineParams.kv_cache_dtype in vllm_engine_load, and exposed as --kv-cache-dtype in vllm-cli. NULL or "auto" (the zero-initialized default) is byte-identical to before. "fp8" or "fp8_e4m3" stores 1-byte fp8-e4m3 K/V. Served-model gate (Qwen3.5-4B Q4_K_M, RX 7900 XTX, ROCm 10.0.0): bf16 baseline vs fp8 KV + VT_ATTN_DECODE_GQA4=1, greedy, 256 tokens. First ~180 tokens byte-identical; divergence only at mid-stream near-tie moves (the reduction-order difference the spec documents). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:glm-5-2 [OMP]
Four gates for PR mudler#2168 review 5070827757: G6 (tests/vt/test_rocm_fp8_kv_cache.cpp): exact-geometry reach test for PagedAttnDecodeGqaF32Q with fp8 KV. Uses hq=16, kv=4, d=128, f32 q/out, fp8-e4m3 K/V — the exact guard geometry. G4/G4b (hq=2, H=1, D=16) never reaches the GQA4 dispatch, so deleting it left them green. A file-scope setenv initializer activates VT_ATTN_DECODE_GQA4=1 before any PagedAttentionKernelRocm call; G4/G4b's geometry does not match the guard, so they are not rerouted. C ABI v24 (tests/capi/test_capi.cpp): kv_cache_dtype defaults to NULL, and a valid string reaches FromModelDir (fails with MODEL_LOAD, not INVALID_ARGUMENT). The ABI version pin moves to >= 24. vllm-cli (tests/vllm/entrypoints/test_cli_kv_cache_dtype.cpp): --kv-cache-dtype is accepted and reaches the loader (the load fails on a missing checkpoint, not on an unrecognised argument). Registered unconditionally; exits 77 without VLLM_CPP_BUILD_EXAMPLES. ROCM_ATTN (tests/vllm/v1/attention/test_attn_backend_registry.cpp): supported_kv_cache_dtypes includes fp8 and fp8_e4m3, and supports_kv_cache_dtype binds the declaration to selection. e5m2 is NOT advertised. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-2 [OMP]
USAGE.md: replace the 'Not on the C ABI yet' section with the C ABI usage (vllm_model_params.kv_cache_dtype, ABI v24) and a vllm-cli example. The field defaults to NULL (auto), byte-identical to before. FEATURES.md: the fp8 KV cache row no longer says 'the C ABI are refused' — it now names the ABI v24 field. Add a C-ABI capability table row for KV-cache storage dtype. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-2 [OMP]
Row: GFX1100-TG200 (campaign, fork issue #5). The KV-FP8 engine-matrix row owns the store/read correctness surface; this row owns the decode performance arm on top of it. Issue #7, spec, and PR body now all reference the same owning row. ## Now updated to record the review-driven gates: G6 reach test, C ABI v24 gate, vllm-cli gate, ROCM_ATTN acceptance test, public docs, and the history rebuild that put the spec before implementation. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-2 [OMP]
…kend.h:587 The supported_kv_cache_dtypes insertion in backend.h (commit 0734605) shifted TritonMLABackend::get_kv_cache_shape from :580 to :587. The KERNEL-ATTN-MLA-SPARSE row still cited :580, which is now an empty line, making the anchor stale (29 > baseline 28). Advance the citation to :587 where the declaration actually lives. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-2 [OMP]
bb588ea to
e93c0a9
Compare
…he prefill fast path Row: `GFX1100-TG200` `PagedAttnPrefillSharedK` is templated on the KV element type, the query type and the output type, so an fp8 KV cache no longer falls off the fast prefill path. The fp8 tile load dequantizes each byte with `F8E4M3ToF32Dev(byte) * scale` into the same bf16 shared-memory tile the bf16 arm fills, leaving the QK dot product, the online softmax and the V accumulation untouched. This is the prefill half of the arm whose decode half landed in #2168. What this tree could check, and what it could not. No CI lane here compiles a .hip translation unit, so the kernel itself rests on the contributor's own gfx1100 evidence. Everything the lanes do build passed. The red lanes on this head are all the tree's, and each was checked rather than assumed: `sanitize-cpu (address,undefined)` reports eight failures, and all eight are the misaligned-borrow class, not this change. Every one of the job's eight `load of misaligned address` reports comes from `cpu_matmul_elem.cpp:577` (six) or `cpu_exl3_kernels.cpp:130` (two) -- the exact consumer sites #2581 repairs, in files this branch does not touch. Note the baseline for that lane grew from seven tests to eight while this pull request was in flight, because `main` gained `test_qwen3_dflash2_exl3`, which hits the same defect; that test passes on #2581's head. Counting the failures would have mis-attributed this. Reading their signatures does not. `windows-msvc-cpu` and `windows-msvc-vulkan` are red tree-wide. The add/add conflict this branch carried on `.agents/specs/rocm-fp8-kv-decode-attn.md` -- #2168 landed that spec while this branch still carried the pre-implementation copy of it -- was resolved by the contributor's own rebase before this merge. Its prefill spec, `.agents/specs/rocm-fp8-prefill-fastpath.md`, was never contested. Merged with --no-ff so the contributor's commits keep their authorship. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:GLM-5-2 [OMP] Assisted-by: claude-code:claude-opus-5 [Claude Code]
3 participants
Row:
GFX1100-TG200What changed
Add
PagedAttnDecodeGqaF32Q— a warp-strided decode attention kernel forthe f32-query x bf16/fp8-KV x f32-output dtype combination that the GGUF
dense path uses. Without this arm, f32-query decode falls through to
PagedAttnOnline, which walks each context token with a per-key__syncthreads()reduction (~593us/call, 8 calls/token on Qwen3.5-4B).The kernel fuses QG=4 query heads per KV group, walks the sequence
warp-strided, and does register online softmax with no per-key block
sync. EPL = d/32 (8 at d=256, 4 at d=128). Dispatch is gated behind
VT_ATTN_DECODE_GQA4=1for the exact geometry (f32 q/out, bf16 or fp8k/v, d=128 or 256, hq=16, kv=4).
fp8 KV cache support
--kv-cache-dtype fp8set the KV cache toDType::kI8, which failedthe
kBF16guard on every optimized decode kernel inrocm_paged_attn.hip. The dispatch fell through toPagedAttnOnline— the reference kernel with per-key
__syncthreadsreduction, no GQAamortization, no vectorized loads. Up to 7.5x decode regression at long
context on gfx1100.
Fix: template
PagedAttnDecodeGqaF32QonTKV(default__hip_bfloat16) withk_scale/v_scaleparams. AddLoadRowEplFp8<EPL>(vectorized uint8_t loads +F8E4M3ToF32Devdequant x scale) and
LoadRowEplKv<EPL, TKV>(if constexprdispatchbetween bf16 and fp8). Widen the dispatch guard:
kBF16 || (kI8 && kv_cache_dtype != kAuto).The dequant arithmetic is identical to the existing
LoadKv(uint8_t *, ...)at line 176 — the fast kernel now calls the same dequant thefallback already used.
Benchmark
Qwen3.5-4B Q4_K_M, RX 7900 XTX, ROCm 7.14.0, 128-token greedy decode,
4 reps:
Correctness
3/3 simple prompts produce token-exact identical output vs the
PagedAttnOnlinefallback. Divergence at thinking-tag near-ties is thereduction-order risk the bf16 GQA4 path already carries — same policy
class as
VT_ATTN_DECODE_D128. The arm ships opt-in(
VT_ATTN_DECODE_GQA4=1); near-tie adjudication before any default-ONflip remains owed.
Relationship with #2080 and #1936
This PR is the fp8 decode-attention split from #1936, per the reviewer's
request. #1936 carried two units of work: the GFX1100-TG200 decode
campaign and this fp8 KV decode-attention arm. This PR contains only the
fp8 arm.
This PR builds on the ROCm fp8-e4m3 KV cache store and read that landed
in #2080 (
191f64608). The store side is onmain; this PR adds thedecode-attention read side. The
08595f335fix (#2161) for the G1btest predicate is also in the base.
The TG200 campaign remains on #1936 and will be sent separately after
this lands.
How to verify
cmake -B /tmp/build -DVLLM_CPP_HIP=ON make -j4 vllm test_rocm_fp8_kv_cache test_attn_backend_registry test_capi vllm-cli test_cli_kv_cache_dtype ctest -R 'rocm_fp8|attn_backend|capi|cli_kv_cache'Verified on gfx1100 (RX 7900 XTX), ROCm 7.14.0:
test_rocm_fp8_kv_cache: 8/8 cases, 32 assertions (G6 exact-geometry reach test passes on GPU)test_attn_backend_registry: 21/21 cases, 135 assertions (ROCM_ATTN FP8 cache config acceptance)test_capi: 68/68 cases, 676 assertions (ABI v24 kv_cache_dtype default + acceptance)test_cli_kv_cache_dtype: 2/2 cases, 12 assertions (--kv-cache-dtype reaches the loader)test_rocm_backend: 9/9 cases, 1065 assertionstest_ops_quant_dot: 28/28 cases, 210138 assertionsReview-driven updates (review 5070827757)
test_rocm_fp8_kv_cache.cppG6 useshq=16, kv=4, d=128, f32 q/out, fp8 KV — the exact GQA4 guard geometry.
G4/G4b (hq=2, H=1, D=16) never reaches the dispatch; G6 does.
test_capi.cppv24kv_cache_dtypedefault-NULL and acceptance;test_cli_kv_cache_dtype.cpp--kv-cache-dtypereaches the loader.test_attn_backend_registry.cppasserts
supported_kv_cache_dtypesincludes fp8/fp8_e4m3 andsupports_kv_cache_dtypebinds the declaration to selection.USAGE.mdandFEATURES.mdupdated forABI v24
kv_cache_dtypefield.b8a859e1e(spec) precedes7bc79e87b(implementation).GFX1100-TG200.Out of scope
fp8_e5m2 compute, per-attention-head scales, the Metal arm, fast-path
(tensor-core/rocWMMA) fp8 attention kernels, the bf16 decode-opt
kernels, the prefill fp8 path, non-gfx1100 architectures.
Spec:
.agents/specs/rocm-fp8-kv-decode-attn.mdEvidence:
docs/bench-evidence/gfx1100-tg200-t3a-20260823.mdTracking issue: ghazni101#7
FOLLOWING_AGENTS_PROTOCOL
Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:GLM-5-2 [OMP]