From 57729fd92d4d36da7e6f86b312b6313eda7b738b Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Tue, 14 Jul 2026 21:30:55 +0800 Subject: [PATCH] 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 --- astrai/extension/__init__.py | 8 + astrai/extension/ops.py | 161 ++++++++++++--- csrc/kernels/attn_common.h | 22 +- csrc/kernels/attn_decode.cu | 63 ++---- csrc/kernels/attn_decode_split_kv.cuh | 38 +++- csrc/kernels/attn_decode_split_kv_mma.cuh | 26 ++- csrc/kernels/attn_entry_utils.cuh | 195 ++++++++++++++++-- csrc/kernels/attn_mma_utils.cuh | 25 ++- csrc/kernels/attn_paged_decode.cu | 103 ++------- csrc/kernels/attn_paged_decode_split_kv.cuh | 19 +- .../attn_paged_decode_split_kv_mma.cuh | 17 +- csrc/kernels/attn_prefill.cu | 37 +--- csrc/kernels/attn_prefill_split_q.cuh | 32 ++- csrc/kernels/attn_prefill_split_q_mma.cuh | 41 ++-- csrc/tests/attn_decode_test.cu | 7 +- csrc/tests/attn_paged_decode_test.cu | 8 +- csrc/tests/attn_prefill_test.cu | 8 +- csrc/tests/test_utils.cuh | 31 ++- 18 files changed, 558 insertions(+), 283 deletions(-) diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index 573e643..14548b9 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -5,6 +5,14 @@ Public API: - ``attn_prefill`` — multi-query prefill attention - ``attn_paged_decode`` — paged decode attention (direct page-table access) +Interface (shared by all wrappers): + causal_offset: -1 = non-causal; >=0 = absolute position of first Q token + mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True = keep) + scale: 0.0 = auto (1/sqrt(head_dim)); >0 = explicit + layout: "bhld" (default) or "blhd" + +Causal and mask can coexist — both are applied simultaneously. + Each wrapper dispatches to its compiled CUDA kernel (``astrai.extension.attn_*``) when available, otherwise falls back to ``torch.nn.functional.scaled_dot_product_attention``. """ diff --git a/astrai/extension/ops.py b/astrai/extension/ops.py index 0dc4d9c..4d33ede 100644 --- a/astrai/extension/ops.py +++ b/astrai/extension/ops.py @@ -3,15 +3,43 @@ Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when available, otherwise falls back to ``torch`` SDPA. +Interface (all functions): + causal_offset: -1 = non-causal; >=0 = absolute position of first Q token + mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool) + scale: 0.0 = auto (1/sqrt(head_dim)); >0 = explicit + layout: "bhld" (default) or "blhd" + Add new kernel wrappers here; split into per-variant files only if this file grows large. """ +import math + import torch import torch.nn.functional as F from astrai.extension.loader import _available, _modules +_LAYOUT_CODES: dict[str, int] = {"bhld": 0, "blhd": 1} + + +def _parse_layout(layout: str | int) -> int: + if isinstance(layout, int): + return layout + code = _LAYOUT_CODES.get(layout.lower()) + if code is None: + raise ValueError( + f"unknown layout '{layout}', expected one of {list(_LAYOUT_CODES)}" + ) + return code + + +def _to_bhld(t: torch.Tensor, layout: int) -> torch.Tensor: + """Normalize to b h l d view. Zero-copy transpose if layout==1 (b l h d).""" + if layout == 1: + return t.transpose(1, 2) + return t + def _expand_kv_heads( k: torch.Tensor, v: torch.Tensor, q_head: int @@ -26,20 +54,84 @@ def _expand_kv_heads( return k, v +def _build_attn_mask( + q: torch.Tensor, + k: torch.Tensor, + mask: torch.Tensor | None, + causal_offset: int, + scale: float, +) -> tuple[torch.Tensor | None, float]: + """Build SDPA-compatible attn_mask + resolved scale. + + q and k must already be in b h l d layout. + Causal and mask can coexist: causal sets -inf above the diagonal, mask + sets -inf for padded positions. Both are OR'd into a single bool mask. + """ + q_len = q.size(2) + kv_len = k.size(2) + head_dim = q.size(3) + resolved_scale = scale if scale and scale > 0 else 1.0 / math.sqrt(head_dim) + + attn_mask = None + + if mask is not None: + if mask.dim() == 2: + # [batch, kv_len] → [batch, 1, 1, kv_len] + attn_mask = mask[:, None, None, :] + elif mask.dim() == 3: + # [batch, q_len, kv_len] → [batch, 1, q_len, kv_len] + attn_mask = mask[:, None, :, :] + else: + raise ValueError(f"mask must be 2D or 3D, got {mask.dim()}D") + + if causal_offset >= 0: + batch = q.size(0) + # q row i attends to kv cols 0..(causal_offset + i) + q_idx = torch.arange(q_len, device=q.device).unsqueeze(1) # [q_len, 1] + kv_idx = torch.arange(kv_len, device=q.device).unsqueeze(0) # [1, kv_len] + causal_bool = kv_idx > (causal_offset + q_idx) # True = masked out + causal_mask = causal_bool.unsqueeze(0).expand( + batch, -1, -1 + ) # [batch, q_len, kv_len] + causal_mask = causal_mask[:, None, :, :] # [batch, 1, q_len, kv_len] + + if attn_mask is not None: + attn_mask = attn_mask | causal_mask + else: + attn_mask = causal_mask + + return attn_mask, resolved_scale + + def _torch_fallback( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None, - is_causal: bool, - scale: float | None, + causal_offset: int, + scale: float, + q_layout: int, + kv_layout: int | None = None, ) -> torch.Tensor: - """Reference attention via ``scaled_dot_product_attention``.""" + """Reference attention via ``scaled_dot_product_attention``. + + q_layout / kv_layout: 0 = b h l d, 1 = b l h d. + If kv_layout is None, uses q_layout (Q and K/V share the same layout). + """ + if kv_layout is None: + kv_layout = q_layout + q = _to_bhld(q, q_layout) + k = _to_bhld(k, kv_layout) + v = _to_bhld(v, kv_layout) k, v = _expand_kv_heads(k, v, q.size(1)) - attn_mask = mask[:, None, None, :] if mask is not None else None - return F.scaled_dot_product_attention( - q, k, v, attn_mask=attn_mask, is_causal=is_causal and mask is None, scale=scale + attn_mask, resolved_scale = _build_attn_mask(q, k, mask, causal_offset, scale) + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=False, scale=resolved_scale ) + # Restore Q's original layout + if q_layout == 1: + out = out.transpose(1, 2) + return out def _gather_kv_from_pages( @@ -56,23 +148,22 @@ def _gather_kv_from_pages( k_cache : [n_pages, page_size, n_kv_heads, head_dim] v_cache : same as k_cache Returns: - k, v : [batch, n_kv_heads, kv_len, head_dim] + k, v : [batch, kv_len, n_kv_heads, head_dim] (b l h d) """ batch, max_pages = page_table.shape - n_pages, ps, n_kv_heads, head_dim = k_cache.shape + _, ps, n_kv_heads, head_dim = k_cache.shape if ps != page_size: raise ValueError(f"k_cache page_size mismatch: {ps} vs {page_size}") - k = k_cache.new_empty(batch, n_kv_heads, kv_len, head_dim) - v = v_cache.new_empty(batch, n_kv_heads, kv_len, head_dim) + # Vectorized gather: build physical page + offset indices, then advanced-index + positions = torch.arange(kv_len, device=page_table.device) + logical_pages = positions // page_size # [kv_len] + page_offsets = positions % page_size # [kv_len] - for b in range(batch): - for pos in range(kv_len): - log_pg = pos // page_size - pg_off = pos % page_size - phys = int(page_table[b, log_pg].item()) - k[b, :, pos, :] = k_cache[phys, pg_off, :, :] - v[b, :, pos, :] = v_cache[phys, pg_off, :, :] + phys_pages = page_table[:, logical_pages] # [batch, kv_len] + # k_cache[phys_pages, page_offsets] → [batch, kv_len, n_kv_heads, head_dim] (b l h d) + k = k_cache[phys_pages, page_offsets] + v = v_cache[phys_pages, page_offsets] return k, v @@ -81,21 +172,22 @@ def attn_decode( k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None, - is_causal: bool = False, - causal_offset: int = 0, - scale: float | None = None, + causal_offset: int = -1, + scale: float = 0.0, + layout: str = "bhld", ) -> torch.Tensor: + li = _parse_layout(layout) if _available["attn_decode"]: return _modules["attn_decode"].attn_decode( q, k, v, mask=mask, - is_causal=is_causal, causal_offset=causal_offset, scale=scale, + layout=li, ) - return _torch_fallback(q, k, v, mask, is_causal, scale) + return _torch_fallback(q, k, v, mask, causal_offset, scale, q_layout=li) def attn_prefill( @@ -103,21 +195,22 @@ def attn_prefill( k: torch.Tensor, v: torch.Tensor, mask: torch.Tensor | None = None, - is_causal: bool = False, - causal_offset: int = 0, - scale: float | None = None, + causal_offset: int = -1, + scale: float = 0.0, + layout: str = "bhld", ) -> torch.Tensor: + li = _parse_layout(layout) if _available["attn_prefill"]: return _modules["attn_prefill"].attn_prefill( q, k, v, mask=mask, - is_causal=is_causal, causal_offset=causal_offset, scale=scale, + layout=li, ) - return _torch_fallback(q, k, v, mask, is_causal, scale) + return _torch_fallback(q, k, v, mask, causal_offset, scale, q_layout=li) def attn_paged_decode( @@ -128,10 +221,11 @@ def attn_paged_decode( page_size: int, kv_len: int, mask: torch.Tensor | None = None, - is_causal: bool = False, - causal_offset: int = 0, - scale: float | None = None, + causal_offset: int = -1, + scale: float = 0.0, + layout: str = "bhld", ) -> torch.Tensor: + li = _parse_layout(layout) if _available["attn_paged_decode"]: return _modules["attn_paged_decode"].attn_paged_decode( q, @@ -141,9 +235,12 @@ def attn_paged_decode( page_size, kv_len, mask=mask, - is_causal=is_causal, causal_offset=causal_offset, scale=scale, + layout=li, ) + # Gathered K/V are always b l h d k, v = _gather_kv_from_pages(page_table, k_cache, v_cache, page_size, kv_len) - return _torch_fallback(q, k, v, mask, is_causal, scale) + return _torch_fallback( + q, k, v, mask, causal_offset, scale, q_layout=li, kv_layout=1 + ) diff --git a/csrc/kernels/attn_common.h b/csrc/kernels/attn_common.h index f99c38d..b4b9546 100644 --- a/csrc/kernels/attn_common.h +++ b/csrc/kernels/attn_common.h @@ -10,11 +10,19 @@ struct AttentionParams { int kv_len; int head_dim; int use_mask; - int is_causal; - int causal_offset; + int causal_offset; // -1 = non-causal; >=0 = absolute position of first Q token int num_splits; float scale; + // Q strides (element offsets for each dim — layout-agnostic) + int q_stride_b, q_stride_h, q_stride_l, q_stride_d; + // KV strides (K and V share the same layout — only base pointers differ) + int kv_stride_b, kv_stride_h, kv_stride_l, kv_stride_d; + + // Mask: 2D [batch, kv_len] (mask_q_stride=0) or 3D [batch, q_len, kv_len] + int mask_b_stride; // = kv_len (both 2D and 3D) + int mask_q_stride; // 2D: 0 (all q rows share); 3D: kv_len + const T* __restrict__ q; const T* __restrict__ k; const T* __restrict__ v; @@ -34,7 +42,6 @@ struct PagedAttentionParams { int kv_len; int head_dim; int use_mask; - int is_causal; int causal_offset; float scale; @@ -42,12 +49,19 @@ struct PagedAttentionParams { int page_size; int max_pages; + // Q strides (layout-agnostic) + int q_stride_b, q_stride_h, q_stride_l, q_stride_d; + + // Mask strides (2D or 3D) + int mask_b_stride; + int mask_q_stride; + const T* __restrict__ q; const T* __restrict__ k_cache; const T* __restrict__ v_cache; const bool* __restrict__ mask; const int64_t* __restrict__ page_table; - + T* __restrict__ o; AT* __restrict__ o_part; AT* __restrict__ ml_part; diff --git a/csrc/kernels/attn_decode.cu b/csrc/kernels/attn_decode.cu index 862364f..8fd3fb0 100644 --- a/csrc/kernels/attn_decode.cu +++ b/csrc/kernels/attn_decode.cu @@ -5,24 +5,12 @@ #include "attn_decode_split_kv_mma.cuh" #endif -static int 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))); -} - // Scalar fallback: one warp per query head, split-KV across grid.z. static void launch_scalar_decode(AttentionParams& p) { int group_size = p.q_head / p.kv_head; int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK; - p.num_splits = 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(); - p.ml_part = ml_part.data_ptr(); + p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total); + alloc_split_partials(p); size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16); attn_decode_split_kv_kernel<<>>(p); @@ -38,13 +26,8 @@ static void launch_scalar_decode(AttentionParams& p) { template static void launch_mma_decode(AttentionParams& p) { int tiles_total = (p.kv_len + BC - 1) / BC; - p.num_splits = 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(); - p.ml_part = ml_part.data_ptr(); + p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total); + alloc_split_partials(p); attn_decode_split_kv_mma_kernel<<>>(p); attn_decode_combine_kernel<<>>(p); @@ -68,35 +51,21 @@ torch::Tensor attn_decode( torch::Tensor k, torch::Tensor v, c10::optional mask, - bool is_causal = false, - int64_t causal_offset = 0, - c10::optional scale = c10::nullopt + int64_t causal_offset, + double scale, + int64_t layout ) { AttentionParams p; - attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p); + attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p); TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1"); TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32"); - auto O = torch::empty_like(q); - p.o = (bf16*)O.data_ptr(); + // O matches Q's original layout + 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(); - switch (p.head_dim) { - case 32: - dispatch_decode<32>(p); - break; - case 64: - dispatch_decode<64>(p); - break; - case 128: - dispatch_decode<128>(p); - break; - case 256: - dispatch_decode<256>(p); - break; - default: - TORCH_CHECK(false, "decode: unsupported head_dim ", p.head_dim, - " (supported: 32, 64, 128, 256)"); - } + dispatch_head_dim(p.head_dim, [&]() { dispatch_decode(p); }); return O; } @@ -106,8 +75,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("k"), py::arg("v"), 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, "GQA decode (tensor-core head-packing on sm_80+, scalar fallback)"); } diff --git a/csrc/kernels/attn_decode_split_kv.cuh b/csrc/kernels/attn_decode_split_kv.cuh index e790313..b0ed646 100644 --- a/csrc/kernels/attn_decode_split_kv.cuh +++ b/csrc/kernels/attn_decode_split_kv.cuh @@ -21,13 +21,16 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams p) { int lane = threadIdx.x; int hd_per_thread = p.head_dim / 32; + // Q: [batch, q_head, q_len=1, head_dim] — stride-based float q_reg[8]; - int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread; + int q_off = batch * p.q_stride_b + q_head * p.q_stride_h + + lane * hd_per_thread * p.q_stride_d; for (int i = 0; i < hd_per_thread; i++) - q_reg[i] = __bfloat162float(p.q[q_off + i]); + q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]); - int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * p.head_dim; - int mask_base = batch * p.kv_len; + // KV: [batch, kv_head, kv_len, head_dim] — stride-based base + int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h; + int mask_base = batch * p.mask_b_stride; float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f}; @@ -43,9 +46,15 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams p) { int chunk_start = ci * DC_CHUNK; int this_chunk = min(DC_CHUNK, p.kv_len - chunk_start); + // Load K into shared memory (gather from strided global) int total = this_chunk * p.head_dim; - for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) - k_smem[i] = p.k[kv_base + chunk_start * p.head_dim + i]; + for (int i = threadIdx.y * 32 + lane; i < total; i += blockDim.x * blockDim.y) { + int s = i / p.head_dim; + int d_dim = i % p.head_dim; + int kv_idx = chunk_start + s; + int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d; + k_smem[i] = p.k[g_off]; + } __syncthreads(); for (int s = 0; s < this_chunk; s++) { @@ -54,9 +63,10 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams p) { partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]); partial = warp_reduce_sum(partial) * p.scale; - if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s]) + int kv_idx = chunk_start + s; + if (p.use_mask && p.mask && !p.mask[mask_base + kv_idx]) partial = -FLT_MAX; - if (p.is_causal && (chunk_start + s) > p.causal_offset) + if (p.causal_offset >= 0 && kv_idx > p.causal_offset) partial = -FLT_MAX; float new_m = fmaxf(m, partial); @@ -64,9 +74,10 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams p) { float beta = expf(partial - new_m); d = d * alpha + beta; - int v_off = kv_base + (chunk_start + s) * p.head_dim + lane * hd_per_thread; + // V: stride-based read + int v_off = kv_base + kv_idx * p.kv_stride_l + lane * hd_per_thread * p.kv_stride_d; for (int i = 0; i < hd_per_thread; i++) - acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i]) * beta; + acc_reg[i] = acc_reg[i] * alpha + __bfloat162float(p.v[v_off + i * p.kv_stride_d]) * beta; m = new_m; } __syncthreads(); @@ -94,6 +105,9 @@ __global__ void attn_decode_combine_kernel(AttentionParams p) { int d = threadIdx.x; if (d >= p.head_dim) return; + int batch = bh / p.q_head; + int q_head = bh % p.q_head; + size_t split_base = (size_t)bh * p.num_splits; const float* mlp = p.ml_part + split_base * 2; const float* op = p.o_part + split_base * p.head_dim; @@ -112,5 +126,7 @@ __global__ void attn_decode_combine_kernel(AttentionParams p) { } float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f; - p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv); + // Stride-based output write (q_len=1 for decode, so stride_l not needed) + int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d; + p.o[o_off] = __float2bfloat16(acc * inv); } diff --git a/csrc/kernels/attn_decode_split_kv_mma.cuh b/csrc/kernels/attn_decode_split_kv_mma.cuh index c1075bb..07ba795 100644 --- a/csrc/kernels/attn_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_decode_split_kv_mma.cuh @@ -68,7 +68,8 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { // ---- Load Q directly from global into mma A-operand registers ---- // Same layout as prefill: frag[0]/[2] = row gid, frag[1]/[3] = row gid+8 // cols kt*16 + tid4*2 + {0,1} / +{8,9}. pau[0]=cols c,c+1; pau[4]=c+8,c+9. - const int q_base = (batch * p.q_head + q_head0) * HEAD_DIM; + // Stride-based: Q is [batch, q_head, q_len=1, head_dim] + const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h; const int qra = gid; const int qrb = gid + 8; const bool va = qra < G, vb = qrb < G; @@ -77,9 +78,9 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { for (int kt = 0; kt < KD; kt++) { int c = kt * 16 + tid4 * 2; const unsigned* pau = reinterpret_cast( - &p.q[q_base + qra * HEAD_DIM + c]); + &p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]); const unsigned* pbu = reinterpret_cast( - &p.q[q_base + qrb * HEAD_DIM + c]); + &p.q[q_base + qrb * p.q_stride_h + c * p.q_stride_d]); Qa[kt][0] = va ? pau[0] : 0u; Qa[kt][1] = vb ? pbu[0] : 0u; Qa[kt][2] = va ? pau[4] : 0u; @@ -92,8 +93,9 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f; float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f; - const int kv_base = (batch * p.kv_head + kv_head) * p.kv_len * HEAD_DIM; - const int mask_base = batch * p.kv_len; + // KV: stride-based base — [batch, kv_head, kv_len, head_dim] + const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h; + const int mask_batch_base = batch * p.mask_b_stride; const int tiles_total = (p.kv_len + BC - 1) / BC; const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits; const int ti_begin = split * tiles_per_split; @@ -111,8 +113,10 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { int kc = kv0 + r; bool valid = kc < p.kv_len; int off = r * LD + swiz_col(d, r, SWIZ_MASK); - cp_async_16_pred(&dK[off], &p.k[kv_base + kc * HEAD_DIM + d], valid); - cp_async_16_pred(&dV[off], &p.v[kv_base + kc * HEAD_DIM + d], valid); + // KV stride-based: contiguous within head_dim (stride_d == 1 typically) + int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d; + cp_async_16_pred(&dK[off], &p.k[g_off], valid); + cp_async_16_pred(&dV[off], &p.v[g_off], valid); } cp_async_commit(); }; @@ -148,9 +152,13 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) { Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale, Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale; - int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; + // Decode: q_len=1, so qrow0=qrow1=0, mask_q_stride irrelevant + int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; mma_softmax_tile(kv0, maxc, maxc, - mask_base, p.mask, has_mask, + 0, 0, + mask_batch_base, 0, + batch, + p.mask, has_mask, Sacc, Oacc, m0, m1, l0, l1, lane); mma_pv_accumulate(Sacc, bV, LD, SWIZ_MASK, lane, Oacc); diff --git a/csrc/kernels/attn_entry_utils.cuh b/csrc/kernels/attn_entry_utils.cuh index 7e83730..edcc286 100644 --- a/csrc/kernels/attn_entry_utils.cuh +++ b/csrc/kernels/attn_entry_utils.cuh @@ -1,43 +1,214 @@ #pragma once #include +#include #include "attn_common.h" +using bf16 = __nv_bfloat16; + +// --------------------------------------------------------------------------- +// Shared dispatch helpers — eliminates duplication across .cu entry files. +// --------------------------------------------------------------------------- + +inline int compute_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))); +} + +// Dispatch head_dim to a generic lambda FN (callable as FN.operator()()). +template +inline void dispatch_head_dim(int hd, Fn&& fn) { + switch (hd) { + case 32: fn.template operator()<32>(); break; + case 64: fn.template operator()<64>(); break; + case 128: fn.template operator()<128>(); break; + case 256: fn.template operator()<256>(); break; + default: + TORCH_CHECK(false, "unsupported head_dim ", hd, + " (supported: 32, 64, 128, 256)"); + } +} + +// Allocate split-KV partial buffers and wire into params. +template +inline void alloc_split_partials(P& p) { + 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 = (float*)o_part.data_ptr(); + p.ml_part = (float*)ml_part.data_ptr(); +} + +// --------------------------------------------------------------------------- +// Param packing — fills AttentionParams from torch tensors with validation. +// --------------------------------------------------------------------------- + template inline void attn_pack_params( torch::Tensor q, torch::Tensor k, torch::Tensor v, c10::optional mask, - bool is_causal, int64_t causal_offset, - c10::optional scale, + double scale, + int64_t layout, // 0 = b h l d, 1 = b l h d AttentionParams& p ) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda()); TORCH_CHECK(q.dtype() == torch::kBFloat16); TORCH_CHECK(k.dtype() == torch::kBFloat16); TORCH_CHECK(v.dtype() == torch::kBFloat16); + TORCH_CHECK(k.sizes() == v.sizes(), "K and V must have identical shapes"); + TORCH_CHECK(q.dim() == 4 && k.dim() == 4, "Q/K/V must be 4D"); + + // Normalize to b h l d view (zero-copy transpose if user passed b l h d) + if (layout == 1) { + q = q.transpose(1, 2); + k = k.transpose(1, 2); + v = v.transpose(1, 2); + } p.batch = (int)q.size(0); p.q_head = (int)q.size(1); - p.kv_head = (int)k.size(1); p.q_len = (int)q.size(2); - p.kv_len = (int)k.size(2); p.head_dim = (int)q.size(3); - p.use_mask = mask.has_value() ? 1 : 0; - p.is_causal = is_causal ? 1 : 0; + p.kv_head = (int)k.size(1); + p.kv_len = (int)k.size(2); + TORCH_CHECK(k.size(3) == p.head_dim, "K/V head_dim must match Q"); + + // Strides (layout-agnostic: works for b h l d and b l h d) + p.q_stride_b = (int)q.stride(0); + p.q_stride_h = (int)q.stride(1); + p.q_stride_l = (int)q.stride(2); + p.q_stride_d = (int)q.stride(3); + p.kv_stride_b = (int)k.stride(0); + p.kv_stride_h = (int)k.stride(1); + p.kv_stride_l = (int)k.stride(2); + p.kv_stride_d = (int)k.stride(3); + p.causal_offset = (int)causal_offset; - p.scale = scale.has_value() ? (float)scale.value() : 1.0f / sqrtf((float)p.head_dim); + p.use_mask = mask.has_value() ? 1 : 0; + p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim); + p.q = (const T*)q.data_ptr(); p.k = (const T*)k.data_ptr(); p.v = (const T*)v.data_ptr(); + p.o = nullptr; + p.o_part = nullptr; + p.ml_part = nullptr; + if (p.use_mask) { - TORCH_CHECK(mask.value().dtype() == torch::kBool); - TORCH_CHECK(mask.value().dim() == 2); - TORCH_CHECK(mask.value().size(0) == p.batch); - TORCH_CHECK(mask.value().size(1) == p.kv_len); - p.mask = mask.value().data_ptr(); + auto m = mask.value(); + TORCH_CHECK(m.is_cuda(), "mask must be on CUDA"); + TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool"); + TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch"); + TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch"); + if (m.dim() == 2) { + p.mask_b_stride = (int)m.stride(0); + p.mask_q_stride = 0; + } else if (m.dim() == 3) { + TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch"); + p.mask_b_stride = (int)m.stride(0); + p.mask_q_stride = (int)m.stride(1); + } else { + TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]"); + } + p.mask = m.data_ptr(); } else { p.mask = nullptr; + p.mask_b_stride = 0; + p.mask_q_stride = 0; + } +} + +// --------------------------------------------------------------------------- +// Param packing for paged attention. +// --------------------------------------------------------------------------- + +template +inline void attn_pack_paged_params( + torch::Tensor q, + torch::Tensor page_table, + torch::Tensor k_cache, + torch::Tensor v_cache, + int64_t page_size, + int64_t kv_len, + c10::optional mask, + int64_t causal_offset, + double scale, + int64_t layout, // 0 = b h l d, 1 = b l h d + PagedAttentionParams& p +) { + const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + + 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(k_cache.sizes() == v_cache.sizes(), "k_cache and v_cache must have identical shapes"); + + // Normalize Q to b h l d view if user passed b l h d + if (layout == 1) { + q = q.transpose(1, 2); + } + + p.batch = (int)q.size(0); + p.q_head = (int)q.size(1); + p.q_len = (int)q.size(2); + p.head_dim = (int)q.size(3); + p.kv_head = (int)k_cache.size(2); + p.kv_len = (int)kv_len; + p.page_size = (int)page_size; + p.max_pages = (int)page_table.size(1); + + TORCH_CHECK(q.size(2) == 1, "Q seq_len must be 1 (decode)"); + TORCH_CHECK(p.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); + + // Q strides + p.q_stride_b = (int)q.stride(0); + p.q_stride_h = (int)q.stride(1); + p.q_stride_l = (int)q.stride(2); + p.q_stride_d = (int)q.stride(3); + + p.causal_offset = (int)causal_offset; + p.use_mask = (mask.has_value() && mask.value().defined()) ? 1 : 0; + p.scale = (scale > 0.0) ? (float)scale : 1.0f / sqrtf((float)p.head_dim); + + p.page_table = page_table.data_ptr(); + p.k_cache = (const T*)k_cache.data_ptr(); + p.v_cache = (const T*)v_cache.data_ptr(); + p.q = (const T*)q.data_ptr(); + p.o = nullptr; + p.o_part = nullptr; + p.ml_part = nullptr; + + if (p.use_mask) { + auto m = mask.value(); + TORCH_CHECK(m.is_cuda(), "mask must be on CUDA"); + TORCH_CHECK(m.dtype() == torch::kBool, "mask must be bool"); + TORCH_CHECK(m.size(0) == p.batch, "mask batch mismatch"); + TORCH_CHECK(m.size(m.dim() - 1) == p.kv_len, "mask kv_len mismatch"); + if (m.dim() == 2) { + p.mask_b_stride = (int)m.stride(0); + p.mask_q_stride = 0; + } else if (m.dim() == 3) { + TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch"); + p.mask_b_stride = (int)m.stride(0); + p.mask_q_stride = (int)m.stride(1); + } else { + TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]"); + } + p.mask = m.data_ptr(); + } else { + p.mask = nullptr; + p.mask_b_stride = 0; + p.mask_q_stride = 0; } } diff --git a/csrc/kernels/attn_mma_utils.cuh b/csrc/kernels/attn_mma_utils.cuh index f7307d9..47f95e2 100644 --- a/csrc/kernels/attn_mma_utils.cuh +++ b/csrc/kernels/attn_mma_utils.cuh @@ -122,7 +122,9 @@ template __device__ inline void mma_compute_scores( const unsigned Qa[KD][4], const bf16* __restrict__ sK, - int LD, int SWIZ_MASK, int lane, + int LD, + int SWIZ_MASK, + int lane, float Sacc[NC8][4]) { #pragma unroll @@ -142,13 +144,20 @@ __device__ inline void mma_compute_scores( // Online softmax + Oacc rescale for one K/V tile. // maxc0/maxc1: per-row KV column bounds (prefill: per-query-row causal limits; // decode: same value for both rows since q_len==1). +// qrow0/qrow1: query row indices (for 3D mask indexing; decode passes 0). +// mask_b_stride/mask_q_stride: mask layout (2D: mask_q_stride=0; 3D: =kv_len). // Reads Sacc (Q@K^T scores), applies causal/mask, computes P = exp(S - nm), // rescales Oacc by exp(m_old - nm), and updates m/l — all in place. template __device__ inline void mma_softmax_tile( int kv0, - int maxc0, int maxc1, - int mask_base, + int maxc0, + int maxc1, + int qrow0, + int qrow1, + int mask_b_stride, + int mask_q_stride, + int mask_batch, const bool* __restrict__ mask, bool has_mask, float Sacc[NC8][4], @@ -162,14 +171,16 @@ __device__ inline void mma_softmax_tile( // Mask out-of-bounds / masked columns: set -FLT_MAX so expf → 0 downstream // without per-element sentinel checks. Compute tile-local row maxima. float rmax0 = -FLT_MAX, rmax1 = -FLT_MAX; + int mask_base0 = mask_batch * mask_b_stride + qrow0 * mask_q_stride; + int mask_base1 = mask_batch * mask_b_stride + qrow1 * mask_q_stride; #pragma unroll for (int n8 = 0; n8 < NC8; n8++) { int cc = kv0 + n8 * 8 + 2 * tid4; int c1 = cc + 1; - bool b0 = (cc >= maxc0) || (has_mask && !mask[mask_base + cc]); - bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base + c1]); - bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base + cc]); - bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base + c1]); + bool b0 = (cc >= maxc0) || (has_mask && !mask[mask_base0 + cc]); + bool b1 = (c1 >= maxc0) || (has_mask && !mask[mask_base0 + c1]); + bool b2 = (cc >= maxc1) || (has_mask && !mask[mask_base1 + cc]); + bool b3 = (c1 >= maxc1) || (has_mask && !mask[mask_base1 + c1]); float s0 = b0 ? -FLT_MAX : Sacc[n8][0]; float s1 = b1 ? -FLT_MAX : Sacc[n8][1]; float s2 = b2 ? -FLT_MAX : Sacc[n8][2]; diff --git a/csrc/kernels/attn_paged_decode.cu b/csrc/kernels/attn_paged_decode.cu index 01f4c70..0af4c8c 100644 --- a/csrc/kernels/attn_paged_decode.cu +++ b/csrc/kernels/attn_paged_decode.cu @@ -3,26 +3,13 @@ #include "attn_paged_decode_split_kv_mma.cuh" #endif -#include -#include - -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& 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(); - p.ml_part = ml_part.data_ptr(); + 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& p) { template static void launch_paged_mma_decode(PagedAttentionParams& 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(); - p.ml_part = ml_part.data_ptr(); + p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total); + alloc_split_partials(p); paged_attn_decode_split_kv_mma_kernel<<>>(p); paged_attn_decode_combine_kernel<<>>(p); @@ -68,68 +50,19 @@ torch::Tensor attn_paged_decode( int64_t page_size, int64_t kv_len, c10::optional mask, - bool is_causal = false, - int64_t causal_offset = 0, - c10::optional scale = c10::nullopt + int64_t causal_offset, + double scale, + int64_t layout ) { - const at::cuda::OptionalCUDAGuard device_guard(device_of(q)); + PagedAttentionParams 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(scale.value()) - : 1.0f / std::sqrt(static_cast(head_dim)); - - auto O = torch::empty_like(q); - - PagedAttentionParams p; - p.batch = batch; - p.q_head = q_head; - p.kv_head = kv_head; - p.q_len = static_cast(q.size(2)); - p.kv_len = static_cast(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(causal_offset); - p.scale = scale_val; - p.page_size = static_cast(page_size); - p.max_pages = max_pages; - p.page_table = page_table.data_ptr(); - p.k_cache = reinterpret_cast(k_cache.data_ptr()); - p.v_cache = reinterpret_cast(v_cache.data_ptr()); - p.q = reinterpret_cast(q.data_ptr()); - p.mask = p.use_mask ? mask.value().data_ptr() : nullptr; - p.o = reinterpret_cast(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, [&]() { dispatch_paged_decode(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."); } diff --git a/csrc/kernels/attn_paged_decode_split_kv.cuh b/csrc/kernels/attn_paged_decode_split_kv.cuh index a54f467..7a66fc0 100644 --- a/csrc/kernels/attn_paged_decode_split_kv.cuh +++ b/csrc/kernels/attn_paged_decode_split_kv.cuh @@ -22,11 +22,13 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams p) int lane = threadIdx.x; int hd_per_thread = p.head_dim / 32; + // Q: stride-based [batch, q_head, q_len=1, head_dim] float q_reg[8]; - int q_off = ((batch * p.q_head + q_head) * 1) * p.head_dim + lane * hd_per_thread; + int q_off = batch * p.q_stride_b + q_head * p.q_stride_h + + lane * hd_per_thread * p.q_stride_d; #pragma unroll for (int i = 0; i < hd_per_thread; i++) - q_reg[i] = __bfloat162float(p.q[q_off + i]); + q_reg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]); float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f}; @@ -37,7 +39,7 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams p) int ch_begin = split * chunks_per_split; int ch_end = min(chunks_total, ch_begin + chunks_per_split); - const int mask_base = batch * p.kv_len; + const int mask_base = batch * p.mask_b_stride; for (int ci = ch_begin; ci < ch_end; ci++) { int chunk_start = ci * PDC_CHUNK; @@ -70,9 +72,10 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams p) partial += q_reg[i] * __bfloat162float(k_smem[s * p.head_dim + lane * hd_per_thread + i]); partial = paged_warp_reduce_sum(partial) * p.scale; - if (p.use_mask && p.mask && !p.mask[mask_base + chunk_start + s]) + int kv_idx = chunk_start + s; + if (p.use_mask && p.mask && !p.mask[mask_base + kv_idx]) partial = -FLT_MAX; - if (p.is_causal && (chunk_start + s) > p.causal_offset) + if (p.causal_offset >= 0 && kv_idx > p.causal_offset) partial = -FLT_MAX; float new_m = fmaxf(m, partial); @@ -118,6 +121,9 @@ __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams p) { int d = threadIdx.x; if (d >= p.head_dim) return; + int batch = bh / p.q_head; + int q_head = bh % p.q_head; + size_t split_base = (size_t)bh * p.num_splits; const float* mlp = p.ml_part + split_base * 2; const float* op = p.o_part + split_base * p.head_dim; @@ -136,5 +142,6 @@ __global__ void paged_attn_decode_combine_kernel(PagedAttentionParams p) { } float inv = (l > 1e-20f) ? (1.0f / l) : 0.0f; - p.o[(size_t)bh * p.head_dim + d] = __float2bfloat16(acc * inv); + int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + d * p.q_stride_d; + p.o[o_off] = __float2bfloat16(acc * inv); } diff --git a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh index f8c30f8..64f0d65 100644 --- a/csrc/kernels/attn_paged_decode_split_kv_mma.cuh +++ b/csrc/kernels/attn_paged_decode_split_kv_mma.cuh @@ -42,7 +42,8 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams __shared__ __align__(16) bf16 sV[STAGES * BC * LD]; // ---- Load Q directly from global into mma A-operand registers ---- - const int q_base = (batch * p.q_head + q_head0) * HEAD_DIM; + // Q stride-based: [batch, q_head, q_len=1, head_dim] + const int q_base = batch * p.q_stride_b + q_head0 * p.q_stride_h; const int qra = gid; const int qrb = gid + 8; const bool va = qra < G, vb = qrb < G; @@ -51,9 +52,9 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams for (int kt = 0; kt < KD; kt++) { int c = kt * 16 + tid4 * 2; const unsigned* pau = reinterpret_cast( - &p.q[q_base + qra * HEAD_DIM + c]); + &p.q[q_base + qra * p.q_stride_h + c * p.q_stride_d]); const unsigned* pbu = reinterpret_cast( - &p.q[q_base + qrb * HEAD_DIM + c]); + &p.q[q_base + qrb * p.q_stride_h + c * p.q_stride_d]); Qa[kt][0] = va ? pau[0] : 0u; Qa[kt][1] = vb ? pbu[0] : 0u; Qa[kt][2] = va ? pau[4] : 0u; @@ -66,7 +67,7 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f; float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f; - const int mask_base = batch * p.kv_len; + const int mask_batch_base = batch * p.mask_b_stride; const int tiles_total = (p.kv_len + BC - 1) / BC; const int tiles_per_split = (tiles_total + p.num_splits - 1) / p.num_splits; const int ti_begin = split * tiles_per_split; @@ -130,9 +131,13 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale, Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale; - int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; + // Decode: q_len=1, so qrow0=qrow1=0, mask_q_stride irrelevant + int maxc = (p.causal_offset >= 0) ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; mma_softmax_tile(kv0, maxc, maxc, - mask_base, p.mask, has_mask, + 0, 0, + mask_batch_base, 0, + batch, + p.mask, has_mask, Sacc, Oacc, m0, m1, l0, l1, lane); mma_pv_accumulate(Sacc, bV, LD, SWIZ_MASK, lane, Oacc); diff --git a/csrc/kernels/attn_prefill.cu b/csrc/kernels/attn_prefill.cu index bae21b0..7d621e6 100644 --- a/csrc/kernels/attn_prefill.cu +++ b/csrc/kernels/attn_prefill.cu @@ -35,34 +35,19 @@ torch::Tensor attn_prefill( torch::Tensor k, torch::Tensor v, c10::optional mask, - bool is_causal = false, - int64_t causal_offset = 0, - c10::optional scale = c10::nullopt + int64_t causal_offset, + double scale, + int64_t layout ) { AttentionParams p; - attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p); + attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p); TORCH_CHECK(p.head_dim % 16 == 0, "head_dim must be multiple of 16"); - auto O = torch::empty_like(q); - p.o = (bf16*)O.data_ptr(); + 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(); - switch (p.head_dim) { - case 32: - dispatch_prefill<32>(p); - break; - case 64: - dispatch_prefill<64>(p); - break; - case 128: - dispatch_prefill<128>(p); - break; - case 256: - dispatch_prefill<256>(p); - break; - default: - TORCH_CHECK(false, "prefill: unsupported head_dim ", p.head_dim, - " (supported: 32,64,128,256)"); - } + dispatch_head_dim(p.head_dim, [&]() { dispatch_prefill(p); }); return O; } @@ -72,8 +57,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("k"), py::arg("v"), 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, "GQA prefill (tensor-core mma on sm_80+, scalar fallback)"); } diff --git a/csrc/kernels/attn_prefill_split_q.cuh b/csrc/kernels/attn_prefill_split_q.cuh index d84e08b..6faa6a1 100644 --- a/csrc/kernels/attn_prefill_split_q.cuh +++ b/csrc/kernels/attn_prefill_split_q.cuh @@ -50,12 +50,14 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { __shared__ __align__(16) bf16 sK[P_BC * HEAD_DIM]; __shared__ __align__(16) bf16 sV[P_BC * HEAD_DIM]; + // Q: stride-based load [batch, q_head, q_len, head_dim] float qreg[DPT]; if (q_row < p.q_len) { - int q_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT; + int q_off = batch * p.q_stride_b + q_head * p.q_stride_h + + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d; #pragma unroll for (int i = 0; i < DPT; i++) - qreg[i] = __bfloat162float(p.q[q_off + i]) * p.scale; + qreg[i] = __bfloat162float(p.q[q_off + i * p.q_stride_d]) * p.scale; } float m = -FLT_MAX, l = 0.0f; @@ -64,7 +66,9 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { for (int i = 0; i < DPT; i++) acc[i] = 0.0f; - int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM; + // KV: stride-based base + int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h; + int mask_batch_base = batch * p.mask_b_stride; int tiles = (p.kv_len + P_BC - 1) / P_BC; int tt = G * ROWS; int lid = row * G + gpos; @@ -79,15 +83,19 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { int kv0 = ti * P_BC; int tlen = min(P_BC, p.kv_len - kv0); + // Load K/V into shared memory from strided global for (int i = lid; i < tlen * HEAD_DIM; i += tt) { - int gidx = kv_base + (kv0 + i / HEAD_DIM) * HEAD_DIM + (i % HEAD_DIM); - sK[i] = p.k[gidx]; - sV[i] = p.v[gidx]; + int s = i / HEAD_DIM; + int d_dim = i % HEAD_DIM; + int kv_idx = kv0 + s; + int g_off = kv_base + kv_idx * p.kv_stride_l + d_dim * p.kv_stride_d; + sK[i] = p.k[g_off]; + sV[i] = p.v[g_off]; } __syncthreads(); int lim = tlen; - if (p.is_causal && q_row < p.q_len) { + if (p.causal_offset >= 0 && q_row < p.q_len) { int ep = q_row + p.causal_offset + 1; if (kv0 >= ep) lim = 0; @@ -95,6 +103,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { lim = ep - kv0; } + int mask_row_base = mask_batch_base + q_row * p.mask_q_stride; for (int s = 0; s < lim; s++) { const bf16* kr = sK + s * HEAD_DIM + gpos * DPT; float part = 0.0f; @@ -108,7 +117,8 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { } float dot = group_reduce_sum(part, gmask); - if (p.use_mask && p.mask && !p.mask[batch * p.kv_len + kv0 + s]) + int kv_idx = kv0 + s; + if (p.use_mask && p.mask && !p.mask[mask_row_base + kv_idx]) dot = -FLT_MAX; float nm = fmaxf(m, dot); @@ -131,10 +141,12 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams p) { } if (q_row < p.q_len) { - int o_off = ((batch * p.q_head + q_head) * p.q_len + q_row) * HEAD_DIM + gpos * DPT; + // O: stride-based write + int o_off = batch * p.q_stride_b + q_head * p.q_stride_h + + q_row * p.q_stride_l + gpos * DPT * p.q_stride_d; float rl = (l > 1e-10f) ? (1.0f / l) : 0.0f; #pragma unroll for (int i = 0; i < DPT; i++) - p.o[o_off + i] = __float2bfloat16(acc[i] * rl); + p.o[o_off + i * p.q_stride_d] = __float2bfloat16(acc[i] * rl); } } diff --git a/csrc/kernels/attn_prefill_split_q_mma.cuh b/csrc/kernels/attn_prefill_split_q_mma.cuh index 628d170..edc62db 100644 --- a/csrc/kernels/attn_prefill_split_q_mma.cuh +++ b/csrc/kernels/attn_prefill_split_q_mma.cuh @@ -72,7 +72,8 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { // registers across the tile loop. // frag[0]/[2]: row = qrow0 + gid ; frag[1]/[3]: row = qrow0 + gid + 8 // frag[0]/[1]: cols kt*16 + tid4*2 + {0,1} ; frag[2]/[3]: + 8 - const int q_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM; + // Q stride-based: [batch, q_head, q_len, head_dim] + const int q_base = batch * p.q_stride_b + q_head * p.q_stride_h; const int qra = qrow0 + gid; const int qrb = qrow0 + gid + 8; const bool va = qra < p.q_len, vb = qrb < p.q_len; @@ -81,9 +82,9 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { for (int kt = 0; kt < KD; kt++) { int c = kt * 16 + tid4 * 2; const unsigned* pau = reinterpret_cast( - &p.q[q_base + qra * HEAD_DIM + c]); + &p.q[q_base + qra * p.q_stride_l + c * p.q_stride_d]); const unsigned* pbu = reinterpret_cast( - &p.q[q_base + qrb * HEAD_DIM + c]); + &p.q[q_base + qrb * p.q_stride_l + c * p.q_stride_d]); Qa[kt][0] = va ? pau[0] : 0u; Qa[kt][1] = vb ? pbu[0] : 0u; Qa[kt][2] = va ? pau[4] : 0u; @@ -96,18 +97,19 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { Oacc[j][0] = Oacc[j][1] = Oacc[j][2] = Oacc[j][3] = 0.0f; float m0 = -FLT_MAX, m1 = -FLT_MAX, l0 = 0.0f, l1 = 0.0f; - const int kv_base = ((batch * p.kv_head + kv_head) * p.kv_len) * HEAD_DIM; + // KV: stride-based base + const int kv_base = batch * p.kv_stride_b + kv_head * p.kv_stride_h; const int tiles = (p.kv_len + BC - 1) / BC; const int qr0 = qrow0 + gid; // row for c0/c1 const int qr1 = qrow0 + gid + 8; // row for c2/c3 - // Causal tile-skip bounds (no-op when is_causal == 0) - const int use_skip = p.is_causal; + // Causal tile-skip bounds (no-op when causal_offset < 0) + const int use_skip = (p.causal_offset >= 0) ? 1 : 0; const int max_kv = qrow0 + BR - 1 + p.causal_offset; const int block_max_kv = blockIdx.x * WARPS * BR + WARPS * BR - 1 + p.causal_offset; const int has_mask = p.use_mask && p.mask; - const int mb = batch * p.kv_len; + const int mask_batch_base = batch * p.mask_b_stride; // Last active tile: block-level causal bound (all warps in the block share // the K/V load, so the prefetch range is the block max, not per-warp). @@ -132,8 +134,9 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { int kc = kv0 + r; bool valid = kc < p.kv_len; int off = r * LD + swiz_col(d, r, SWIZ_MASK); - cp_async_16_pred(&dK[off], &p.k[kv_base + kc * HEAD_DIM + d], valid); - cp_async_16_pred(&dV[off], &p.v[kv_base + kc * HEAD_DIM + d], valid); + int g_off = kv_base + kc * p.kv_stride_l + d * p.kv_stride_d; + cp_async_16_pred(&dK[off], &p.k[g_off], valid); + cp_async_16_pred(&dV[off], &p.v[g_off], valid); } cp_async_commit(); }; @@ -170,12 +173,15 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale, Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale; - int maxc0 = p.is_causal ? min(p.kv_len, qr0 + p.causal_offset + 1) - : p.kv_len; - int maxc1 = p.is_causal ? min(p.kv_len, qr1 + p.causal_offset + 1) - : p.kv_len; + int maxc0 = (p.causal_offset >= 0) ? min(p.kv_len, qr0 + p.causal_offset + 1) + : p.kv_len; + int maxc1 = (p.causal_offset >= 0) ? min(p.kv_len, qr1 + p.causal_offset + 1) + : p.kv_len; mma_softmax_tile(kv0, maxc0, maxc1, - mb, p.mask, has_mask, + qr0, qr1, + mask_batch_base, p.mask_q_stride, + batch, + p.mask, has_mask, Sacc, Oacc, m0, m1, l0, l1, lane); mma_pv_accumulate(Sacc, bV, LD, SWIZ_MASK, lane, Oacc); @@ -186,19 +192,20 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams p) { // halves store count and removes the uncoalesced scalar-store penalty) float rl0 = (l0 > 1e-20f) ? (1.0f / l0) : 0.0f; float rl1 = (l1 > 1e-20f) ? (1.0f / l1) : 0.0f; - const int o_base = ((batch * p.q_head + q_head) * p.q_len) * HEAD_DIM; + // O: stride-based write + const int o_base = batch * p.q_stride_b + q_head * p.q_stride_h; #pragma unroll for (int dn8 = 0; dn8 < DN8; dn8++) { int d = dn8 * 8 + 2 * tid4; if (qr0 < p.q_len) { __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][0] * rl0, Oacc[dn8][1] * rl0); - *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr0 * HEAD_DIM + d]) = v; + *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr0 * p.q_stride_l + d * p.q_stride_d]) = v; } if (qr1 < p.q_len) { __nv_bfloat162 v = __floats2bfloat162_rn(Oacc[dn8][2] * rl1, Oacc[dn8][3] * rl1); - *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr1 * HEAD_DIM + d]) = v; + *reinterpret_cast<__nv_bfloat162*>(&p.o[o_base + qr1 * p.q_stride_l + d * p.q_stride_d]) = v; } } } diff --git a/csrc/tests/attn_decode_test.cu b/csrc/tests/attn_decode_test.cu index 549669e..780dcb3 100644 --- a/csrc/tests/attn_decode_test.cu +++ b/csrc/tests/attn_decode_test.cu @@ -98,8 +98,9 @@ static void bench() { AttentionParams p; p.batch = B; p.q_head = Hq; p.kv_head = Hk; p.q_len = 1; p.kv_len = sl; - p.head_dim = D; p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0; + p.head_dim = D; p.use_mask = 0; p.causal_offset = -1; p.scale = 1.0f / sqrtf((float)D); + set_default_strides(p); p.q = dQ; p.k = dK; p.v = dV; p.mask = nullptr; p.o = dO; DecodeScratch sc; @@ -160,7 +161,7 @@ int main() { AttentionParams p; p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=1; p.kv_len=sl; p.head_dim=D; - p.use_mask=0; p.is_causal=0; p.causal_offset=0; + p.use_mask=0; p.causal_offset=-1; p.scale=1.0f/sqrtf((float)D); p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO; @@ -180,7 +181,7 @@ int main() { cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost); float* ref=new float[nQ]; - cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, 0, 0); + cpu_attention_ref(hQ, hK, hV, hMask, ref, B, Hq, Hk, 1, sl, D, -1); float max_err=0; for (size_t i=0;i p; p.batch = B; p.q_head = Hq; p.kv_head = Hkv; p.q_len = 1; p.kv_len = kv_len; p.head_dim = HEAD_DIM; - p.use_mask = 0; p.is_causal = 0; p.causal_offset = 0; + p.use_mask = 0; p.causal_offset = -1; + set_default_strides(p); p.num_splits = 1; p.scale = scale_val; p.page_size = page_size; p.max_pages = max_pages; p.page_table = d_pt; @@ -272,7 +273,8 @@ static void bench_config(int B, int Hq, int Hkv, int kv_len, int page_size) { PagedAttentionParams pa; pa.batch = B; pa.q_head = Hq; pa.kv_head = Hkv; pa.q_len = 1; pa.kv_len = kv_len; pa.head_dim = HEAD_DIM; - pa.use_mask = 0; pa.is_causal = 0; pa.causal_offset = 0; + pa.use_mask = 0; pa.causal_offset = -1; + set_default_paged_strides(pa); pa.num_splits = 1; pa.scale = scale_val; pa.page_size = page_size; pa.max_pages = max_pages; pa.page_table = d_pt; diff --git a/csrc/tests/attn_prefill_test.cu b/csrc/tests/attn_prefill_test.cu index 6e315a5..30bfe40 100644 --- a/csrc/tests/attn_prefill_test.cu +++ b/csrc/tests/attn_prefill_test.cu @@ -75,7 +75,8 @@ static void bench() { AttentionParams p; p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D; - p.use_mask=0; p.is_causal=causal; p.causal_offset=0; + p.use_mask=0; p.causal_offset=causal?0:-1; + set_default_strides(p); p.scale=1.0f/sqrtf((float)D); p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO; @@ -143,7 +144,8 @@ int main() { AttentionParams p; p.batch=B; p.q_head=Hq; p.kv_head=Hk; p.q_len=ql; p.kv_len=kl; p.head_dim=D; - p.use_mask=0; p.is_causal=causal; p.causal_offset=0; + p.use_mask=0; p.causal_offset=causal?0:-1; + set_default_strides(p); p.scale=1.0f/sqrtf((float)D); p.q=dQ; p.k=dK; p.v=dV; p.mask=nullptr; p.o=dO; @@ -158,7 +160,7 @@ int main() { cudaMemcpy(hOut,dO,nQ*2,cudaMemcpyDeviceToHost); float* ref=new float[nQ]; - cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal, 0); + cpu_attention_ref(hQ, hK, hV, nullptr, ref, B, Hq, Hk, ql, kl, D, causal ? 0 : -1); float max_err=0; for (size_t i=0;i::call(head_dim, fn); } +// Set default strides for contiguous b h l d layout on AttentionParams. +template +inline void set_default_strides(P& p) { + p.q_stride_b = p.q_head * p.q_len * p.head_dim; + p.q_stride_h = p.q_len * p.head_dim; + p.q_stride_l = p.head_dim; + p.q_stride_d = 1; + p.kv_stride_b = p.kv_head * p.kv_len * p.head_dim; + p.kv_stride_h = p.kv_len * p.head_dim; + p.kv_stride_l = p.head_dim; + p.kv_stride_d = 1; + p.mask_b_stride = p.kv_len; + p.mask_q_stride = 0; +} + +// Set default Q strides for contiguous b h l d layout on PagedAttentionParams. +template +inline void set_default_paged_strides(P& p) { + p.q_stride_b = p.q_head * p.q_len * p.head_dim; + p.q_stride_h = p.q_len * p.head_dim; + p.q_stride_l = p.head_dim; + p.q_stride_d = 1; + p.mask_b_stride = p.kv_len; + p.mask_q_stride = 0; +} + // Generic CPU reference for multi-query / grouped-query attention. // Tensor shapes (all float*): // Q : [B, Hq, q_len, D] @@ -108,10 +134,11 @@ void dispatch_by_head_dim(int head_dim, Fn&& fn) { // V : [B, Hk, kv_len, D] // O : [B, Hq, q_len, D] // mask: if q_len == 1, shape is [B, kv_len]; otherwise mask is not supported. +// causal_offset: -1 = non-causal; >=0 = absolute position of first Q token. static void cpu_attention_ref( const float* Q, const float* K, const float* V, const bool* mask, float* O, int B, int Hq, int Hk, int q_len, int kv_len, int D, - int is_causal, int causal_offset + int causal_offset ) { float scale = 1.0f / sqrtf((float)D); int n_rep = Hq / Hk; @@ -122,7 +149,7 @@ static void cpu_attention_ref( float mv = -INFINITY, sv = 0.0f; float accum[256] = {0.0f}; int lim = kv_len; - if (is_causal) { + if (causal_offset >= 0) { int c = qi + causal_offset + 1; lim = (c < kv_len) ? c : kv_len; }