From d8832757c723607c07075c7676d22481fb5e37cf Mon Sep 17 00:00:00 2001 From: Han-Yin Chang Date: Wed, 29 Jul 2026 01:32:53 +0800 Subject: [PATCH] [Bugfix] Harden top_k_per_row histogram path against NaN logits and under-filled output The histogram path of topKPerRowJob (top_k_per_row_prefill / top_k_per_row_decode) mishandles NaN logits (issue #49896): - extractBinIdx bins positive-NaN bit patterns above +inf, so NaN-dominated rows put the threshold bin on a NaN bin. - The insertion-sort final pass compares NaN logits (all comparisons false), so every threshold-bin element computes outIndex 0: one slot is written repeatedly and the remaining smemOutput slots keep uninitialized shared memory, which the unconditional final store loop then emits as indices (observed: 2143289344 = 0x7FC00000, a float32 NaN bit pattern reinterpreted as an int32 index). Downstream consumers gather block_table with these indices -> illegal memory access (DeepSeek-V4 long prefill on SM12x). Fix (data-driven, arch-independent): - Exclude NaN logits from selection in both histogram phases, via a bit-pattern test that is robust under --use_fast_math. - Pre-fill smemOutput with -1 (and -FLT_MAX for the logits half in the multiple-blocks-per-row variant) so under-filled rows produce the documented -1 padding of the short-row path instead of leaking uninitialized shared memory. - Default the per-step threshold state to "emit every matching element" so rows with fewer than topK non-NaN candidates terminate cleanly instead of reading an unset threshold bin. - Initialize the radix-sort padding indices to -1 so padding lanes can never write garbage into smemOutput. - Preserve -1 padding in the strided final store (do not rebase the sentinel by rowStart). Adds NaN-input regression tests covering the prefill insertion-sort and radix-sort variants and all three decode paths (insertion, radix, split+merge). Validated on NVIDIA H200 (SM90): the 25-line canary from #49896 fails before the fix (garbage indices -> device-side gather asserts) and passes after; the existing test_top_k_per_row.py suite passes unchanged. Per the issue, the defect is data-driven (NaN input) rather than SM120-specific. Signed-off-by: Han-Yin Chang --- csrc/libtorch_stable/sampler.cu | 42 +++++- tests/kernels/test_top_k_per_row.py | 209 ++++++++++++++++++++++++++++ 2 files changed, 250 insertions(+), 1 deletion(-) diff --git a/csrc/libtorch_stable/sampler.cu b/csrc/libtorch_stable/sampler.cu index 519e213281cd..326c62b52d85 100644 --- a/csrc/libtorch_stable/sampler.cu +++ b/csrc/libtorch_stable/sampler.cu @@ -43,6 +43,11 @@ __global__ void apply_repetition_penalties_kernel( } } +// Bit-pattern NaN test: robust regardless of --use_fast_math semantics. +__device__ __forceinline__ bool isNaNLogit(float x) { + return (__float_as_uint(x) & 0x7fffffffu) > 0x7f800000u; +} + __device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t { uint32_t bits = __float_as_uint(x); return (bits & 0x80000000) ? bits : ~bits & 0x7fffffff; @@ -165,6 +170,16 @@ __device__ bool processHistogramStep( smemFinal.histo.data[idx] = 0; } + // Reset the per-step threshold state. If fewer than topK candidates match + // the pattern (e.g. NaN-dominated rows, where NaN logits are excluded from + // selection below), no threshold bin exists: default to kNumBins so that + // every matching element is emitted directly and the remaining output slots + // keep the -1 padding pre-filled by topKPerRowJob. + if (threadIdx.x == 0) { + smemThresholdBinIdx[0] = kNumBins; + smemFinalBinSize[0] = 0; + } + // Make sure the histogram is ready. __syncthreads(); @@ -179,6 +194,11 @@ __device__ bool processHistogramStep( } auto distributeToBins = [&](float logit, int /* idx */ = 0) { + // NaN logits are never candidates: their bit patterns would otherwise be + // binned above +inf and poison the selection (see issue #49896). + if (isNaNLogit(logit)) { + return; + } if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); atomicAdd(&smemFinal.histo.data[binIdx], 1); @@ -254,6 +274,11 @@ __device__ bool processHistogramStep( thresholdBinIdx = smemThresholdBinIdx[0]; auto processBins = [&](float logit, int idx) { + // Keep in sync with the NaN exclusion in distributeToBins: the histogram + // counts and the prefix sums must describe the same candidate set. + if (isNaNLogit(logit)) { + return; + } if (isPartialMatch(logit, logitPattern)) { uint32_t binIdx = extractBinIdx(logit); // Only write elements with binIdx < thresholdBinIdx when: @@ -410,6 +435,16 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, smemFinalDstIdx[0] = 0; smemFoundTopKValues[0] = 0; } + // Pre-fill the output slots with -1 (and the logits, when present, with + // -FLT_MAX) so that rows with fewer than topK selectable candidates (e.g. + // NaN-dominated logits) produce -1 padding like the short-row path above + // instead of leaking uninitialized shared memory as indices. + for (int i = threadIdx.x; i < topK; i += kNumThreadsPerBlock) { + smemOutput[i] = -1; + if constexpr (multipleBlocksPerRow) { + reinterpret_cast(smemOutput + topK)[i] = -FLT_MAX; + } + } __syncthreads(); int thresholdBinIdx = -1; uint32_t logitPattern = 0; @@ -464,6 +499,9 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, #pragma unroll for (int ii = 0; ii < kNumFinalItemsPerThread; ++ii) { finalLogits[ii] = -FLT_MAX; + // Padding slots must carry the -1 sentinel: when fewer than topK + // candidates exist they may be copied to smemOutput below. + finalIndices[ii] = -1; } // Read the elements from SMEM. @@ -535,7 +573,9 @@ static __device__ void topKPerRowJob(const int* indices, const float* logits, // the rowStart. outIndices[i] = smemOutput[i]; } else { - outIndices[i] = smemOutput[i] - rowStart; + // Keep the -1 padding intact; only real indices are rebased. + int idx = smemOutput[i]; + outIndices[i] = idx < 0 ? idx : idx - rowStart; } } } diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index 3a1ad0f0d23c..bf3215cd7262 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -271,6 +271,215 @@ def test_top_k_per_row( ), "CUDA top_k_per_row_prefill results don't match torch.topk" +def _check_nan_hardened_topk( + logits: torch.Tensor, + indices: torch.Tensor, + row_starts: torch.Tensor, + row_ends: torch.Tensor, + top_k: int, + kernel_name: str, + check_rows: range | None = None, +) -> None: + """Validate the NaN-input contract of the top-k kernels (issue #49896). + + For rows longer than top_k (histogram path) the kernel must never emit + garbage: every output slot is either a valid in-range index of a non-NaN + logit or -1, the -1 padding forms a suffix, and the selected values match + a NaN-excluded torch.topk reference. Rows with row_len <= top_k keep the + documented short-row behavior (all candidate positions, then -1 padding). + """ + out = indices.cpu() + logits_cpu = logits.cpu() + row_starts_cpu = row_starts.cpu() + row_ends_cpu = row_ends.cpu() + rows = check_rows if check_rows is not None else range(indices.shape[0]) + for i in rows: + row_start = int(row_starts_cpu[i]) + row_end = int(row_ends_cpu[i]) + row_len = row_end - row_start + row = out[i] + # 1. No garbage: every slot is -1 padding or a valid in-range index. + assert ((row >= -1) & (row < row_len)).all(), ( + f"{kernel_name} row {i}: out-of-range index " + f"(min={int(row.min())}, max={int(row.max())}, row_len={row_len})" + ) + # 2. -1 padding is a suffix. + num_valid = int((row >= 0).sum()) + assert (row[num_valid:] == -1).all(), ( + f"{kernel_name} row {i}: -1 padding is not a suffix" + ) + valid = row[:num_valid].long() + # 3. No duplicate indices. + assert valid.unique().numel() == num_valid, ( + f"{kernel_name} row {i}: duplicate indices emitted" + ) + if row_len <= top_k: + # Short-row path: all candidate positions are selected. + assert num_valid == row_len, ( + f"{kernel_name} row {i}: short row expected {row_len} valid " + f"indices, got {num_valid}" + ) + continue + window = logits_cpu[i, row_start:row_end] + # 4. Histogram path: NaN logits are never selected. + selected_vals = window[valid] + assert not torch.isnan(selected_vals).any(), ( + f"{kernel_name} row {i}: NaN logit selected" + ) + # 5. Exactly min(top_k, #non-NaN) candidates are returned. + finite_count = int((~torch.isnan(window)).sum()) + expected_valid = min(top_k, finite_count) + assert num_valid == expected_valid, ( + f"{kernel_name} row {i}: expected {expected_valid} valid indices, " + f"got {num_valid}" + ) + # 6. The selected values match a NaN-excluded torch.topk reference. + if num_valid > 0: + ref_vals = window.nan_to_num(float("-inf")).topk(num_valid)[0] + got_vals = selected_vals.sort(descending=True)[0] + assert torch.equal(got_vals, ref_vals), ( + f"{kernel_name} row {i}: selected values do not match " + f"torch.topk reference" + ) + + +@pytest.mark.parametrize("top_k,vocab_size", [(512, 1243), (2048, 4507)]) +@pytest.mark.parametrize("nan_fraction", [0.5, 0.83, 1.0]) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@torch.inference_mode() +def test_top_k_per_row_prefill_nan_logits( + top_k: int, + vocab_size: int, + nan_fraction: float, +) -> None: + """Regression test for issue #49896: NaN-dominated logits must not make + top_k_per_row_prefill emit uninitialized shared memory or NaN bit + patterns as indices (histogram path).""" + set_random_seed(0) + torch.set_default_device("cuda:0") + + num_rows = 4869 + row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") + # Clamp to the number of columns (mirrors the issue #49896 repro). + row_ends = torch.clamp( + torch.arange(1, num_rows + 1, dtype=torch.int32, device="cuda"), + max=vocab_size, + ) + logits = torch.randn(num_rows, vocab_size, dtype=torch.float32, device="cuda") + if nan_fraction >= 1.0: + logits[:] = float("nan") + else: + nan_mask = torch.rand(num_rows, vocab_size, device="cuda") < nan_fraction + logits[nan_mask] = float("nan") + + # -7 canary: pre-fix, under-filled rows leaked uninitialized shared memory. + indices = torch.full((num_rows, top_k), -7, dtype=torch.int32, device="cuda") + torch.ops._C.top_k_per_row_prefill( + logits, + row_starts, + row_ends, + indices, + num_rows, + logits.stride(0), + logits.stride(1), + top_k, + ) + torch.accelerator.synchronize() + + _check_nan_hardened_topk( + logits, indices, row_starts, row_ends, top_k, "top_k_per_row_prefill" + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@torch.inference_mode() +def test_top_k_per_row_prefill_nan_logits_radix_rows() -> None: + """NaN hardening for the radix-sort variant of top_k_per_row_prefill + (rows beyond the first 12288 use the radix-sort final pass).""" + set_random_seed(0) + torch.set_default_device("cuda:0") + + num_rows, vocab_size, top_k = 12544, 1243, 512 + row_starts = torch.zeros(num_rows, dtype=torch.int32, device="cuda") + row_ends = torch.full((num_rows,), vocab_size, dtype=torch.int32, device="cuda") + logits = torch.randn(num_rows, vocab_size, dtype=torch.float32, device="cuda") + nan_mask = torch.rand(num_rows, vocab_size, device="cuda") < 0.83 + logits[nan_mask] = float("nan") + + indices = torch.full((num_rows, top_k), -7, dtype=torch.int32, device="cuda") + torch.ops._C.top_k_per_row_prefill( + logits, + row_starts, + row_ends, + indices, + num_rows, + logits.stride(0), + logits.stride(1), + top_k, + ) + torch.accelerator.synchronize() + + # Check a slice of insertion-sort rows and all radix-sort rows. + _check_nan_hardened_topk( + logits, + indices, + row_starts, + row_ends, + top_k, + "top_k_per_row_prefill (radix rows)", + check_rows=range(12160, num_rows), + ) + + +@pytest.mark.parametrize( + "vocab_size", [8192, 16384, 250000] +) # insertion sort / radix sort / split+merge +@pytest.mark.parametrize("nan_fraction", [0.5, 0.83, 1.0]) +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@torch.inference_mode() +def test_top_k_per_row_decode_nan_logits( + vocab_size: int, + nan_fraction: float, +) -> None: + """NaN hardening for all three top_k_per_row_decode code paths.""" + set_random_seed(0) + torch.set_default_device("cuda:0") + + top_k = 2048 + batch_size = 4 + seq_lens = torch.randint( + low=top_k + 1, + high=vocab_size, + size=(batch_size,), + dtype=torch.int32, + device="cuda", + ) + row_starts = torch.zeros(batch_size, dtype=torch.int32, device="cuda") + logits = torch.randn(batch_size, vocab_size, dtype=torch.float32, device="cuda") + if nan_fraction >= 1.0: + logits[:] = float("nan") + else: + nan_mask = torch.rand(batch_size, vocab_size, device="cuda") < nan_fraction + logits[nan_mask] = float("nan") + + indices = torch.full((batch_size, top_k), -7, dtype=torch.int32, device="cuda") + torch.ops._C.top_k_per_row_decode( + logits, + 1, + seq_lens, + indices, + batch_size, + logits.stride(0), + logits.stride(1), + top_k, + ) + torch.accelerator.synchronize() + + _check_nan_hardened_topk( + logits, indices, row_starts, seq_lens, top_k, "top_k_per_row_decode" + ) + + def _run_top_k_per_row_decode_test( top_k: int, batch_size: int,