diff --git a/python/sglang/jit_kernel/csrc/moe/grouped_topk.cuh b/python/sglang/jit_kernel/csrc/moe/grouped_topk.cuh index 19677e0e7..69c8d3b57 100644 --- a/python/sglang/jit_kernel/csrc/moe/grouped_topk.cuh +++ b/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. -// Since sigmoid + positive bias yields non-negative scores, this works correctly. +// Transform IEEE 754 bits into an unsigned ordering that is monotonic for the +// 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 >> 31) | 0x80000000u); + val_bits ^= (val_bits & 0x80000000u) ? 0xffffffffu : 0x80000000u; // Use (65535 - idx) so that smaller indices win ties uint32_t idx_bits = static_cast(65535 - idx); return (static_cast(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(packed & 0xFFFFFFFF); idx = static_cast(65535 - idx_bits); uint32_t val_bits = static_cast(packed >> 32); - // Undo the sign-bit flip - val_bits ^= (~(val_bits >> 31) | 0x80000000u); + val_bits ^= (val_bits & 0x80000000u) ? 0x80000000u : 0xffffffffu; 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; } } diff --git a/python/sglang/jit_kernel/tests/test_grouped_topk.py b/python/sglang/jit_kernel/tests/test_grouped_topk.py new file mode 100644 index 000000000..5c3289094 --- /dev/null +++ b/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"])) diff --git a/python/sglang/srt/models/nemotron_h.py b/python/sglang/srt/models/nemotron_h.py index 1b564dd97..b858950c9 100644 --- a/python/sglang/srt/models/nemotron_h.py +++ b/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)( diff --git a/test/registered/models/test_nvidia_nemotron_3_nano.py b/test/registered/models/test_nvidia_nemotron_3_nano.py index 4834265a2..b3bfc24bd 100644 --- a/test/registered/models/test_nvidia_nemotron_3_nano.py +++ b/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 = [