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_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;