From f21f88a2010e1dc735c23109fb1be23e33c7e840 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 02:06:54 +0000 Subject: [PATCH 1/3] [None][fix] Self-sampling top-k: keep +inf in the register-family top-k An in-window +inf drives the register-family hint-free bracket max to +inf, so the bracket width GMAX-Tv is +inf and SC=rcp(+inf)=0 folds every value into bin 0. The whole-bin emit then ranks by position and drops the +inf from the top-k (DKG issue #58; oracle and degenerate hints fail identically, so this is in the selection, not the hint). The sentinel fallback did not help because its own width (SENT_HI-SENT_LO) also overflows to +inf. The collapse guard now rejects an infinite bracket width, and the register kernel forces the count-crossing escape whenever the bracket is degenerate (okc=0). The escape ranks the full row in key space independent of the bracket, where fkey(+inf) is the maximum key, so +inf is selected. The finite +/-3e38 and -inf paths are unchanged; NaN ordering stays implementation-specific per the documented limitation. Verified on B200 sm100: the issue repro (in-/out-of-window +inf across N in {4K,16K,64K}, cr in {1,4}) passes, the -inf duplicate-index sweep (3 seed-dependent configs) stays fixed, and the full test_gvr_selfsampling_topk suite is 82 passed. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 12 +++++++-- .../parallel/test_gvr_selfsampling_topk.py | 25 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index a3d0d2f9d8ac..fa3dd9c36119 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -3713,10 +3713,13 @@ def kern( Tv = invkey(lmin) GMAX = invkey(lmax) - # ---- collapse guard, NaN-safe + # ---- collapse guard, NaN-safe. w_ < 3.5e38 (> FLT_MAX) also + # rejects an infinite bracket width: an in-window +inf (GMAX=+inf) + # or -inf (Tv=-inf) makes SC=0 and folds every value into bin 0. okc = cutlass.Int32(0) if Tv < GMAX: - if (GMAX - Tv) > cutlass.Float32(1e-30): + w_ = GMAX - Tv + if w_ > cutlass.Float32(1e-30) and w_ < cutlass.Float32(3.5e38): okc = cutlass.Int32(1) if okc == cutlass.Int32(0): Tv = cutlass.Float32(SENT_LO) @@ -3805,6 +3808,11 @@ def kern( esc = cutlass.Int32(1) if tot < k: esc = cutlass.Int32(1) + # A degenerate bracket (okc==0) collapses the histogram into + # bin 0; escape to the bracket-independent key-space rank, where + # fkey(+inf) is the maximum key. + if okc == cutlass.Int32(0): + esc = cutlass.Int32(1) if esc == cutlass.Int32(1): if tid == cutlass.Int32(0): s_cnt[0] = cutlass.Int32(0) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index ea19bcf8b6f4..f19ccb42ad7a 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -292,6 +292,31 @@ def test_selfsampling_topk_neginf_tail_completeness(): _check_exact(logits, indices, n_valid, ref_vals) +@pytest.mark.parametrize("pos", [1000, 3000], ids=["in_window", "out_of_window"]) +def test_selfsampling_topk_posinf_completeness(pos): + """A +inf in the register-family fold window drives the bracket max to + +inf, so the bracket width GMAX-Tv=+inf and SC=rcp(+inf)=0 fold every + value into bin 0; the whole-bin emit then drops the +inf from the top-k + (regression, DKG issue #58). The infinite-width bracket must be rejected + by the collapse guard and take the key-space escape, where fkey(+inf) is + the maximum key. Both an in-window and an out-of-window +inf are pinned; + N=4096 k=1024 keeps the register 'reg' family (not the streaming tiers, + which never collapse this way).""" + top_k = 1024 + n_valid = 4096 + gen = torch.Generator(device=_DEV).manual_seed(top_k + n_valid) + logits = torch.randn((1, n_valid), generator=gen, dtype=torch.float32, device=_DEV) * 2.0 + logits[0, pos] = float("inf") + ref_vals, _ = torch.topk(logits, top_k, dim=1) + indices = torch.full((1, top_k), -7, dtype=torch.int32, device=_DEV) + kv = torch.full((1,), n_valid, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(logits, kv, indices, max_seq_len=n_valid) + torch.cuda.synchronize() + assert int((indices == -7).sum()) == 0, "unwritten output slots" + assert torch.isinf(logits[0][indices[0].long()]).any(), "+inf dropped from the top-k" + _check_exact(logits, indices, n_valid, ref_vals) + + def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): """Build a per-row-poisoned varlen batch, run run_varlen, verify every row against its own n_r (production formula) — short rows included.""" From 9f9203312f0a3b52e953a27f9f30a7542462ed9f Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:06:23 +0000 Subject: [PATCH 2/3] [None][fix] Handle +inf in clustered register GVR Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../top_k/gvr_topk_decode_self_sampling.py | 13 +++++++++++-- .../parallel/test_gvr_selfsampling_topk.py | 18 ++++++++++++++++++ 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py index fa3dd9c36119..a07ba529eb6d 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py @@ -6136,10 +6136,13 @@ def kern( Tv = invkey(lmin) GMAX = invkey(lmax) - # ---- collapse guard, NaN-safe + # ---- collapse guard, NaN-safe. Reject infinite width as well: + # otherwise SC becomes zero and can collapse +inf into a finite + # histogram bin. okc = cutlass.Int32(0) if Tv < GMAX: - if (GMAX - Tv) > cutlass.Float32(1e-30): + w_ = GMAX - Tv + if w_ > cutlass.Float32(1e-30) and w_ < cutlass.Float32(3.5e38): okc = cutlass.Int32(1) if okc == cutlass.Int32(0): Tv = cutlass.Float32(SENT_LO) @@ -6201,6 +6204,12 @@ def kern( degen = cutlass.Int32(0) if m > cutlass.Int32(CS * CMPC): degen = cutlass.Int32(1) + # The sentinel bracket also has infinite width. Bypass its + # collapsed histogram and enter the exact whole-row key-space + # fallback on rank 0, where fkey(+inf) is the maximum key. + if okc == cutlass.Int32(0): + whole = cutlass.Int32(0) + degen = cutlass.Int32(1) for z in cutlass.range_constexpr(NB__regclus // self.blk): i = tid + cutlass.Int32(z * self.blk) s_mrg[i] = s_mrg[i] + s_hoff[i] # global cursor diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index f19ccb42ad7a..06d3a22fa29e 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -317,6 +317,24 @@ def test_selfsampling_topk_posinf_completeness(pos): _check_exact(logits, indices, n_valid, ref_vals) +def test_selfsampling_topk_posinf_regclus_completeness(): + """A collapsed reg_clus bracket must use its whole-row key-space + fallback instead of emitting from the lossy float-space histogram.""" + batch_size, n_valid, top_k = 4, 32768, 1024 + assert ss_host.route(batch_size, n_valid, n_valid, top_k)["kernel"] == "reg_clus" + gen = torch.Generator(device=_DEV).manual_seed(top_k + n_valid) + logits = torch.randn((batch_size, n_valid), generator=gen, dtype=torch.float32, device=_DEV) + logits[:, 1000] = float("inf") + ref_vals, _ = torch.topk(logits, top_k, dim=1) + indices = torch.full((batch_size, top_k), -7, dtype=torch.int32, device=_DEV) + kv = torch.full((batch_size,), n_valid, dtype=torch.int32, device=_DEV) + ss_host.run_varlen(logits, kv, indices, max_seq_len=n_valid) + torch.cuda.synchronize() + assert int((indices == -7).sum()) == 0, "unwritten output slots" + assert torch.isposinf(logits.gather(1, indices.long())).any(dim=1).all() + _check_exact(logits, indices, n_valid, ref_vals) + + def _run_varlen_case(kv, next_n, cr, top_k, seed, with_values=False): """Build a per-row-poisoned varlen batch, run run_varlen, verify every row against its own n_r (production formula) — short rows included.""" From 73af66531c6dbd6cba2027e35cdd9bd4875acaa0 Mon Sep 17 00:00:00 2001 From: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> Date: Thu, 3 Sep 2026 04:35:30 +0000 Subject: [PATCH 3/3] [None][test] Annotate GVR +inf regression test Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com> --- .../unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py index 06d3a22fa29e..3408dd325fd6 100644 --- a/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py +++ b/tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py @@ -317,7 +317,7 @@ def test_selfsampling_topk_posinf_completeness(pos): _check_exact(logits, indices, n_valid, ref_vals) -def test_selfsampling_topk_posinf_regclus_completeness(): +def test_selfsampling_topk_posinf_regclus_completeness() -> None: """A collapsed reg_clus bracket must use its whole-row key-space fallback instead of emitting from the lossy float-space histogram.""" batch_size, n_valid, top_k = 4, 32768, 1024