From e2d68404cdfa719c6d7098221661b5bce6ae012e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 12 Aug 2026 20:33:33 +0000 Subject: [PATCH] feat(MODEL-NEMOTRON-H W2): the non-gated relu^2 MoE expert -- one GEMM, not a merged pair (#517) Every grouped-MoE path in this tree is SwiGLU-shaped: a merged gate+up pair with a silu(gate)*up epilogue. NemotronH's expert has no gate half at all -- `ckpt_names=("up_proj", "down_proj", "")` (nemotron_h.py:220 @ 555967922), the empty third entry being the absent gate -- so the expert is h = up_proj(x); h = relu(h)^2; y = down_proj(h) SEAM VERDICT: this is NOT a new merged pair and gets no `MergedGemmGroup` descriptor. `MergedGemmGroup` describes N GEMMs SHARING operand A collapsed into one launch; with N == 1 there is nothing to merge and no launch to save, so an arity-1 descriptor would name a fusion that does not exist. `MlpGateUpMethodBase` is likewise a merged [2I,H] gate_up seam with no pair to hold. The arm is therefore the EXISTING grouped projection plus the activation we did not have -- exactly the shape the gated bf16 archs had before their pair was folded (kMoeGroupedGemmBf16 + kMoeSiluMul). The reasoning is recorded next to the seam it excludes, in merged_gemm.h. No parallel MoE path was added. up : kMoeGroupedGemmBf16 | kMoeGroupedGemmNvfp4Marlin (W4A16 g16) act : kMoeRelu2 <- the only new kernel down : kMoeGroupedGemmBf16 | kMoeGroupedGemmNvfp4Marlin (W4A16 g16) comb : kMoeCombine(..., routed_scale) vt::MoeRelu2 (OpId::kMoeRelu2, appended before kCount; CPU + CUDA) mirrors ReLUSquaredActivation (layers/activation.py:609-628) as the fused-MoE path reaches it: activation_without_mul("relu2") -> MoEActivation.RELU2_NO_MUL (layers/fused_moe/activation.py:33,98) -> `F.relu(input, inplace=True); torch.square(input, out=output)`. The DTYPE ORDER is the mirrored part, not an implementation detail: upstream's kernel (csrc/libtorch_stable/ activation_kernels.cu:673-678) widens to f32, clamps at zero in f32, squares in f32 and rounds ONCE on the store. No new f32 buffer is introduced -- the op reads and writes the caller's dtype and only its arithmetic is f32. routed_scaling_factor goes on the OUTPUT (apply_routed_scale_to_output=True, nemotron_h.py:234), so vt::MoeCombine gained a trailing `routed_scale` (default 1.0f -- every landed caller stays byte-identical, proven by a memcmp test) that multiplies the routed sum BEFORE the shared term is added. That is literally moe_runner.py:389-406 (`fused_output *= routed_scaling_factor`, `shared_output` untouched) then :722-725 (`shared_output + fused_output`). Upstream forces the ROUTER's factor to 1.0 in exactly this case (layer.py:291-300), which is the opposite polarity from Laguna, which folds the same factor into the router weights by linearity (laguna_ops.h:48). group_size=16 NVFP4 -- the spec's named risk -- is SUPPORTED, not emulated: MoeMarlinArgs already defaults to group_size=16 / mxfp4=false, and cuda_moe_marlin.cu:7,115-129 consumes exactly that (group_blocks=1, s_type=kFE4M3fn, num_groups=size_k/group_size); 32 is reachable only via the MXFP4 branch. A test pins the default so a later widening cannot silently re-point these experts. RED first: the new test failed to build on `vt::MoeRelu2 is not a member of vt` and `too many arguments to vt::MoeCombine`. Green after: focused 10/10 cases, 71/71 assertions, Status SUCCESS; clean-tree Release -Werror rebuild 395/395 ctest; Debug arm green on the MoE + op-parity suites (Release is NDEBUG). Mutations executed and caught (restored and re-proven green after each): relu instead of relu^2 (5 cases red), silu instead of relu^2 (5 red), the square narrowed through bf16 before the store (2 red -- the bf16-in/f32-out arm is what sees it; a bf16-out-only test absorbs it), routed scale dropped (2 red), routed scale applied to the combined output including the shared term (2 red), routed scale folded into the router LOGITS (3 cases / 498 assertions red in test_ops_moe_router_grouped), NVFP4 group_size default moved to 32 (1 red). OWED, not claimed: this worktree has no GPU (nvcc absent), so the CUDA arms -- kMoeRelu2 on kCUDA and kMoeGroupedGemmNvfp4Marlin on the real g16 tensors -- are compiled-and-reviewed only. The spec's W2 note records them as owed to a GB10 run. W3/W4 (loader, model file) still own wiring this into NemotronH itself. FOLLOWING_AGENTS_PROTOCOL Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code] --- .agents/specs/nemotron-h-model.md | 72 ++++ include/vt/merged_gemm.h | 15 + include/vt/ops.h | 49 ++- src/vt/cpu/cpu_ops.cpp | 26 +- src/vt/cuda/cuda_moe.cu | 86 +++- src/vt/ops.cpp | 15 +- tests/CMakeLists.txt | 1 + tests/vt/test_ops_moe_nongated_relu2.cpp | 485 +++++++++++++++++++++++ 8 files changed, 730 insertions(+), 19 deletions(-) create mode 100644 tests/vt/test_ops_moe_nongated_relu2.cpp diff --git a/.agents/specs/nemotron-h-model.md b/.agents/specs/nemotron-h-model.md index cbccc7948..bb1ca552d 100644 --- a/.agents/specs/nemotron-h-model.md +++ b/.agents/specs/nemotron-h-model.md @@ -189,6 +189,78 @@ RAM on GB10 and has OOM-rebooted the box. non-attention layers and may not move tokens on short prompts. Gate with a long-prompt arm, not only a 6-token one. +## 6a. W2 note — the non-gated `relu²` expert, as built + +**Seam verdict: the non-gated expert is NOT a merged pair, and does not get a +`MergedGemmGroup` descriptor.** `MergedGemmGroup` describes N GEMMs *sharing +operand A* collapsed into one launch (`merged_gemm.h:1-22`). NemotronH's expert +has exactly one projection — `ckpt_names=("up_proj", "down_proj", "")` +(`nemotron_h.py:220`, the empty third entry being the absent gate) — so with +N == 1 there is nothing to merge and no launch to save; an arity-1 descriptor +would name a fusion that does not exist. `MlpGateUpMethodBase` +(`linear.h:82-86`) is likewise a *merged `[2I,H]` gate_up* seam and has no pair +to hold either. + +The arm is therefore the **existing** grouped projection plus the activation we +did not have — exactly the shape the gated bf16 archs had before their pair was +folded (`kMoeGroupedGemmBf16` + `kMoeSiluMul`): + +``` +up : kMoeGroupedGemmBf16 (bf16) | kMoeGroupedGemmNvfp4Marlin (W4A16 g16) +act : kMoeRelu2 <- NEW, the only new kernel +down : kMoeGroupedGemmBf16 (bf16) | kMoeGroupedGemmNvfp4Marlin (W4A16 g16) +comb : kMoeCombine(..., routed_scale) <- routed scale on the OUTPUT +``` + +No parallel MoE path was added. The reasoning is recorded next to the seam it +excludes (`merged_gemm.h`, the note after the bf16-sibling block). + +**`vt::MoeRelu2` (`OpId::kMoeRelu2`, CPU + CUDA).** Mirrors +`ReLUSquaredActivation` (`layers/activation.py:609-628`) as the fused-MoE path +reaches it: `activation_without_mul("relu2")` → `MoEActivation.RELU2_NO_MUL` +(`layers/fused_moe/activation.py:33,98`) → `apply_moe_activation`'s +`F.relu(input, inplace=True); torch.square(input, out=output)`. The **dtype +order is the mirrored part**: upstream's kernel +(`csrc/libtorch_stable/activation_kernels.cu:673-678`) widens to f32, clamps at +zero in f32, squares in f32 and rounds ONCE on the store. No new f32 buffer is +introduced — the op reads and writes the caller's dtype and only its arithmetic +is f32, which is what `LoadF32`/`StoreF32` already are elsewhere in `vt`. + +**`routed_scaling_factor` is applied to the OUTPUT** +(`apply_routed_scale_to_output=True`, `nemotron_h.py:234`). `vt::MoeCombine` +gained a trailing `routed_scale` (default `1.0f`, so every landed caller is +byte-identical) which multiplies the routed sum *before* the shared term is +added — literally `moe_runner.py:389-406` (`fused_output *= routed_scaling_factor`, +`shared_output` untouched) followed by `:722-725` (`shared_output + fused_output`). +Upstream forces the ROUTER's factor to `1.0` in exactly this case +(`layer.py:291-300`), so `MoeRouterTopKArgs::routed_scaling_factor` stays 1.0 on +this path. Note this is the *opposite* polarity from Laguna, which folds the same +factor into the router weights by linearity (`laguna_ops.h:48`); NemotronH takes +the literal upstream form. + +**`group_size=16` NVFP4 — SUPPORTED, risk closed by source.** `MoeMarlinArgs` +already defaults to `group_size = 16` with `mxfp4 = false` (`ops.h`), and +`cuda_moe_marlin.cu:7,115-129` documents and consumes exactly that +(`group_blocks=1`, `s_type = kFE4M3fn`, `num_groups = size_k / group_size`); 32 +is reachable only via the MXFP4 branch. It is the configuration the landed +NVFP4 MoE archs (Laguna, Qwen3.5) already run. A unit test pins the default so a +later widening cannot silently re-point these experts. **Not run here**: this +worktree has no GPU (`nvcc` absent), so the CUDA arms — `kMoeRelu2` on kCUDA, +`kMoeGroupedGemmNvfp4Marlin` on the real g16 tensors — are compiled-and-reviewed +only and remain owed to a GB10 run (W6, or an earlier GPU-host spot check). + +**Evidence.** `tests/vt/test_ops_moe_nongated_relu2.cpp` (10 cases): the +activation against hand-computed exact values, the `relu`/`silu` mis-ports, a +bf16-in/f32-out arm that catches narrowing the square, a bf16-out raw-bit arm, +the routed scale on the routed sum only, the 1.0 default being byte-identical to +the landed call, and the whole expert `up → relu² → down → scaled combine` against +an independently-written scalar reference. Mutations executed and caught: +`relu` (5 cases red), `silu` (5 red), square narrowed through bf16 (2 red), +routed scale dropped (2 red), routed scale applied to the combined output +including the shared term (2 red), routed scale folded into the router logits +(3 cases / 498 assertions red in `test_ops_moe_router_grouped`), NVFP4 +`group_size` default changed to 32 (1 red). + ## 7. Now **State at this commit:** spec committed, implementation **not started**. The diff --git a/include/vt/merged_gemm.h b/include/vt/merged_gemm.h index 7e1d64d6c..313310c64 100644 --- a/include/vt/merged_gemm.h +++ b/include/vt/merged_gemm.h @@ -76,6 +76,21 @@ inline constexpr MergedGemmGroup kKeepQuantGateUpSwiGLU = { // op vt::MoeGroupedGemmBf16GateUpSilu / OpId::kMoeGroupedGemmBf16GateUpSilu — the // bf16 twin of kMoeGateUpSwiGLUGrouped, BIT-IDENTICAL to {2x MoeGroupedGemmBf16 + // MoeSiluMul}. Same family, distinct weight-marshaling seam. +// +// NON-GATED experts are NOT in this family at all, and deliberately get no +// descriptor. NemotronH's expert (models/nemotron_h.py:126-256 @ 555967922) has +// NO gate half — `ckpt_names=("up_proj", "down_proj", "")` (:220), the empty +// third entry being the absent gate — so the expert is +// h = up_proj(x); h = relu(h)^2; y = down_proj(h) +// with `activation_without_mul(config.mlp_hidden_act)` (:227). A MergedGemmGroup +// describes N GEMMs SHARING operand A collapsed into one launch; with N == 1 +// there is nothing to merge and no launch to save, so an arity-1 descriptor would +// name a fusion that does not exist. The non-gated arm is therefore realized as +// the EXISTING single grouped GEMM plus the activation — kMoeGroupedGemmBf16 (or +// kMoeGroupedGemmNvfp4Marlin for the W4A16 group-16 arm) followed by +// OpId::kMoeRelu2 — which is exactly the shape the gated bf16 archs had before +// their pair was folded. See vt::MoeRelu2 (ops.h) and +// .agents/specs/nemotron-h-model.md §4 W2. // ── Dispatch ───────────────────────────────────────────────────────────────── // Run a merged-GEMM group. For an arity-2 kSiluMulClamp group over keep-quant diff --git a/include/vt/ops.h b/include/vt/ops.h index e2395fca0..8875f7fd3 100644 --- a/include/vt/ops.h +++ b/include/vt/ops.h @@ -371,6 +371,13 @@ enum class OpId : uint8_t { // participates in plan/algo selection. Appended before kCount so no existing // op's id shifts. kMatmulFp8CublasLtAlphaVec, + // The NON-GATED MoE activation: out = relu(x)^2, the whole epilogue of a + // NemotronH expert (models/nemotron_h.py:227 activation_without_mul("relu2") + // -> MoEActivation.RELU2_NO_MUL). Sibling of kMoeSiluMul with ONE input + // instead of two, because a non-gated expert has no gate half to multiply by + // (nemotron_h.py:220 ckpt_names=("up_proj","down_proj","")). See vt::MoeRelu2. + // Appended before kCount so no existing op's id shifts. + kMoeRelu2, kCount }; @@ -860,6 +867,8 @@ using MarlinDenseGemmFn = const Tensor& /*b_scales*/, const Tensor& /*global_scale*/, Tensor& /*workspace*/, const MarlinDenseArgs&); using MoeSiluMulFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&); +// kMoeRelu2: out[i] = relu(x[i])^2 — the NON-GATED MoE activation (one input). +using MoeRelu2Fn = void (*)(Queue&, Tensor&, const Tensor&); // --- Qwen3.6 elementwise "glue" ops (M0.9 forward). These replace host-side // loops so the decode step can run entirely on-device (CUDA-graph capture). // All math in f32; dims are inferred from the tensor shapes (no args structs). @@ -958,8 +967,10 @@ using IndexSelectFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&); using IndexCopyFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&); using MoeRouterTopKFn = void (*)(Queue&, Tensor&, Tensor&, const Tensor&, const MoeRouterTopKArgs&, const Tensor*); +// The trailing float is `routed_scale` — the routed_scaling_factor applied to +// the ROUTED sum before the shared term is added (see vt::MoeCombine). using MoeCombineFn = - void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor*); + void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor*, float); using MoeCombineGateFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor&, const Tensor&); using AttentionFn = void (*)(Queue&, Tensor&, const Tensor&, const Tensor&, const Tensor&, @@ -1576,6 +1587,28 @@ void MarlinDenseGemm(Queue& q, Tensor& c, const Tensor& a, const Tensor& b_q_wei // projections so no concat/copy is needed. CPU + CUDA. void MoeSiluMul(Queue& q, Tensor& out, const Tensor& gate, const Tensor& up); +// out[R,I] = relu(x[R,I])^2 — the NON-GATED MoE activation, and the whole +// epilogue of a NemotronH expert. Mirror of vLLM's `ReLUSquaredActivation` +// (layers/activation.py:609-628, forward_native = torch.square(F.relu(x))) as +// reached through the fused-MoE path: `activation_without_mul("relu2")` -> +// `MoEActivation.RELU2_NO_MUL` -> `apply_moe_activation`'s +// `F.relu(input, inplace=True); torch.square(input, out=output)` +// (layers/fused_moe/activation.py:33,98 and its RELU2_NO_MUL branch). +// +// Why this is NOT a MergedGemmGroup epilogue: a NON-gated expert has no gate +// half to merge with (nemotron_h.py:220 `ckpt_names=("up_proj","down_proj","")` +// — the empty third entry IS the absent gate). There is exactly ONE projection, +// so the expert is the EXISTING grouped GEMM plus this activation, exactly as +// the gated bf16 archs are kMoeGroupedGemmBf16 + kMoeSiluMul. See +// merged_gemm.h's note on the non-gated family. +// +// DTYPE/ROUNDING ORDER is the mirrored part, not an implementation detail: +// upstream's kernel (csrc/libtorch_stable/activation_kernels.cu:673-678) +// widens to f32, clamps at zero in f32, squares in f32 and rounds ONCE on the +// store — so a bf16 input with an f32 output keeps the FULL f32 square. x f32 +// or bf16, out f32/bf16. CPU + CUDA. +void MoeRelu2(Queue& q, Tensor& out, const Tensor& x); + // out[T,H] = x[T,H] / sqrt(mean(x^2) + eps) * w (or *(1+w) when gemma); // out f32 or bf16 (computed in f32, rounded on store). // With residual != nullptr (f32 OR bf16 [T,H]): residual += x first (new residual @@ -2158,7 +2191,7 @@ void MoeRouterTopK(Queue& q, Tensor& weights, Tensor& indices, const Tensor& log const Tensor* e_score_correction_bias = nullptr); // Weighted scatter-combine of the per-expert outputs (moe-semantics.md §4/§6). -// out[t,:] = sum_j weights[t,j] * expert_out[t,j,:] (f32 accumulation) +// out[t,:] = routed_scale * sum_j weights[t,j] * expert_out[t,j,:] (f32 accum) // + shared[t,:] (when shared != nullptr) // expert_out [T,K,H] any float dtype (the K per-slot expert MLP outputs for // token t), weights [T,K] f32 (router weights, §3), optional shared [T,H] any @@ -2167,8 +2200,18 @@ void MoeRouterTopK(Queue& q, Tensor& weights, Tensor& indices, const Tensor& log // (§6 combine order: shared_output + routed_output). The activation-dtype // rounding of the routed sum before the shared add is carried by the caller // materializing expert_out/shared in the activation dtype. +// +// `routed_scale` is upstream's `apply_routed_scale_to_output=True` arm +// (layers/fused_moe/runner/moe_runner.py:389-406 `fused_output *= +// routed_scaling_factor`, then :722-725 `result = shared_output + fused_output`). +// It multiplies the ROUTED sum ONLY — the shared-expert term is added unscaled, +// which is the whole point of the flag and the error a token gate catches late. +// The DEFAULT 1.0f is the `apply_routed_scale_to_output=False` polarity every +// landed caller uses, where the factor is instead folded into the router weights +// by MoeRouterTopKArgs::routed_scaling_factor (layer.py:291-300 forces the +// router's factor to 1.0 exactly when this one is not). void MoeCombine(Queue& q, Tensor& out, const Tensor& expert_out, const Tensor& weights, - const Tensor* shared = nullptr); + const Tensor* shared = nullptr, float routed_scale = 1.0f); // --- Fused MoE combine + shared-expert gate (MoE glue fusion). Equivalent to // SharedExpertGate(shared=bf16(sigmoid(gl)*sd)) followed by MoeCombine(...,shared), diff --git a/src/vt/cpu/cpu_ops.cpp b/src/vt/cpu/cpu_ops.cpp index 9e62a4771..f31fafbd6 100644 --- a/src/vt/cpu/cpu_ops.cpp +++ b/src/vt/cpu/cpu_ops.cpp @@ -460,6 +460,23 @@ void MoeSiluMulKernel(Queue&, Tensor& out, const Tensor& gate, const Tensor& up) }); } +// The NON-GATED MoE activation: out[i] = relu(x[i])^2, the whole epilogue of a +// NemotronH expert (nemotron_h.py:227 -> MoEActivation.RELU2_NO_MUL). Mirrors +// vLLM's relu_squared_kernel (csrc/libtorch_stable/activation_kernels.cu:673-678) +// EXACTLY in dtype order: widen to f32, clamp at zero in f32, square in f32, and +// round ONCE on the store. LoadF32/StoreF32 are that widen/round pair, so a bf16 +// input with an f32 output keeps the full f32 square (no intermediate narrowing). +void MoeRelu2Kernel(Queue&, Tensor& out, const Tensor& x) { + const int64_t n = out.Numel(); + ForRows(n, [&](int64_t r0, int64_t r1) { + for (int64_t i = r0; i < r1; ++i) { + const float f = LoadF32(x, i); + const float v = f > 0.0f ? f : 0.0f; + StoreF32(out, i, v * v); + } + }); +} + // --- TRUE W4A4 (fp4xfp4) helpers + kernels (notes §7). Self-contained fp8/fp4 // codec (vt does not depend on vllm), bit-matching vllm::F8E4M3ToF32 / // F32ToF8E4M3 / CastToFp4 / kE2M1Lut so the op equals vllm::RunNvfp4Emulation. @@ -1975,8 +1992,12 @@ void MoeRouterTopKKernel(Queue&, Tensor& weights, Tensor& indices, const Tensor& // §4/§6 weighted scatter-combine: out[t,:] = sum_j w[t,j]*expert_out[t,j,:] // (f32 accumulation) + shared[t,:] (optional). Stored at out's dtype. +// `routed_scale` multiplies the ROUTED sum only, BEFORE the shared term is added +// — upstream's apply_routed_scale_to_output arm (moe_runner.py:389-406 scales +// `fused_output`, leaves `shared_output` alone, then :722-725 adds them). The +// default 1.0f is the fold-into-router-weights polarity every landed caller uses. void MoeCombineKernel(Queue&, Tensor& out, const Tensor& expert_out, const Tensor& weights, - const Tensor* shared) { + const Tensor* shared, float routed_scale) { const int64_t t = out.shape[0], h = out.shape[1], k = weights.shape[1]; ForRows(t, [&](int64_t r0, int64_t r1) { for (int64_t row = r0; row < r1; ++row) { @@ -1985,6 +2006,7 @@ void MoeCombineKernel(Queue&, Tensor& out, const Tensor& expert_out, const Tenso for (int64_t j = 0; j < k; ++j) acc += weights.Ptr()[row * k + j] * LoadF32(expert_out, (row * k + j) * h + col); + if (routed_scale != 1.0f) acc *= routed_scale; if (shared != nullptr) acc += LoadF32(*shared, row * h + col); StoreF32(out, row * h + col, acc); } @@ -2540,6 +2562,8 @@ struct Registrar { reinterpret_cast(static_cast(&SoftCapKernel))); RegisterOp(OpId::kMoeSiluMul, DeviceType::kCPU, reinterpret_cast(static_cast(&MoeSiluMulKernel))); + RegisterOp(OpId::kMoeRelu2, DeviceType::kCPU, + reinterpret_cast(static_cast(&MoeRelu2Kernel))); RegisterOp(OpId::kScaledFp4Quant, DeviceType::kCPU, reinterpret_cast(static_cast(&ScaledFp4QuantKernel))); RegisterOp(OpId::kSiluMulFp4Quant, DeviceType::kCPU, diff --git a/src/vt/cuda/cuda_moe.cu b/src/vt/cuda/cuda_moe.cu index 1e6537418..1480077a2 100644 --- a/src/vt/cuda/cuda_moe.cu +++ b/src/vt/cuda/cuda_moe.cu @@ -469,9 +469,16 @@ void MoeRouterTopKKernelCuda(Queue& q, Tensor& weights, Tensor& indices, const T // Upstream counterpart: layers/fused_moe/ (moe_sum reduction over the topk // weighted w2 outputs) — M2.2 replaces this correctness-grade path. +// `routed_scale` multiplies the ROUTED sum only, BEFORE the shared term is added +// — upstream's apply_routed_scale_to_output arm (layers/fused_moe/runner/ +// moe_runner.py:389-406 scales `fused_output` and leaves `shared_output` alone, +// then :722-725 adds them). Applied in the same f32 accumulator the CPU +// reference (cpu_ops.cpp MoeCombineKernel) uses, in the same order, so CPU and +// CUDA stay bit-for-bit equal. Default 1.0f == the landed fold-into-weights arm. template __global__ void MoeCombineKernel(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; @@ -481,6 +488,7 @@ __global__ void MoeCombineKernel(Tout* out, const Teo* expert_out, const float* float acc = 0.0f; for (int j = 0; j < k; ++j) acc += weights[row * k + j] * Load(expert_out, (row * k + j) * h + col); + if (routed_scale != 1.0f) acc *= routed_scale; if (shared != nullptr) acc += Load(shared, idx); Store(out, idx, acc); } @@ -488,36 +496,37 @@ __global__ void MoeCombineKernel(Tout* out, const Teo* expert_out, const float* template void LaunchCombine(cudaStream_t s, Tensor& out, const Tensor& expert_out, const Tensor& weights, - const Tensor* shared, int64_t t, int64_t h, int k) { + const Tensor* shared, int64_t t, int64_t h, int k, float routed_scale) { MoeCombineKernel<<>>( 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); Check(cudaGetLastError(), "moe_combine launch"); } // Dispatch shared dtype (or the no-shared path, where Tsh is unused). template void DispatchShared(cudaStream_t s, Tensor& out, const Tensor& expert_out, const Tensor& weights, - const Tensor* shared, int64_t t, int64_t h, int k) { + const Tensor* shared, int64_t t, int64_t h, int k, float routed_scale) { if (shared == nullptr || shared->dtype == DType::kF32) { - LaunchCombine(s, out, expert_out, weights, shared, t, h, k); + LaunchCombine(s, out, expert_out, weights, shared, t, h, k, routed_scale); } else { - LaunchCombine(s, out, expert_out, weights, shared, t, h, k); + LaunchCombine(s, out, expert_out, weights, shared, t, h, k, + routed_scale); } } template void DispatchOut(cudaStream_t s, Tensor& out, const Tensor& expert_out, const Tensor& weights, - const Tensor* shared, int64_t t, int64_t h, int k) { + const Tensor* shared, int64_t t, int64_t h, int k, float routed_scale) { if (out.dtype == DType::kF32) { - DispatchShared(s, out, expert_out, weights, shared, t, h, k); + DispatchShared(s, out, expert_out, weights, shared, t, h, k, routed_scale); } else { - DispatchShared(s, out, expert_out, weights, shared, t, h, k); + DispatchShared(s, out, expert_out, weights, shared, t, h, k, routed_scale); } } void MoeCombineKernelCuda(Queue& q, Tensor& out, const Tensor& expert_out, const Tensor& weights, - const Tensor* shared) { + const Tensor* shared, float routed_scale) { VT_CHECK(expert_out.dtype == DType::kF32 || expert_out.dtype == DType::kBF16, "cuda moe_combine: unsupported expert_out dtype (f32/bf16 only)"); VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, @@ -528,9 +537,11 @@ void MoeCombineKernelCuda(Queue& q, Tensor& out, const Tensor& expert_out, const if (t == 0 || h == 0) return; cudaStream_t s = AsStream(q); if (expert_out.dtype == DType::kF32) { - DispatchOut(s, out, expert_out, weights, shared, t, h, static_cast(k)); + DispatchOut(s, out, expert_out, weights, shared, t, h, static_cast(k), + routed_scale); } else { - DispatchOut<__nv_bfloat16>(s, out, expert_out, weights, shared, t, h, static_cast(k)); + DispatchOut<__nv_bfloat16>(s, out, expert_out, weights, shared, t, h, static_cast(k), + routed_scale); } } @@ -672,6 +683,55 @@ void MoeSiluMulKernelCuda(Queue& q, Tensor& out, const Tensor& gate, const Tenso } } +// --------------------------------------------------------------------------- +// moe_relu2: out[i] = relu(x[i])^2, the NON-GATED MoE activation (NemotronH's +// expert epilogue — nemotron_h.py:227 activation_without_mul("relu2") -> +// MoEActivation.RELU2_NO_MUL). Sibling of moe_silu_mul with ONE input, because a +// non-gated expert has no gate half. Dtype order is upstream's relu_squared_kernel +// (csrc/libtorch_stable/activation_kernels.cu:673-678) verbatim: widen to f32, +// clamp at zero in f32, square in f32, ONE round on the store — so a bf16 input +// with an f32 output keeps the full f32 square. Byte-identical to the CPU +// reference (cpu_ops.cpp MoeRelu2Kernel): both are exact f32 ops, no expf. +template +__global__ void MoeRelu2Kernel(Tout* out, const Tx* x, int64_t n) { + const int64_t step = static_cast(gridDim.x) * blockDim.x; + for (int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n; i += step) { + const float f = Load(x, i); + const float v = f > 0.0f ? f : 0.0f; + Store(out, i, v * v); + } +} + +template +void LaunchRelu2(cudaStream_t s, Tensor& out, const Tensor& x, int64_t n) { + MoeRelu2Kernel<<>>(out.Ptr(), x.Ptr(), n); + Check(cudaGetLastError(), "moe_relu2 launch"); +} + +template +void Relu2ByOut(cudaStream_t s, Tensor& out, const Tensor& x, int64_t n) { + if (out.dtype == DType::kF32) { + LaunchRelu2(s, out, x, n); + } else { + LaunchRelu2(s, out, x, n); + } +} + +void MoeRelu2KernelCuda(Queue& q, Tensor& out, const Tensor& x) { + VT_CHECK(x.dtype == DType::kF32 || x.dtype == DType::kBF16, + "cuda moe_relu2: unsupported x dtype (f32/bf16 only)"); + VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16, + "cuda moe_relu2: unsupported out dtype (f32/bf16 only)"); + const int64_t n = out.Numel(); + if (n == 0) return; + cudaStream_t s = AsStream(q); + if (x.dtype == DType::kF32) { + Relu2ByOut(s, out, x, n); + } else { + Relu2ByOut<__nv_bfloat16>(s, out, x, n); + } +} + // Registers the CUDA MoE kernels during static init (pre-main, like the M0.6 // ops in cuda_ops.cu). Filling the op table is harmless on machines without a // GPU: the kCUDA backend never registers there, so no CUDA queue can dispatch. @@ -685,6 +745,8 @@ struct Registrar { reinterpret_cast(static_cast(&MoeCombineGateKernelCuda))); RegisterOp(OpId::kMoeSiluMul, DeviceType::kCUDA, reinterpret_cast(static_cast(&MoeSiluMulKernelCuda))); + RegisterOp(OpId::kMoeRelu2, DeviceType::kCUDA, + reinterpret_cast(static_cast(&MoeRelu2KernelCuda))); } } registrar; diff --git a/src/vt/ops.cpp b/src/vt/ops.cpp index e4b533fc3..773aed2f4 100644 --- a/src/vt/ops.cpp +++ b/src/vt/ops.cpp @@ -923,6 +923,15 @@ void MoeSiluMul(Queue& q, Tensor& out, const Tensor& gate, const Tensor& up) { reinterpret_cast(GetOp(OpId::kMoeSiluMul, q.device.type))(q, out, gate, up); } +void MoeRelu2(Queue& q, Tensor& out, const Tensor& x) { + VT_CHECK(x.Numel() == out.Numel(), "moe_relu2: out/x must have the same element count"); + VT_CHECK(IsFloat(x.dtype) && IsOutFloat(out.dtype), "moe_relu2: float x, f32/bf16 out"); + VT_CHECK(out.IsContiguous() && x.IsContiguous(), "moe_relu2: contiguous tensors required"); + VT_CHECK(out.device == q.device && x.device == q.device, + "moe_relu2: device mismatch (out/x/queue)"); + reinterpret_cast(GetOp(OpId::kMoeRelu2, q.device.type))(q, out, x); +} + void RmsNorm(Queue& q, Tensor& out, const Tensor& x, const Tensor& weight, const RmsNormArgs& args, Tensor* residual) { VT_CHECK(x.rank == 2 && out.rank == 2 && weight.rank == 1, "rmsnorm: x/out rank-2, w rank-1"); @@ -2303,7 +2312,7 @@ void MoeRouterTopK(Queue& q, Tensor& weights, Tensor& indices, const Tensor& log } void MoeCombine(Queue& q, Tensor& out, const Tensor& expert_out, const Tensor& weights, - const Tensor* shared) { + const Tensor* shared, float routed_scale) { VT_CHECK(expert_out.rank == 3 && weights.rank == 2 && out.rank == 2, "moe_combine: expert_out [T,K,H], weights [T,K], out [T,H]"); const int64_t t = out.shape[0], h = out.shape[1], k = weights.shape[1]; @@ -2324,8 +2333,8 @@ void MoeCombine(Queue& q, Tensor& out, const Tensor& expert_out, const Tensor& w shared->device == q.device, "moe_combine: shared must be float [T,H] contiguous on the queue device"); } - reinterpret_cast(GetOp(OpId::kMoeCombine, q.device.type))(q, out, expert_out, - weights, shared); + reinterpret_cast(GetOp(OpId::kMoeCombine, q.device.type))( + q, out, expert_out, weights, shared, routed_scale); } void MoeCombineGate(Queue& q, Tensor& out, const Tensor& expert_out, const Tensor& weights, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b9c5dc0c1..af225ee81 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1129,6 +1129,7 @@ vllm_cpp_add_test(test_ops_moe_grouped vt/test_ops_moe_grouped.cpp) vllm_cpp_add_test(test_ops_moe_grouped_bf16 vt/test_ops_moe_grouped_bf16.cpp) vllm_cpp_add_test(test_ops_moe_grouped_bf16_gate_up_silu vt/test_ops_moe_grouped_bf16_gate_up_silu.cpp) +vllm_cpp_add_test(test_ops_moe_nongated_relu2 vt/test_ops_moe_nongated_relu2.cpp) vllm_cpp_add_test(test_ops_rmsnorm vt/test_ops_rmsnorm.cpp) vllm_cpp_add_test(test_ops_fused_chain vt/test_ops_fused_chain.cpp) vllm_cpp_add_test(test_ops_layernorm vt/test_ops_layernorm.cpp) diff --git a/tests/vt/test_ops_moe_nongated_relu2.cpp b/tests/vt/test_ops_moe_nongated_relu2.cpp new file mode 100644 index 000000000..d72907d11 --- /dev/null +++ b/tests/vt/test_ops_moe_nongated_relu2.cpp @@ -0,0 +1,485 @@ +// NemotronH's NON-GATED relu^2 MoE expert (row MODEL-TEXT-nemotron-h, W2). +// Spec: .agents/specs/nemotron-h-model.md §4 W2. Issue #517. +// +// Upstream mirror, all @ 5559679229bc961848b121ccdeaa8fa5d79bec98 (vLLM 0.26.0.dev0): +// models/nemotron_h.py:126-256 `NemotronHMoE` +// :220 ckpt_names=("up_proj", "down_proj", "") — the EMPTY third entry is the +// absent gate half. The expert is up_proj -> relu^2 -> down_proj; there +// is no gate_proj tensor anywhere in the checkpoint. +// :227 activation=activation_without_mul(config.mlp_hidden_act) ("relu2" -> +// "relu2_no_mul") +// :234 apply_routed_scale_to_output=True, :232 routed_scaling_factor +// layers/fused_moe/activation.py:98 `activation_without_mul`, :33 +// `MoEActivation.RELU2_NO_MUL`, and `apply_moe_activation`'s RELU2_NO_MUL +// branch: F.relu(input, inplace=True); torch.square(input, out=output). +// layers/activation.py:609-628 `ReLUSquaredActivation` +// (forward_native = torch.square(F.relu(x))). +// csrc/libtorch_stable/activation_kernels.cu:672-678 `relu_squared_kernel` — +// the DTYPE/ROUNDING ORDER this file pins: widen to f32, clamp at 0 in f32, +// square in f32, then ONE round back to the store dtype. +// layers/fused_moe/runner/moe_runner.py:389-406 +// `_maybe_apply_routed_scale_to_output` — `fused_output *= routed_scaling_factor` +// with the SHARED output left UNSCALED, then :722-725 `result = shared_output +// + fused_output`. +// layers/fused_moe/layer.py:291-300 — with apply_routed_scale_to_output the +// ROUTER's routed_scaling_factor is forced to 1.0 ("so it ends up being a nop"), +// which is why the scale is NOT visible in the router weights here. +// +// CPU-only by construction: the two things W2 adds (vt::MoeRelu2 and the +// routed-output scale on vt::MoeCombine) are both registered on kCPU, and the +// non-gated expert composite below runs through vt::MatmulBT — the same +// per-expert reference loop the CPU/GGUF MoE path already uses. The CUDA arms +// (kMoeGroupedGemmBf16 / kMoeGroupedGemmNvfp4Marlin) are registered for kCUDA +// only and are exercised where a GPU exists. +#include + +#include +#include +#include +#include +#include + +#include "vt/dtype.h" +#include "vt/ops.h" + +namespace { + +using vt::Device; +using vt::DeviceType; +using vt::DType; +using vt::Queue; +using vt::Tensor; + +Device Cpu() { return Device{DeviceType::kCPU, 0}; } +Queue Q() { return Queue{Cpu(), nullptr}; } + +Tensor F32_1(std::vector& v) { + return Tensor::Contiguous(v.data(), DType::kF32, Cpu(), {static_cast(v.size())}); +} +Tensor F32_2(std::vector& v, int64_t a, int64_t b) { + return Tensor::Contiguous(v.data(), DType::kF32, Cpu(), {a, b}); +} +Tensor F32_3(std::vector& v, int64_t a, int64_t b, int64_t c) { + return Tensor::Contiguous(v.data(), DType::kF32, Cpu(), {a, b, c}); +} +Tensor Bf16_1(std::vector& v) { + return Tensor::Contiguous(v.data(), DType::kBF16, Cpu(), {static_cast(v.size())}); +} + +// doctest::Approx carries a 1.19e-5 ABSOLUTE floor (its `scale` defaults to 1.0), +// which silently accepts a dropped term on small values. Every non-exact +// comparison here goes through this explicit relative+absolute comparator. +bool Close(float a, float b, float rel = 1e-6f, float abs_tol = 1e-30f) { + const float d = std::fabs(a - b); + return d <= abs_tol || d <= rel * std::fmax(std::fabs(a), std::fabs(b)); +} + +float SiluRef(float x) { return x / (1.0f + std::exp(-x)); } + +} // namespace + +// --------------------------------------------------------------------------- +// 1. The activation itself — ReLUSquaredActivation (activation.py:609-628). +// --------------------------------------------------------------------------- + +// torch.square(F.relu(x)) on exactly-representable inputs, so the expected +// values are exact and the check needs no tolerance at all. +TEST_CASE("moe relu2: mirrors torch.square(F.relu(x)) on the f32 arm") { + std::vector x = {-3.0f, -0.5f, -0.0f, 0.0f, 0.5f, 1.0f, 2.0f, 3.0f}; + std::vector out(x.size(), -1.0f); + Tensor xt = F32_1(x); + Tensor ot = F32_1(out); + Queue q = Q(); + vt::MoeRelu2(q, ot, xt); + + const std::vector want = {0.0f, 0.0f, 0.0f, 0.0f, 0.25f, 1.0f, 4.0f, 9.0f}; + for (size_t i = 0; i < want.size(); ++i) { + CHECK(out[i] == want[i]); + } + // The negative half is exactly ZERO, not a small negative (a `x*|x|` or + // `x*x*sign(x)` mis-port would keep the sign). + CHECK(std::signbit(out[0]) == false); +} + +// The two mis-ports a token gate would catch late: relu (forgot the square) and +// silu (took the GATED family's activation). Pinned at a single input where all +// three differ by a wide margin. +TEST_CASE("moe relu2: is neither relu nor silu") { + std::vector x = {2.0f, 3.0f, -1.5f}; + std::vector out(x.size(), -1.0f); + Tensor xt = F32_1(x); + Tensor ot = F32_1(out); + Queue q = Q(); + vt::MoeRelu2(q, ot, xt); + + CHECK(out[0] == 4.0f); // relu(2)^2 + CHECK(out[0] != 2.0f); // NOT relu(2) + CHECK(!Close(out[0], SiluRef(2.0f))); // NOT silu(2) == 1.7615942 + CHECK(out[1] == 9.0f); // relu(3)^2 + CHECK(out[1] != 3.0f); // NOT relu(3) + CHECK(!Close(out[1], SiluRef(3.0f))); // NOT silu(3) == 2.8577223 + CHECK(out[2] == 0.0f); // relu(-1.5)^2 + CHECK(!Close(out[2], SiluRef(-1.5f))); // silu(-1.5) == -0.27440965, NOT zero +} + +// The dtype/rounding order of relu_squared_kernel (activation_kernels.cu:673-678): +// the square is computed in FP32 and stored ONCE. bf16 in, f32 out is the arm +// that catches a kernel narrowing the product back through bf16 before the store +// — the exact defect a bf16-out-only test absorbs (bf16 rounds it away). +// x = 1 + 1/128 = 1.0078125 (exactly representable in bf16) +// x^2 = 16641/16384 = 1.01568603515625 (exact in f32) +// bf16(x^2) = 1 + 2/128 = 1.015625 (what a narrowed square would store) +TEST_CASE("moe relu2: bf16 in / f32 out keeps the square in f32 (no narrowing)") { + const float x0 = 1.0078125f; + std::vector x = {vt::F32ToBF16(x0)}; + REQUIRE(vt::BF16ToF32(x[0]) == x0); // the input itself is exact in bf16 + std::vector out(1, -1.0f); + Tensor xt = Bf16_1(x); + Tensor ot = F32_1(out); + Queue q = Q(); + vt::MoeRelu2(q, ot, xt); + + CHECK(out[0] == x0 * x0); // 1.01568603515625, the f32 square + CHECK(out[0] != vt::BF16ToF32(vt::F32ToBF16(x0 * x0))); // NOT the narrowed 1.015625 +} + +// The bf16 STORE arm: exactly one round-to-nearest-even of the f32 square. Raw +// bit comparison, not a tolerance — a second rounding step is invisible to any +// bf16 tolerance wide enough to be meaningful. +TEST_CASE("moe relu2: bf16 out rounds the f32 square exactly once") { + const float x0 = 1.0078125f; + std::vector x = {vt::F32ToBF16(x0), vt::F32ToBF16(-2.5f)}; + std::vector out(2, 0xFFFF); + Tensor xt = Bf16_1(x); + Tensor ot = Bf16_1(out); + Queue q = Q(); + vt::MoeRelu2(q, ot, xt); + + CHECK(out[0] == vt::F32ToBF16(x0 * x0)); + CHECK(out[1] == vt::F32ToBF16(0.0f)); +} + +TEST_CASE("moe relu2: rejects a shape/dtype/device contract violation") { + std::vector x(8, 1.0f); + std::vector small(4, 0.0f); + Tensor xt = F32_1(x); + Tensor st = F32_1(small); + Queue q = Q(); + CHECK_THROWS_AS(vt::MoeRelu2(q, st, xt), std::runtime_error); + + std::vector ints(8, 0); + Tensor it = Tensor::Contiguous(ints.data(), DType::kI32, Cpu(), {8}); + CHECK_THROWS_AS(vt::MoeRelu2(q, it, xt), std::runtime_error); +} + +// --------------------------------------------------------------------------- +// 2. routed_scaling_factor on the OUTPUT (apply_routed_scale_to_output=True). +// --------------------------------------------------------------------------- + +// moe_runner.py:400-406 scales `fused_output` and leaves `shared_output` alone; +// :722-725 then adds them. So the combine is +// out = routed_scale * sum_j w[t,j]*expert_out[t,j] + shared[t] +// and NOT routed_scale * (routed + shared), and NOT (routed + shared). +TEST_CASE("moe combine: routed_scale multiplies the ROUTED sum, not the shared term") { + const int64_t T = 2, K = 2, H = 3; + std::vector expert_out = { + // t=0 + 1.0f, 2.0f, 3.0f, // slot 0 + -1.0f, 0.5f, 4.0f, // slot 1 + // t=1 + 2.0f, -2.0f, 1.0f, // slot 0 + 0.25f, 1.5f, -3.0f, // slot 1 + }; + std::vector weights = {0.75f, 0.25f, 0.5f, 0.5f}; + std::vector shared = {10.0f, 20.0f, 30.0f, -1.0f, -2.0f, -3.0f}; + std::vector out(static_cast(T * H), 0.0f); + + Tensor eo = F32_3(expert_out, T, K, H); + Tensor wt = F32_2(weights, T, K); + Tensor sh = F32_2(shared, T, H); + Tensor ot = F32_2(out, T, H); + Queue q = Q(); + const float scale = 2.5f; + vt::MoeCombine(q, ot, eo, wt, &sh, scale); + + for (int64_t t = 0; t < T; ++t) { + for (int64_t h = 0; h < H; ++h) { + float routed = 0.0f; + for (int64_t j = 0; j < K; ++j) { + routed += weights[static_cast(t * K + j)] * + expert_out[static_cast((t * K + j) * H + h)]; + } + const float sv = shared[static_cast(t * H + h)]; + const size_t i = static_cast(t * H + h); + CHECK(Close(out[i], scale * routed + sv)); + // The two misplacements: scaling the shared term too, and dropping the + // scale entirely. Both are wrong by a wide margin on every element here. + CHECK(!Close(out[i], scale * (routed + sv))); + CHECK(!Close(out[i], routed + sv)); + } + } +} + +// The default keeps every landed caller byte-identical: the 5-argument call with +// routed_scale == 1.0f must produce the same bits as the 4-argument call. +TEST_CASE("moe combine: routed_scale defaults to 1.0 (landed callers unchanged)") { + const int64_t T = 2, K = 3, H = 4; + std::vector expert_out(static_cast(T * K * H)); + for (size_t i = 0; i < expert_out.size(); ++i) { + expert_out[i] = 0.125f * static_cast(i) - 1.5f; + } + std::vector weights = {0.5f, 0.25f, 0.25f, 0.125f, 0.375f, 0.5f}; + std::vector shared(static_cast(T * H)); + for (size_t i = 0; i < shared.size(); ++i) shared[i] = 0.5f - 0.25f * static_cast(i); + + std::vector a(static_cast(T * H), 0.0f); + std::vector b(static_cast(T * H), 0.0f); + Tensor eo = F32_3(expert_out, T, K, H); + Tensor wt = F32_2(weights, T, K); + Tensor sh = F32_2(shared, T, H); + Tensor at = F32_2(a, T, H); + Tensor bt = F32_2(b, T, H); + Queue q = Q(); + vt::MoeCombine(q, at, eo, wt, &sh); // landed 4-arg form + vt::MoeCombine(q, bt, eo, wt, &sh, 1.0f); // explicit no-op scale + CHECK(std::memcmp(a.data(), b.data(), a.size() * sizeof(float)) == 0); +} + +// --------------------------------------------------------------------------- +// 3. The whole non-gated expert, end to end on the shared ops. +// --------------------------------------------------------------------------- + +namespace { + +// An independent scalar reference for ONE NemotronH MoE block. Deliberately +// written from the upstream formula rather than from any vt op, so it cannot +// agree with the implementation by sharing a helper. +// h = x @ W_up[e]^T ; h = relu(h)^2 ; y = h @ W_down[e]^T +// out = routed_scale * sum_j w[t,j] * y[t,j] + shared[t] +std::vector NonGatedExpertRef(const std::vector& x, int64_t T, int64_t H, + int64_t I, int64_t E, const std::vector& w_up, + const std::vector& w_down, + const std::vector& ids, + const std::vector& weights, int64_t K, + const std::vector& shared, float routed_scale) { + std::vector out(static_cast(T * H), 0.0f); + (void)E; + for (int64_t t = 0; t < T; ++t) { + std::vector acc(static_cast(H), 0.0f); + for (int64_t j = 0; j < K; ++j) { + const int64_t e = ids[static_cast(t * K + j)]; + std::vector hbuf(static_cast(I), 0.0f); + for (int64_t i = 0; i < I; ++i) { + float s = 0.0f; + for (int64_t k = 0; k < H; ++k) { + s += x[static_cast(t * H + k)] * + w_up[static_cast((e * I + i) * H + k)]; + } + const float r = s > 0.0f ? s : 0.0f; + hbuf[static_cast(i)] = r * r; // relu^2, NOT relu, NOT silu + } + for (int64_t h = 0; h < H; ++h) { + float s = 0.0f; + for (int64_t i = 0; i < I; ++i) { + s += hbuf[static_cast(i)] * + w_down[static_cast((e * H + h) * I + i)]; + } + acc[static_cast(h)] += weights[static_cast(t * K + j)] * s; + } + } + for (int64_t h = 0; h < H; ++h) { + out[static_cast(t * H + h)] = + routed_scale * acc[static_cast(h)] + shared[static_cast(t * H + h)]; + } + } + return out; +} + +float Synth(int64_t a, int64_t b, float k) { + return std::sin(static_cast(a) * 0.7f + static_cast(b) * 0.13f) * k; +} + +} // namespace + +// The W2 deliverable: the expert has NO gate half, so it is the existing SINGLE +// grouped projection plus the relu^2 activation — up_proj -> relu^2 -> down_proj +// -> weighted combine with the routed scale on the OUTPUT. Every step is a +// shared vt:: op; nothing here is a NemotronH-specific MoE path. +TEST_CASE("nemotron-h non-gated expert: up -> relu^2 -> down -> scaled combine") { + const int64_t T = 3, H = 6, I = 4, E = 4, K = 2; + const float routed_scale = 2.5f; // config.routed_scaling_factor + + std::vector x(static_cast(T * H)); + for (int64_t t = 0; t < T; ++t) { + for (int64_t h = 0; h < H; ++h) x[static_cast(t * H + h)] = Synth(t, h, 0.9f); + } + // up_proj weight [E, I, H] and down_proj weight [E, H, I] — torch Linear + // (out, in) orientation, which is exactly vt::MatmulBT's `b [N, K]`. + std::vector w_up(static_cast(E * I * H)); + for (int64_t e = 0; e < E; ++e) { + for (int64_t i = 0; i < I; ++i) { + for (int64_t h = 0; h < H; ++h) { + w_up[static_cast((e * I + i) * H + h)] = Synth(e * I + i, h, 0.4f); + } + } + } + std::vector w_down(static_cast(E * H * I)); + for (int64_t e = 0; e < E; ++e) { + for (int64_t h = 0; h < H; ++h) { + for (int64_t i = 0; i < I; ++i) { + w_down[static_cast((e * H + h) * I + i)] = Synth(e * H + h + 3, i, 0.3f); + } + } + } + const std::vector ids = {0, 1, 2, 3, 1, 0}; + const std::vector weights = {0.6f, 0.4f, 0.7f, 0.3f, 0.5f, 0.5f}; + std::vector shared(static_cast(T * H)); + for (size_t i = 0; i < shared.size(); ++i) shared[i] = 0.05f * static_cast(i) - 0.2f; + + Queue q = Q(); + // Per (token, slot) expert projection through the shared GEMM op, exactly as + // the CPU/GGUF MoE reference loop does for the gated archs. The ONLY structural + // difference from a SwiGLU expert is that there is one projection, not a merged + // pair, and the epilogue is relu^2 instead of silu*up. + std::vector expert_out(static_cast(T * K * H), 0.0f); + for (int64_t t = 0; t < T; ++t) { + for (int64_t j = 0; j < K; ++j) { + const int64_t e = ids[static_cast(t * K + j)]; + Tensor xt = Tensor::Contiguous(&x[static_cast(t * H)], DType::kF32, Cpu(), {1, H}); + std::vector hbuf(static_cast(I), 0.0f); + Tensor ht = F32_2(hbuf, 1, I); + Tensor wu = Tensor::Contiguous(&w_up[static_cast(e * I * H)], DType::kF32, Cpu(), + {I, H}); + vt::MatmulBT(q, ht, xt, wu); + std::vector act(static_cast(I), 0.0f); + Tensor at = F32_2(act, 1, I); + vt::MoeRelu2(q, at, ht); + Tensor yt = Tensor::Contiguous(&expert_out[static_cast((t * K + j) * H)], + DType::kF32, Cpu(), {1, H}); + Tensor wd = Tensor::Contiguous(&w_down[static_cast(e * H * I)], DType::kF32, Cpu(), + {H, I}); + vt::MatmulBT(q, yt, at, wd); + } + } + + std::vector out(static_cast(T * H), 0.0f); + Tensor eo = F32_3(expert_out, T, K, H); + std::vector wcopy = weights; + Tensor wt = F32_2(wcopy, T, K); + Tensor sh = F32_2(shared, T, H); + Tensor ot = F32_2(out, T, H); + vt::MoeCombine(q, ot, eo, wt, &sh, routed_scale); + + const std::vector want = + NonGatedExpertRef(x, T, H, I, E, w_up, w_down, ids, weights, K, shared, routed_scale); + for (size_t i = 0; i < want.size(); ++i) { + CHECK(Close(out[i], want[i], 1e-5f, 1e-6f)); + } + + // The same composite with the GATED family's activation is a DIFFERENT answer: + // proves the block is genuinely sensitive to the activation choice and is not + // dominated by the combine. + const std::vector silu_ref = [&] { + std::vector o(static_cast(T * H), 0.0f); + for (int64_t t = 0; t < T; ++t) { + for (int64_t j = 0; j < K; ++j) { + const int64_t e = ids[static_cast(t * K + j)]; + std::vector hb(static_cast(I), 0.0f); + for (int64_t i = 0; i < I; ++i) { + float s = 0.0f; + for (int64_t k = 0; k < H; ++k) { + s += x[static_cast(t * H + k)] * + w_up[static_cast((e * I + i) * H + k)]; + } + hb[static_cast(i)] = SiluRef(s); + } + for (int64_t h = 0; h < H; ++h) { + float s = 0.0f; + for (int64_t i = 0; i < I; ++i) { + s += hb[static_cast(i)] * + w_down[static_cast((e * H + h) * I + i)]; + } + o[static_cast(t * H + h)] += weights[static_cast(t * K + j)] * s; + } + } + } + for (size_t i = 0; i < o.size(); ++i) o[i] = routed_scale * o[i] + shared[i]; + return o; + }(); + bool any_differs = false; + for (size_t i = 0; i < want.size(); ++i) { + if (!Close(out[i], silu_ref[i], 1e-3f, 1e-6f)) any_differs = true; + } + CHECK(any_differs); +} + +// The scale lives on the OUTPUT, so the ROUTER runs with routed_scaling_factor +// 1.0 (layer.py:291-300) and its weights are the plain renormalized sigmoid +// scores. Scaling the LOGITS instead is a different answer entirely — sigmoid is +// non-linear, so it moves the weights (and can move the SELECTION). +TEST_CASE("nemotron-h routed scale: on the output, not on the router logits") { + const int64_t T = 1, E = 6, K = 3; + std::vector logits = {0.2f, -0.4f, 1.1f, 0.05f, -1.3f, 0.7f}; + std::vector scaled_logits(logits.size()); + const float routed_scale = 2.5f; + for (size_t i = 0; i < logits.size(); ++i) scaled_logits[i] = routed_scale * logits[i]; + + vt::MoeRouterTopKArgs args; + args.top_k = static_cast(K); + args.renormalize = true; // config.norm_topk_prob + args.scoring_func = vt::MoeScoringFunc::kSigmoid; + args.num_expert_group = 1; // config.n_group + args.topk_group = 1; // config.topk_group + // apply_routed_scale_to_output=True => the ROUTER's factor is a nop (1.0). + args.routed_scaling_factor = 1.0f; + std::vector bias(static_cast(E), 0.0f); + bias[2] = -5.0f; // e_score_correction_bias: biases SELECTION only + + std::vector w(static_cast(T * K), 0.0f); + std::vector ids(static_cast(T * K), -1); + Tensor lt = F32_2(logits, T, E); + Tensor wt = F32_2(w, T, K); + Tensor it = Tensor::Contiguous(ids.data(), DType::kI32, Cpu(), {T, K}); + Tensor bt = F32_1(bias); + Queue q = Q(); + vt::MoeRouterTopK(q, wt, it, lt, args, &bt); + + // Renormalized => the weights sum to 1: the routed scale is NOT in them. + float sum = 0.0f; + for (float v : w) sum += v; + CHECK(Close(sum, 1.0f, 1e-6f)); + for (float v : w) CHECK(v <= 1.0f); + + // Scaling the logits produces DIFFERENT weights, so a mis-port that folds the + // scale into the router input cannot pass the block-level gate above. + std::vector w2(static_cast(T * K), 0.0f); + std::vector ids2(static_cast(T * K), -1); + Tensor lt2 = F32_2(scaled_logits, T, E); + Tensor wt2 = F32_2(w2, T, K); + Tensor it2 = Tensor::Contiguous(ids2.data(), DType::kI32, Cpu(), {T, K}); + vt::MoeRouterTopK(q, wt2, it2, lt2, args, &bt); + bool weights_differ = false; + for (size_t i = 0; i < w.size(); ++i) { + if (!Close(w[i], w2[i], 1e-4f, 1e-6f)) weights_differ = true; + } + CHECK(weights_differ); +} + +// --------------------------------------------------------------------------- +// 4. The NVFP4 W4A16 arm's group size (spec §6 named risk). +// --------------------------------------------------------------------------- + +// NemotronH's routed experts are W4A16_NVFP4 with group_size=16. The grouped +// Marlin path takes the group size as an explicit argument whose DEFAULT is the +// NVFP4 one, so 16 is the configuration it already runs (the alternative, 32, +// only comes with mxfp4=true). Pinned here so a later widening of the default +// cannot silently re-point NemotronH's experts at a group size the checkpoint +// does not carry. +TEST_CASE("nvfp4 grouped moe: group_size 16 is the NVFP4 default NemotronH needs") { + vt::MoeMarlinArgs args; + CHECK(args.group_size == 16); + CHECK(args.mxfp4 == false); + vt::MarlinDenseArgs dense; + CHECK(dense.group_size == 16); + CHECK(dense.mxfp4 == false); +}