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

.agents/specs/gemma4-rocm-fp8-moe.md

@@ -0,0 +1,60 @@
+# SPIKE / SPEC — Gemma-4 MoE ROCm FP8 + SharedK-WMMA (#317)
+
+**Status:** ACTIVE lab path (dual R9700 gfx1201) · PR tip `feat/gemma4-rocm-fp8-split`  
+**Claim class:** ROCm serve correctness + decode/prefill knobs (not CUDA SACRED parity)  
+**Related PRs:** #234 sampler · #227 KV fail-fast · #316 SSE · #328 bare toolcall · **this**
+
+## Problem
+Gemma-4-26B MoE FP8 on consumer RDNA4 needs: resident dual-GPU experts, peer mix, SharedK-WMMA prefill, decode multi-CTA KV splits, and portable `vt::` seams so `models/` never names HIP (device-leakage).
+
+## In scope (this PR)
+| Area | Files (representative) | Gate |
+|------|------------------------|------|
+| Portable fused_ops / backend hooks | `include/vt/fused_ops.h`, `src/vt/fused_ops.cpp`, `backend.h` | CPU link + device-leakage |
+| FP8 ExpertGeGLU / channel GEMV | `rocm_fp8_channel_gemv.hip`, `rocm_gemma4_experts.hip` | Lab A/B + quality (Paris/arith) |
+| MoE host policy | `gemma4_moe.cpp` / `.h` | Lab load + decode |
+| Prefill SharedK-WMMA | `rocm_paged_attn.hip` (prefill path) | Prefill eng ~2k class @11–12k |
+| Decode KV splits / slide | `rocm_paged_attn.hip` | `VT_ATTN_DECODE_*` recipe |
+| Env surface | `docs/ENVIRONMENT.md` `VT_GEMMA4_*` / `VT_ATTN_*` | check-env-doc |
+
+## Out of scope
+- SSE keepalives (#316), ROCm sampler (#234), bare tool parser (#328)
+- CUDA SACRED gemma4-E4B text golden (dgx) — unchanged by ROCm HIP paths when `HasRocm` off
+- Full MoE token-exact vs vLLM on RDNA4 (no AMD CI runners)
+
+## Env recipe (lab binding)
+See `docs/ENVIRONMENT.md` and lab `gemma4-fp8-recipe.env`:
+`VT_GEMMA4_FP8_HW_CVT`, `VT_ATTN_DECODE_KV_SPLITS`, `VT_ATTN_DECODE_SLIDE_*`, resident experts, prefill GEMM_M / PEER_ACT / SharedK-WMMA.
+
+## Automated coverage in-tree
+1. **Device-leakage** — models/ must not call `vt::rocm::*` (CI `check-device-leakage`).
+2. **CPU unit** — `test_gemma4_rocm_fp8_seams` (this PR): fused_ops symbols resolve; env knobs parse inert defaults on CPU.
+3. **Existing** — `test_gemma4_paged_engine` / registry e2e remain CUDA/dgx optional skips.
+4. **Lab (not CI)** — exclusive decode_depth_curve + Paris/READY on dual R9700; evidence in contributor lab notes, not forced into STATUS ratchet growth.
+
+## Residuals (named, not silent)
+- HIP guards polish in dense `gemma4.cpp` paths
+- Broader parity matrix rows when AMD CI exists
+- Further decode micro-opts stay lab-recipe until KEEP+quality
+
+## Merge criteria (maintainer)
+- Spec present (this file)
+- Device-leakage OK; env documented
+- CPU seam test green
+- No SSE/sampler/toolcall scope creep
+- Sanitize ambient reds treated as main baseline unless unique product fail
+
+## Scope note (2026-08-11 rebuild on main)
+
+This tip is **ROCm FP8 MoE + SharedK-WMMA + fused_ops + expert LRU/prewarm**.
+
+**Explicitly deferred** (localai-bot hold on prior tip):
+- `ForwardGemma4Layers` extract / layer-loop restructure in `gemma4.cpp`
+- `Gemma4DecodeGraph` / pure-decode hipGraph driver
+- Any unguarded CUDA-path forward refactor requiring GB10 token-exact golden
+
+Rationale: those changes touch the gate model forward on all backends. ROCm kernels
+and env-gated MoE paths do not. Revisit as a **separate** PR with CUDA golden.
+
+`VT_GEMMA4_MLP_MOE_PARALLEL` is documented but **not** wired in this tip (was only
+in the deferred layer-loop path).

docs/ENVIRONMENT.md

@@ -176,10 +176,23 @@ on CUDA/CPU builds beyond the documented behavior.
 
 | Variable | Default | What it does |
 |---|---|---|
-| `VT_GEMMA4_EXPERT_VRAM_MB` | unlimited (unset or `0`) | A positive MiB value caps the device expert-cache LRU; unset or `0` leaves the cache unlimited |
+| `VT_GEMMA4_EXPERT_VRAM_MB` | off (unset/`0`) | `N>0` enables device expert LRU with N MiB fill budget; unset/`0` = device expert LRU **off** (resident path uses `VT_GEMMA4_RESIDENT_*`). Evict only with `VT_GEMMA4_EXPERT_EVICT=1` |
 | `VT_GEMMA4_RESIDENT_EXPERTS` | unset | `=1` preloads the Gemma-4 MoE experts resident on the GPU(s) after the first use instead of streaming them per step (discrete-ROCm optimization). No-op (with a stderr note) on a binary built without `-DVLLM_CPP_HIP` |
 | `VT_GEMMA4_RESIDENT_GPUS` | `2` | Number of GPUs across which resident Gemma-4 experts are spread; clamped to the ROCm device count. Read only when `VT_GEMMA4_RESIDENT_EXPERTS=1` |
 | `VT_GEMMA4_RESIDENT_MAX_LAYERS` | (all) | Caps how many MoE layers get resident-preloaded, to fit a smaller VRAM budget. Read only when `VT_GEMMA4_RESIDENT_EXPERTS=1` |
+| `VT_GEMMA4_RESIDENT_BF16` | unset | Force BF16 resident expert packs when set |
+| `VT_GEMMA4_RESIDENT_NATIVE` | unset | Prefer native FP8 resident packs when set |
+| `VT_ATTN_PREFILL_FLASH_SHAREDK` | on | ROCm Gemma-4 prefill: scoreless SharedK flash path (prod default). `0` forces DecodeGqa baseline for A/B |
+| `VT_ATTN_PREFILL_FLASH_WMMA` | off | Lab-only ROCm flash WMMA QK path; leave off for prod |
+| `VT_ATTN_PREFILL_SHAREDK_WMMA` | on | ROCm SharedK fused with rocWMMA QK on sliding d=256 (and global-Q d=512). `0` forces scalar SharedK |
+| `VT_GEMMA4_PREFILL_PEER_ACT` | on | Prefill MoE: run ExpertGeGLU on the expert GPU and peer activations only (not full weight PeerCopy). `0` restores weight PeerCopy |
+| `VT_GEMMA4_GPU0_HEADROOM_GB` | `12` | GiB kept free on GPU0 when packing resident experts (decode vs long-prefill trade). Lab dual R9700 + 49k KV: `8` survives 16k+ prefill; `6` OOMs ~11k |
+| `VT_GEMMA4_PREFILL_BATCH_MOE` | auto / `1` in lab recipe | `=1` group-by-expert prefill GEMM for `T>=64`; `=0` serial M=1 (slow). Unset = auto |
+| `VT_GEMMA4_MLP_MOE_PARALLEL` | off | `=1` run Gemma4 MLP and MoE on two HIP streams (lab; wall ~flat on R9700). Not wired in this PR tip (decode-graph-free split) |
+| `VT_ATTN_PREFILL_FLASH` | off | `=1` SGLang-style BM×BN GQA flash prefill (lab A/B) |
+| `VT_GEMMA4_PREFILL_GEMM_M` | `256` | Tokens per expert in prefill-batch GEMM chunks (`16..2048`). Larger M → fewer launches; lab `512` ~+37% prefill vs `64` |
+| `VT_GEMMA4_HOST_EXPERT_MB` | `512` | Host-side expert staging budget (MiB) for non-resident paths |
+| `VT_GEMMA4_LAYER_TRACE` | off | `=1` layer GPU-synced phase timers; `=2` per-layer heartbeats |
 | `VLLM_CPP_HTTP_FIXED_POOL` | `1` (fixed) | `=0` reverts the HTTP worker pool to the legacy dynamic mode. Production uses the capacity-derived fixed pool; the opt-out exists for same-binary A/B attribution |
 | `VT_ROCM_ATTN_CPU_REF` | unset | `=1` routes ROCm paged attention through the CPU reference kernel instead of the HIP kernel — a correctness A/B for the ROCm attention bring-up |
 | `VT_DEBUG_SAMPLED` | unset | `=1` prints the per-step sampled token id(s) to stderr (sampling-loop debug). Read-only; does not change output. Read once per token, so it does not stall the hot loop |

docs/FEATURES.md

@@ -316,7 +316,7 @@ CPU elementwise GEMM (f32/f16/bf16) runs AVX2 and AVX-512 tiers on x86 where the
 | XPU, TPU | Not started | CUDA, CPU, Metal and Vulkan are the built backends |
 | Custom logits processors on CUDA | Open, not root-caused | Segfaults in a CUDA build, 232/232 green on CPU |
 | Memory budgeting (`ROAD-V1-MEM`, #83) | M1+M2 landed (absolute bytes) | `--kv-cache-memory` sizes the KV pool from an absolute byte budget (ABI v16, group-aware divisor); `--num-blocks` overrides; `--gpu-memory-utilization` needs the M3 profile run (dgx-gated). See `specs/kv-sizing.md` |
-| Gemma4 MoE ROCm fused helpers + V1 sampler | Partial | `vt::fused_ops` seam; ROCm registers full V1 sample ops (temp/top-p/masks/penalties). Public: `VT_GEMMA4_EXPERT_VRAM_MB`, `VT_SERVER_MAX_{PROMPT_CHARS,NEW_TOKENS}` |
+| Gemma4 MoE ROCm FP8 + SharedK-WMMA | Partial | Dual-GPU FP8 resident experts, SharedK-WMMA prefill (RDNA4); decode-graph and forward extract deferred. Env `VT_GEMMA4_*`/`VT_ATTN_*`, seam `test_gemma4_rocm_fp8_seams`. [spec](../.agents/specs/gemma4-rocm-fp8-moe.md) |
 
 ## How to read this page
 

docs/USAGE.md

@@ -1438,3 +1438,10 @@ waiting on the engine (long prefill / TTFT). Interval is `VT_SERVER_SSE_PING_S`
 (default 15s; `0` disables). Comment frames are not `data:` events and do not
 carry tokens. Token streaming still uses a timed wait on the request collector
 so deltas are not collapsed by a poll loop.
+
+## Gemma4 FP8 on ROCm (RDNA4)
+
+Dual-GPU resident FP8 MoE and SharedK-WMMA prefill are controlled via
+ENVIRONMENT.md (`VT_GEMMA4_RESIDENT_*`, `VT_ATTN_*`). Defaults stay safe off RDNA4.
+This PR does **not** restructure the Gemma-4 layer loop or enable decode hipGraph
+(those stay lab-only until a CUDA token-exact gate can land them).

include/vllm/model_executor/models/gemma4_moe.h

@@ -49,6 +49,13 @@ struct Gemma4FusedExperts {
   // Optional device-resident BF16 fused stacks after Prepare.
   mutable void* gate_up_dev = nullptr;
   mutable void* down_dev = nullptr;
+  // Native FP8 layer packs (preferred for is_fp8 resident). Per-expert
+  // dev_fp8_* / dev_s_* point into these bases; free bases on teardown only.
+  mutable void* fp8_gu_base = nullptr;   // u8 [E, 2I, H]
+  mutable void* fp8_dn_base = nullptr;   // u8 [E, H, I]
+  mutable void* fp8_sgu_base = nullptr;  // bf16 [E, 2I]
+  mutable void* fp8_sdn_base = nullptr;  // bf16 [E, H]
+  mutable bool fp8_native_resident = false;
   mutable int dev_id = -1;
   bool Empty() const { return gate_up.Empty() && fp8.empty(); }
 };
@@ -80,15 +87,48 @@ size_t UploadGemma4ExpertsResident(std::vector<Gemma4MoeLayerWeights>& layers,
                                    int num_gpus);
 size_t UploadGemma4ExpertsResidentForWeights(Gemma4Weights& weights, int num_gpus);
 
-// Peer-copy one resident expert (fused stacks on src_dev) into dst buffers on
+// Peer-copy one resident expert (fused BF16 stacks on src_dev) into dst buffers on
 // compute_dev. Returns false if peer path unavailable.
 bool PeerCopyGemma4ExpertSlice(int src_dev, const void* gate_up_base,
                                const void* down_base, int expert_id, int64_t I,
                                int64_t H, int compute_dev, void* gate_up_dst,
                                void* down_dst);
 
+// Peer-copy one native FP8 expert (weights+scales) into compute_dev dsts.
+bool PeerCopyGemma4Fp8ExpertSlice(int src_dev, const void* fp8_gu, const void* fp8_dn,
+                                  const void* s_gu, const void* s_dn, int64_t I, int64_t H,
+                                  int compute_dev, void* fp8_gu_dst, void* fp8_dn_dst,
+                                  void* s_gu_dst, void* s_dn_dst);
+
+// Decode T=1: run top-k FP8 ExpertGeGLU on expert_dev (weights stay put).
+// Peer-copies x (H bf16) to expert_dev and ysum back onto compute_q's device —
+// not the expert weights. Async peer + events when possible.
+// Returns false → caller falls back to weight peer-copy path.
+bool RunGemma4Fp8TopKOnExpertDevice(vt::Queue& compute_q, int expert_dev, void* ysum_compute,
+                                    const void* x_compute, const void* const* fp8_gu,
+                                    const void* const* s_gu, const void* const* fp8_dn,
+                                    const void* const* s_dn, const float* wts, int G, int I,
+                                    int H);
+
+// Prefill batch: one expert GeGLU for M token rows. Weights stay on expert_dev;
+// peer only x/y activations (M×H bf16). Sticky FP8→BF16 dequant on expert.
+// Returns false → caller falls back to weight PeerCopy path.
+bool RunGemma4Fp8ExpertGeGLUPrefillOnExpertDevice(vt::Queue& compute_q, int expert_dev,
+                                                  void* y_compute, const void* x_compute,
+                                                  const void* fp8_gu, const void* s_gu,
+                                                  const void* fp8_dn, const void* s_dn, int M,
+                                                  int I, int H);
+// Same, but contiguous FP8 bases + device idx/wts (no host pointer gather).
+bool RunGemma4Fp8TopKIndexedOnExpertDevice(vt::Queue& compute_q, int expert_dev, void* ysum_compute,
+                                           const void* x_compute, const void* gu_base,
+                                           const void* dn_base, const void* sgu_base,
+                                           const void* sdn_base, const int32_t* idx_compute,
+                                           const float* wts_compute, int G, int I, int H);
+
 // hipHostRegister BF16 expert cache for faster H2D (no-op if already pinned).
 void PinGemma4Fp8ExpertHostCache(const Gemma4Fp8ExpertMats& ex);
+// hipHostUnregister before dropping host BF16 cache (no-op if not pinned).
+void UnpinGemma4Fp8ExpertHostCache(const Gemma4Fp8ExpertMats& ex);
 
 // Dequant one FP8 expert into host BF16 gate_up[2I,H] and down[H,I] (caller-owned).
 // Fills permanent host cache (decode path). Prefer Ephemeral for bulk upload.

include/vt/backend.h

@@ -75,6 +75,13 @@ class Backend {
   // device pointer to a host memcpy and segfaults.
   virtual bool DeviceMemoryIsHostAddressable() const { return false; }
 
+  // Optional device free/total VRAM probe (bytes). Default false = unknown.
+  // ROCm/CUDA override with hipMemGetInfo/cudaMemGetInfo so model code can
+  // size LRU caches without including vendor headers (device-leakage).
+  virtual bool DeviceMemoryInfo(size_t* /*free_bytes*/, size_t* /*total_bytes*/) const {
+    return false;
+  }
+
   // --- Device compute capability (BACKEND-CUDA-ARCH-ADDITIVITY seam-gap #4) ---
   // The architecture the backend is actually running on, as the familiar
   // `(major, minor)` pair (GB10/sm_121 -> {12, 1}). Before this, the capability

include/vt/fused_ops.h

@@ -27,7 +27,30 @@ void MatmulBTAlphaBeta(Queue& q, void* out, const void* a, const void* b, int M,
 void MatmulBTFp8Channel(Queue& q, void* out, const void* a, const void* b_fp8,
                         const void* scale_bf16, int M, int N, int K, float alpha, float beta);
 
+// Device FP8 E4M3 + BF16 channel scale → BF16 weights [N,K] (prefill hipBLAS path).
+void DequantFp8ChannelBf16(Queue& q, void* out_bf16, const void* fp8, const void* scale_bf16,
+                           int N, int K);
+
 bool ExpertGeGLUBf16TopKM1(Queue& q, void* ysum, const void* x, const void* const* w_gu,
                            const void* const* w_dn, const float* wts, int G, int I, int H);
 
+// Fused FP8 expert GeGLU top-k (T=1). Uses hipBLASLt FP8 when available, else fast HIP.
+bool ExpertGeGLUFp8TopKM1(Queue& q, void* ysum, const void* x, const void* const* fp8_gu,
+                          const void* const* s_gu, const void* const* fp8_dn,
+                          const void* const* s_dn, const float* wts, int G, int I, int H);
+bool ExpertGeGLUFp8TopKIndexed(Queue& q, void* ysum, const void* x, const void* gu_base,
+                               const void* dn_base, const void* sgu_base, const void* sdn_base,
+                               const int32_t* idx_dev, const float* wts_dev, int G, int I, int H);
+void ApplyExpertScaleRw(Queue& q, float* rw_dev, const int32_t* ri_dev, const float* escale_dev,
+                        int G, int E);
+// Pre-alloc ExpertGeGLU scratch on `dev` (call after resident expert upload).
+bool PrewarmExpertGeGLUFp8TopK(int dev, int G, int I, int H);
+
+// Prefill MoE: GPU gather / weighted scatter (no host accumulation).
+void MoeGatherRows(Queue& q, void* out_bf16, const void* in_bf16, const int32_t* token_ids_dev,
+                   int n, int H);
+void MoeWeightedScatterAdd(Queue& q, void* acc_bf16, const void* y_bf16,
+                           const int32_t* token_ids_dev, const float* weights_dev, int n, int H);
+void MoeZeroBf16(Queue& q, void* buf_bf16, int64_t nelem);
+
 }  // namespace vt

include/vt/rocm/rocm_device_bind.h

@@ -0,0 +1,31 @@
+// Capture-aware current-device bind for ROCm.
+// hipSetDevice is illegal during hipGraph capture; stream launches carry device.
+// Outside capture, bind so peer-MoE can leave current device on the expert GPU.
+#pragma once
+
+#include <hip/hip_runtime.h>
+
+#include "vt/device.h"
+
+namespace vt::rocm {
+
+inline bool StreamIsCapturing(hipStream_t st) {
+  if (st == nullptr) return false;
+  hipStreamCaptureStatus status = hipStreamCaptureStatusNone;
+  if (hipStreamIsCapturing(st, &status) != hipSuccess) return false;
+  return status != hipStreamCaptureStatusNone;
+}
+
+// Bind process current device to q.device when not capturing. No-op if already set
+// or if the stream is mid-capture (graph-safe).
+inline void EnsureQueueDevice(const Queue& q) {
+  const int dev = q.device.index;
+  if (dev < 0) return;
+  hipStream_t st = static_cast<hipStream_t>(q.handle);
+  if (StreamIsCapturing(st)) return;
+  int cur = -1;
+  if (hipGetDevice(&cur) == hipSuccess && cur == dev) return;
+  (void)hipSetDevice(dev);
+}
+
+}  // namespace vt::rocm

include/vt/rocm/rocm_matmul_batch.h

@@ -33,4 +33,34 @@ void MatmulBTFp8ChannelRocm(Queue& q, void* out, const void* a, const void* b_fp
                             const void* scale_bf16, int M, int N, int K, float alpha,
                             float beta);
 
+// out_bf16[N,K] = scale[n] * f8_e4m3(w[n,k])  (device; for hipBLAS prefill)
+void DequantFp8ChannelBf16Rocm(Queue& q, void* out_bf16, const void* fp8,
+                               const void* scale_bf16, int N, int K);
+
+// Fused Expert GeGLU FP8 decode (T=1). Faster than 3× MatmulBTFp8Channel.
+bool ExpertGeGLUFp8M1Rocm(Queue& q, void* y, const void* x, const void* fp8_gu, const void* s_gu,
+                          const void* fp8_dn, const void* s_dn, int I, int H, float alpha,
+                          float beta);
+bool ExpertGeGLUFp8TopKM1Rocm(Queue& q, void* ysum, const void* x, const void* const* fp8_gu,
+                              const void* const* s_gu, const void* const* fp8_dn,
+                              const void* const* s_dn, const float* wts, int G, int I, int H);
+// Contiguous resident FP8 packs + device idx/wts (decode T=1, no host gather).
+bool ExpertGeGLUFp8TopKIndexedRocm(Queue& q, void* ysum, const void* x, const void* gu_base,
+                                   const void* dn_base, const void* sgu_base, const void* sdn_base,
+                                   const int32_t* idx_dev, const float* wts_dev, int G, int I,
+                                   int H);
+void ApplyExpertScaleRwRocm(Queue& q, float* rw_dev, const int32_t* ri_dev, const float* escale_dev,
+                            int G, int E);
+bool PrewarmExpertGeGLUFp8TopKIndexedRocm(int dev, int G, int I, int H);
+
+// Prefill MoE helpers (GPU-only gather / weighted scatter — no host hacc).
+// out[n,H] = in[token_ids[i], H] for i in [0,n)
+void MoeGatherRowsRocm(Queue& q, void* out_bf16, const void* in_bf16, const int32_t* token_ids,
+                       int n, int H);
+// acc[token_ids[i], :] += weight[i] * y[i, :]  (bf16 acc + bf16 y, float weights)
+void MoeWeightedScatterAddRocm(Queue& q, void* acc_bf16, const void* y_bf16,
+                               const int32_t* token_ids, const float* weights, int n, int H);
+// Zero bf16 buffer [rows*H]
+void MoeZeroBf16Rocm(Queue& q, void* buf_bf16, int64_t nelem);
+
 }  // namespace vt::rocm