From e2d68404cdfa719c6d7098221661b5bce6ae012e Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 12 Aug 2026 20:33:33 +0000 Subject: [PATCH 1/5] 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); +} From dd7a6477d915669fbf7347fbb865d26258b049d7 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 12 Aug 2026 22:47:06 +0000 Subject: [PATCH 2/5] fix(MODEL-NEMOTRON-H W2): pin the routed-scale PLACEMENT -- the fold survived green (#517) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Repairs the six findings from the fresh review of `row/MODEL-NEMOTRON-H-W2` @ `e2d68404`. The verdict was PASS; finding 1 is the one that mattered, and it is on exactly the error class the spec names as its top risk. FINDING 1 (the real one). `vt::MoeCombine` scales the ASSEMBLED routed sum: `acc = routed_scale * sum_j w[j]*eo[j]`. The alternative -- folding the factor into each router weight, `acc += (routed_scale*w[j])*eo[j]` -- is equal in exact arithmetic and a DIFFERENT f32 value, because it rounds K times inside the reduction instead of once at the end. Upstream scales the finished tensor (`moe_runner.py:402-406`, `fused_output *= routed_scaling_factor`). That fold SURVIVED the landed suite green -- 10/10 cases, 71/71 assertions -- because every landed case compared with a tolerance and the two forms differ only in the last ulps. It is also the single most likely W4 mistake: Laguna performs exactly that fold (`laguna_ops.h:48`), legally, because Laguna passes no `shared`. Now pinned bitwise, on decimal-grid data whose f32 products carry full mantissas so the orderings actually separate (rows differ by 10 and 4 ULP). A `REQUIRE` asserts the data separates the two forms BEFORE the CHECKs assert which one we implement, so the green cannot be vacuous. The comparison is exact and portable because the project pins `-ffp-contract=off` for every C++ TU (`CMakeLists.txt:42-56`) -- no FMA contraction can make reference and kernel disagree. FINDING 2. The spec claimed the CUDA arms were "compiled-and-reviewed"; they were never compiled by the implementer, whose own handoff says so. Replaced with what is now true and WHO ran it: the fresh REVIEWER compiled and GPU-verified them on dgx.casa (GB10, nvcc 13.0.88) from a `git archive` of `e2d68404` -- Release CUDA build exit 0, 671/671 targets, zero warnings, `cuda_moe.cu.o` 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. Still OWED, and now recorded as owed: `kMoeGroupedGemmNvfp4Marlin` on the real g16 tensors, and the end-to-end MoE block on GB10. FINDING 3. Spec §7 `## Now` said implementation "not started" and "dispatch implementers for W1 and W2" in the same file whose §6a documents W2 as built. Updated to the real state. FINDING 4. Anchor drift, re-verified against the pinned oracle rather than taken on faith -- and the check found one MORE than the review reported: - `apply_routed_scale_to_output=True` :246 -> :234 (spec §2) - `routed_scaling_factor` :232 -> :233 (test header) - `MoEActivation.RELU2_NO_MUL` :33 -> :34 (NOT in the review) - `_maybe_apply_routed_scale_to_output` :389-406 -> :390-407, branch :402-406 `:220`, `:227`, `:98`, `:609-628`, `:672-678`, `:722-725` and `layer.py:291-300` re-checked and correct; left alone. FINDING 5. Upstream's fp16 arm (`:403-406`, divide `shared_output` instead) is genuinely unreachable -- `MoeCombine` gates `out.dtype` through `IsOutFloat` (`ops.cpp:22`), which admits f32/bf16 only. Recorded next to the op AND pinned by a test, so "unreachable" cannot quietly become false: widening `IsOutFloat` to admit `kF16` now REDs instead of silently making an unmirrored upstream branch reachable. FINDING 6. The case named ".../device contract violation" never exercised the device `VT_CHECK`. It does now, in both operand positions. Runnable without a GPU: the wrapper validates devices before `GetOp` dispatches, so a tensor merely LABELLED kCUDA is rejected host-side and never dereferenced. MUTATIONS (Release, each restored and md5-verified afterwards). All 10 RED: M1 relu, square dropped RED 5 cases / 27 assertions M2 silu, the gated family's activation RED 5 / 35 M3 square narrowed through bf16 RED 2 / 20 M4 routed_scale dropped RED 3 / 32 M5 routed_scale on routed + shared RED 2 / 29 M6 routed_scale folded into each weight RED 1 / 4 <- was GREEN before M7 routed scale folded into router logits RED 3 / 498 (router suite) M8 NVFP4 group_size default 16 -> 32 RED 1 / 1 M9 MoeRelu2 device VT_CHECK dropped RED 1 / 2 M10 IsOutFloat widened to admit kF16 RED 1 / 2 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. GATES. Clean Release `-Werror`: exit 0, 1207/1207 targets, zero warnings; full `ctest` 401 tests, 400 pass + 1 skip, `test_engine_core_proc` failed under `-j 8` and passes serially (known parallel flake). Debug arm (NDEBUG absent, asserts live): exit 0, 1207/1207, zero warnings, full `ctest` 401/401 pass. `test_ops_moe_nongated_relu2` 10 cases/71 assertions -> 12 cases/81 assertions. `agent-preflight.sh --staged` exit 0. `test_cpu_x86_llamacpp_floor` first failed with `NO_QUIET_WINDOW after 30s (busy=141%)` -- this box was running ~10 concurrent `cc1plus` from other sessions. `scripts/` and `tests/scripts/` are byte-identical to `origin/main` on this branch, so the gate's inputs are untouched by it; re-run on a quiet box (load 5.18) it is 10/10 OK. Contention, not a defect. No CUDA on this host, so the CUDA edits here are comment-only and the CUDA arms remain covered by the reviewer's GB10 run recorded in §6a. FOLLOWING_AGENTS_PROTOCOL Refs #517. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code] --- .agents/specs/nemotron-h-model.md | 105 +++++++++++++---- include/vt/ops.h | 2 +- src/vt/cpu/cpu_ops.cpp | 13 ++- src/vt/cuda/cuda_moe.cu | 8 +- tests/vt/test_ops_moe_nongated_relu2.cpp | 139 +++++++++++++++++++++-- 5 files changed, 230 insertions(+), 37 deletions(-) diff --git a/.agents/specs/nemotron-h-model.md b/.agents/specs/nemotron-h-model.md index bb1ca552d..1b5a84b23 100644 --- a/.agents/specs/nemotron-h-model.md +++ b/.agents/specs/nemotron-h-model.md @@ -88,7 +88,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`) | @@ -230,7 +230,7 @@ is f32, which is what `LoadF32`/`StoreF32` already are elsewhere in `vt`. (`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`, +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 @@ -244,32 +244,87 @@ already defaults to `group_size = 16` with `mxfp4 = false` (`ops.h`), and (`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). +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 -**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. - -**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 W2 are built and under review.** The Mamba2 +SSD kernel work W1 landed on `main` at `47960a009` (#496), which is merged into +this branch — the `include/vt/ops.h` enum carries main's +`kMamba2ChunkScan`/`kMamba2StateUpdate`/`kRmsNormGatedGroup` first and appends +`kMoeRelu2` after them, so no existing op id shifted. W2 (the non-gated `relu²` +expert, §6a) was reviewed PASS at `e2d68404` and this branch is the repair pass +for that review's six findings. + +The row stays `INVENTORIED`; this commit changes no lifecycle state, so it owes +no `STATUS`/`BENCHMARKS` write. The checkpoint is staged on the NAS and the +oracle smoke run is still queued behind the GPU lock. + +**Next action:** land W2 after a fresh scoped re-review, then dispatch **W3**. +Carry forward the two OWED GPU items named in §6a +(`kMoeGroupedGemmNvfp4Marlin` on the real g16 tensors, and the end-to-end +NemotronH MoE block on GB10). ## 8. Stop conditions diff --git a/include/vt/ops.h b/include/vt/ops.h index ac11122b1..ee9ca2182 100644 --- a/include/vt/ops.h +++ b/include/vt/ops.h @@ -2381,7 +2381,7 @@ void MoeRouterTopK(Queue& q, Tensor& weights, Tensor& indices, const Tensor& log // 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 *= +// (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. diff --git a/src/vt/cpu/cpu_ops.cpp b/src/vt/cpu/cpu_ops.cpp index 1358c516c..a8122d6ab 100644 --- a/src/vt/cpu/cpu_ops.cpp +++ b/src/vt/cpu/cpu_ops.cpp @@ -2493,9 +2493,20 @@ 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 +// — 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, float routed_scale) { const int64_t t = out.shape[0], h = out.shape[1], k = weights.shape[1]; diff --git a/src/vt/cuda/cuda_moe.cu b/src/vt/cuda/cuda_moe.cu index 1480077a2..6810252de 100644 --- a/src/vt/cuda/cuda_moe.cu +++ b/src/vt/cuda/cuda_moe.cu @@ -471,10 +471,16 @@ void MoeRouterTopKKernelCuda(Queue& q, Tensor& weights, Tensor& indices, const T // `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, +// moe_runner.py:390-407 (:402-406) scales `fused_output`, 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. +// Like the CPU reference it scales the ASSEMBLED sum, not each router weight +// (:404 `fused_output *= factor` is one multiply on the finished tensor); the +// fold is equal in exact arithmetic and a different f32 value. Upstream's fp16 +// arm (:403-406, divide `shared_output` instead) is unreachable here — `out` is +// gated to f32/bf16 by `IsOutFloat` (ops.cpp:22). See cpu_ops.cpp for the full +// note. template __global__ void MoeCombineKernel(Tout* out, const Teo* expert_out, const float* weights, const Tsh* shared, int64_t t, int64_t h, int k, diff --git a/tests/vt/test_ops_moe_nongated_relu2.cpp b/tests/vt/test_ops_moe_nongated_relu2.cpp index d72907d11..562675cfa 100644 --- a/tests/vt/test_ops_moe_nongated_relu2.cpp +++ b/tests/vt/test_ops_moe_nongated_relu2.cpp @@ -8,19 +8,22 @@ // 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). +// :234 apply_routed_scale_to_output=True, :233 routed_scaling_factor +// layers/fused_moe/activation.py:98 `activation_without_mul`, :34 +// `MoEActivation.RELU2_NO_MUL`, and :184 `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/runner/moe_runner.py:390-407 +// `_maybe_apply_routed_scale_to_output` — :402-406, `fused_output *= +// routed_scaling_factor` on the ASSEMBLED tensor with the SHARED output left +// UNSCALED, then :722-725 `result = shared_output + fused_output`. The +// :403-406 fp16 arm (divide `shared_output` instead) is unreachable here; +// see the f16-out refusal case below. // 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. @@ -76,6 +79,16 @@ bool Close(float a, float b, float rel = 1e-6f, float abs_tol = 1e-30f) { float SiluRef(float x) { return x / (1.0f + std::exp(-x)); } +// Raw f32 bit pattern. The routed-scale PLACEMENT test below distinguishes two +// expressions that agree in exact arithmetic and differ only in where they +// round, so it can only be gated bitwise — any tolerance loose enough to be a +// tolerance accepts both. +uint32_t Bits(float v) { + uint32_t u = 0; + std::memcpy(&u, &v, sizeof(u)); + return u; +} + } // namespace // --------------------------------------------------------------------------- @@ -170,13 +183,27 @@ TEST_CASE("moe relu2: rejects a shape/dtype/device contract violation") { 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); + + // The DEVICE half of the contract, which the name promises. Runnable without + // a GPU: `vt::MoeRelu2` validates `out.device == q.device && x.device == + // q.device` BEFORE `GetOp` dispatches, so a tensor merely LABELLED kCUDA is + // rejected on the host and its data pointer is never dereferenced. Both + // operand positions are checked — a wrapper that validated only `out` would + // pass the first and fail the second. + std::vector host(8, 1.0f); + Tensor cuda_out = Tensor::Contiguous(host.data(), DType::kF32, Device{DeviceType::kCUDA, 0}, {8}); + Tensor cuda_x = Tensor::Contiguous(host.data(), DType::kF32, Device{DeviceType::kCUDA, 0}, {8}); + std::vector out8(8, 0.0f); + Tensor ot8 = F32_1(out8); + CHECK_THROWS_AS(vt::MoeRelu2(q, cuda_out, xt), std::runtime_error); // out on the wrong device + CHECK_THROWS_AS(vt::MoeRelu2(q, ot8, cuda_x), std::runtime_error); // x on the wrong device } // --------------------------------------------------------------------------- // 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; +// moe_runner.py:402-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). @@ -220,6 +247,71 @@ TEST_CASE("moe combine: routed_scale multiplies the ROUTED sum, not the shared t } } +// The mis-port the case above CANNOT see, and the one most likely to be made: +// folding the factor into each router weight — `acc += (scale*w[j])*eo[j]` — +// instead of scaling the assembled sum — `acc = scale * sum_j w[j]*eo[j]`. +// The two agree in exact arithmetic, so every wide-margin check above passes +// under the fold; they differ in f32 because the fold rounds K times INSIDE the +// reduction instead of once at the end. +// +// Upstream scales the ASSEMBLED tensor: `_maybe_apply_routed_scale_to_output` +// (moe_runner.py:402-404) runs `fused_output *= self.routed_scaling_factor` on +// the finished `fused_output` — one multiply per output element, AFTER the +// K-term reduction that produced it. Laguna folds the same factor into the +// router weights by linearity (`laguna_ops.h:48`) and is entitled to; NemotronH +// takes the literal upstream form, so the fold is a defect here. +// +// The comparison is bitwise and portable: the project pins -ffp-contract=off +// for every C++ TU (CMakeLists.txt:42-56), so both expressions below are IEEE +// as-written, with no FMA contraction to make the reference disagree with the +// kernel. The REQUIRE is what stops this being a vacuous green — it proves the +// chosen data actually SEPARATES the two forms before the CHECKs assert which +// one we implement. +TEST_CASE("moe combine: routed_scale scales the ASSEMBLED sum, not each router weight") { + const int64_t T = 2, K = 4, H = 1; + const float scale = 2.5f; + // Decimal-grid values: their f32 products carry full mantissas, so the K-term + // reduction genuinely rounds and the two orderings separate (row 0 by 10 ULP, + // row 1 by 4 ULP). A dyadic grid like 1/8 makes every product exact and both + // forms bit-identical — which is exactly why the wide-margin case above, on + // such data, cannot catch this. + std::vector expert_out = {2.2f, -1.15f, -1.18f, -1.52f, // t=0, slots 0..3 + 2.72f, -2.55f, 2.06f, -2.21f}; // t=1, slots 0..3 + std::vector weights = {0.95f, 0.32f, 0.37f, 0.72f, + 0.80f, 0.45f, 0.86f, 0.59f}; + 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 ot = F32_2(out, T, H); + Queue q = Q(); + // No shared term: this case isolates the placement of the scale WITHIN the + // routed reduction. The shared-term placement is the case above. + vt::MoeCombine(q, ot, eo, wt, nullptr, scale); + + for (int64_t t = 0; t < T; ++t) { + CAPTURE(t); + // Upstream form: reduce the K terms, THEN scale once. + float acc = 0.0f; + for (int64_t j = 0; j < K; ++j) { + acc += weights[static_cast(t * K + j)] * + expert_out[static_cast(t * K + j)]; + } + const float want = acc * scale; + // Folded form: scale each router weight, then reduce. + float folded = 0.0f; + for (int64_t j = 0; j < K; ++j) { + folded += (scale * weights[static_cast(t * K + j)]) * + expert_out[static_cast(t * K + j)]; + } + // The data separates the two forms — without this the CHECKs below prove + // nothing. + REQUIRE(Bits(want) != Bits(folded)); + CHECK(Bits(out[static_cast(t)]) == Bits(want)); + CHECK(Bits(out[static_cast(t)]) != Bits(folded)); + } +} + // 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)") { @@ -245,6 +337,35 @@ TEST_CASE("moe combine: routed_scale defaults to 1.0 (landed callers unchanged)" CHECK(std::memcmp(a.data(), b.data(), a.size() * sizeof(float)) == 0); } +// The ONE arm of `_maybe_apply_routed_scale_to_output` this port does not +// mirror: moe_runner.py:403-406 keys on `fused_output.dtype == torch.float16` +// and, when a `shared_output` exists, divides the SHARED term by the factor +// instead of multiplying the fused one (fp16 overflow protection; the decoder +// layer compensates). It is not mirrored because it is UNREACHABLE, not because +// it was missed: the analogue of `fused_output` here is `out`, and MoeCombine +// gates `out.dtype` through `IsOutFloat` (src/vt/ops.cpp:22), which admits f32 +// and bf16 ONLY. This case pins that gate, so the "unreachable" claim next to +// the kernel cannot quietly become false if `IsOutFloat` is ever widened to +// kF16 — that widening makes an unmirrored upstream branch reachable and must +// fail here rather than in a token gate. +TEST_CASE("moe combine: f16 out is refused, which is why upstream's fp16 arm is unreachable") { + const int64_t T = 1, K = 2, H = 2; + std::vector expert_out = {1.0f, 2.0f, 3.0f, 4.0f}; + std::vector weights = {0.5f, 0.5f}; + std::vector out_f16(static_cast(T * H), 0); + std::vector shared(static_cast(T * H), 1.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 = Tensor::Contiguous(out_f16.data(), DType::kF16, Cpu(), {T, H}); + Queue q = Q(); + CHECK_THROWS_AS(vt::MoeCombine(q, ot, eo, wt, &sh, 2.5f), std::runtime_error); + // ... and with no shared term either: the refusal is on the OUT dtype, which + // is what makes the fp16 branch unreachable regardless of `shared`. + CHECK_THROWS_AS(vt::MoeCombine(q, ot, eo, wt, nullptr, 2.5f), std::runtime_error); +} + // --------------------------------------------------------------------------- // 3. The whole non-gated expert, end to end on the shared ops. // --------------------------------------------------------------------------- From 190cb6aa5ac08185df20c23f2aeb8a6680de41c0 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 12 Aug 2026 23:22:44 +0000 Subject: [PATCH 3/5] =?UTF-8?q?record(MODEL-NEMOTRON-H):=20reconcile=20?= =?UTF-8?q?=C2=A77=20after=20main=20closed=20oracle=20gateability=20(#517)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `main` gained §5a and §5b in `af8170154` while this repair branch was gating: the pinned oracle now demonstrably LOADS AND RUNS Nemotron-3.5-Lightning on GB10, with three greedy goldens committed, and the GGUF k-quant arm is recorded as OWED and tracked as W7. The merge took both sections cleanly, but left §7 `## Now` saying "the oracle smoke run is still queued behind the GPU lock" -- false as of `af8170154`, and the exact defect finding 3 of this repair was raised to remove. Fixing the same class twice in one branch would be a poor look; leaving it would be worse. Per the record rule, main's version was taken wholesale and only the scoped §7 sentence reapplied. §7 now states oracle gateability is CLOSED and carries forward three open items rather than two: the two OWED GPU items in §6a (`kMoeGroupedGemmNvfp4Marlin` on the real g16 tensors, the end-to-end MoE block on GB10) and the OWED GGUF arm (W7, §5b). No lifecycle state changes; the row stays `INVENTORIED`, so this owes no `STATUS`/`BENCHMARKS` write. FOLLOWING_AGENTS_PROTOCOL Refs #517. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code] --- .agents/specs/nemotron-h-model.md | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/.agents/specs/nemotron-h-model.md b/.agents/specs/nemotron-h-model.md index ccacd7c6a..ae42afcf1 100644 --- a/.agents/specs/nemotron-h-model.md +++ b/.agents/specs/nemotron-h-model.md @@ -371,13 +371,14 @@ expert, §6a) was reviewed PASS at `e2d68404` and this branch is the repair pass for that review's six findings. The row stays `INVENTORIED`; this commit changes no lifecycle state, so it owes -no `STATUS`/`BENCHMARKS` write. The checkpoint is staged on the NAS and the -oracle smoke run is still queued behind the GPU lock. +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 **W3**. -Carry forward the two OWED GPU items named in §6a -(`kMoeGroupedGemmNvfp4Marlin` on the real g16 tensors, and the end-to-end -NemotronH MoE block on GB10). +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 From 2b96eaf74aaa919e23a5d4e3d40a2b0fd45b3488 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Wed, 12 Aug 2026 23:23:22 +0000 Subject: [PATCH 4/5] fix(parity): main does not build -- Nemotron35LightningSnapshot calls HfSnapshot before it is declared (#546) `origin/main` @ `8b00f79f2` does not compile. `tests/parity/hf_snapshot.h` defines `parity::Nemotron35LightningSnapshot()` at `:51-57`, which calls `HfSnapshot(...)` at `:52`, but `HfSnapshot` is not declared until `:62`. Introduced by `af8170154` (#517), which inserted the new helper ABOVE `HfSnapshot` instead of below it. Every other snapshot helper in the file sits after the `HfSnapshot` definition, so this is a placement slip, not a design question -- the fix restores the file's own convention. Found while re-gating `row/MODEL-NEMOTRON-H-W2-FIX` after merging `main`. It is NOT this branch's defect, and the check that proves it is one line: the header is byte-identical to `origin/main` here, and it fails to compile ON ITS OWN, with no build system involved: printf '#include "parity/hf_snapshot.h"\nint main(){return 0;}\n' > /tmp/h.cpp g++ -std=c++20 -I tests -fsyntax-only /tmp/h.cpp tests/parity/hf_snapshot.h:52:10: error: 'HfSnapshot' was not declared in this scope 52 | return HfSnapshot("models--nvidia--NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", | ^~~~~~~~~~ In a normal configure it surfaces as a failed TU rather than as a header problem, which is what makes it easy to misattribute to whichever branch hits it first: FAILED: tests/CMakeFiles/test_qwen36_weights.dir/vllm/test_qwen36_weights.cpp.o Every target including `parity/hf_snapshot.h` is affected; a full `cmake --build` stops at ~413/1207. Fixed by MOVING the helper and its comment block below the `HfSnapshot` definition. Pure relocation -- 13 insertions, 13 deletions, no semantic change, no revision string or env-override spelling touched. After it, a clean Release `-Werror` build is exit 0, 1207/1207 targets, zero warnings, and full `ctest` is 401/401. SCOPE NOTE, flagged for the reviewer: `tests/parity/hf_snapshot.h` is outside the authority this repair task was given. It is repaired here rather than deferred because AGENTS.md requires a bug found mid-flow to be filed AND fixed in the same flow, and because no gate on this branch -- or any other -- can be run while `main` does not build. Issue #546 was filed before the fix. A reviewer who would rather see this land on its own row should say so and it will be split out. FOLLOWING_AGENTS_PROTOCOL Fixes #546. Refs #517. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code] --- tests/parity/hf_snapshot.h | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/tests/parity/hf_snapshot.h b/tests/parity/hf_snapshot.h index 7d4018616..1f6a02aeb 100644 --- a/tests/parity/hf_snapshot.h +++ b/tests/parity/hf_snapshot.h @@ -41,19 +41,6 @@ inline constexpr const char* kQwen27NvfP4Revision = inline constexpr const char* kNemotron35LightningNvfP4Revision = "29f2d1746d8f41e316523194b19018707749b1b1"; -// The Nemotron-3.5-Lightning gate model (#517). Unlike the Qwen pins above, -// this one is NOT in the HF cache: it is staged on the NAS as a `local_dir` -// snapshot at `$CHECKPOINT_ROOT/nemotron-3.5-lightning-30b-nvfp4`, so there is -// no `models--org--name/snapshots/` layout to resolve. The env override is -// therefore the ONLY reachable path, and the cache spelling below exists so the -// revision still names what the golden belongs to. Absent env var => "" => the -// caller emits its loud SKIP, which is the intended behavior off the gate host. -inline std::string Nemotron35LightningSnapshot() { - return HfSnapshot("models--nvidia--NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", - kNemotron35LightningNvfP4Revision, - "VT_NEMOTRON35_SNAPSHOT"); -} - // Snapshot directory for `` at `revision`, or "" when it is not cached // (the caller then emits its loud SKIP). `env_override`, when set and non-empty, // names an explicit snapshot directory for a deliberate different-checkpoint @@ -78,6 +65,19 @@ inline std::string HfSnapshot(const char* repo_dir, const char* revision, return snap.string(); } +// The Nemotron-3.5-Lightning gate model (#517). Unlike the Qwen pins below, +// this one is NOT in the HF cache: it is staged on the NAS as a `local_dir` +// snapshot at `$CHECKPOINT_ROOT/nemotron-3.5-lightning-30b-nvfp4`, so there is +// no `models--org--name/snapshots/` layout to resolve. The env override is +// therefore the ONLY reachable path, and the cache spelling below exists so the +// revision still names what the golden belongs to. Absent env var => "" => the +// caller emits its loud SKIP, which is the intended behavior off the gate host. +inline std::string Nemotron35LightningSnapshot() { + return HfSnapshot("models--nvidia--NVIDIA-Nemotron-3.5-Lightning-30B-A3B-NVFP4", + kNemotron35LightningNvfP4Revision, + "VT_NEMOTRON35_SNAPSHOT"); +} + // The 27B NVFP4 gate model, pinned to the goldens' revision. inline std::string Qwen27NvfP4Snapshot() { return HfSnapshot("models--unsloth--Qwen3.6-27B-NVFP4", From f2faefb920ab1c78477ad7125384dfa6a0f9ff18 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 13 Aug 2026 07:01:14 +0000 Subject: [PATCH 5/5] record(MODEL-NEMOTRON-H W2): the four PASS-with-findings repairs, all record and comment (#517) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fresh review of PR #586 (`row/MODEL-NEMOTRON-H-W2-LAND` @ `f6a7f8709`) returned PASS -- the code is correct and the previously-unreviewed repair delta changes zero executable lines. Its four findings were record and comment accuracy. Each was re-verified against the file or the measurement before being changed; none was taken on the reviewer's word. F3 -- the `:33` -> `:34` anchor fix had landed in 1 of 3 places. At the pinned oracle (`5559679229bc`, `vllm/model_executor/layers/fused_moe/activation.py`, verified byte-identical to the pin with `git diff 5559679229bc -- `), `:34` is `RELU2_NO_MUL` and `:33` is `GELU_TANH_NO_MUL`. Fixed at `include/vt/ops.h:1669` and in §6a of the spec; `:98` (`activation_without_mul`) and `:184` (`apply_moe_activation`'s branch) were re-checked and are correct, and `tests/vt/test_ops_moe_nongated_relu2.cpp:12` already carried the right numbers, so it is unchanged. F4 -- §6a's M7 evidence did not reproduce. Recorded `RED 3 / 498`; re-measured here as **3 cases / 525 assertions** failing out of 14 / 941, agreeing with the review's and the operator's independent reproductions. Third identical measurement, so the count is stable and 498 was a transcription slip, not drift: `tests/vt/test_ops_moe_router_grouped.cpp` and the grouped-router kernel have not moved since `e2d68404`. Mutation: fold `routed_scaling_factor` into the logits fed to sigmoid/softmax (`cpu_ops.cpp:2325-2339`) and disable the post-renormalize weight scale (`:2430-2434`); tree restored and md5-verified (`cd409b9465c00834be373cf3ecfb4c1d`), rebuilt, back to 14/941 SUCCESS. F5 -- §4 contradicted §6a. §4's W2 row still described the arm as going "through `MlpGateUpMethodBase` / `vt::MergedGemmGroup`" -- the plan §6a correctly refutes, since a non-gated expert has no pair to merge and an arity-1 descriptor would name a fusion that does not exist. §4's Content cell now records that the planned routing was refuted during implementation and points at §6a as the authority; its Gate column was already satisfied and is untouched. F6 -- an overstated bit-identity comment. `cuda_moe.cu` claimed the combine "stays bit-for-bit equal" between CPU and CUDA. Nothing sets nvcc `--fmad=false`: `CMakeLists.txt` pins `-ffp-contract=off` for CXX (`:55`), HIP (`:393`) and OBJCXX (`:467`) only, and `:41-56` carves CUDA out deliberately ("GPU parity tests compare GPU-vs-GPU"). Device code may therefore contract the `acc += w * Load(...)` reduction into an FMA where the host may not -- this project has already observed nvcc doing exactly that (`.agents/benchmark-record.md:532`). The comment is narrowed to what holds: the `routed_scale` step is one standalone multiply on the finished accumulator with nothing to contract into, and `MoeRelu2` is a compare-select plus one multiply, so both are bit-identical by construction; the reduction is not covered. The flag gap itself is repo-wide and belongs to **#591**, not to this row; the pre-existing `cuda_moe.cu:467` claim about that same reduction is main's line and is left for #591. Also corrected in the record: four of this branch's land-prep merge messages inferred "so the CUDA arm GPU-verified at `e2d68404` is unchanged" from "`vt/ops.h` is untouched by this merge". Each per-merge half is true (`git diff --shortstat ^1 -- include/vt/ops.h` is empty for `57e4301489`, `40908db153`, `6ac8eed85e`, `721f44d38d`, `f6a7f87090`) but the conclusion is false: since `e2d68404`, `include/vt/ops.h` has moved **+180/-1** and `src/vt/cuda/cuda_moe.cu` **+7/-1**, the header movement being `d3cd442e6`, the Mamba2 SSD W1 merge, which touched `ops.h` by +179 lines against its first parent. §6a now records that, and names the PR's `cuda-fat-build` CI job as the standing evidence for the CUDA arm at this head. No dgx recompile was attempted: `dgx.casa` is unreachable. EXECUTABLE-LINE DIFF IS EMPTY. `git diff -U0 -- include src`, filtered of comment and blank lines, produces nothing; the whole delta is comments and the spec. The ten mutations therefore cover unchanged code and were not re-run, except M7, which is the finding. Gates (CPU box, no GPU; `df -h` 67G avail throughout, so no silent `No space left on device`): - clean-tree Release rebuild from an empty build dir: exit 0, **414 targets, ZERO warnings** (`-Werror` on by default per `cmake/CompilerWarnings.cmake`). - `test_ops_moe_nongated_relu2` 12 cases / 81 assertions, Status SUCCESS. - `test_ops_moe_router_grouped` 14 / 941, Status SUCCESS. - `test_hf_snapshot_pinning` 11 / 42, Status SUCCESS. - Full `ctest -j4`: **100% tests passed, 0 failed out of 404** (skipped: `test_modelopt_mixed_precision_checkpoint`, `test_voxtral_e2e` -- neither has its asset here). Re-run on the clean rebuild, same result. - `scripts/agent-preflight.sh --staged`: every gate ok except `commit-trailers` on `656b04dba7fd`, which is the pre-existing merge-commit trailer shape this branch inherits and which the squash-merge resolves (review finding F2, the operator's landing decision, deliberately not repaired here). FOLLOWING_AGENTS_PROTOCOL Refs #517. Refs #591. Following-Agents-Protocol: true AI-Assisted: true Assisted-by: AGENT:claude-opus-5 [Claude Code] --- .agents/specs/nemotron-h-model.md | 74 +++++++++++++++++++++++++++++-- include/vt/ops.h | 4 +- src/vt/cuda/cuda_moe.cu | 15 +++++-- 3 files changed, 85 insertions(+), 8 deletions(-) diff --git a/.agents/specs/nemotron-h-model.md b/.agents/specs/nemotron-h-model.md index 461736a8c..9c1a85f67 100644 --- a/.agents/specs/nemotron-h-model.md +++ b/.agents/specs/nemotron-h-model.md @@ -179,7 +179,7 @@ cannot. | W | Content | Gate | Depends on | |---|---|---|---| | **W1** | ModelOpt `MIXED_PRECISION` resolver: parse `quantization_config`, resolve per-module `quant_algo` (direct then shard-prefix), expose it to weight loading. Refuse an unknown algo by name | unit tests on the REAL `config.json` (committed as a fixture, weights not needed): every one of the 5981 entries resolves, the `ignore` list resolves to unquantized, an unknown algo refuses | — | -| **W2** | Non-gated `relu²` grouped MoE through `MlpGateUpMethodBase` / `vt::MergedGemmGroup`; bf16 arm then NVFP4 W4A16 g16 | byte/tolerance tests vs a host reference; `relu²` mutation caught; routed scale applied to the OUTPUT, not the logits | — | +| **W2** | Non-gated `relu²` grouped MoE; bf16 arm then NVFP4 W4A16 g16. **As built it is NOT a merged pair** — the planned `MlpGateUpMethodBase` / `vt::MergedGemmGroup` routing was refuted during implementation and the arm is the EXISTING grouped GEMM plus a new `vt::MoeRelu2`; §6a is the authority on the seam | byte/tolerance tests vs a host reference; `relu²` mutation caught; routed scale applied to the OUTPUT, not the logits | — | | **W3** | `nemotron_h_weights.cpp` + `_registry.cpp`: `layers_block_type` dispatch, `backbone.` prefix, het-KV group construction (1 Mamba group + 1 full-attn group over 6 layers), enumeration gate vs the released index | enumeration: every tensor in `model.safetensors.index.json` is claimed or explicitly refused; KV spec shapes match `mamba2_state_shape` | #496 W1 | | **W4** | `nemotron_h.cpp` forward: hybrid layer loop, Mamba2 mixer wiring, 6 attention layers, MoE layers | CPU forward runs; per-layer activations vs a dumped oracle reference | #496 W1, W2, W3 | | **W5** | MTP head (`mtp.layers.0`, `eh_proj`/`enorm`/`hnorm`) on the existing spec-decode seam | draft acceptance non-zero; spec-off and spec-on token-identical | W4 | @@ -908,7 +908,8 @@ 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 +(`layers/fused_moe/activation.py:34`; `activation_without_mul` is `:98`, the +enumerator at `:33` is `GELU_TANH_NO_MUL`) → `apply_moe_activation`'s (`:184`) `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 @@ -954,6 +955,47 @@ a `git archive` of `e2d68404`: - Branch tests on the GPU box: `test_ops_moe_nongated_relu2` 10/10, `test_ops_moe` 9/9 with 33451 assertions. +**That GPU build is NOT a proof about this head, and the merge messages that +implied otherwise were wrong (corrected 2026-08-13).** Four of this branch's +land-prep merge commits carry a sentence of the form "`src/vt/cuda/cuda_moe.cu` +and `vt/ops.h` are untouched by this merge, so the CUDA arm the fresh reviewer +compiled and GPU-verified on GB10 at `e2d68404` is unchanged." Each half is +individually true — `git diff --shortstat ^1 -- include/vt/ops.h` +is empty for `57e4301489`, `40908db153`, `6ac8eed85e`, `721f44d38d` and +`f6a7f87090` — but the CONCLUSION does not follow, because the compile inputs +moved in commits those merges do not cover. Measured on this branch: + +| Since `e2d68404` (the GPU-verified tree) | Delta | +|---|---| +| `include/vt/ops.h` | **+180 / −1** (`git diff --numstat e2d68404 HEAD`) | +| `src/vt/cuda/cuda_moe.cu` | **+7 / −1** | + +The header movement is `d3cd442e6`, the merge that brought Mamba2 SSD W1 +(`fe81bd3b5`, #496) — a merge that DID touch `include/vt/ops.h`, by +179 lines +against its first parent (`git show --stat d3cd442e6`) — and the `.cu` movement +is this row's own `dd7a6477d` repair. So `cuda_moe.cu` has been compiled at +`e2d68404` and nowhere since, and no local GPU host is reachable to redo it +(`dgx.casa` is down). The standing evidence for the CUDA arm at THIS head is +therefore the PR's `cuda-fat-build` CI job, which compiles `cuda_moe.cu` +against the merged `include/vt/ops.h` for 10 architectures under +`-Werror=all-warnings`; the GB10 run remains the evidence for the arm's +RUNTIME behaviour at `e2d68404`, which the eight-line `.cu` delta since is a +comment change plus the `routed_scale` parameter the reviewer's own GPU parity +test exercised. + +**Scope of the "bit-for-bit" claim.** The reviewer's GPU parity result is a +MEASUREMENT on GB10, not a guarantee the build flags provide: `-ffp-contract=off` +is pinned for CXX, HIP and OBJCXX (`CMakeLists.txt:55`, `:393`, `:467`) but +nothing passes nvcc `--fmad=false`, and `CMakeLists.txt:41-56` carves CUDA out +deliberately ("GPU parity tests compare GPU-vs-GPU"). `MoeRelu2` is safe by +construction — a compare-select and one multiply, with no multiply-add to +contract — and the `routed_scale` step is one standalone multiply on the +finished accumulator, which is why §6a's evidence is stated for exactly those +two. `MoeCombine`'s `acc += w * Load(...)` reduction is the contractable one +and is NOT covered; the comment in `cuda_moe.cu` claiming it was has been +narrowed to what holds. The flag gap itself is repo-wide and tracked as **#591** +(found independently in `cuda_mamba2_ssd.cuh`); it is not repaired here. + **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 @@ -981,7 +1023,7 @@ and md5-verified afterwards): | 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 | +| M7 | routed scale folded into the router **logits** | `test_ops_moe_router_grouped` | RED 3 / **525** (was recorded 498 — see below) | | 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 | @@ -997,6 +1039,22 @@ 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. +**M7's assertion count was recorded wrong (repaired 2026-08-13).** The table and +the `dd7a6477d` commit message both said `RED 3 / 498`. Three independent +re-measurements — the fresh review's, the operator's, and this repair's — all +read **3 cases / 525 assertions** failing, out of the unchanged 14 / 941 +baseline. The CASE count was right; only the assertion count was wrong, and 498 +is reproducible by nobody. `tests/vt/test_ops_moe_router_grouped.cpp` and the +grouped-router kernel it exercises have not moved since `e2d68404` +(`git log --oneline e2d68404..HEAD -- tests/vt/test_ops_moe_router_grouped.cpp` +is empty), so the count is stable and the original entry was a transcription +slip, not drift. Re-measured here by folding +`args.routed_scaling_factor` into the logits fed to sigmoid/softmax in +`MoeRouterGroupedTopKKernel` (`cpu_ops.cpp:2325-2339`) and disabling the +post-renormalize weight scale (`:2430-2434`); Release, `-ffp-contract=off`; +tree restored and md5-verified (`cd409b9465c00834be373cf3ecfb4c1d`), rebuilt, +back to 14/941 SUCCESS. + ## 7. Now **State at this commit:** **W1 and W3 have LANDED on `main`; W2 is in @@ -1009,7 +1067,15 @@ byte for byte. The Mamba2 SSD kernel work W1 landed earlier at `47960a009` `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. +`origin/main` and fully re-gated. A second fresh review (PR #586 @ `f6a7f8709`) +returned **PASS** — it established that the repair delta changes zero executable +lines — with four RECORD findings, all repaired on +`row/MODEL-NEMOTRON-H-W2-RECORDS`: two stale `activation.py:33` anchors +(`include/vt/ops.h`, this spec), M7's assertion count (§6a), §4's W2 seam text +contradicting §6a, and an overstated bit-identity comment in `cuda_moe.cu` +(#591). That repair touches comments and this spec only; no executable line +moved. `tests/vt/test_ops_moe_nongated_relu2.cpp:12` already carried the +corrected anchors and needed no change. The row stays `INVENTORIED`; this commit changes no lifecycle state, so it owes no `STATUS`/`BENCHMARKS` write. **Oracle gateability is CLOSED** — §5a records diff --git a/include/vt/ops.h b/include/vt/ops.h index ee9ca2182..a3b7df619 100644 --- a/include/vt/ops.h +++ b/include/vt/ops.h @@ -1666,7 +1666,9 @@ void MoeSiluMul(Queue& q, Tensor& out, const Tensor& gate, const Tensor& up); // 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). +// (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","")` diff --git a/src/vt/cuda/cuda_moe.cu b/src/vt/cuda/cuda_moe.cu index 6810252de..3d2991356 100644 --- a/src/vt/cuda/cuda_moe.cu +++ b/src/vt/cuda/cuda_moe.cu @@ -473,8 +473,14 @@ void MoeRouterTopKKernelCuda(Queue& q, Tensor& weights, Tensor& indices, const T // — upstream's apply_routed_scale_to_output arm (layers/fused_moe/runner/ // moe_runner.py:390-407 (:402-406) scales `fused_output`, 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. +// reference (cpu_ops.cpp MoeCombineKernel) uses, in the same order. The scale +// itself is ONE standalone f32 multiply on the finished accumulator, with +// nothing adjacent to contract into, so THIS step is bit-identical to the CPU +// reference. That does not extend to the `acc += w * Load(...)` reduction above +// it: only CXX/HIP/OBJCXX carry -ffp-contract=off (CMakeLists.txt:55, :393, +// :467) and nothing passes nvcc --fmad=false, so device code may contract that +// multiply-add into an FMA where the host may not. See #591, which tracks that +// repo-wide flag gap. Default 1.0f == the landed fold-into-weights arm. // Like the CPU reference it scales the ASSEMBLED sum, not each router weight // (:404 `fused_output *= factor` is one multiply on the finished tensor); the // fold is equal in exact arithmetic and a different f32 value. Upstream's fp16 @@ -697,7 +703,10 @@ void MoeSiluMulKernelCuda(Queue& q, Tensor& out, const Tensor& gate, const Tenso // (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. +// reference (cpu_ops.cpp MoeRelu2Kernel): both are exact f32 ops, no expf — +// and unlike the combine reduction above there is no multiply-add here for +// nvcc to contract into an FMA (a compare-select and ONE multiply), so this +// one holds without --fmad=false. See #591 for the flag gap itself. template __global__ void MoeRelu2Kernel(Tout* out, const Tx* x, int64_t n) { const int64_t step = static_cast(gridDim.x) * blockDim.x;