feat(rocm): MoE combine/gate ops — SharedExpertGate, MoeCombine, MoeCombineGate (issue #41)

CMakeLists.txt

@@ -1363,6 +1363,7 @@ if(VLLM_CPP_HIP)
     src/vt/rocm/rocm_gemma4_expert_geglu.hip
     src/vt/rocm/rocm_fp8_channel_gemv.hip
     src/vt/rocm/rocm_moe_router.hip
+    src/vt/rocm/rocm_moe_chain.hip
     src/vt/rocm/rocm_sample.hip
     src/vt/rocm/rocm_gdn_state.hip
     src/vt/rocm/rocm_gdn_conv.hip
@@ -1383,6 +1384,8 @@ if(VLLM_CPP_HIP)
       src/vt/rocm/rocm_gemma4_expert_geglu.hip
       src/vt/rocm/rocm_fp8_channel_gemv.hip
       src/vt/rocm/rocm_moe_router.hip
+      src/vt/rocm/rocm_moe_chain.hip
+    src/vt/rocm/rocm_moe_chain.hip
       src/vt/rocm/rocm_sample.hip
       src/vt/rocm/rocm_gdn_state.hip
       src/vt/rocm/rocm_gdn_conv.hip

docs/USAGE.md

@@ -50,11 +50,11 @@ portable scan), and the norm-gate/preamble ops (`kRmsNormGated`,
 GDN-hybrid models call. Compressed conv/SSM state (bf16, the vLLM
 `mamba_cache_dtype` default) is advertised via the
 `SupportsCompressedConvState`/`SupportsCompressedGdnState` backend probes.
-MoE-path coverage is partial: `MoeRouterTopK` (f32/bf16 logits, ungrouped
-softmax, no bias) and `MoeSiluMul` are native; the remaining chain
-(`kSharedExpertGate`, `kMoeCombine`/`kMoeCombineGate`, and the grouped quant
-expert GEMM) is not registered yet, so MoE-bearing models still throw on
-those ops. On a
+MoE-path coverage: `MoeRouterTopK` (f32/bf16 logits, ungrouped softmax, no
+bias), `MoeSiluMul`, `SharedExpertGate`, `MoeCombine`, and `MoeCombineGate`
+are native. The grouped quant expert GEMM (`kMatmulBTQuantGrouped`) is not
+registered yet, so models whose experts run keep-quant grouped GEMMs still
+throw there. On a
 discrete card there is no CPU fallback tier, so a model whose layers call an op
 that is not registered yet fails loudly with `vt: no kernel for op N on device
 type 5` — that is the memory-safety design working, not a crash. Run with

src/vt/rocm/rocm_moe_chain.hip

@@ -0,0 +1,176 @@
+// ROCm MoE combine/gate ops (BACKEND-ROCM; the model-path MoE chain, #41).
+// Hand-translations from src/vt/cuda/cuda_moe.cu (MoeCombineKernel :473,
+// MoeCombineGateKernel :555) and the SharedExpertGate CPU oracle
+// (cpu_ops.cpp:2387), readable side by side against the donors. All f32 math;
+// bf16/f32 dtype arms via the Ld/St boundary conversions.
+//
+//   SharedExpertGate: out[t,c] = sigmoid(gl[t]) * sd[t,c]
+//   MoeCombine:       out[r,c] = sum_j w[r,j]*expert_out[(r*k+j),c] (+ shared[r,c])
+//   MoeCombineGate:   MoeCombine + the shared-expert sigmoid gate folded in
+//                     (rounded through bf16 exactly as the donor).
+
+#include <hip/hip_bf16.h>
+#include <hip/hip_runtime.h>
+
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+
+#include "vt/ops.h"
+
+namespace vt::rocm {
+namespace {
+
+constexpr int kBlock = 256;
+
+inline void Check(hipError_t err, const char* what) {
+  if (err != hipSuccess)
+    throw std::runtime_error(std::string("vt rocm moe_chain: ") + what + ": " +
+                             hipGetErrorString(err));
+}
+inline hipStream_t AsStream(const Queue& q) { return static_cast<hipStream_t>(q.handle); }
+inline unsigned GridFor(int64_t n) {
+  if (n <= 0) return 1;
+  const int64_t g = (n + kBlock - 1) / kBlock;
+  return static_cast<unsigned>(g > 65535 ? 65535 : g);
+}
+__device__ inline float Ld(const float* p, int64_t i) { return p[i]; }
+__device__ inline float Ld(const __hip_bfloat16* p, int64_t i) {
+  return __bfloat162float(p[i]);
+}
+__device__ inline void St(float* p, int64_t i, float v) { p[i] = v; }
+__device__ inline void St(__hip_bfloat16* p, int64_t i, float v) {
+  p[i] = __float2bfloat16(v);
+}
+__device__ inline float SigmoidF(float x) { return 1.0f / (1.0f + expf(-x)); }
+
+// SharedExpertGate (cpu_ops.cpp:2387): out[t,c] = sigmoid(gl[t]) * sd[t,c].
+// out bf16 [T,H]; sd f32 [T,H]; gl f32 [T].
+template <typename Tsd, typename Tout>
+__global__ void SharedExpertGateK(Tout* out, const Tsd* sd, const float* gl, int64_t t,
+                                  int64_t h) {
+  const int64_t n = t * h;
+  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       idx < n; idx += step) {
+    const int64_t row = idx / h;
+    St(out, idx, SigmoidF(gl[row]) * Ld(sd, idx));
+  }
+}
+
+// MoeCombine (cuda_moe.cu:473).
+template <typename Teo, typename Tsh, typename Tout>
+__global__ void MoeCombineK(Tout* out, const Teo* expert_out, const float* weights,
+                            const Tsh* shared, int64_t t, int64_t h, int k) {
+  const int64_t n = t * h;
+  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
+       idx += step) {
+    const int64_t row = idx / h;
+    const int64_t col = idx % h;
+    float acc = 0.0f;
+    for (int j = 0; j < k; ++j)
+      acc += weights[row * k + j] * Ld(expert_out, (row * k + j) * h + col);
+    if (shared != nullptr) acc += Ld(shared, idx);
+    St(out, idx, acc);
+  }
+}
+
+// MoeCombineGate (cuda_moe.cu:555): MoeCombine + shared-expert sigmoid gate
+// folded in, the shared term rounded through bf16 exactly as the donor.
+template <typename Teo, typename Tsd, typename Tout>
+__global__ void MoeCombineGateK(Tout* out, const Teo* expert_out, const float* weights,
+                                const Tsd* sd, const float* gl, int64_t t, int64_t h,
+                                int k) {
+  const int64_t n = t * h;
+  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
+       idx += step) {
+    const int64_t row = idx / h;
+    const int64_t col = idx % h;
+    float acc = 0.0f;
+    for (int j = 0; j < k; ++j)
+      acc += weights[row * k + j] * Ld(expert_out, (row * k + j) * h + col);
+    const float sv = SigmoidF(gl[row]) * Ld(sd, idx);
+    acc += __bfloat162float(__float2bfloat16(sv));
+    St(out, idx, acc);
+  }
+}
+
+}  // namespace
+
+void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const Tensor& gl) {
+  const int64_t t = out.shape[0], h = out.shape[1];
+  if (t == 0 || h == 0) return;
+  hipStream_t s = AsStream(q);
+  const int64_t n = t * h;
+  const bool obf = out.dtype == DType::kBF16, sbf = sd.dtype == DType::kBF16;
+  if (obf && sbf)
+    SharedExpertGateK<<<GridFor(n), kBlock, 0, s>>>(out.Ptr<__hip_bfloat16>(),
+                                                    sd.Ptr<__hip_bfloat16>(),
+                                                    gl.Ptr<float>(), t, h);
+  else if (obf)
+    SharedExpertGateK<<<GridFor(n), kBlock, 0, s>>>(out.Ptr<__hip_bfloat16>(), sd.Ptr<float>(),
+                                                    gl.Ptr<float>(), t, h);
+  else if (sbf)
+    SharedExpertGateK<<<GridFor(n), kBlock, 0, s>>>(out.Ptr<float>(), sd.Ptr<__hip_bfloat16>(),
+                                                    gl.Ptr<float>(), t, h);
+  else
+    SharedExpertGateK<<<GridFor(n), kBlock, 0, s>>>(out.Ptr<float>(), sd.Ptr<float>(),
+                                                    gl.Ptr<float>(), t, h);
+  Check(hipGetLastError(), "shared_expert_gate launch");
+}
+
+void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out,
+                          const Tensor& weights, const Tensor* shared) {
+  const int64_t t = out.shape[0], h = out.shape[1];
+  const int k = static_cast<int>(weights.shape[1]);
+  const int64_t n = t * h;
+  if (n == 0) return;
+  hipStream_t s = AsStream(q);
+  auto launch = [&](auto eo, auto sh, auto ot) {
+    using Teo = decltype(eo); using Tsh = decltype(sh); using Tout = decltype(ot);
+    MoeCombineK<Teo, Tsh, Tout><<<GridFor(n), kBlock, 0, s>>>(
+        out.Ptr<Tout>(), expert_out.Ptr<Teo>(), weights.Ptr<float>(),
+        shared != nullptr ? shared->Ptr<Tsh>() : nullptr, t, h, k);
+  };
+  auto by_shared = [&](auto eo, auto ot) {
+    if (shared != nullptr && shared->dtype == DType::kBF16) launch(eo, __hip_bfloat16{}, ot);
+    else launch(eo, float{}, ot);
+  };
+  auto by_out = [&](auto eo) {
+    if (out.dtype == DType::kBF16) by_shared(eo, __hip_bfloat16{});
+    else by_shared(eo, float{});
+  };
+  if (expert_out.dtype == DType::kBF16) by_out(__hip_bfloat16{});
+  else by_out(float{});
+  Check(hipGetLastError(), "moe_combine launch");
+}
+
+void MoeCombineGateKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out,
+                              const Tensor& weights, const Tensor& sd, const Tensor& gl) {
+  const int64_t t = out.shape[0], h = out.shape[1];
+  const int k = static_cast<int>(weights.shape[1]);
+  const int64_t n = t * h;
+  if (n == 0) return;
+  hipStream_t s = AsStream(q);
+  auto launch = [&](auto eo, auto sd_t, auto ot) {
+    using Teo = decltype(eo); using Tsd = decltype(sd_t); using Tout = decltype(ot);
+    MoeCombineGateK<Teo, Tsd, Tout><<<GridFor(n), kBlock, 0, s>>>(
+        out.Ptr<Tout>(), expert_out.Ptr<Teo>(), weights.Ptr<float>(), sd.Ptr<Tsd>(),
+        gl.Ptr<float>(), t, h, k);
+  };
+  auto by_sd = [&](auto eo, auto ot) {
+    if (sd.dtype == DType::kBF16) launch(eo, __hip_bfloat16{}, ot);
+    else launch(eo, float{}, ot);
+  };
+  auto by_out = [&](auto eo) {
+    if (out.dtype == DType::kBF16) by_sd(eo, __hip_bfloat16{});
+    else by_sd(eo, float{});
+  };
+  if (expert_out.dtype == DType::kBF16) by_out(__hip_bfloat16{});
+  else by_out(float{});
+  Check(hipGetLastError(), "moe_combine_gate launch");
+}
+
+}  // namespace vt::rocm

