MODEL-NEMOTRON-H W2: the non-gated relu² MoE expert — one GEMM, not a merged pair (#517)

.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

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

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),

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<float>()[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<void*>(static_cast<SoftCapFn>(&SoftCapKernel)));
     RegisterOp(OpId::kMoeSiluMul, DeviceType::kCPU,
                reinterpret_cast<void*>(static_cast<MoeSiluMulFn>(&MoeSiluMulKernel)));
+    RegisterOp(OpId::kMoeRelu2, DeviceType::kCPU,
+               reinterpret_cast<void*>(static_cast<MoeRelu2Fn>(&MoeRelu2Kernel)));
     RegisterOp(OpId::kScaledFp4Quant, DeviceType::kCPU,
                reinterpret_cast<void*>(static_cast<ScaledFp4QuantFn>(&ScaledFp4QuantKernel)));
     RegisterOp(OpId::kSiluMulFp4Quant, DeviceType::kCPU,