feat(gemma4/rocm): FP8 resident MoE + SharedK-WMMA prefill (RDNA4)

CMakeLists.txt

@@ -1310,6 +1310,7 @@ if(VLLM_CPP_HIP)
     src/vt/rocm/rocm_gemma4_experts.hip
     src/vt/rocm/rocm_gemma4_fused_experts.hip
     src/vt/rocm/rocm_gemma4_expert_geglu.hip
+    src/vt/rocm/rocm_sample.hip
     src/vt/rocm/rocm_fp8_channel_gemv.hip
     src/vt/rocm/rocm_moe_router.hip
     src/vt/rocm/rocm_ops.hip)
@@ -1324,6 +1325,7 @@ if(VLLM_CPP_HIP)
       src/vt/rocm/rocm_gemma4_experts.hip
       src/vt/rocm/rocm_gemma4_fused_experts.hip
       src/vt/rocm/rocm_gemma4_expert_geglu.hip
+    src/vt/rocm/rocm_sample.hip
       src/vt/rocm/rocm_fp8_channel_gemv.hip
       src/vt/rocm/rocm_moe_router.hip
       src/vt/rocm/rocm_ops.hip

docs/FEATURES.md

@@ -300,7 +300,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 (`vt::fused_ops`) | Partial | Portable ROCm seam. Public: `VT_GEMMA4_EXPERT_VRAM_MB` (positive-MiB LRU cap; unset/0 unlimited) + `VT_SERVER_MAX_{PROMPT_CHARS,NEW_TOKENS}` (200000/4096; 0 disables). Nine tuning vars internal; defaults unchanged |
+| Gemma4 MoE ROCm fused helpers (`vt::fused_ops`) | partial | ROCm SharedK-WMMA prefill + PEER_ACT MoE + V1 sampler. Public: `VT_GEMMA4_*`, `VT_SERVER_MAX_*` (see ENVIRONMENT). |
 
 ## How to read this page
 

docs/USAGE.md

@@ -472,6 +472,27 @@ generation then runs to the token budget. The ids still count toward
 `min_tokens` masking either way, so `min_tokens` cannot be satisfied by emitting
 a stop token early.
 
