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..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 @@ -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 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..3408dd325fd6 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,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."""