perf(gemma4): dual-GPU FP8 resident without host OOM + peer mix

[bakon11]

Summary

Follow-up to merged #140 (ROCm gfx1201 + Gemma-4-26B MoE BF16/FP8).

This PR hardens dual-GPU FP8 expert residency and decode mix so the path fits ~30 GB host RAM and keeps weights on GPU:

1. Pack GPU0 first, spill to GPU1 (same-device fast path maximized)
2. Stream per-expert FP8→BF16→H2D — no permanent host BF16 cache (was OOM ~45 GiB on 30 G host)
3. 10 GiB compute headroom on GPU0 so dense/KV still allocate
4. Peer-copy off-device resident expert slices into compute scratch (no FP8 re-dequant)
5. Device MulScalar + Add for expert weighted sum (default; VT_GEMMA4_HOST_AXPY=1 fallback)

Lab evidence (2× R9700 gfx1201, ROCm 7.2.x)

Check	Result
Full resident upload	30/30 layers, 42.54 GiB experts
Packing	GPU0 ~14 layers, GPU1 ~16 layers
Peak VRAM	~28.6 GB
Gate	<bos>The capital of France is → **Paris** EXIT=0
Wall (CLI load+resident+12 tok)	~191–198 s after peer + device mix (was ~294 s pre-peer)

Host OOM root cause was resident upload calling cache-filling dequant for every expert every layer.

Usage

HIP_VISIBLE_DEVICES=0,1 \
VT_GEMMA4_RESIDENT_EXPERTS=1 \
VT_GEMMA4_RESIDENT_GPUS=2 \
./build-hip/examples/vllm-cli \
  --model /path/to/Firworks/gemma-4-26B-A4B-it-fp8 \
  --prompt '<bos>The capital of France is' --max-tokens 12 --temperature 0

Optional: VT_GEMMA4_HOST_AXPY=1, VT_GEMMA4_RESIDENT_MAX_LAYERS=N

Test plan

• HIP build gfx1201
• Full dual-GPU resident Paris (lab)
• CI (no AMD GPU expected)
• Load-once server multi-req tok/s (follow-up; avoid host thrash)

Notes

• Does not claim pip-vLLM oracle token-exact or production tok/s yet.
• CLI wall still dominated by load + 42 GiB H2D; warm decode needs server load-once benches.

[localai-bot]

Thanks for the follow-up to #140 — the dual-GPU FP8 resident direction is right, and the peer-mix idea is one we want. Two things to sort before this can land, one mechanical and one structural. They have the same fix.

1. The CPU, CUDA and Vulkan builds do not link

build-test-cpu and build-test-vulkan are failing for a real reason, not a flake:

gemma4_moe.cpp: undefined reference to `vllm::PinGemma4Fp8ExpertHostCache(...)`
gemma4_moe.cpp: undefined reference to `vllm::PeerCopyGemma4ExpertSlice(...)`
gemma4_moe.cpp: undefined reference to `vt::rocm::MatmulBTFp8ChannelRocm(...)`
gemma4_moe.cpp: undefined reference to `vt::rocm::ExpertGeGLUBf16TopKM1Rocm(...)`
gemma4.cpp:     undefined reference to `vt::rocm::RmsNormPlusAddRocm(...)`
gemma4.cpp:     undefined reference to `vt::rocm::DualRmsNormPlusResRocm(...)`
gemma4.cpp:     undefined reference to `vt::rocm::GeluMulSeparateRocm(...)`

These are declared in include/vllm/model_executor/models/gemma4_moe.h but defined only in src/vt/rocm/rocm_gemma4_experts.hip, and that file is added to the build only under the ROCm branch (CMakeLists.txt:1148). Every non-ROCm configuration compiles the calls and then has nothing to link them against.

2. Backend names in the device-agnostic layer

The deeper reason it breaks is that src/vllm/model_executor/models/ is the shared, device-agnostic layer, and this PR puts 11 vt::rocm:: call sites into it (4 in gemma4.cpp, 7 in gemma4_moe.cpp). Today main has exactly zero — the only vt::rocm:: reference outside src/vt/ is in src/vllm/platforms/rocm.cpp, which is the right place for it. This is what check-device-leakage guards, and it is a rule we hold ourselves to: I had to fix one of my own violations in #159 this week for the same reason.

