Skip to content

spec(BACKEND-ROCM): the ROCm head_dim=128 decode arm, the ROCm half of #382 - #564

Merged
localai-bot merged 4 commits into
mudler:mainfrom
joral:row/ROCM-DECODE-ATTN-D128-SPEC
Aug 14, 2026
Merged

spec(BACKEND-ROCM): the ROCm head_dim=128 decode arm, the ROCm half of #382#564
localai-bot merged 4 commits into
mudler:mainfrom
joral:row/ROCM-DECODE-ATTN-D128-SPEC

Conversation

@joral

@joral joral commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

What this is

Spec and records only — no source changes. One spec, its claim, and one row
in the roadmap issue table. The implementation follows in its own PR, against
this spec.

Issue: #382 — "decode-opt attention
kernel is head_dim-256 only; head_dim 128 (Qwen3-dense, Llama, Mistral) falls to
the block kernel." This is the ROCm half. The CUDA half already merged as
#425 (66399617); #382 stays open because only CUDA was covered and the flip to
default-ON is still owed on both backends.

Three-place agreement: .agents/roadmap_v1.md:92 (#382KERNEL-ATTN-PAGED),
the spec's header block, and this PR body.

 .agents/specs/rocm-decode-attn-d128.md        | 269 +
 .agents/claims/CLAIM-ROCM-DECODE-ATTN-D128.md |   5 +
 .agents/roadmap_v1.md                         |   1 +

The gap it specs

The ROCm dispatch gated every fast decode kernel it has to
d == 256 || d == 512:

const bool bf16_decode_opt =
    decode_opt && (d == 256 || d == 512) && query.dtype == DType::kBF16 && ...

head_dim=128 is what Qwen3-dense, Llama and Mistral use, so every one of those
models fell straight through to PagedAttnOnline. Not a gfx1200-specific gap
and not a WMMA-specific one — the file had no fast decode path at all for the
head size the models actually use.

Cross-checked against the pinned oracle (555967922): vLLM's own
CALL_CUSTOM_LAUNCHER_BLK_HEAD switches on head_size with cases 64 and
128
(csrc/rocm/attention.cu:3607-3618), including on RDNA4 via the
is_navi_gpu() launcher variant. Upstream's fast kernel covers d=128; ours
didn't.

The change it specs is correspondingly small: PagedAttnDecodeOptBf16T<EPL> and
PagedAttnDecodeGqaBf16<QG,EPL,NWARPS> are already generic over head_dim via
EPL. The only hardcoded constraint is in the vectorized load/store helpers
(static_assert(EPL == 8 || EPL == 16)), and EPL=4 needs an 8-byte uint2
load instead of a 16-byte uint4 — the same pattern, one size down. No new
algorithm, tiling or synchronization.

It is a mirror of merged work, not new design

§1 is a reconciliation section, written because the re-verify-before-claiming
rule changed the shape of the change. #382 — filed by an outside contributor —
already names this defect exactly, on the CUDA file, and proposes the same fix.
PR #425 already landed it for CUDA.

So this adopts the merged arm's flag (VT_ATTN_DECODE_D128), default (OFF)
and stated reason verbatim rather than inventing new ones: the arm is
correctness-complete but not byte-exact, because warp-strided online softmax
reduces the KV sequence in a different order than the kernel it replaces, so
a greedy anchor can move at an exact bf16 tie. Shipping OFF keeps every existing
golden byte-identical.

The part worth reviewing closely: §5's cross-architecture reversal

#382 measured this same EPL=4 arm on sm_110 / Jetson AGX Thor as 1.6x
slower
than the kernel it replaces (81.6 → 131.0 ms at c=1, worse at c=8).
gfx1200 measures it 3.53x faster.

Both can be true — different kernels, different fallbacks (CUDA's generic block
kernel vs ROCm's PagedAttnOnline), different occupancy and LDS budgets,
different memory systems. The spec records the disagreement rather than
averaging it away, and draws the conclusion that follows: a shared default is
forbidden.
Any future flip to default-ON must be argued per backend with
per-backend measurement. The ROCm win is not evidence for the CUDA arm, and
#382's sm_110 regression is not evidence against this one.

That is what keeps #382 open, and it is the single most important claim in this
spec to agree or disagree with.

One deviation from the original commit, and one reviewer's call

  1. §7's link to the rocWMMA spec was de-linked. The rocWMMA d=128 arm is a
    separate spike with its own claim and its own (unfiled) issue, and it is
    deliberately not in this PR. That left §7 pointing a markdown link at a file
    this branch does not contain — which check-agent-record catches as
    dangling link rocm-decode-attn-d128-wmma.md, a hard gate failure, not a
    cosmetic one. §7 now names the file in backticks as "landing separately"
    instead. That is the only change from the originally committed spec; the
    rest of the file and the claim are byte-identical. §4's decode_wmma
    disjunct and §6's "the WMMA spec" are prose forward-references to the same
    thing and needed no change.

  2. The spec already carries an ## Outcome section. AGENTS.md §"Spec before
    code" says the spec is committed before implementation, never written up
    afterwards. The commit order satisfies that, but the timestamps show spec and
    implementation were authored 37 seconds apart, so it was not a real sequence —
    the Outcome was written from work already done. Recording it here as a
    reviewer's call rather than deciding it in my own favour.

What is deliberately not here

  • The implementation. src/vt/rocm/rocm_paged_attn.hip (+474), a new bf16
    Qwen3-geometry cross-device test (+93), two flag-on ctest registrations, and
    one scripts/env-doc-allowlist.txt entry. Follows in its own PR against this
    spec.
  • The rocWMMA d=128 spike. Separate spec, separate claim, separate issue.
  • The flip to default-ON, on either backend. Owes the near-tie razor, a
    distributional gate and golden regen under the ratified-tie rule — and per §5
    argued per backend, not once.
  • qg=4/qg=8 GQA fusion at any d (PagedAttnDecodeGqaBf16's fused
    condition only ever covered qg==2 and qg==8 && d==512; real, but
    head-dim-independent), and prefill at d=128.

Gate status

scripts/agent-preflight.sh on this branch: 11 gates fail, and the set is
identical to main's own.
Confirmed rather than asserted — the same script was
run in a clean detached worktree at mudler/vllm.cpp main (fafa16f0, this
PR's base), which fails exactly these: check-release-workflow,
check-test-registration, audit-live-rows (stale ACTIVE row
ENG-FORGE-COAUTHOR, unrelated), test_release_{archive,pipeline,metadata,accelerator_metadata,macos_metadata},
test_agent_onboard (asserts the default branch is master, gets main),
test_check_test_registration, test_audit_live_rows. This PR adds no gate
failure.

A twelfth was ours and is fixed: check-commit-trailers — a gate that landed
on main after this branch's original base — rejected
Assisted-by: Claude Sonnet 5 [Claude Code] as malformed against
ASSISTED_BY = <identity>:<model> [tool]. Now
Assisted-by: Claude:claude-sonnet-5 [Claude Code], and the checker reports
OK: commit trailer contract.

The commit SHA was rebuilt. The branch was originally cut at bbc482a2, 167
commits back, and carried that malformed trailer. It was rebased onto current
main (fafa16f0) rather than merged, so there is no merge commit and the
trailer fix reaches the commit. Apart from the one de-linked line above, the
spec and claim are byte-identical to the original commit, and the roadmap
issue-table edit is exactly the one #488 row with every other key untouched —
a keyed record reapplied scope-wise, not three-way merged.

Implementation re-gated on this base, gfx1200, GPU lock held

The spec's evidence was taken at the old base. Since this PR carries no source,
this is context for the follow-on rather than this PR's own gate — but the
implementation branch was rebuilt and re-gated on fafa16f0 so the spec's
Outcome is not quoting a tree that no longer exists:

result
cmake --build build-hip (gfx1200, ROCm 7.2.3, Release) 783/783, clean
ctest -R 'rocm|cross_device' 6/6, incl. both flag-on registrations
new case under VT_ATTN_DECODE_D128=1 1 case, 6 assertions, NMSE 2.97e-11 vs tol 5e-4
full ctest 8 failures, all pre-existing (below)

Those 8 are attributed, not waved past. Every one throws the same root cause
vt: no kernel for op 63 on device type 5 (src/vt/op_provider.cpp:254).
OpId(63) is kSharedExpertGate, a MoE shared-expert gate op with no ROCm
kernel registered; the ROCm implementation of it is open and unmerged as PR
#509
("ROCm MoE combine/gate ops — SharedExpertGate, MoeCombine,
MoeCombineGate", issue #41). Nothing to do with paged attention.

Note the spec's own figures say ctest 385/393; the tree is now 399 tests with
the same 8 failures. Same failures, more tests on main. The spec was left as
the record of what was measured when, rather than silently retconned.

The performance figures in §5 — 3.53x TPOT, the +42.7%/+25.0%/+17.8% decode
throughput table — were taken at the original base and have not been
re-measured. They belong to the implementation PR. §5 already carries the caveat
that the throughput table is single-run on a board that may also drive a
display: indicative, not the 2-3x-idle-reproduced standard.

…alf of mudler#382

FOLLOWING_AGENTS_PROTOCOL

Spec and claim only -- no source changes. The implementation follows in its
own PR, against this spec.

Reconciles the record before implementing, per the re-verify-before-claiming
rule. Issue mudler#382 already names this exact defect ("decode-opt attention kernel
is head_dim-256 only; head_dim 128 falls to the block kernel") and PR mudler#425
already merged the CUDA half of it. This spec is therefore the ROCm MIRROR of
landed work, not new design, and it adopts that arm's flag
(VT_ATTN_DECODE_D128), default (OFF) and stated reason verbatim: the arm is
correctness-complete but not byte-exact, because warp-strided online softmax
reduces the KV sequence in a different order than the kernel it replaces, so a
greedy anchor can move at an exact bf16 tie.

mudler#488 -- the ROCm 8.1x per-call decode-attention measurement -- is recorded as
the motivating observation, not the issue: it explicitly asserts no cause, and
this change supplies one without closing it. Adds mudler#488 to the roadmap issue
table, where it was missing.

The finding worth carrying: mudler#382 measured this same EPL=4 arm 1.6x SLOWER on
sm_110, where gfx1200 measures it 3.53x faster. Recorded, not reconciled. It is
why the flip to default-ON must be argued per backend rather than once, and it
is what keeps mudler#382 open.

Row: BACKEND-ROCM (the ACTIVE row whose code this edits; mudler#382 itself is filed
against the cross-backend kernel row, which check-agent-record will not accept
as an active claim's row -- the two differ deliberately and the claim says so).
Issue: mudler#382

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: Claude:claude-sonnet-5 [Claude Code]
@joral
joral marked this pull request as ready for review August 13, 2026 01:27
…OWER than the CUDA arm it mirrors (mudler#382)

FOLLOWING_AGENTS_PROTOCOL

Spec only -- no source changes.

Found by reading the merged CUDA arm (6639961) hunk by hunk against this one.
CUDA added TWO EPL=4 loaders, not one: LoadRowN<4, __nv_bfloat16> AND
LoadRowN<4, float>. The second exists because CUDA's d == 32*4 launch branch
carries no dtype gate at all, and cuda_paged_attn.cu:2796 dispatches
LaunchDecode<TQ, TKV, float> for out.dtype == kF32 -- so f32 reaches the
decode-opt kernel on CUDA.

ROCm's bf16_decode_opt requires query, k_cache, v_cache and out to ALL be
kBF16, and LoadRowEplBf16/StoreRowEplBf16 are bf16 by construction. The
fallback dispatch supports FIVE dtype combinations, so four of the five still
fall to PagedAttnOnline at d=128 -- the exact fallback this arm exists to get
off. f32 is live, exercised surface on ROCm, not dead code.

The limitation is PRE-EXISTING, not introduced: ROCm's decode-opt has been
bf16-only at every head_dim, so the d=128 arm inherits the gap rather than
creating it. Recording it anyway, because §1 calls this change a mirror of
merged work and on dtype coverage it is NOT one -- and because an arm that is
not implemented is owed a record, not left to be discovered later. Not waived:
an f32 d=128 decode is silently slow rather than refused, which is the weaker
of the two failure modes.

§1 gains the same qualifier so a reader of the reconciliation section is not
told "mirror" without the exception.

Row: BACKEND-ROCM
Issue: mudler#382

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: Claude:claude-opus-5 [Claude Code]
@localai-org-maint-bot

Copy link
Copy Markdown
Contributor

@mudler @richiejp This spec-only PR looks good to merge from my review. It cleanly separates the ROCm implementation follow-up, preserves the CUDA arm’s default-OFF correctness policy, records the gfx1200/sm_110 performance reversal instead of generalizing across backends, and its agent-record gates pass. The current Linux/Windows reds are documented as baseline failures rather than changes introduced by these three record files.

…, and lacked Risks and Stop conditions (mudler#382)

Review repairs on top of joral's spec, pushed to the PR branch so the authorship
of the work itself stays with them. Spec-only; no source file is touched.

The blocking one: `## Outcome (2026-08-12)` opened with "Landed the ROCm `d=128`
decode arm". Nothing landed. On main, `git log -S'VT_ATTN_DECODE_D128' --
src/vt/rocm/` is EMPTY and `rocm_paged_attn.hip:1684` still reads
`d == 256 || d == 512`. AGENTS.md scopes `## Outcome` to a row reaching DONE,
and `BACKEND-ROCM` is ACTIVE with the code unmerged. Merging that text would put
a false "landed" on main, where the next agent greps for the flag, finds
nothing, and cannot tell "never merged" from "reverted" from "renamed" -- the
exact ambiguity this project has already lost time to.

Renamed to `## Result on the implementation branch`, with a banner stating
outright that it is not landed, quoting the two commands that show it, and
naming what has to be true before it becomes `## Outcome`. Section 4 moves from
"What changed" to "What will change" for the same reason: it enumerates five
edits to `rocm_paged_attn.hip` that are not in the tree.

Added the two sections AGENTS.md requires and this spec did not have:

- `## 8. Risks and decisions` -- the bf16-tie reduction-order risk that is the
  reason the arm ships default OFF, the single-board provenance of the 3.53x,
  the sm_110 1.6x-SLOWER reversal recorded but deliberately unreconciled, the
  spec-before-code consequence, the dangling WMMA flag, and the unquantified
  residual mudler#488 gap.
- `## 9. Stop conditions` -- stop if the gate moves with the flag OFF, stop if
  the flag-ON arm cannot be shown to REACH the new kernel (a green gate that
  never entered the kernel proves nothing), stop before flipping the default,
  stop if a second board contradicts the first, and do not extend this row to
  the rocWMMA arm.

`## 6. What this does not claim` was left where it is; it is a non-claims list,
not a stop condition, and it reads correctly beside the new sections.

Three smaller corrections, each verified rather than assumed:

- `**Base:**` said `bbc482a2`; the branch was rebased and its own claim file
  records `fafa16f0`. The spec and the claim now agree.
- Section 4 specifies a gate referencing `VT_ATTN_DECODE_WMMA`. That flag exists
  nowhere in `src/` or `tests/`, and `rocm-decode-attn-d128-wmma.md` is neither
  in this PR nor on main, so an implementer following section 4 verbatim writes
  a reference to an undefined symbol. Marked as a forward reference.
- The upstream anchor `attention.cu:3607-3618` is `3609-3620` at the pin. The
  content it names is exact; the line numbers had drifted by two.

Verified on this tree: check-doc-checkpoint OK, check-agent-record OK
(ENGINE=152 MODEL=362 QUANT=82 KERNEL=51 BACKEND=81).

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: Claude:claude-opus-5 [Claude Code]
@localai-bot

Copy link
Copy Markdown
Collaborator

I've pushed a commit to this branch (de49041) rather than sending you round again — the fixes were all in the spec text and mechanical enough that a round trip would have cost more than it was worth. Your two commits are untouched and still yours; mine sits on top.

The one that was blocking. ## Outcome (2026-08-12) opened with "Landed the ROCm d=128 decode arm". It hasn't landed — on main, git log -S'VT_ATTN_DECODE_D128' -- src/vt/rocm/ is empty and rocm_paged_attn.hip:1684 still reads d == 256 || d == 512. Merging that text would put a false "landed" on main, and the next person greps for the flag, finds nothing, and can't tell "never merged" from "reverted" from "renamed".

Renamed to ## Result on the implementation branch, with a banner saying outright that it is not landed, the two commands that show it, and what has to be true before it becomes ## Outcome. ## 4. What changed## 4. What will change for the same reason — it lists five edits to rocm_paged_attn.hip that aren't in the tree.

Added the two sections AGENTS.md requires and the spec didn't have: ## 8. Risks and decisions and ## 9. Stop conditions. I drew both from material already in your spec rather than inventing positions — the bf16-tie reduction-order risk that's the reason for default OFF, the single-board provenance of the 3.53x, the sm_110 reversal recorded-but-unreconciled, the dangling WMMA flag, and the unquantified residual #488 gap. Read them and push back if I've put words in your mouth anywhere. I left ## 6. What this does not claim alone; it's a non-claims list rather than stop conditions, and it reads fine beside them.

Three smaller ones: **Base:** said bbc482a2 where your own claim file records fafa16f0 after the rebase, so they now agree; the VT_ATTN_DECODE_WMMA gate in §4 is marked as a forward reference, since that flag exists nowhere in src/ or tests/ and the sibling spec is neither here nor on main; and the upstream anchor attention.cu:3607-3618 is 3609-3620 at the pin — content exact, line numbers drifted by two.

Worth saying about the review itself: the technical content here checked out better than most specs I've reviewed in this repo. Every source anchor you cite resolves exactly — bf16_decode_opt at :1684 matches your quoted snippet verbatim, the static_assert(EPL == 8 || EPL == 16) at 268 and 289, the launch switches at 1875-1904 with only 8/16 instantiations, the five-combination fallback dispatch at 1937-1947, and the CUDA counter-claim at cuda_paged_attn.cu:321/329/2796. The upstream anchor is right at the pinned oracle too. And §5 labels its own throughput table "indicative, not the 2-3x-idle-reproduced standard" and states plainly that no post-change per-call oracle re-measure was run — that's the spec being honest about a gate it couldn't reach, which is exactly right.

On CI: the build-test-cpu red is a flake, not yours — test_cpu_threadpool asserts a wall-clock ratio on a shared runner and measured 111 against a 100 threshold. Your PR changes only Markdown, so it can't cause a C++ timing failure. Tracked as #631. Windows is the known-broken arm (#584). And pr-size/agent-record red on fork PRs because the checker demands the base SHA be an ancestor of head, which it never is once main moves — I hit it verifying your trailers and it's now filed as a real bug on our side; from the merge base your trailers check OK.

Merging once CI settles. Thanks for this one — it was the least work of the twelve to get right.

@joral

joral commented Aug 13, 2026

Copy link
Copy Markdown
Contributor Author

I still need to figure out my process a little. When I completed the graph capture changes with almost no performance impact, I felt the need to find a reason. Went down the rabbit hole here and then into the wmma changes. Then tried to figure out how best to untangle the two separate issues into the requested spec and implementation changes, and how to apply them in the right order.

Will make sure the implementation is cleaned up and ready for PR tomorrow, then the WMMA changes.

Brings the branch current so the checkers can compute a range: they require the
base SHA to be an ancestor of head, which stops being true once main moves.
Merging rather than rebasing keeps joral's commits byte-identical.

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: Claude:claude-opus-5 [Claude Code]

# Conflicts:
#	.agents/roadmap_v1.md
@localai-bot
localai-bot merged commit 373aa12 into mudler:main Aug 14, 2026
20 of 23 checks passed
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.

4 participants