obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-22 21:50:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Bugfix][Kernel] Fix persistent top-k histogram reuse after short rows (#49139)

Browse Source 
Signed-off-by: fxfxfxfxfxfxfxfx <[email protected]>
Co-authored-by: Michael Goin <[email protected]>
This commit is contained in:
xiao feng
2026-08-12 12:54:45 -07:00
committed by GitHub
co-authored by Michael Goin
parent 7f7a32cfec
commit e62abc37d4
2 changed files with 47 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/persistent_topk.cuh
+5 -3
View File
@@ -661,7 +661,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
uint32_t* shared_scalars, uint32_t* shared_ordered,
RadixRowState* state, uint32_t cta_in_group,
uint32_t ctas_per_group, int& barrier_phase,
uint32_t iter, uint32_t tx) {
uint32_t radix_iter, uint32_t tx) {
const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len)
? my_chunk_start + chunk_size
: seq_len;
@@ -718,7 +718,7 @@ __device__ void radix_topk(const float* __restrict__ row_input,
// -- Stage 2: 4 rounds of radix select --
for (uint32_t round = 0; round < 4; round++) {
const uint32_t global_round = iter * 4 + round;
const uint32_t global_round = radix_iter * 4 + round;
const uint32_t shift = 24 - round * 8;
const uint32_t prefix = shared_scalars[0];
const uint32_t remaining_k = shared_scalars[1];
@@ -898,6 +898,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
RadixRowState* state = &params.row_states[group_id];
int barrier_phase = 0;
uint32_t radix_iter = 0;
const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups;
for (uint32_t iter = 0; iter < total_iters; iter++) {
@@ -930,7 +931,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 2)
radix_topk<TopK, VEC_SIZE>(
row_input, row_output, seq_len, my_chunk_start, chunk_size,
local_histogram, suffix_sum, shared_scalars, shared_ordered, state,
cta_in_group, ctas_per_group, barrier_phase, iter, tx);
cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx);
radix_iter++;
}
}
tests/kernels/test_top_k_per_row.py
+42
View File
@@ -918,6 +918,48 @@ def test_workspace_topk(test_config: dict, top_k: int, backend: str) -> None:
)
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
@torch.inference_mode()
def test_persistent_topk_reused_group_after_short_row() -> None:
"""A short row must not advance a group's radix histogram ring."""
torch.set_default_device("cuda:0")
set_random_seed(0)
top_k = 2048
long_seq_len = 32769
radix = 256
fixed_smem = ((radix + radix + 5) * 4 + 15) & ~15
props = torch.cuda.get_device_properties(0)
max_smem = props.shared_memory_per_block_optin
if max_smem <= props.shared_memory_per_multiprocessor // 2:
pytest.skip("Cannot force one persistent_topk CTA per SM")
max_chunk = ((max_smem - fixed_smem) // 4 // 4) * 4
ctas_per_group = max(
(props.multi_processor_count - 1 + 9) // 10,
(long_seq_len + max_chunk - 1) // max_chunk,
)
if ctas_per_group >= props.multi_processor_count:
pytest.skip("Not enough SMs to construct a reused CTA group")
stride = ctas_per_group * max_chunk
num_groups = max(1, (props.multi_processor_count - 1) // ctas_per_group)
num_rows = 3 * num_groups
lengths = torch.full((num_rows,), top_k, dtype=torch.int32, device="cuda")
target_row = 2 * num_groups
lengths[0] = long_seq_len
lengths[num_groups] = long_seq_len - 1
lengths[target_row] = long_seq_len
logits = torch.randn(num_rows, stride, dtype=torch.float32, device="cuda")
indices = torch.empty((num_rows, top_k), dtype=torch.int32, device="cuda")
_run_topk_backend("persistent_topk", logits, lengths, indices, top_k, stride)
torch.accelerator.synchronize()
expected = logits[target_row, :long_seq_len].topk(top_k).indices
assert set(indices[target_row].cpu().tolist()) == set(expected.cpu().tolist())
@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
@pytest.mark.parametrize("top_k", [512, 2048])
@pytest.mark.parametrize("backend", WORKSPACE_TOPK_BACKENDS)