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

feat: unify attention backend with multi-dim mask support

Browse Source 
- Add attention() functional entry delegating to active backend
- GQA/MLA forward calls attention() instead of inline cache/SDPA
- CUDA kernels support 2D/3D/4D mask via mask_h_stride field
- CudaBackend.fwd_decode builds 2D padding mask for mixed seq_lens
- KVCache.max_len precomputed in bind_tasks to avoid GPU sync
- batch==1 decode short-circuits mask=None
- Split tests into conftest, test_backend, test_backend_equivalence, test_kernel_mask
- 440 tests pass, L20 decode 1.44-1.60x speedup vs torch native
This commit is contained in:
ViperEkura
2026-07-30 20:38:34 +08:00
parent 97114b95a4
commit 3067a8e1a6
19 changed files with 438 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/__init__.py
+2
View File
@@ -20,6 +20,7 @@ from astrai.extension.attention_backend import (
AttentionBackend,
CudaBackend,
TorchNativeBackend,
attention,
attn_backend,
get_backend,
)
@@ -35,6 +36,7 @@ __all__ = [
"AttentionBackend",
"CudaBackend",
"TorchNativeBackend",
"attention",
"attn_backend",
"get_backend",
"attn_decode",
astrai/extension/attention_backend.py
+40 -3
View File
@@ -113,6 +113,37 @@ def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
)
def attention(
q: Tensor,
k: Tensor,
v: Tensor,
kv_cache: Optional[KVCache] = None,
layer_id: int = 0,
attn_mask: Optional[Tensor] = None,
is_causal: bool = False,
) -> Tensor:
"""Functional attention entry point — mirrors ``F.scaled_dot_product_attention``.
Delegates to the active backend (set via ``with attn_backend(...)``).
Handles KV cache I/O, GQA head expansion, and causal masking so the
caller only needs to provide projected q/k/v.
Args:
q: [batch, q_len, n_heads, head_dim] (blhd)
k: [batch, q_len, n_kv_heads, head_dim] (blhd)
v: [batch, q_len, n_kv_heads, head_dim] (blhd)
kv_cache: cache dataclass, or None for training (no cache).
layer_id: transformer layer index for buffer access.
attn_mask: pre-built attention mask (SDPA-compatible).
is_causal: whether to apply causal masking.
Returns:
[batch, q_len, n_heads * head_dim]
"""
backend = get_backend()
return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
class AttentionBackend(ABC):
"""Abstract base for attention computation strategies.
@@ -307,13 +338,19 @@ class CudaBackend(AttentionBackend):
kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
max_len = kv_cache.seq_lens.max().item()
seq_lens = kv_cache.seq_lens
max_len = kv_cache.max_len
page_table = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
k_cache = kv_cache.k_buffer[layer_id].unsqueeze(1)
v_cache = kv_cache.v_buffer[layer_id].unsqueeze(1)
if q.size(0) == 1:
mask = None
else:
mask = torch.arange(max_len, device=q.device)[None, :] < seq_lens[:, None]
out = attn_paged_decode(
q,
page_table,
@@ -321,7 +358,7 @@ class CudaBackend(AttentionBackend):
v_cache,
page_size=1,
kv_len=max_len,
mask=None,
mask=mask,
is_causal=is_causal,
)
@@ -354,7 +391,7 @@ class CudaBackend(AttentionBackend):
kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
max_len = kv_cache.seq_lens.max()
max_len = kv_cache.max_len
indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
pos_mask = (
torch.arange(max_len, device=q.device)[None, :] < kv_cache.seq_lens[:, None]
astrai/inference/core/cache.py
+3
View File
@@ -202,6 +202,7 @@ class KVCache:
req_pool_indices: [batch_size] — row indices into req_to_token
seq_lens: [batch_size] — per-request total sequence lengths
out_cache_loc: [batch, new_seq_len] or [batch, 1] — write indices
max_len: max(seq_lens) as Python int — avoids GPU sync in decode
"""
k_buffer: Tensor
@@ -210,6 +211,7 @@ class KVCache:
req_pool_indices: Tensor
seq_lens: Tensor
out_cache_loc: Tensor
max_len: int = 0
class PagePool:
@@ -439,6 +441,7 @@ class PagePool:
req_pool_indices=req_pool_indices,
seq_lens=seq_lens_t,
out_cache_loc=out_cache_loc,
max_len=max(seq_lens),
)
# ---- internals ----
astrai/model/components/__init__.py
+1 -2
View File
@@ -1,4 +1,4 @@
from astrai.model.components.attention import GQA, MLA, repeat_kv
from astrai.model.components.attention import GQA, MLA
from astrai.model.components.decoder_block import DecoderBlock
from astrai.model.components.embedding import Embedding
from astrai.model.components.linear import Linear
@@ -21,5 +21,4 @@ __all__ = [
"RotaryEmbedding",
"apply_rotary_emb",
"get_rotary_emb",
"repeat_kv",
]
astrai/model/components/attention.py
+3 -56
View File
@@ -5,6 +5,7 @@ import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from astrai.extension import attention
from astrai.factory import BaseFactory
from astrai.inference.core.cache import KVCache
from astrai.model.components.linear import Linear
@@ -12,17 +13,6 @@ from astrai.model.components.norm import RMSNorm
from astrai.model.components.rope import apply_rotary_emb
def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
bs, slen, n_heads, head_dim = x.shape
if n_rep == 1:
return x
return (
x[:, :, :, None, :]
.expand(bs, slen, n_heads, n_rep, head_dim)
.reshape(bs, slen, n_heads * n_rep, head_dim)
)
class AttnFactory(BaseFactory[nn.Module]):
pass
@@ -86,29 +76,7 @@ class GQA(nn.Module):
if self.use_qk_norm:
q, k = self.q_norm(q), self.k_norm(k)
if kv_cache is not None:
kv_cache.k_buffer[self.layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[self.layer_id, kv_cache.out_cache_loc] = v
max_len = kv_cache.seq_lens.max()
indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
pos_mask = (
torch.arange(max_len, device=x.device)[None, :]
< kv_cache.seq_lens[:, None]
)
indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
k = kv_cache.k_buffer[self.layer_id, indices]
v = kv_cache.v_buffer[self.layer_id, indices]
k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
sdqa_out = (
F.scaled_dot_product_attention(q, k, v, attn_mask, is_causal=is_causal)
.permute(0, 2, 1, 3)
.contiguous()
.flatten(2)
)
sdqa_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal)
if self.use_gated_attention:
sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
@@ -203,28 +171,7 @@ class MLA(nn.Module):
q = self.q_norm(q)
k = self.k_norm(k)
if kv_cache is not None:
kv_cache.k_buffer[self.layer_id, kv_cache.out_cache_loc] = k
kv_cache.v_buffer[self.layer_id, kv_cache.out_cache_loc] = v
max_len = kv_cache.seq_lens.max()
indices = kv_cache.req_to_token[kv_cache.req_pool_indices, :max_len]
pos_mask = (
torch.arange(max_len, device=x.device)[None, :]
< kv_cache.seq_lens[:, None]
)
indices = torch.where(pos_mask, indices, torch.zeros_like(indices))
k = kv_cache.k_buffer[self.layer_id, indices]
v = kv_cache.v_buffer[self.layer_id, indices]
q = q.permute(0, 2, 1, 3)
k = k.permute(0, 2, 1, 3)
v = v.permute(0, 2, 1, 3)
attn_out = F.scaled_dot_product_attention(
q, k, v, attn_mask, is_causal=is_causal
)
attn_out = attn_out.permute(0, 2, 1, 3).contiguous().flatten(2)
attn_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal)
if self.use_gated_attention:
attn_out = attn_out * F.sigmoid(self.gate(x))