150 lines
4.2 KiB
Python
150 lines
4.2 KiB
Python
"""GQA attention wrapper functions — one entry point per compiled kernel.
|
|
|
|
Each wrapper dispatches to its CUDA kernel (loaded in ``loader.py``) when
|
|
available, otherwise falls back to ``torch`` SDPA.
|
|
|
|
Add new kernel wrappers here; split into per-variant files only if this file
|
|
grows large.
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from astrai.extension.loader import _available, _modules
|
|
|
|
|
|
def _expand_kv_heads(
|
|
k: torch.Tensor, v: torch.Tensor, q_head: int
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Expand K/V heads to match Q heads for GQA fallback."""
|
|
kv_head = k.size(1)
|
|
if kv_head == q_head:
|
|
return k, v
|
|
group = q_head // kv_head
|
|
k = k.repeat_interleave(group, dim=1)
|
|
v = v.repeat_interleave(group, dim=1)
|
|
return k, v
|
|
|
|
|
|
def _torch_fallback(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
mask: torch.Tensor | None,
|
|
is_causal: bool,
|
|
scale: float | None,
|
|
) -> torch.Tensor:
|
|
"""Reference attention via ``scaled_dot_product_attention``."""
|
|
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
|
|
)
|
|
|
|
|
|
def _gather_kv_from_pages(
|
|
page_table: torch.Tensor,
|
|
k_cache: torch.Tensor,
|
|
v_cache: torch.Tensor,
|
|
page_size: int,
|
|
kv_len: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Gather contiguous K/V from paged cache for torch SDPA fallback.
|
|
|
|
Shapes:
|
|
page_table : [batch, max_pages] (int64)
|
|
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]
|
|
"""
|
|
batch, max_pages = page_table.shape
|
|
n_pages, 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)
|
|
|
|
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, :, :]
|
|
return k, v
|
|
|
|
|
|
def attn_decode(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
mask: torch.Tensor | None = None,
|
|
is_causal: bool = False,
|
|
causal_offset: int = 0,
|
|
scale: float | None = None,
|
|
) -> torch.Tensor:
|
|
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,
|
|
)
|
|
return _torch_fallback(q, k, v, mask, is_causal, scale)
|
|
|
|
|
|
def attn_prefill(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
mask: torch.Tensor | None = None,
|
|
is_causal: bool = False,
|
|
causal_offset: int = 0,
|
|
scale: float | None = None,
|
|
) -> torch.Tensor:
|
|
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,
|
|
)
|
|
return _torch_fallback(q, k, v, mask, is_causal, scale)
|
|
|
|
|
|
def attn_paged_decode(
|
|
q: torch.Tensor,
|
|
page_table: torch.Tensor,
|
|
k_cache: torch.Tensor,
|
|
v_cache: torch.Tensor,
|
|
page_size: int,
|
|
kv_len: int,
|
|
mask: torch.Tensor | None = None,
|
|
is_causal: bool = False,
|
|
causal_offset: int = 0,
|
|
scale: float | None = None,
|
|
) -> torch.Tensor:
|
|
if _available["attn_paged_decode"]:
|
|
return _modules["attn_paged_decode"].attn_paged_decode(
|
|
q,
|
|
page_table,
|
|
k_cache,
|
|
v_cache,
|
|
page_size,
|
|
kv_len,
|
|
mask=mask,
|
|
is_causal=is_causal,
|
|
causal_offset=causal_offset,
|
|
scale=scale,
|
|
)
|
|
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)
|