From 782c4458320558eac2ba61b29acbfe5bedd800db Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Sun, 19 Jul 2026 22:37:45 -0700 Subject: [PATCH] Fix OOB lanes in varlen indexer top-k Signed-off-by: Hongxiao Bai --- .../indexer_top_k_varlen_util.py | 94 ++++++++++--------- .../fe_api/dsa/test_DSA_indexer_top_k.py | 54 +++++++++++ 2 files changed, 106 insertions(+), 42 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 dbbe6cce2..b3e503ad0 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, ) + 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, - ) + relative_idx = cutlass.Int32(cur_tXcX[i // vec_size][1] + i % vec_size) + # The predicated copy fills lanes outside aligned_size with + # -inf, but those lanes are not part of the input row and + # must not contribute to the radix histogram. + if relative_idx < 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 + relative_idx = cutlass.Int32(cur_tXcX[i // vec_size][1] + i % vec_size) + if relative_idx < 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(relative_idx + vec_start) + s_indices[pos] = idx # for initial scalar load part. for j in range(tidx, prologue_elems, self.num_threads_per_cta): @@ -648,42 +656,44 @@ def indexer_topk_kernel_per_row( ) for i in cutlass.range(cute.size(tXrX), unroll_full=True): - 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 + relative_idx = cutlass.Int32(cur_tXcX[i // vec_size][1] + i % vec_size) + if relative_idx < aligned_size: + raw_input = tXrX[i] + bin_val = self.to_coarse_key(raw_input) + idx = self.index_type(relative_idx + 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..c167d49b2 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,57 @@ def test_DSA_indexer_top_k_wrapper( values, return_val, ) + + +@pytest.mark.L0 +@torch_fork_set_rng(seed=0) +def test_DSA_indexer_top_k_wrapper_ignores_vector_padding_with_negative_infinity(): + """OOB vector lanes must not join a real -inf threshold bin.""" + try: + from cudnn import DSA + except ImportError: + pytest.skip("Environment not supported: cudnn[cutedsl] not installed") + + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("Indexer top-k requires compute capability 9.0 or newer") + + num_rows = 633 + num_cols = 768 + seq_len = 633 + finite_values = 475 + top_k = 512 + + input_values = torch.full( + (num_rows, num_cols), + float("-inf"), + dtype=torch.float32, + device="cuda", + ) + input_values[:, :finite_values] = torch.randn( + num_rows, + finite_values, + dtype=torch.float32, + device="cuda", + ) + seq_lens = torch.full((num_rows,), seq_len, dtype=torch.int32, device="cuda") + + result = DSA.indexer_top_k_wrapper( + input_values, + seq_lens, + top_k=top_k, + next_n=1, + return_val=False, + ) + torch.cuda.synchronize() + + indices = result["indices"] + assert torch.all((indices >= 0) & (indices < seq_len)).item() + + selected_values = torch.gather(input_values, 1, indices.to(torch.int64)) + expected_values = torch.topk(input_values[:, :seq_len], top_k, dim=1).values + torch.testing.assert_close( + torch.sort(selected_values, dim=1).values, + torch.sort(expected_values, dim=1).values, + atol=0.0, + rtol=0.0, + )