refactor: stride-based attn interface with layout and causal mask
- Replace is_causal + causal_offset with unified causal_offset (-1 = off, >=0 = first Q pos)
- Causal and mask can now coexist (was mutually exclusive)
- Add stride-based addressing for Q/KV/O (layout-agnostic, zero-copy)
- Add layout param ("bhld"/"blhd") parsed in Python, passed as int to C++
- Support 2D [batch, kv_len] and 3D [batch, q_len, kv_len] mask
- Vectorize paged KV gather in Python fallback (was per-token Python loop)
- Extract shared helpers: compute_num_splits, alloc_split_partials, dispatch_head_dim
- Unify paged_decode entry via attn_pack_paged_params
- Update mma_softmax_tile for 3D mask with per-row qrow indexing
This commit is contained in:
@@ -3,26 +3,13 @@
|
||||
#include "attn_paged_decode_split_kv_mma.cuh"
|
||||
#endif
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
|
||||
static int paged_decode_num_splits(int base_blocks, int tiles_total) {
|
||||
int sm_count = 0;
|
||||
cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
|
||||
int n = (2 * sm_count + base_blocks - 1) / base_blocks;
|
||||
return std::max(1, std::min(n, std::min(tiles_total, 32)));
|
||||
}
|
||||
#include "attn_entry_utils.cuh"
|
||||
|
||||
static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
|
||||
int group_size = p.q_head / p.kv_head;
|
||||
int chunks_total = (p.kv_len + PDC_CHUNK - 1) / PDC_CHUNK;
|
||||
p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, chunks_total);
|
||||
|
||||
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
|
||||
auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
|
||||
p.o_part = o_part.data_ptr<float>();
|
||||
p.ml_part = ml_part.data_ptr<float>();
|
||||
p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
|
||||
alloc_split_partials(p);
|
||||
|
||||
size_t smem = PDC_CHUNK * p.head_dim * sizeof(bf16);
|
||||
dim3 grid = dim3(p.batch * p.kv_head, 1, p.num_splits);
|
||||
@@ -35,13 +22,8 @@ static void launch_paged_scalar_decode(PagedAttentionParams<bf16>& p) {
|
||||
template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
|
||||
static void launch_paged_mma_decode(PagedAttentionParams<bf16>& p) {
|
||||
int tiles_total = (p.kv_len + BC - 1) / BC;
|
||||
p.num_splits = paged_decode_num_splits(p.batch * p.kv_head, tiles_total);
|
||||
|
||||
auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
|
||||
auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
|
||||
auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
|
||||
p.o_part = o_part.data_ptr<float>();
|
||||
p.ml_part = ml_part.data_ptr<float>();
|
||||
p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
|
||||
alloc_split_partials(p);
|
||||
|
||||
paged_attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
|
||||
paged_attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
|
||||
@@ -68,68 +50,19 @@ torch::Tensor attn_paged_decode(
|
||||
int64_t page_size,
|
||||
int64_t kv_len,
|
||||
c10::optional<torch::Tensor> mask,
|
||||
bool is_causal = false,
|
||||
int64_t causal_offset = 0,
|
||||
c10::optional<double> scale = c10::nullopt
|
||||
int64_t causal_offset,
|
||||
double scale,
|
||||
int64_t layout
|
||||
) {
|
||||
const at::cuda::OptionalCUDAGuard device_guard(device_of(q));
|
||||
PagedAttentionParams<bf16> p;
|
||||
attn_pack_paged_params(q, page_table, k_cache, v_cache,
|
||||
page_size, kv_len, mask, causal_offset, scale, layout, p);
|
||||
|
||||
int batch = q.size(0);
|
||||
int q_head = q.size(1);
|
||||
int head_dim = q.size(3);
|
||||
int kv_head = k_cache.size(2);
|
||||
int max_pages = page_table.size(1);
|
||||
|
||||
TORCH_CHECK(q.is_cuda() && page_table.is_cuda() && k_cache.is_cuda() && v_cache.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");
|
||||
TORCH_CHECK(page_table.dtype() == torch::kLong, "page_table must be int64");
|
||||
TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)");
|
||||
TORCH_CHECK(head_dim % 32 == 0, "head_dim must be multiple of 32");
|
||||
TORCH_CHECK(k_cache.size(1) == page_size,
|
||||
"k_cache dim 1 must equal page_size, got ",
|
||||
k_cache.size(1), " vs ", page_size);
|
||||
TORCH_CHECK(k_cache.size(0) >= 0, "k_cache must have at least 0 pages");
|
||||
|
||||
float scale_val = scale.has_value()
|
||||
? static_cast<float>(scale.value())
|
||||
: 1.0f / std::sqrt(static_cast<float>(head_dim));
|
||||
|
||||
auto O = torch::empty_like(q);
|
||||
|
||||
PagedAttentionParams<bf16, float> p;
|
||||
p.batch = batch;
|
||||
p.q_head = q_head;
|
||||
p.kv_head = kv_head;
|
||||
p.q_len = static_cast<int>(q.size(2));
|
||||
p.kv_len = static_cast<int>(kv_len);
|
||||
p.head_dim = head_dim;
|
||||
p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0;
|
||||
p.is_causal = is_causal ? 1 : 0;
|
||||
p.causal_offset = static_cast<int>(causal_offset);
|
||||
p.scale = scale_val;
|
||||
p.page_size = static_cast<int>(page_size);
|
||||
p.max_pages = max_pages;
|
||||
p.page_table = page_table.data_ptr<int64_t>();
|
||||
p.k_cache = reinterpret_cast<const bf16*>(k_cache.data_ptr());
|
||||
p.v_cache = reinterpret_cast<const bf16*>(v_cache.data_ptr());
|
||||
p.q = reinterpret_cast<const bf16*>(q.data_ptr());
|
||||
p.mask = p.use_mask ? mask.value().data_ptr<bool>() : nullptr;
|
||||
p.o = reinterpret_cast<bf16*>(O.data_ptr());
|
||||
p.o_part = nullptr;
|
||||
p.ml_part = nullptr;
|
||||
|
||||
switch (p.head_dim) {
|
||||
case 32: dispatch_paged_decode<32>(p); break;
|
||||
case 64: dispatch_paged_decode<64>(p); break;
|
||||
case 128: dispatch_paged_decode<128>(p); break;
|
||||
case 256: dispatch_paged_decode<256>(p); break;
|
||||
default:
|
||||
TORCH_CHECK(false, "paged_decode: unsupported head_dim ", p.head_dim,
|
||||
" (supported: 32, 64, 128, 256)");
|
||||
}
|
||||
auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
|
||||
auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
|
||||
p.o = (bf16*)O_view.data_ptr();
|
||||
|
||||
dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_paged_decode<D>(p); });
|
||||
return O;
|
||||
}
|
||||
|
||||
@@ -142,8 +75,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
py::arg("page_size"),
|
||||
py::arg("kv_len"),
|
||||
py::arg("mask") = py::none(),
|
||||
py::arg("is_causal") = false,
|
||||
py::arg("causal_offset") = 0,
|
||||
py::arg("scale") = py::none(),
|
||||
py::arg("causal_offset") = -1,
|
||||
py::arg("scale") = 0.0,
|
||||
py::arg("layout") = 0,
|
||||
"Paged GQA decode — split-KV with direct page-table access.");
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user