#include #include #include #include #include #include #include #include #if defined(__HIP_PLATFORM_AMD__) static constexpr unsigned long long kWarpSyncMask = 0xFFFFFFFFFFFFFFFFull; #else #include static constexpr unsigned int kWarpSyncMask = 0xFFFFFFFFu; #endif namespace { // Block top-k selection over a per-(head, batch) row of block scores, run by one // CTA of TopKTrait::kCTASize threads. Picks the `topk` highest-scoring block ids // (k_eff = min(topk, num_blocks)). Three size regimes, chosen by num_blocks: // * <= kSmallThreshold : O(n^2) rank-by-compare (no radix). // * <= kCTASize : 4-pass 8-bit radix, one element per thread in a reg. // * <= kMaxNumBlocks : 4-pass 8-bit radix, kIters elements per thread cached // in registers (row read from global exactly once); // liveness is a uint32_t bitmask, selection is an // in-loop scatter -- nothing is cached in shared memory. // The trivial case num_blocks <= topk (every block selected) is handled by the // kernels below, outside the Trait. struct TopKTrait { static constexpr uint32_t kMaxTopK = 32; static constexpr uint32_t kCTASize = 512; static constexpr uint32_t kNumWarps = kCTASize / device::kWarpThreads; static constexpr uint32_t kMaxNumBlocks = 4096; // block topk static constexpr uint32_t kSmallThreshold = 8 * kNumWarps; static constexpr uint32_t kRadixBits = 8; static constexpr uint32_t kRadixSize = 1 << kRadixBits; static constexpr float kNegInf = -std::numeric_limits::infinity(); struct Smem { uint32_t warp_sum[kNumWarps]; alignas(128) uint32_t counter; alignas(128) uint32_t counter_final; alignas(128) uint32_t threshold_bin; uint32_t equal_count; uint32_t above_count; uint32_t histogram[2][kRadixSize]; // 8 bit radix float small_scores[kSmallThreshold]; // small (O(n^2)) path only }; SGL_DEVICE static void forward( const float* __restrict__ scores, const uint32_t num_blocks, int32_t* __restrict__ topk_out, const uint32_t topk, Smem* smem) { using namespace device; const auto tx = threadIdx.x; __builtin_assume(tx < kCTASize); const auto warp_id = tx / kWarpThreads; const auto lane_id = tx % kWarpThreads; constexpr auto is_greater = [](float x, float y, int32_t delta) { return (x > y) || ((x == y) && delta < 0); // lower block id wins }; constexpr auto warp_inclusive_sum = [](uint32_t lane_id, uint32_t val) { #pragma unroll for (uint32_t offset = 1; offset < 32; offset *= 2) { uint32_t n = __shfl_up_sync(kWarpSyncMask, val, offset, 32); if (lane_id >= offset) val += n; } return val; }; constexpr auto clip_nan = [](float x) { return x != x ? kNegInf : x; }; constexpr auto score_to_key = [](float x) { uint32_t b = __float_as_uint(x); return (b & 0x80000000u) ? ~b : (b | 0x80000000u); }; // Find the radix bin holding the topk_remain-th largest of `total_active` // elements currently counted in `histogram`. Writes threshold_bin (the bin), // above_count (elements strictly above it), equal_count (elements in it). const auto find_threshold = [&](uint32_t* histogram, uint32_t total_active, uint32_t topk_remain) { using namespace device; uint32_t hist_val = 0; uint32_t warp_inc = 0; if (tx < kRadixSize) { hist_val = histogram[tx]; warp_inc = warp_inclusive_sum(lane_id, hist_val); if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc; } __syncthreads(); if (tx < kRadixSize) { const auto inter = warp::reduce_sum(lane_id < warp_id ? smem->warp_sum[lane_id] : 0); const auto prefix = inter + warp_inc; // count in bins [0, tx] const auto above = total_active - prefix; // count in bins ABOVE tx if (above < topk_remain && above + hist_val >= topk_remain) { smem->threshold_bin = tx; smem->above_count = above; smem->equal_count = hist_val; } } __syncthreads(); }; if (num_blocks <= kSmallThreshold) { // O(n^2) compare: each block's rank = #blocks that outrank it; the ones // with rank < topk are selected (rank is its position in topk_out). static_assert(kSmallThreshold <= kCTASize); if (tx < num_blocks) smem->small_scores[tx] = clip_nan(scores[tx]); __syncthreads(); constexpr uint32_t kNumCandidates = kSmallThreshold / kNumWarps; constexpr uint32_t kNumTargets = kSmallThreshold / kWarpThreads; float candidates[kNumCandidates]; float target[kNumTargets]; #pragma unroll for (uint32_t i = 0; i < kNumTargets; ++i) { const auto idx = lane_id + i * kWarpThreads; target[i] = (idx < num_blocks) ? smem->small_scores[idx] : kNegInf; } #pragma unroll for (uint32_t i = 0; i < kNumCandidates; ++i) { const auto idx = warp_id + i * kNumWarps; candidates[i] = (idx < num_blocks) ? smem->small_scores[idx] : kNegInf; } #pragma unroll for (uint32_t i = 0; i < kNumCandidates; ++i) { const int32_t idx = warp_id + i * kNumWarps; if (idx >= static_cast(num_blocks)) break; uint32_t rank = 0; #pragma unroll for (uint32_t j = 0; j < kNumTargets; ++j) { const int32_t delta = lane_id + j * kWarpThreads - idx; // partial rank = how many of this lane's targets outrank the candidate rank += is_greater(target[j], candidates[i], delta); } // full rank = sum of the per-lane partial ranks across the warp rank = warp::reduce_sum(rank); if (rank < topk) topk_out[rank] = idx; } } else if (num_blocks <= kCTASize) { // 4-pass 8-bit radix select, one element per thread held in a register. bool active = tx < num_blocks; const auto value = active ? clip_nan(scores[tx]) : kNegInf; const auto key = score_to_key(value); uint32_t topk_remain = topk; uint32_t write_pos = topk; // sentinel: not selected if (tx < kRadixSize) smem->histogram[0][tx] = 0; if (tx == kRadixSize) smem->counter = smem->counter_final = 0; __syncthreads(); uint32_t total_active = num_blocks; #pragma unroll for (int round = 0; round < 4; round++) { const uint32_t shift = 24 - round * 8; const uint32_t bin = (key >> shift) & 0xFFu; const auto hist_idx = round % 2; const auto histogram = smem->histogram[hist_idx]; if (active) atomicAdd(&histogram[bin], 1); if (round < 3 && tx < kRadixSize) smem->histogram[hist_idx ^ 1][tx] = 0; __syncthreads(); find_threshold(histogram, total_active, topk_remain); const auto threshold_bin = smem->threshold_bin; const auto above_count = smem->above_count; const auto equal_count = smem->equal_count; if (round < 3) total_active = equal_count; topk_remain -= above_count; // scatter: above -> selected now; equal at the last pass -> keep the rest if (active) { if (bin > threshold_bin) { write_pos = atomicAdd(&smem->counter, 1); active = false; } else if (bin < threshold_bin) { active = false; } else if (round == 3) { write_pos = topk - topk_remain + atomicAdd(&smem->counter_final, 1); } // bin == threshold && round < 3: stay active for the next pass } if (round == 3 || topk_remain == 0) break; } if (write_pos < topk) topk_out[write_pos] = tx; } else { // num_blocks in (kCTASize, kMaxNumBlocks]: each thread caches its (up to // kIters) slice of the row in registers -- read from global exactly ONCE -- // then runs the same 4-pass radix select as the single-element path looped // over those slots. Liveness is a uint32_t bitmask (bit i = slot i still in // the running set), so there is no per-element flag array; selection is an // in-loop scatter, so there is no per-element position array. Nothing is // cached in shared memory beyond the histogram. constexpr uint32_t kIters = kMaxNumBlocks / kCTASize; static_assert(kIters <= 32, "active liveness is packed into a uint32_t"); uint32_t key[kIters]; uint32_t active = 0; #pragma unroll for (uint32_t i = 0; i < kIters; ++i) { const uint32_t idx = i * kCTASize + tx; if (idx < num_blocks) { key[i] = score_to_key(clip_nan(scores[idx])); active |= 1u << i; } } if (tx < kRadixSize) smem->histogram[0][tx] = 0; if (tx == kRadixSize) smem->counter = smem->counter_final = 0; __syncthreads(); uint32_t topk_remain = topk; uint32_t total_active = num_blocks; #pragma unroll for (int round = 0; round < 4; ++round) { const uint32_t shift = 24 - round * 8; const auto hb = round & 1; #pragma unroll for (uint32_t i = 0; i < kIters; ++i) if (active & (1u << i)) atomicAdd(&smem->histogram[hb][(key[i] >> shift) & 0xFFu], 1); if (round < 3 && tx < kRadixSize) smem->histogram[hb ^ 1][tx] = 0; __syncthreads(); find_threshold(smem->histogram[hb], total_active, topk_remain); const auto threshold_bin = smem->threshold_bin; const auto above_count = smem->above_count; const auto equal_count = smem->equal_count; if (round < 3) total_active = equal_count; topk_remain -= above_count; #pragma unroll for (uint32_t i = 0; i < kIters; ++i) { if (active & (1u << i)) { const uint32_t bin = (key[i] >> shift) & 0xFFu; if (bin > threshold_bin) { topk_out[atomicAdd(&smem->counter, 1)] = i * kCTASize + tx; active &= ~(1u << i); } else if (bin < threshold_bin) { active &= ~(1u << i); } else if (round == 3) { const auto pos = topk - topk_remain + atomicAdd(&smem->counter_final, 1); if (pos < topk) topk_out[pos] = i * kCTASize + tx; } // bin == threshold && round < 3: slot stays live for the next pass } } if (round == 3 || topk_remain == 0) break; } } } }; // ------------------------------------------------------------------------- // Kernels: one CTA (kCTASize threads) per (head, batch) row. The trivial case // num_blocks <= topk (every block selected) is special-judged here, outside the // Trait; otherwise the Trait selects the top-k block ids. // ------------------------------------------------------------------------- // Block-id output: topk_idx[h, b, 0:k_eff) = selected block ids (front-packed, // unordered), [k_eff:topk) = -1. template __global__ void minimax_decode_topk_block_kernel( const float* __restrict__ score, const SeqLenT* __restrict__ seq_lens, int32_t* __restrict__ topk_idx, int batch, int num_heads, int max_seqblock, int block_size, int topk) { const int b = blockIdx.x; // grid.x = batch const int h = blockIdx.y; // grid.y = num_heads const int tx = threadIdx.x; // seq_lens is from an earlier kernel; prefetch it (and the cheap setup) before // waiting on the score producer so the prologue overlaps its tail (PDL). const int64_t seq_len = static_cast(seq_lens[b]); const int num_blocks_raw = static_cast((seq_len + block_size - 1) / block_size); // Never scan past the materialized score columns. const int num_blocks = num_blocks_raw < max_seqblock ? num_blocks_raw : max_seqblock; int32_t* __restrict__ out = topk_idx + (static_cast(h) * batch + b) * topk; device::PDLWaitPrimary(); if (num_blocks <= topk) { // trivial: identity, -1 padded for (int i = tx; i < topk; i += TopKTrait::kCTASize) out[i] = (i < num_blocks) ? i : -1; return; } const float* __restrict__ row = score + (static_cast(h) * batch + b) * max_seqblock; __shared__ TopKTrait::Smem smem; TopKTrait::forward(row, static_cast(num_blocks), out, static_cast(topk), &smem); } // Page-table output: for each (batch b, kv-head h) pseudo-request emit the // trtllm/fa3 page table -- selected blocks sorted ascending (so the final partial // block's pages land last), each expanded to its ppb = block_size/page_size pages // via req_to_token -- plus the effective KV length seq_lens_out. // // DP attention (num_kv_heads > 1): each kv head selects its OWN blocks, so the // per-request page table can't be shared across heads. We flatten (b, h) into // num_heads*batch pseudo-requests laid out batch-major (row = b*num_heads + h, // matching q.view(bs, nkv, gqa, d).reshape(bs*nkv, gqa, d)). seq_lens / slot_ids / // req_to_token are per-batch (head-independent: a token's cache slot is the same // for every head). The page index is head-encoded (head-minor) as // base_page*num_heads + h, which is exactly the page index into an HND cache // [num_pages, nkv, page_size, D] reshaped to [num_pages*nkv, 1, page_size, D] (a // free view when the cache is contiguous HND). num_heads == 1 (h == 0) reproduces // the single-kv-head TP>=4 behavior (page index == base_page). template __global__ void minimax_decode_topk_page_table_kernel( const float* __restrict__ score, const SeqLenT* __restrict__ seq_lens, const int32_t* __restrict__ req_to_token, const int64_t* __restrict__ slot_ids, int32_t* __restrict__ page_table, int32_t* __restrict__ seq_lens_out, int batch, int num_heads, int max_seqblock, int block_size, int topk, int page_size, int r2t_stride, int max_kv_len, int max_sparse_pages) { const int b = blockIdx.x; // grid.x = batch const int h = blockIdx.y; // grid.y = num_heads (kv head) const int tx = threadIdx.x; // Prefetch seq_lens / slot_ids (from earlier kernels) and the cheap setup // before waiting on the score producer, so the prologue overlaps its tail (PDL). const int64_t seq_len = static_cast(seq_lens[b]); const int num_blocks_raw = static_cast((seq_len + block_size - 1) / block_size); const int num_blocks = num_blocks_raw < max_seqblock ? num_blocks_raw : max_seqblock; const int ppb = block_size / page_size; const int64_t out_row = static_cast(b) * num_heads + h; // flattened pseudo-request int32_t* __restrict__ pt_row = page_table + out_row * max_sparse_pages; const int64_t r2t_base = static_cast(slot_ids[b]) * r2t_stride; device::PDLWaitPrimary(); if (num_blocks <= topk) { // trivial: every block selected, all tokens valid if (tx == 0) seq_lens_out[out_row] = static_cast(seq_len); // block id == ascending slot, so the partial final block's pages land last const int total = num_blocks * ppb; for (int e = tx; e < total; e += TopKTrait::kCTASize) { const int slot = e / ppb; const int pp = e % ppb; int tok = slot * block_size + pp * page_size; if (tok >= max_kv_len) tok = max_kv_len - 1; pt_row[e] = req_to_token[r2t_base + tok] / page_size * num_heads + h; } return; } const int k_eff = topk; // num_blocks > topk const float* __restrict__ row = score + (static_cast(h) * batch + b) * max_seqblock; // head-major score __shared__ TopKTrait::Smem smem; __shared__ int32_t s_topk[TopKTrait::kMaxTopK]; TopKTrait::forward(row, static_cast(num_blocks), s_topk, static_cast(topk), &smem); __syncthreads(); // s_topk fully written before the transform reads it // Sort the selected block ids ascending (k_eff <= kMaxTopK is tiny) so the // partial final block lands last, accumulating the effective KV length in the // same pass: each selected block contributes min(block_size, seq_len - c*block) // valid tokens (only the final block can be partial). __shared__ int32_t s_sorted[TopKTrait::kMaxTopK]; __shared__ int s_eff_kv; if (tx == 0) s_eff_kv = 0; __syncthreads(); for (int slot = tx; slot < k_eff; slot += TopKTrait::kCTASize) { const int32_t v = s_topk[slot]; int rank = 0; for (int j = 0; j < k_eff; ++j) rank += (s_topk[j] < v); s_sorted[rank] = v; const int rem = static_cast(seq_len - static_cast(v) * block_size); atomicAdd(&s_eff_kv, rem < block_size ? rem : block_size); } __syncthreads(); if (tx == 0) seq_lens_out[out_row] = s_eff_kv; // Parallel page emit: one thread per output page. const int total = k_eff * ppb; for (int e = tx; e < total; e += TopKTrait::kCTASize) { const int slot = e / ppb; const int pp = e % ppb; int tok = s_sorted[slot] * block_size + pp * page_size; if (tok >= max_kv_len) tok = max_kv_len - 1; pt_row[e] = req_to_token[r2t_base + tok] / page_size * num_heads + h; } } // ------------------------------------------------------------------------- // Launchers // ------------------------------------------------------------------------- template void minimax_decode_topk( tvm::ffi::TensorView score, // [H, B, S] fp32 tvm::ffi::TensorView seq_lens, // [B] int32/int64 tvm::ffi::TensorView topk_idx, // [H, B, T] int32 int64_t block_size, int64_t topk) { using namespace host; SymbolicSize H = {"num_heads"}; SymbolicSize B = {"batch"}; SymbolicSize S = {"max_seqblock"}; SymbolicSize T = {"topk"}; SymbolicDevice device_; device_.set_options(); TensorMatcher({H, B, S}).with_dtype().with_device(device_).verify(score); TensorMatcher({B}).with_dtype().with_device(device_).verify(seq_lens); TensorMatcher({H, B, T}).with_dtype().with_device(device_).verify(topk_idx); const int num_heads = static_cast(H.unwrap()); const int batch = static_cast(B.unwrap()); const int max_seqblock = static_cast(S.unwrap()); const int topk_i = static_cast(T.unwrap()); const DLDevice device = device_.unwrap(); RuntimeCheck( static_cast(topk_i) == topk, "minimax_decode_topk: topk arg (", topk, ") must match topk_idx last dim (", topk_i, ")"); RuntimeCheck(block_size > 0, "block_size must be > 0, got ", block_size); if (batch == 0 || num_heads == 0) return; const dim3 grid(static_cast(batch), static_cast(num_heads)); LaunchKernel(grid, TopKTrait::kCTASize, device, 0) .enable_pdl(kUsePDL)( minimax_decode_topk_block_kernel, static_cast(score.data_ptr()), static_cast(seq_lens.data_ptr()), static_cast(topk_idx.data_ptr()), batch, num_heads, max_seqblock, static_cast(block_size), topk_i); } // Page-table variant: emit the per-(batch, kv-head) paged page table consumed by // the dense backend (trtllm_mha / fa3) plus the effective KV length, instead of // block ids. For DP attention (num_kv_heads > 1) each kv head selects its own // blocks, so (b, h) pseudo-requests are flattened batch-major into the output // (B*num_heads rows); num_heads == 1 is the TP>=4 single-kv-head case. The page // index is head-encoded (head-minor) as base_page*num_heads + h -- the index into // an HND cache [num_pages, nkv, ps, D] reshaped to [num_pages*nkv, 1, ps, D]. // page_table and seq_lens_out are allocated by the caller. template void minimax_decode_topk_page_table( tvm::ffi::TensorView score, // [H, B, S] fp32 (H = num_kv_heads) tvm::ffi::TensorView seq_lens, // [B] int32/int64 tvm::ffi::TensorView req_to_token, // [max_reqs, max_kv_len] int32 tvm::ffi::TensorView slot_ids, // [B] int64 (req_pool_indices) tvm::ffi::TensorView page_table, // [B*H, max_sparse_pages] int32 (out) tvm::ffi::TensorView seq_lens_out, // [B*H] int32 (effective KV length, out) int64_t block_size, int64_t topk, int64_t page_size) { using namespace host; SymbolicSize H = {"num_heads"}; SymbolicSize B = {"batch"}; SymbolicSize S = {"max_seqblock"}; SymbolicSize R = {"max_reqs"}; SymbolicSize KV = {"max_kv_len"}; SymbolicSize BH = {"batch_heads"}; SymbolicSize P = {"max_sparse_pages"}; SymbolicDevice device_; device_.set_options(); TensorMatcher({H, B, S}).with_dtype().with_device(device_).verify(score); TensorMatcher({B}).with_dtype().with_device(device_).verify(seq_lens); TensorMatcher({R, KV}).with_dtype().with_device(device_).verify(req_to_token); TensorMatcher({B}).with_dtype().with_device(device_).verify(slot_ids); TensorMatcher({BH, P}).with_dtype().with_device(device_).verify(page_table); TensorMatcher({BH}).with_dtype().with_device(device_).verify(seq_lens_out); const int num_heads = static_cast(H.unwrap()); const int batch = static_cast(B.unwrap()); const int max_seqblock = static_cast(S.unwrap()); const int max_kv_len = static_cast(KV.unwrap()); const int max_sparse_pages = static_cast(P.unwrap()); const int r2t_stride = static_cast(req_to_token.stride(0)); const DLDevice device = device_.unwrap(); RuntimeCheck( BH.unwrap() == static_cast(batch) * num_heads, "page_table rows (", BH.unwrap(), ") must equal batch*num_heads (", static_cast(batch) * num_heads, ")"); RuntimeCheck( block_size > 0 && page_size > 0 && block_size % page_size == 0, "block_size must be a positive multiple of page_size"); RuntimeCheck(topk <= static_cast(TopKTrait::kMaxTopK), "topk exceeds kMaxTopK for page-table mode"); if (batch == 0 || num_heads == 0) return; const dim3 grid(static_cast(batch), static_cast(num_heads)); LaunchKernel(grid, TopKTrait::kCTASize, device, 0) .enable_pdl(kUsePDL)( minimax_decode_topk_page_table_kernel, static_cast(score.data_ptr()), static_cast(seq_lens.data_ptr()), static_cast(req_to_token.data_ptr()), static_cast(slot_ids.data_ptr()), static_cast(page_table.data_ptr()), static_cast(seq_lens_out.data_ptr()), batch, num_heads, max_seqblock, static_cast(block_size), static_cast(topk), static_cast(page_size), r2t_stride, max_kv_len, max_sparse_pages); } } // namespace