src/vt/rocm/rocm_ops.hip

@@ -54,6 +54,13 @@ void ApplyLogitBiasKernelRocm(Queue& q, Tensor& logits, const Tensor& rows, cons
                               const Tensor& biases);
 // Companion MoE-path op (same TU): elementwise silu(gate)*up.
 void MoeSiluMulKernelRocm(Queue& q, Tensor& out, const Tensor& gate, const Tensor& up);
+// MoE-path combine/gate ops (rocm_moe_chain.hip): shared-expert gate and the
+// weighted expert combinations.
+void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const Tensor& gl);
+void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out,
+                          const Tensor& weights, const Tensor* shared);
+void MoeCombineGateKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out,
+                              const Tensor& weights, const Tensor& sd, const Tensor& gl);
 // BACKEND-ROCM-GDN-KERNELS family 1 (rocm_gdn_state.hip): the indexed state I/O
 // pair `IndexedGdnOpsNative()` requires (issue #41, spec rocm-gdn-kernels.md).
 void GdnStateGatherKernelRocm(Queue& q, Tensor& working, const Tensor& cache,
@@ -181,6 +188,14 @@ struct Registrar {
                    static_cast<ApplyTokenMaskFn>(&ApplyTokenMaskKernelRocm)));
     RegisterOp(OpId::kMoeSiluMul, DeviceType::kROCM,
                reinterpret_cast<void*>(static_cast<MoeSiluMulFn>(&MoeSiluMulKernelRocm)));
