From 1bcd8f53ab4be142b1bd52c99a2fa732b82aba99 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sun, 16 Aug 2026 23:32:46 +0800 Subject: [PATCH] perf: precompute ragged Q tile scheduling --- astrai/extension/backend/attention.py | 2 + astrai/extension/ops/attention.py | 6 + astrai/inference/cache/buffer.py | 2 + astrai/inference/cache/pool.py | 20 +- astrai/inference/workspace.py | 8 + csrc/kernels/attn_common.h | 5 +- csrc/kernels/attn_decode_split_kv.cuh | 2 +- csrc/kernels/attn_decode_split_kv_mma.cuh | 2 +- csrc/kernels/attn_dispatchers.cuh | 35 ++-- csrc/kernels/attn_entry_utils.cuh | 14 ++ ...kv_source.cuh => attn_layout_policies.cuh} | 190 ++++++++---------- csrc/kernels/attn_paged_prefill.cu | 9 +- csrc/kernels/attn_prefill_split_q.cuh | 14 +- csrc/kernels/attn_prefill_split_q_mma.cuh | 15 +- csrc/tests/attn_paged_test.cu | 40 ++++ docs/developer/cuda_kernels.md | 50 ++++- tests/inference/test_cache.py | 10 + 17 files changed, 283 insertions(+), 141 deletions(-) rename csrc/kernels/{attn_kv_source.cuh => attn_layout_policies.cuh} (55%) diff --git a/astrai/extension/backend/attention.py b/astrai/extension/backend/attention.py index 7921bd0..bba16d3 100644 --- a/astrai/extension/backend/attention.py +++ b/astrai/extension/backend/attention.py @@ -610,6 +610,8 @@ class CudaBackend(AttentionBackend): kv_cache.req_pool_indices, kv_cache.kv_indptr, kv_cache.qo_indptr, + kv_cache.q_tile_to_batch, + kv_cache.q_tile_to_index, attn_mask, is_causal=is_causal, ) diff --git a/astrai/extension/ops/attention.py b/astrai/extension/ops/attention.py index 47a8d55..b3345e7 100644 --- a/astrai/extension/ops/attention.py +++ b/astrai/extension/ops/attention.py @@ -150,6 +150,8 @@ def attn_paged_prefill( req_pool_indices: torch.Tensor, kv_indptr: torch.Tensor, qo_indptr: torch.Tensor, + q_tile_to_batch: torch.Tensor, + q_tile_to_index: torch.Tensor, mask: Optional[torch.Tensor] = None, is_causal: bool = False, ) -> torch.Tensor: @@ -167,6 +169,8 @@ def attn_paged_prefill( req_pool_indices: [batch] (int32) kv_indptr: [batch+1] (int32) — prefix sum of per-request kv_lens qo_indptr: [batch+1] (int32) — prefix sum of per-request q_lens + q_tile_to_batch: [num_q_tiles] (int32) — request index per Q tile + q_tile_to_index: [num_q_tiles] (int32) — local Q tile index per request mask: 4D [batch, 1, q_len, kv_len] (bool, True=keep) or None is_causal: apply causal mask @@ -183,6 +187,8 @@ def attn_paged_prefill( req_pool_indices, kv_indptr, qo_indptr, + q_tile_to_batch, + q_tile_to_index, mask, causal_offset=causal_offset, ) diff --git a/astrai/inference/cache/buffer.py b/astrai/inference/cache/buffer.py index 30a8d83..de89892 100644 --- a/astrai/inference/cache/buffer.py +++ b/astrai/inference/cache/buffer.py @@ -99,6 +99,8 @@ class KVCache: max_len: int = 0 kv_indptr: Optional[Tensor] = None qo_indptr: Optional[Tensor] = None + q_tile_to_batch: Optional[Tensor] = None + q_tile_to_index: Optional[Tensor] = None decode_o_part: Optional[Tensor] = None decode_ml_part: Optional[Tensor] = None decode_out: Optional[Tensor] = None diff --git a/astrai/inference/cache/pool.py b/astrai/inference/cache/pool.py index 72ec88b..3f377b8 100644 --- a/astrai/inference/cache/pool.py +++ b/astrai/inference/cache/pool.py @@ -25,7 +25,7 @@ from astrai.inference.cache.strategy import ( RadixCache, TaskCacheState, ) -from astrai.inference.workspace import InferenceWorkspace +from astrai.inference.workspace import Q_TILE_ROWS, InferenceWorkspace # Re-export everything so existing ``from astrai.inference.cache import ...`` # continues to work unchanged after the file split. @@ -217,6 +217,21 @@ class PagePool: torch.tensor(q_lens, dtype=torch.int32, device=device).cumsum(0) ) qo_indptr = workspace.qo_indptr[: b + 1] + tile_batches = [] + tile_indices = [] + for batch, q_len in enumerate(q_lens): + n_tiles = (q_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS + tile_batches.extend([batch] * n_tiles) + tile_indices.extend(range(n_tiles)) + n_tiles = len(tile_batches) + workspace.q_tile_to_batch[:n_tiles].copy_( + torch.tensor(tile_batches, dtype=torch.int32, device=device) + ) + workspace.q_tile_to_index[:n_tiles].copy_( + torch.tensor(tile_indices, dtype=torch.int32, device=device) + ) + q_tile_to_batch = workspace.q_tile_to_batch[:n_tiles] + q_tile_to_index = workspace.q_tile_to_index[:n_tiles] decode_o_part = decode_ml_part = decode_out = None else: # ---- decode: out_cache_loc is a single column (last position) ---- @@ -226,6 +241,7 @@ class PagePool: out_cache_loc = ocl_buf[:b].reshape(-1) workspace.qo_indptr[: b + 1].copy_(inc_buf[: b + 1]) qo_indptr = workspace.qo_indptr[: b + 1] + q_tile_to_batch = q_tile_to_index = None decode_o_part = getattr(workspace, "decode_o_part", None) decode_ml_part = getattr(workspace, "decode_ml_part", None) decode_out = getattr(workspace, "decode_out", None) @@ -240,6 +256,8 @@ class PagePool: max_len=max(seq_lens), kv_indptr=kv_indptr, qo_indptr=qo_indptr, + q_tile_to_batch=q_tile_to_batch, + q_tile_to_index=q_tile_to_index, decode_o_part=decode_o_part, decode_ml_part=decode_ml_part, decode_out=decode_out, diff --git a/astrai/inference/workspace.py b/astrai/inference/workspace.py index ca2ea0d..eaf4287 100644 --- a/astrai/inference/workspace.py +++ b/astrai/inference/workspace.py @@ -10,6 +10,7 @@ import torch from torch import Tensor _MAX_SPLITS = 32 +Q_TILE_ROWS = 64 class InferenceWorkspace: @@ -83,6 +84,13 @@ class InferenceWorkspace: self.qo_indptr = torch.empty( (max_batch_size + 1,), dtype=torch.int32, device=device ) + max_q_tiles = max_batch_size * ((max_seq_len + Q_TILE_ROWS - 1) // Q_TILE_ROWS) + self.q_tile_to_batch = torch.empty( + (max_q_tiles,), dtype=torch.int32, device=device + ) + self.q_tile_to_index = torch.empty( + (max_q_tiles,), dtype=torch.int32, device=device + ) self.inc = torch.arange(max_batch_size + 1, dtype=torch.int32, device=device) self.out_cache_loc = torch.empty( (max_batch_size, 1), dtype=torch.int32, device=device diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attn_common.h index c29a4b3..1929679 100644 --- a/csrc/kernels/attn_common.h +++ b/csrc/kernels/attn_common.h @@ -13,7 +13,7 @@ enum TensorLayout : int { // - Contiguous K/V: dense [batch, kv_head, kv_len, head_dim] tensors (k/v). // - Paged (SGLang-style): flat pool [size, kv_head, head_dim] + req_to_token. // Each kernel selects the addressing via a KVSource policy (see -// attn_kv_source.cuh); a given call only touches the fields of one mode, so +// attn_layout_policies.cuh); a given call only touches the fields of one mode, so // this is a POD shared by both paths rather than two parallel structs that // drift out of sync. template @@ -59,6 +59,9 @@ struct AttentionParams { const int* __restrict__ req_pool_indices; // [batch] const int* __restrict__ kv_indptr; // [batch + 1] const int* __restrict__ qo_indptr; // [batch + 1] or nullptr for decode + const int* __restrict__ q_tile_to_batch; // [num_q_tiles], prefill only + const int* __restrict__ q_tile_to_index; // [num_q_tiles], prefill only + int num_q_tiles; int max_context_len; // req_to_token stride (dim 1) // Decode split-KV workspace diff --git a/csrc/kernels/attn_decode_split_kv.cuh b/csrc/kernels/attn_decode_split_kv.cuh index caee3bf..a840b6e 100644 --- a/csrc/kernels/attn_decode_split_kv.cuh +++ b/csrc/kernels/attn_decode_split_kv.cuh @@ -2,7 +2,7 @@ #include #include #include "attn_common.h" -#include "attn_kv_source.cuh" +#include "attn_layout_policies.cuh" #include "attn_warp_utils.cuh" constexpr int DC_CHUNK = 64; diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index 203659a..7718fa0 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -2,7 +2,7 @@ #include #include #include "attn_common.h" -#include "attn_kv_source.cuh" +#include "attn_layout_policies.cuh" #include "attn_mma_utils.cuh" #include "attn_warp_utils.cuh" diff --git a/csrc/kernels/attn_dispatchers.cuh b/csrc/kernels/attn_dispatchers.cuh index 949f4f2..84ad175 100644 --- a/csrc/kernels/attn_dispatchers.cuh +++ b/csrc/kernels/attn_dispatchers.cuh @@ -3,7 +3,7 @@ // No torch dependency; pure CUDA. // // The paged and contiguous kernels are unified by the KVSource policy -// (ContigKV / PagedKV from attn_kv_source.cuh), so each launcher struct +// (ContigKV / PagedKV from attn_layout_policies.cuh), so each launcher struct // below is templated on KV and the paged dispatch is just the same launcher // instantiated with PagedKV. Only the grid/split math differs, and that is // covered by KV::host_q_len / KV::host_kv_len. @@ -11,7 +11,7 @@ #include #include #include "attn_warp_utils.cuh" -#include "attn_kv_source.cuh" +#include "attn_layout_policies.cuh" #include "attn_prefill_split_q.cuh" #include "attn_decode_split_kv.cuh" #ifndef ASTRAI_NO_MMA @@ -80,31 +80,32 @@ template <> struct PrefillConfigMap<128, true> : PrefillKernelConfig<32> {}; template <> struct PrefillConfigMap<256, false> : PrefillKernelConfig<16> {}; template <> struct PrefillConfigMap<256, true> : PrefillKernelConfig<16> {}; -template +template struct PrefillLauncherMMA { template static void launch(AttentionParams& p, cudaStream_t stream) { using Config = PrefillConfigMap; using Traits = KernelTraits; constexpr int ROWS = Traits::BR * Config::WARPS; - dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head, - KV::kPaged ? 1 : p.batch); + dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head, + QSchedule::host_grid_batch(p)); dim3 block(Traits::NUM_THREADS); - attn_prefill_split_q_mma_kernel + attn_prefill_split_q_mma_kernel <<>>(p); } }; #endif -template +template struct PrefillLauncherScalar { template static void launch(AttentionParams& p, cudaStream_t stream) { - constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 32, P_BC = 32; - dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head, - KV::kPaged ? 1 : p.batch); + constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 64, P_BC = 32; + dim3 grid(QSchedule::host_q_blocks(p, ROWS), p.q_head, + QSchedule::host_grid_batch(p)); dim3 block(G, ROWS); - attn_prefill_split_q_kernel_t + attn_prefill_split_q_kernel_t <<>>(p); } }; @@ -115,12 +116,14 @@ static inline void dispatch_prefill(AttentionParams& p, cudaStream_t strea bool has_mask = (p.use_mask && p.mask); #ifndef ASTRAI_NO_MMA + using Launcher = PrefillLauncherMMA; DISPATCH_CAUSAL_MASK(is_causal, has_mask, - PrefillLauncherMMA::template launch, + Launcher::template launch, HEAD_DIM, p, stream); #else + using Launcher = PrefillLauncherScalar; DISPATCH_CAUSAL_MASK(is_causal, has_mask, - PrefillLauncherScalar::template launch, + Launcher::template launch, HEAD_DIM, p, stream); #endif } @@ -131,12 +134,14 @@ static inline void dispatch_paged_prefill(AttentionParams& p, cudaStream_t bool has_mask = (p.use_mask && p.mask); #ifndef ASTRAI_NO_MMA + using Launcher = PrefillLauncherMMA; DISPATCH_CAUSAL_MASK(is_causal, has_mask, - PrefillLauncherMMA::template launch, + Launcher::template launch, HEAD_DIM, p, stream); #else + using Launcher = PrefillLauncherScalar; DISPATCH_CAUSAL_MASK(is_causal, has_mask, - PrefillLauncherScalar::template launch, + Launcher::template launch, HEAD_DIM, p, stream); #endif } diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 0e56604..5ec9efc 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -227,6 +227,8 @@ inline void attn_pack_paged_prefill_params( torch::Tensor req_pool_indices, torch::Tensor kv_indptr, torch::Tensor qo_indptr, + torch::Tensor q_tile_to_batch, + torch::Tensor q_tile_to_index, c10::optional mask, int64_t causal_offset, double scale, @@ -237,6 +239,7 @@ inline void attn_pack_paged_prefill_params( TORCH_CHECK(q.is_cuda() && k_cache.is_cuda() && v_cache.is_cuda()); TORCH_CHECK(req_to_token.is_cuda() && req_pool_indices.is_cuda()); TORCH_CHECK(kv_indptr.is_cuda() && qo_indptr.is_cuda()); + TORCH_CHECK(q_tile_to_batch.is_cuda() && q_tile_to_index.is_cuda()); TORCH_CHECK(q.dtype() == torch::kBFloat16, "q must be bf16"); TORCH_CHECK(k_cache.dtype() == torch::kBFloat16, "k_cache must be bf16"); TORCH_CHECK(v_cache.dtype() == torch::kBFloat16, "v_cache must be bf16"); @@ -245,6 +248,10 @@ inline void attn_pack_paged_prefill_params( "req_pool_indices must be int32"); TORCH_CHECK(kv_indptr.dtype() == torch::kInt32, "kv_indptr must be int32"); TORCH_CHECK(qo_indptr.dtype() == torch::kInt32, "qo_indptr must be int32"); + TORCH_CHECK(q_tile_to_batch.dtype() == torch::kInt32, + "q_tile_to_batch must be int32"); + TORCH_CHECK(q_tile_to_index.dtype() == torch::kInt32, + "q_tile_to_index must be int32"); TORCH_CHECK(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must match"); TORCH_CHECK(k_cache.dim() == 3, "k_cache must be 3D [size, kv_head, head_dim]"); TORCH_CHECK(q.dim() == 3, "q must be 3D [total_q, q_head, head_dim]"); @@ -261,6 +268,10 @@ inline void attn_pack_paged_prefill_params( TORCH_CHECK(p.q_head % p.kv_head == 0, "q_head must be divisible by kv_head"); TORCH_CHECK(kv_indptr.size(0) == p.batch + 1, "kv_indptr must be [batch+1]"); TORCH_CHECK(qo_indptr.size(0) == p.batch + 1, "qo_indptr must be [batch+1]"); + TORCH_CHECK(q_tile_to_batch.dim() == 1 && q_tile_to_index.dim() == 1, + "Q tile mappings must be 1D"); + TORCH_CHECK(q_tile_to_batch.size(0) == q_tile_to_index.size(0), + "Q tile mappings must have equal length"); p.q_l_stride = (int)q.stride(0); p.q_h_stride = (int)q.stride(1); @@ -273,6 +284,9 @@ inline void attn_pack_paged_prefill_params( p.req_pool_indices = req_pool_indices.data_ptr(); p.kv_indptr = kv_indptr.data_ptr(); p.qo_indptr = qo_indptr.data_ptr(); + p.q_tile_to_batch = q_tile_to_batch.data_ptr(); + p.q_tile_to_index = q_tile_to_index.data_ptr(); + p.num_q_tiles = (int)q_tile_to_batch.size(0); p.max_context_len = (int)req_to_token.size(1); p.causal_offset = (int)causal_offset; diff --git a/csrc/kernels/attn_kv_source.cuh b/csrc/kernels/attn_layout_policies.cuh similarity index 55% rename from csrc/kernels/attn_kv_source.cuh rename to csrc/kernels/attn_layout_policies.cuh index 678def9..99646e2 100644 --- a/csrc/kernels/attn_kv_source.cuh +++ b/csrc/kernels/attn_layout_policies.cuh @@ -3,13 +3,10 @@ #include "attn_common.h" // ============================================================================ -// KVSource policies — the single dimension along which the paged and -// non-paged attention kernels differ. Each kernel is templated on one of -// these (ContigKV / PagedKV) and stays fully generic: the policy owns every -// place where "where does K/V live" and "what is this request's seq_len" -// are answered. All methods are __host__ __device__ so the same policy -// serves both the device kernels (addressing, seq_len) and the host-side -// launchers (grid / split computation). +// Attention layout policies keep Q scheduling independent from K/V storage. +// DenseQSchedule / PackedQSchedule map blocks to Q tiles; ContigKV / PagedKV +// resolve logical K/V positions to physical addresses. This lets the shared +// kernels compose Q layout and K/V storage without coupling the two concerns. // // ContigKV: K/V are dense [batch, kv_head, kv_len, head_dim] tensors. // Params fields used: k, v, kv_stride_*, kv_len, q_len, @@ -25,11 +22,76 @@ // (e.g. the req_pool_indices global read) element-by-element. // ============================================================================ -// Every policy method is static + callable from both host and device code. +#define HOST_FORCEINLINE static __host__ __forceinline__ +#define DEVICE_FORCEINLINE static __device__ __forceinline__ #define HOST_DEV_FORCEINLINE static __host__ __device__ __forceinline__ using bf16 = __nv_bfloat16; +// ============================================================================ +// Q scheduling policies +// +// Map CUDA blocks to request-local Q tiles independently of K/V storage. +// Dense tensors encode the request in blockIdx.z; packed ragged tensors use +// a compact precomputed work map indexed by blockIdx.x. +// ============================================================================ + +struct DenseQSchedule { + HOST_FORCEINLINE int host_q_blocks( + const AttentionParams& p, int rows) { + return (p.q_len + rows - 1) / rows; + } + + HOST_FORCEINLINE int host_grid_batch( + const AttentionParams& p) { + return p.batch; + } + + DEVICE_FORCEINLINE void map_block( + const AttentionParams&, int& batch, int& q_tile) { + batch = blockIdx.z; + q_tile = blockIdx.x; + } + + DEVICE_FORCEINLINE int q_len( + const AttentionParams& p, int) { + return p.q_len; + } + + DEVICE_FORCEINLINE int q_base( + const AttentionParams& p, int batch, int q_head) { + return batch * p.q_b_stride + q_head * p.q_h_stride; + } +}; + +struct PackedQSchedule { + HOST_FORCEINLINE int host_q_blocks( + const AttentionParams& p, int) { + return p.num_q_tiles; + } + + HOST_FORCEINLINE int host_grid_batch( + const AttentionParams&) { + return 1; + } + + DEVICE_FORCEINLINE void map_block( + const AttentionParams& p, int& batch, int& q_tile) { + batch = p.q_tile_to_batch[blockIdx.x]; + q_tile = p.q_tile_to_index[blockIdx.x]; + } + + DEVICE_FORCEINLINE int q_len( + const AttentionParams& p, int batch) { + return p.qo_indptr[batch + 1] - p.qo_indptr[batch]; + } + + DEVICE_FORCEINLINE int q_base( + const AttentionParams& p, int batch, int q_head) { + return p.qo_indptr[batch] * p.q_l_stride + q_head * p.q_h_stride; + } +}; + // Hoisted per-(batch, kv_head) addressing context. struct KVContext { int kv_base; // contig: batch*kv_b_stride + kv_head*kv_h_stride @@ -56,59 +118,40 @@ struct KVAddr { struct ContigKV { static constexpr bool kPaged = false; - // host-side length hooks (grid + split computation in the launchers) - HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams& p, int rows) { - return (p.q_len + rows - 1) / rows; - } - template - HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams&, - int flat_tile, int grid_batch, - int& batch, int& q_tile) { - batch = grid_batch; - q_tile = flat_tile; - return true; - } - HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams& p) { + HOST_FORCEINLINE int host_kv_len(const AttentionParams& p) { return p.kv_len; } - // prefill: element offset of the request's Q rows (kernel adds qrow*q_l_stride) - HOST_DEV_FORCEINLINE int q_base( - const AttentionParams& p, int batch, int q_head) { - return batch * p.q_b_stride + q_head * p.q_h_stride; - } // decode: same offset (q_len == 1, so there is no row stride component) - HOST_DEV_FORCEINLINE int q_decode_base( + DEVICE_FORCEINLINE int q_decode_base( const AttentionParams& p, int batch, int q_head) { return batch * p.q_b_stride + q_head * p.q_h_stride; } - HOST_DEV_FORCEINLINE int kv_len(const AttentionParams& p, int batch) { + DEVICE_FORCEINLINE int kv_len(const AttentionParams& p, int) { return p.kv_len; } - HOST_DEV_FORCEINLINE int q_len(const AttentionParams& p, int batch) { - return p.q_len; - } - HOST_DEV_FORCEINLINE int causal_offset(const AttentionParams& p, int batch) { + DEVICE_FORCEINLINE int causal_offset( + const AttentionParams& p, int, int) { return p.causal_offset; } // decode: exclusive bound of the single query's attend range - HOST_DEV_FORCEINLINE int decode_attend_len(const AttentionParams& p, int batch) { + DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams& p, int) { return (p.kv_len < p.causal_offset + 1) ? p.kv_len : (p.causal_offset + 1); } template - HOST_DEV_FORCEINLINE KVContext make_ctx( + DEVICE_FORCEINLINE KVContext make_ctx( const AttentionParams& p, int batch, int kv_head) { KVContext c = {}; c.kv_base = batch * p.kv_b_stride + kv_head * p.kv_h_stride; return c; } - HOST_DEV_FORCEINLINE int resolve_token( + DEVICE_FORCEINLINE int resolve_token( const AttentionParams& p, const KVContext& c, int kc, bool valid) { return valid ? kc : -1; } - HOST_DEV_FORCEINLINE KVAddr kv_addr_from_token( + DEVICE_FORCEINLINE KVAddr kv_addr_from_token( const AttentionParams& p, const KVContext& c, int token, int d) { const bool valid = token >= 0; const int safe_token = valid ? token : 0; @@ -123,58 +166,30 @@ struct ContigKV { struct PagedKV { static constexpr bool kPaged = true; - HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams& p, int rows) { - // sum(ceil(q_len[b] / rows)) <= ceil(total_q / rows) + batch - 1. - return (p.q_len + rows - 1) / rows + p.batch - 1; - } - template - HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams& p, - int flat_tile, int, - int& batch, int& q_tile) { - int tile_base = 0; - for (int b = 0; b < p.batch; ++b) { - int len = p.qo_indptr[b + 1] - p.qo_indptr[b]; - int tiles = (len + ROWS - 1) / ROWS; - if (flat_tile < tile_base + tiles) { - batch = b; - q_tile = flat_tile - tile_base; - return true; - } - tile_base += tiles; - } - return false; - } - HOST_DEV_FORCEINLINE int host_kv_len(const AttentionParams& p) { + HOST_FORCEINLINE int host_kv_len(const AttentionParams& p) { return p.max_context_len; } - // prefill: Q rows start at qo_indptr[batch] (ragged batch base) - HOST_DEV_FORCEINLINE int q_base( - const AttentionParams& p, int batch, int q_head) { - return p.qo_indptr[batch] * p.q_l_stride + q_head * p.q_h_stride; - } // decode: Q is [batch, q_head, head_dim], so batch is the outer row - HOST_DEV_FORCEINLINE int q_decode_base( + DEVICE_FORCEINLINE int q_decode_base( const AttentionParams& p, int batch, int q_head) { return batch * p.q_l_stride + q_head * p.q_h_stride; } - HOST_DEV_FORCEINLINE int kv_len(const AttentionParams& p, int batch) { + DEVICE_FORCEINLINE int kv_len(const AttentionParams& p, int batch) { return p.kv_indptr[batch + 1] - p.kv_indptr[batch]; } - HOST_DEV_FORCEINLINE int q_len(const AttentionParams& p, int batch) { - return p.qo_indptr[batch + 1] - p.qo_indptr[batch]; - } - HOST_DEV_FORCEINLINE int causal_offset(const AttentionParams& p, int batch) { - return kv_len(p, batch) - q_len(p, batch); + DEVICE_FORCEINLINE int causal_offset( + const AttentionParams& p, int batch, int q_len) { + return kv_len(p, batch) - q_len; } // decode: the query is the last token, so [0, seq_len) IS its causal range - HOST_DEV_FORCEINLINE int decode_attend_len(const AttentionParams& p, int batch) { + DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams& p, int batch) { return kv_len(p, batch); } template - HOST_DEV_FORCEINLINE KVContext make_ctx( + DEVICE_FORCEINLINE KVContext make_ctx( const AttentionParams& p, int batch, int kv_head) { KVContext c = {}; c.req_idx = p.req_pool_indices[batch]; @@ -183,11 +198,11 @@ struct PagedKV { c.head_off = (int64_t)kv_head * HEAD_DIM; return c; } - HOST_DEV_FORCEINLINE int resolve_token( + DEVICE_FORCEINLINE int resolve_token( const AttentionParams& p, const KVContext& c, int kc, bool valid) { return valid ? p.req_to_token[c.req_idx * c.rtt_stride + kc] : -1; } - HOST_DEV_FORCEINLINE KVAddr kv_addr_from_token( + DEVICE_FORCEINLINE KVAddr kv_addr_from_token( const AttentionParams& p, const KVContext& c, int slot, int d) { const bool valid = slot >= 0; const int safe_slot = valid ? slot : 0; @@ -195,30 +210,3 @@ struct PagedKV { return {&p.k_ptr[gmem_off], &p.v_ptr[gmem_off], valid}; } }; - -// ---- Q-block mapping ---- -// Contiguous grids map directly to (batch, q_tile). Paged grids flatten the -// ragged Q tiles, so one thread resolves the request and broadcasts it. -template -__device__ __forceinline__ bool map_q_block( - const AttentionParams& p, int& batch, int& q_tile) { - if constexpr (!KV::kPaged) { - batch = blockIdx.z; - q_tile = blockIdx.x; - return true; - } else { - __shared__ int mapped_batch; - __shared__ int mapped_q_tile; - - if ((threadIdx.x | threadIdx.y) == 0) { - mapped_batch = -1; - KV::template map_q_tile( - p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile); - } - __syncthreads(); - - batch = mapped_batch; - q_tile = mapped_q_tile; - return batch >= 0; - } -} diff --git a/csrc/kernels/attn_paged_prefill.cu b/csrc/kernels/attn_paged_prefill.cu index 99d40d7..b30f89b 100644 --- a/csrc/kernels/attn_paged_prefill.cu +++ b/csrc/kernels/attn_paged_prefill.cu @@ -9,6 +9,8 @@ torch::Tensor attn_paged_prefill( torch::Tensor req_pool_indices, torch::Tensor kv_indptr, torch::Tensor qo_indptr, + torch::Tensor q_tile_to_batch, + torch::Tensor q_tile_to_index, c10::optional mask, int64_t causal_offset, double scale @@ -18,8 +20,9 @@ torch::Tensor attn_paged_prefill( AttentionParams p; attn_pack_paged_prefill_params(q, k_cache, v_cache, - req_to_token, req_pool_indices, - kv_indptr, qo_indptr, mask, + req_to_token, req_pool_indices, + kv_indptr, qo_indptr, + q_tile_to_batch, q_tile_to_index, mask, causal_offset, scale, p); auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options()); @@ -39,6 +42,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("req_pool_indices"), py::arg("kv_indptr"), py::arg("qo_indptr"), + py::arg("q_tile_to_batch"), + py::arg("q_tile_to_index"), py::arg("mask") = py::none(), py::arg("causal_offset") = -1, py::arg("scale") = 0.0, diff --git a/csrc/kernels/attn_prefill_split_q.cuh b/csrc/kernels/attn_prefill_split_q.cuh index 9de9f8f..3157512 100644 --- a/csrc/kernels/attn_prefill_split_q.cuh +++ b/csrc/kernels/attn_prefill_split_q.cuh @@ -2,7 +2,7 @@ #include #include #include "attn_common.h" -#include "attn_kv_source.cuh" +#include "attn_layout_policies.cuh" using bf16 = __nv_bfloat16; @@ -32,13 +32,13 @@ __device__ __forceinline__ void ld8(const bf16* p, float* o) { } } -template +template __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { constexpr int DPT = HEAD_DIM / G; int batch, q_tile; - if (!map_q_block(p, batch, q_tile)) - return; + QSchedule::map_block(p, batch, q_tile); int q_head = blockIdx.y; int gpos = threadIdx.x; // 0..G-1 (which d-chunk) @@ -47,8 +47,8 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr). const int seq_len = KV::kv_len(p, batch); - const int q_len = KV::q_len(p, batch); - const int causal_off = KV::causal_offset(p, batch); + const int q_len = QSchedule::q_len(p, batch); + const int causal_off = KV::causal_offset(p, batch, q_len); const int kv_head = q_head / (p.q_head / p.kv_head); const KVContext kctx = KV::template make_ctx(p, batch, kv_head); @@ -56,7 +56,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM]; // Q: stride-based load [batch, q_head, q_len, head_dim] - const int q_base = KV::q_base(p, batch, q_head); + const int q_base = QSchedule::q_base(p, batch, q_head); float qreg[DPT]; if (q_row < q_len) { int q_off = q_base + q_row * p.q_l_stride + gpos * DPT * p.q_d_stride; diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index e27097d..c3c5122 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -2,7 +2,7 @@ #include #include #include "attn_common.h" -#include "attn_kv_source.cuh" +#include "attn_layout_policies.cuh" #include "attn_mma_utils.cuh" // Tensor-core prefill flash attention (raw mma.sync PTX), unified across @@ -16,7 +16,7 @@ // dead branches in the inner compute loop (FA2-style). // // Traits = KernelTraits. -template +template __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { const int warp = threadIdx.x / 32; const int lane = threadIdx.x % 32; @@ -25,15 +25,14 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { const int q_head = blockIdx.y; int batch, q_tile; - if (!map_q_block(p, batch, q_tile)) - return; + QSchedule::map_block(p, batch, q_tile); const int kv_head = q_head / (p.q_head / p.kv_head); const int qrow0 = (q_tile * Traits::WARPS + warp) * Traits::BR; // Per-request dims (from KV policy — paged reads kv_indptr/qo_indptr). const int seq_len = KV::kv_len(p, batch); - const int q_len = KV::q_len(p, batch); - const int causal_off = KV::causal_offset(p, batch); + const int q_len = QSchedule::q_len(p, batch); + const int causal_off = KV::causal_offset(p, batch, q_len); const KVContext kctx = KV::template make_ctx(p, batch, kv_head); // Static shared memory: double-buffered K/V (no sQ — Q goes direct @@ -42,7 +41,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { __shared__ __align__(16) bf16 sV[Traits::STAGES * Traits::BC * Traits::LD]; // Load Q fragments straight from global into mma A-operand layout. - const int q_base = KV::q_base(p, batch, q_head); + const int q_base = QSchedule::q_base(p, batch, q_head); const int qra = qrow0 + gid; const int qrb = qrow0 + gid + 8; const bool va = qra < q_len, vb = qrb < q_len; @@ -138,7 +137,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { // ---- write output: packed bf16x2 stores ---- float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f; float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f; - const int o_base = KV::q_base(p, batch, q_head); + const int o_base = QSchedule::q_base(p, batch, q_head); #pragma unroll for (int dn8 = 0; dn8 < Traits::DN8; dn8++) { int d = dn8 * 8 + 2 * tid4; diff --git a/csrc/tests/attn_paged_test.cu b/csrc/tests/attn_paged_test.cu index 57cb659..25c7f7c 100644 --- a/csrc/tests/attn_paged_test.cu +++ b/csrc/tests/attn_paged_test.cu @@ -12,6 +12,26 @@ struct PagedDecodeDispatch { AttentionParams& p; template void operator()() { dispatch_paged_decode(p, 0); } }; struct PagedPrefillDispatch { AttentionParams& p; template void operator()() { dispatch_paged_prefill(p, 0); } }; +static int make_q_tile_mapping(const std::vector& q_lens, + int** d_batch, int** d_tile) { + constexpr int ROWS = 64; + std::vector h_batch; + std::vector h_tile; + for (int b = 0; b < (int)q_lens.size(); ++b) { + int n_tiles = (q_lens[b] + ROWS - 1) / ROWS; + for (int tile = 0; tile < n_tiles; ++tile) { + h_batch.push_back(b); + h_tile.push_back(tile); + } + } + size_t bytes = h_batch.size() * sizeof(int); + cudaMalloc(d_batch, bytes); + cudaMalloc(d_tile, bytes); + cudaMemcpy(*d_batch, h_batch.data(), bytes, cudaMemcpyHostToDevice); + cudaMemcpy(*d_tile, h_tile.data(), bytes, cudaMemcpyHostToDevice); + return (int)h_batch.size(); +} + // ---- CPU reference: paged decode with variable seq_lens ---- // Q: [B, Hq, D], K/V pool: [pool_size, Hkv, D] // req_to_token: [num_reqs, max_ctx_len], req_pool_indices: [B] @@ -483,6 +503,9 @@ static int run_prefill_test(int B, int Hq, int Hkv, nullptr, 0, 0, B, Hq, Hkv, HEAD_DIM, max_ctx, causal, h_o_ref); + int *d_qtb, *d_qti; + int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti); + // Kernel launch AttentionParams p; p.batch = B; p.q_head = Hq; p.kv_head = Hkv; @@ -497,6 +520,8 @@ static int run_prefill_test(int B, int Hq, int Hkv, p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool; p.req_to_token = d_rtt; p.req_pool_indices = d_rpi; p.kv_indptr = d_kvi; p.qo_indptr = d_qoi; + p.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti; + p.num_q_tiles = num_q_tiles; p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr; dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p}); @@ -523,6 +548,7 @@ static int run_prefill_test(int B, int Hq, int Hkv, free(h_o_ref); free(h_o_bf); free(h_o_got); cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi); + cudaFree(d_qtb); cudaFree(d_qti); return pass ? 0 : 1; } @@ -619,6 +645,10 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) { h_mask, q_len, q_len, B, Hq, Hkv, HEAD_DIM, max_ctx, 0, h_o_ref); + std::vector q_lens(B, q_len); + int *d_qtb, *d_qti; + int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti); + AttentionParams p; p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.head_dim = HEAD_DIM; @@ -632,6 +662,8 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) { p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool; p.req_to_token = d_rtt; p.req_pool_indices = d_rpi; p.kv_indptr = d_kvi; p.qo_indptr = d_qoi; + p.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti; + p.num_q_tiles = num_q_tiles; p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr; dispatch_by_head_dim(HEAD_DIM, PagedPrefillDispatch{p}); @@ -659,6 +691,7 @@ static int run_prefill_mask_test(int Hq, int Hkv, int q_len, int seed) { cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi); cudaFree(d_mask); + cudaFree(d_qtb); cudaFree(d_qti); return pass ? 0 : 1; } @@ -786,6 +819,10 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau for (int b = 0; b < B; b++) h_qoi[b + 1] = h_qoi[b] + q_len; cudaMemcpy(d_qoi, h_qoi, sz_qoi, cudaMemcpyHostToDevice); + std::vector q_lens(B, q_len); + int *d_qtb, *d_qti; + int num_q_tiles = make_q_tile_mapping(q_lens, &d_qtb, &d_qti); + AttentionParams p; p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.head_dim = HEAD_DIM; @@ -798,6 +835,8 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau p.q_ptr = d_q; p.k_ptr = d_k_pool; p.v_ptr = d_v_pool; p.req_to_token = d_rtt; p.req_pool_indices = d_rpi; p.kv_indptr = d_kvi; p.qo_indptr = d_qoi; + p.q_tile_to_batch = d_qtb; p.q_tile_to_index = d_qti; + p.num_q_tiles = num_q_tiles; p.o_ptr = d_o; p.o_part = nullptr; p.ml_part = nullptr; auto launch = [&]() { @@ -827,6 +866,7 @@ static void bench_prefill(int B, int Hq, int Hkv, int q_len, int kv_len, int cau free(tmp); free(h_rtt); free(h_rpi); free(h_kvi); free(h_qoi); cudaFree(d_q); cudaFree(d_o); cudaFree(d_k_pool); cudaFree(d_v_pool); cudaFree(d_rtt); cudaFree(d_rpi); cudaFree(d_kvi); cudaFree(d_qoi); + cudaFree(d_qtb); cudaFree(d_qti); } int main() { diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 70047bb..e3833d7 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -19,9 +19,10 @@ Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Ac | Split-KV MMA decode | `attn_decode_split_kv_mma.cuh` | Split KV across warps + MMA (sm_80+) | | Split-Q MMA prefill | `attn_prefill_split_q_mma.cuh` | Split Q across warps + MMA (sm_80+) | -> The paged and non-paged paths are ONE kernel templated on a `KVSource` -> policy (`ContigKV` / `PagedKV` in `attn_kv_source.cuh`); there are no -> separate `attn_paged_*.cuh` files anymore. +> The paged and non-paged paths share one kernel body. Prefill is templated on +> an independent Q schedule (`DenseQSchedule` / `PackedQSchedule`) and KV +> source (`ContigKV` / `PagedKV`); decode only needs the KV source. There are +> no separate `attn_paged_*.cuh` files. ### Rotary Embedding Kernel @@ -238,6 +239,47 @@ mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep) Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`. +### Q Scheduling and KV Addressing + +Prefill separates Q work scheduling from KV storage: + +- `DenseQSchedule` maps a rectangular grid directly with + `batch = blockIdx.z` and `q_tile = blockIdx.x`. +- `PackedQSchedule` consumes a compact work map for a packed + `[total_q, q_heads, head_dim]` tensor. +- `ContigKV` and `PagedKV` only provide KV lengths and translate logical KV + positions into physical addresses. They do not schedule Q blocks. + +For ragged Q lengths `[70, 10, 130]` and 64 rows per Q tile, cache binding +builds: + +```text +qo_indptr = [0, 70, 80, 210] +q_tile_to_batch = [0, 0, 1, 2, 2, 2] +q_tile_to_index = [0, 1, 0, 0, 1, 2] +``` + +Paged prefill launches: + +```text +grid.x = num_q_tiles # 6, exactly the valid ragged work items +grid.y = q_heads +grid.z = 1 +``` + +Each block resolves its request and request-local tile in O(1): + +```cpp +batch = q_tile_to_batch[blockIdx.x]; +q_tile = q_tile_to_index[blockIdx.x]; +``` + +The kernel then uses `qo_indptr[batch]` for the packed Q base and adjacent +`qo_indptr` / `kv_indptr` entries for that request's Q and KV lengths. This +avoids the previous per-block linear scan over the batch, shared-memory +broadcast, mapping barrier, and upper-bound grid with potentially invalid +blocks. + ## Standalone Testing Each `csrc/tests/*.cu` file has the `nvcc` compile command in its header comment. Example: @@ -285,7 +327,7 @@ csrc/ │ ├── attn_decode_split_kv_mma.cuh # Split-KV + MMA variant (contig + paged) │ ├── attn_prefill_split_q.cuh # Split-Q variant (contig + paged via KVSource) │ ├── attn_prefill_split_q_mma.cuh # Split-Q + MMA variant (contig + paged) -│ ├── attn_kv_source.cuh # KVSource policies (ContigKV / PagedKV) +│ ├── attn_layout_policies.cuh # Q schedules and KVSource policies │ ├── attn_dispatchers.cuh # Kernel dispatch macros + KV-templated launchers │ ├── attn_entry_utils.cuh # Entry point helpers │ ├── attn_mma_utils.cuh # MMA utilities diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py index ed277aa..7435026 100644 --- a/tests/inference/test_cache.py +++ b/tests/inference/test_cache.py @@ -286,6 +286,16 @@ def test_page_pool_contiguous_bind_tasks_prefill(): assert kv.req_pool_indices.dtype == torch.int32 +def test_page_pool_bind_tasks_builds_compact_q_tile_mapping(): + pool = _make_contiguous_pool(max_batch_size=3, max_seq_len=256) + + kv = pool.bind_tasks([0, 1, 2], [70, 10, 130], _ws(pool), start_pos=0) + + assert kv.qo_indptr.tolist() == [0, 70, 80, 210] + assert kv.q_tile_to_batch.tolist() == [0, 0, 1, 2, 2, 2] + assert kv.q_tile_to_index.tolist() == [0, 1, 0, 0, 1, 2] + + def test_page_pool_contiguous_bind_tasks_decode(): pool = _make_contiguous_pool() task_cache = _make_task_cache(pool)