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); +}