+    RegisterOp(OpId::kSharedExpertGate, DeviceType::kROCM,
+               reinterpret_cast<void*>(
+                   static_cast<SharedExpertGateFn>(&SharedExpertGateKernelRocm)));
+    RegisterOp(OpId::kMoeCombine, DeviceType::kROCM,
+               reinterpret_cast<void*>(static_cast<MoeCombineFn>(&MoeCombineKernelRocm)));
+    RegisterOp(OpId::kMoeCombineGate, DeviceType::kROCM,
+               reinterpret_cast<void*>(
+                   static_cast<MoeCombineGateFn>(&MoeCombineGateKernelRocm)));
     RegisterOp(OpId::kGdnStateGather, DeviceType::kROCM,
                reinterpret_cast<void*>(
                    static_cast<GdnStateGatherFn>(&GdnStateGatherKernelRocm)));

tests/vt/test_backend_cross_device.cpp

@@ -1979,6 +1979,96 @@ TEST_CASE("MoeRouterTopK matches the CPU oracle (f32 and bf16 logits)") {
   }
 }
 
+// Scalar bf16 RNE round-trip helpers for host-side oracles (the MoE combine
+// gate reference rounds the shared term through bf16 exactly like the kernel).
+static uint16_t F32ToBf16Rne(float f) {
+  uint32_t u;
+  std::memcpy(&u, &f, 4);
+  return static_cast<uint16_t>((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16);
+}
+static float Bf16ToF32(uint16_t b) {
+  uint32_t u = static_cast<uint32_t>(b) << 16;
+  float f;
+  std::memcpy(&f, &u, 4);
+  return f;
+}
+
+TEST_CASE("MoE combine/gate ops match the CPU oracle") {
+  // SharedExpertGate (sigmoid*mul), MoeCombine (weighted expert sum +/-
+  // shared), MoeCombineGate (combine + folded shared gate). f32 and bf16 arms.
+  constexpr int64_t T = 5, H = 64, K = 3;
+  const size_t en = static_cast<size_t>(T) * K * H, on = static_cast<size_t>(T) * H;
+  const std::vector<float> eo = RandomVec(en, 911);
+  const std::vector<float> w = RandomVec(static_cast<size_t>(T) * K, 912, 0.0f, 1.0f);
+  const std::vector<float> sd = RandomVec(on, 913);
+  const std::vector<float> gl = RandomVec(static_cast<size_t>(T), 914);
+
+  for (DeviceType dt : RegisteredDevices()) {
+    vt::Backend& dev = vt::GetBackend(dt);
+    Queue q = dev.CreateQueue();
+    const Device d{dt, 0};
+    // CPU oracle for all three, f32.
+    std::vector<uint16_t> ref_sg_b(on, 0);
+  std::vector<float> ref_c(on), ref_cg(on);
+    {
+      vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU);
+      Queue cq = cpu.CreateQueue();
+      const Device cd{DeviceType::kCPU, 0};
+      std::vector<float> csd = sd, cgl = gl, ceo = eo, cw = w;
+      Tensor tout = Tensor::Contiguous(ref_sg_b.data(), DType::kBF16, cd, {T, H});
+      Tensor tsd = T2(csd.data(), cd, T, H);
+      Tensor tgl = T1(cgl.data(), cd, T);
+      if (OpAvailable(vt::OpId::kSharedExpertGate, DeviceType::kCPU))
+        vt::SharedExpertGate(cq, tout, tsd, tgl);
+      Tensor teo = Tensor::Contiguous(ceo.data(), DType::kF32, cd, {T, K, H});
+      Tensor tw = T2(cw.data(), cd, T, K);
+      Tensor to2 = T2(ref_c.data(), cd, T, H);
+      if (OpAvailable(vt::OpId::kMoeCombine, DeviceType::kCPU))
+        vt::MoeCombine(cq, to2, teo, tw, &tsd);
+      cpu.DestroyQueue(cq);
+      // MoeCombineGate has no CPU op registration; the oracle is the composite
+      // computed on host: MoeCombine (no shared) + bf16-round(sigmoid(gl)*sd).
+      for (int64_t r = 0; r < T; ++r) {
+        const float g = 1.0f / (1.0f + std::exp(-gl[static_cast<size_t>(r)]));
+        for (int64_t c2 = 0; c2 < H; ++c2) {
+          float acc = 0.0f;
+          for (int64_t j = 0; j < K; ++j)
+            acc += w[static_cast<size_t>(r * K + j)] * eo[static_cast<size_t>((r * K + j) * H + c2)];
+          const float sv = g * sd[static_cast<size_t>(r * H + c2)];
+          const uint16_t svb = F32ToBf16Rne(sv);
+          acc += Bf16ToF32(svb);
+          ref_cg[static_cast<size_t>(r * H + c2)] = acc;
+        }
+      }
+    }
+    // device
+    DevBuf deo(dev, q, en), dw(dev, q, T * K), dsd(dev, q, on), dgl(dev, q, T), dout(dev, q, on);
+    DevBufBytes doutb(dev, q, on * 2);
+    deo.Upload(eo); dw.Upload(w); dsd.Upload(sd); dgl.Upload(gl);
+    Tensor teo = Tensor::Contiguous(deo.ptr(), DType::kF32, d, {T, K, H});
+    Tensor tw = T2(dw.ptr(), d, T, K);
+    Tensor tsd = T2(dsd.ptr(), d, T, H);
+    Tensor tgl = T1(dgl.ptr(), d, T);
+    Tensor tout = T2(dout.ptr(), d, T, H);
+    if (OpAvailable(vt::OpId::kSharedExpertGate, dt)) {
+      Tensor toutb = Tensor::Contiguous(doutb.ptr(), DType::kBF16, d, {T, H});
+      vt::SharedExpertGate(q, toutb, tsd, tgl);
+      std::vector<uint16_t> gotb(on);
+      doutb.Download(gotb.data());
+      CHECK(gotb == ref_sg_b);  // both sides store bf16: bit-exact
+    }
+    if (OpAvailable(vt::OpId::kMoeCombine, dt)) {
+      vt::MoeCombine(q, tout, teo, tw, &tsd);
+      CHECK(Nmse(ref_c, dout.Download()) <= kNmseTol);
+    }
+    if (OpAvailable(vt::OpId::kMoeCombineGate, dt)) {
+      vt::MoeCombineGate(q, tout, teo, tw, tsd, tgl);
+      CHECK(Nmse(ref_cg, dout.Download()) <= kNmseTol);
+    }
+    dev.DestroyQueue(q);
+  }
+}
+
 TEST_CASE("reference tier: an op with no native kernel matches the CPU oracle (unified only)") {
   constexpr int64_t kRows = 7, kCols = 48;
   constexpr size_t kN = kRows * kCols;