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/4] 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/4] 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/4] =?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/4] 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",