A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 19:45 UTC
test: add nan/inf clamp regression test for fused_topk_bias (#40553)
Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com>
Jhao-Ting Chen committed
Apr 21, 2026 at 17:46 UTC
46794958f0c60bc3a4f30562032e991222ab5d56
1 file changed
+69
tests/kernels/moe/test_fused_topk.py
+69
@@ -202,3 +202,72 @@ def test_fused_topk_nan_inf_clamp(
202
f"Row {row} has non-finite weights {topk_weights[row].tolist()} "
203
f"(bad_value={bad_value}, scoring_func={scoring_func})"
204
)
205
+
206
+
207
+@pytest.mark.skipif(
208
+ not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
209
+)
210
+@pytest.mark.parametrize("num_experts", [6, 8, 16])
211
+@pytest.mark.parametrize("topk", [3, 4])
212
+@pytest.mark.parametrize("scoring_func", ["softmax", "sigmoid"])
213
+@pytest.mark.parametrize("bad_value", [float("nan"), float("inf")])
214
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32])
215
+def test_fused_topk_bias_nan_inf_clamp(
216
+ num_experts: int,
217
+ topk: int,
218
+ scoring_func: str,
219
+ bad_value: float,
220
+ dtype: torch.dtype,
221
+):
222
+ """Regression test: NaN/Inf in gating logits must not produce duplicate
223
+ expert IDs or non-finite weights when e_score_correction_bias is present.
224
+
225
+ Same scenario as test_fused_topk_nan_inf_clamp but exercising the bias
226
+ path (fused_topk_bias) so the fix in topk_softmax_kernels.cu is covered
227
+ for that entry point as well.
228
+ """
229
+ torch.manual_seed(0)
230
+ num_tokens = 4
231
+ hidden_size = 1024
232
+ hidden_states = torch.randn((num_tokens, hidden_size), dtype=dtype, device="cuda")
233
+ e_score_correction_bias = torch.randn(
234
+ (num_experts,), dtype=torch.float32, device="cuda"
235
+ )
236
+
237
+ gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda")
238
+ gating_output[1:, :] = bad_value
239
+
240
+ topk_weights, topk_ids = fused_topk_bias(
241
+ hidden_states=hidden_states,
242
+ gating_output=gating_output,
243
+ e_score_correction_bias=e_score_correction_bias,
244
+ topk=topk,
245
+ renormalize=False,
246
+ scoring_func=scoring_func,
247
+ )
248
+
249
+ # Normal row must still match the torch reference.
250
+ ref_weights, ref_ids = torch_topk(
251
+ gating_output=gating_output[:1],
252
+ topk=topk,
253
+ renormalize=False,
254
+ e_score_correction_bias=e_score_correction_bias,
255
+ scoring_func=scoring_func,
256
+ )
257
+ torch.testing.assert_close(
258
+ ref_weights.to(torch.float32), topk_weights[:1], atol=1e-2, rtol=1e-2
259
+ )
260
+ torch.testing.assert_close(ref_ids.to(torch.int32), topk_ids[:1], atol=0, rtol=0)
261
+
262
+ # Poisoned rows: IDs must be unique (no duplicates) and weights must be
263
+ # finite (no NaN/Inf propagation into downstream MoE kernels).
264
+ for row in range(1, num_tokens):
265
+ row_ids = topk_ids[row]
266
+ assert row_ids.unique().numel() == topk, (
267
+ f"Row {row} has duplicate expert IDs {row_ids.tolist()} "
268
+ f"(bad_value={bad_value}, scoring_func={scoring_func})"
269
+ )
270
+ assert torch.isfinite(topk_weights[row]).all(), (
271
+ f"Row {row} has non-finite weights {topk_weights[row].tolist()} "
272
+ f"(bad_value={bad_value}, scoring_func={scoring_func})"
273
+ )