fix(rocm/gemma4): #838 widen indexed MoE gate to T<=63 + retirement-safe dispatch

[bakon11]

fix(rocm/gemma4): #838 widen indexed MoE gate to T<=63 + retirement-safe dispatch

Fixes #838.

Problem

T=2..63 MoE tokens are observed on the serial M1 path, and the Gemma-4 FP8
device-indexed gate admitted only T == 1. Everything wider fell through to the
legacy host-gather path, which copies the router weights and indices back to the
host every step. Dispatch retirement was also not scope-safe.

What changed

• The indexed gate widens from T == 1 to T <= 63, behind
  VT_GEMMA4_DECODE_INDEXED_MAX_T (unset → 63, =1 restores the old T=1-only
  behavior, clamped to [1,63]). The bound is not arbitrary: Gemma4IndexedOkT
  also requires T < kGemma4PrefillBatchMinT, the same constant that admits the
  prefill-batch path at 64, so the two gates cannot both claim a token count.
• A latent double-scale is fixed along the way. The old code applied the
  per-expert scale to the router weights rw in place before attempting the
  indexed arm, so a fall-through to the host-gather path scaled them a second
  time and squared the factor. The scale now lands on a copy, and
  Gemma4ApplyHostExpertScaleOnce takes an already_scaled argument so the
  fallback cannot repeat it. gemma4 indexed-max-t: fallback scale is once, not s^2 pins that: it asserts the result is orig * 3 and not orig * 9.
• The T-loop, arm selection, scratch choice and fail-path ordering move into
  gemma4_indexed_gate.h as injectable host seams, so the retire-before-release
  ordering is testable without a GPU. Peer and compute work is retired while the
  scratch buffer is still owned; a buffer whose retirement was not observed is
  quarantined rather than returned to the pool.
• Gemma4IndexedCall's product instantiations are named explicitly
  (<uint16_t, uint16_t>) so the product lambdas compile against current main.
• Host indexed tests: 10 cases, 1101 assertions.

Verification

Built and smoked on gfx1201, shipping ROCm 7.2.4 (clang f58b06d), as the
combined #837+#838+#839 stack on d1b0ea3a: gfx1201 build BUILD_RC=0, all 30
layers resident and bound, ready in ~20 s; batched-MoE prefill deadlock-free
(T=2012 done, PEER_ACT active, no hang, no HIP error); Paris/63 quality PASS.

Stated plainly, because it bounds what this PR alone demonstrates: that smoke
covers the three-PR stack, not this branch in isolation. What is isolated
here is the host-side evidence — the 10 cases above, which run on
build-test-cpu and include real RED mutants for the double-scale, the
release-before-retire ordering, the argument packing and the T-loop striding.

No throughput measurement accompanies this change. #838 is filed as perf
and the change flips a product default for T=2..63, so docs/BENCHMARKS.md
still owes a before/after against VT_GEMMA4_DECODE_INDEXED_MAX_T=1 on the same
host. Correctness comes first there: the token-exact result for T=2..63 is owed
before any throughput ratio is accepted. Both need the gfx1201 pair. They are
recorded under ## Open on gfx1201 hardware in
.agents/specs/rocm-gemma4-indexed-max-t.md rather than left implicit.

Maintainer changes on top

Three repairs were applied while landing, none touching product behavior:

• The donor evidence moved from .agents/evidence/ — a path
  scripts/check-pr-size.py cannot classify, so the gate refused the change
  outright — to .agents/specs/rocm-gemma4-indexed-max-t-donor.{md,log}, which
  match SPEC and SPEC_EVIDENCE at check-pr-size.py:188-189. The bytes are
  unchanged; both slice sha256s still match the manifest table.
