ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: flatten paged prefill tile dispatch

Browse Source 
- remove the host-provided max_q_len argument
- dispatch only the ragged prefill tile upper bound
- validate the rebuilt CUDA backend end to end
This commit is contained in:
ViperEkura
2026-08-09 23:12:53 +08:00
parent cd31f1f62f
commit c5fba9c238
10 changed files with 69 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_common.h
+1 -3
View File
@@ -23,7 +23,7 @@ struct AttentionParams {
int q_head;
int kv_head;
int head_dim;
int q_len; // Contiguous mode; paged mode uses qo_indptr.
int q_len; // Per-request in contiguous mode; total_q in paged mode.
int kv_len; // Contiguous mode; paged mode uses kv_indptr.
// Attention behavior
@@ -66,6 +66,4 @@ struct AttentionParams {
AT* __restrict__ o_part;
AT* __restrict__ ml_part;
// Host-provided paged prefill grid bound
int max_q_len;
};
csrc/kernels/attn_dispatchers.cuh
+5 -5
View File
@@ -66,9 +66,9 @@ struct PrefillLauncherMMA {
constexpr int WARPS = 4;
constexpr int BC = (HEAD_DIM <= 128) ? 32 : 16;
using Traits = KernelTraits<HEAD_DIM, BC, WARPS, 2>;
int q_len = KV::host_q_len(p);
dim3 grid((q_len + Traits::BR * WARPS - 1) / (Traits::BR * WARPS),
p.q_head, p.batch);
constexpr int ROWS = Traits::BR * WARPS;
dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head,
KV::kPaged ? 1 : p.batch);
dim3 block(Traits::NUM_THREADS);
attn_prefill_split_q_mma_kernel<Traits, KV, IsCausal, HasMask>
<<<grid, block, 0, stream>>>(p);
@@ -81,8 +81,8 @@ struct PrefillLauncherScalar {
template <int HEAD_DIM, bool IsCausal, bool HasMask>
static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
constexpr int G = (HEAD_DIM == 32) ? 4 : 8, ROWS = 32, P_BC = 32;
int q_len = KV::host_q_len(p);
dim3 grid((q_len + ROWS - 1) / ROWS, p.q_head, p.batch);
dim3 grid(KV::host_q_blocks(p, ROWS), p.q_head,
KV::kPaged ? 1 : p.batch);
dim3 block(G, ROWS);
attn_prefill_split_q_kernel_t<HEAD_DIM, KV, G, ROWS, P_BC, IsCausal, HasMask>
<<<grid, block, 0, stream>>>(p);
csrc/kernels/attn_entry_utils.cuh
+2 -4
View File
@@ -189,7 +189,6 @@ inline void attn_pack_paged_decode_params(
p.kv_indptr = kv_indptr.data_ptr<int>();
p.qo_indptr = nullptr;
p.max_context_len = (int)req_to_token.size(1);
p.max_q_len = 1;
p.causal_offset = (int)causal_offset;
p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
@@ -228,7 +227,6 @@ inline void attn_pack_paged_prefill_params(
torch::Tensor kv_indptr,
torch::Tensor qo_indptr,
c10::optional<torch::Tensor> mask,
int64_t max_q_len,
int64_t causal_offset,
double scale,
AttentionParams<T>& p
@@ -251,6 +249,7 @@ inline void attn_pack_paged_prefill_params(
p.q_head = (int)q.size(1);
p.head_dim = (int)q.size(2);
p.q_len = (int)q.size(0);
p.kv_head = (int)k_cache.size(1);
p.batch = (int)req_pool_indices.size(0);
TORCH_CHECK(k_cache.size(2) == p.head_dim, "k_cache head_dim mismatch");
@@ -273,7 +272,6 @@ inline void attn_pack_paged_prefill_params(
p.kv_indptr = kv_indptr.data_ptr<int>();
p.qo_indptr = qo_indptr.data_ptr<int>();
p.max_context_len = (int)req_to_token.size(1);
p.max_q_len = (int)max_q_len;
p.causal_offset = (int)causal_offset;
p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
@@ -288,7 +286,7 @@ inline void attn_pack_paged_prefill_params(
p.mask_l_stride = 0;
} else if (m.dim() == 4) {
TORCH_CHECK(m.size(1) == 1 || m.size(1) == p.q_head, "mask head mismatch");
TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.max_q_len, "mask q_len mismatch");
TORCH_CHECK(m.size(2) > 0 && m.size(2) <= p.q_len, "mask q_len mismatch");
TORCH_CHECK(m.size(3) <= p.max_context_len, "mask kv_len mismatch");
p.mask_b_stride = (int)m.stride(0);
p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
csrc/kernels/attn_kv_source.cuh
+30 -4
View File
@@ -57,8 +57,16 @@ struct ContigKV {
static constexpr bool kPaged = false;
// host-side length hooks (grid + split computation in the launchers)
HOST_DEV_FORCEINLINE int host_q_len(const AttentionParams<bf16>& p) {
return p.q_len;
HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams<bf16>& p, int rows) {
return (p.q_len + rows - 1) / rows;
}
template <int ROWS>
HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams<bf16>&,
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<bf16>& p) {
return p.kv_len;
@@ -107,8 +115,26 @@ struct ContigKV {
struct PagedKV {
static constexpr bool kPaged = true;
HOST_DEV_FORCEINLINE int host_q_len(const AttentionParams<bf16>& p) {
return p.max_q_len;
HOST_DEV_FORCEINLINE int host_q_blocks(const AttentionParams<bf16>& 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 <int ROWS>
HOST_DEV_FORCEINLINE bool map_q_tile(const AttentionParams<bf16>& 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<bf16>& p) {
return p.max_context_len;
csrc/kernels/attn_paged_prefill.cu
+1 -3
View File
@@ -10,7 +10,6 @@ torch::Tensor attn_paged_prefill(
torch::Tensor kv_indptr,
torch::Tensor qo_indptr,
c10::optional<torch::Tensor> mask,
int64_t max_q_len,
int64_t causal_offset,
double scale
) {
@@ -21,7 +20,7 @@ torch::Tensor attn_paged_prefill(
attn_pack_paged_prefill_params(q, k_cache, v_cache,
req_to_token, req_pool_indices,
kv_indptr, qo_indptr, mask,
max_q_len, causal_offset, scale, p);
causal_offset, scale, p);
auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
p.o_ptr = (bf16*)O.data_ptr();
@@ -41,7 +40,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
py::arg("kv_indptr"),
py::arg("qo_indptr"),
py::arg("mask") = py::none(),
py::arg("max_q_len"),
py::arg("causal_offset") = -1,
py::arg("scale") = 0.0,
"SGLang-style paged prefill: flat KV pool + ragged batch.");
csrc/kernels/attn_prefill_split_q.cuh
+13 -2
View File
@@ -36,9 +36,20 @@ template <int HEAD_DIM, typename KV, int G, int ROWS, int P_BC, bool IsCausal, b
__global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
constexpr int DPT = HEAD_DIM / G;
int q_tile = blockIdx.x;
__shared__ int mapped_batch;
__shared__ int mapped_q_tile;
if (threadIdx.x == 0 && threadIdx.y == 0) {
mapped_batch = -1;
KV::template map_q_tile<ROWS>(
p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile);
}
__syncthreads();
if (mapped_batch < 0)
return;
int q_tile = mapped_q_tile;
int q_head = blockIdx.y;
int batch = blockIdx.z;
int batch = mapped_batch;
int gpos = threadIdx.x; // 0..G-1 (which d-chunk)
int row = threadIdx.y; // 0..ROWS-1
int q_row = q_tile * ROWS + row;
csrc/kernels/attn_prefill_split_q_mma.cuh
+14 -3
View File
@@ -24,9 +24,20 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
const int tid4 = lane & 3; // 0..3
const int q_head = blockIdx.y;
const int batch = blockIdx.z;
__shared__ int mapped_batch;
__shared__ int mapped_q_tile;
if (threadIdx.x == 0) {
mapped_batch = -1;
KV::template map_q_tile<Traits::BR * Traits::WARPS>(
p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile);
}
__syncthreads();
if (mapped_batch < 0)
return;
const int batch = mapped_batch;
const int q_tile = mapped_q_tile;
const int kv_head = q_head / (p.q_head / p.kv_head);
const int qrow0 = (blockIdx.x * Traits::WARPS + warp) * Traits::BR;
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);
@@ -61,7 +72,7 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
// Causal tile-skip bounds (dead code when IsCausal == false)
const int max_kv = qrow0 + Traits::BR - 1 + causal_off;
const int block_max_kv =
blockIdx.x * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
q_tile * Traits::WARPS * Traits::BR + Traits::WARPS * Traits::BR - 1
+ causal_off;
int t_end = tiles - 1;