[Bugfix][Kernel] Fix persistent top-k histogram reuse after short rows (#49139)
Signed-off-by: fxfxfxfxfxfxfxfx <227935476@qq.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
xiao feng committed
Aug 13, 2026 at 03:54 UTC
e62abc37d4c341b6026b945d7aff7e7537cce75d
2 files changed
+47
-3
csrc/libtorch_stable/persistent_topk.cuh
+5
-3
@@ -661,7 +661,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
661
uint32_t* shared_scalars, uint32_t* shared_ordered,
662
RadixRowState* state, uint32_t cta_in_group,
663
uint32_t ctas_per_group, int& barrier_phase,
664
- uint32_t iter, uint32_t tx) {
664
+ uint32_t radix_iter, uint32_t tx) {
665
const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len)
666
? my_chunk_start + chunk_size
667
: seq_len;
@@ -718,7 +718,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
718
719
// -- Stage 2: 4 rounds of radix select --
720
for (uint32_t round = 0; round < 4; round++) {
721
- const uint32_t global_round = iter * 4 + round;
721
+ const uint32_t global_round = radix_iter * 4 + round;
722
const uint32_t shift = 24 - round * 8;
723
const uint32_t prefix = shared_scalars[0];
724
const uint32_t remaining_k = shared_scalars[1];
@@ -898,6 +898,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
898
RadixRowState* state = ¶ms.row_states[group_id];
899
900
int barrier_phase = 0;
901
+ uint32_t radix_iter = 0;
902
const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups;
903
904
for (uint32_t iter = 0; iter < total_iters; iter++) {
@@ -930,7 +931,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
931
radix_topk<TopK, VEC_SIZE>(
932
row_input, row_output, seq_len, my_chunk_start, chunk_size,
933
local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
933
- cta_in_group, ctas_per_group, barrier_phase, iter, tx);
934
+ cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx);
935
+ radix_iter++;
936
}
937
}
938
tests/kernels/test_top_k_per_row.py
+42
@@ -918,6 +918,48 @@ def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None:
918
)
919
920
921
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
922
+@torch.inference_mode()
923
+def test_persistent_topk_reused_group_after_short_row() -> None:
924
+ """A short row must not advance a group's radix histogram ring."""
925
+ torch.set_default_device("cuda:0")
926
+ set_random_seed(0)
927
+
928
+ top_k = 2048
929
+ long_seq_len = 32769
930
+ radix = 256
931
+ fixed_smem = ((radix + radix + 5) * 4 + 15) & ~15
932
+ props = torch.cuda.get_device_properties(0)
933
+ max_smem = props.shared_memory_per_block_optin
934
+ if max_smem <= props.shared_memory_per_multiprocessor // 2:
935
+ pytest.skip("Cannot force one persistent_topk CTA per SM")
936
+
937
+ max_chunk = ((max_smem - fixed_smem) // 4 // 4) * 4
938
+ ctas_per_group = max(
939
+ (props.multi_processor_count - 1 + 9) // 10,
940
+ (long_seq_len + max_chunk - 1) // max_chunk,
941
+ )
942
+ if ctas_per_group >= props.multi_processor_count:
943
+ pytest.skip("Not enough SMs to construct a reused CTA group")
944
+
945
+ stride = ctas_per_group * max_chunk
946
+ num_groups = max(1, (props.multi_processor_count - 1) // ctas_per_group)
947
+ num_rows = 3 * num_groups
948
+ lengths = torch.full((num_rows,), top_k, dtype=torch.int32, device="cuda")
949
+ target_row = 2 * num_groups
950
+ lengths[0] = long_seq_len
951
+ lengths[num_groups] = long_seq_len - 1
952
+ lengths[target_row] = long_seq_len
953
+
954
+ logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda")
955
+ indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
956
+ _run_topk_backend("persistent_topk", logits, lengths, indices, top_k, stride)
957
+ torch.accelerator.synchronize()
958
+
959
+ expected = logits[target_row, :long_seq_len].topk(top_k).indices
960
+ assert set(indices[target_row].cpu().tolist()) == set(expected.cpu().tolist())
961
+
962
+
963
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
964
@pytest.mark.parametrize("top_k", [512, 2048])
965
@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)