(Do rerun CI after rebasing, by the way — the device-leakage failure you are seeing right now is partly our bug, not yours. It reports kcuda=1 from a DeviceType::kCUDA I left in runner.cpp. That is fixed on main in #159.)

The fix for both

Route these through the ops/backend seam rather than calling the ROCm entry points by name:

• Kernels (RmsNormPlusAddRocm, DualRmsNormPlusResRocm, GeluMulSeparateRocm, MatmulBTAlphaBetaRocm, MatmulBTFp8ChannelRocm, ExpertGeGLUBf16TopKM1Rocm): register them as OpId implementations for DeviceType::kROCm and call the portable vt:: entry point. The model file then reads the same on every backend, and CPU keeps its reference implementation.
• Policy/capability (PeerCopyGemma4ExpertSlice, PinGemma4Fp8ExpertHostCache): these are "can this backend do peer copy / pinned expert staging", so they belong on vt::Backend as virtuals with a conservative base implementation, in the shape of SupportsAuxStream (include/vt/backend.h:122) and SupportsAsyncSampledTokenReadback (:141). Non-ROCm backends inherit the base and the peer path simply does not engage.

That keeps your ROCm fast path exactly as fast, and makes every other build link again.

Rebase onto current main while you are in there — the merge base here is c05cee1d, several commits back — and agent-record / documentation-checkpoint should be satisfied per AGENTS.md (a record under .agents/ bound to this row, plus the docs/ checkpoint files).

Happy to help with the ops-registration wiring if it would be useful — say the word.

[bakon11]

Thanks for the clear CHANGES_REQUESTED review — fully agree on both points.

Ack

1. Non-ROCm link failures — real. The new symbols live only under the ROCm CMake arm while gemma4.cpp / gemma4_moe.cpp call them unconditionally. That will be fixed.

2. Device-agnostic layer — correct rule. Putting vt::rocm:: call sites in src/vllm/model_executor/models/ was the wrong seam. We'll route kernels through OpId + RegisterOp(..., DeviceType::kROCm) portable vt:: entry points, and peer/pin capabilities through vt::Backend virtuals (same shape as SupportsAuxStream / SupportsAsyncSampledTokenReadback), with conservative base defaults so CPU/CUDA/Vulkan simply no-op the peer path.

3. fix(vt,hooks): async readback is a Backend capability + unbreak the pre-push sandbox #159 — thanks for flagging the runner kCUDA leakage on our side of CI; we've merged current main (including TP-W1 and the async-readback/backend work).

Process / PR hygiene

• Split PRs: generic serve/engine fixes that we found while Hermes-testing (e.g. LLMEngine::generate waiting on global unfinished requests, chat-template enable_thinking kwargs) will move to small standalone PRs so they are not blocked on the ROCm MoE campaign and benefit other models.
• perf(gemma4): dual-GPU FP8 resident without host OOM + peer mix #154 stays the Gemma4 MoE / ROCm resident+stream expert work, rewritten to the ops/Backend seam + docs checkpoint / agent-record / PR-size cleanup (PR_DRAFT_*.md etc.).
• spec(gemma4): repair #140 integration regressions #158 — noted; we'll re-inspect gemma4_moe.cpp against the feat(rocm): gfx1201 hipBLAS ops + Gemma-4-26B-A4B MoE (BF16/FP8) #140 fusion-family repair plan so we don't walk over that work.

Rebasing/reworking the ROCm registration path next. Happy to take a pointer if there's a preferred OpId naming pattern for the fused RmsNorm/GeGLU helpers beyond matching existing MatmulBT-style registrations.

FOLLOWING_AGENTS_PROTOCOL
Assisted-by: Hermes:grok-4.5 [Hermes]

[bakon11]

Progress update (post-review)

Addressed the structural review items we could land without a full PR split:

1. No vt::rocm:: in model files — portable vt::fused_ops seam (include/vt/fused_ops.h).
2. check-device-leakage green locally and on CI.
3. build-test-cpu + build-test-vulkan PASS on the post-seam run (link errors fixed).
4. Rebased onto current main.
5. Peer/pin non-HIP stubs.

Still red (honest)

Gate	Why
pr-size	~3400 product lines vs 900 budget — needs split PRs or waiver
documentation-checkpoint	Per-commit POL-DOC on historical commits in this branch
agent-record	Also fails on clean main (stale ACTIVE rows) — repo-wide

Split already done

Generic serve hang fix is #163 (not blocked on this MoE campaign).

Next for #154 when size policy is the priority: squash or carve into ≤900-line PRs (ops seam first, then expert path, then kernels).

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: Hermes:grok-4.5 [Hermes]

[localai-bot]

Re-reviewed at e086435e — the defect I flagged is fixed, and fixed the right way.

Verified rather than taken on trust:

• git grep -c "vt::rocm::" -- src/vllm/model_executor/ on your head returns zero (was 11 across gemma4.cpp and gemma4_moe.cpp).
• The new include/vt/fused_ops.h + src/vt/fused_ops.cpp seam carries exactly the right contract in its own header comment: "Model code MUST call these — never vt::rocm::* — so CPU/CUDA/Vulkan link." Non-HIP stubs sit under the #endif // VLLM_CPP_HIP guard.
• CI on that SHA: build-test-cpu ✅, build-test-vulkan ✅, device-leakage ✅, both sanitizers ✅. The three link failures I quoted are gone.

That is the ops/backend-seam route rather than a #ifdef around the call sites, which is what keeps the device-agnostic layer honest. Nice.

Remaining red is process only — pr-size, agent-record, documentation-checkpoint. Per operator instruction I am merging and carrying those myself rather than bouncing the PR back to you again.

Thanks for the ROCm campaign generally — the gfx1201 evidence in #41 is the strongest data point we have on the discrete lane.