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

perf: precompute ragged Q tile scheduling

Browse Source
This commit is contained in:
ViperEkura
2026-08-16 23:32:46 +08:00
parent 0d0dc64884
commit 1bcd8f53ab
17 changed files with 283 additions and 141 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/attn_common.h
+4 -1
View File
@@ -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<typename T, typename AT = float>
@@ -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
csrc/kernels/attn_decode_split_kv.cuh
+1 -1
View File
@@ -2,7 +2,7 @@
#include <cuda_bf16.h>
#include <float.h>
#include "attn_common.h"
#include "attn_kv_source.cuh"
#include "attn_layout_policies.cuh"
#include "attn_warp_utils.cuh"
constexpr int DC_CHUNK = 64;
csrc/kernels/attn_decode_split_kv_mma.cuh
+1 -1
View File
@@ -2,7 +2,7 @@
#include <cfloat>
#include <cuda_bf16.h>
#include "attn_common.h"
#include "attn_kv_source.cuh"
#include "attn_layout_policies.cuh"
#include "attn_mma_utils.cuh"
#include "attn_warp_utils.cuh"
csrc/kernels/attn_dispatchers.cuh
+20 -15
View File
@@ -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 <cuda_runtime.h>
#include <algorithm>
#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 <typename KV>
template <typename QSchedule, typename KV>
struct PrefillLauncherMMA {
template <int HEAD_DIM, bool IsCausal, bool HasMask>
static void launch(AttentionParams<bf16>& p, cudaStream_t stream) {
using Config = PrefillConfigMap<HEAD_DIM, IsCausal>;
using Traits = KernelTraits<HEAD_DIM, Config::BC, Config::WARPS, Config::STAGES>;
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<Traits, KV, IsCausal, HasMask>
attn_prefill_split_q_mma_kernel<Traits, QSchedule, KV, IsCausal, HasMask>
<<<grid, block, 0, stream>>>(p);
}
};
#endif
template <typename KV>
template <typename QSchedule, typename KV>
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;
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<HEAD_DIM, KV, G, ROWS, P_BC, IsCausal, HasMask>
attn_prefill_split_q_kernel_t<HEAD_DIM, QSchedule, KV, G, ROWS, P_BC,
IsCausal, HasMask>
<<<grid, block, 0, stream>>>(p);
}
};
@@ -115,12 +116,14 @@ static inline void dispatch_prefill(AttentionParams<bf16>& p, cudaStream_t strea
bool has_mask = (p.use_mask && p.mask);
#ifndef ASTRAI_NO_MMA
using Launcher = PrefillLauncherMMA<DenseQSchedule, ContigKV>;
DISPATCH_CAUSAL_MASK(is_causal, has_mask,
PrefillLauncherMMA<ContigKV>::template launch,
Launcher::template launch,
HEAD_DIM, p, stream);
#else
using Launcher = PrefillLauncherScalar<DenseQSchedule, ContigKV>;
DISPATCH_CAUSAL_MASK(is_causal, has_mask,
PrefillLauncherScalar<ContigKV>::template launch,
Launcher::template launch,
HEAD_DIM, p, stream);
#endif
}
@@ -131,12 +134,14 @@ static inline void dispatch_paged_prefill(AttentionParams<bf16>& p, cudaStream_t
bool has_mask = (p.use_mask && p.mask);
#ifndef ASTRAI_NO_MMA
using Launcher = PrefillLauncherMMA<PackedQSchedule, PagedKV>;
DISPATCH_CAUSAL_MASK(is_causal, has_mask,
PrefillLauncherMMA<PagedKV>::template launch,
Launcher::template launch,
HEAD_DIM, p, stream);
#else
using Launcher = PrefillLauncherScalar<PackedQSchedule, PagedKV>;
DISPATCH_CAUSAL_MASK(is_causal, has_mask,
PrefillLauncherScalar<PagedKV>::template launch,
Launcher::template launch,
HEAD_DIM, p, stream);
#endif
}
csrc/kernels/attn_entry_utils.cuh
+14
View File
@@ -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<torch::Tensor> 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<int>();
p.kv_indptr = kv_indptr.data_ptr<int>();
p.qo_indptr = qo_indptr.data_ptr<int>();
p.q_tile_to_batch = q_tile_to_batch.data_ptr<int>();
p.q_tile_to_index = q_tile_to_index.data_ptr<int>();
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;
csrc/kernels/attn_kv_source.cuh → csrc/kernels/attn_layout_policies.cuh
+89 -101
View File
@@ -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<bf16>& p, int rows) {
return (p.q_len + rows - 1) / rows;
}
HOST_FORCEINLINE int host_grid_batch(
const AttentionParams<bf16>& p) {
return p.batch;
}
DEVICE_FORCEINLINE void map_block(
const AttentionParams<bf16>&, int& batch, int& q_tile) {
batch = blockIdx.z;
q_tile = blockIdx.x;
}
DEVICE_FORCEINLINE int q_len(
const AttentionParams<bf16>& p, int) {
return p.q_len;
}
DEVICE_FORCEINLINE int q_base(
const AttentionParams<bf16>& 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<bf16>& p, int) {
return p.num_q_tiles;
}
HOST_FORCEINLINE int host_grid_batch(
const AttentionParams<bf16>&) {
return 1;
}
DEVICE_FORCEINLINE void map_block(
const AttentionParams<bf16>& 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<bf16>& p, int batch) {
return p.qo_indptr[batch + 1] - p.qo_indptr[batch];
}
DEVICE_FORCEINLINE int q_base(
const AttentionParams<bf16>& 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<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) {
HOST_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& 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<bf16>& 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<bf16>& 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<bf16>& p, int batch) {
DEVICE_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int) {
return p.kv_len;
}
HOST_DEV_FORCEINLINE int q_len(const AttentionParams<bf16>& p, int batch) {
return p.q_len;
}
HOST_DEV_FORCEINLINE int causal_offset(const AttentionParams<bf16>& p, int batch) {
DEVICE_FORCEINLINE int causal_offset(
const AttentionParams<bf16>& 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<bf16>& p, int batch) {
DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int) {
return (p.kv_len < p.causal_offset + 1) ? p.kv_len : (p.causal_offset + 1);
}
template <int HEAD_DIM>
HOST_DEV_FORCEINLINE KVContext make_ctx(
DEVICE_FORCEINLINE KVContext make_ctx(
const AttentionParams<bf16>& 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<bf16>& 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<bf16>& 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<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) {
HOST_FORCEINLINE int host_kv_len(const AttentionParams<bf16>& 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<bf16>& 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<bf16>& 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<bf16>& p, int batch) {
DEVICE_FORCEINLINE int kv_len(const AttentionParams<bf16>& p, int batch) {
return p.kv_indptr[batch + 1] - p.kv_indptr[batch];
}
HOST_DEV_FORCEINLINE int q_len(const AttentionParams<bf16>& p, int batch) {
return p.qo_indptr[batch + 1] - p.qo_indptr[batch];
}
HOST_DEV_FORCEINLINE int causal_offset(const AttentionParams<bf16>& p, int batch) {
return kv_len(p, batch) - q_len(p, batch);
DEVICE_FORCEINLINE int causal_offset(
const AttentionParams<bf16>& 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<bf16>& p, int batch) {
DEVICE_FORCEINLINE int decode_attend_len(const AttentionParams<bf16>& p, int batch) {
return kv_len(p, batch);
}
template <int HEAD_DIM>
HOST_DEV_FORCEINLINE KVContext make_ctx(
DEVICE_FORCEINLINE KVContext make_ctx(
const AttentionParams<bf16>& 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<bf16>& 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<bf16>& 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 <int ROWS, typename KV>
__device__ __forceinline__ bool map_q_block(
const AttentionParams<bf16>& 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<ROWS>(
p, blockIdx.x, blockIdx.z, mapped_batch, mapped_q_tile);
}
__syncthreads();
batch = mapped_batch;
q_tile = mapped_q_tile;
return batch >= 0;
}
}
csrc/kernels/attn_paged_prefill.cu
+7 -2
View File
@@ -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<torch::Tensor> mask,
int64_t causal_offset,
double scale
@@ -18,8 +20,9 @@ torch::Tensor attn_paged_prefill(
AttentionParams<bf16> 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,
csrc/kernels/attn_prefill_split_q.cuh
+7 -7
View File
@@ -2,7 +2,7 @@
#include <cfloat>
#include <cuda_bf16.h>
#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 <int HEAD_DIM, typename KV, int G, int ROWS, int P_BC, bool IsCausal, bool HasMask>
template <int HEAD_DIM, typename QSchedule, typename KV, int G, int ROWS, int P_BC,
bool IsCausal, bool HasMask>
__global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
constexpr int DPT = HEAD_DIM / G;
int batch, q_tile;
if (!map_q_block<ROWS, KV>(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<bf16> 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<HEAD_DIM>(p, batch, kv_head);
@@ -56,7 +56,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> 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;
csrc/kernels/attn_prefill_split_q_mma.cuh
+7 -8
View File
@@ -2,7 +2,7 @@
#include <cfloat>
#include <cuda_bf16.h>
#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<HEAD_DIM, BC, WARPS=4, STAGES=2>.
template <typename Traits, typename KV, bool IsCausal, bool HasMask>
template <typename Traits, typename QSchedule, typename KV, bool IsCausal, bool HasMask>
__global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> 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<bf16> p) {
const int q_head = blockIdx.y;
int batch, q_tile;
if (!map_q_block<Traits::BR * Traits::WARPS, KV>(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<Traits::HEAD_DIM>(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<bf16> 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<bf16> 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;