diff --git a/csrc/libtorch_stable/cooperative_topk.cuh b/csrc/libtorch_stable/cooperative_topk.cuh index f34664a66f7d..561dc71220d8 100644 --- a/csrc/libtorch_stable/cooperative_topk.cuh +++ b/csrc/libtorch_stable/cooperative_topk.cuh @@ -598,6 +598,11 @@ constexpr size_t kSmemSize2 = kSmemSize4; constexpr size_t kSmemSize8 = sizeof(SmemFused) + sizeof(int32_t) * 2048 + 128; +static_assert(sizeof(hist4096::Histogram4096Smem) <= kSmemSize4, + "cooperative smem must cover histogram_4096_topk"); +static_assert(sizeof(hist4096::Histogram4096Smem) <= kSmemSize8, + "cooperative smem must cover histogram_4096_topk"); + } // namespace cooperative } // namespace vllm diff --git a/csrc/libtorch_stable/persistent_topk.cuh b/csrc/libtorch_stable/persistent_topk.cuh index 9ae0fd4bcbaa..46d33c6393a2 100644 --- a/csrc/libtorch_stable/persistent_topk.cuh +++ b/csrc/libtorch_stable/persistent_topk.cuh @@ -172,13 +172,26 @@ __device__ __noinline__ void histogram_2048_topk( constexpr int MAX_ITEMS_PER_THREAD = (HIST2048_THRESHOLD + kThreadsPerBlock - 1) / kThreadsPerBlock; - enum : int { sTHR = 0, sOUT = 1, sREF = 2, sFIN = 3, sBUF0 = 4, sBUF1 = 5 }; + enum : int { + sTHR = 0, + sOUT = 1, + sREF = 2, + sFIN = 3, + sBUF0 = 4, + sBUF1 = 5, + sPOP = 6 + }; // ---- Initialize scalars (prevents stale data from prior rows) ---- if (tx < 8) { decode_smem[SBASE + tx] = 0; } + // Unfilled selection slots surface as -1 pads, not stale indices. + for (int i = tx; i < TopK; i += kThreadsPerBlock) { + output_indices[i] = -1; + } + // ---- Phase 1: Build 2048-bin histogram with float4 vectorized loads ---- int* histo = decode_smem; uint16_t reg_bins[MAX_ITEMS_PER_THREAD]; @@ -233,68 +246,184 @@ __device__ __noinline__ void histogram_2048_topk( if (pair_suffix >= TopK && (pair_suffix - h0) < TopK) { decode_smem[SBASE + sTHR] = 2 * tx; + decode_smem[SBASE + sPOP] = h0; } { const int right_suf = pair_suffix - h0; const int next_suf = pair_suffix - pair_sum; if (right_suf >= TopK && next_suf < TopK) { decode_smem[SBASE + sTHR] = 2 * tx + 1; + decode_smem[SBASE + sPOP] = pair_sum - h0; } } __syncthreads(); const int threshold = decode_smem[SBASE + sTHR]; + const uint32_t uthr = static_cast(threshold); + // Threshold-bin population from the finder's register. histo[threshold] + // aliases bufs[0] (written below in this same barrier interval). + const int bin_pop = decode_smem[SBASE + sPOP]; // ---- Phase 2: Collection with warp-aggregated atomicAdds ---- int* bufs[2] = {decode_smem + BOFF, decode_smem + BOFF + DBUF}; const int sOUT_abs = SBASE + sOUT; const int sBUF0_abs = SBASE + sBUF0; - { - const uint32_t uthr = static_cast(threshold); - int item = 0; - const int n_vec_iters = (n_vec + kThreadsPerBlock - 1) / kThreadsPerBlock; + int* refine[2] = {decode_smem, decode_smem + RHIST}; + auto compute_suffix_sum = [&]() { +#pragma unroll 8 + for (int i = 0; i < 8; ++i) { + if (tx < RADIX) { + const int stride = 1 << i; + const int s = i & 1; + const int d = s ^ 1; + int value = refine[s][tx]; + if (tx < RADIX - stride) value += refine[s][tx + stride]; + refine[d][tx] = value; + } + __syncthreads(); + } + }; + + if (bin_pop <= DBUF) { + { + const uint32_t uthr = static_cast(threshold); + int item = 0; + const int n_vec_iters = (n_vec + kThreadsPerBlock - 1) / kThreadsPerBlock; - for (int iter = 0; iter < n_vec_iters; iter++) { - const int i = tx + iter * kThreadsPerBlock; - const bool vec_valid = (i < n_vec); - const int base_idx = i << 2; + for (int iter = 0; iter < n_vec_iters; iter++) { + const int i = tx + iter * kThreadsPerBlock; + const bool vec_valid = (i < n_vec); + const int base_idx = i << 2; #pragma unroll 4 - for (int sub = 0; sub < 4; sub++) { - const int elem_idx = base_idx + sub; - uint32_t bin = 0; - if (vec_valid) bin = reg_bins[item++]; - const bool is_above = vec_valid && (bin > uthr); - const bool is_equal = vec_valid && (bin == uthr); - - const uint32_t above_mask = __ballot_sync(0xffffffff, is_above); - if (above_mask) { - const int above_count = __popc(above_mask); - const int above_rank = __popc(above_mask & ((1u << lane) - 1)); - int above_base; - if (lane == 0) { - above_base = atomicAdd(&decode_smem[sOUT_abs], above_count); + for (int sub = 0; sub < 4; sub++) { + const int elem_idx = base_idx + sub; + uint32_t bin = 0; + if (vec_valid) bin = reg_bins[item++]; + const bool is_above = vec_valid && (bin > uthr); + const bool is_equal = vec_valid && (bin == uthr); + + const uint32_t above_mask = __ballot_sync(0xffffffff, is_above); + if (above_mask) { + const int above_count = __popc(above_mask); + const int above_rank = __popc(above_mask & ((1u << lane) - 1)); + int above_base; + if (lane == 0) { + above_base = atomicAdd(&decode_smem[sOUT_abs], above_count); + } + above_base = __shfl_sync(0xffffffff, above_base, 0); + if (is_above && above_base + above_rank < TopK) { + output_indices[above_base + above_rank] = elem_idx; + } } - above_base = __shfl_sync(0xffffffff, above_base, 0); - if (is_above) { - output_indices[above_base + above_rank] = elem_idx; + + const uint32_t equal_mask = __ballot_sync(0xffffffff, is_equal); + if (equal_mask) { + const int equal_count = __popc(equal_mask); + const int equal_rank = __popc(equal_mask & ((1u << lane) - 1)); + int equal_base; + if (lane == 0) { + equal_base = atomicAdd(&decode_smem[sBUF0_abs], equal_count); + } + equal_base = __shfl_sync(0xffffffff, equal_base, 0); + if (is_equal && + __builtin_expect(equal_base + equal_rank < DBUF, 1)) { + bufs[0][equal_base + equal_rank] = elem_idx; + } } } - - const uint32_t equal_mask = __ballot_sync(0xffffffff, is_equal); - if (equal_mask) { - const int equal_count = __popc(equal_mask); - const int equal_rank = __popc(equal_mask & ((1u << lane) - 1)); - int equal_base; - if (lane == 0) { - equal_base = atomicAdd(&decode_smem[sBUF0_abs], equal_count); + } + } + } else { + // Overflow path: the threshold bin exceeds the stash. Descend the + // remaining FP32 key bytes until it fits; clip only exact ties. + int p1 = -1, p2 = -1, p3 = -1; + const auto key_participates = [&](int idx, int level) -> bool { + if (decode_bin(logits[idx]) != uthr) return false; + const uint32_t key = convert_to_uint32_v2(logits[idx]); + if (level >= 2 && static_cast((key >> 24) & 0xFF) != p1) + return false; + if (level >= 3 && static_cast((key >> 16) & 0xFF) != p2) + return false; + if (level >= 4 && static_cast((key >> 8) & 0xFF) != p3) return false; + return true; + }; + // Fill the definite members above the coarse threshold bin. + for (int idx = tx; idx < seq_len; idx += kThreadsPerBlock) { + if (static_cast(decode_bin(logits[idx])) > uthr) { + const int pos = atomicAdd(&decode_smem[sOUT_abs], 1); + if (pos < TopK) output_indices[pos] = idx; + } + } + __syncthreads(); + int rem = TopK - decode_smem[SBASE + sOUT]; + for (int level = 1; level <= 4; ++level) { + const int shift = 24 - 8 * (level - 1); + for (int i = tx; i < RHIST; i += kThreadsPerBlock) { + refine[0][i] = 0; + } + __syncthreads(); + for (int idx = tx; idx < seq_len; idx += kThreadsPerBlock) { + if (key_participates(idx, level)) { + const uint32_t key = convert_to_uint32_v2(logits[idx]); + atomicAdd(&refine[0][(key >> shift) & 0xFF], 1); + } + } + __syncthreads(); + compute_suffix_sum(); + if (tx == 0) { + decode_smem[SBASE + sTHR] = -1; + } + __syncthreads(); + if (tx < RADIX && refine[0][tx] > rem && refine[0][tx + 1] <= rem) { + decode_smem[SBASE + sTHR] = tx; + } + __syncthreads(); + const int thr = decode_smem[SBASE + sTHR]; + const int above = thr < 0 ? 0 : refine[0][thr + 1]; + const int pop = thr < 0 ? 0 : refine[0][thr] - above; + // Fill this level's definite members. + for (int idx = tx; idx < seq_len; idx += kThreadsPerBlock) { + if (key_participates(idx, level)) { + const uint32_t key = convert_to_uint32_v2(logits[idx]); + if (static_cast((key >> shift) & 0xFF) > thr) { + const int pos = atomicAdd(&decode_smem[sOUT_abs], 1); + if (pos < TopK) output_indices[pos] = idx; } - equal_base = __shfl_sync(0xffffffff, equal_base, 0); - if (is_equal && __builtin_expect(equal_base + equal_rank < DBUF, 1)) { - bufs[0][equal_base + equal_rank] = elem_idx; + } + } + __syncthreads(); + rem -= above; + if (rem == 0) { + // Already filled above this bin. + return; + } + if (pop <= DBUF || level == 4) { + // Terminal bin: stash it; a capacity clip only drops exact ties. + if (tx == 0) { + decode_smem[sBUF0_abs] = 0; + } + __syncthreads(); + for (int idx = tx; idx < seq_len; idx += kThreadsPerBlock) { + if (!key_participates(idx, level)) continue; + const uint32_t key = convert_to_uint32_v2(logits[idx]); + if (static_cast((key >> shift) & 0xFF) == thr) { + const int bp = atomicAdd(&decode_smem[sBUF0_abs], 1); + if (__builtin_expect(bp < DBUF, 1)) { + bufs[0][bp] = idx; + } } } + __syncthreads(); + break; + } + if (level == 1) { + p1 = thr; + } else if (level == 2) { + p2 = thr; + } else if (level == 3) { + p3 = thr; } } } @@ -309,14 +438,13 @@ __device__ __noinline__ void histogram_2048_topk( const int nb = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; const int base = decode_smem[SBASE + sOUT]; for (int i = tx; i < nb; i += kThreadsPerBlock) { - output_indices[base + i] = bufs[0][i]; + if (base + i < TopK) output_indices[base + i] = bufs[0][i]; } __syncthreads(); return; } // ---- Phase 3: Deferred refinement (rare path) ---- - int* refine[2] = {decode_smem, decode_smem + RHIST}; const int num_buf0 = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; for (int i = tx; i < RHIST; i += kThreadsPerBlock) { @@ -330,21 +458,6 @@ __device__ __noinline__ void histogram_2048_topk( } __syncthreads(); - auto compute_suffix_sum = [&]() { -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { - if (tx < RADIX) { - const int stride = 1 << i; - const int s = i & 1; - const int d = s ^ 1; - int value = refine[s][tx]; - if (tx < RADIX - stride) value += refine[s][tx + stride]; - refine[d][tx] = value; - } - __syncthreads(); - } - }; - #pragma unroll 4 for (int pass = 0; pass < 4; ++pass) { const int src = pass & 1; @@ -355,16 +468,22 @@ __device__ __noinline__ void histogram_2048_topk( compute_suffix_sum(); + // Reset the stash target and selection scalars independently of the finder. + if (tx == 0) { + decode_smem[SBASE + sREF] = -1; + decode_smem[SBASE + sBUF0 + dst] = 0; + decode_smem[SBASE + sFIN] = 0; + } + __syncthreads(); if (tx < RADIX && refine[0][tx] > remaining_k && refine[0][tx + 1] <= remaining_k) { decode_smem[SBASE + sREF] = tx; - decode_smem[SBASE + sBUF0 + dst] = 0; decode_smem[SBASE + sFIN] = remaining_k - refine[0][tx + 1]; } __syncthreads(); const int ref_thr = decode_smem[SBASE + sREF]; - remaining_k -= refine[0][ref_thr + 1]; + remaining_k -= ref_thr < 0 ? 0 : refine[0][ref_thr + 1]; const int bit_offset = 24 - pass * 8; if (remaining_k == 0) { @@ -373,7 +492,7 @@ __device__ __noinline__ void histogram_2048_topk( const uint32_t fp32 = convert_to_uint32_v2(logits[idx]); if (((fp32 >> bit_offset) & 0xFF) > static_cast(ref_thr)) { const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); - output_indices[pos] = idx; + if (pos < TopK) output_indices[pos] = idx; } } __syncthreads(); @@ -392,11 +511,11 @@ __device__ __noinline__ void histogram_2048_topk( if (bin > ref_thr) { const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); - output_indices[pos] = idx; + if (pos < TopK) output_indices[pos] = idx; } else if (bin == ref_thr) { if (pass == 3) { const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1); - if (slot > 0) output_indices[TopK - slot] = idx; + if (slot > 0 && slot <= TopK) output_indices[TopK - slot] = idx; } else { const int bp = atomicAdd(&decode_smem[SBASE + sBUF0 + dst], 1); if (__builtin_expect(bp < DBUF, 1)) { @@ -442,6 +561,15 @@ __device__ __noinline__ void histogram_256_topk( const int thread_id = threadIdx.x; int remaining_k = TopK; + static_assert(TopK <= MAX_BUFFERED_ITEMS, + "selection must fit one stash round"); + + // Unfilled selection slots surface as -1 pads, not stale indices. + for (int i = thread_id; i < TopK; i += kThreadsPerBlock) { + output_indices[i] = -1; + } + __syncthreads(); + if (thread_id < RADIX + 1) { shared_histogram[0][thread_id] = 0; } @@ -472,52 +600,167 @@ __device__ __noinline__ void histogram_256_topk( compute_cumulative_sum(); + // Init selection scalars unconditionally: a no-finder round must not + // consume stale state from a previous row (threshold_bin < 0 -> no-op). + if (thread_id == 0) { + shared_threshold_bin = -1; + shared_buffered_count[0] = 0; + shared_output_count = 0; + shared_final_k = 0; + } + __syncthreads(); if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && shared_histogram[0][thread_id + 1] <= remaining_k) { shared_threshold_bin = thread_id; - shared_buffered_count[0] = 0; - shared_output_count = 0; } __syncthreads(); const int threshold_bin = shared_threshold_bin; - remaining_k -= shared_histogram[0][threshold_bin + 1]; + const int above_coarse = + threshold_bin < 0 ? 0 : shared_histogram[0][threshold_bin + 1]; + const int bin_pop = + threshold_bin < 0 ? 0 : shared_histogram[0][threshold_bin] - above_coarse; + remaining_k -= above_coarse; if (remaining_k == 0) { for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { const int bin = convert_to_uint8(logits[idx + logits_offset]); if (bin > threshold_bin) { const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; + if (output_pos < TopK) output_indices[output_pos] = idx; } } __syncthreads(); return; } - __syncthreads(); - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; - } - __syncthreads(); + if (bin_pop <= MAX_BUFFERED_ITEMS) { + // Fast path: the threshold bin fits the stash whole. + __syncthreads(); + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const float logit_value = logits[idx + logits_offset]; - const int bin = convert_to_uint8(logit_value); - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } else if (bin == threshold_bin) { - const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); - if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { - buffered_indices[0][buffer_pos] = idx; - const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); - const int next_bin = (fp32_bits >> 24) & 0xFF; - atomicAdd(&shared_histogram[0][next_bin], 1); + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const float logit_value = logits[idx + logits_offset]; + const int bin = convert_to_uint8(logit_value); + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + if (output_pos < TopK) output_indices[output_pos] = idx; + } else if (bin == threshold_bin) { + const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); + if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { + buffered_indices[0][buffer_pos] = idx; + const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); + const int next_bin = (fp32_bits >> 24) & 0xFF; + atomicAdd(&shared_histogram[0][next_bin], 1); + } + } + } + __syncthreads(); + } else { + // Overflow path: the threshold bin exceeds the stash. Descend the + // remaining FP32 key bytes until it fits; clip only exact ties. + int p1 = -1, p2 = -1, p3 = -1; + const auto key_participates = [&](float logit_value, int level) -> bool { + if (static_cast(convert_to_uint8(logit_value)) != threshold_bin) + return false; + const uint32_t key = convert_to_uint32_v2(logit_value); + if (level >= 2 && static_cast((key >> 24) & 0xFF) != p1) + return false; + if (level >= 3 && static_cast((key >> 16) & 0xFF) != p2) + return false; + if (level >= 4 && static_cast((key >> 8) & 0xFF) != p3) return false; + return true; + }; + // Fill the definite members above the coarse threshold bin. + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + if (static_cast(convert_to_uint8(logits[idx + logits_offset])) > + threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + if (output_pos < TopK) output_indices[output_pos] = idx; + } + } + __syncthreads(); + for (int level = 1; level <= 4; ++level) { + const int shift = 24 - 8 * (level - 1); + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const float logit_value = logits[idx + logits_offset]; + if (key_participates(logit_value, level)) { + const uint32_t key = convert_to_uint32_v2(logit_value); + atomicAdd(&shared_histogram[0][(key >> shift) & 0xFF], 1); + } + } + __syncthreads(); + compute_cumulative_sum(); + if (thread_id == 0) { + shared_threshold_bin = -1; + } + __syncthreads(); + if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && + shared_histogram[0][thread_id + 1] <= remaining_k) { + shared_threshold_bin = thread_id; + } + __syncthreads(); + const int thr = shared_threshold_bin; + const int above = thr < 0 ? 0 : shared_histogram[0][thr + 1]; + const int pop = thr < 0 ? 0 : shared_histogram[0][thr] - above; + // Fill this level's definite members. + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const float logit_value = logits[idx + logits_offset]; + if (key_participates(logit_value, level)) { + const uint32_t key = convert_to_uint32_v2(logit_value); + if (static_cast((key >> shift) & 0xFF) > thr) { + const int output_pos = atomicAdd(&shared_output_count, 1); + if (output_pos < TopK) output_indices[output_pos] = idx; + } + } + } + __syncthreads(); + remaining_k -= above; + if (remaining_k == 0) { + // Already filled above this bin. + return; + } + if (pop <= MAX_BUFFERED_ITEMS || level == 4) { + // Terminal bin: stash it for the refine rounds (entry state matches). + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + if (thread_id == 0) { + shared_buffered_count[0] = 0; + } + __syncthreads(); + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const float logit_value = logits[idx + logits_offset]; + if (!key_participates(logit_value, level)) continue; + const uint32_t key = convert_to_uint32_v2(logit_value); + if (static_cast((key >> shift) & 0xFF) == thr) { + const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); + if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { + buffered_indices[0][buffer_pos] = idx; + atomicAdd(&shared_histogram[0][(key >> 24) & 0xFF], 1); + } + } + } + __syncthreads(); + break; + } + if (level == 1) { + p1 = thr; + } else if (level == 2) { + p2 = thr; + } else if (level == 3) { + p3 = thr; } } } - __syncthreads(); #pragma unroll 4 for (int pass = 0; pass < 4; ++pass) { @@ -529,16 +772,23 @@ __device__ __noinline__ void histogram_256_topk( compute_cumulative_sum(); + // Reset the stash target and selection scalars independently of the finder. + if (thread_id == 0) { + shared_threshold_bin = -1; + shared_buffered_count[dst_buffer] = 0; + shared_final_k = 0; + } + __syncthreads(); if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && shared_histogram[0][thread_id + 1] <= remaining_k) { shared_threshold_bin = thread_id; - shared_buffered_count[dst_buffer] = 0; shared_final_k = remaining_k - shared_histogram[0][thread_id + 1]; } __syncthreads(); const int threshold_bin = shared_threshold_bin; - remaining_k -= shared_histogram[0][threshold_bin + 1]; + remaining_k -= + threshold_bin < 0 ? 0 : shared_histogram[0][threshold_bin + 1]; const int bit_offset = 24 - pass * 8; if (remaining_k == 0) { @@ -549,7 +799,7 @@ __device__ __noinline__ void histogram_256_topk( const int bin = (fp32_bits >> bit_offset) & 0xFF; if (bin > threshold_bin) { const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; + if (output_pos < TopK) output_indices[output_pos] = idx; } } __syncthreads(); @@ -569,11 +819,11 @@ __device__ __noinline__ void histogram_256_topk( const int bin = (fp32_bits >> bit_offset) & 0xFF; if (bin > threshold_bin) { const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; + if (output_pos < TopK) output_indices[output_pos] = idx; } else if (bin == threshold_bin) { if (pass == 3) { const int slot = atomicAdd(&shared_final_k, -1); - if (slot > 0) { + if (slot > 0 && slot <= TopK) { output_indices[TopK - slot] = idx; } } else { @@ -1029,6 +1279,10 @@ constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = 16 * 1024; // 16K indices per buffer constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB +// The dynamic allocation doubles as the hist4096 short path's smem. +static_assert( + FILTERED_TOPK_SMEM_DYNAMIC >= sizeof(hist4096::Histogram4096Smem<2048, 12>), + "FilteredTopK dynamic smem must cover the histogram_4096_topk path"); /*! * \brief Filtered Top-K kernel for ragged sequences. @@ -1095,6 +1349,11 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) using Traits = FilteredTopKTraits; int topk = top_k; + // Unfilled selection slots surface as -1 pads, not stale reads. + for (int i = tx; i < static_cast(top_k); i += BLOCK_SIZE) { + s_indices[i] = -1; + } + // Stage 1: 8-bit coarse histogram with vectorized loads if (tx < RADIX + 1) s_histogram[tx] = 0; __syncthreads(); @@ -1137,15 +1396,24 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) }; run_cumsum(); - if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { - s_threshold_bin_id = tx; + // Init selection scalars unconditionally (threshold_bin < 0 -> no-op). + if (tx == 0) { + s_threshold_bin_id = -1; s_num_input[0] = 0; s_counter = 0; } __syncthreads(); + if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { + s_threshold_bin_id = tx; + } + __syncthreads(); const auto threshold_bin = s_threshold_bin_id; - topk -= s_histogram[threshold_bin + 1]; + const int above_coarse = + threshold_bin < 0 ? 0 : s_histogram[threshold_bin + 1]; + const int bin_pop = + threshold_bin < 0 ? 0 : s_histogram[threshold_bin] - above_coarse; + topk -= above_coarse; constexpr int NUM_ROUNDS = Traits::NUM_REFINE_ROUNDS; constexpr int FIRST_SHIFT = Traits::FIRST_REFINE_SHIFT; @@ -1161,7 +1429,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const auto bin = static_cast(Traits::ToCoarseKey(score_vec[j])); if (bin > threshold_bin) { const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = base + j; + if (pos < static_cast(top_k)) s_indices[pos] = base + j; } } } @@ -1170,48 +1438,147 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const auto bin = static_cast(Traits::ToCoarseKey(score[i])); if (bin > threshold_bin) { const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = i; + if (pos < static_cast(top_k)) s_indices[pos] = i; } } __syncthreads(); } else { - __syncthreads(); - if (tx < RADIX + 1) s_histogram[tx] = 0; - __syncthreads(); + if (bin_pop <= static_cast(SMEM_INPUT_SIZE)) { + // Fast path: the threshold bin fits the stash whole. + __syncthreads(); + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); - // Filter + histogram for refinement - auto filter_and_add_to_histogram = [&](auto raw_input, int index) { - const auto bin = static_cast(Traits::ToCoarseKey(raw_input)); - if (bin > threshold_bin) { - const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = index; - } else if (bin == threshold_bin) { - const auto pos = atomicAdd(&s_num_input[0], 1); - if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { - s_input_idx[0][pos] = index; - const auto ordered = Traits::ToOrdered(raw_input); - const auto sub_bin = (ordered >> FIRST_SHIFT) & 0xFF; - atomicAdd(&s_histogram[sub_bin], 1); + // Filter + histogram for refinement + auto filter_and_add_to_histogram = [&](auto raw_input, int index) { + const auto bin = static_cast(Traits::ToCoarseKey(raw_input)); + if (bin > threshold_bin) { + const auto pos = atomicAdd(&s_counter, 1); + if (pos < static_cast(top_k)) s_indices[pos] = index; + } else if (bin == threshold_bin) { + const auto pos = atomicAdd(&s_num_input[0], 1); + if (__builtin_expect(pos < int(SMEM_INPUT_SIZE), 1)) { + s_input_idx[0][pos] = index; + const auto ordered = Traits::ToOrdered(raw_input); + const auto sub_bin = (ordered >> FIRST_SHIFT) & 0xFF; + atomicAdd(&s_histogram[sub_bin], 1); + } } - } - }; + }; #pragma unroll 2 - for (int base = tx * VEC_SIZE; base < aligned_length; - base += BLOCK_SIZE * VEC_SIZE) { - score_vec.cast_load(&score[base]); + for (int base = tx * VEC_SIZE; base < aligned_length; + base += BLOCK_SIZE * VEC_SIZE) { + score_vec.cast_load(&score[base]); #pragma unroll - for (int j = 0; j < VEC_SIZE; ++j) { - filter_and_add_to_histogram(score_vec[j], base + j); + for (int j = 0; j < VEC_SIZE; ++j) { + filter_and_add_to_histogram(score_vec[j], base + j); + } + } + // Handle tail + for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { + filter_and_add_to_histogram(score[i], i); + } + __syncthreads(); + } else { + // Overflow path: the threshold bin exceeds the stash. Descend the + // remaining FP32 key bytes until it fits; clip only exact ties. + int p1 = -1, p2 = -1, p3 = -1; + const auto key_participates = [&](auto raw_input, int level) -> bool { + if (static_cast(Traits::ToCoarseKey(raw_input)) != threshold_bin) + return false; + const auto key = Traits::ToOrdered(raw_input); + if (level >= 2 && static_cast((key >> 24) & 0xFF) != p1) + return false; + if (level >= 3 && static_cast((key >> 16) & 0xFF) != p2) + return false; + if (level >= 4 && static_cast((key >> 8) & 0xFF) != p3) + return false; + return true; + }; + // Fill the definite members above the coarse threshold bin. + for (int i = tx; i < length; i += BLOCK_SIZE) { + if (static_cast(Traits::ToCoarseKey(score[i])) > threshold_bin) { + const auto pos = atomicAdd(&s_counter, 1); + if (pos < static_cast(top_k)) s_indices[pos] = i; + } + } + __syncthreads(); + for (int level = 1; level <= 4; ++level) { + const int shift = FIRST_SHIFT - 8 * (level - 1); + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + for (int i = tx; i < length; i += BLOCK_SIZE) { + const auto raw_input = score[i]; + if (key_participates(raw_input, level)) { + const auto key = Traits::ToOrdered(raw_input); + atomicAdd(&s_histogram[(key >> shift) & 0xFF], 1); + } + } + __syncthreads(); + run_cumsum(); + if (tx == 0) { + s_threshold_bin_id = -1; + } + __syncthreads(); + if (tx < RADIX && s_histogram[tx] > topk && + s_histogram[tx + 1] <= topk) { + s_threshold_bin_id = tx; + } + __syncthreads(); + const int thr = s_threshold_bin_id; + const int above = thr < 0 ? 0 : s_histogram[thr + 1]; + const int pop = thr < 0 ? 0 : s_histogram[thr] - above; + // Fill this level's definite members. + for (int i = tx; i < length; i += BLOCK_SIZE) { + const auto raw_input = score[i]; + if (key_participates(raw_input, level)) { + const auto key = Traits::ToOrdered(raw_input); + if (static_cast((key >> shift) & 0xFF) > thr) { + const auto pos = atomicAdd(&s_counter, 1); + if (pos < static_cast(top_k)) s_indices[pos] = i; + } + } + } + __syncthreads(); + topk -= above; + if (topk == 0) { + // Already filled above this bin. + goto output_phase; + } + if (pop <= static_cast(SMEM_INPUT_SIZE) || level == 4) { + // Terminal bin: stash it for the refine rounds (entry state matches). + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + if (tx == 0) { + s_num_input[0] = 0; + } + __syncthreads(); + for (int i = tx; i < length; i += BLOCK_SIZE) { + const auto raw_input = score[i]; + if (!key_participates(raw_input, level)) continue; + const auto key = Traits::ToOrdered(raw_input); + if (static_cast((key >> shift) & 0xFF) == thr) { + const auto pos = atomicAdd(&s_num_input[0], 1); + if (__builtin_expect(pos < int(SMEM_INPUT_SIZE), 1)) { + s_input_idx[0][pos] = i; + atomicAdd(&s_histogram[(key >> 24) & 0xFF], 1); + } + } + } + __syncthreads(); + break; + } + if (level == 1) { + p1 = thr; + } else if (level == 2) { + p2 = thr; + } else if (level == 3) { + p3 = thr; + } } } - // Handle tail - for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { - filter_and_add_to_histogram(score[i], i); - } - __syncthreads(); - // Stage 2: refine with 8bit radix passes -#pragma unroll +#pragma unroll 4 for (int round = 0; round < NUM_ROUNDS; ++round) { __shared__ int s_last_remain; const auto r_idx = round % 2; @@ -1221,15 +1588,22 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; run_cumsum(); + // Reset the stash target and selection scalars independently of the + // finder. + if (tx == 0) { + s_threshold_bin_id = -1; + s_num_input[r_idx ^ 1] = 0; + s_last_remain = 0; + } + __syncthreads(); if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { s_threshold_bin_id = tx; - s_num_input[r_idx ^ 1] = 0; s_last_remain = topk - s_histogram[tx + 1]; } __syncthreads(); const auto threshold = s_threshold_bin_id; - topk -= s_histogram[threshold + 1]; + topk -= threshold < 0 ? 0 : s_histogram[threshold + 1]; const int offset = FIRST_SHIFT - round * 8; const bool is_last_round = (round == NUM_ROUNDS - 1); @@ -1240,7 +1614,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const auto bin = (Traits::ToOrdered(score[idx]) >> offset) & 0xFF; if (static_cast(bin) > threshold) { const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = idx; + if (pos < static_cast(top_k)) s_indices[pos] = idx; } } __syncthreads(); @@ -1255,11 +1629,11 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) const auto bin = (Traits::ToOrdered(raw_input) >> offset) & 0xFF; if (static_cast(bin) > threshold) { const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = idx; + if (pos < static_cast(top_k)) s_indices[pos] = idx; } else if (static_cast(bin) == threshold) { if (is_last_round) { const auto pos = atomicAdd(&s_last_remain, -1); - if (pos > 0) { + if (pos > 0 && pos <= static_cast(top_k)) { s_indices[top_k - pos] = idx; } } else { @@ -1278,6 +1652,7 @@ __global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) } } +output_phase: // Output phase - mode-specific #pragma unroll 2 for (int base = tx; base < static_cast(top_k); base += BLOCK_SIZE) { diff --git a/csrc/libtorch_stable/topk_histogram_4096.cuh b/csrc/libtorch_stable/topk_histogram_4096.cuh index 5f9f823a3399..0b1c8ce8085e 100644 --- a/csrc/libtorch_stable/topk_histogram_4096.cuh +++ b/csrc/libtorch_stable/topk_histogram_4096.cuh @@ -321,6 +321,14 @@ struct Histogram4096Smem { alignas(128) uint32_t counter_eq; MatchBin match; uint32_t warp_sum[kNumWarps]; + // Oversized-bin descent state: dlevel = descended key-byte count (0 = + // threshold bin fit the tie stash), dthr[i] = chosen key byte 24 - 8*i, + // dabove/dpop/dlvl_above = cumulative above, terminal pop, level above. + uint32_t dthr[4]; + uint32_t dlevel; + uint32_t dabove; + uint32_t dpop; + uint32_t dlvl_above; union { uint32_t histogram[HIST_BINS]; Tie tie_buffer[TIE_CAPACITY]; @@ -442,49 +450,185 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, __syncthreads(); - // Phase 3: Scatter from registers + // Phase 2.5: oversized threshold bin (> TopK tie stash). Stashing an + // arrival-order subset would drop top-K candidates, so descend the FP32 + // key bytes until the bin fits; clip only exact ties. const auto [thr_bin, num_above, num_equal] = smem->match; - const bool need_tie = (num_equal + num_above > TopK); + if (tx == 0) { + smem->dlevel = 0; + smem->dabove = num_above; + smem->dpop = num_equal; + smem->dlvl_above = 0; + } + __syncthreads(); - done = false; + if (num_equal > TopK) { + for (uint32_t level = 1; level <= 4; ++level) { + const uint32_t sh = 24 - 8 * (level - 1); + const uint32_t remain = TopK - smem->dabove; + if (tx < RADIX) { + smem->histogram[tx] = 0; + } + __syncthreads(); + done = false; #pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { - const float* elems = reinterpret_cast(&vecs[v]); + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); #pragma unroll - for (uint32_t e = 0; e < 4 && !done; e++) { - const uint32_t idx = (tx + v * kBlockSize) * 4 + e; - if (idx >= length) { - done = true; - } else { - const uint32_t bin = extract_coarse_bin_N(elems[e]); - if (bin > thr_bin) { - output[atomicAdd(&smem->counter_gt, 1)] = - idx; // above -> output directly - } else if (bin == thr_bin) { - const auto pos = atomicAdd(&smem->counter_eq, 1); - if (!need_tie) { - if (pos + num_above < TopK) { - output[pos + num_above] = idx; // all fit + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else if (extract_coarse_bin_N(elems[e]) == thr_bin) { + const uint32_t key = convert_to_uint32_v2(elems[e]); + bool participates = true; +#pragma unroll + for (uint32_t l = 0; l < 4; ++l) { + if (l >= level - 1) break; + if (((key >> (24 - 8 * l)) & 0xFF) != smem->dthr[l]) { + participates = false; + break; + } } - } else { - if (pos < TopK) { - smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement + if (participates) { + atomicAdd(&smem->histogram[(key >> sh) & 0xFF], 1); } } } - // else: bin < thr_bin - discard (not in top-k) } + __syncthreads(); + + // Suffix scan + threshold finder (same pattern as tie_handle). + uint32_t hv = 0, wi2 = 0; + if (tx < RADIX) { + hv = smem->histogram[tx]; + wi2 = warp_inclusive_sum(lane_id, hv); + if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = wi2; + } + __syncthreads(); + if (tx < RADIX) { + const auto tmp = + (lane_id < RADIX / kWarpSize) ? smem->warp_sum[lane_id] : 0; + const auto tot = warp_reduce_sum_full(tmp); + const auto inter = warp_reduce_sum_full(lane_id < warp_id ? tmp : 0); + const auto above = tot - (inter + wi2); + if (above < remain && above + hv >= remain) { + smem->dthr[level - 1] = tx; + smem->dpop = hv; + smem->dlvl_above = above; + } + } + __syncthreads(); + if (tx == 0) { + smem->dabove += smem->dlvl_above; + if (smem->dpop <= TopK || level == 4) { + smem->dlevel = level; + } + } + __syncthreads(); + if (smem->dlevel != 0) break; } } + const bool descended = smem->dlevel != 0; + const uint32_t eff_above = descended ? smem->dabove : num_above; + const uint32_t eff_equal = descended ? smem->dpop : num_equal; - // Phase 4: Tie-breaking - if (!need_tie) return; - __syncthreads(); + // Phase 3: Scatter from registers + const bool need_tie = (eff_equal + eff_above > TopK); + + if (descended) { + const uint32_t lvls = smem->dlevel; + done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + const uint32_t bin = extract_coarse_bin_N(elems[e]); + if (bin > thr_bin) { + const auto pos = atomicAdd(&smem->counter_gt, 1); + if (pos < TopK) output[pos] = idx; + } else if (bin == thr_bin) { + // First diverged key byte decides membership; full match = + // terminal bin. + const uint32_t key = convert_to_uint32_v2(elems[e]); + int cls = 3; +#pragma unroll + for (uint32_t l = 0; l < 4; ++l) { + if (l >= lvls) break; + const uint32_t b = (key >> (24 - 8 * l)) & 0xFF; + if (b > smem->dthr[l]) { + cls = 1; + break; + } + if (b < smem->dthr[l]) { + cls = 2; + break; + } + } + if (cls == 1) { + const auto pos = atomicAdd(&smem->counter_gt, 1); + if (pos < TopK) output[pos] = idx; + } else if (cls == 3) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (!need_tie) { + if (pos + eff_above < TopK) output[pos + eff_above] = idx; + } else if (pos < TopK) { + smem->tie_buffer[pos] = {idx, elems[e]}; + } + } + } + } + } + } + if (!need_tie) return; + __syncthreads(); + } else { + done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + const uint32_t bin = extract_coarse_bin_N(elems[e]); + if (bin > thr_bin) { + output[atomicAdd(&smem->counter_gt, 1)] = + idx; // above -> output directly + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (!need_tie) { + if (pos + num_above < TopK) { + output[pos + num_above] = idx; // all fit + } + } else { + if (pos < TopK) { + smem->tie_buffer[pos] = {idx, + elems[e]}; // store for refirement + } + } + } + // else: bin < thr_bin - discard (not in top-k) + } + } + } + + // Phase 4: Tie-breaking + if (!need_tie) return; + __syncthreads(); + } // Fast warp-ballot tie-breaking for small tie counts - const uint32_t num_ties = min(num_equal, static_cast(TopK)); + const uint32_t num_ties = min(eff_equal, static_cast(TopK)); const uint32_t topk_remain = - TopK - num_above; // pick exactly remaining elements to fill topK + TopK - eff_above; // pick exactly remaining elements to fill topK auto is_greater = [](const Tie& a, const Tie& b) { return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); @@ -506,7 +650,7 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, const auto rank = static_cast( __popc(__ballot_sync(mask, pred))); // count how many are greater if (lane_id == 0 && rank < topk_remain) { - output[num_above + rank] = target.idx; // place at correct position + output[eff_above + rank] = target.idx; // place at correct position } } else if (num_ties <= kWarpSize * @@ -526,7 +670,7 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, const auto r1 = __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); if (lane_id == 0 && r0 + r1 < topk_remain) - output[num_above + r0 + r1] = target.idx; + output[eff_above + r0 + r1] = target.idx; } if (warp1 < num_ties) { const auto target = smem->tie_buffer[warp1]; @@ -535,14 +679,14 @@ __device__ void histogram_4096_topk(const float* __restrict__ scores, const auto r1 = __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); if (lane_id == 0 && r0 + r1 < topk_remain) - output[num_above + r0 + r1] = target.idx; + output[eff_above + r0 + r1] = target.idx; } } else { // Large tie count: fall back to 4-round radix-256 sort if constexpr (TopK <= kBlockSize) { - tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); + tie_handle(smem->tie_buffer, num_ties, eff_above, output, smem); } else { - tie_handle_large(smem->tie_buffer, num_ties, num_above, output, + tie_handle_large(smem->tie_buffer, num_ties, eff_above, output, smem); } } diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py index a1e27d2e50a3..969a1d9a0b0b 100644 --- a/tests/kernels/test_top_k_per_row.py +++ b/tests/kernels/test_top_k_per_row.py @@ -735,6 +735,35 @@ def test_deepseek_workspace_topk( ) +def _clusters_separating_at_bit_depth(length: int, depth: int) -> torch.Tensor: + """Two clusters of positive floats that first differ at key bit ``depth``. + + All values share the coarse bin (raw bits 0x3F000000..0x3F01FFFF); the + cluster with bit ``depth`` set holds every true top-k member. + """ + per_side = length // 2 + noise = torch.randint(0, 1 << depth, (2 * per_side,), device="cuda") + keys = ( + torch.full((2 * per_side,), 0x3F000000, dtype=torch.int64, device="cuda") + + noise + ) + keys[:per_side] += 1 << depth + keys = keys[torch.randperm(2 * per_side, device="cuda")] + if 2 * per_side < length: + keys = torch.cat( + [ + keys, + torch.full( + (length - 2 * per_side,), + 0x3F000000, + dtype=torch.int64, + device="cuda", + ), + ] + ) + return keys.to(torch.int32).view(torch.float32) + + def run_large_context_topk_test( batch_size: int, seq_lens: list[int], @@ -742,6 +771,7 @@ def run_large_context_topk_test( data_type: str = "random", seed: int = 42, backend: str = "cooperative_topk", + tolerance: float = 1e-4, ) -> None: """ Helper to run a top-k backend test with given parameters. @@ -753,6 +783,8 @@ def run_large_context_topk_test( data_type: Type of test data to generate seed: Random seed for reproducibility backend: Top-k backend to test + tolerance: Value tolerance for the reference comparison; clustered + rows need exact 0.0 """ torch.set_default_device("cuda:0") set_random_seed(seed) @@ -803,6 +835,32 @@ def run_large_context_topk_test( for i, length in enumerate(seq_lens): if length < max_len: logits[i, length:] = float("-inf") + elif data_type == "tight_cluster": + # 1.0 + [0, 1e-3): one coarse bin for the whole row; rows longer + # than the stash overflow the stage-1 threshold bin. + logits = torch.empty(num_rows, max_len, dtype=torch.float32, device="cuda") + for i, length in enumerate(seq_lens): + logits[i, :length] = 1.0 + torch.rand(length, device="cuda") * 1e-3 + if length < max_len: + logits[i, length:] = float("-inf") + elif data_type.startswith("cluster_at_depth_"): + depth = int(data_type.rsplit("_", 1)[-1]) + logits = torch.empty(num_rows, max_len, dtype=torch.float32, device="cuda") + for i, length in enumerate(seq_lens): + logits[i, :length] = _clusters_separating_at_bit_depth(length, depth) + if length < max_len: + logits[i, length:] = float("-inf") + elif data_type == "tight_cluster_then_random": + # Overflow row first, plain rows after: selection state must not leak + # across rows of one persistent launch. + logits = torch.empty(num_rows, max_len, dtype=torch.float32, device="cuda") + for i, length in enumerate(seq_lens): + if i == 0: + logits[i, :length] = 1.0 + torch.rand(length, device="cuda") * 1e-3 + else: + logits[i, :length] = torch.randn(length, device="cuda") + if length < max_len: + logits[i, length:] = float("-inf") else: raise ValueError(f"Unknown data_type: {data_type}") @@ -863,7 +921,7 @@ def run_large_context_topk_test( max_non_topk = non_topk_vals.max() # Allow small tolerance for floating point errors - assert min_cuda_val >= max_non_topk - 1e-4, ( + assert min_cuda_val >= max_non_topk - tolerance, ( f"Row {i}: CUDA top-k contains values smaller than non-top-k. " f"Min CUDA: {min_cuda_val}, Max non-top-k: {max_non_topk}, " f"Length: {length}, k: {k_i}, CUDA indices: {sorted(cuda_set)[:10]}..., " # noqa: E501 @@ -874,8 +932,8 @@ def run_large_context_topk_test( assert torch.allclose( cuda_vals.sort(descending=True)[0], torch_vals.sort(descending=True)[0], - rtol=1e-4, - atol=1e-4, + rtol=tolerance, + atol=tolerance, ), f"""Row {i}: Top-k values don't match. CUDA: {cuda_vals.sort(descending=True)[0][:10]}, Torch: {torch_vals.sort(descending=True)[0][:10]}""" @@ -1315,3 +1373,77 @@ def test_workspace_topk_padded_stride(top_k: int, backend: str) -> None: f"Row {i}: {backend} with padded stride doesn't match. " f"seq_len={sl}, stride={padded_stride}" ) + + +# Regression guard: stage-1 threshold bins whose population exceeds the +# kernel's shared-memory stash (tied/tightly clustered scores). Comparison +# is exact (tolerance=0); one case per kernel path and terminal-clip mode. +OVERSIZED_BIN_ROW_CASES = [ + pytest.param("tight_cluster", [5000], id="tight_5000_decode2048"), + pytest.param("tight_cluster", [9407], id="tight_9407"), + pytest.param("cluster_at_depth_16", [17802], id="depth16_17802"), + pytest.param("cluster_at_depth_0", [8193], id="depth0_8193_exact_ties"), + pytest.param("all_same", [17802], id="all_same_17802_exact_ties"), + pytest.param( + "tight_cluster_then_random", [17802, 3000], id="state_reuse_after_overflow" + ), +] + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("top_k", [512, 2048]) +@pytest.mark.parametrize("data_type,seq_lens", OVERSIZED_BIN_ROW_CASES) +@torch.inference_mode() +def test_persistent_topk_oversized_threshold_bin( + top_k: int, data_type: str, seq_lens: list[int] +) -> None: + run_large_context_topk_test( + batch_size=len(seq_lens), + seq_lens=seq_lens, + top_k=top_k, + data_type=data_type, + seed=1234, + backend="persistent_topk", + tolerance=0.0, + ) + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA") +@pytest.mark.parametrize("top_k", [512, 2048]) +@pytest.mark.parametrize( + "seq_len", [9000, 40000], ids=["hist4096_short", "radix_16k_stash"] +) +@torch.inference_mode() +def test_persistent_topk_oversized_threshold_bin_large_batch( + top_k: int, seq_len: int +) -> None: + # >32 rows dispatch to FilteredTopK on GPUs with >=128KB smem: 9000 + # hits histogram_4096_topk, 40000 the 256-bin + 16K-stash path. + run_large_context_topk_test( + batch_size=40, + seq_lens=[seq_len] * 40, + top_k=top_k, + data_type="tight_cluster", + seed=1234, + backend="persistent_topk", + tolerance=0.0, + ) + + +@pytest.mark.skipif( + not _has_device_capability(90), reason="cooperative_topk requires SM90+" +) +@pytest.mark.parametrize("top_k", [512, 2048]) +@torch.inference_mode() +def test_cooperative_topk_oversized_threshold_bin(top_k: int) -> None: + # sl <= kHist4096MaxLen routes to the shared histogram_4096_topk + # short-medium path inside the cluster kernel (stride must be %4 for TMA). + run_large_context_topk_test( + batch_size=8, + seq_lens=[9000] * 8, + top_k=top_k, + data_type="tight_cluster", + seed=1234, + backend="cooperative_topk", + tolerance=0.0, + )