feat(rocm): gfx1201 hipBLAS ops + Gemma-4-26B-A4B MoE (BF16/FP8)

CMakeLists.txt

@@ -604,6 +604,7 @@ add_library(vllm STATIC
   src/vllm/model_executor/models/mistral_weights.cpp
   src/vllm/model_executor/models/gemma4_registry.cpp
   src/vllm/model_executor/models/gemma4_weights.cpp
+  src/vllm/model_executor/models/gemma4_moe.cpp
   src/vllm/model_executor/models/gemma4.cpp
   src/vllm/model_executor/models/gemma4_mm.cpp
   src/vllm/model_executor/models/gemma3_registry.cpp
@@ -1080,11 +1081,21 @@ if(VLLM_CPP_HIP)
     src/vllm/platforms/rocm.cpp
     src/vt/rocm/rocm_backend.hip
     src/vt/rocm/rocm_rmsnorm.hip
+    src/vt/rocm/rocm_embedding.hip
+    src/vt/rocm/rocm_dense_basic.hip
+    src/vt/rocm/rocm_matmul_hipblaslt.hip
+    src/vt/rocm/rocm_paged_attn.hip
+    src/vt/rocm/rocm_gemma4_experts.hip
     src/vt/rocm/rocm_ops.hip)
   if(VLLM_CPP_HIP_ARCHITECTURES)
     set_source_files_properties(
       src/vt/rocm/rocm_backend.hip
       src/vt/rocm/rocm_rmsnorm.hip
+      src/vt/rocm/rocm_embedding.hip
+      src/vt/rocm/rocm_dense_basic.hip
+      src/vt/rocm/rocm_matmul_hipblaslt.hip
+      src/vt/rocm/rocm_paged_attn.hip
+      src/vt/rocm/rocm_gemma4_experts.hip
       src/vt/rocm/rocm_ops.hip
       PROPERTIES HIP_ARCHITECTURES "${VLLM_CPP_HIP_ARCHITECTURES}")
   endif()
@@ -1097,6 +1108,27 @@ if(VLLM_CPP_HIP)
   else()
     target_link_libraries(vllm PUBLIC amdhip64)
   endif()
+  # hipBLAS + hipBLASLt — dense GEMM (ROCm math libraries).
+  find_path(VLLM_CPP_HIPBLAS_INCLUDE hipblas/hipblas.h
+    PATHS ${ROCM_PATH}/include /opt/rocm/include)
+  find_library(VLLM_CPP_HIPBLAS_LIB NAMES hipblas
+    PATHS ${ROCM_PATH}/lib ${ROCM_PATH}/lib64 /opt/rocm/lib /opt/rocm/lib64)
+  find_path(VLLM_CPP_HIPBLASLT_INCLUDE hipblaslt/hipblaslt.h
+    PATHS ${ROCM_PATH}/include /opt/rocm/include)
+  find_library(VLLM_CPP_HIPBLASLT_LIB NAMES hipblaslt
+    PATHS ${ROCM_PATH}/lib ${ROCM_PATH}/lib64 /opt/rocm/lib /opt/rocm/lib64)
+  if(VLLM_CPP_HIPBLAS_INCLUDE AND VLLM_CPP_HIPBLAS_LIB)
+    target_include_directories(vllm PUBLIC ${VLLM_CPP_HIPBLAS_INCLUDE})
+    target_link_libraries(vllm PUBLIC ${VLLM_CPP_HIPBLAS_LIB})
+    message(STATUS "ROCm hipBLAS: ${VLLM_CPP_HIPBLAS_LIB}")
+  else()
+    message(FATAL_ERROR "VLLM_CPP_HIP=ON requires hipBLAS (libhipblas + hipblas/hipblas.h).")
+  endif()
+  if(VLLM_CPP_HIPBLASLT_INCLUDE AND VLLM_CPP_HIPBLASLT_LIB)
+    target_include_directories(vllm PUBLIC ${VLLM_CPP_HIPBLASLT_INCLUDE})
+    target_link_libraries(vllm PUBLIC ${VLLM_CPP_HIPBLASLT_LIB})
+    message(STATUS "ROCm hipBLASLt: ${VLLM_CPP_HIPBLASLT_LIB}")
+  endif()
 else()
   # BIT-ROT GUARD (BACKEND-ROCM W0). The ROCm platform leg is deliberately plain
   # C++ with no HIP header, which means it CAN be compiled on a machine with no

