minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix failing test_nvidia_nemotron_3_nano by fixing test_grouped_topk (#23874)

Browse Source
This commit is contained in:
Khoa Pham
2026-04-28 15:03:58 -07:00
committed by GitHub
parent 345fecc547
commit ddcacaf1bd
4 changed files with 223 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/jit_kernel/csrc/moe/grouped_topk.cuh
+11 -18
View File
@@ -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<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 >> 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;
}
}
python/sglang/jit_kernel/tests/test_grouped_topk.py
+210
View File
@@ -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
+2
View File
@@ -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)(