[None][fix] Self-sampling top-k drops +inf from the top-k (register family)

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)
@@ -6128,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)
@@ -6193,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

tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py

@@ -292,6 +292,49 @@ 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 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
+    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."""