include/vllm/model_executor/model_loader/nvfp4_dequant.h

@@ -76,4 +76,12 @@ void DequantNvfp4ToBf16(const uint8_t* packed, const uint8_t* weight_scale_fp8,
 void DequantFp8ToBf16(const uint8_t* weight_f8, float weight_scale,
                       int64_t numel, uint16_t* out_bf16);
 
+// Channel-wise FP8 (compressed-tensors / llm-compressor FP8_DYNAMIC weights):
+//   weight_f8     [N, K]  F8_E4M3
+//   scale_bf16    [N] or [N,1]  bf16 per-output-channel scale
+//   out_bf16      [N, K]
+//   out[n,k] = bf16( f8(w[n,k]) * bf16_to_f32(scale[n]) )
+void DequantFp8ChannelToBf16(const uint8_t* weight_f8, const uint16_t* scale_bf16,
+                             int64_t N, int64_t K, uint16_t* out_bf16);
+
 }  // namespace vllm

include/vllm/model_executor/models/gemma4.h

@@ -53,6 +53,7 @@
 #include "vllm/model_executor/models/model_registry.h"
 #include "vllm/model_executor/models/qwen3_5.h"          // PagedKvCache, ForwardLogits
 #include "vllm/model_executor/models/qwen3_5_weights.h"  // OwnedTensor
+#include "vllm/model_executor/models/gemma4_moe.h"       // MoE AWQ layer weights
 #include "vllm/transformers_utils/hf_config.h"
 #include "vllm/v1/attention/backend.h"  // CommonAttentionMetadata
 #include "vllm/v1/kv_cache_interface.h"
@@ -97,9 +98,13 @@ struct Gemma4LayerWeights {
   OwnedTensor layer_scalar;               // bf16 [1]  (learned per-layer scalar)
   Gemma4AttnWeights attn;
   Gemma4MlpWeights mlp;
+  // Parallel MoE (26B-A4B): empty when enable_moe_block=false (12B dense).
+  Gemma4MoeLayerWeights moe;
   bool is_full_attention = false;  // layer_type == "full_attention"
   bool is_kv_shared = false;       // layer_idx >= num_layers - num_kv_shared_layers
+  bool k_eq_v = false;             // no v_proj; V shares K (attention_k_eq_v)
   int64_t head_dim = 0;            // 512 full / 256 sliding
+  int64_t num_kv_heads = 0;        // may differ full vs sliding (global vs local GQA)
   int64_t kv_target_layer = -1;    // for shared layers: source of K/V (-1 = self)
 };
 
@@ -115,15 +120,19 @@ struct Gemma4Weights {
   OwnedTensor final_norm;               // bf16 [H]  (model.norm, plain RMSNorm)
   OwnedTensor lm_head;                  // bf16 [H, vocab] Matmul-B; EMPTY when tied
   std::vector<Gemma4LayerWeights> layers;
+  // Keeps safetensors mmaps alive for borrowed fused expert tensors (26B MoE).
+  std::shared_ptr<const void> shards_keepalive;
 };
 
-// Load the `Gemma4ForConditionalGeneration` text backbone (unsloth/gemma-4-E4B-it)
-// safetensors into Gemma4Weights. Name map strips the mm wrapper's
-// `model.language_model.` prefix; the vision/audio towers + embed_vision/
-// embed_audio projectors are SKIPPED (handled by the future G2/G3 towers).
+// Dense / small loads (shards need not outlive weights).
 Gemma4Weights LoadGemma4ForConditionalGenerationWeights(
     const std::vector<SafetensorsFile>& shards, const HfConfig& config);
 
