92 lines
2.4 KiB
Python
92 lines
2.4 KiB
Python
import importlib
|
|
import logging
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_available: dict[str, bool] = {}
|
|
_modules: dict[str, object] = {}
|
|
|
|
for _name in ["gqa_decode_attn", "gqa_prefill_attn"]:
|
|
try:
|
|
_mod = importlib.import_module(f".{_name}", package=__package__)
|
|
_available[_name] = True
|
|
_modules[_name] = _mod
|
|
except ImportError:
|
|
_available[_name] = False
|
|
_modules[_name] = None
|
|
|
|
|
|
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:
|
|
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 gqa_decode_attn(
|
|
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["gqa_decode_attn"]:
|
|
return _modules["gqa_decode_attn"].gqa_decode_attn(
|
|
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 gqa_prefill_attn(
|
|
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["gqa_prefill_attn"]:
|
|
return _modules["gqa_prefill_attn"].gqa_prefill_attn(
|
|
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)
|