MODEL-NEMOTRON-H W2 land-prep: the non-gated relu² MoE expert, re-merged onto main and fully re-gated (#517)

.agents/specs/nemotron-h-model.md

@@ -95,7 +95,7 @@ memory format against the oracle explicitly.
 | MoE | `nemotron_h.py:126-256` (`NemotronHMoE`), decoder layer `:317` |
 | non-gated activation | `activation_without_mul(config.mlp_hidden_act)` -> `ReLUSquaredActivation` (`layers/activation.py`) |
 | expert ckpt naming | `ckpt_names=("up_proj", "down_proj", "")` (`nemotron_h.py:220`) |
-| routed scale applied to OUTPUT | `apply_routed_scale_to_output=True` (`nemotron_h.py:246`) |
+| routed scale applied to OUTPUT | `apply_routed_scale_to_output=True` (`nemotron_h.py:234`), factor `:233` |
 | router dtype | `GateLinear(..., out_dtype=torch.float32, force_fp32_compute=True)` (`nemotron_h.py:150-156`) |
 | state shape / dtype | `mamba_utils.py:174-199`, `:73-81` |
 | MTP | `models/nemotron_h_mtp.py::NemotronHMTP` (`registry.py:638`) |
@@ -879,15 +879,148 @@ Read `Status:`.
   non-attention layers and may not move tokens on short prompts. Gate with a
   long-prompt arm, not only a 6-token one.
 
-## 7. Now
+## 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`):
 
-**State at this commit:** spec committed, implementation **not started**. The
-row stays `INVENTORIED`; this commit changes no lifecycle state. The checkpoint
-is staged on the NAS and the oracle smoke run is queued behind the GPU lock.
+```
+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:390-407` (`:402-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.
+
+**CUDA arms — what was actually run, and by whom.** The implementer did NOT
+compile them: their worktree had no `nvcc`, so at `e2d68404` the CUDA arms were
+*written and reviewed*, never built, and the earlier wording here
+("compiled-and-reviewed") overstated it. They have since been compiled and
+GPU-verified **by the fresh reviewer**, on `dgx.casa` (GB10, nvcc 13.0.88), from
+a `git archive` of `e2d68404`:
+
+- Release `-DVLLM_CPP_CUDA=ON -DVLLM_CPP_CUDA_ARCHITECTURES=121a
+  -DVLLM_CPP_CUTLASS_DIR=$HOME/cutlass-4.5.0 -DVLLM_CPP_TRITON=ON` exited 0 with
+  **671/671 targets and zero warnings**; `cuda_moe.cu.o` compiled under
+  `-Werror=all-warnings`.
+- A reviewer-authored GPU parity test proved `MoeRelu2` CUDA == CPU
+  **bit-for-bit** over 4097 elements in all four dtype arms; that CUDA
+  `routed_scale` scales the routed sum only; and that the `1.0f` default is
+  byte-identical to the landed 4-arg call across all 8 dtype combinations.
+- Branch tests on the GPU box: `test_ops_moe_nongated_relu2` 10/10,
+  `test_ops_moe` 9/9 with 33451 assertions.
+
+**Still OWED** (no GPU in the implementer/repair worktrees, and not covered by
+the above): `kMoeGroupedGemmNvfp4Marlin` exercised on the real NemotronH g16
+tensors, and the end-to-end NemotronH MoE block on GB10. Both remain owed to W6
+or an earlier GPU-host spot check. The `group_size` unit test pins the default
+only — it is not a run of the Marlin arm.
+
+**Evidence.** `tests/vt/test_ops_moe_nongated_relu2.cpp` (**12 cases / 81
+assertions**): 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 shape/dtype/**device** contract refusals, the routed
+scale on the routed sum only, the routed scale on the **assembled sum rather than
+each router weight** (bitwise), the f16-out refusal that makes upstream's fp16
+arm unreachable, 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 (Release, `-ffp-contract=off`; every one restored
+and md5-verified afterwards):
+
+| # | Mutation | Target | Result |
+|---|---|---|---|
+| M1 | `relu` (square dropped) | `test_ops_moe_nongated_relu2` | RED 5 cases / 27 assertions |
+| M2 | `silu` (the gated family's activation) | same | RED 5 / 35 |
+| M3 | square narrowed through bf16 | same | RED 2 / 20 |
+| M4 | `routed_scale` dropped | same | RED 3 / 32 |
+| M5 | `routed_scale` applied to routed **+ shared** | same | RED 2 / 29 |
+| M6 | `routed_scale` **folded into each router weight** | same | RED 1 / 4 |
+| M7 | routed scale folded into the router **logits** | `test_ops_moe_router_grouped` | RED 3 / 498 |
+| M8 | NVFP4 `group_size` default 16 → 32 | `test_ops_moe_nongated_relu2` | RED 1 / 1 |
+| M9 | `MoeRelu2` device `VT_CHECK` dropped | same | RED 1 / 2 |
+| M10 | `IsOutFloat` widened to admit `kF16` | same | RED 1 / 2 |
+
+M6 is the one this repair added. At `e2d68404` it **survived green** (10/10
+cases, 71/71 assertions): the landed cases all compared with a tolerance, and the
+fold is exact-arithmetic-equal, so nothing could see it. It is also the most
+likely W4 mistake, because Laguna performs exactly that fold
+(`laguna_ops.h:48`) — legally, since Laguna passes no `shared`. The new case
+pins it bitwise on decimal-grid data whose f32 products carry full mantissas
+(rows separate by 10 and 4 ULP), with a `REQUIRE` that the data separates the two
+forms so the green cannot be vacuous. M7 does NOT red the NemotronH file by
+design — this path forces the router factor to 1.0 (`layer.py:291-300`), so the
+router's own suite is where that defect is visible.
+
+## 7. Now
 
-**Next action:** dispatch fresh implementers for **W1** and **W2** (both
-independent of #496) as soon as `row/KERNEL-SSM-MAMBA-SSD-W1` clears review,
-so the ops-header churn does not collide.
+**State at this commit:** **W1 and W3 have LANDED on `main`; W2 is in
+re-review.** The `MIXED_PRECISION` resolver landed at `1bc5ef82c` (#561) and the
+W3 scaffold at `c6b240edd` (#576); both are merged into this branch, and their
+spec sections (§4's three W1 subsections, §5c/§5d/§5e) are main's, carried here
+byte for byte. The Mamba2 SSD kernel work W1 landed earlier at `47960a009`
+(#496), so `include/vt/ops.h` carries main's
+`kMamba2ChunkScan`/`kMamba2StateUpdate`/`kRmsNormGatedGroup` first and appends
+`kMoeRelu2` after them; no existing op id shifted. W2 (the non-gated `relu²`
+expert, §6a) was reviewed PASS at `e2d68404`, repaired for that review's six
+findings at `dd7a6477d`, and this branch is its land-prep: re-merged onto
+`origin/main` and fully re-gated.
+
+The row stays `INVENTORIED`; this commit changes no lifecycle state, so it owes
+no `STATUS`/`BENCHMARKS` write. **Oracle gateability is CLOSED** — §5a records
+the pinned oracle loading and running the checkpoint on GB10 with three greedy
+goldens committed, so W6 has a denominator whenever it is reached.
+
+**Next action:** land W2 after a fresh scoped re-review, then dispatch **W4**
+(W3 having landed). Three things are carried forward, not resolved here: the two
+OWED GPU items in §6a (`kMoeGroupedGemmNvfp4Marlin` on the real g16 tensors, and
+the end-to-end NemotronH MoE block on GB10), and the OWED GGUF k-quant arm
+tracked as W7 (§5b).
 
 ## 8. Stop conditions
 

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,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
@@ -2337,7 +2370,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 +2379,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),