+// MoE BF16: experts mmap-borrowed; pass FromSafetensorsOwned shared_ptr.
+Gemma4Weights LoadGemma4ForConditionalGenerationWeightsOwned(
+    std::shared_ptr<const std::vector<SafetensorsFile>> shards,
+    const HfConfig& config);
+
 // The Gemma-4 text backbone forward. Per decoder layer (gemma4.py:709-767):
 //   r=h; input_layernorm(h) -> attn(Q/K/V plain-norm, V weight-less norm, dual
 //   proportional/standard RoPE, per-layer head_dim, YOCO KV-share) ->

include/vllm/model_executor/models/gemma4_moe.h

@@ -0,0 +1,81 @@
+// Gemma-4 MoE experts: BF16 fused (Google) or FP8 per-expert (Firworks) + resident.
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "vllm/model_executor/models/qwen3_5_weights.h"
+#include "vt/device.h"
+#include "vt/tensor.h"
+
+namespace vllm {
+
+struct Gemma4Weights;
+
+// One FP8 expert (compressed-tensors channel scales). Host mmap borrows.
+struct Gemma4Fp8ExpertMats {
+  OwnedTensor gate_w;  // F8 as I8 [I,H]
+  OwnedTensor gate_s;  // BF16 [I] or [I,1]
+  OwnedTensor up_w;
+  OwnedTensor up_s;
+  OwnedTensor down_w;  // F8 [H,I]
+  OwnedTensor down_s;  // BF16 [H]
+  // Lazy host BF16 cache after first dequant (decode reuse).
+  mutable std::vector<uint16_t> cached_gu;  // [2I,H]
+  mutable std::vector<uint16_t> cached_dn;  // [H,I]
+  // Lazy device BF16 copy on compute GPU (avoids H2D every token).
+  mutable void* dev_gu = nullptr;  // [2I,H] bf16
+  mutable void* dev_dn = nullptr;  // [H,I] bf16
+};
+
+struct Gemma4FusedExperts {
+  // Google BF16 fused stacks (optional).
+  OwnedTensor gate_up;  // bf16 [E, 2I, H]
+  OwnedTensor down;     // bf16 [E, H, I]
+  // Firworks FP8 per-expert (optional). size()==E when is_fp8.
+  bool is_fp8 = false;
+  std::vector<Gemma4Fp8ExpertMats> fp8;
+  int64_t num_experts = 0;
+  int64_t intermediate = 0;
+  int64_t hidden = 0;
+  // Optional device-resident BF16 fused stacks after Prepare.
+  mutable void* gate_up_dev = nullptr;
+  mutable void* down_dev = nullptr;
+  mutable int dev_id = -1;
+  bool Empty() const { return gate_up.Empty() && fp8.empty(); }
+};
+
+struct Gemma4MoeLayerWeights {
+  bool enabled = false;
+  OwnedTensor router_scale;
+  OwnedTensor router_proj;
+  OwnedTensor router_proj_fused;
+  OwnedTensor per_expert_scale;
+  OwnedTensor pre_feedforward_layernorm_2;
+  OwnedTensor post_feedforward_layernorm_1;
+  OwnedTensor post_feedforward_layernorm_2;
+  Gemma4FusedExperts experts;
+  int top_k = 8;
+  int64_t moe_intermediate = 0;
+};
+
+struct Gemma4MoeScratch {
+  vt::Tensor tensor;
+  std::shared_ptr<void> storage;
+};
+
+Gemma4MoeScratch RunGemma4Moe(vt::Queue& q, const Gemma4MoeLayerWeights& moe,
+                              const vt::Tensor& router_in, const vt::Tensor& expert_in,
+                              int64_t T, int64_t H, float rms_eps);
+
+size_t UploadGemma4ExpertsResident(std::vector<Gemma4MoeLayerWeights>& layers,
+                                   int num_gpus);
+size_t UploadGemma4ExpertsResidentForWeights(Gemma4Weights& weights, int num_gpus);
+
+// Dequant one FP8 expert into host BF16 gate_up[2I,H] and down[H,I] (caller-owned).
+void DequantGemma4Fp8ExpertToBf16(const Gemma4Fp8ExpertMats& ex, int64_t I, int64_t H,
+                                  uint16_t* gate_up_out, uint16_t* down_out);
+
+}  // namespace vllm

src/vllm/model_executor/model_loader/nvfp4_dequant.cpp

