record(MODEL-NEMOTRON-H W2): the non-gated relu^2 expert, with #586's four review findings repaired (#517)

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

@@ -384,6 +384,13 @@ enum class OpId : uint8_t {
   kMamba2ChunkScan,
   kMamba2StateUpdate,
   kRmsNormGatedGroup,
+  // 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
 };
 
@@ -908,6 +915,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).
@@ -1031,8 +1040,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&,
@@ -1649,6 +1660,30 @@ 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:34 `RELU2_NO_MUL`, :98
+// `activation_without_mul`, and the :184 RELU2_NO_MUL branch; `:33` is
+// GELU_TANH_NO_MUL, the neighbouring enumerator).
+//
+// 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
@@ -2337,7 +2372,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
@@ -2346,8 +2381,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:390-407, :402-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

@@ -461,6 +461,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.
@@ -2475,8 +2492,23 @@ 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:390-407, :402-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.
+// It scales the ASSEMBLED sum, not each router weight: upstream's
+// `fused_output *= factor` (:404) is one multiply on the finished tensor, so the
+// scale rounds ONCE after the K-term reduction. Folding it into `weights[j]` is
+// equal in exact arithmetic and a different f32 value (it rounds K times inside
+// the sum); Laguna is entitled to that fold (`laguna_ops.h:48`, no `shared`),
+// this path is not. Pinned bitwise in test_ops_moe_nongated_relu2.cpp.
+// NOT MIRRORED, UNREACHABLE: upstream's fp16 arm (:403-406) instead divides
+// `shared_output` by the factor to dodge an fp16 overflow. That branch is keyed
+// on `fused_output.dtype == torch.float16`; the analogue here is `out`, whose
+// dtype `MoeCombine` gates through `IsOutFloat` (ops.cpp:22 — f32/bf16 only, no
+// kF16), so no caller can reach it. Pinned by the f16-out refusal test.
 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) {
@@ -2485,6 +2517,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);
     }
@@ -3040,6 +3073,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,