From 5cb18d321e2577012e9cb0c7ee09b91e4611f1e6 Mon Sep 17 00:00:00 2001 From: Jack Rao Date: Mon, 27 Jul 2026 22:28:39 -0700 Subject: [PATCH] DSA indexer_top_k: exclude OOB tile-fill lanes from radix selection The CuTe-DSL radix top-k reads each row in fixed-width vector tiles and fills the final partial tile's out-of-bounds lanes with -inf so the predicated copy is safe, but the histogram and candidate-collection loops then iterate every fragment lane without a bounds check, counting those phantom -inf lanes as real elements. This is harmless while the row's top-k threshold sits above the phantom bin: phantom -inf lanes sort last and are dropped. But the coarse radix bin is derived from the fp16 conversion of the fp32 score (to_coarse_key), so every score below fp16's -65504 minimum collapses into the fp16 -inf bin -- the same bin the phantom lanes occupy. When a row's top-k threshold lands inside that bin (fewer than top_k values above ~-65504), the threshold bin's candidate list becomes real_count + up-to-tile-width phantom lanes. In the large_occupancy compile (>148 rows) the per-row candidate buffers (512-entry smem + num_cols gmem) overflow, producing out-of-bounds shared/global writes and cudaErrorIllegalAddress; in any configuration phantom lanes can be selected as winners, yielding silently out-of-range indices. Observed in production on a GLM-5.2 context-parallel (CP32) training run: deterministic, data-dependent illegal memory access (Xid 43) on B200 and B300, reproduced from the captured score tensor; isolated via CUDA_LAUNCH_BLOCKING and compute-sanitizer (invalid 2-byte shared writes); trigger confirmed as the value distribution, not the shape. Fix: skip lanes whose tile column coordinate is out of bounds in the three vectorized histogram/collection loops, reusing the same aligned_size bound the predicated copy already computes. The scalar prologue/leftover loops are exact-bounded and unchanged. Testing: new L0 regression test test_DSA_indexer_top_k_oob_tile_lanes deterministically drives a row into the flood regime in both compile-time variants (phantom-dominated candidate lists with identical and with distinct collapsed values). On v1.26.0 with this patch applied (B200/B300, SM90+/SM100): the previously-crashing inputs complete with in-range indices, compute-sanitizer reports no invalid accesses, and 12/12 randomized parity trials against torch.topk pass across row counts 37-862, widths 4097-51720, and fp16-overflow/exact-tie distributions. Kernel microbench on (862, 51720) k=2048: 0.115 -> 0.129 ms/call. Co-Authored-By: Claude Fable 5 --- .../indexer_top_k_varlen_util.py | 91 +++++++++++-------- .../fe_api/dsa/test_DSA_indexer_top_k.py | 69 ++++++++++++++ 2 files changed, 120 insertions(+), 40 deletions(-) diff --git a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py index af9d49b3f..def1820fe 100644 --- a/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py +++ b/python/cudnn/deepseek_sparse_attention/indexer_top_k/indexer_top_k_varlen_util.py @@ -525,12 +525,18 @@ def indexer_topk_kernel_per_row( -tXrX.element_type.inf, ) + # Skip out-of-bounds lanes: the last tile can extend past + # aligned_size, and its -inf-filled lanes must not be counted + # as real elements. + cur_tXcX = tXcX[None, None, None, tile_idx] for i in cutlass.range(cute.size(tXrX), unroll_full=True): - bin_val = self.to_coarse_key(tXrX[i]) - atomicAdd( - s_histogram.iterator + cutlass.Int32(bin_val), - val_one, - ) + col = cur_tXcX[i // vec_size][1] + i % vec_size + if col < aligned_size: + bin_val = self.to_coarse_key(tXrX[i]) + atomicAdd( + s_histogram.iterator + cutlass.Int32(bin_val), + val_one, + ) # for initial scalar load part. for j in range(tidx, prologue_elems, self.num_threads_per_cta): @@ -594,11 +600,13 @@ def indexer_topk_kernel_per_row( ) for i in cutlass.range(cute.size(tXrX), unroll_full=True): cur_tXcX = tXcX[None, None, None, tile_idx] - bin_val = self.to_coarse_key(tXrX[i]) - if bin_val < threshold_bin: - pos = atomicAdd(s_counter.iterator, val_one) - idx = self.index_type(cur_tXcX[i // vec_size][1] + i % vec_size + vec_start) - s_indices[pos] = idx + col = cur_tXcX[i // vec_size][1] + i % vec_size + if col < aligned_size: + bin_val = self.to_coarse_key(tXrX[i]) + if bin_val < threshold_bin: + pos = atomicAdd(s_counter.iterator, val_one) + idx = self.index_type(col + vec_start) + s_indices[pos] = idx # for initial scalar load part. for j in range(tidx, prologue_elems, self.num_threads_per_cta): @@ -651,39 +659,42 @@ def indexer_topk_kernel_per_row( raw_input = tXrX[i] bin_val = self.to_coarse_key(raw_input) cur_tXcX = tXcX[None, None, None, tile_idx] - idx = self.index_type(cur_tXcX[i // vec_size][1] + i % vec_size + vec_start) - if bin_val < threshold_bin: - pos = atomicAdd(s_counter.iterator, val_one) - s_indices[pos] = idx - elif bin_val == threshold_bin: - # pos = atomicAdd(s_num_input[0], 1) - pos = atomicAdd(s_num_input.iterator, val_one) - if cutlass.const_expr(self.enable_gmem_store): - if pos < self.indexer_topk_smem_input_size: - s_input_idx[0, pos] = idx + col = cur_tXcX[i // vec_size][1] + i % vec_size + # Skip out-of-bounds lanes filled with -inf by _fill_oob. + if col < aligned_size: + idx = self.index_type(col + vec_start) + if bin_val < threshold_bin: + pos = atomicAdd(s_counter.iterator, val_one) + s_indices[pos] = idx + elif bin_val == threshold_bin: + # pos = atomicAdd(s_num_input[0], 1) + pos = atomicAdd(s_num_input.iterator, val_one) + if cutlass.const_expr(self.enable_gmem_store): + if pos < self.indexer_topk_smem_input_size: + s_input_idx[0, pos] = idx + else: + buffer_pos = atomicAdd( + g_num_input.iterator, + val_one, + ) + buffer[0, buffer_pos] = cutlass.Int32(cutlass.Uint32(idx)) + ordered = self.to_ordered(raw_input) + sub_bin = (ordered >> self.first_refine_shift) & 0xFF + # atomicAdd(s_histogram[sub_bin], 1) + atomicAdd( + s_histogram.iterator + cutlass.Int32(sub_bin), + val_one, + ) else: - buffer_pos = atomicAdd( - g_num_input.iterator, + if pos < self.indexer_topk_smem_input_size: + s_input_idx[0, pos] = idx + ordered = self.to_ordered(raw_input) + sub_bin = (ordered >> self.first_refine_shift) & 0xFF + # atomicAdd(s_histogram[sub_bin], 1) + atomicAdd( + s_histogram.iterator + cutlass.Int32(sub_bin), val_one, ) - buffer[0, buffer_pos] = cutlass.Int32(cutlass.Uint32(idx)) - ordered = self.to_ordered(raw_input) - sub_bin = (ordered >> self.first_refine_shift) & 0xFF - # atomicAdd(s_histogram[sub_bin], 1) - atomicAdd( - s_histogram.iterator + cutlass.Int32(sub_bin), - val_one, - ) - else: - if pos < self.indexer_topk_smem_input_size: - s_input_idx[0, pos] = idx - ordered = self.to_ordered(raw_input) - sub_bin = (ordered >> self.first_refine_shift) & 0xFF - # atomicAdd(s_histogram[sub_bin], 1) - atomicAdd( - s_histogram.iterator + cutlass.Int32(sub_bin), - val_one, - ) # for initial scalar load part. for j in range(tidx, prologue_elems, self.num_threads_per_cta): diff --git a/test/python/fe_api/dsa/test_DSA_indexer_top_k.py b/test/python/fe_api/dsa/test_DSA_indexer_top_k.py index 0f427bfe0..29e160d62 100644 --- a/test/python/fe_api/dsa/test_DSA_indexer_top_k.py +++ b/test/python/fe_api/dsa/test_DSA_indexer_top_k.py @@ -143,3 +143,72 @@ def test_DSA_indexer_top_k_wrapper( values, return_val, ) + + +@pytest.mark.L0 +@pytest.mark.parametrize("trigger_dist", ["identical", "distinct"]) +def test_DSA_indexer_top_k_oob_tile_lanes(trigger_dist): + """Regression test: -inf-filled OOB lanes of a row's final vector tile must + not participate in the radix top-k. + + The kernel reads each row in fixed-width vector tiles and fills the final + partial tile's out-of-bounds lanes with -inf. Those phantom lanes were + counted as real elements by the histogram and candidate-collection passes. + When a row's top-k threshold lands in the fp16 -inf coarse bin (fewer than + top_k values above ~-65504, since to_coarse_key converts fp32 scores to + fp16 first), the phantom lanes flood the threshold bin's candidate list and + overflow the per-row candidate buffers in the large-occupancy compile + (>148 rows), producing out-of-bounds shared/global writes + (cudaErrorIllegalAddress) or silently out-of-range selected indices. + + This test's batch deterministically places row 0 in that regime: 149 rows + (>148 -> large_occupancy), a 4310-wide matrix, and a 4122-long row whose + values below the fp16 negative-overflow point plus the phantom lanes + outnumber the candidate-buffer capacity. + """ + try: + from cudnn import DSA + from cuda.bindings import driver as cuda + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("indexer top-k kernel requires SM90+") + + num_rows, num_cols, top_k, trigger_len, num_above = 149, 4310, 2048, 4122, 1122 + scores = torch.empty(num_rows, num_cols, dtype=torch.float32) + scores[1:] = torch.arange(num_cols, dtype=torch.float32).expand(num_rows - 1, -1) + row0 = torch.empty(num_cols, dtype=torch.float32) + row0[:num_above] = torch.arange(num_above, dtype=torch.float32) + if trigger_dist == "identical": + row0[num_above:trigger_len] = -100000.0 + else: + row0[num_above:trigger_len] = torch.linspace(-66000.0, -200000.0, trigger_len - num_above) + row0[trigger_len:] = float("-inf") + scores[0] = row0 + seq_lens = torch.full((num_rows,), num_cols, dtype=torch.int32) + seq_lens[0] = trigger_len + + input_values = scores.cuda() + seq_lens = seq_lens.cuda() + stream = cuda.CUstream(torch.cuda.current_stream().cuda_stream) + result = DSA.indexer_top_k_wrapper( + input_values, + seq_lens, + top_k, + next_n=1, + return_val=True, + stream=stream, + ) + torch.cuda.synchronize() + + indices = result["indices"].cpu() + values = result["values"].cpu() + for r in range(num_rows): + L = int(seq_lens[r].cpu()) + k = min(top_k, L) + valid = indices[r] >= 0 + assert int(valid.sum()) == k, f"row {r}: expected {k} valid indices, got {int(valid.sum())}" + assert int(indices[r].max()) < L, f"row {r}: out-of-range index {int(indices[r].max())} >= seq_len {L}" + got = values[r][valid].sort(descending=True).values + ref = torch.topk(scores[r, :L], k).values.sort(descending=True).values + torch.testing.assert_close(got, ref, atol=1e-6, rtol=1e-5)