@@ -85,4 +85,18 @@ void DequantFp8ToBf16(const uint8_t* weight_f8, float weight_scale,
   }
 }
 
+void DequantFp8ChannelToBf16(const uint8_t* weight_f8, const uint16_t* scale_bf16,
+                             int64_t N, int64_t K, uint16_t* out_bf16) {
+  VT_CHECK(weight_f8 != nullptr && scale_bf16 != nullptr && out_bf16 != nullptr,
+           "fp8 channel dequant: null");
+  VT_CHECK(N > 0 && K > 0, "fp8 channel dequant: dims");
+  for (int64_t n = 0; n < N; ++n) {
+    const float s = vt::BF16ToF32(scale_bf16[n]);
+    const uint8_t* wr = weight_f8 + n * K;
+    uint16_t* orow = out_bf16 + n * K;
+    for (int64_t k = 0; k < K; ++k)
+      orow[k] = vt::F32ToBF16(F8E4M3ToF32(wr[k]) * s);
+  }
+}
+
 }  // namespace vllm

src/vllm/model_executor/models/gemma4.cpp

@@ -44,7 +44,8 @@
 #include "vllm/model_executor/layers/linear.h"             // UnquantizedMlpGateUpGeluMethod seam
 #include "vllm/model_executor/models/dense_attn_block.h"  // Dev/DBuf/glue
 #include "vllm/model_executor/models/device_pool.h"       // Pool
-#include "vllm/model_executor/models/qwen3_5_common.h"     // HostLogits
+#include "vllm/model_executor/models/gemma4_moe.h"
+#include "vllm/model_executor/models/qwen3_5_common.h"  // HostLogits
 #include "vt/backend.h"
 #include "vt/dtype.h"
 #include "vt/ops.h"