+
+### Gemma-4 MoE on ROCm (long-prompt / dual-GPU)
+
+Stop ids come from both `config.json` and sibling `generation_config.json`
+(see “Which token ids stop a generation” above — #262). Prefer GPU-resident
+**native FP8** experts when VRAM allows:
+
+```sh
+export VT_GEMMA4_RESIDENT_EXPERTS=1
+export VT_GEMMA4_RESIDENT_GPUS=2
+export VT_GEMMA4_PREFILL_BATCH_MOE=1
+export HIP_VISIBLE_DEVICES=0,1
+build-hip/examples/vllm-server --model /path/to/gemma-4-26B-A4B-it-fp8 \
+  --host 127.0.0.1 --port 8010 --max-model-len 49152 --num-blocks 1536 \
+  --max-num-seqs 1 --no-enable-thinking --verbose
+```
+
+Default resident packs are **native FP8** (fused ExpertGeGLU decode; prefill uses
+GPU FP8 channel GEMM for T≥64 + device scatter). `VT_GEMMA4_RESIDENT_BF16=1`
+forces BF16 expand for A/B. See [ENVIRONMENT.md](ENVIRONMENT.md).
+
 ### Server flags
 
 | Flag | Default | Meaning |

include/vllm/entrypoints/openai/request_logger.h

@@ -38,9 +38,11 @@ void LogRequestReceived(const std::string& request_id, const std::string& endpoi
 void LogRequestStage(const std::string& request_id, const std::string& stage);
 
 // Completion of a request.
+// prefill_sec >= 0: also log decode_tok_s = completion / (elapsed - prefill).
 void LogRequestFinished(const std::string& request_id, int prompt_tokens,
                         int completion_tokens, const std::string& finish_reason,
-                        double elapsed_sec, const std::string& output_text);
+                        double elapsed_sec, const std::string& output_text,
+                        double prefill_sec = -1.0);
 
 // Errors.
 void LogRequestError(const std::string& request_id, const std::string& endpoint,

include/vllm/entrypoints/openai/serving_utils.h

@@ -36,6 +36,15 @@ namespace vllm::entrypoints::openai {
 // result is always well-formed UTF-8 and therefore safe for nlohmann json dump.
 std::string SanitizeUtf8(const std::string& s);
 
+// SSE comment keepalives while the stream is silent (long prefill / TTFT).
+// Mirrors llama.cpp server `--sse-ping-interval` (comment body ":\n\n").
+// Clients (Hermes inactivity_timeout, undici, proxies) treat no body bytes as
+// dead even when TCP is up; OpenAI data: frames are not required for liveness.
+// VT_SERVER_SSE_PING_S: seconds between pings; default 15; <=0 disables.
+int SsePingIntervalSec();
+// SSE comment frame (not a data: event). Safe for OpenAI SSE clients to ignore.
+inline constexpr const char kSsePingFrame[] = ":\n\n";
+
 // Ported from: vllm/entrypoints/serve/utils/api_utils.py:276-289
 // (should_include_usage). Force mode enables final and continuous usage;
 // request-level continuous stats never take effect without include_usage.

include/vllm/model_executor/models/gemma4.h

@@ -48,6 +48,7 @@
 
 #include <cstdint>
 #include <memory>
+#include <optional>
 #include <vector>
 
 #include "vllm/model_executor/models/model_registry.h"
@@ -224,6 +225,30 @@ std::unique_ptr<LoadedModel> BorrowGemma4LoadedModel(const Gemma4Weights& weight
 // config.raw by the loader/forward. No-op hook.
 void ParseGemma4ForConditionalGenerationConfig(const HfConfig& config);
 
+// Opt-in pure-decode hipGraph driver (VLLM_CPP_GEMMA4_DECODE_GRAPH=1).
+// Default OFF. Fail-soft: capture errors disable graph for the process.
+class Gemma4DecodeGraph {
+ public:
+  Gemma4DecodeGraph(const Gemma4Weights& weights, const HfConfig& config, vt::Queue queue);
+  ~Gemma4DecodeGraph();
+  Gemma4DecodeGraph(const Gemma4DecodeGraph&) = delete;
+  Gemma4DecodeGraph& operator=(const Gemma4DecodeGraph&) = delete;
+  ForwardLogits Step(const std::vector<int32_t>& token_ids,
+                     const std::vector<int32_t>& positions,
+                     const v1::CommonAttentionMetadata& attn_meta,
+                     const std::vector<PagedKvCache>& attn_kv);
+  bool captured() const;
+  int64_t replay_count() const;
+
+ private:
+  struct Impl;
+  std::unique_ptr<Impl> impl_;
+};
+bool Gemma4DecodeGraphEnabled();
+std::optional<ForwardLogits> Gemma4DecodeGraphForward(
+    std::unique_ptr<Gemma4DecodeGraph>& graph, const Gemma4Weights& weights,
+    const ModelForwardInput& input);
+
 // KV-cache spec builder. Emits TWO groups reflecting the true topology: sliding
 // layers (head_dim 256) and full-attention layers (head_dim 512). NOTE: the
 // current runner reads a single uniform head_dim, so consuming a two-head-dim

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/vllm/v1/engine/input_processor.h

@@ -31,7 +31,9 @@
 //   - update_from_generation_config: sets eos_token_id, adds the eos id(s) to
 //     all_stop_token_ids (for MinTokens masking) and merges the SECONDARY eos
 //     ids into stop_token_ids, matching sampling_params.py:627-655 (ROAD-V1-C7
-//     wired all_stop_token_ids; it was previously dropped).
+//     wired all_stop_token_ids; it was previously dropped). The eos list comes
+//     from HfConfig.raw["eos_token_id"], which LoadHfConfig unions with sibling
+//     generation_config.json (e.g. Gemma-4 [1,106] U [1,106,50]).
 //   - update_from_tokenizer tokenizes bad_words into bad_words_token_ids
 //     (sampling_params.py:659-698), ROAD-V1-C7 (previously a no-op stub while
 //     bad_words was deferred on SamplingParams).

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

src/vllm/entrypoints/openai/request_logger.cpp

@@ -92,7 +92,8 @@ void LogRequestStage(const std::string& request_id, const std::string& stage) {
 
 void LogRequestFinished(const std::string& request_id, int prompt_tokens,
                         int completion_tokens, const std::string& finish_reason,
-                        double elapsed_sec, const std::string& output_text) {
+                        double elapsed_sec, const std::string& output_text,
+                        double prefill_sec) {
   if (!g_cfg.enable_log_requests) return;
   std::ostringstream os;
   os << "Finished request " << request_id << " prompt_tokens=" << prompt_tokens
@@ -102,6 +103,13 @@ void LogRequestFinished(const std::string& request_id, int prompt_tokens,
   if (completion_tokens > 0 && elapsed_sec > 0.001) {
     os << " gen_tok_s=" << (static_cast<double>(completion_tokens) / elapsed_sec);
   }
+  if (prefill_sec >= 0.0 && completion_tokens > 0) {
+    const double dec = elapsed_sec - prefill_sec;
+    if (dec > 0.001) {
+      os << " prefill_s=" << prefill_sec
+         << " decode_tok_s=" << (static_cast<double>(completion_tokens) / dec);
+    }
+  }
   if (g_cfg.enable_log_outputs && !output_text.empty()) {
     os << " output: '" << LogPreview(output_text, g_cfg.max_log_len) << "'";
   }