feat(gemma4/rocm): FP8 resident MoE + SharedK-WMMA (split from #228)

[bakon11]

Summary

ROCm / Gemma4 half of the #228 split. #228 is closed — land this instead.

Related	Role
#234	ROCm V1 sampler + parallel RandomSample
#227	KV fail-fast
#316	SSE keepalives (timed wait) — not here
#228	Original megapr — closed, superseded

In

Dual-GPU FP8 resident MoE, peer mix, SharedK-WMMA prefill, decode KV-split, fused_ops seam, ENVIRONMENT VT_GEMMA4_* / VT_ATTN_*.

Out

SSE/serving (#316), rocm_sample (#234), support_static_graph_mode flip.

Review fixes in this tip

• HIPBLAS_V2 restored
• No platform graph-mode flip
• Host #if VT_ROCWMMA_OK around SharedK-WMMA launches

Honest follow-ups

HIP guards in models/gemma4.cpp, MoE policy guards, peer weight_id/hscale, tests/spec. Large diff (~5k prod lines) — may need another slice or maintainer path like #154 if review asks.

Lab: dual R9700 gfx1201 / start-gemma4-fp8-8010.sh.

Land order

1. feat(rocm): V1 sampling ops for Hermes EngineCore #234 → 2. fix(server): fail-fast unschedulable KV waits #227 → 3. fix(server): SSE keepalives via collector timed wait #316 → 4. This PR (feat(gemma4/rocm): FP8 resident MoE + SharedK-WMMA (split from #228) #317)

[bakon11]

CI hygiene 2026-08-11: rebased onto origin/main (60e71a0). Local gates green (doc-checkpoint, env-doc, agent-record, pr-size, device-leakage). Land order remains #234 → #227 → #316 → #317.

[bakon11]

Rebased onto origin/main c70f42b9 (main moved +2 during CI). Local gates green. Land order unchanged: #234 → #227 → #316 → #317. PR bodies refreshed (RandomSample note on #234; #228 closed in land path).

[localai-org-maint-bot]

@bakon11 This is not reviewable/mergeable with the PR body itself listing tests/spec as honest follow-ups. The ~5k-line production change needs the project-required spike/spec, focused automated parity/regression coverage for the new resident MoE/peer/SharedK/decode paths, and the same-change STATUS/BENCHMARKS/matrix/state evidence anchors required by AGENTS.md. Please land those in this PR (or split the stack further) before asking for merge.

[bakon11]

@localai-org-maint-bot Addressed the mergeability blockers you named:

1. Spike/spec landed: .agents/specs/gemma4-rocm-fp8-moe.md (scope, gates, residuals, merge criteria).
2. Focused automated coverage: new test_gemma4_rocm_fp8_seams (CPU) — portable vt::ExpertGeGLUFp8* / FP8 channel symbols link, recipe env knobs parse inertly. No AMD CI runners exist; ROCm kernel numerical parity remains lab evidence (dual R9700) as documented in the spec.
3. FEATURES row now anchors the spec + seam test (AGENTS public-doc surface).
4. PR body no longer lists tests/spec as vague follow-ups — residuals are named.

Local: test_gemma4_rocm_fp8_seams green. Device-leakage / env-doc / trailer checks OK on tip.

[localai-bot]

Reviewed while landing the external-PR queue. Holding this one — not rejecting it; the ROCm half looks good and the levers are honestly gated. The blocker is that it is not ROCm-only.

Why it is held

672 of the added lines are in src/vllm/model_executor/models/gemma4.cpp and 974 in gemma4_moe.cpp, and exactly one of them mentions ROCm. Gemma-4 is a gate model, so those lines are on the CUDA path too, unguarded.

The specific change I cannot sign off without a run is ForwardGemma4Layers: the layer loop moves out of the monolithic forward and its intermediates become a static thread_local LayerTls, re-emplaced whenever (dev, T, H, I, ple) changes. The commit says "Behavior is intentionally identical to the previous monolithic loop", and I believe that is the intent — but "intentionally identical" is a claim about a gate model's forward, and this repo settles those with a token-exact run, not a reading.

What would close it

1. Gemma-4 token-exact gate on GB10, on this branch, against the committed golden. Byte-identical tokens is the whole claim; if the refactor is really behavior-preserving this is cheap and it closes the question permanently.
2. The TLS-vs-graph-capture question, explicitly. down_out is caller-owned (layer TLS) so hipGraph capture keeps stable pointers is right about stability, but the buffers are destroyed and reallocated when the shape key changes. A graph captured at shape A holds device pointers into buffer set A; running shape B frees them; replaying A then reads freed device memory. This repo has been bitten by exactly this class before (capture baking addresses of function-local upload temporaries). Either show that a shape change invalidates every captured graph, or say plainly that capture is off on this path. VLLM_CPP_GEMMA4_DECODE_GRAPH defaults off, which helps, but VT_ROCM_GRAPH defaults on.
3. A word on static thread_local lifetime — one full set of [T,H]/[T,I] bf16 buffers per thread, never released. Worth a sentence on the expected thread count.

What I checked and am not worried about

• The new levers in gemma4_moe.cpp are env-gated and default-off (VT_GEMMA4_EXPERT_VRAM_MB unset = off, VT_GEMMA4_BATCH_EXPERTS off, eviction opt-in).
• The ones that default on (VT_ATTN_PREFILL_FLASH_SHAREDK, VT_ATTN_PREFILL_SHAREDK_WMMA, VT_GEMMA4_PREFILL_PEER_ACT, VT_ROCM_GRAPH) are all ROCm-path, so CUDA is untouched by them.
• Backend::DeviceMemoryInfo is a clean additive seam with a false default — model code sizing an LRU without vendor headers is the right shape.

Splitting the unconditional gemma4.cpp restructure into its own PR would also unblock the rest of this immediately, if a GB10 window is hard to get.

Reviewed with Claude Code while landing #309/#316/#326/#328/#334-#348. Nothing here is a request to change the ROCm kernels.

[bakon11]

Rebuilt on current main — addresses hold (CUDA forward restructure)

Tip: a5bd7ad0 on origin/main c43092f9 (ancestor YES, behind=0).

What changed vs prior tip

Prior hold called out unguarded gemma4.cpp / gemma4_moe.cpp forward work and ForwardGemma4Layers + layer TLS + Gemma4DecodeGraph (CUDA path / graph lifetime).

This rebuild drops all of that from the PR:

Deferred (lab-only / later PR)	Still in #317
ForwardGemma4Layers extract	ROCm FP8 ExpertGeGLU + channel GEMV
Gemma4DecodeGraph / decode hipGraph	SharedK + SharedK-WMMA prefill
MLP∥MoE side-stream wiring	vt::fused_ops portable seams
Any CUDA-path layer-loop rewrite	Expert LRU / host cache / prewarm
GB10 token-exact requirement for this tip	DeviceMemoryInfo backend seam
	CPU seam test test_gemma4_rocm_fp8_seams
	Spec + ENVIRONMENT/FEATURES/USAGE

gemma4.cpp / gemma4.h are unchanged vs main in this tip. Registry only adds ExpertGeGLU prewarm after resident upload (no graph hook).

Bot checklist

1. Token-exact GB10 — N/A for this tip (no forward restructure). Happy to land graph/extract later with golden.
2. TLS vs graph capture — N/A (no Gemma4 decode graph in PR).
3. thread_local lifetime — only existing main layer TLS remains; no new capture driver.

Local gates (green)

• check-device-leakage
• check-env-doc
• check-doc-checkpoint
• check-agent-record
• check-pr-size (feat(gemma4/rocm): FP8 resident MoE + SharedK-WMMA (split from #228) #317)

Sanitize-cpu ASan/TSan ambient noise may still red; not specific to this payload.

Please re-review when CI settles on a5bd7ad0.

[localai-bot]

Hardware-checked on dgx (GB10) — this PR is clear of the concern I raised, and the check found an unrelated bug on main instead.

Two things, both useful to you:

1. The CUDA-forward concern is resolved. Your a5bd7ad0 ("no CUDA forward restructure") does exactly what was needed — ForwardGemma4Layers is gone and gemma4.cpp is untouched. I verified it at the object level rather than by reading the diff: gemma4.cpp.o compiles byte-identical between main and main+this PR. That also retires the TLS-vs-graph-capture question, since nothing is moved into shape-keyed thread-local storage any more. Thanks for turning that around.

2. The Gemma-4 SACRED gate was already RED on main, and it is not yours. Running it against your branch produced:

vt::GeluMulSeparate: ROCm-only fast path in this build

…and the baseline arm failed identically. Root cause was 0c2827c1 (perf(gemma4): dual-GPU FP8 resident, #154), which replaced the portable vt::GeluAndMul at gemma4.cpp's per-layer-embedding call site with the ROCm-only vt::GeluMulSeparate — from the shared forward, so Gemma-4 aborted on layer one on CUDA/CPU/Metal/Vulkan, in the server as well as the test. Filed as #377 and now fixed on main: the gate is 32/32 token-exact vs the vLLM golden, green on CUDA for the first time since 2026-08-08.

Please rebase — the gate can actually run for you now.

What this does and does not tell us about #317. It says the PR does not regress the shipped Gemma-4 CUDA path. It says nothing about the FP8 resident MoE work itself, and cannot: unsloth/gemma-4-E4B-it is dense bf16 with enable_moe_block: False, and all nine Gemma checkpoints cached on that box are the same, so RunGemma4Moe, EnsureGemma4Fp8*, ExpertGeGLU* and the guarded registry prewarm are unreachable there. GB10 is also the wrong silicon for a hipBLASLt FP8 path. Validating the feature needs your ROCm box with a MoE checkpoint — that evidence has to come from you, and a note in the PR body saying which board and which checkpoint produced it would be enough.

The remaining review comment on the earlier head still applies to what is left: gemma4_moe.cpp is shared code even though the paths are guarded, so a line on what a CUDA build does with these changes (nothing, as far as I can tell — every new lever is either env-gated off or behind moe.enabled) would close it out.

Two smaller things while I was in there: VT_GEMMA4_PREFILL_PEER_ACT, VT_ATTN_PREFILL_FLASH_SHAREDK, VT_ATTN_PREFILL_SHAREDK_WMMA and VT_ROCM_GRAPH default on — fine, since they are all ROCm-path, but worth being explicit in the PR body that they change default behavior on ROCm boards. And Backend::DeviceMemoryInfo is a clean additive seam with a false default; no objection to that at all.

Checked with Claude Code. dgx GPU work was serialized behind flock $HOME/gpu.lock alongside another session's campaign; nothing of theirs was touched.

[bakon11]

Rebased on current main (post-#377 gate fix)

Tip: d1c4a656 on current origin/main (clean rebase of prior a5bd7ad0; behind=0).

Thanks for the GB10 hardware check — confirmed:

1. CUDA forward concern closed (no ForwardGemma4Layers / no new shape-keyed TLS; gemma4.cpp untouched vs main).
2. SACRED red was main’s GeluMulSeparate (Gemma-4 text generation throws on every CUDA build: shared forward calls the ROCm-only vt::GeluMulSeparate (SACRED gate RED on main since 0c2827c1) #377) — not this PR. Rebased so the gate can run on the fixed baseline.

What CUDA / non-ROCm builds do with this tip

• gemma4.cpp / gemma4.h: identical to main → no CUDA path change.
• gemma4_moe.cpp: shared file, but every new lever is env-gated default-off or behind moe.enabled / ROCm ExpertGeGLU dispatch. On CUDA/CPU/Metal/Vulkan without the ROCm FP8 MoE path: no new default behavior; resident FP8 / ExpertGeGLU / peer-act do not activate.
• GB10 dense bf16 Gemma (enable_moe_block: false): MoE/FP8 ExpertGeGLU / registry prewarm remain unreachable — expected. Feature validation is ROCm + MoE FP8 ckpt only.

ROCm lab evidence (this box — not GB10)

Board	2× AMD Radeon AI PRO R9700 (gfx1201 / RDNA4), ROCm 7.x
Checkpoint	gemma-4-26B-A4B-it-fp8 (HF FP8 MoE)
Serve	~/llms/scripts/vllm/start-gemma4-fp8-8010.sh → :8010
Measured class	Prefill eng ~2k @11k (SharedK-WMMA + PEER_ACT + GEMM_M=2048); short decode ~45–52 t/s; long pad quality Paris OK under BF16 KV default

Default-on ROCm knobs (explicit)

These change ROCm default behavior when unset (documented in ENVIRONMENT); non-ROCm unaffected:

• VT_GEMMA4_PREFILL_PEER_ACT default on
• VT_ATTN_PREFILL_FLASH_SHAREDK default on
• VT_ATTN_PREFILL_SHAREDK_WMMA default on
• VT_ROCM_GRAPH default on (no Gemma4 decode-graph driver in this PR)

KV cache dtype (separate from weight FP8)

Serve selection is env today (llama.cpp-style intent, different surface):

• default bf16 KV with every weight quant
• VT_KV_CACHE_DTYPE=fp8|bf16|… or VT_KV_CACHE_FP8=1 opt-in
  CLI --kv-cache-dtype twin not in this PR (follow-up).

Local gates on d1c4a656

device-leakage / env-doc / doc-checkpoint / agent-record / pr-size / commit-trailers — green.

Please re-run CI + SACRED on the new tip when ready.

[localai-bot]

Rebased for you and re-gated on hardware — the CUDA concern is now measured and closed.

I merged current main into this branch and pushed it here (d1c4a656..75ed6a17), so no action needed on your side. main now carries the #377 fix, which is what was making the Gemma-4 gate throw for this branch and for bare main.

Gate result — dgx GB10, both arms in one flock window

arm	tree	result
gmain	main	gemma4-E4B STRICT: 32/32 token-exact vs vLLM 0.25.0 golden
g317	main + this PR	32/32 token-exact

Same box, same checkpoint, same minute. This PR does not regress the shipped Gemma-4 CUDA path — that is the concern I raised originally, and it is now settled by measurement rather than by reading the diff. gemma4.cpp also remains byte-identical to main after the merge.

One fix carried in

The pre-push hook refused my merge, correctly: this branch's docs/FEATURES.md row was 254 chars against the 220-char cell budget that check-public-doc-tables enforces in CI, so the PR would have gone red there. Trimmed to 219 keeping every fact — dual-GPU FP8 resident experts, SharedK-WMMA prefill on RDNA4, the deferred decode-graph/forward extract, the env prefixes, the CPU seam test — with the spec link carrying the detail. Commit is on the branch.

What is still owed, and it is not much

The gate above proves non-regression. It cannot prove the feature, and no run on my side ever will: unsloth/gemma-4-E4B-it is dense bf16 with enable_moe_block: False, all nine Gemma checkpoints cached on that box are the same, and the FP8 path is hipBLASLt regardless — so RunGemma4Moe, EnsureGemma4Fp8*, ExpertGeGLU* and the guarded registry prewarm are simply unreachable on GB10.

So the last thing this PR needs is from your ROCm box:

1. Which board and which MoE checkpoint you ran, and the result. A short paragraph in the PR body is enough — it does not need to be a formal gate.
2. A line confirming what a CUDA build does with these changes. From my reading it is nothing (every new lever is env-gated off or behind moe.enabled), and the gate above is consistent with that, but you know the code better.

Worth stating explicitly in the body too: VT_GEMMA4_PREFILL_PEER_ACT, VT_ATTN_PREFILL_FLASH_SHAREDK, VT_ATTN_PREFILL_SHAREDK_WMMA and VT_ROCM_GRAPH default on. That is fine — they are all ROCm-path — but they change default behavior on ROCm boards and a reader should not have to diff docs/ENVIRONMENT.md to find out.

Backend::DeviceMemoryInfo is a clean additive seam with a false default; no objection there at all.

Thanks for turning the CUDA forward restructure around so quickly — that is what made this checkable.

Rebased and gated with Claude Code. dgx GPU work serialized behind flock $HOME/gpu.lock alongside another session's campaign; nothing of theirs was touched.

[localai-bot]

Landed on main as 69d0ac5a via row/GEMMA4-ROCM-FP8-SPLIT, merged --no-ff so your commits keep their authorship. Thanks — this took several rounds and the rebuild you did was the right call.

The hold is closed by measurement, not by reading the diff. Both arms gated on dgx GB10 in one flock window: gmain (main) and g317 (main + this PR) each returned gemma4-E4B STRICT 32/32 token-exact vs the vLLM 0.25.0 golden, and gemma4.cpp is byte-identical to main.

Your cuda-fat-build red was never yours. The gencode audit failed on seven files, every one of them Marlin, while this branch touches ROCm and gemma4_moe.cpp and no Marlin gencode at all. That was the drift #407 fixed, and this branch predated it — the same stale-base pattern that also explained #383 and #393 tonight.

Accepted with a residual I want on the record, because it is a product decision rather than a gate result. This flips four knobs default-ON for ROCm — VT_GEMMA4_PREFILL_PEER_ACT, VT_ATTN_PREFILL_FLASH_SHAREDK, VT_ATTN_PREFILL_SHAREDK_WMMA, VT_ROCM_GRAPH. Non-ROCm builds are unaffected and each knob is individually reversible by env, but there is no AMD CI runner, so those defaults rest entirely on your lab evidence (2× R9700 gfx1201, ROCm 7.x, gemma-4-26B-A4B-it-fp8). And feature correctness stays unprovable on GB10, honestly so: every cached Gemma checkpoint there is dense bf16 with enable_moe_block: false, so RunGemma4Moe / ExpertGeGLU* are unreachable. The GB10 gate proves non-regression, never the feature.

Gate: scripts/agent-preflight.sh rc=0 on the merged tree, all record checkers green. CI is queue-blocked (#274), so per AGENTS.md the operator gate run is the authority — worth saying plainly rather than implying this went green in CI.