diff --git a/.agents/claims/CLAIM-ROCM-DECODE-ATTN-D128.md b/.agents/claims/CLAIM-ROCM-DECODE-ATTN-D128.md index b962d0dc4..554d7b7bf 100644 --- a/.agents/claims/CLAIM-ROCM-DECODE-ATTN-D128.md +++ b/.agents/claims/CLAIM-ROCM-DECODE-ATTN-D128.md @@ -2,4 +2,4 @@ | Claim | Row IDs | Agent | Worktree / remote dir | Branch | Owned scope | State | Last update | |---|---|---|---|---|---|---|---| -| `CLAIM-ROCM-DECODE-ATTN-D128` | `BACKEND-ROCM` (`ACTIVE`) | Claude Code (sonnet-5), helper role | worktree `rdna3-kernel-porting-b9ec47`, real gfx1200 hardware (AMD Radeon RX 9060 XT, RDNA4, 32 CU), `$GPU_LOCK` respected | `row/ROCM-DECODE-ATTN-D128-SPEC` (this spec; the implementation follows on `row/ROCM-DECODE-ATTN-D128-IMPL`, stacked), base `main` `fafa16f0`; issue [#382](https://github.com/mudler/vllm.cpp/issues/382) (the ROCm half; the CUDA half landed as [PR #425](https://github.com/mudler/vllm.cpp/pull/425), `66399617`), motivated by [#488](https://github.com/mudler/vllm.cpp/issues/488). NOTE: #382 is filed against the cross-backend kernel row (state `ANCHOR-BACKFILL`), while this claim's Row ID is the `ACTIVE` backend row whose code it edits — `check-agent-record` requires an active claim to name a `SPIKE`/`ACTIVE` row, so the two deliberately differ | Owns ONLY: the `LoadRowEplBf16`/`StoreRowEplBf16` `EPL=4` case, the `VT_ATTN_DECODE_D128` gate (default OFF, same flag/default/reason as the merged CUDA arm), the `bf16_decode_opt`/`decode_gqa` gate extensions and the two `d==128` launch-dispatch branches in `src/vt/rocm/rocm_paged_attn.hip`; the new "Qwen3 geometry (bf16, GQA 2, head_dim 128)" case in `tests/vt/test_backend_cross_device.cpp` and its two flag-on ctest registrations in `tests/CMakeLists.txt`; `.agents/specs/rocm-decode-attn-d128.md` and this claim file. **NON-COLLISION:** disjoint from `CLAIM-ROCM-SKINNY-GEMM-GFX1200` (different files: `rocm_skinny_gemm.hip`/`rocm_matmul_hipblaslt.hip` vs `rocm_paged_attn.hip`), not stacked on any other branch. EXCLUDED: **the flip to default-ON on either backend** (owes the near-tie razor + distributional gate + golden regen, and per the spec §5 cross-arch reversal must be argued per backend — this is what keeps #382 open), rocWMMA for `d=128` (separate claim, separate spec, separate issue), `qg=4`/`qg=8` GQA fusion at any `d` (pre-existing, board-independent gap), any `d=128` prefill path, and the 8 pre-existing unrelated `ctest` failures (`vt: no kernel for op 63 on device type 5`) | `ACTIVE` | 2026-08-12 — **reconciled against the existing record before landing**, per the re-verify-before-claiming rule: #382 already named this exact defect and PR #425 had already merged the CUDA half, so this became a mirror of merged work rather than new design, and was re-gated from default-ON to **default OFF behind `VT_ATTN_DECODE_D128`** — the merged arm's own flag, default and stated reason (warp-strided online softmax reduces the KV sequence in a different order, so a greedy anchor can move at an exact bf16 tie; OFF keeps every golden byte-identical). gfx1200-verified: `ctest -R 'rocm\|cross_device'` **6/6** including two new flag-on registrations (verified non-vacuous: 1 case, 6 assertions, not zero); full `ctest` 385/393 with the 8 failures independently confirmed pre-existing. Gate exercised **both directions on one binary** — Qwen3-0.6B @1024 ctx TPOT 44.82/44.82 ms OFF vs 12.80/12.60 ms ON = **3.53x**; decode throughput +42.7% / +25.0% / +17.8% on 0.6B / 1.7B / 4B. **Carried finding:** #382 measured this same `EPL=4` arm **1.6x slower** on sm_110 where we measure it 3.5x faster — recorded, not reconciled; it is why the default-ON flip must be argued per backend. Rebased from `bbc482a2` onto `main` `fafa16f0` (167 commits), which required reformatting `Assisted-by` for the `check-commit-trailers` gate that landed in between, and de-linking §7's forward reference to the rocWMMA spec — that spec now lands on its own branch, so a markdown link to it fails `check-agent-record` as a dangling link. Spec content otherwise byte-identical. Re-gated on the new base, gfx1200: build 783/783, `ctest -R 'rocm\|cross_device'` 6/6, the new case non-vacuous under both flags (1 case, 6 assertions), full `ctest` with 8 pre-existing `kSharedExpertGate` (`OpId(63)`) failures owed to unmerged PR #509. `agent-preflight` fails 11, set-identical to a clean `fafa16f0` baseline. Spec PR open; implementation PR follows. | +| `CLAIM-ROCM-DECODE-ATTN-D128` | `BACKEND-ROCM` (`ACTIVE`) | Claude Code (sonnet-5), helper role | worktree `rdna3-kernel-porting-b9ec47`, real gfx1200 hardware (AMD Radeon RX 9060 XT, RDNA4, 32 CU), `$GPU_LOCK` respected | `row/ROCM-DECODE-ATTN-D128-IMPL` (the implementation; its spec landed from `row/ROCM-DECODE-ATTN-D128-SPEC` as [PR #564](https://github.com/mudler/vllm.cpp/pull/564), squashed to `373aa125`), rebased off the now-merged spec commits onto `main` `2784dd7b`; issue [#382](https://github.com/mudler/vllm.cpp/issues/382) (the ROCm half; the CUDA half landed as [PR #425](https://github.com/mudler/vllm.cpp/pull/425), `66399617`), motivated by [#488](https://github.com/mudler/vllm.cpp/issues/488). NOTE: #382 is filed against the cross-backend kernel row (state `ANCHOR-BACKFILL`), while this claim's Row ID is the `ACTIVE` backend row whose code it edits — `check-agent-record` requires an active claim to name a `SPIKE`/`ACTIVE` row, so the two deliberately differ | Owns ONLY: the `LoadRowEplBf16`/`StoreRowEplBf16` `EPL=4` case, the `VT_ATTN_DECODE_D128` gate (default OFF, same flag/default/reason as the merged CUDA arm), the `bf16_decode_opt`/`decode_gqa` gate extensions and the two `d==128` launch-dispatch branches in `src/vt/rocm/rocm_paged_attn.hip`; the new "Qwen3 geometry (bf16, GQA 2, head_dim 128)" case in `tests/vt/test_backend_cross_device.cpp` and its `VT_ATTN_DECODE_D128` flag-on ctest registration in `tests/CMakeLists.txt` (the second, `VT_ATTN_DECODE_WMMA`, moved to the rocWMMA branch with the arm it gates); `.agents/specs/rocm-decode-attn-d128.md` and this claim file. **NON-COLLISION:** disjoint from `CLAIM-ROCM-SKINNY-GEMM-GFX1200` (different files: `rocm_skinny_gemm.hip`/`rocm_matmul_hipblaslt.hip` vs `rocm_paged_attn.hip`), not stacked on any other branch. EXCLUDED: **the flip to default-ON on either backend** (owes the near-tie razor + distributional gate + golden regen, and per the spec §5 cross-arch reversal must be argued per backend — this is what keeps #382 open), rocWMMA for `d=128` (separate claim, separate spec, separate issue), `qg=4`/`qg=8` GQA fusion at any `d` (pre-existing, board-independent gap), any `d=128` prefill path, and the pre-existing unrelated `ctest` failures (`vt: no kernel for op SharedExpertGate` on ROCm, plus a missing `shellcheck`, an mmap-RSS assertion and a JSON type error) | `ACTIVE` | 2026-08-12 — **reconciled against the existing record before landing**, per the re-verify-before-claiming rule: #382 already named this exact defect and PR #425 had already merged the CUDA half, so this became a mirror of merged work rather than new design, and was re-gated from default-ON to **default OFF behind `VT_ATTN_DECODE_D128`** — the merged arm's own flag, default and stated reason (warp-strided online softmax reduces the KV sequence in a different order, so a greedy anchor can move at an exact bf16 tie; OFF keeps every golden byte-identical). gfx1200-verified: `ctest -R 'rocm\|cross_device'` **6/6** including two new flag-on registrations (verified non-vacuous: 1 case, 6 assertions, not zero); full `ctest` 385/393 with the 8 failures independently confirmed pre-existing. Gate exercised **both directions on one binary** — Qwen3-0.6B @1024 ctx TPOT 44.82/44.82 ms OFF vs 12.80/12.60 ms ON = **3.53x**; decode throughput +42.7% / +25.0% / +17.8% on 0.6B / 1.7B / 4B. **Carried finding:** #382 measured this same `EPL=4` arm **1.6x slower** on sm_110 where we measure it 3.5x faster — recorded, not reconciled; it is why the default-ON flip must be argued per backend. Rebased from `bbc482a2` onto `main` `fafa16f0` (167 commits), which required reformatting `Assisted-by` for the `check-commit-trailers` gate that landed in between, and de-linking §7's forward reference to the rocWMMA spec — that spec now lands on its own branch, so a markdown link to it fails `check-agent-record` as a dangling link. Spec content otherwise byte-identical. Re-gated on the new base, gfx1200: build 783/783, `ctest -R 'rocm\|cross_device'` 6/6, the new case non-vacuous under both flags (1 case, 6 assertions), full `ctest` with 8 pre-existing `kSharedExpertGate` (`OpId(63)`) failures owed to unmerged PR #509. `agent-preflight` fails 11, set-identical to a clean `fafa16f0` baseline. **2026-08-14 — spec LANDED as PR #564 (`373aa125`); this claim now tracks the implementation.** Rebased off the two now-squashed spec commits onto `main` `2784dd7b`; the commit is source-only (3 files) and carries no forward reference to the rocWMMA flag, so it stands alone. Re-gated on that base, gfx1200, `$GPU_LOCK` held: build 1220/1220; `ctest -R 'rocm\|cross_device'` **5/5** (5 not 6 — the `VT_ATTN_DECODE_WMMA` registration left with its arm); flag A/B on ONE binary re-measured **3.47x** (Qwen3-0.6B @1024 ctx, 45.15/45.13 ms OFF vs 12.97/13.04 ms ON), holding the 3.53x from the old base across 76 commits of drift. Full `ctest` 448/455 with **7** failures, and those 7 are now PROVEN pre-existing rather than argued: a clean `main` `2784dd7b` worktree, built from source with none of this code, fails the identical set (only `test_op_parity`'s index shifts 403→404, from the added registration). `agent-preflight` fails 9, a strict SUBSET of that same baseline's 10 (differing only by `role-undeclared`). Note `origin/main` (the `joral` fork) is 75 commits behind `upstream/main`, so preflight's range gates grade 76 commits of which 75 are other people's — `check-commit-trailers` and `check-doc-checkpoint` both pass against `upstream/main`, the base the spec actually merged to. **Fresh evidence 2026-08-14, SUPERSEDING the `+42.7% / +25.0% / +17.8%` figures above** — those came from a 128-token-context stash-based A/B; every number here is 1024-token synthetic prompt, 128 generated, greedy, seed 0, one binary, `$GPU_LOCK` held, 2 reps per cell agreeing within ~1%. Four-model TPOT OFF→ON: Qwen3-0.6B 42.53→11.78 ms (**3.61x**, `qg=2` fused), Qwen3-1.7B 52.85→21.93 ms (**2.41x**, `qg=2` fused), Qwen3-4B 81.89→39.22 ms (**2.09x**, `qg=4` per-head — no GQA fusion at any `d`, so this isolates the `EPL` widening from the fusion), Qwen3.5-0.8B 23.76→23.55 ms (**1.01x**). The last is the **NEGATIVE CONTROL** and it earned its keep: its `head_dim` is 256, so the `d == 128` gate provably cannot reach it, yet its first OFF rep landed a 33% outlier at 31.14 ms — a blind 2-rep average would have reported a bogus ~1.2x "win" for a model the flag cannot affect. Re-run 3x it gives 23.86/23.75/23.68 against ON's 23.52/23.57. End-to-end output throughput rises less than TPOT on the same runs (0.6B 2.48x, 1.7B 2.05x, 4B 2.02x) because these carry a 1024-token prefill the flag does not touch; TPOT isolates decode, throughput dilutes it. **Qwen3-1.7B concurrency sweep** (`--num-prompts` = 2x concurrency), throughput tok/s OFF→ON (ratio): c1 12.89→24.66 (1.91x), c2 23.27→47.45 (2.04x), c4 39.10→86.86 (2.22x), c8 58.97→147.35 (**2.50x**), c16 78.43→227.08 (**2.90x**); TPOT ratio over the same points 2.40x→3.18x. **The advantage GROWS with concurrency rather than compressing** — the opposite of the prediction made before the run, which reasoned that a tiny c1 grid flatters the fast kernel. The dominant effect is the reverse: from c8 to c16 the fallback scales only 1.33x against the arm's 1.54x, and scaling efficiency at c16 relative to perfect-linear-from-c1 is **38% OFF against 58% ON**. `PagedAttnOnline` is therefore the batch-scaling bottleneck, not merely slow per call, and the win is largest in the regime a server actually runs in. The c1 row reproduces the independent four-model sweep to within ~1% (52.85/21.93 vs 53.40/22.26), a cross-check on run-to-run stability. **Caveats:** single board; `--input-len` builds synthetic tokens, so all of the above is a decode-path A/B and not a serving benchmark. Implementation PR not yet opened. | diff --git a/.agents/specs/rocm-decode-attn-d128.md b/.agents/specs/rocm-decode-attn-d128.md index 8b964aa57..14d2bd5ac 100644 --- a/.agents/specs/rocm-decode-attn-d128.md +++ b/.agents/specs/rocm-decode-attn-d128.md @@ -99,7 +99,9 @@ synchronization. This is the same observation #382 made about the CUDA file. `decode_opt`/`decode_gqa` flags — the same env var, default and rationale as the merged CUDA arm. 3. `bf16_decode_opt` gate: `d == 256 || d == 512` → `(d == 128 && (decode_d128 - || decode_wmma)) || d == 256 || d == 512`. The `decode_wmma` disjunct is + || decode_wmma)) || d == 256 || d == 512`. **As landed the gate omits the + `decode_wmma` disjunct**, because that flag does not exist — see the forward + reference below. The `decode_wmma` disjunct is deliberate: the rocWMMA arm (separate spec) is a second, independently opt-in kernel for the same head size, and without it a bare `VT_ATTN_DECODE_WMMA=1` would be a silent no-op. **Forward reference:** @@ -142,10 +144,19 @@ bf16 correctness coverage in this suite. `tests/CMakeLists.txt`: because the arm ships OFF **and** its flag is read into a `static const bool` — once per process — the default registration only ever -gates the `PagedAttnOnline` fallback. Two extra ctest registrations re-run the -same binary filtered to this case with `VT_ATTN_DECODE_D128=1` and -`VT_ATTN_DECODE_WMMA=1`, so the arms that actually run the new kernels are -gated. Same shape as the existing `test_dense_gateup_fused_marlin_off_*` pair. +gates the `PagedAttnOnline` fallback. **As landed there is ONE extra ctest +registration**, `VT_ATTN_DECODE_D128=1`; the planned second, +`VT_ATTN_DECODE_WMMA=1`, does not exist because that flag does not. Same shape +as the existing `test_dense_gateup_fused_marlin_off_*` pair. + +That registration does NOT by itself prove the new kernel ran. It re-runs the +same case with the env set, and the case's only backend assertion is +`declines == 0`, which `OpProviderStats` reports at PROVIDER granularity — +identical with the flag set and unset. On any non-ROCm machine the case runs +1 test case and **0 assertions** and exits 0, so the registration is green on +nothing everywhere this project has hardware. Closing §9's stop condition 2 +needs a kernel-selection counter in `rocm_paged_attn.hip` asserted to DIFFER +between the two registrations. Verified non-vacuous (the trap `SKIP_RETURN_CODE 77` exists for, issue #463): the filter resolves to `test cases: 1 | 1 passed`, `assertions: 6 | 6 passed`, not zero. @@ -191,8 +202,10 @@ Any future flip to default-ON must be argued per backend with per-backend measurement; the fact that the ROCm arm is a large win is not evidence for the CUDA arm, and #382's sm_110 regression is not evidence against this one. -**Correctness**, gfx1200, real hardware: `ctest -R 'rocm|cross_device'` **6/6 -pass**, including both new flag-on registrations. Full `ctest` 393 tests, +**Correctness**, gfx1200, real hardware: `ctest -R 'rocm|cross_device'` **5/5 +pass** as landed, including the one new flag-on registration. (An earlier draft +of this section said 6/6 "including both new flag-on registrations", from the +two-registration plan above that did not land.) Full `ctest` 393 tests, 385 passed / 8 failed; all 8 reproduce identically (same tests, same root cause `vt: no kernel for op 63 on device type 5`, an unrelated pre-existing ROCm op-registration gap) on an isolated build of this branch **without** this @@ -301,12 +314,22 @@ flock "$HOME/gpu.lock" -c ' ## Result on the implementation branch (2026-08-12) -> **Not landed.** This section records what was built and measured on the -> unmerged implementation branch. No `VT_ATTN_DECODE_D128` exists in -> `src/vt/rocm/` on `main` — `git log -S'VT_ATTN_DECODE_D128' -- src/vt/rocm/` -> is empty, and `rocm_paged_attn.hip` still gates on `d == 256 || d == 512`. -> This spec is committed BEFORE its implementation, per AGENTS.md; the section -> becomes `## Outcome` when the code merges and the row reaches `DONE`. +> **Landed by PR #767**, which carries this correction. The banner this +> paragraph replaced said "Not landed" and offered +> `git log -S'VT_ATTN_DECODE_D128' -- src/vt/rocm/` as proof — a command that +> returns the opposite once the code is in, which is how a record starts +> disagreeing with the tree. +> +> The section stays `## Result` rather than becoming `## Outcome`: `BACKEND-ROCM` +> remains `ACTIVE`, and `## Outcome` is scoped to a row reaching `DONE`. The +> arm ships **default OFF**, so nothing here is a shipped-behaviour claim. +> +> **Still owed, and NOT discharged by this landing:** the flag-ON arm has no +> proof it REACHES the new kernel. `OpProviderStats` counts at provider +> granularity, so `declines == 0` is identical with the flag set and unset, and +> the ctest registration runs 0 assertions on every non-ROCm machine. §9's stop +> condition 2 — "stop if the flag-ON arm cannot be shown to reach the new +> kernel; confirm selection counts, not just tokens" — is therefore still open. **Built the ROCm `d=128` decode arm, default OFF, mirroring the merged CUDA arm of the same issue.** Root cause for the ROCm decode-attention gap #488 diff --git a/src/vt/rocm/rocm_paged_attn.hip b/src/vt/rocm/rocm_paged_attn.hip index f73993197..e0e658bfe 100644 --- a/src/vt/rocm/rocm_paged_attn.hip +++ b/src/vt/rocm/rocm_paged_attn.hip @@ -262,32 +262,48 @@ __global__ void PagedAttnOnline(TO* out, const TQ* query, const TKV* k_cache, constexpr int kDecWarps = 8; constexpr int kDecEpl = 8; // 32 lanes * 8 = 256 head_dim -// EPL elems/lane: 8 → d=256, 16 → d=512 (32 lanes). +// EPL elems/lane: 4 → d=128 (Qwen3/Llama-class GQA), 8 → d=256, 16 → d=512 +// (32 lanes). d=128 (#488, ROCM-DECODE-ATTN-D128): the kernel body below is +// already generic in EPL; only the 4/8/16-element vectorized load/store and +// the host dispatch needed a new case, not a new kernel. template __device__ inline void LoadRowEplBf16(const __hip_bfloat16* p, int64_t base, int lane, float r[EPL]) { - static_assert(EPL == 8 || EPL == 16, "EPL"); - const uint4* u = reinterpret_cast(p + base); - if constexpr (EPL == 8) { - const uint4 w = u[lane]; + static_assert(EPL == 4 || EPL == 8 || EPL == 16, "EPL"); + if constexpr (EPL == 4) { + const uint2* u = reinterpret_cast(p + base); + const uint2 w = u[lane]; const __hip_bfloat16* h = reinterpret_cast(&w); #pragma unroll - for (int i = 0; i < 8; ++i) r[i] = __bfloat162float(h[i]); + for (int i = 0; i < 4; ++i) r[i] = __bfloat162float(h[i]); } else { - const uint4 w0 = u[lane]; - const uint4 w1 = u[lane + 32]; - const __hip_bfloat16* h0 = reinterpret_cast(&w0); - const __hip_bfloat16* h1 = reinterpret_cast(&w1); + const uint4* u = reinterpret_cast(p + base); + if constexpr (EPL == 8) { + const uint4 w = u[lane]; + const __hip_bfloat16* h = reinterpret_cast(&w); #pragma unroll - for (int i = 0; i < 8; ++i) r[i] = __bfloat162float(h0[i]); + for (int i = 0; i < 8; ++i) r[i] = __bfloat162float(h[i]); + } else { + const uint4 w0 = u[lane]; + const uint4 w1 = u[lane + 32]; + const __hip_bfloat16* h0 = reinterpret_cast(&w0); + const __hip_bfloat16* h1 = reinterpret_cast(&w1); +#pragma unroll + for (int i = 0; i < 8; ++i) r[i] = __bfloat162float(h0[i]); #pragma unroll - for (int i = 0; i < 8; ++i) r[i + 8] = __bfloat162float(h1[i]); + for (int i = 0; i < 8; ++i) r[i + 8] = __bfloat162float(h1[i]); + } } } template __device__ inline void StoreRowEplBf16(__hip_bfloat16* p, int64_t base, int lane, const float r[EPL]) { - static_assert(EPL == 8 || EPL == 16, "EPL"); - if constexpr (EPL == 8) { + static_assert(EPL == 4 || EPL == 8 || EPL == 16, "EPL"); + if constexpr (EPL == 4) { + __hip_bfloat16 tmp[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) tmp[i] = __float2bfloat16(r[i]); + reinterpret_cast(p + base)[lane] = *reinterpret_cast(tmp); + } else if constexpr (EPL == 8) { __hip_bfloat16 tmp[8]; #pragma unroll for (int i = 0; i < 8; ++i) tmp[i] = __float2bfloat16(r[i]); @@ -1638,7 +1654,6 @@ __global__ void PagedAttnPrefillSharedKWmma(__hip_bfloat16* out, const __hip_bfl #endif // VT_ROCWMMA_OK } - } // namespace void PagedAttentionKernelRocm(Queue& q, Tensor& out, const Tensor& query, const Tensor& k_cache, @@ -1681,8 +1696,26 @@ void PagedAttentionKernelRocm(Queue& q, Tensor& out, const Tensor& query, const return true; // default ON when decode_opt }(); + // head_dim-128 arm (#382) -- the Qwen3-dense / Llama / Mistral head size. + // Same kernels as d==256/512, EPL=4 instead of 8/16; before this, bf16 decode + // at d==128 fell all the way to the generic PagedAttnOnline (#488 measured + // that fallback at 41.1us/call against vLLM's 5.10us on gfx1200). + // + // DEFAULT OFF, opt in with VT_ATTN_DECODE_D128=1 -- the SAME flag, default and + // reason as the merged CUDA arm (cuda_paged_attn.cu, DecodeD128Enabled, #382 / + // PR #425). It is correctness-complete but NOT byte-exact against the kernel it + // replaces: warp-strided online softmax reduces the KV sequence in a different + // ORDER than PagedAttnOnline's per-tile loop, so a greedy anchor can move at an + // exact bf16 tie. Shipping OFF keeps every existing golden byte-identical. The + // flip owes the near-tie razor + distributional gate + regen under the ratified- + // tie rule, on both backends, and is a separate change. + static const bool decode_d128 = [] { + const char* e = std::getenv("VT_ATTN_DECODE_D128"); + return e != nullptr && e[0] == '1'; + }(); const bool bf16_decode_opt = - decode_opt && (d == 256 || d == 512) && query.dtype == DType::kBF16 && + decode_opt && ((d == 128 && decode_d128) || d == 256 || d == 512) && + query.dtype == DType::kBF16 && k_cache.dtype == DType::kBF16 && v_cache.dtype == DType::kBF16 && out.dtype == DType::kBF16 && kc_hd % 8 == 0; @@ -1690,7 +1723,8 @@ void PagedAttentionKernelRocm(Queue& q, Tensor& out, const Tensor& query, const const int64_t qg = hq / num_kv_heads; // GQA fuse: sliding QG=2. Global QG=8: tile as DecodeGqa QG=2 × z=4 // (proven kernel; halves K/V reloads vs DecodeOpt's 8×). QG=4/8 CTAs lost A/B. - if (decode_gqa && ((qg == 2 && (d == 256 || d == 512)) || (qg == 8 && d == 512))) { + if (decode_gqa && + ((qg == 2 && (d == 128 || d == 256 || d == 512)) || (qg == 8 && d == 512))) { // Prefill SharedK (scoreless). Default ON for long single-req prefill (lab win // ~1.09× @11k vs DecodeGqa-z4). Set VT_ATTN_PREFILL_FLASH_SHAREDK=0 to disable. // SHAREDK_WMMA=1: rocWMMA QK + online V (lab; keep only if beats scalar SharedK). @@ -1871,7 +1905,14 @@ void PagedAttentionKernelRocm(Queue& q, Tensor& out, const Tensor& query, const const unsigned z_groups = (qg == 8 && d == 512) ? 4u : 1u; dim3 grid(static_cast(total_q), static_cast(num_kv_heads), z_groups); dim3 block(static_cast(nwarps * 32)); - if (d == 256) { + if (d == 128) { + PagedAttnDecodeGqaBf16<2, 4, 8><<>>( + out.Ptr<__hip_bfloat16>(), query.Ptr<__hip_bfloat16>(), k_cache.Ptr<__hip_bfloat16>(), + v_cache.Ptr<__hip_bfloat16>(), block_table.Ptr(), seq_lens.Ptr(), + query_start_loc.Ptr(), num_reqs, hq, num_kv_heads, d, block_size, bt_row, + bt_col, kc_blk, kc_pg, kc_hd, vc_blk, vc_pg, vc_hd, args.scale, args.logits_soft_cap, + args.causal, window_left, window_right); + } else if (d == 256) { PagedAttnDecodeGqaBf16<2, 8, 8><<>>( out.Ptr<__hip_bfloat16>(), query.Ptr<__hip_bfloat16>(), k_cache.Ptr<__hip_bfloat16>(), v_cache.Ptr<__hip_bfloat16>(), block_table.Ptr(), seq_lens.Ptr(), @@ -1893,7 +1934,14 @@ void PagedAttentionKernelRocm(Queue& q, Tensor& out, const Tensor& query, const sizeof(float) * (static_cast(kDecWarps) * static_cast(d) + 2 * static_cast(kDecWarps)); dim3 grid(static_cast(total_q), static_cast(hq)); - if (d == 256) { + if (d == 128) { + PagedAttnDecodeOptBf16T<4><<>>( + out.Ptr<__hip_bfloat16>(), query.Ptr<__hip_bfloat16>(), k_cache.Ptr<__hip_bfloat16>(), + v_cache.Ptr<__hip_bfloat16>(), block_table.Ptr(), seq_lens.Ptr(), + query_start_loc.Ptr(), num_reqs, hq, num_kv_heads, d, block_size, bt_row, bt_col, + kc_blk, kc_pg, kc_hd, vc_blk, vc_pg, vc_hd, args.scale, args.logits_soft_cap, args.causal, + window_left, window_right); + } else if (d == 256) { PagedAttnDecodeOptBf16T<8><<>>( out.Ptr<__hip_bfloat16>(), query.Ptr<__hip_bfloat16>(), k_cache.Ptr<__hip_bfloat16>(), v_cache.Ptr<__hip_bfloat16>(), block_table.Ptr(), seq_lens.Ptr(), diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 97009945a..7e7df75f4 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1183,6 +1183,16 @@ add_test(NAME test_vulkan_loader COMMAND test_vulkan_loader) # so it is meaningful on a CUDA box, on the Metal M4, and inert (CPU-only bit- # exactness checks) on a plain CPU build. vllm_cpp_add_test(test_backend_cross_device vt/test_backend_cross_device.cpp) +# The head_dim-128 decode arm (#382) ships default OFF, and its flag is read into +# a `static const bool` -- once per process. The default registration above +# therefore only ever gates the PagedAttnOnline fallback; the arm that runs the new +# kernel needs its own INVOCATION with the flag set. Same shape as the +# test_dense_gateup_fused_marlin_off_* lever pair above. The filter avoids commas: +# doctest reads them as filter separators and the case name contains three. +add_test(NAME test_backend_cross_device_vt_attn_decode_d128 + COMMAND test_backend_cross_device "--test-case=*head_dim 128*") +set_tests_properties(test_backend_cross_device_vt_attn_decode_d128 PROPERTIES + ENVIRONMENT "VT_ATTN_DECODE_D128=1" SKIP_RETURN_CODE 77) # The ROCm GCN-arch capability parse (BACKEND-ROCM W0). UNCONDITIONAL by design: # include/vt/rocm/rocm_arch.h is deliberately free of HIP headers so the one part # of that skeleton carrying a DECISION is gated on a CPU-only CI runner with no diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index cd3820c63..fc5d0e61b 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -1129,6 +1129,99 @@ std::vector Bf16Bits(const std::vector& src) { return out; } +TEST_CASE("paged attention at Qwen3 geometry (bf16, GQA 2, head_dim 128) matches the CPU oracle") { + // #488 / ROCM-DECODE-ATTN-D128: bf16 decode at head_dim==128 (Qwen3/Llama- + // class GQA) fell all the way to the generic PagedAttnOnline on ROCm -- + // every "fast" decode kernel was gated to d==256/512 only. Mirrors the + // Metal "Qwen3 geometry" test's shape (nblocks/bsz/hq/hkv/dh, mixed + // prefill+decode across 2 requests) so a bf16, GQA=2, d=128 case exists + // for every registered device, not just Metal. + constexpr int64_t kNBlocks = 24, kBsz = 16, kHq = 16, kHkv = 8, kDh = 128; + constexpr int64_t kNumReqs = 2; + const std::vector qsl{0, 40, 45}; // req0: 40 new (prefill); req1: 5 new + const std::vector slens{40, 71}; // req1 carries 66 context tokens + const int64_t t_total = qsl.back(); + constexpr int64_t kMaxBlocks = 6; + std::vector btab(static_cast(kNumReqs * kMaxBlocks)); + for (int64_t r = 0; r < kNumReqs; ++r) { + for (int64_t c = 0; c < kMaxBlocks; ++c) { + btab[static_cast(r * kMaxBlocks + c)] = static_cast(r * kMaxBlocks + c); + } + } + + const size_t cache_elems = static_cast(kNBlocks * kBsz * kHkv * kDh); + const std::vector qf = RandomVec(static_cast(t_total * kHq * kDh), 811, -1.5f, 1.5f); + const std::vector kf = RandomVec(cache_elems, 812, -1.5f, 1.5f); + const std::vector vf = RandomVec(cache_elems, 813, -1.5f, 1.5f); + const std::vector qb = Bf16Bits(qf), kb = Bf16Bits(kf), vb = Bf16Bits(vf); + + vt::PagedAttentionArgs args; + args.scale = 1.0f / std::sqrt(static_cast(kDh)); + args.causal = true; + args.query_start_loc_host = qsl.data(); + args.max_seq_len = 71; + + std::vector ref(qb.size(), 0); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector cq_v = qb, ckc = kb, cvc = vb; + std::vector cbt = btab, csl = slens, cqsl = qsl; + Tensor tq = Tensor::Contiguous(cq_v.data(), DType::kBF16, cd, {t_total, kHq, kDh}); + Tensor tkc = Tensor::Contiguous(ckc.data(), DType::kBF16, cd, {kNBlocks, kBsz, kHkv, kDh}); + Tensor tvc = Tensor::Contiguous(cvc.data(), DType::kBF16, cd, {kNBlocks, kBsz, kHkv, kDh}); + Tensor tbt = Tensor::Contiguous(cbt.data(), DType::kI32, cd, {kNumReqs, kMaxBlocks}); + Tensor tsl = Tensor::Contiguous(csl.data(), DType::kI32, cd, {kNumReqs}); + Tensor tqsl = Tensor::Contiguous(cqsl.data(), DType::kI32, cd, {kNumReqs + 1}); + Tensor to = Tensor::Contiguous(ref.data(), DType::kBF16, cd, {t_total, kHq, kDh}); + vt::PagedAttention(cq, to, tq, tkc, tvc, tbt, tsl, tqsl, args); + cpu.DestroyQueue(cq); + } + std::vector reff(ref.size()); + for (size_t i = 0; i < ref.size(); ++i) reff[i] = vt::BF16ToF32(ref[i]); + + for (DeviceType dt : RegisteredDevices()) { + if (!OpAvailable(vt::OpId::kPagedAttention, dt)) continue; + CAPTURE(DeviceName(dt)); + vt::Backend& dev = vt::GetBackend(dt); + Queue q = dev.CreateQueue(); + const Device d{dt, 0}; + + DevBufBytes dq(dev, q, qb.size() * 2), dkc(dev, q, kb.size() * 2), dvc(dev, q, vb.size() * 2), + dout(dev, q, qb.size() * 2); + dq.Upload(qb.data()); + dkc.Upload(kb.data()); + dvc.Upload(vb.data()); + DevBufI32 dbt(dev, q, btab.size()), dsl(dev, q, slens.size()), dqsl(dev, q, qsl.size()); + dbt.Upload(btab); + dsl.Upload(slens); + dqsl.Upload(qsl); + dev.Synchronize(q); + + Tensor tq = Tensor::Contiguous(dq.ptr(), DType::kBF16, d, {t_total, kHq, kDh}); + Tensor tkc = Tensor::Contiguous(dkc.ptr(), DType::kBF16, d, {kNBlocks, kBsz, kHkv, kDh}); + Tensor tvc = Tensor::Contiguous(dvc.ptr(), DType::kBF16, d, {kNBlocks, kBsz, kHkv, kDh}); + Tensor tbt = Tensor::Contiguous(dbt.ptr(), DType::kI32, d, {kNumReqs, kMaxBlocks}); + Tensor tsl = Tensor::Contiguous(dsl.ptr(), DType::kI32, d, {kNumReqs}); + Tensor tqsl = Tensor::Contiguous(dqsl.ptr(), DType::kI32, d, {kNumReqs + 1}); + Tensor to = Tensor::Contiguous(dout.ptr(), DType::kBF16, d, {t_total, kHq, kDh}); + + vt::ResetOpProviderStats(vt::OpId::kPagedAttention, dt); + vt::PagedAttention(q, to, tq, tkc, tvc, tbt, tsl, tqsl, args); + dev.Synchronize(q); + CHECK(vt::GetOpProviderStats(vt::OpId::kPagedAttention, dt).declines == 0); + + std::vector got(qb.size()); + dout.Download(got.data()); + std::vector gotf(got.size()); + for (size_t i = 0; i < got.size(); ++i) gotf[i] = vt::BF16ToF32(got[i]); + CHECK(Nmse(reff, gotf) <= kNmseTol); + + dev.DestroyQueue(q); + } +} + // Rank-3 padded-row view [T, H, D] over a [T, row_stride] f32 buffer — the // merged-qkvz slice shape the GDN/attention glue ops consume in the model. Tensor T3PaddedF32(void* p, Device d, int64_t t, int64_t h, int64_t w,