From ac2920040fd476acf75c71a2fc4c80d499adcb56 Mon Sep 17 00:00:00 2001 From: Lori Ren Date: Thu, 20 Aug 2026 03:56:47 +0000 Subject: [PATCH 1/2] [https://nvbugs/6550099][fix] Raise no-top-k equivalence tolerance above top-p renorm fp32 precision Signed-off-by: Lori Ren --- tests/integration/test_lists/waives.txt | 1 - .../hw_agnostic/test_advanced_sampling_mode.py | 8 ++++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 1265c7b21161..3c0778f7bb7c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -362,7 +362,6 @@ unittest/_torch/multi_gpu/test_linear.py::test_row_linear_norm_fusion[2-hidden:1 unittest/_torch/sampler -k "not test_speculative_d2h_parity_real_predictor" SKIP (https://nvbugs/6619882) unittest/_torch/sampler/test_beam_search.py::test_beam_search_e2e[multi_process-TRTLLMSampler-cuda_graph_and_overlap-None-1-1-True-True-False] SKIP (https://nvbugs/6463819) unittest/_torch/sampler/test_trtllm_sampler.py::test_trtllm_sampler_best_of_with_logprobs SKIP (https://nvbugs/6487837) -unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py::test_no_topk_matches_full[0.9] SKIP (https://nvbugs/6550099) unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b[False] SKIP (https://nvbugs/6535767) unittest/_torch/speculative/hw_agnostic/test_dflash.py::test_dflash_qwen3_5_4b[True] SKIP (https://nvbugs/6535767) unittest/_torch/speculative/hw_agnostic/test_ngram.py::test_llama_ngram[True-True-TRTLLM] SKIP (https://nvbugs/6507102) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py b/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py index 9c760526ff79..beecf2bcc9c5 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py @@ -85,9 +85,9 @@ def test_resolve_advanced_sampling_filters(mode, expect_top_k_none, expect_top_p def test_no_topk_matches_full(top_p_val): """With top_k disabled, NO_TOPK skips the top_k mask kernel (a no-op at k=vocab) and yields the same sampling distribution as FULL. We compare the resulting - probability distributions rather than the sampled tokens: the flashinfer top_k mask - at k=vocab injects ~1e-8 fp noise that leaves the distribution unchanged but can flip - an individual sampled token across GPU archs, so exact-token equality is not portable. + probability distributions rather than the sampled tokens, and with a tolerance above + the top-p renorm's fp32 precision: at the nucleus cutoff that precision decides + whether one token is kept, worth up to ~8e-5 of mass for this batch. A real (non-no-op) filter would move mass by orders of magnitude, far above atol.""" dev = "cuda" torch.manual_seed(0) @@ -105,7 +105,7 @@ def test_no_topk_matches_full(top_p_val): ) probs_full = su.compute_probs_from_logits(logits.clone(), temperatures, ek_full, ep_full) probs_no_topk = su.compute_probs_from_logits(logits.clone(), temperatures, ek_nt, ep_nt) - assert torch.allclose(probs_full, probs_no_topk, atol=1e-5, rtol=0) + assert torch.allclose(probs_full, probs_no_topk, atol=1e-4, rtol=0) @pytest.mark.skipif( From f6e1319c1dcaa5f2aa80e38fc88c8939dfd3b650 Mon Sep 17 00:00:00 2001 From: Lori Ren Date: Mon, 24 Aug 2026 03:31:55 +0000 Subject: [PATCH 2/2] [https://nvbugs/6550099][fix] Bound the no-top-k equivalence check by total mass, not atol Signed-off-by: Lori Ren --- .../hw_agnostic/test_advanced_sampling_mode.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py b/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py index beecf2bcc9c5..d88ba5777351 100644 --- a/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py +++ b/tests/unittest/_torch/speculative/hw_agnostic/test_advanced_sampling_mode.py @@ -84,11 +84,13 @@ def test_resolve_advanced_sampling_filters(mode, expect_top_k_none, expect_top_p @pytest.mark.parametrize("top_p_val", [1.0, 0.9]) def test_no_topk_matches_full(top_p_val): """With top_k disabled, NO_TOPK skips the top_k mask kernel (a no-op at k=vocab) - and yields the same sampling distribution as FULL. We compare the resulting - probability distributions rather than the sampled tokens, and with a tolerance above - the top-p renorm's fp32 precision: at the nucleus cutoff that precision decides - whether one token is kept, worth up to ~8e-5 of mass for this batch. - A real (non-no-op) filter would move mass by orders of magnitude, far above atol.""" + and yields the same sampling distribution as FULL. The k=vocab mask is a bit-exact + identity, so both modes feed the top-p renorm identical inputs; but that kernel is + not run-to-run reproducible right at the nucleus cutoff, so a boundary token can + flip between calls. We therefore bound the per-row total probability-mass + difference rather than compare pointwise: a boundary flip costs at most one + token's mass, while a real (non-no-op) filter moves O(0.1) of mass -- orders of + magnitude apart.""" dev = "cuda" torch.manual_seed(0) batch, vocab = 64, 32000 @@ -105,7 +107,8 @@ def test_no_topk_matches_full(top_p_val): ) probs_full = su.compute_probs_from_logits(logits.clone(), temperatures, ek_full, ep_full) probs_no_topk = su.compute_probs_from_logits(logits.clone(), temperatures, ek_nt, ep_nt) - assert torch.allclose(probs_full, probs_no_topk, atol=1e-4, rtol=0) + l1_diff = (probs_full - probs_no_topk).abs().sum(-1) + assert l1_diff.max().item() < 1e-3 @pytest.mark.skipif(