| 1 | // Persistent TopK kernel for DeepSeek V3 sparse attention indexer. |
| 2 | // See persistent_topk.cuh for kernel implementation. |
| 3 | |
| 4 | #include <cuda_runtime.h> |
| 5 | #include <algorithm> |
| 6 | |
| 7 | #include "ops.h" |
| 8 | #include "torch_utils.h" |
| 9 | |
| 10 | #ifndef USE_ROCM |
| 11 | #include "persistent_topk.cuh" |
| 12 | #include "sampled_topk.cuh" |
| 13 | #endif |
| 14 | |
| 15 | namespace { |
| 16 | |
| 17 | #ifndef USE_ROCM |
| 18 | template <int TopK> |
| 19 | void launch_persistent_topk(const torch::stable::Tensor& logits, |
| 20 | const torch::stable::Tensor& lengths, |
| 21 | torch::stable::Tensor& output, |
| 22 | torch::stable::Tensor& workspace, |
| 23 | int64_t max_seq_len) { |
| 24 | namespace P = vllm::persistent; |
| 25 | |
| 26 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 27 | logits.get_device_index()); |
| 28 | const int64_t num_rows = logits.size(0); |
| 29 | const int64_t stride = logits.stride(0); |
| 30 | const cudaStream_t stream = get_current_cuda_stream(); |
| 31 | |
| 32 | static int num_sms = 0; |
| 33 | static int max_smem_per_block = 0; |
| 34 | if (num_sms == 0) { |
| 35 | const cudaDeviceProp* device_prop = get_device_prop(); |
| 36 | num_sms = device_prop->multiProcessorCount; |
| 37 | max_smem_per_block = device_prop->sharedMemPerBlockOptin; |
| 38 | } |
| 39 | |
| 40 | // Allow static fallback storage in addition to the 128 KiB dynamic buffer. |
| 41 | if (num_rows > 64 && |
| 42 | max_seq_len >= vllm::sampled_topk::kMinSampledLength<TopK> && |
| 43 | max_smem_per_block >= 144 * 1024) { |
| 44 | auto kernel = vllm::sampled_topk::sampled_topk_kernel<TopK>; |
| 45 | constexpr size_t smem_size = |
| 46 | vllm::filtered_topk::FILTERED_TOPK_SMEM_DYNAMIC; |
| 47 | cudaError_t status = cudaFuncSetAttribute( |
| 48 | kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); |
| 49 | STD_TORCH_CHECK(status == cudaSuccess, |
| 50 | "sampled_topk smem failed: ", cudaGetErrorString(status)); |
| 51 | kernel<<<num_rows, vllm::sampled_topk::kThreads, smem_size, stream>>>( |
| 52 | logits.const_data_ptr<float>(), lengths.const_data_ptr<int32_t>(), |
| 53 | output.mutable_data_ptr<int32_t>(), stride, |
| 54 | static_cast<int>(std::min(max_seq_len, logits.size(1)))); |
| 55 | } else if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { |
| 56 | cudaError_t status = |
| 57 | vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>( |
| 58 | logits.const_data_ptr<float>(), output.mutable_data_ptr<int32_t>(), |
| 59 | lengths.const_data_ptr<int32_t>(), static_cast<uint32_t>(num_rows), |
| 60 | static_cast<uint32_t>(TopK), static_cast<uint32_t>(stride), stream); |
| 61 | STD_TORCH_CHECK(status == cudaSuccess, |
| 62 | "FilteredTopK failed: ", cudaGetErrorString(status)); |
| 63 | } else { |
| 64 | STD_TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); |
| 65 | STD_TORCH_CHECK( |
| 66 | workspace.scalar_type() == torch::headeronly::ScalarType::Byte, |
| 67 | "workspace must be uint8"); |
| 68 | |
| 69 | int effective_max_smem; |
| 70 | if (num_rows <= 4) { |
| 71 | effective_max_smem = |
| 72 | std::min(max_smem_per_block, static_cast<int>(P::kSmemMedium)); |
| 73 | } else if (num_rows <= 8) { |
| 74 | constexpr int kSmemCapMedium = 48 * 1024; |
| 75 | effective_max_smem = std::min(max_smem_per_block, kSmemCapMedium); |
| 76 | } else { |
| 77 | effective_max_smem = max_smem_per_block; |
| 78 | } |
| 79 | |
| 80 | size_t available_for_ordered = |
| 81 | static_cast<size_t>(effective_max_smem) - P::kFixedSmemLarge; |
| 82 | uint32_t max_chunk_elements = |
| 83 | static_cast<uint32_t>(available_for_ordered / sizeof(uint32_t)); |
| 84 | |
| 85 | uint32_t vec_size = 1; |
| 86 | if (stride % 4 == 0) |
| 87 | vec_size = 4; |
| 88 | else if (stride % 2 == 0) |
| 89 | vec_size = 2; |
| 90 | |
| 91 | max_chunk_elements = (max_chunk_elements / vec_size) * vec_size; |
| 92 | uint32_t min_chunk = vec_size * P::kThreadsPerBlock; |
| 93 | if (max_chunk_elements < min_chunk) max_chunk_elements = min_chunk; |
| 94 | |
| 95 | uint32_t ctas_per_group = |
| 96 | (static_cast<uint32_t>(stride) + max_chunk_elements - 1) / |
| 97 | max_chunk_elements; |
| 98 | uint32_t chunk_size = |
| 99 | (static_cast<uint32_t>(stride) + ctas_per_group - 1) / ctas_per_group; |
| 100 | chunk_size = ((chunk_size + vec_size - 1) / vec_size) * vec_size; |
| 101 | if (chunk_size > max_chunk_elements) chunk_size = max_chunk_elements; |
| 102 | |
| 103 | size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t); |
| 104 | if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium; |
| 105 | |
| 106 | // Query occupancy for the instantiation that will actually launch; |
| 107 | // overestimating it deadlocks the cooperative barrier. |
| 108 | int occupancy = 1; |
| 109 | cudaError_t occ_err = cudaSuccess; |
| 110 | if (vec_size == 4) { |
| 111 | occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( |
| 112 | &occupancy, P::persistent_topk_kernel<TopK, 4>, P::kThreadsPerBlock, |
| 113 | smem_size); |
| 114 | } else if (vec_size == 2) { |
| 115 | occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( |
| 116 | &occupancy, P::persistent_topk_kernel<TopK, 2>, P::kThreadsPerBlock, |
| 117 | smem_size); |
| 118 | } else { |
| 119 | occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( |
| 120 | &occupancy, P::persistent_topk_kernel<TopK, 1>, P::kThreadsPerBlock, |
| 121 | smem_size); |
| 122 | } |
| 123 | STD_TORCH_CHECK(occ_err == cudaSuccess, |
| 124 | "persistent_topk occupancy query failed: ", |
| 125 | cudaGetErrorString(occ_err)); |
| 126 | if (occupancy < 1) occupancy = 1; |
| 127 | |
| 128 | // The cooperative spin-wait barrier only runs when at least one row hits |
| 129 | // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs |
| 130 | // early-exit, so oversubscription can't deadlock and headroom is wasted. |
| 131 | const bool needs_cooperative = |
| 132 | static_cast<uint32_t>(max_seq_len) > P::RADIX_THRESHOLD; |
| 133 | |
| 134 | const uint32_t hw_resident_cap = |
| 135 | static_cast<uint32_t>(num_sms) * static_cast<uint32_t>(occupancy); |
| 136 | uint32_t max_resident_ctas = hw_resident_cap; |
| 137 | if (needs_cooperative) { |
| 138 | // Reserve one CTA per SM when occupancy allows; fall back to a single |
| 139 | // CTA when occupancy == 1 (the most deadlock-prone case — any straggler |
| 140 | // kernel that takes the only slot on one SM hangs the barrier). Never |
| 141 | // drop below one full group's worth. |
| 142 | uint32_t headroom = (occupancy > 1) ? static_cast<uint32_t>(num_sms) : 1u; |
| 143 | if (max_resident_ctas >= headroom + ctas_per_group) { |
| 144 | max_resident_ctas -= headroom; |
| 145 | } |
| 146 | } |
| 147 | uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group, |
| 148 | static_cast<uint32_t>(num_rows)); |
| 149 | if (num_groups == 0) num_groups = 1; |
| 150 | uint32_t total_ctas = num_groups * ctas_per_group; |
| 151 | |
| 152 | // If the cooperative launch wouldn't fit, use the generic decode kernel on |
| 153 | // low-smem devices or FilteredTopK where its 128 KiB requirement is met. |
| 154 | if (needs_cooperative && total_ctas > hw_resident_cap) { |
| 155 | if (max_smem_per_block < 128 * 1024) { |
| 156 | const int64_t next_n = lengths.dim() == 2 ? lengths.size(1) : 1; |
| 157 | top_k_per_row_decode(logits, next_n, lengths, output, num_rows, |
| 158 | logits.stride(0), logits.stride(1), TopK); |
| 159 | return; |
| 160 | } |
| 161 | cudaError_t status = |
| 162 | vllm::FilteredTopKRaggedTransform<float, int32_t, TopK>( |
| 163 | logits.const_data_ptr<float>(), |
| 164 | output.mutable_data_ptr<int32_t>(), |
| 165 | lengths.const_data_ptr<int32_t>(), |
| 166 | static_cast<uint32_t>(num_rows), static_cast<uint32_t>(TopK), |
| 167 | static_cast<uint32_t>(stride), stream); |
| 168 | STD_TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ", |
| 169 | cudaGetErrorString(status)); |
| 170 | return; |
| 171 | } |
| 172 | |
| 173 | size_t state_bytes = num_groups * sizeof(P::RadixRowState); |
| 174 | STD_TORCH_CHECK(workspace.size(0) >= static_cast<int64_t>(state_bytes), |
| 175 | "workspace too small, need ", state_bytes, " bytes"); |
| 176 | |
| 177 | // Zero the per-group RadixRowState region before launch. |
| 178 | // |
| 179 | // Issued UNCONDITIONALLY so the memset is captured as its own node in |
| 180 | // the cudagraph (a separate cudaMemsetAsync node, sequenced before the |
| 181 | // persistent_topk_kernel launch on the same stream). The previous |
| 182 | // host-side guard `if (needs_cooperative)` was evaluated at capture time; |
| 183 | // when capture-time max_seq_len <= RADIX_THRESHOLD (always true under |
| 184 | // FULL_DECODE_ONLY with max_model_len < 32 K) the memset would NOT be |
| 185 | // captured, leaving the workspace state to accumulate across replays. |
| 186 | // That's a latent correctness bug if the runtime data ever takes the |
| 187 | // radix path, and removes one variable while debugging hangs in the |
| 188 | // decode/medium paths. |
| 189 | // |
| 190 | // Cost is sub-microsecond: state_bytes = num_groups * sizeof(RadixRowState) |
| 191 | // is ~3 KB per group, ~100 KB for the largest grids on this hardware. |
| 192 | // |
| 193 | // Why the memset is required (regardless of which path the kernel takes): |
| 194 | // 1. arrival_counter accumulates within a launch and is never reset, |
| 195 | // so a prior call leaves it at a large positive value. Without this |
| 196 | // reset, the very first wait_ge in the next call sees counter >> |
| 197 | // target and returns instantly, breaking the barrier. |
| 198 | // 2. The previous in-kernel init only ran in CTA-0 with intra-CTA |
| 199 | // __syncthreads(), so it had no happens-before edge to CTA-1+'s |
| 200 | // first red_release. cudaMemsetAsync is stream-ordered: the zero |
| 201 | // is globally visible before any CTA runs. |
| 202 | { |
| 203 | cudaError_t mz_err = cudaMemsetAsync( |
| 204 | workspace.mutable_data_ptr<uint8_t>(), 0, state_bytes, stream); |
| 205 | STD_TORCH_CHECK(mz_err == cudaSuccess, |
| 206 | "row_states memset failed: ", cudaGetErrorString(mz_err)); |
| 207 | } |
| 208 | |
| 209 | P::PersistentTopKParams params; |
| 210 | params.input = logits.const_data_ptr<float>(); |
| 211 | params.output = output.mutable_data_ptr<int32_t>(); |
| 212 | params.lengths = lengths.const_data_ptr<int32_t>(); |
| 213 | params.num_rows = static_cast<uint32_t>(num_rows); |
| 214 | params.stride = static_cast<uint32_t>(stride); |
| 215 | params.top_k = static_cast<uint32_t>(TopK); |
| 216 | params.chunk_size = chunk_size; |
| 217 | params.row_states = reinterpret_cast<P::RadixRowState*>( |
| 218 | workspace.mutable_data_ptr<uint8_t>()); |
| 219 | params.ctas_per_group = ctas_per_group; |
| 220 | params.max_seq_len = static_cast<uint32_t>(max_seq_len); |
| 221 | |
| 222 | #define LAUNCH_PERSISTENT(TOPK_VAL, VS) \ |
| 223 | do { \ |
| 224 | auto kernel = &P::persistent_topk_kernel<TOPK_VAL, VS>; \ |
| 225 | cudaError_t err = cudaFuncSetAttribute( \ |
| 226 | kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \ |
| 227 | STD_TORCH_CHECK(err == cudaSuccess, \ |
| 228 | "Failed to set smem: ", cudaGetErrorString(err)); \ |
| 229 | kernel<<<total_ctas, P::kThreadsPerBlock, smem_size, stream>>>(params); \ |
| 230 | } while (0) |
| 231 | |
| 232 | if (vec_size == 4) { |
| 233 | LAUNCH_PERSISTENT(TopK, 4); |
| 234 | } else if (vec_size == 2) { |
| 235 | LAUNCH_PERSISTENT(TopK, 2); |
| 236 | } else { |
| 237 | LAUNCH_PERSISTENT(TopK, 1); |
| 238 | } |
| 239 | #undef LAUNCH_PERSISTENT |
| 240 | } |
| 241 | |
| 242 | cudaError_t err = cudaGetLastError(); |
| 243 | STD_TORCH_CHECK(err == cudaSuccess, |
| 244 | "persistent_topk failed: ", cudaGetErrorString(err)); |
| 245 | } |
| 246 | #endif |
| 247 | |
| 248 | } // anonymous namespace |
| 249 | |
| 250 | void persistent_topk(const torch::stable::Tensor& logits, |
| 251 | const torch::stable::Tensor& lengths, |
| 252 | torch::stable::Tensor& output, |
| 253 | torch::stable::Tensor& workspace, int64_t k, |
| 254 | int64_t max_seq_len) { |
| 255 | #ifndef USE_ROCM |
| 256 | STD_TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); |
| 257 | STD_TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); |
| 258 | STD_TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); |
| 259 | STD_TORCH_CHECK(logits.scalar_type() == torch::headeronly::ScalarType::Float, |
| 260 | "Only float32 supported"); |
| 261 | STD_TORCH_CHECK(lengths.scalar_type() == torch::headeronly::ScalarType::Int, |
| 262 | "lengths must be int32"); |
| 263 | STD_TORCH_CHECK(output.scalar_type() == torch::headeronly::ScalarType::Int, |
| 264 | "output must be int32"); |
| 265 | STD_TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); |
| 266 | STD_TORCH_CHECK(lengths.dim() == 1 || lengths.dim() == 2, |
| 267 | "lengths must be 1D or 2D"); |
| 268 | STD_TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); |
| 269 | STD_TORCH_CHECK(output.dim() == 2, "output must be 2D"); |
| 270 | |
| 271 | const int64_t num_rows = logits.size(0); |
| 272 | |
| 273 | STD_TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); |
| 274 | STD_TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, |
| 275 | "output size mismatch"); |
| 276 | STD_TORCH_CHECK( |
| 277 | k == 512 || k == 1024 || k == 2048, |
| 278 | "persistent_topk supports k=512, k=1024, or k=2048, got k=", k); |
| 279 | |
| 280 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 281 | logits.get_device_index()); |
| 282 | |
| 283 | if (k == 512) { |
| 284 | launch_persistent_topk<512>(logits, lengths, output, workspace, |
| 285 | max_seq_len); |
| 286 | } else if (k == 1024) { |
| 287 | launch_persistent_topk<1024>(logits, lengths, output, workspace, |
| 288 | max_seq_len); |
| 289 | } else { |
| 290 | launch_persistent_topk<2048>(logits, lengths, output, workspace, |
| 291 | max_seq_len); |
| 292 | } |
| 293 | #else |
| 294 | STD_TORCH_CHECK(false, "persistent_topk is not supported on ROCm"); |
| 295 | #endif |
| 296 | } |