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))
csrc/kernels/attn_common.h
+7 -4
View File
@@ -19,9 +19,11 @@ struct AttentionParams {
// 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
// Mask: 2D [batch, kv_len], 3D [batch, q_len, kv_len],
// or 4D [batch, n_heads, q_len, kv_len] (head dim broadcasts when stride=0)
int mask_b_stride; // batch stride
int mask_h_stride; // head stride (0 = broadcast across heads)
int mask_q_stride; // q stride (0 = all q rows share)
const T* __restrict__ q;
const T* __restrict__ k;
@@ -52,8 +54,9 @@ struct PagedAttentionParams {
// Q strides (layout-agnostic)
int q_stride_b, q_stride_h, q_stride_l, q_stride_d;
// Mask strides (2D or 3D)
// Mask strides (2D, 3D, or 4D)
int mask_b_stride;
int mask_h_stride;
int mask_q_stride;
const T* __restrict__ q;
csrc/kernels/attn_decode_split_kv.cuh
+1 -1
View File
@@ -24,7 +24,7 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
// 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;
int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
csrc/kernels/attn_decode_split_kv_mma.cuh
+2 -2
View File
@@ -109,8 +109,8 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
0, 0,
p.mask_b_stride, 0,
batch,
p.mask_b_stride, 0, 0,
batch, 0,
p.mask,
Sacc, Oacc, m0, m1, l0, l1, lane);
csrc/kernels/attn_entry_utils.cuh
+14 -3
View File
@@ -44,6 +44,9 @@ inline void extract_q_dims_and_strides(torch::Tensor& q, int64_t layout, P& p) {
}
// ---- Shared mask packing ----
// Accepts 2D [batch, kv_len], 3D [batch, q_len, kv_len],
// or 4D [batch, n_heads, q_len, kv_len].
// Head/q dimensions with size 1 broadcast (stride set to 0).
template <typename P>
inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
if (p.use_mask) {
@@ -54,18 +57,26 @@ inline void pack_mask(const c10::optional<torch::Tensor>& mask, P& p) {
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_h_stride = 0;
p.mask_q_stride = 0;
} else if (m.dim() == 3) {
TORCH_CHECK(m.size(1) == p.q_len, "mask q_len mismatch");
TORCH_CHECK(m.size(1) == 1 || 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);
p.mask_h_stride = 0;
p.mask_q_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
} else if (m.dim() == 4) {
TORCH_CHECK(m.size(2) == 1 || m.size(2) == p.q_len, "mask q_len mismatch");
p.mask_b_stride = (int)m.stride(0);
p.mask_h_stride = (m.size(1) == 1) ? 0 : (int)m.stride(1);
p.mask_q_stride = (m.size(2) == 1) ? 0 : (int)m.stride(2);
} else {
TORCH_CHECK(false, "mask must be 2D [batch, kv_len] or 3D [batch, q_len, kv_len]");
TORCH_CHECK(false, "mask must be 2D, 3D, or 4D");
}
p.mask = m.data_ptr<bool>();
} else {
p.mask = nullptr;
p.mask_b_stride = 0;
p.mask_h_stride = 0;
p.mask_q_stride = 0;
}
}
csrc/kernels/attn_mma_utils.cuh
+4 -4
View File
@@ -192,8 +192,8 @@ __device__ inline void mma_softmax_tile(
int kv0,
int maxc0, int maxc1,
int qrow0, int qrow1,
int mask_b_stride, int mask_q_stride,
int mask_batch,
int mask_b_stride, int mask_h_stride, int mask_q_stride,
int mask_batch, int mask_head,
const bool* __restrict__ mask,
float Sacc[Traits::NC8][4],
float Oacc[Traits::DN8][4],
@@ -204,8 +204,8 @@ __device__ inline void mma_softmax_tile(
int tid4 = lane & 3;
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;
int mask_base0 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow0 * mask_q_stride;
int mask_base1 = mask_batch * mask_b_stride + mask_head * mask_h_stride + qrow1 * mask_q_stride;
#pragma unroll
for (int n8 = 0; n8 < Traits::NC8; n8++) {
int cc = kv0 + n8 * 8 + 2 * tid4;
csrc/kernels/attn_paged_decode_split_kv.cuh
+1 -1
View File
@@ -31,7 +31,7 @@ __global__ void paged_attn_decode_split_kv_kernel(PagedAttentionParams<bf16> 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.mask_b_stride;
const int mask_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
for (int ci = ch_begin; ci < ch_end; ci++) {
int chunk_start = ci * PDC_CHUNK;
csrc/kernels/attn_paged_decode_split_kv_mma.cuh
+2 -2
View File
@@ -110,8 +110,8 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
int maxc = IsCausal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
mma_softmax_tile<Traits, HasMask>(kv0, maxc, maxc,
0, 0,
p.mask_b_stride, 0,
batch,
p.mask_b_stride, 0, 0,
batch, 0,
p.mask,
Sacc, Oacc, m0, m1, l0, l1, lane);
csrc/kernels/attn_prefill_split_q.cuh
+1 -1
View File
@@ -64,7 +64,7 @@ __global__ void attn_prefill_split_q_kernel_t(AttentionParams<bf16> p) {
// 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 mask_batch_base = batch * p.mask_b_stride + q_head * p.mask_h_stride;
int tiles = (p.kv_len + P_BC - 1) / P_BC;
int tt = G * ROWS;
int lid = row * G + gpos;
csrc/kernels/attn_prefill_split_q_mma.cuh
+2 -2
View File
@@ -114,8 +114,8 @@ __global__ void attn_prefill_split_q_mma_kernel(AttentionParams<bf16> p) {
: p.kv_len;
mma_softmax_tile<Traits, HasMask>(kv0, maxc0, maxc1,
qr0, qr1,
p.mask_b_stride, p.mask_q_stride,
batch,
p.mask_b_stride, p.mask_h_stride, p.mask_q_stride,
batch, q_head,
p.mask,
Sacc, Oacc, m0, m1, l0, l1, lane);
csrc/tests/test_utils.cuh
+2
View File
@@ -103,6 +103,7 @@ inline void set_default_strides(P& p) {
p.kv_stride_l = p.head_dim;
p.kv_stride_d = 1;
p.mask_b_stride = p.kv_len;
p.mask_h_stride = 0;
p.mask_q_stride = 0;
}
@@ -114,6 +115,7 @@ inline void set_default_paged_strides(P& p) {
p.q_stride_l = p.head_dim;
p.q_stride_d = 1;
p.mask_b_stride = p.kv_len;
p.mask_h_stride = 0;
p.mask_q_stride = 0;
}
tests/extension/conftest.py
+35
View File
@@ -0,0 +1,35 @@
"""Shared fixtures for extension tests."""
import pytest
import torch
from astrai.config.model_config import AutoRegressiveLMConfig
from astrai.extension import is_available
from astrai.model.transformer import AutoRegressiveLM
CUDA_AVAILABLE = torch.cuda.is_available() and is_available("attn_paged_decode")
skip_no_cuda = pytest.mark.skipif(
not CUDA_AVAILABLE, reason="CUDA not available or kernels not built"
)
D = 64
CFG = dict(
vocab_size=1000,
hidden_size=128,
num_attention_heads=2,
num_key_value_heads=1,
intermediate_size=256,
max_position_embeddings=64,
num_hidden_layers=2,
rms_norm_eps=1e-5,
attn_type="gqa",
ffn_type="mlp",
)
@pytest.fixture
def cuda_model():
config = AutoRegressiveLMConfig(**CFG)
model = AutoRegressiveLM(config).to(device="cuda", dtype=torch.bfloat16)
model.eval()
return model, config
tests/extension/test_backend.py
+45
View File
@@ -0,0 +1,45 @@
"""Backend selection and context-manager switching tests.
These tests do not require CUDA — they only check that the active
backend is correctly set and restored.
"""
import pytest
from astrai.extension import (
ATTN_BACKEND,
CudaBackend,
TorchNativeBackend,
attn_backend,
get_backend,
)
def test_default_backend_is_torch_native():
backend = get_backend()
assert isinstance(backend, TorchNativeBackend)
def test_attn_backend_context_with_enum():
with attn_backend(ATTN_BACKEND.CUDA):
assert isinstance(get_backend(), CudaBackend)
assert isinstance(get_backend(), TorchNativeBackend)
def test_attn_backend_context_with_class():
with attn_backend(CudaBackend):
assert isinstance(get_backend(), CudaBackend)
assert isinstance(get_backend(), TorchNativeBackend)
def test_attn_backend_context_with_instance():
custom = CudaBackend()
with attn_backend(custom):
assert get_backend() is custom
assert isinstance(get_backend(), TorchNativeBackend)
def test_cudabackend_is_context_manager():
with CudaBackend():
assert isinstance(get_backend(), CudaBackend)
assert isinstance(get_backend(), TorchNativeBackend)
tests/extension/test_backend_equivalence.py
+199
View File
@@ -0,0 +1,199 @@
"""Numerical equivalence between TorchNativeBackend and CudaBackend.
Covers training forward, inference prefill, inference decode (mixed
seq_lens with padding mask), and end-to-end scheduler.run_batch.
"""
import torch
from astrai.extension import ATTN_BACKEND, attn_backend
from astrai.inference.core.cache import PagePool
from tests.extension.conftest import D, skip_no_cuda
@skip_no_cuda
def test_training_forward_matches_torch(cuda_model):
"""Training forward (kv_cache=None) should produce identical logits."""
model, _ = cuda_model
input_ids = torch.randint(0, 1000, (2, 16), device="cuda")
with torch.no_grad():
out_torch = model(input_ids)
with attn_backend(ATTN_BACKEND.CUDA):
with torch.no_grad():
out_cuda = model(input_ids)
diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
assert diff == 0.0, f"Training forward diff {diff} should be 0"
@skip_no_cuda
def test_prefill_with_kv_cache_matches_torch(cuda_model):
"""Inference prefill with KV cache should match torch backend."""
model, _ = cuda_model
prompt_ids = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 11, 12, 13, 14, 15]]
max_len = max(len(p) for p in prompt_ids)
batch = len(prompt_ids)
device = "cuda"
input_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
input_mask = torch.zeros(batch, max_len, dtype=torch.bool, device=device)
position_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
for i, p in enumerate(prompt_ids):
input_ids[i, : len(p)] = torch.tensor(p, device=device)
input_mask[i, : len(p)] = True
position_ids[i, : len(p)] = torch.arange(len(p), device=device)
cache = PagePool(
n_layers=2,
n_kv_heads=1,
head_dim=D,
max_batch_size=4,
max_seq_len=64,
device=device,
dtype=torch.bfloat16,
)
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv1 = cache.bind_tasks(
["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
)
with torch.inference_mode():
out_torch = model(
input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids
)
cache.task_free("t1")
cache.task_free("t2")
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv2 = cache.bind_tasks(
["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
)
with attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode():
out_cuda = model(
input_ids,
input_mask=input_mask,
kv_cache=kv2,
position_ids=position_ids,
)
for i, p in enumerate(prompt_ids):
d = (
(
out_torch["logits"][i, : len(p)].float()
- out_cuda["logits"][i, : len(p)].float()
)
.abs()
.max()
.item()
)
assert d == 0.0, f"Prefill diff for sample {i}: {d}"
@skip_no_cuda
def test_decode_mixed_seq_lens_matches_torch(cuda_model):
"""Decode with mixed seq_lens in batch — padding mask must produce correct output."""
model, _ = cuda_model
device = "cuda"
prompt_ids = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 11, 12, 13, 14, 15]]
cache = PagePool(
n_layers=2,
n_kv_heads=1,
head_dim=D,
max_batch_size=4,
max_seq_len=64,
device=device,
dtype=torch.bfloat16,
)
# Prefill to populate cache
max_len = max(len(p) for p in prompt_ids)
batch = len(prompt_ids)
input_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
input_mask = torch.zeros(batch, max_len, dtype=torch.bool, device=device)
position_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device)
for i, p in enumerate(prompt_ids):
input_ids[i, : len(p)] = torch.tensor(p, device=device)
input_mask[i, : len(p)] = True
position_ids[i, : len(p)] = torch.arange(len(p), device=device)
cache.task_alloc("t1", prompt_ids[0])
cache.task_alloc("t2", prompt_ids[1])
kv = cache.bind_tasks(
["t1", "t2"], [len(prompt_ids[0]), len(prompt_ids[1])], device, start_pos=0
)
with torch.inference_mode():
model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids)
# Decode step — seq_lens are 9 and 7 (after extending)
dec_ids = torch.tensor([[99], [98]], dtype=torch.long, device=device)
dec_pos = torch.tensor([[8], [6]], dtype=torch.long, device=device)
total_len = 9
dec_mask = dec_pos[:, None, None] >= torch.arange(total_len, device=device)
kv_t = cache.bind_tasks(["t1", "t2"], [9, 7], device)
with torch.inference_mode():
out_torch = model(
dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
)
kv_c = cache.bind_tasks(["t1", "t2"], [9, 7], device)
with attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode():
out_cuda = model(
dec_ids, input_mask=dec_mask, kv_cache=kv_c, position_ids=dec_pos
)
diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
@skip_no_cuda
def test_run_batch_cuda_matches_torch_greedy(cuda_model):
"""Greedy decode (temperature=0) should produce identical tokens."""
from astrai.inference.core.scheduler import InferenceScheduler
from tests.helpers import FakeTokenizer
model, _ = cuda_model
tokenizer = FakeTokenizer()
prompts = [[1, 2, 3, 4, 5], [10, 11, 12, 13, 14, 15, 16]]
sched = InferenceScheduler(
model=model,
tokenizer=tokenizer,
max_batch_size=4,
max_seq_len=64,
device="cuda",
dtype=torch.bfloat16,
)
out_torch = sched.run_batch(prompts, max_tokens=5, temperature=0.0)
sched.stop()
cache_cuda = PagePool(
n_layers=2,
n_kv_heads=1,
head_dim=D,
max_batch_size=4,
max_seq_len=64,
device="cuda",
dtype=torch.bfloat16,
)
sched2 = InferenceScheduler(
model=model,
tokenizer=tokenizer,
max_batch_size=4,
max_seq_len=64,
device="cuda",
dtype=torch.bfloat16,
cache=cache_cuda,
)
with attn_backend(ATTN_BACKEND.CUDA):
out_cuda = sched2.run_batch(prompts, max_tokens=5, temperature=0.0)
sched2.stop()
assert out_torch == out_cuda, f"Torch={out_torch} != CUDA={out_cuda}"
tests/extension/test_kernel_mask.py
+74
View File
@@ -0,0 +1,74 @@
"""Kernel-level mask dimension support (2D, 3D, 4D)."""
import torch
from tests.extension.conftest import D, skip_no_cuda
@skip_no_cuda
def test_kernel_accepts_2d_mask():
"""Kernel should accept 2D mask [batch, kv_len]."""
from astrai.extension.attention_ops import attn_prefill
batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
kv_len = 8
q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
mask = torch.ones(batch, kv_len, dtype=torch.bool, device="cuda")
mask[:, 4:] = False
out = attn_prefill(q, k, v, mask=mask, is_causal=False)
assert out.shape == (batch, q_len, n_heads, D)
@skip_no_cuda
def test_kernel_accepts_3d_mask():
"""Kernel should accept 3D mask [batch, q_len, kv_len]."""
from astrai.extension.attention_ops import attn_prefill
batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
kv_len = 8
q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
mask = torch.ones(batch, q_len, kv_len, dtype=torch.bool, device="cuda")
out = attn_prefill(q, k, v, mask=mask, is_causal=False)
assert out.shape == (batch, q_len, n_heads, D)
@skip_no_cuda
def test_kernel_accepts_4d_mask():
"""Kernel should accept 4D mask [batch, n_heads, q_len, kv_len]."""
from astrai.extension.attention_ops import attn_prefill
batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
kv_len = 8
q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
mask = torch.ones(batch, 1, q_len, kv_len, dtype=torch.bool, device="cuda")
mask[:, :, :, 4:] = False
out = attn_prefill(q, k, v, mask=mask, is_causal=False)
assert out.shape == (batch, q_len, n_heads, D)
@skip_no_cuda
def test_4d_mask_matches_no_mask_when_all_true():
"""A 4D all-True mask should produce the same output as no mask."""
from astrai.extension.attention_ops import attn_prefill
batch, q_len, n_heads, n_kv_heads = 1, 8, 4, 1
kv_len = 8
q = torch.randn(batch, q_len, n_heads, D, device="cuda", dtype=torch.bfloat16)
k = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
v = torch.randn(batch, kv_len, n_kv_heads, D, device="cuda", dtype=torch.bfloat16)
out_no_mask = attn_prefill(q, k, v, mask=None, is_causal=False)
mask = torch.ones(batch, 1, q_len, kv_len, dtype=torch.bool, device="cuda")
out_with_mask = attn_prefill(q, k, v, mask=mask, is_causal=False)
diff = (out_no_mask.float() - out_with_mask.float()).abs().max().item()
assert diff == 0.0, f"4D all-True mask diff: {diff}"