@@ -201,7 +202,7 @@ DBuf Gemma4AttnBlock(Dev d, const Gemma4LayerWeights& w, const Gemma4Layout& g,
                      double rope_theta_sliding) {
   const int64_t H = g.hidden;
   const int64_t Hq = g.num_q_heads;
-  const int64_t Hkv = g.num_kv_heads;
+  const int64_t Hkv = w.num_kv_heads > 0 ? w.num_kv_heads : g.num_kv_heads;
   const float eps = 1e-6f;  // rms_norm_eps (E4B)
   const int64_t qdim = Hq * Dh, kdim = Hkv * Dh;
   const DType adt = DType::kBF16;
@@ -408,12 +409,14 @@ DBuf ForwardBody(Dev d, const std::vector<int32_t>& token_ids,
   }
 
   // --- Per-Layer Embeddings precompute (gemma4.py:845-898). ---
-  // ple_emb = embed_tokens_per_layer(ids) * sqrt(ple)         -> [T, L*ple]
-  // ple_proj = (per_layer_model_projection @ tok) * hidden^-0.5 -> [T, L*ple]
-  //            reshape [T,L,ple]; RMSNorm(plain) over ple.
-  // ple_input = (ple_proj + ple_emb) * rsqrt(2)               -> [T, L, ple]
-  DBuf ple_input(d, DType::kBF16, {T, L, ple});
-  {
+  // Skipped when hidden_size_per_layer_input==0 (google/gemma-4-12B-it dense).
+  DBuf ple_input(d, DType::kBF16, ple > 0 ? std::vector<int64_t>{T, L, ple}
+                                          : std::vector<int64_t>{1, 1, 1});
+  if (ple > 0) {
+    // ple_emb = embed_tokens_per_layer(ids) * sqrt(ple)         -> [T, L*ple]
+    // ple_proj = (per_layer_model_projection @ tok) * hidden^-0.5 -> [T, L*ple]
+    //            reshape [T,L,ple]; RMSNorm(plain) over ple.
+    // ple_input = (ple_proj + ple_emb) * rsqrt(2)               -> [T, L, ple]
     const int64_t LP = L * ple;
     DBuf ple_emb(d, DType::kBF16, {T, LP});
     {
@@ -479,19 +482,42 @@ DBuf ForwardBody(Dev d, const std::vector<int32_t>& token_ids,
     vt::Add(d.q, h1.t(), attn_n.t(), hidden.t());
 
     // dh2 = pre_feedforward_layernorm(h1); mlp; post_feedforward_layernorm; +h1
+    // MoE (26B-A4B): parallel dense MLP + MoE on residual (sglang gemma4_causal).
     Tensor w_pf = ResidentWeight(d, w.pre_feedforward_layernorm, {H});
     DBuf dh2(d, DType::kBF16, {T, H});
     vt::RmsNorm(d.q, dh2.t(), h1.t(), w_pf, plain);
     DBuf mlp = Gemma4MlpBlock(d, w.mlp, H, I, dh2.t(), T);
-    Tensor w_pff = ResidentWeight(d, w.post_feedforward_layernorm, {H});
-    DBuf mlp_n(d, DType::kBF16, {T, H});
-    vt::RmsNorm(d.q, mlp_n.t(), mlp.t(), w_pff, plain);
+
     DBuf h2(d, DType::kBF16, {T, H});
-    vt::Add(d.q, h2.t(), mlp_n.t(), h1.t());
+    if (w.moe.enabled) {
+      // residual for router = h1 (post-attn residual stream)
+      DBuf moe_in(d, DType::kBF16, {T, H});
+      Tensor w_pf2 = ResidentWeight(d, w.moe.pre_feedforward_layernorm_2, {H});
+      vt::RmsNorm(d.q, moe_in.t(), h1.t(), w_pf2, plain);
+      Gemma4MoeScratch moe_out =
+          RunGemma4Moe(d.q, w.moe, /*router_in=*/h1.t(), /*expert_in=*/moe_in.t(), T, H, eps);
+      Tensor w_p1 = ResidentWeight(d, w.moe.post_feedforward_layernorm_1, {H});
+      Tensor w_p2 = ResidentWeight(d, w.moe.post_feedforward_layernorm_2, {H});
+      DBuf n1(d, DType::kBF16, {T, H});
+      DBuf n2(d, DType::kBF16, {T, H});
+      vt::RmsNorm(d.q, n1.t(), mlp.t(), w_p1, plain);
+      vt::RmsNorm(d.q, n2.t(), moe_out.tensor, w_p2, plain);
+      DBuf sum(d, DType::kBF16, {T, H});
+      vt::Add(d.q, sum.t(), n1.t(), n2.t());
+      Tensor w_pff = ResidentWeight(d, w.post_feedforward_layernorm, {H});
+      DBuf n3(d, DType::kBF16, {T, H});
+      vt::RmsNorm(d.q, n3.t(), sum.t(), w_pff, plain);
+      vt::Add(d.q, h2.t(), n3.t(), h1.t());
+    } else {
+      Tensor w_pff = ResidentWeight(d, w.post_feedforward_layernorm, {H});
+      DBuf mlp_n(d, DType::kBF16, {T, H});
+      vt::RmsNorm(d.q, mlp_n.t(), mlp.t(), w_pff, plain);
+      vt::Add(d.q, h2.t(), mlp_n.t(), h1.t());
+    }
 
     // --- PLE (gemma4.py:753-761): gate = gelu(gate_lin(h2)); gated = gate *
     // ple_l; contrib = post_per_layer_input_norm(proj(gated)); h2 += contrib. ---
-    {
+    if (ple > 0) {
       Tensor wg = ResidentWeight(d, w.per_layer_input_gate, {ple, H});
       DBuf gate_lin(d, DType::kBF16, {T, ple});
       vt::MatmulBT(d.q, gate_lin.t(), h2.t(), wg);
@@ -523,8 +549,10 @@ DBuf ForwardBody(Dev d, const std::vector<int32_t>& token_ids,
     }
 
     // Per-layer learned scalar (gemma4.py:707,765).
-    const double scalar = static_cast<double>(ReadBf16Scalar(w.layer_scalar));
-    vt::MulScalar(d.q, h2.t(), h2.t(), scalar);
+    if (!w.layer_scalar.Empty()) {
+      const double scalar = static_cast<double>(ReadBf16Scalar(w.layer_scalar));
+      vt::MulScalar(d.q, h2.t(), h2.t(), scalar);
+    }
 
     hidden = std::move(h2);
   }