Fix failing test_nvidia_nemotron_3_nano by fixing test_grouped_topk (#23874)

python/sglang/jit_kernel/csrc/moe/grouped_topk.cuh

@@ -23,12 +23,11 @@ static constexpr int WARP_SIZE = 32;
 static constexpr int MAX_TOPK = 8;
 
 // Pack (value, index) into a single uint64_t for warp-level max reduction.
-// Uses IEEE 754 bit-trick: float bits are order-preserving for positive values.
+// Transform IEEE 754 bits into an unsigned ordering that is monotonic for the
-// Since sigmoid + positive bias yields non-negative scores, this works correctly.
+// full float range; correction bias can make sigmoid(score) + bias negative.
 __device__ __forceinline__ uint64_t pack_val_idx(float val, int32_t idx) {
   uint32_t val_bits = __float_as_uint(val);
-  // Flip sign bit so that comparison works for all floats
+  val_bits ^= (val_bits & 0x80000000u) ? 0xffffffffu : 0x80000000u;
-  val_bits ^= ((val_bits >> 31) | 0x80000000u);
   // Use (65535 - idx) so that smaller indices win ties
   uint32_t idx_bits = static_cast<uint32_t>(65535 - idx);
   return (static_cast<uint64_t>(val_bits) << 32) | idx_bits;
@@ -38,8 +37,7 @@ __device__ __forceinline__ void unpack_val_idx(uint64_t packed, float& val, int3
   uint32_t idx_bits = static_cast<uint32_t>(packed & 0xFFFFFFFF);
   idx = static_cast<int32_t>(65535 - idx_bits);
   uint32_t val_bits = static_cast<uint32_t>(packed >> 32);
-  // Undo the sign-bit flip
+  val_bits ^= (val_bits & 0x80000000u) ? 0x80000000u : 0xffffffffu;
-  val_bits ^= (~(val_bits >> 31) | 0x80000000u);
   val = __uint_as_float(val_bits);
 }
 
@@ -157,20 +155,15 @@ __global__ void grouped_topk_single_group_kernel(
     __syncwarp();
   }
 
-  // Phase 3: renormalize and write output
+  // Phase 3: renormalize and write output. All lanes named by the full-warp
+  // shuffle mask must execute warp_sum_f32 together; inactive lanes contribute
+  // the additive identity.
+  float weight = (lane_id < topk) ? selected_weights[lane_id] : 0.0f;
+  float divisor = renormalize ? warp_sum_f32(weight) + 1e-20f : 1.0f;
+
   if (lane_id < topk) {
-    float weight = selected_weights[lane_id];
-    float final_weight = weight * scaling_factor;
-
-    if (renormalize) {
-      // Warp-level sum of selected weights (only lanes < topk contribute)
-      float partial = (lane_id < topk) ? weight : 0.0f;
-      float total = warp_sum_f32(partial);
-      final_weight = weight * scaling_factor / (total + 1e-20f);
-    }
-
     out_ids[lane_id] = selected_ids[lane_id];
-    out_vals[lane_id] = final_weight;
+    out_vals[lane_id] = weight * scaling_factor / divisor;
   }
 }
 

python/sglang/jit_kernel/tests/test_grouped_topk.py

@@ -0,0 +1,210 @@
+import itertools
+import sys
+
+import pytest
+import torch
+
+from sglang.jit_kernel.grouped_topk import grouped_topk as jit_grouped_topk
+from sglang.jit_kernel.utils import get_ci_test_range
+from sglang.srt.layers.moe.topk import biased_grouped_topk_impl
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, suite="stage-b-kernel-unit-1-gpu-large")
+register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True)
+
+
+CORRECTNESS_CASES = get_ci_test_range(
+    full_range=list(
+        itertools.product(
+            [1, 17, 128],
+            [16, 32, 64, 128, 192, 256, 384, 512],
+            [1, 2, 3, 4, 5, 6, 7, 8],
+        )
+    ),
+    ci_range=[
+        (1, 16, 3),  # smallest non-power-of-two topk
+        (17, 128, 6),  # Nemotron-3-Nano shape that exposed the bug
+        (128, 192, 8),  # Hunyuan-3 shape, power-of-two topk sanity case
+        (33, 512, 7),  # largest expert-count tier with non-power-of-two topk
+    ],
+)
+
+
+def _make_inputs(num_tokens: int, num_experts: int, seed: int):
+    torch.manual_seed(seed)
+    hidden_states = torch.empty((num_tokens, 1), dtype=torch.float32, device="cuda")
+    gating_output = torch.randn(
+        (num_tokens, num_experts), dtype=torch.float32, device="cuda"
+    )
+    correction_bias = torch.randn(num_experts, dtype=torch.float32, device="cuda") * 0.1
+    return hidden_states, gating_output, correction_bias
+
+
+def _scatter_by_expert(
+    weights: torch.Tensor, ids: torch.Tensor, num_experts: int
+) -> torch.Tensor:
+    dense = torch.zeros(
+        (weights.shape[0], num_experts), dtype=torch.float32, device=weights.device
+    )
+    dense.scatter_(1, ids.long(), weights)
+    return dense
+
+
+@pytest.mark.parametrize("num_tokens,num_experts,topk", CORRECTNESS_CASES)
+def test_grouped_topk_renormalize_matches_reference(
+    num_tokens: int, num_experts: int, topk: int
+) -> None:
+    hidden_states, gating_output, correction_bias = _make_inputs(
+        num_tokens, num_experts, seed=1000 + num_experts * 10 + topk
+    )
+    scaling_factor = 2.826 if (num_experts, topk) == (192, 8) else 1.0
+
+    topk_weights, topk_ids = jit_grouped_topk(
+        gating_output,
+        correction_bias,
+        1,
+        1,
+        topk,
+        True,
+        scaling_factor,
+    )
+    ref_weights, ref_ids = biased_grouped_topk_impl(
+        hidden_states,
+        gating_output,
+        correction_bias,
+        topk,
+        True,
+        1,
+        1,
+        routed_scaling_factor=scaling_factor,
+        apply_routed_scaling_factor_on_output=True,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(
+        _scatter_by_expert(topk_weights, topk_ids, num_experts),
+        _scatter_by_expert(ref_weights, ref_ids, num_experts),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+    torch.testing.assert_close(
+        topk_weights.sum(dim=-1),
+        torch.full((num_tokens,), scaling_factor, dtype=torch.float32, device="cuda"),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+
+@pytest.mark.parametrize("topk", [3, 5, 6, 7])
+def test_grouped_topk_non_power_of_two_renormalize(topk: int) -> None:
+    hidden_states, gating_output, correction_bias = _make_inputs(
+        num_tokens=64, num_experts=128, seed=2000 + topk
+    )
+
+    topk_weights, topk_ids = jit_grouped_topk(
+        gating_output,
+        correction_bias,
+        1,
+        1,
+        topk,
+        True,
+        1.0,
+    )
+    ref_weights, ref_ids = biased_grouped_topk_impl(
+        hidden_states,
+        gating_output,
+        correction_bias,
+        topk,
+        True,
+        1,
+        1,
+        routed_scaling_factor=1.0,
+        apply_routed_scaling_factor_on_output=True,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(
+        _scatter_by_expert(topk_weights, topk_ids, 128),
+        _scatter_by_expert(ref_weights, ref_ids, 128),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+    torch.testing.assert_close(
+        topk_weights.sum(dim=-1),
+        torch.ones((64,), dtype=torch.float32, device="cuda"),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+
+def test_grouped_topk_negative_choice_scores_match_reference() -> None:
+    hidden_states, gating_output, correction_bias = _make_inputs(
+        num_tokens=64, num_experts=128, seed=23758
+    )
+    correction_bias.fill_(-2.0)
+
+    topk_weights, topk_ids = jit_grouped_topk(
+        gating_output,
+        correction_bias,
+        1,
+        1,
+        6,
+        True,
+        1.0,
+    )
+    ref_weights, ref_ids = biased_grouped_topk_impl(
+        hidden_states,
+        gating_output,
+        correction_bias,
+        6,
+        True,
+        1,
+        1,
+        routed_scaling_factor=1.0,
+        apply_routed_scaling_factor_on_output=True,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(
+        _scatter_by_expert(topk_weights, topk_ids, 128),
+        _scatter_by_expert(ref_weights, ref_ids, 128),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+
+def test_grouped_topk_without_renormalize_matches_reference() -> None:
+    hidden_states, gating_output, correction_bias = _make_inputs(
+        num_tokens=64, num_experts=128, seed=3006
+    )
+
+    topk_weights, topk_ids = jit_grouped_topk(
+        gating_output,
+        correction_bias,
+        1,
+        1,
+        6,
+        False,
+        1.0,
+    )
+    ref_weights, ref_ids = biased_grouped_topk_impl(
+        hidden_states,
+        gating_output,
+        correction_bias,
+        6,
+        False,
+        1,
+        1,
+    )
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(
+        _scatter_by_expert(topk_weights, topk_ids, 128),
+        _scatter_by_expert(ref_weights, ref_ids, 128),
+        rtol=1e-5,
+        atol=1e-6,
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

python/sglang/srt/models/nemotron_h.py

@@ -923,6 +923,8 @@ def nemotron_mamba2_with_output(
 
     # Copy result back; output may be larger (padded) so only fill actual tokens
     output[:num_actual_tokens].view(ret.shape).copy_(ret)
+    if output.shape[0] != num_actual_tokens:
+        output[num_actual_tokens:].zero_()
 
 
 breakable_nemotron_mamba2_with_output = eager_on_graph(True)(

test/registered/models/test_nvidia_nemotron_3_nano.py

@@ -7,7 +7,6 @@ from sglang.test.server_fixtures.default_fixture import DefaultServerBase
 register_cuda_ci(
     est_time=564,
     suite="stage-b-test-2-gpu-large",
-    disabled="Temporarily disabled; failing on main.",
 )
 
 NEMOTRON_3_NANO_THINKING_ARGS = [