feat(gemma4/rocm): FP8 resident MoE + SharedK-WMMA prefill (RDNA4) - #228
feat(gemma4/rocm): FP8 resident MoE + SharedK-WMMA prefill (RDNA4)#228bakon11 wants to merge 5 commits into
Conversation
|
Update after live Hermes proof (RDNA4 / dual R9700): Full agent path works end-to-end (SOUL ~38k prefill → streamed decode → real reply). The blocking EngineCore fatals after prefill were missing ROCm V1 sampler ops ( That fix is split out as a small tip-of-main PR for fast merge: This #228 remains the larger FP8 resident / fused decode / GPU prefill-MoE package. Prefer landing #234 first, then rebase this head. |
TESTERS WANTED — clear split: FP8 (lab) vs BF16/FP16 MoE (community)What this PR isFP8-first Gemma-4 MoE on ROCm: native FP8 resident experts, fused FP8 decode helpers, GPU prefill-batch MoE, etc. Lab-proven on 2× R9700 gfx1201 with Hermes-scale prompts (slow decode at 38k, but correct). What it is notA claim that every kernel here is the BF16 path. FP8 weight GEMM / FP8 resident upload do not apply to BF16 weights. What does still matter for BF16/FP16 MoE testersIf you run BF16 or FP16 MoE on ROCm, you still want:
Optional on this PR after those: any shared MoE control flow you can exercise in BF16 (group-by-expert prefill batch, dual-GPU placement) — report whether BF16 hits the same batch path or stays on serial expert loops. Please do not treat as MoE validation
Ideal tester hardware
Report template
Prefer merge order: #234 → #227 → rebase this FP8 MoE package. Thanks — FP8 is the consumer-RDNA4 daily driver; BF16/FP16 MoE needs wider silicon than this lab can soak alone. |
localai-org-maint-bot
left a comment
There was a problem hiding this comment.
The lab evidence is useful, but this introduces a large FP8 resident/prefill execution surface without automated kernel/model parity tests; the only added tests cover HF config parsing. Please add focused parity coverage for resident expert placement, fused decode, and grouped prefill (including multi-device behavior or a testable dispatch seam), then rebase after #234/#227 as proposed and clear the current repository gates. The current single-hardware smoke is not sufficient protection for this execution-path change.
|
Reviewed. The ROCm work looks solid and it builds clean on CUDA — but I'm asking for this to be split before it lands, and the reason is reviewability rather than anything wrong with the code. Verified so far (GB10 sm_121a, rebased on current
|
…stant Three open PRs each grew their own workaround for the same defect, with three different answers: #227 clamps to 16, #234 and #228 to 4096. All three would conflict with each other, and all three truncate. This fixes it once, in the one place that owns the translation. Hermes and some OpenAI clients send `max_tokens: -1` for "no client-side limit". `SamplingParams::PostInit` rejects anything below 1, which is what made a clamp look necessary — but a constant is the wrong answer, because "unlimited" already has a correct representation here: UNSET. The engine then generates to `max_model_len - seq_len` (`v1/engine/input_processor.cpp:175-180`, a 1:1 mirror of vllm `input_processor.py:317-321`), which is exactly what the client asked for. With a constant, a client asking for unlimited against a 32k model gets 4096 (or 16) tokens back with `finish_reason=length`, indistinguishable from a real limit it set itself. It also hardcodes a policy number into the protocol layer, where `VT_SERVER_MAX_NEW_TOKENS` is already the tunable. Both request paths are covered: `/v1/completions` (`max_tokens`) and `/v1/chat/completions` (`max_completion_tokens` preferred over `max_tokens`, matching chat_completion/serving.py:299). A positive value is honoured unchanged, and a serving-resolved default still wins over an unset request value. RED before, GREEN after, same test and command, on dgx with the new test built against UNPATCHED main: before: test case THREW exception: max_tokens must be at least 1, got -1 29 cases | 28 passed | 1 failed; 175 assertions | 2 failed after: 29 cases | 29 passed | 0 failed; 179 assertions | 0 failed `test_openai_protocol` gains six subcases pinning the semantics: -1 and 0 unset on both paths, a positive value unchanged, and -1 yielding to a serving default when one is supplied. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Claude:claude-opus-5 [ClaudeCode]
…ans UNSET Three open PRs each grew a workaround for the same defect with three different answers: #227 clamps to 16, #234 and #228 to 4096. They would conflict with each other and all three truncate. Fixed once, in the layer that owns the translation, so those PRs can drop it and rebase. Hermes and some OpenAI clients send max_tokens=-1 for "no client-side limit"; PostInit rejects <1, which is what made a clamp look necessary. But "unlimited" already has a correct representation: UNSET, which routes to max_model_len - seq_len (input_processor.cpp:175-180, mirroring vllm input_processor.py:317-321). A constant silently truncates exactly the long-context request that asked to be left unlimited. RED before / GREEN after on dgx, the new test built against UNPATCHED main: before: THREW 'max_tokens must be at least 1, got -1'; 28/29 cases, 2 failed after: 29/29 cases, 179/179 assertions Covers both request paths and documents the semantics in docs/USAGE.md. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Claude:claude-opus-5 [ClaudeCode]
|
Heads-up: the A non-positive RED before / GREEN after on dgx, the new test built against unpatched main:
Everything else in your PR stands — this is only about the shared hunk. |
|
Reviewed at That was the open question from my last pass, so it is now answered and off the table. I can't review or run the ROCm kernels (no AMD GPU on this side), but the ~1600 lines under The
|
fix(parity): read generation_config.json eos ids We only ever read config.json. `generation_config_eos_ids_` was named for generation_config.json but populated from config.json's `eos_token_id`, and nothing in the LLM path opened generation_config.json at all — only Parakeet's loader did. Upstream reads both (try_get_generation_config for the default --generation-config auto, then sampling_params.py:645-655 merges the ids into stop_token_ids), and the two files routinely disagree: Gemma-4-26B ships config.json [1, 106] against generation_config.json [1, 106, 50]. Id 50 never stopped a generation, so chat models ran to the length cap. Found via #228, which hit it on Gemma-4 but fixed it in raw["eos_token_id"] — the field that supplies the PRIMARY eos, which our existing 1:1 port already gates on ignore_eos. These are SECONDARY stop ids. This carries them in their own field, leaves raw untouched, and unions by append so the primary keeps its position. Red first: both suites fail 2 assertions each without the fix. test_qwen3_paged_engine 184/184 test_qwen27_paged_engine 235/235 SACRED test_qwen36_paged_engine 315/315 SACRED test_hf_config 194/194 test_input_processor 64/64 Both SACRED gates matter here: honoring new stop ids could have shortened gate-model generations and drifted the goldens. They did not move. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Claude:claude-opus-5 [ClaudeCode]
|
Follow-up: the corrected eos fix is now on main as |
…config.json `InputProcessor::generation_config_eos_ids_` is named for generation_config.json but was populated entirely from config.json's `eos_token_id`. Nothing in the LLM path read generation_config.json at all (only Parakeet's loader did), so upstream's "secondary eos ids" were, in our port, just a second copy of the primary list. Upstream loads both. ModelConfig.try_get_generation_config (vllm/config/ model.py) reads generation_config.json whenever --generation-config is "auto" (the default) or "vllm", and SamplingParams.update_from_generation_config (vllm/sampling_params.py:645-655) merges its ids into stop_token_ids. The two files routinely disagree: Gemma-4-26B ships config.json eos: [1, 106] generation_config.json eos: [1, 106, 50] so id 50 — a turn-level stop — never terminated a generation for us. The symptom is a chat model that runs past its stop token to the length cap. Found via mudler#228, which hit this on Gemma-4 and fixed it by merging the ids into `cfg.raw["eos_token_id"]`. That lands them at the wrong layer: raw's `eos_token_id` supplies the PRIMARY eos id, which our existing 1:1 port of update_from_generation_config already gates on `ignore_eos`, while these are SECONDARY stop ids. Rewriting raw would make them unconditional, and — since the primary is `generation_config_eos_ids_.front()` on the list path — a re-sorted union could silently displace which id is primary. So this carries the field separately (`HfConfig::generation_config_eos_ids`), leaves `raw` untouched, and unions by append so the primary keeps its position. The existing ignore_eos gates then apply unchanged: the ids reach stop_token_ids only when ignore_eos is false, and all_stop_token_ids (min_tokens masking) regardless — which is exactly sampling_params.py:653. Red first: both new suites fail 2 assertions each against the current implementation and pass with it. A malformed or absent generation_config.json is a silent no-op, mirroring try_get_generation_config returning {} rather than raising. test_hf_config 194/194 (+5 subcases) test_input_processor 64/64 (+2 subcases) FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Claude:claude-opus-5 [ClaudeCode]
Status (lab + PR hygiene)Land order: #234 (ROCm V1 sampler, rebuilt tip-of-main, clamp dropped) → #227 (KV fail-fast, rebuilt tip-of-main, prefix probe + clamp dropped) → then this #228 rebase. Maintainer notes absorbed for the next #228 push:
Lab prefill (not blocking this PR thread): SharedK-WMMA + |
Rebuild mudler#228 onto current main. Drop workarounds that already landed: - max_tokens clamp / -1→4096 (main 2d2bdd0 SERVE-MAXTOKENS-UNSET) - generation_config eos merge (main fb2dfac / mudler#262) Keep operator VT_SERVER_MAX_NEW_TOKENS for *positive* oversized requests only. Non-positive remains UNSET → max_model_len - seq_len. Includes lab wins: SharedK-WMMA, PREFILL_GEMM_M=2048, PEER_ACT default ON, SSE keepalives + deferred add_request for long Hermes prefill, prefill progress tok_s, rocm_sample.hip for V1 sampler ops. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Hermes:grok-4.5 [Hermes]
9a009eb to
e6c8c4e
Compare
FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Hermes:grok-4.5 [Hermes]
Rebased on tip-of-main — clamp + eos droppedPer maintainer note (
Head: Still in this PR: FP8 resident MoE, SharedK-WMMA prefill, PEER_ACT default ON, GEMM_M=2048, SSE keepalives + deferred Land order preference unchanged: #234 sampler → #227 KV fail-fast → this #228 (sampler TU also present here so either order can work after a final rebase). Local: |
Tip CMake had dropped main's DSpark sources and HIP -O1 floor (mudler#132). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Hermes:grok-4.5 [Hermes]
|
Follow-up on the rebuild tip: CMakeLists restored from main (DSpark + HIP -O1 floor kept) with only |
|
Thanks for this — there's real engineering in here, and the FP8 resident-expert design plus the SharedK/WMMA prefill is exactly the direction the ROCm lane needs. I reviewed it in depth rather than landing it, and I'm asking for a split instead of a merge. Reasons below, with the evidence for each. I want to be upfront about a limit that shapes everything: no CI job compiles ROCm and there is no AMD GPU on the maintainer box. So the three green build checks prove only that x86 CPU still builds — roughly 4,400 lines of HIP in this PR have never been compiled by anyone but you. That's why the ROCm half needs your hardware to land, and why I'm not going to "fix it up and merge" blind. 1. The three red checks have two root causes, and neither is a race
2. The SSE rewrite regresses streaming for every model, on every backendSame-binary A/B, identical tree, only the six serving files swapped:
Three distinct defects:
The keepalive/SSE idea is legitimate and worth having — it just needs its own PR, against the streaming tests it currently breaks. 3. Two changes reach much further than "gemma4/ROCm"
4. ROCm-tuned policy applied to CPU and CUDA, unguarded
5. Kernel-level findings I could not run, only readThese need your gfx1201 box to confirm or refute:
Credit where it's due, verified clean: FP8 scale semantics are correct throughout (per-output-channel, never per-K, gate/up halves get separate scales, fp32 accumulation); the e4m3fn bit decode matches OCP apart from NaN mapping to 0.0; 6. ProcessNo Worth noting What I'd suggestSplit it in two. The A/B shows that reverting just the six serving files makes this PR green on everything CI can see:
I'd rather have this land properly than land fast, and the ROCm half genuinely needs your hardware to gate — that's the one thing I can't do from here. Happy to review either half quickly as you split them, and to pre-check anything CPU-side. |
Status — hold merge; split next (back on track)Land order in progress:
Lab decode WIP (BPE heap, structural-tag default-off, FP8 HW cvt, KV-split/slide, etc.) lives on bakon11 |
Superseded by #317This megapr head is behind main and carries mixed serve/docs/history noise. Land path (split):
Closing #228 in favor of #317 to avoid dual-tracking. Lab decode stack stays on bakon11 WIP branches, not this PR. |
ROCm/Gemma4 half of the mudler#228 split. No OpenAI SSE/serving changes (see mudler#316). No V1 sampler registration (see mudler#234). Includes dual-GPU FP8 resident experts, peer mix helpers, SharedK-WMMA prefill, decode KV-split path, fused_ops seam. Review fixes absorbed in this tip: - Restore HIPBLAS_V2 (ROCm 6.x mudler#201/mudler#132) — was clobbered in megapr - Do NOT flip platforms/rocm support_static_graph_mode (stays main) - Host-side #if VT_ROCWMMA_OK around SharedK-WMMA launches (no empty kernel launch on non-gfx1200/1201 slices) Still follow-ups (called out in PR body): HIP guards in models/, MoE policy FreeBytes/FP8_NATIVE defaults, peer weight_id / hscale, tests/spec. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Hermes:grok-4.5 [Hermes]
ROCm/Gemma4 half of the mudler#228 split. No OpenAI SSE/serving changes (see mudler#316). No V1 sampler registration (see mudler#234). Includes dual-GPU FP8 resident experts, peer mix helpers, SharedK-WMMA prefill, decode KV-split path, fused_ops seam. Review fixes absorbed in this tip: - Restore HIPBLAS_V2 (ROCm 6.x mudler#201/mudler#132) — was clobbered in megapr - Do NOT flip platforms/rocm support_static_graph_mode (stays main) - Host-side #if VT_ROCWMMA_OK around SharedK-WMMA launches (no empty kernel launch on non-gfx1200/1201 slices) Still follow-ups (called out in PR body): HIP guards in models/, MoE policy FreeBytes/FP8_NATIVE defaults, peer weight_id / hscale, tests/spec. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Hermes:grok-4.5 [Hermes]
ROCm/Gemma4 half of the mudler#228 split. No OpenAI SSE/serving changes (see mudler#316). No V1 sampler registration (see mudler#234). Includes dual-GPU FP8 resident experts, peer mix helpers, SharedK-WMMA prefill, decode KV-split path, fused_ops seam. Review fixes absorbed in this tip: - Restore HIPBLAS_V2 (ROCm 6.x mudler#201/mudler#132) — was clobbered in megapr - Do NOT flip platforms/rocm support_static_graph_mode (stays main) - Host-side #if VT_ROCWMMA_OK around SharedK-WMMA launches (no empty kernel launch on non-gfx1200/1201 slices) Still follow-ups (called out in PR body): HIP guards in models/, MoE policy FreeBytes/FP8_NATIVE defaults, peer weight_id / hscale, tests/spec. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Hermes:grok-4.5 [Hermes]
ROCm/Gemma4 half of the mudler#228 split. No OpenAI SSE/serving changes (see mudler#316). No V1 sampler registration (see mudler#234). Includes dual-GPU FP8 resident experts, peer mix helpers, SharedK-WMMA prefill, decode KV-split path, fused_ops seam. Review fixes absorbed in this tip: - Restore HIPBLAS_V2 (ROCm 6.x mudler#201/mudler#132) — was clobbered in megapr - Do NOT flip platforms/rocm support_static_graph_mode (stays main) - Host-side #if VT_ROCWMMA_OK around SharedK-WMMA launches (no empty kernel launch on non-gfx1200/1201 slices) Still follow-ups (called out in PR body): HIP guards in models/, MoE policy FreeBytes/FP8_NATIVE defaults, peer weight_id / hscale, tests/spec. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: Hermes:grok-4.5 [Hermes]
…dK-WMMA (#41) Lands bakon11's work from PR #317, split out of #228. Merged --no-ff so the contributor commits keep their authorship rather than being re-authored to whoever pressed the button. FP8 resident MoE experts and a SharedK-WMMA prefill path for Gemma-4 on RDNA4 ROCm, plus a CPU seam test that pins the portable vt::ExpertGeGLUFp8* symbols link and the recipe env knobs parse inertly. The concern this PR was originally held on is CLOSED BY MEASUREMENT, not by reading the diff. It was held because ~1,600 added lines sat in gemma4.cpp and gemma4_moe.cpp and Gemma-4 is a gate model. bakon11 rebuilt it to drop the CUDA forward restructure entirely, and both arms were then 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. gemma4.cpp is byte-identical to main. Its cuda-fat-build red was never its own: 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 the branch predated it -- the same stale-base class that also explained #383 and #393 tonight. ACCEPTED WITH A NAMED RESIDUAL, recorded because it is a product decision and not 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 and VT_ROCM_GRAPH. Non-ROCm builds are unaffected -- every new lever is env-gated off or behind moe.enabled / ROCm ExpertGeGLU dispatch -- and each knob is individually reversible by env. But there is NO AMD CI runner, so those defaults rest entirely on the contributor's lab evidence (2x AMD Radeon AI PRO R9700 gfx1201 / RDNA4, ROCm 7.x, gemma-4-26B-A4B-it-fp8; prefill ~2k @11k, short decode ~45-52 t/s). Feature correctness is unprovable on GB10 and honestly so: all nine cached Gemma checkpoints there are dense bf16 with enable_moe_block false, so RunGemma4Moe, EnsureGemma4Fp8* and ExpertGeGLU* are simply unreachable. The GB10 gate proves NON-REGRESSION, never the feature. Gate: scripts/agent-preflight.sh rc=0 on the merged tree; check-agent-record, check-doc-checkpoint, check-public-doc-tables, check-now-current, check-env-doc and check-commit-trailers all green. CI is queue-blocked (#274), so the operator's own gate run is the authority, as AGENTS.md prescribes. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code]
Summary
Gemma-4-26B MoE ROCm path for dual-GPU (gfx1201 lab):
config.json∪generation_config.json([1,50,106])vt::fused_opsseams — novt::rocmfrommodels/Lab (2× R9700, ROCm 7.2.4)
hipBLASLt W8A8 microbench still slower than BF16 GemmEx on gfx1201 for pure matmul — not default.
Size
Over the ~900 product-line ratchet (MoE + HIP kernels). Happy to split further if maintainers prefer (decode/resident vs prefill batch).
Supersedes
Supersedes the bulk of #205 content (rebased to current
mainas a single tip commit). Please close #205 in favor of this + the serve reliability PR.Test plan
check-device-leakagecheck-env-doc/ FEATURES / USAGE / ENVIRONMENT:8010dual FP8 resident smoke (Paris/READY/prefill rates)