• docs/FEATURES.md is left as main has it. The branch had rewritten the
  Gemma4 row, dropping the spec link, the VT_GEMMA4_*/VT_ATTN_* env pointer
  and the test_gemma4_rocm_fp8_seams seam name. Restoring those and adding the
  ROCm Gemma-4 hyp A: T=2..63 MoE is observed on serial M1; widen indexed gate #838 sentence does not fit — check-public-doc-tables caps a cell at 220
  characters and main's is already 219 — and check-doc-checkpoint does not ask
  for a FEATURES edit here, because feature_surface now keys off a change to
  the set of REGISTER_VLLM_MODEL registrations (doc-checkpoint makes docs/FEATURES.md a LOCK: 2 of 5 concurrent LTX-2.5 PRs conflict there and nowhere else #595, fix(#595): doc-checkpoint asks whether the registry moved, not whether a file did #1086) and this branch
  changes none. The ROCm Gemma-4 hyp A: T=2..63 MoE is observed on serial M1; widen indexed gate #838 detail stays where the branch already put it in full:
  docs/USAGE.md, docs/ENVIRONMENT.md and the spec.
• The source-invariant case sliced rocm_gemma4_experts.hip between
  RetireGemma4Fp8TopKIndexedPeer and the next symbol, then asserted
  return true; was absent. That is a slice of the gap between two
  functions, not of the function, so an unrelated definition added after the
  closing brace lands inside it. A git merge-tree of this branch with open PR
  fix(rocm/gemma4): #839 retirement-safe prefill peer GeGLU Launch/Finish #1047 (3df512a11) grows that gap from 16 lines to 176 and adds three
  return true; inside it, which would have turned main red on
  build-test-cpu once both landed — this case is registered unconditionally
  and reads the source from disk. It now brace-matches the function's own body
  and gains two assertions, that the slice is bounded and that the function does
  synchronize the compute stream, so the guarantee is stated more completely
  than before rather than relaxed. Proved against the real test binary, built
  standalone and pointed at a scratch source root: the repaired case reds on a
  mutant that discards the sync status and returns unconditional success (3
  assertions fail), and stays green on the actual merged tree, where the
  original case reds. The scratch copy was restored and its sha256 compared
  against the tree.

The branch was rebuilt by rebase rather than merge, so it carries no untrailered
merge commit; all six original commits keep their authorship. The docs/USAGE.md
conflict against #837's landed text was resolved as a union — both the GetBlas
dual-slot paragraph and this row's env description survive.

Known-unrelated CI

windows-msvc-cpu and windows-msvc-vulkan are red on every open PR from a
break predating this branch (#503, #584).

FOLLOWING_AGENTS_PROTOCOL

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

[localai-bot]

Thanks — the core of this is right, and the double-scale fix inside it is a real bug caught in passing. I rebased the branch onto current main and pushed the maintainer repairs so the gates go green, same as I did on #1045. I have not merged it, because two of the findings below need the gfx1201 pair and I would rather leave them open than paper over them.

What I changed on your branch

Your six commits are preserved with their authorship, rebased (not merged, so no untrailered merge commit). One repair commit sits on top.

1. .agents/evidence/ is a path check-pr-size.py cannot classify, so the gate refused the change outright:

   ERROR: PR size check could not classify the change: unclassified repository path
   '.agents/evidence/rocm-gemma4-indexed-max-t/MANIFEST.md'
   

   The classifiable homes are .agents/specs/*.md (SPEC) and .agents/specs/*.{patch,json,log} (SPEC_EVIDENCE), at check-pr-size.py:188-189. The three files are now rocm-gemma4-indexed-max-t-donor.md, -donor-gate-911-933.log and -donor-ok-1088-1110.log, with the manifest table and the spec's two references updated. The bytes did not change: both slice sha256s still match the table, which is exactly what that table exists to check. This is the same move fix(rocm/gemma4): #837 GetBlas dual-slot TLS + host lifetime seam #1045 needed — worth knowing for fix(rocm/gemma4): #839 retirement-safe prefill peer GeGLU Launch/Finish #1047 and ROCm Gemma-4 hyp C: prefill peer GeGLU helper is monolithic; retirement-safe Launch/Finish #839.

2. docs/FEATURES.md is back to what main has. Your rewrite of the Gemma4 row dropped the spec link, the VT_GEMMA4_*/VT_ATTN_* env pointer and the test_gemma4_rocm_fp8_seams seam name. I could not restore those and add the ROCm Gemma-4 hyp A: T=2..63 MoE is observed on serial M1; widen indexed gate #838 sentence: check-public-doc-tables caps a table cell at 220 characters and main's cell is already 219. And it turns out nothing asks for the edit — as of fix(#595): doc-checkpoint asks whether the registry moved, not whether a file did #1086/doc-checkpoint makes docs/FEATURES.md a LOCK: 2 of 5 concurrent LTX-2.5 PRs conflict there and nowhere else #595, check-doc-checkpoint's feature_surface class keys off a change to the set of REGISTER_VLLM_MODEL(...) registrations rather than off the path src/vllm/model_executor/models/, and this branch changes none. Your ROCm Gemma-4 hyp A: T=2..63 MoE is observed on serial M1; widen indexed gate #838 detail already lives in full in docs/USAGE.md, docs/ENVIRONMENT.md and the spec, so nothing was lost.

3. documentation-checkpoint goes green on the rebase, and the reason is worth knowing. It was not your change — the fix in (2), commit 8fa405bb7, is not an ancestor of your old base d1b0ea3a but is an ancestor of current main. Rebasing picked up the repaired checker. If any of your other branches show that red, rebasing is the whole fix.

4. docs/USAGE.md conflicted against ROCm Gemma-4 hyp B: GetBlas single TLS destroys hipBLAS handle on peer-MoE device hop #837's landed text in all four of your commits that touch it. I resolved it as a union each time, so both the GetBlas dual-slot paragraph and this row's VT_GEMMA4_DECODE_INDEXED_MAX_T description survive. I also dropped main's now-false "No new env." sentence, since this row adds one.

5. The source-invariant case would have turned main red once fix(rocm/gemma4): #839 retirement-safe prefill peer GeGLU Launch/Finish #1047 landed. This is the one I would most like you to look at, because it is a shape rather than a typo:

   const auto ret = hip.find("bool RetireGemma4Fp8TopKIndexedPeer");
   const auto ret_end = hip.find("RunGemma4Fp8ExpertGeGLUPrefillOnExpertDevice", ret);
   const std::string retire = hip.substr(ret, ret_end - ret);
   CHECK(retire.find("return true;") == std::string::npos);

   That slices the gap between two symbols, not the function. Anything added after RetireGemma4Fp8TopKIndexedPeer's closing brace lands inside the slice, and this case is registered unconditionally and reads the .hip from disk — so it runs on build-test-cpu.

   I checked this against the real merged tree rather than reasoning about it. git merge-tree of this branch with fix(rocm/gemma4): #839 retirement-safe prefill peer GeGLU Launch/Finish #1047's head 3df512a11 grows that gap from 16 lines to 176, and the added region carries three return true; of its own. Running your test binary against that exact merged file:

   [doctest] test cases:   10 |    9 passed | 1 failed | 0 skipped
   [doctest] assertions: 1099 | 1098 passed | 1 failed |
   [doctest] Status: FAILURE!
   

   So main goes red once both land, with a failure pointing at a function neither PR touched.

   It now brace-matches the function's own body via a FunctionBody() helper (which skips a forward declaration by requiring the first { before the first ;), and gains two assertions: that the slice is bounded, and that the function does synchronize the compute stream. The intent is stated more completely than before, not relaxed.

   I built the real test TU standalone against a scratch source root and ran both versions three ways:

   source under test	your case	repaired case
   as committed	SUCCESS, 1099 assertions	SUCCESS, 1101 assertions
   mutant: sync status discarded, return true;	FAILURE, 3 failed	FAILURE, 3 failed
   real merge-tree of this branch + fix(rocm/gemma4): #839 retirement-safe prefill peer GeGLU Launch/Finish #1047	FAILURE, 1 failed	SUCCESS

   The middle row is the one that matters most: the repaired case still catches the defect it exists for, so the intent is intact rather than relaxed. The bottom row is the landmine. The scratch copy was restored and its sha256 compared against the tree afterwards; nothing in the worktree was mutated.

What I could not verify

I have no RDNA4 hardware and did not build the project (this host is at 95% disk, and a false ENOSPC failure reads as a code verdict). The test above was compiled and run as a standalone TU — real doctest, real binary, real assertions — and the full test file passes -fsyntax-only against the project headers. It has not been through the project's CMake build or a HIP compile.

What needs you

These are recorded under ## Open on gfx1201 hardware in .agents/specs/rocm-gemma4-indexed-max-t.md so they do not get lost. The first two are why I have not merged.

A. rw_idx breaks the invariant Scope item 3 exists to protect. At gemma4_moe.cpp:760 the T=1 arm now reads its router weights from a per-call pooled DBuf where it previously used the caller-owned rw. A pooled buffer's address moves between calls, so it is not capture-stable — while the comment three lines below still reads "Stable T=1 acc for hipGraph (do not pool-Release)", and your spec's Scope item 3 says "T=1: keep hipGraph-stable TLS acc". Decode hipGraph is lab-only today, which is exactly why nothing observes this; it is latent, not absent. I did not touch it because you own the graph-capture context and the hardware to check it. Either hold rw_idx in the same TLS that holds acc_fast, or say in the spec that T=1 has stopped being capture-stable and why that is acceptable.

B. No measurement. #838 is filed as perf, and this flips a product default across T=2..63. Nothing in docs/BENCHMARKS.md or the spec records a before/after. Correctness first: the token-exact result for T=2..63 against VT_GEMMA4_DECODE_INDEXED_MAX_T=1 on the same host, then the throughput ratio. Both need your box.

C. The per-expert scale costs one kernel launch per token — up to 63 of them. gemma4_moe.cpp:766-773 calls vt::ApplyExpertScaleRw inside the T-loop. ApplyExpertScaleRwKernel (rocm_fp8_channel_gemv.hip:635-641) is <<<1, G>>> with g = threadIdx.x, and it indexes rw[g], ri[g] and escale[ri[g]] purely linearly — no reduction, no shared memory, no cross-thread interaction — and ri holds global expert ids. So one call with G = T*top_k is element-wise identical to the loop, and Gemma4IndexedOkT bounds top_k <= 8 and T <= 63, giving G <= 504, comfortably inside a 1024-thread block.

I deliberately did not make this change, for a reason I think is worth stating: collapsing the loop removes the structural guarantee that G stays small. Today G = top_k <= 8 by construction; after the collapse G scales with T, and ApplyExpertScaleRwRocm checks neither G against the block limit nor the launch result. Widen the T bound later and the launch fails silently, the scale never applies, and the tokens are quietly wrong. So the one-call version is worth doing and should ship with that guard — on hardware that can run it, with (B) already answered so the win is a number rather than a claim.

D. Six test-only helpers ship in a production header. Gemma4IndexedHostApplyToken, Gemma4IndexedHostSerialRef, Gemma4IndexedOracleClose, Gemma4IndexedArgsEq, Gemma4IndexedMayReleaseToPool and Gemma4IndexedFailPathRetireThenMaybeRelease (gemma4_indexed_gate.h:113-160) have no production caller. Moving them into tests/ is mechanical, and I left it to you only because it churns includes across your other open branches — say the word and I will do it.

While you are there: the "tensor oracle" computes s = (e+1) * rw[g], which models none of the FP8 GeGLU arithmetic. What it actually gates is the T-loop's striding and the argument packing — genuinely worth gating, and the tests around it do catch real mutants. But the name promises a numerical oracle it is not, and someone will eventually read a green run as evidence the kernel math is right.

E. Smaller, all yours:

• Gemma4IndexedHelperHits() does a global std::atomic fetch_add per token on the decode hot path purely so a host test can observe it.
• Gemma4IndexedScratchValidForT cannot return false in the T>1 branch that calls it — Gemma4IndexedScratchKindFor(T) supplies the kind, y_base is non-null and elems == T*H, so all three conditions hold by construction. Dead guard, or the call site is missing a case.
• The new RestoreComputeDev (rocm_gemma4_experts.hip:572) adds a second hipSetDevice on a success path that already ends with one at :659 — and whose own comment at :639 says "no hipSetDevice between stream ops (graph-prep + less driver tax)".
• RetireGemma4Fp8TopKIndexedPeer returns ok = true while skipping the expert-stream drain entirely when tls.edev != expert_dev. That is a real retirement that did not happen, reported as one that did.
• One comment string is load-bearing: the source-invariant case asserts retire-before-acc_idx-dtor appears in gemma4_moe.cpp. Reordering the statements it names while keeping the comment passes; deleting the comment while keeping the order fails. I left it because fixing it properly means asserting on the statement order, which is a bigger change than this landing pass should make.

One heads-up for #1047

The same merge-tree shows docs/USAGE.md conflicting between the two branches — you have four commits here rewriting the "Gemma4 FP8 on ROCm (RDNA4)" paragraph, and #837 landed its own text into it. I resolved it as a union on this branch; #1047 will want the same treatment rather than taking one side, or the GetBlas dual-slot sentences disappear.

---

Happy to take A and C myself if you send me a trace, or just confirm the graph-capture behavior — I mainly need someone with the hardware to say which way A should go. Once A and B are answered this is ready; everything else here is either landed or filed.

[localai-bot]

@bakon11 — I re-verified this against current main and did not push anything. It is still clean and I did not want to reset its CI approval for no gain.

State as of cff626f93:

• Merges cleanly into main at both affc2a7fd and dd8a3b0e1 (git merge-tree, no conflict, including the .agents/issue-index.md union).
• Rebased locally onto current main and ran the full preflight: all gates green, pr-size OK, check-doc-checkpoint OK. That last one is worth spelling out, because running it against the branch's own stale copy of the script reports a failure: the branch is 12 commits behind and its check-doc-checkpoint.py predates doc-checkpoint makes docs/FEATURES.md a LOCK: 2 of 5 concurrent LTX-2.5 PRs conflict there and nowhere else #595, which narrowed feature_surface from "any file under models/" to "a changed REGISTER_VLLM_MODEL set". CI checks out the merge ref, so it uses main's checker and passes — which is why reverting docs/FEATURES.md to main's row was the right call and stays right.
• The test_gemma4_indexed_max_t suite is 10/10, 1101 assertions on the rebased tree.
• The only non-green checks are windows-msvc-cpu and windows-msvc-vulkan, which are the pre-existing break (main does not compile under MSVC, and main-baseline.py reports GREEN because windows-msvc-* are PR-only jobs it never runs #503, test_openai_api_server crashes on Windows with STATUS_STACK_BUFFER_OVERRUN (0xC0000409), unmasked by the #512 fix #584) — the failure is test_openai_api_server.exe exiting 0xC0000409, nothing to do with a ROCm MoE gate. The one verify (vulkan) red was GitHub's codeload returning 429/503/502 while downloading docker/setup-buildx-action; I re-ran it.

Your FunctionBody() brace-matcher is confirmed against #1047 on a real merged tree, not by argument. I merged the two rebased branches and measured:

slice	#838 alone	#838 + #839 merged
FunctionBody("bool RetireGemma4Fp8TopKIndexedPeer")	13 lines, no return true;	13 lines, no return true;
the old symbol-to-symbol slice	17 lines, no return true;	177 lines, contains return true;

rocm_gemma4_experts.hip auto-merges clean, and the whole suite is 10/10 / 1101 assertions against the merged sources. So the old assertion really would have red-ed and the repair is load-bearing.

---

Two things still hold this open, and both are yours.

1. src/vllm/model_executor/models/gemma4_moe.cpp:760 puts a pooled DBuf on the hipGraph-capture path. When escale_ptr is set you now copy the router weights into a per-call rw_idx and pass helper_rw into Gemma4IndexedDispatchTokens on the TlsT1 arm — three lines above the comment "Stable T=1 acc for hipGraph (do not pool-Release)", and Scope item 3 in the spec says the T=1 path stays capture-stable. aft.acc is still TLS-stable, but a pooled allocation's address moves between calls and kernel arguments are baked at capture, so the rw pointer in a captured graph can go stale. main scaled rw in place and passed rw.ptr(), which was caller-owned and therefore stable.

Your reason for not mutating in place is good and I am not asking you to undo it — "fallback must still see unscaled weights" is a real correctness fix for T>1. The shape that keeps both is a TLS-stable rw_idx on the T=1 arm, keyed the way AccFastTls is keyed (device, plus T*top_k), so the scaled copy has a fixed address for the life of the capture. That is a small change, but it is on the decode path on your hardware and I am not going to write it blind.

It is latent today — decode hipGraph is lab-only — but it is silent, which is the part I do not want to leave undocumented.

2. No measurement on a perf default flip across T=2..63. This widens the device-indexed gate from T == 1 to T <= 63 by default, which is the right direction, and the host-gather path it replaces copies router weights and indices back to the host every step. But nothing in the PR, the spec or docs/BENCHMARKS.md says what it bought, and the interesting part is the shape — whether the win holds at T=2 and T=8, not just at T=63. Only you have the box.

Everything else here is ready.

[bakon11]

Both items accepted — the TLS-stable rw_idx fix (your suggested AccFastTls-keyed shape) and the T={2,8,63} gate-flip measurements are queued on the 2×R9700 now; fix lands on this branch (on top of 9d4a16c5), numbers go to docs/BENCHMARKS.md. Thanks for confirming the brace-matcher on a real merged tree.

[bakon11]

Both owed items delivered (pushed cff626f9 → ed1c0fe6, stacked on your rebase, no force):

1. hipGraph capture-stability — e2ebf084: TLS-stable T=1 rw_idx keyed like AccFastTls (device + T·top_k), exactly the shape you suggested — the scaled router-weight copy now has a fixed address for the life of a capture, and the T>1 fallback still sees unscaled weights. Plus 4616f440/ed1c0fe6: explicit current-device in the indexed peer helper/retire (found while validating on the 2×R9700). Keyed identity GREEN, natural-EOS identity GREEN.

2. The T=2..63 default-flip measurement (2×R9700, isolated serve, PREFIX_CACHE=0, profile off, 4-process A(idx=63)→B(1)→B(1)→A(63), 5 bursts/leg, 32 tok/request API-usage denominator, independently reviewed raw):

T	indexed T≤63 vs host-gather	verdict
2	+14.0% (1.1403× median, both order pairings favor A)	GREEN
8	+2–4% directional	held (within dispersion — reported honestly, not promoted)
63	+24.6% (1.2455× median, both pairings favor A)	GREEN

So the flip pays at both ends of the range and the win grows with T; the shape question is answered — it's not a T=63-only artifact. rocprof arm-verification confirmed each arm dispatches the distinct kernel path before timing was accepted. Numbers + protocol are in docs/BENCHMARKS.md on the branch.

[localai-org-maint-bot]

This branch is currently conflicting with main, and its latest commit merges main into the feature branch. Please rebase the feature commits onto current main instead, then rerun the ROCm-specific evidence; the current conflict prevents a meaningful merge review.

[bakon11]

Conflict resolved — merged current main into the branch (dd19c8e3), matching the merge style already on this branch rather than rewriting your rebase. Sole conflict was docs/USAGE.md: resolved as a union — this PR's VT_GEMMA4_DECODE_INDEXED_MAX_T paragraph kept, and the Prefill peer (#839) sentence that landed via #1047 preserved verbatim. Dropped only main's stale "No new env." clause, which stops being true once this PR documents a new env knob. Payload unchanged (gemma4_indexed_gate.h, gemma4_moe.cpp); 0 behind main.

[bakon11]

Rebased as asked — 9f0ace5a, linear, zero merge commits, 0 behind main. My earlier merge-into-branch is gone; sorry for the noise, your convention is clearly the right one here.

What I did: replayed the 11 feature commits onto current main (your own land #838 commit keeps its authorship). Two conflicts, both additive unions:

• tests/CMakeLists.txt — kept both registrations: test_gemma4_moe_device_arm_guard (from main) and test_gemma4_indexed_max_t (this PR).
• docs/USAGE.md — kept this PR's VT_GEMMA4_DECODE_INDEXED_MAX_T paragraph and preserved the Prefill peer (#839) sentence that landed via fix(rocm/gemma4): #839 retirement-safe prefill peer GeGLU Launch/Finish #1047 verbatim; dropped only the stale "No new env." clause, which stops being true once this PR documents a knob.

Verified on the new head: gemma4_indexed_gate.h + test_gemma4_indexed_max_t.cpp present, the TLS-stable T=1 rw_idx for hipGraph fix you asked for is intact, docs carry both sentences.

ROCm evidence re-run is on our 2×R9700 queue and I'll post it here; the previously reported numbers (indexed T≤63 vs host-gather: +14.0% @t=2, +24.6% @t=63, T=8 within dispersion) were measured on the same payload.