Skip to content
Open
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -377,11 +377,17 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP")
"csrc/libtorch_stable/attention/paged_attention_v1.cu"
"csrc/libtorch_stable/attention/paged_attention_v2.cu"
"csrc/libtorch_stable/cache_kernels.cu"
"csrc/libtorch_stable/cache_kernels.cu"
"csrc/libtorch_stable/cache_kernels_fused.cu"
"csrc/libtorch_stable/custom_all_reduce.cu"
"csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu")

if(VLLM_GPU_LANG STREQUAL "CUDA")
# Raw PTX (ld.global.nc / st.global.cg) in the row copy: CUDA-only.
# The op declarations and registrations are guarded by USE_ROCM.
list(APPEND VLLM_STABLE_EXT_SRC
"csrc/libtorch_stable/hisparse_kernels.cu")
endif()

if(VLLM_GPU_LANG STREQUAL "CUDA" AND
DEFINED CMAKE_CUDA_COMPILER_VERSION AND
CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.0)
Expand Down
10 changes: 10 additions & 0 deletions csrc/libtorch_stable/cooperative_topk.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -452,6 +452,16 @@ __device__ void large_topk(const float* __restrict__ row_input,
if (rank != 0) { // only rank 0 does tie refinement
return;
}
// Ties beyond the per-block collect cap (kMaxTies) are dropped at scatter,
// so the written candidate total can fall short of TopK even though the
// threshold bin guarantees enough true candidates. Pad the unwritten tail
// with -1 ("no token"): consumers read all TopK slots, and a slot left
// holding the previous step's index page-translates to a garbage KV row
// (wrong-row gather or illegal memory access downstream).
for (uint32_t i = s_total_above + s_total_equal + tx; i < TopK;
i += hist4096::kBlockSize) {
row_output[i] = -1u;
}
if (s_total_above + s_total_equal <= TopK) { // no ties to refine
return;
}
Expand Down
Loading
Loading