From 2d1740dc1f80398f323b75d3ec3bc7e112877254 Mon Sep 17 00:00:00 2001 From: Vikash Loomba Date: Wed, 12 Aug 2026 10:43:18 -0700 Subject: [PATCH 1/3] feat(rocm): MoE combine/gate ops (SharedExpertGate, MoeCombine, MoeCombineGate) (#41) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The next links in the generic MoE path after the router/silu-mul. Hand- translated from cuda_moe.cu (MoeCombineKernel :473, MoeCombineGateKernel :555) and the SharedExpertGate CPU oracle (cpu_ops.cpp:2387). Grid-stride, f32 math, bf16/f32 dtype arms via boundary conversions; the combine-gate folds the shared-expert sigmoid gate rounded through bf16 exactly as the donor. Evidence (4x gfx1100, ROCm 7.14, Release): - new MoE combine/gate cross-device case: 9/9 assertions (MoeCombineGate's oracle is the host-computed composite — no CPU op registration exists) - ctest -R 'rocm|cross_device': 4/4 - full ctest: pre-existing failure set shrinks 7 -> 5; test_bench and test_capi now PASS (they failed at op 77 / the router dtype before the chain). test_loaded_engine_dense now fails only on the async-scheduling assertion (a lane capability gap, not a kernel throw). - Named remaining blocker: the grouped quant expert GEMM (kMatmulBTQuantGrouped), the DeepSeek-V4 keep-quant family. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: pi:kimi-k3 [pi] --- CMakeLists.txt | 3 + docs/USAGE.md | 10 +- src/vt/rocm/rocm_moe_chain.hip | 176 +++++++++++++++++++++++++ src/vt/rocm/rocm_ops.hip | 15 +++ tests/vt/test_backend_cross_device.cpp | 90 +++++++++++++ 5 files changed, 289 insertions(+), 5 deletions(-) create mode 100644 src/vt/rocm/rocm_moe_chain.hip diff --git a/CMakeLists.txt b/CMakeLists.txt index a31be7aae..013d0efab 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1421,6 +1421,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 @@ -1441,6 +1442,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 diff --git a/docs/USAGE.md b/docs/USAGE.md index 4f80a43ef..0d3638061 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -66,11 +66,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 diff --git a/src/vt/rocm/rocm_moe_chain.hip b/src/vt/rocm/rocm_moe_chain.hip new file mode 100644 index 000000000..824c88973 --- /dev/null +++ b/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 +#include + +#include +#include +#include + +#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(q.handle); } +inline unsigned GridFor(int64_t n) { + if (n <= 0) return 1; + const int64_t g = (n + kBlock - 1) / kBlock; + return static_cast(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 +__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(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(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 +__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(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(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 +__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(gridDim.x) * blockDim.x; + for (int64_t idx = static_cast(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<<>>(out.Ptr<__hip_bfloat16>(), + sd.Ptr<__hip_bfloat16>(), + gl.Ptr(), t, h); + else if (obf) + SharedExpertGateK<<>>(out.Ptr<__hip_bfloat16>(), sd.Ptr(), + gl.Ptr(), t, h); + else if (sbf) + SharedExpertGateK<<>>(out.Ptr(), sd.Ptr<__hip_bfloat16>(), + gl.Ptr(), t, h); + else + SharedExpertGateK<<>>(out.Ptr(), sd.Ptr(), + gl.Ptr(), 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(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<<>>( + out.Ptr(), expert_out.Ptr(), weights.Ptr(), + shared != nullptr ? shared->Ptr() : 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(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<<>>( + out.Ptr(), expert_out.Ptr(), weights.Ptr(), sd.Ptr(), + gl.Ptr(), 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 diff --git a/src/vt/rocm/rocm_ops.hip b/src/vt/rocm/rocm_ops.hip index aed29660c..08218bf14 100644 --- a/src/vt/rocm/rocm_ops.hip +++ b/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(&ApplyTokenMaskKernelRocm))); RegisterOp(OpId::kMoeSiluMul, DeviceType::kROCM, reinterpret_cast(static_cast(&MoeSiluMulKernelRocm))); + RegisterOp(OpId::kSharedExpertGate, DeviceType::kROCM, + reinterpret_cast( + static_cast(&SharedExpertGateKernelRocm))); + RegisterOp(OpId::kMoeCombine, DeviceType::kROCM, + reinterpret_cast(static_cast(&MoeCombineKernelRocm))); + RegisterOp(OpId::kMoeCombineGate, DeviceType::kROCM, + reinterpret_cast( + static_cast(&MoeCombineGateKernelRocm))); RegisterOp(OpId::kGdnStateGather, DeviceType::kROCM, reinterpret_cast( static_cast(&GdnStateGatherKernelRocm))); diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index 186500e77..a5aa18833 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -1980,6 +1980,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((u + 0x7FFFu + ((u >> 16) & 1u)) >> 16); +} +static float Bf16ToF32(uint16_t b) { + uint32_t u = static_cast(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(T) * K * H, on = static_cast(T) * H; + const std::vector eo = RandomVec(en, 911); + const std::vector w = RandomVec(static_cast(T) * K, 912, 0.0f, 1.0f); + const std::vector sd = RandomVec(on, 913); + const std::vector gl = RandomVec(static_cast(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 ref_sg_b(on, 0); + std::vector 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 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(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(r * K + j)] * eo[static_cast((r * K + j) * H + c2)]; + const float sv = g * sd[static_cast(r * H + c2)]; + const uint16_t svb = F32ToBf16Rne(sv); + acc += Bf16ToF32(svb); + ref_cg[static_cast(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 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; From 5e680135d544e267def6f8e83210be2df39113bd Mon Sep 17 00:00:00 2001 From: Vikash Loomba Date: Fri, 14 Aug 2026 01:23:48 -0700 Subject: [PATCH 2/3] fix(ROCM): MoE combine/gate dtype refusals + production-dtype test arms -- the #509 review rework Review sweep findings (localai-bot, 2026-08-13), all accepted: 1. The donor's dtype refusals were dropped: cuda_moe.cu:520-524/:597-604 open with VT_CHECKs refusing non-f32/bf16; without them an f16 expert_out passes the seam's IsFloat gate and Tensor::Ptr()'s unchecked cast reads 4 bytes per element from a 2-byte allocation. Refusals added to all three ROCm entry points (SharedExpertGate included: its 4-arm dispatch has the same f16 hazard on sd). 2. The case now exercises the PRODUCTION dtype mix, not only f32: the model path runs expert_out bf16 (qwen3_5.cpp DBuf ddown), shared bf16, out bf16. The new bf16 arm is asserted BIT-EXACT against the CPU reference (both sides thread-per-element, same sequential K order, single store rounding, -ffp-contract=off), and the f32 MoeCombine arm is tightened from NMSE to bit-exact per the donor's design comment (cuda_moe.cu:465-468). Writing the bf16 arm caught a construction bug in the first version of it (an f32-typed tensor over a bf16 buffer -- exactly the OOB class the review predicted the missing arm hid); fixed and re-verified against an independent host composite (0/320) plus a raw-hipMalloc scratch replica of both backends. 3. CMakeLists.txt: the mangled duplicate rocm_moe_chain.hip line removed. 4. docs/FEATURES.md op count 44 -> 47 (counted: the registration sites). Gates (gfx1100, flock): test_backend_cross_device 20/20 (bit-exact bf16 + f32 arms). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: pi:kimi-k3 [pi] --- CMakeLists.txt | 1 - docs/FEATURES.md | 4 +-- docs/USAGE.md | 3 +- src/vt/rocm/rocm_moe_chain.hip | 24 +++++++++++++ tests/vt/test_backend_cross_device.cpp | 48 ++++++++++++++++++++++++-- 5 files changed, 73 insertions(+), 7 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 013d0efab..552c9132c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1443,7 +1443,6 @@ if(VLLM_CPP_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 diff --git a/docs/FEATURES.md b/docs/FEATURES.md index 26317660c..5b22fc801 100644 --- a/docs/FEATURES.md +++ b/docs/FEATURES.md @@ -242,7 +242,7 @@ the only ☐ in our column here. | CPU (x86, Arm i8mm; A76 assembly correct/default, llama speed gate open) | ✅ | ◐ | ☐ | ✅ | | Metal (Apple Silicon) | ✅ | ☐ | ☐ | ✅ | | Vulkan | ◐ | ☐ | ☐ | ✅ | -| ROCm | W0 verified on 5 gfx archs; dense and GDN models run all-native. Strict CPU parity is open in the measured near-tie regime (#269) | 44 registered ops including full GDN; ctest-green gfx1151/1103/1100/1201/1200 ([#41](https://github.com/mudler/vllm.cpp/issues/41)). APU managed allocation is unverified. [ROCM.md](ROCM.md) | ✅ | ✅ | +| ROCm | W0 verified on 5 gfx archs; dense and GDN models run all-native. Strict CPU parity is open in the measured near-tie regime (#269) | 47 registered ops including full GDN and MoE combine/gate; ctest-green gfx1151/1103/1100/1201/1200 ([#41](https://github.com/mudler/vllm.cpp/issues/41)). APU managed allocation is unverified. [ROCM.md](ROCM.md) | ✅ | ✅ | | XPU / TPU | ☐ | ✅ | ◐ | ☐ | | Tenstorrent Blackhole | ◐ `ACTIVE`, OPT-125m STRICT 6/6 e2e; Qwen3-0.6B gate wired with device goldens. Full 16x16 rerun and residual-RMS numerics at the rows≥32 device boundary both owed ([spec](../.agents/specs/tenstorrent-backend.md)) | ✅ | ☐ | ☐ | @@ -330,7 +330,7 @@ CPU elementwise GEMM (f32/f16/bf16) runs AVX2 and AVX-512 tiers on x86 where the | LoRA end to end | CPU brick landed | Unwired standalone; not usable through the server | | Multimodal over HTTP | Image request path wired; forward + codec pending | `ROAD-V1-MM` W1-W3 landed (`server_main.cpp:826`). Open: no mm-forward consuming `Request.mm_features`; no image codec vendored (raw RGB only); video/audio/multi-image not started | | Reranking / classify models | Engine side only | Embeddings are LIVE (`LlamaModel`, `vllm_embed`, `/v1/embeddings`); the classify/score heads are landed ops with no registered arch | -| ROCm | W0 community-verified on 5 gfx archs; classic-dense and GDN-hybrid e2e run all-native; correctness gaps remain | 44 registered ops including the GDN state/conv/postconv/recurrence set; APU managed-allocation branch remains unverified. [ROCM.md](ROCM.md) | +| ROCm | W0 community-verified on 5 gfx archs; classic-dense and GDN-hybrid e2e run all-native; correctness gaps remain | 47 registered ops including the GDN state/conv/postconv/recurrence set and MoE combine/gate; APU managed-allocation branch remains unverified. [ROCM.md](ROCM.md) | | 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` | diff --git a/docs/USAGE.md b/docs/USAGE.md index 0d3638061..1adfacf4f 100644 --- a/docs/USAGE.md +++ b/docs/USAGE.md @@ -68,7 +68,8 @@ GDN-hybrid models call. Compressed conv/SSM state (bf16, the vLLM `SupportsCompressedConvState`/`SupportsCompressedGdnState` backend probes. 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 +are native. All three combine/gate ops accept f32 and bf16 operands and +refuse anything else with a named message (f16 is not a supported arm). 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 diff --git a/src/vt/rocm/rocm_moe_chain.hip b/src/vt/rocm/rocm_moe_chain.hip index 824c88973..7e7ddfe1a 100644 --- a/src/vt/rocm/rocm_moe_chain.hip +++ b/src/vt/rocm/rocm_moe_chain.hip @@ -100,6 +100,12 @@ __global__ void MoeCombineGateK(Tout* out, const Teo* expert_out, const float* w } // namespace void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const Tensor& gl) { + // Same f16 hazard as the combine ops (review sweep on #509): the seam admits + // any float dtype and Tensor::Ptr() is unchecked. + VT_CHECK(sd.dtype == DType::kF32 || sd.dtype == DType::kBF16, + "rocm shared_expert_gate: unsupported sd dtype (f32/bf16 only)"); + VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, + "rocm shared_expert_gate: unsupported out dtype (f32/bf16 only)"); const int64_t t = out.shape[0], h = out.shape[1]; if (t == 0 || h == 0) return; hipStream_t s = AsStream(q); @@ -123,6 +129,17 @@ void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const T void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, const Tensor& weights, const Tensor* shared) { + // The donor's dtype refusals (cuda_moe.cu:520-524) — f16 is admitted by the + // seam's IsFloat gate, and Tensor::Ptr() is an unchecked cast, so without + // these an f16 expert_out would be read at 4 bytes/element from a 2-byte + // allocation (review sweep on #509). + VT_CHECK(expert_out.dtype == DType::kF32 || expert_out.dtype == DType::kBF16, + "rocm moe_combine: unsupported expert_out dtype (f32/bf16 only)"); + VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, + "rocm moe_combine: unsupported out dtype (f32/bf16 only)"); + VT_CHECK(shared == nullptr || shared->dtype == DType::kF32 || + shared->dtype == DType::kBF16, + "rocm moe_combine: unsupported shared dtype (f32/bf16 only)"); const int64_t t = out.shape[0], h = out.shape[1]; const int k = static_cast(weights.shape[1]); const int64_t n = t * h; @@ -149,6 +166,13 @@ void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, void MoeCombineGateKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, const Tensor& weights, const Tensor& sd, const Tensor& gl) { + // Same donor refusals as MoeCombine (cuda_moe.cu:597-604). + VT_CHECK(expert_out.dtype == DType::kF32 || expert_out.dtype == DType::kBF16, + "rocm moe_combine_gate: unsupported expert_out dtype (f32/bf16 only)"); + VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, + "rocm moe_combine_gate: unsupported out dtype (f32/bf16 only)"); + VT_CHECK(sd.dtype == DType::kF32 || sd.dtype == DType::kBF16, + "rocm moe_combine_gate: unsupported sd dtype (f32/bf16 only)"); const int64_t t = out.shape[0], h = out.shape[1]; const int k = static_cast(weights.shape[1]); const int64_t n = t * h; diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index a5aa18833..62bb6b9d7 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -1995,13 +1995,20 @@ static float Bf16ToF32(uint16_t b) { } 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(T) * K * H, on = static_cast(T) * H; const std::vector eo = RandomVec(en, 911); const std::vector w = RandomVec(static_cast(T) * K, 912, 0.0f, 1.0f); const std::vector sd = RandomVec(on, 913); + const std::vector eo_bf = Bf16Bits(eo), sd_bf = Bf16Bits(sd); + // SharedExpertGate (sigmoid*mul), MoeCombine (weighted expert sum +/- + // shared), MoeCombineGate (combine + folded shared gate). f32 and bf16 arms, + // PLUS the production dtype mix the model actually runs (review sweep on + // #509): expert_out bf16 (qwen3_5.cpp DBuf ddown), shared bf16, out bf16. + // MoeCombine/MoeCombineGate are thread-per-element with a single store + // rounding and NO cross-lane reduction (cuda_moe.cu:465-468), so both arms + // are asserted BIT-EXACT — the NMSE aggregate would tolerate a few wrong + // elements, which is exactly how a 2x OOB read hides. const std::vector gl = RandomVec(static_cast(T), 914); for (DeviceType dt : RegisteredDevices()) { @@ -2060,12 +2067,47 @@ TEST_CASE("MoE combine/gate ops match the CPU oracle") { } if (OpAvailable(vt::OpId::kMoeCombine, dt)) { vt::MoeCombine(q, tout, teo, tw, &tsd); - CHECK(Nmse(ref_c, dout.Download()) <= kNmseTol); + // Thread-per-element, single store rounding, no cross-lane reduction: + // bit-exact is the achievable and asserted bar (review sweep on #509). + CHECK(dout.Download() == ref_c); } if (OpAvailable(vt::OpId::kMoeCombineGate, dt)) { vt::MoeCombineGate(q, tout, teo, tw, tsd, tgl); CHECK(Nmse(ref_cg, dout.Download()) <= kNmseTol); } + + // The production bf16 arm: expert_out bf16 + shared bf16 + out bf16 + // (qwen3_5.cpp:5463 ddown / :5326 shared). The CPU oracle runs the same + // ops on the same bf16 tensors; both sides thread-per-element with the + // same sequential K order, so the assertion is BIT-EXACT. + DevBufBytes deo_bf(dev, q, en * 2), dsd_bf(dev, q, on * 2), dout_bf(dev, q, on * 2); + deo_bf.Upload(eo_bf.data()); + dsd_bf.Upload(sd_bf.data()); + Tensor teo_b = Tensor::Contiguous(deo_bf.ptr(), DType::kBF16, d, {T, K, H}); + Tensor tsd_b = Tensor::Contiguous(dsd_bf.ptr(), DType::kBF16, d, {T, H}); + Tensor tout_b = Tensor::Contiguous(dout_bf.ptr(), DType::kBF16, d, {T, H}); + if (OpAvailable(vt::OpId::kMoeCombine, dt) && + OpAvailable(vt::OpId::kMoeCombine, DeviceType::kCPU)) { + // CPU reference on bf16 tensors. + std::vector ref_b(on, 0); + { + vt::Backend& cpu = vt::GetBackend(DeviceType::kCPU); + Queue cq = cpu.CreateQueue(); + const Device cd{DeviceType::kCPU, 0}; + std::vector ceo = eo_bf, csd = sd_bf; + std::vector cw = w; + Tensor r = Tensor::Contiguous(ref_b.data(), DType::kBF16, cd, {T, H}); + Tensor teo_c = Tensor::Contiguous(ceo.data(), DType::kBF16, cd, {T, K, H}); + Tensor tw_c = T2(cw.data(), cd, T, K); + Tensor tsd_c = Tensor::Contiguous(csd.data(), DType::kBF16, cd, {T, H}); + vt::MoeCombine(cq, r, teo_c, tw_c, &tsd_c); + cpu.DestroyQueue(cq); + } + vt::MoeCombine(q, tout_b, teo_b, tw, &tsd_b); + std::vector got_b(on); + dout_bf.Download(got_b.data()); + CHECK(got_b == ref_b); + } dev.DestroyQueue(q); } } From 859a666027b969e81a8e57fe04fabb1b485b5aef Mon Sep 17 00:00:00 2001 From: Vikash Loomba Date: Fri, 14 Aug 2026 02:30:21 -0700 Subject: [PATCH 3/3] fix(ROCM): MoeCombine routed_scale (main's signature change at #684) plumbed through the ROCm arm The rebase onto current main brought main's new MoeCombineFn signature (routed_scale, default 1.0f, scaling the ROUTED sum before the shared term -- upstream apply_routed_scale_to_output). The ROCm kernel applies it in the same f32 accumulator in the same order (one standalone multiply on the finished sum, bit-identical to the CPU reference under -ffp-contract=off); the forward declaration in rocm_ops.hip is updated to match. The bf16 test arm now runs at scale 0.7 so the multiply is exercised, not just the 1.0 passthrough. Gates (gfx1100, flock): test_backend_cross_device 20/20 (bit-exact arms unchanged at 1.0, bit-exact at 0.7). FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: pi:kimi-k3 [pi] --- src/vt/rocm/rocm_moe_chain.hip | 12 +++++++++--- src/vt/rocm/rocm_ops.hip | 2 +- tests/vt/test_backend_cross_device.cpp | 8 ++++---- 3 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/vt/rocm/rocm_moe_chain.hip b/src/vt/rocm/rocm_moe_chain.hip index 7e7ddfe1a..c33424e7c 100644 --- a/src/vt/rocm/rocm_moe_chain.hip +++ b/src/vt/rocm/rocm_moe_chain.hip @@ -61,7 +61,8 @@ __global__ void SharedExpertGateK(Tout* out, const Tsd* sd, const float* gl, int // MoeCombine (cuda_moe.cu:473). template __global__ void MoeCombineK(Tout* out, const Teo* expert_out, const float* weights, - const Tsh* shared, int64_t t, int64_t h, int k) { + const Tsh* shared, int64_t t, int64_t h, int k, + float routed_scale) { const int64_t n = t * h; const int64_t step = static_cast(gridDim.x) * blockDim.x; for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < n; @@ -71,6 +72,11 @@ __global__ void MoeCombineK(Tout* out, const Teo* expert_out, const float* weigh float acc = 0.0f; for (int j = 0; j < k; ++j) acc += weights[row * k + j] * Ld(expert_out, (row * k + j) * h + col); + // routed_scale multiplies the ROUTED sum only, BEFORE the shared term is + // added (upstream apply_routed_scale_to_output: fused_output *= factor, + // shared untouched) — one standalone f32 multiply on the finished + // accumulator, bit-identical to the CPU reference under -ffp-contract=off. + acc *= routed_scale; if (shared != nullptr) acc += Ld(shared, idx); St(out, idx, acc); } @@ -128,7 +134,7 @@ void SharedExpertGateKernelRocm(Queue& q, Tensor& out, const Tensor& sd, const T } void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, - const Tensor& weights, const Tensor* shared) { + const Tensor& weights, const Tensor* shared, float routed_scale) { // The donor's dtype refusals (cuda_moe.cu:520-524) — f16 is admitted by the // seam's IsFloat gate, and Tensor::Ptr() is an unchecked cast, so without // these an f16 expert_out would be read at 4 bytes/element from a 2-byte @@ -149,7 +155,7 @@ void MoeCombineKernelRocm(Queue& q, Tensor& out, const Tensor& expert_out, using Teo = decltype(eo); using Tsh = decltype(sh); using Tout = decltype(ot); MoeCombineK<<>>( out.Ptr(), expert_out.Ptr(), weights.Ptr(), - shared != nullptr ? shared->Ptr() : nullptr, t, h, k); + shared != nullptr ? shared->Ptr() : nullptr, t, h, k, routed_scale); }; auto by_shared = [&](auto eo, auto ot) { if (shared != nullptr && shared->dtype == DType::kBF16) launch(eo, __hip_bfloat16{}, ot); diff --git a/src/vt/rocm/rocm_ops.hip b/src/vt/rocm/rocm_ops.hip index 08218bf14..c504f9f49 100644 --- a/src/vt/rocm/rocm_ops.hip +++ b/src/vt/rocm/rocm_ops.hip @@ -58,7 +58,7 @@ void MoeSiluMulKernelRocm(Queue& q, Tensor& out, const Tensor& gate, const Tenso // 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); + const Tensor& weights, const Tensor* shared, float routed_scale); 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 diff --git a/tests/vt/test_backend_cross_device.cpp b/tests/vt/test_backend_cross_device.cpp index 62bb6b9d7..1ce0549ab 100644 --- a/tests/vt/test_backend_cross_device.cpp +++ b/tests/vt/test_backend_cross_device.cpp @@ -2032,7 +2032,7 @@ TEST_CASE("MoE combine/gate ops match the CPU oracle") { 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); + vt::MoeCombine(cq, to2, teo, tw, &tsd, 1.0f); 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). @@ -2066,7 +2066,7 @@ TEST_CASE("MoE combine/gate ops match the CPU oracle") { CHECK(gotb == ref_sg_b); // both sides store bf16: bit-exact } if (OpAvailable(vt::OpId::kMoeCombine, dt)) { - vt::MoeCombine(q, tout, teo, tw, &tsd); + vt::MoeCombine(q, tout, teo, tw, &tsd, 1.0f); // Thread-per-element, single store rounding, no cross-lane reduction: // bit-exact is the achievable and asserted bar (review sweep on #509). CHECK(dout.Download() == ref_c); @@ -2100,10 +2100,10 @@ TEST_CASE("MoE combine/gate ops match the CPU oracle") { Tensor teo_c = Tensor::Contiguous(ceo.data(), DType::kBF16, cd, {T, K, H}); Tensor tw_c = T2(cw.data(), cd, T, K); Tensor tsd_c = Tensor::Contiguous(csd.data(), DType::kBF16, cd, {T, H}); - vt::MoeCombine(cq, r, teo_c, tw_c, &tsd_c); + vt::MoeCombine(cq, r, teo_c, tw_c, &tsd_c, 0.7f); cpu.DestroyQueue(cq); } - vt::MoeCombine(q, tout_b, teo_b, tw, &tsd_b); + vt::MoeCombine(q, tout_b, teo_b, tw, &tsd_b, 0.7f); std::vector got_b(on); dout_bf.Download(got_b.data()); CHECK(got_b == ref_b);