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

[VikashLoomba]

Row

BACKEND-ROCM — the next links in the generic MoE path, after the router/silu-mul (#348). Issue #41. Claim CLAIM-ROCM-GDN-KERNELS continues.

What changed

NEW src/vt/rocm/rocm_moe_chain.hip with three ops, hand-translated from cuda_moe.cu (MoeCombineKernel :473, MoeCombineGateKernel :555) + the SharedExpertGate CPU oracle (cpu_ops.cpp:2387):

• kSharedExpertGate — out[t,c] = sigmoid(gl[t]) * sd[t,c] (bf16 out, f32/bf16 sd)
• kMoeCombine — weighted top-k expert sum + optional shared term
• kMoeCombineGate — combine with the shared-expert sigmoid gate folded in, the shared term rounded through bf16 exactly as the donor

All grid-stride, f32 math, bf16/f32 dtype arms via boundary conversions. New cross-device case gates all three (MoeCombineGate's oracle is the host-computed composite — no CPU op registration exists for it).

Evidence (4× gfx1100, ROCm 7.14, Release)

• MoE combine/gate case: 9/9 assertions, runs not skips
• 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 (runner_supports_async()=false on ROCm), a lane capability gap, not a kernel throw.
• preflight --staged + trailers green

Speed claims

• This PR makes NO speed claim.

Honest gaps

• Named remaining blocker: the grouped quant expert GEMM (kMatmulBTQuantGrouped, the DeepSeek-V4 keep-quant family) — the heavy lift for Qwen3.5-27B-class GDN-MoE models, a proper kernel project of its own.
• The router's grouped/bias/noaux_tc forms still throw by design (same guard as feat(rocm): gfx1201 hipBLAS ops + Gemma-4-26B-A4B MoE (BF16/FP8) #140); the covered path is ungrouped softmax.

[localai-bot]

Reviewed as part of a sweep over the open external PRs. The translations are faithful — I checked all three donor anchors at your base SHA and they resolve exactly (cuda_moe.cu:473 is MoeCombineKernel, :555 is MoeCombineGateKernel, cpu_ops.cpp:2387 is SharedExpertGateKernel), and the kernel bodies match the donors call for call, including the __bfloat162float(__float2bfloat16(sv)) double-round. Registration goes through the existing registrar with no parallel path.

Two things to fix. Both are the same shape as findings on #506 and #523 — the calculation ported cleanly, the guards around it did not — so it is probably worth reading the three together.

1. The donor's dtype refusals were dropped, and f16 is reachable.

MoeCombineKernelCuda (cuda_moe.cu:520-524) and MoeCombineGateKernelCuda (:597-604) each open with VT_CHECKs refusing anything but f32/bf16, with a named message. The ROCm entry points have none, and dispatch is a bare if (expert_out.dtype == DType::kBF16) ... else <float>.

f16 gets through: vt::MoeCombine gates on IsFloat (src/vt/ops.cpp:21 — kF32 || kF16 || kBF16) for both expert_out and shared, and Tensor::Ptr<T>() (include/vt/tensor.h:66-68) is an unchecked static_cast with no dtype assertion. So an f16 expert_out [T,K,H] passes the seam, falls into the else branch, and the kernel reads 4 bytes per element out of a 2-bytes-per-element allocation — T*K*H*2 bytes past the end, garbage out, no error anywhere. AGENTS.md is explicit that an unimplemented arm is refused with a message naming the missing piece.

2. The new test exercises the dtype arm the model never runs.

On the live path expert_out is bf16 (qwen3_5.cpp:5463, DBuf ddown(d, DType::kBF16, {P, H})), dout is bf16, and the unfused arm passes a bf16 shared (:5326). The new case in test_backend_cross_device.cpp only ever constructs DType::kF32 for expert_out, shared, sd and out, so MoeCombine and MoeCombineGate are tested exclusively as <float,float,float>. The comment says "f32 and bf16 arms"; only SharedExpertGate actually exercises a bf16 store.

Consequence: an inverted predicate in by_out/by_shared/by_sd, or a wrong bf16 Ld/St boundary, is a 2x out-of-bounds read on every MoE layer of every token — and the case still reports 9/9 green, because those branches are never entered.

Smaller ones: CMakeLists.txt:1396-1397 has src/vt/rocm/rocm_moe_chain.hip twice with mangled indentation (inert, but it is exactly the shape that merges badly against another additive ROCm PR touching the same list); docs/FEATURES.md:232,320 still say "44 registered ops" where the count is now 47; and the MoeCombine f32 arm is gated at NMSE 5e-4 where the donor asserts bit-exactness by design (cuda_moe.cu:465-468 — single store-rounding, same as the CPU reference) and the ROCm kernel is thread-per-element with no cross-lane reduction, so CHECK(got == ref) is achievable and is the stronger gate. Your SharedExpertGate check already does exactly that.

Credit where due: recomputing the MoeCombineGate oracle from scratch on host rather than routing it through a shared helper is the right call, and your stated reason for it (no CPU registration for kMoeCombineGate) checks out — cpu_ops.cpp registers kMoeCombine only.

No AMD hardware here, so your 9/9 and the ctest results could not be reproduced and I am not disputing them; both findings are static, read from the seam and the model path.