[Bug] Blackwell self-sampling GVR top-k: register family underfills the output when the bottom tie class is -inf

System Info

• CPU architecture: x86_64
• GPU: NVIDIA B200 (sm100a)
• Libraries
  • TensorRT-LLM: main @ 24ab1fab2 (the kernel file is unchanged since 07c5f2145)
  • nvidia-cutlass-dsl 4.7.0, torch 2.11.0+cu130, TensorRT-LLM 1.3.0rc24
• CUDA: driver 590.48.01 (CUDA 13.1)
• OS: Ubuntu 24.04

Who can help?

@longcheng-nv

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

The register family of the Blackwell CuTe DSL self-sampling top-k leaves part of the output unwritten when a row needs only a few members of a bottom tie class that is literally -inf. Through the production varlen entry (this shape routes to reg):

# one B200; python repro.py
import torch
from tensorrt_llm._torch.cute_dsl_kernels.blackwell.top_k import (
    gvr_topk_decode_self_sampling_host as ss_host,
)

torch.cuda.set_device(0)
b, n, k, nfin = 256, 2045, 1024, 600     # nfin finite entries per row, rest in-window -inf
npad = (n + 63) // 64 * 64
g = torch.Generator(device="cuda").manual_seed(20260901)

logits = torch.full((b, npad), float("-inf"), dtype=torch.float32, device="cuda")
pos = torch.arange(nfin, device="cuda")
pos[-1] = n - 1                          # keep one finite value in the tail column
logits[:, pos] = torch.randn((b, nfin), generator=g, dtype=torch.float32, device="cuda") * 2.0
logits[:, n:] = 3e38                     # pad poison, per the family's row-stride contract

pre_idx = torch.arange(k, dtype=torch.int32, device="cuda").repeat(b, 1)
kv_lens = torch.full((b,), n, dtype=torch.int32, device="cuda")
out = torch.full((b, k), -7, dtype=torch.int32, device="cuda")

ss_host.run_varlen(logits, pre_idx, kv_lens, out, next_n=1, compress_ratio=1, max_seq_len=n)
torch.cuda.synchronize()
print("unwritten:", int((out == -7).sum()), "rows:", int((out == -7).any(1).sum()), "/", b)
# unwritten: 108544 rows: 256 / 256      <- 424 = k - nfin slots per row, every row

Every row leaves exactly k - nfin output slots at their pre-call contents. Same shape at nfin=900 gives 124 per row, at nfin=1023 gives 1 per row.

Expected behavior

Every output slot written, with the row's exact tie-interchangeable top-k. The row has n=2045 valid entries and k=1024 requested, so the correct answer includes k - nfin of the -inf entries — -inf is an admissible input value per the family's own contract.

actual behavior

Deterministic silent underfill: out[nfin:k) is never written (stale contents), k - nfin slots per row, identical on every launch. No error, no fallback.

The band, at n=2045, k=1024, tpl (512, 1, 4, 1, True, True, False, 1024):

• clean for nfin <= 514, broken for every nfin in [515, 1023] (deficit 1..509), clean again at nfin >= k;
• k=512 compiles the identical tuple and is clean across nfin 100..512, and k=2048 (vpt=2) is clean across 600..2048. So the trigger is not the compile configuration — it is the k = 2 * BLK emit, where each thread contributes more than one output slot;
• it needs the bottom class to be literally -inf. Ordinary exact ties at the k-th value do not reproduce it: a finite constant (-5.0), exactly 0.0, and a negative denormal (-1e-44) are all clean at the same deficits.

The defect is hint-independent — this cell compiles the DEG bracket arm, which ignores pre_idx entirely. arange(k) (an ordinary duplicate-free hint, above) and a degenerate argmax-anchor hint produce byte-identical underfill counts. It is therefore not avoided by the hint-free mode of #18410: for DEG compiles that PR leaves the generated code unchanged.

additional notes

Root cause, as far as I traced it: the order-preserving key transform (fkey_bits) maps -inf to 0x007fffff, the smallest key it produces over all non-NaN values (key 0 is a negative NaN), so the whole -inf tie class lands in the lowest histogram bin — and the BRL classify arm also clamps out-of-bracket values into that same bin. The histogram total therefore does reach k, which is why the kernel's underfill trigger cannot fire, and the lowest-bin whole-bin emit stops short of k while rejecting the very values it needs.

This is separate from the hint-driven under-bracketing fixed in #18501: I verified the same repro on that branch and the numbers are byte-identical (108544 / 31744 / 256 unwritten), so this one is untouched by it and needs its own fix. Its known-limitations section names this class; this issue is the standalone repro for it.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.