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

refactor: standardize packed 3d inference

Browse Source 
- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch
This commit is contained in:
ViperEkura
2026-08-16 13:24:02 +08:00
parent 0dd9a417b7
commit 3406157431
15 changed files with 301 additions and 248 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/attention_backend.py
+108 -81
View File
@@ -123,6 +123,7 @@ def _backend_supports(
kv_cache: Optional["KVCache"], kv_cache: Optional["KVCache"],
attn_mask: Optional[Tensor], attn_mask: Optional[Tensor],
is_causal: bool, is_causal: bool,
fwd: Optional[str],
) -> bool: ) -> bool:
"""Whether ``backend`` can run this attention call. """Whether ``backend`` can run this attention call.
@@ -131,17 +132,20 @@ def _backend_supports(
""" """
if isinstance(backend, CudaBackend): if isinstance(backend, CudaBackend):
return ( return (
kv_cache is not None fwd in ("prefill", "decode")
and kv_cache is not None
and q.ndim == 3
and q.dtype == torch.bfloat16 and q.dtype == torch.bfloat16
and q.size(-1) in (32, 64, 128, 256) and q.size(-1) in (32, 64, 128, 256)
and is_available(f"attn_paged_{fwd}")
) )
if isinstance(backend, FlashAttnBackend): if isinstance(backend, FlashAttnBackend):
if not flash_attn_available(): if not flash_attn_available():
return False return False
if q.dtype not in (torch.float16, torch.bfloat16): if q.dtype not in (torch.float16, torch.bfloat16):
return False return False
if q.size(1) == 1 and kv_cache is not None: if fwd is not None:
return True return q.ndim == 3 and hasattr(_get_flash_attn(), "flash_attn_varlen_func")
if attn_mask is None or is_causal: if attn_mask is None or is_causal:
return True return True
return attn_mask.dim() == 4 return attn_mask.dim() == 4
@@ -243,13 +247,13 @@ def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
def repeat_kv(x: Tensor, n_rep: int) -> Tensor: def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
"""Expand KV heads to match Q heads for GQA.""" """Expand KV heads to match Q heads for GQA."""
bs, slen, n_heads, head_dim = x.shape
if n_rep == 1: if n_rep == 1:
return x return x
n_heads, head_dim = x.shape[-2:]
return ( return (
x[:, :, :, None, :] x.unsqueeze(-2)
.expand(bs, slen, n_heads, n_rep, head_dim) .expand(*x.shape[:-2], n_heads, n_rep, head_dim)
.reshape(bs, slen, n_heads * n_rep, head_dim) .reshape(*x.shape[:-2], n_heads * n_rep, head_dim)
) )
@@ -283,6 +287,7 @@ def attention(
layer_id: int = 0, layer_id: int = 0,
attn_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None,
is_causal: bool = False, is_causal: bool = False,
fwd: Optional[str] = None,
) -> Tensor: ) -> Tensor:
"""Functional attention entry point — mirrors ``F.scaled_dot_product_attention``. """Functional attention entry point — mirrors ``F.scaled_dot_product_attention``.
@@ -302,9 +307,11 @@ def attention(
Returns: Returns:
[batch, q_len, n_heads * head_dim] [batch, q_len, n_heads * head_dim]
""" """
explicit = get_backend(use_default=False)
backend = get_backend() backend = get_backend()
if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal): if fwd is None and explicit is None:
explicit = get_backend(use_default=False) backend = TorchNativeBackend()
if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal, fwd):
if explicit is not None: if explicit is not None:
raise RuntimeError( raise RuntimeError(
f"Explicitly-set backend {type(backend).__name__} cannot " f"Explicitly-set backend {type(backend).__name__} cannot "
@@ -316,10 +323,10 @@ def attention(
for candidate in _priority_backends(): for candidate in _priority_backends():
if isinstance(candidate, type(backend)): if isinstance(candidate, type(backend)):
continue continue
if _backend_supports(candidate, q, kv_cache, attn_mask, is_causal): if _backend_supports(candidate, q, kv_cache, attn_mask, is_causal, fwd):
backend = candidate backend = candidate
break break
return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal) return backend.forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal, fwd)
class AttentionBackend(ABC): class AttentionBackend(ABC):
@@ -355,6 +362,7 @@ class AttentionBackend(ABC):
layer_id: int, layer_id: int,
attn_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None,
is_causal: bool = False, is_causal: bool = False,
fwd: Optional[str] = None,
) -> Tensor: ) -> Tensor:
"""Dispatch to decode or extend based on q_len. """Dispatch to decode or extend based on q_len.
@@ -370,9 +378,11 @@ class AttentionBackend(ABC):
Returns: Returns:
[batch, q_len, n_heads * head_dim] [batch, q_len, n_heads * head_dim]
""" """
if kv_cache is not None and q.size(1) == 1: if fwd == "decode":
return self.fwd_decode(q, k, v, kv_cache, layer_id, attn_mask, is_causal) return self.fwd_decode(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
return self.fwd_prefill(q, k, v, kv_cache, layer_id, attn_mask, is_causal) if fwd == "prefill" or fwd is None:
return self.fwd_prefill(q, k, v, kv_cache, layer_id, attn_mask, is_causal)
raise ValueError(f"unsupported attention forward mode: {fwd}")
@abstractmethod @abstractmethod
def fwd_decode( def fwd_decode(
@@ -466,23 +476,52 @@ class TorchNativeBackend(AttentionBackend):
attn_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None,
is_causal: bool = False, is_causal: bool = False,
) -> Tensor: ) -> Tensor:
if kv_cache is not None: if q.ndim == 4:
k, v = _write_and_gather_kv(kv_cache, k, v, layer_id, q, attn_mask) n_rep = q.size(2) // k.size(2)
if n_rep > 1:
k = repeat_kv(k, n_rep)
v = repeat_kv(v, n_rep)
return (
F.scaled_dot_product_attention(
q.permute(0, 2, 1, 3),
k.permute(0, 2, 1, 3),
v.permute(0, 2, 1, 3),
attn_mask,
is_causal=is_causal,
)
.permute(0, 2, 1, 3)
.contiguous()
)
n_rep = q.size(2) // k.size(2) if kv_cache is None or kv_cache.qo_indptr is None:
if n_rep > 1: raise ValueError("packed attention requires KV cache metadata")
k = repeat_kv(k, n_rep) kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
v = repeat_kv(v, n_rep) kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
outputs = []
out = F.scaled_dot_product_attention( n_rep = q.size(1) // k.size(1)
q.permute(0, 2, 1, 3), for i in range(kv_cache.req_pool_indices.numel()):
k.permute(0, 2, 1, 3), q_start = int(kv_cache.qo_indptr[i])
v.permute(0, 2, 1, 3), q_end = int(kv_cache.qo_indptr[i + 1])
attn_mask, indices = kv_cache.req_to_token[
is_causal=is_causal, kv_cache.req_pool_indices[i], : kv_cache.seq_lens[i]
) ]
out = out.permute(0, 2, 1, 3).contiguous().flatten(2) k_i = kv_cache.k_buffer[layer_id, indices]
return out v_i = kv_cache.v_buffer[layer_id, indices]
if n_rep > 1:
k_i = repeat_kv(k_i, n_rep)
v_i = repeat_kv(v_i, n_rep)
q_len = q_end - q_start
kv_len = k_i.size(0)
q_pos = torch.arange(kv_len - q_len, kv_len, device=q.device)
causal_mask = q_pos[:, None] >= torch.arange(kv_len, device=q.device)
out = F.scaled_dot_product_attention(
q[q_start:q_end].transpose(0, 1).unsqueeze(0),
k_i.transpose(0, 1).unsqueeze(0),
v_i.transpose(0, 1).unsqueeze(0),
attn_mask=causal_mask,
)
outputs.append(out.squeeze(0).transpose(0, 1))
return torch.cat(outputs)
@AttentionBackendFactory.register(ATTN_BACKEND.CUDA.value) @AttentionBackendFactory.register(ATTN_BACKEND.CUDA.value)
@@ -530,16 +569,14 @@ class CudaBackend(AttentionBackend):
if kv_cache is None: if kv_cache is None:
raise RuntimeError("CudaBackend does not support training (kv_cache=None)") raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
loc = kv_cache.out_cache_loc[:, 0] loc = kv_cache.out_cache_loc
kv_cache.k_buffer[layer_id, loc] = k[:, 0] kv_cache.k_buffer[layer_id, loc] = k
kv_cache.v_buffer[layer_id, loc] = v[:, 0] kv_cache.v_buffer[layer_id, loc] = v
q_3d = q.squeeze(1)
kv_indptr = kv_cache.kv_indptr kv_indptr = kv_cache.kv_indptr
out = attn_paged_decode( out = attn_paged_decode(
q_3d, q,
kv_cache.k_buffer[layer_id], kv_cache.k_buffer[layer_id],
kv_cache.v_buffer[layer_id], kv_cache.v_buffer[layer_id],
kv_cache.req_to_token, kv_cache.req_to_token,
@@ -550,7 +587,7 @@ class CudaBackend(AttentionBackend):
ml_part_buf=kv_cache.decode_ml_part, ml_part_buf=kv_cache.decode_ml_part,
out_buf=kv_cache.decode_out, out_buf=kv_cache.decode_out,
) )
return out.unsqueeze(1).flatten(2) return out
def fwd_prefill( def fwd_prefill(
self, self,
@@ -565,30 +602,22 @@ class CudaBackend(AttentionBackend):
if kv_cache is None: if kv_cache is None:
raise RuntimeError("CudaBackend does not support training (kv_cache=None)") raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
loc = kv_cache.out_cache_loc.reshape(-1) loc = kv_cache.out_cache_loc
kv_cache.k_buffer[layer_id, loc] = k.reshape(-1, k.size(2), k.size(3)) kv_cache.k_buffer[layer_id, loc] = k
kv_cache.v_buffer[layer_id, loc] = v.reshape(-1, v.size(2), v.size(3)) kv_cache.v_buffer[layer_id, loc] = v
b = q.size(0)
q_len = q.size(1)
kv_indptr = kv_cache.kv_indptr
qo_indptr = kv_cache.qo_indptr
q_flat = q.reshape(b * q_len, q.size(2), q.size(3))
out = attn_paged_prefill( out = attn_paged_prefill(
q_flat, q,
kv_cache.k_buffer[layer_id], kv_cache.k_buffer[layer_id],
kv_cache.v_buffer[layer_id], kv_cache.v_buffer[layer_id],
kv_cache.req_to_token, kv_cache.req_to_token,
kv_cache.req_pool_indices, kv_cache.req_pool_indices,
kv_indptr, kv_cache.kv_indptr,
qo_indptr, kv_cache.qo_indptr,
attn_mask, attn_mask,
is_causal=is_causal, is_causal=is_causal,
) )
return out.reshape(b, q_len, q.size(2), q.size(3)).flatten(2) return out
@AttentionBackendFactory.register(ATTN_BACKEND.FLASH.value) @AttentionBackendFactory.register(ATTN_BACKEND.FLASH.value)
@@ -617,7 +646,7 @@ class FlashAttnBackend(AttentionBackend):
attn_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None,
is_causal: bool = False, is_causal: bool = False,
) -> Tensor: ) -> Tensor:
return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal) return self._forward_packed(q, k, v, kv_cache, layer_id)
def fwd_prefill( def fwd_prefill(
self, self,
@@ -629,25 +658,18 @@ class FlashAttnBackend(AttentionBackend):
attn_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None,
is_causal: bool = False, is_causal: bool = False,
) -> Tensor: ) -> Tensor:
return self._forward(q, k, v, kv_cache, layer_id, attn_mask, is_causal) if q.ndim == 3:
return self._forward_packed(q, k, v, kv_cache, layer_id)
return self._forward_dense(q, k, v, attn_mask, is_causal)
def _forward( def _forward_dense(
self, self,
q: Tensor, q: Tensor,
k: Tensor, k: Tensor,
v: Tensor, v: Tensor,
kv_cache: Optional["KVCache"],
layer_id: int,
attn_mask: Optional[Tensor] = None, attn_mask: Optional[Tensor] = None,
is_causal: bool = False, is_causal: bool = False,
) -> Tensor: ) -> Tensor:
if kv_cache is not None:
if q.size(1) == 1 and kv_cache.k_buffer.size(
1
) == kv_cache.req_to_token.size(0) * kv_cache.req_to_token.size(1):
return self._decode_with_kvcache(q, k, v, kv_cache, layer_id)
k, v = _write_and_gather_kv(kv_cache, k, v, layer_id, q, attn_mask)
n_rep = q.size(2) // k.size(2) n_rep = q.size(2) // k.size(2)
if n_rep > 1: if n_rep > 1:
k = repeat_kv(k, n_rep) k = repeat_kv(k, n_rep)
@@ -670,9 +692,9 @@ class FlashAttnBackend(AttentionBackend):
v.contiguous(), v.contiguous(),
causal=is_causal or (attn_mask is not None and attn_mask.dim() == 4), causal=is_causal or (attn_mask is not None and attn_mask.dim() == 4),
) )
return out.contiguous().flatten(2) return out.contiguous()
def _decode_with_kvcache( def _forward_packed(
self, self,
q: Tensor, q: Tensor,
k: Tensor, k: Tensor,
@@ -680,22 +702,27 @@ class FlashAttnBackend(AttentionBackend):
kv_cache: "KVCache", kv_cache: "KVCache",
layer_id: int, layer_id: int,
) -> Tensor: ) -> Tensor:
max_batch = kv_cache.req_to_token.size(0)
max_seq = kv_cache.req_to_token.size(1)
n_kv = k.size(2)
k_cache = kv_cache.k_buffer[layer_id].view(max_batch, max_seq, n_kv, k.size(3))
v_cache = kv_cache.v_buffer[layer_id].view(max_batch, max_seq, n_kv, v.size(3))
fa = _get_flash_attn() fa = _get_flash_attn()
out = fa.flash_attn_with_kvcache( if fa is None or not hasattr(fa, "flash_attn_varlen_func"):
q=q, raise RuntimeError("packed inference requires flash_attn_varlen_func")
k_cache=k_cache, kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
v_cache=v_cache, kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
k=k, page_table = kv_cache.req_to_token[
v=v, kv_cache.req_pool_indices, : kv_cache.max_len
cache_seqlens=(kv_cache.seq_lens - 1).to(torch.int32), ]
cache_batch_idx=kv_cache.req_pool_indices.to(torch.int32), positions = torch.arange(kv_cache.max_len, device=q.device)
indices = page_table[positions.unsqueeze(0) < kv_cache.seq_lens.unsqueeze(1)]
k_flat = kv_cache.k_buffer[layer_id, indices].contiguous()
v_flat = kv_cache.v_buffer[layer_id, indices].contiguous()
out = fa.flash_attn_varlen_func(
q.contiguous(),
k_flat,
v_flat,
kv_cache.qo_indptr,
kv_cache.kv_indptr,
int((kv_cache.qo_indptr[1:] - kv_cache.qo_indptr[:-1]).max()),
int(kv_cache.seq_lens.max()),
dropout_p=0.0,
causal=True, causal=True,
) )
return out.flatten(2) return out
astrai/extension/rotary_backend.py
+1 -1
View File
@@ -26,7 +26,7 @@ def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
dtype = x.dtype dtype = x.dtype
x_ = x.float().reshape(*x.shape[:-1], -1, 2) x_ = x.float().reshape(*x.shape[:-1], -1, 2)
x_complex = torch.view_as_complex(x_) x_complex = torch.view_as_complex(x_)
freqs_cis_complex = torch.complex(cos, sin).unsqueeze(2) freqs_cis_complex = torch.complex(cos, sin).unsqueeze(-2)
x_rotated = x_complex * freqs_cis_complex x_rotated = x_complex * freqs_cis_complex
x_out = torch.view_as_real(x_rotated).flatten(-2) x_out = torch.view_as_real(x_rotated).flatten(-2)
return x_out.to(dtype) return x_out.to(dtype)
astrai/extension/rotary_ops.py
+5 -5
View File
@@ -4,8 +4,8 @@ Calls the compiled CUDA kernel directly. If the kernel is not available,
raises ``RuntimeError``. Fallback to torch complex multiply is the raises ``RuntimeError``. Fallback to torch complex multiply is the
responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``. responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``.
Layout: x is [batch, seq_len, n_heads, head_dim] (bf16, contiguous). Layout: x is packed [tokens, n_heads, head_dim] or dense
freqs_cis is [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs. [batch, seq_len, n_heads, head_dim]. ``freqs_cis`` has matching token axes.
""" """
import torch import torch
@@ -25,11 +25,11 @@ def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
"""Fused rotary embedding kernel. """Fused rotary embedding kernel.
Args: Args:
x: [batch, seq_len, n_heads, head_dim] (bf16, contiguous) x: packed 3D or dense 4D bf16 tensor.
freqs_cis: [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs freqs_cis: matching token axes followed by [head_dim/2, 2].
Returns: Returns:
[batch, seq_len, n_heads, head_dim] (bf16) Tensor with the same shape as ``x``.
""" """
_check_available() _check_available()
if not x.is_contiguous(): if not x.is_contiguous():
astrai/inference/cache/pool.py
Vendored
+18 -10
View File
@@ -200,14 +200,21 @@ class PagePool:
kv_indptr = kvp_buf[: b + 1] kv_indptr = kvp_buf[: b + 1]
if start_pos is not None: if start_pos is not None:
# ---- prefill: out_cache_loc covers prefix range [start_pos:seq_len] ---- # Packed prefill concatenates each request's query tokens.
seq_len = seq_lens[0] q_lens = [seq_len - start_pos for seq_len in seq_lens]
out_cache_loc = self._req_pool.req_to_token[ if any(q_len <= 0 for q_len in q_lens):
req_pool_indices, start_pos:seq_len raise ValueError("prefill sequence lengths must exceed start_pos")
] out_cache_loc = torch.cat(
q_len = seq_len - start_pos [
workspace.qo_indptr[: b + 1].copy_( self._req_pool.req_to_token[
torch.arange(b + 1, dtype=torch.int32, device=device) * q_len req_pool_indices[i], start_pos : seq_lens[i]
]
for i in range(b)
]
)
workspace.qo_indptr[: b + 1].zero_()
workspace.qo_indptr[1 : b + 1].copy_(
torch.tensor(q_lens, dtype=torch.int32, device=device).cumsum(0)
) )
qo_indptr = workspace.qo_indptr[: b + 1] qo_indptr = workspace.qo_indptr[: b + 1]
decode_o_part = decode_ml_part = decode_out = None decode_o_part = decode_ml_part = decode_out = None
@@ -216,8 +223,9 @@ class PagePool:
write_pos = seq_lens_t - 1 write_pos = seq_lens_t - 1
loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1) loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
ocl_buf[:b].copy_(loc) ocl_buf[:b].copy_(loc)
out_cache_loc = ocl_buf[:b] out_cache_loc = ocl_buf[:b].reshape(-1)
qo_indptr = None workspace.qo_indptr[: b + 1].copy_(inc_buf[: b + 1])
qo_indptr = workspace.qo_indptr[: b + 1]
decode_o_part = getattr(workspace, "decode_o_part", None) decode_o_part = getattr(workspace, "decode_o_part", None)
decode_ml_part = getattr(workspace, "decode_ml_part", None) decode_ml_part = getattr(workspace, "decode_ml_part", None)
decode_out = getattr(workspace, "decode_out", None) decode_out = getattr(workspace, "decode_out", None)
astrai/inference/runtime/executor.py
+21 -30
View File
@@ -118,13 +118,13 @@ def _warmup_cuda_graphs(
timed("warmup prefill", logger), timed("warmup prefill", logger),
): ):
kv = task_cache.bind([tid], ws, start_pos=0) kv = task_cache.bind([tid], ws, start_pos=0)
ids_in = torch.arange(warmup_len, device=dev).unsqueeze(0) ids_in = torch.arange(warmup_len, device=dev)
pos_in = ids_in pos_in = ids_in
model( model(
ids_in, ids_in,
input_mask=pos_in.unsqueeze(-1) >= torch.arange(warmup_len, device=dev),
kv_cache=kv, kv_cache=kv,
position_ids=pos_in, position_ids=pos_in,
fwd="prefill",
) )
task_cache.task_free(tid) task_cache.task_free(tid)
@@ -159,15 +159,14 @@ def _warmup_cuda_graphs(
for tid in task_ids: for tid in task_ids:
task_cache.task_extend(tid, seq_pos) task_cache.task_extend(tid, seq_pos)
kv = task_cache.bind(task_ids, ws) kv = task_cache.bind(task_ids, ws)
input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
ids_buf = ws.fill_input_ids([step] * b) ids_buf = ws.fill_input_ids([step] * b)
gctx.forward( gctx.forward(
model, model,
key=(b,), key=(b,),
input_ids=ids_buf.unsqueeze(1), input_ids=ids_buf,
input_mask=input_mask,
kv_cache=kv, kv_cache=kv,
position_ids=ws.position_ids[:b].unsqueeze(1), position_ids=ws.position_ids[:b],
fwd="decode",
) )
for tid in task_ids: for tid in task_ids:
@@ -308,20 +307,15 @@ class Executor:
batch_sz = len(tasks) batch_sz = len(tasks)
input_ids = torch.tensor( input_ids = torch.tensor(
[t.prompt_ids[start_pos:prompt_len] for t in tasks], [token for t in tasks for token in t.prompt_ids[start_pos:prompt_len]],
dtype=torch.long, dtype=torch.long,
device=self.device, device=self.device,
) )
task_ids = [t.task_id for t in tasks] task_ids = [t.task_id for t in tasks]
position_ids = ( position_ids = torch.arange(
torch.arange(start_pos, prompt_len, dtype=torch.long, device=self.device) start_pos, prompt_len, dtype=torch.long, device=self.device
.unsqueeze(0) ).repeat(batch_sz)
.expand(batch_sz, -1)
)
input_mask = position_ids.unsqueeze(-1) >= torch.arange(
prompt_len, device=self.device
)
with ( with (
torch.inference_mode(), torch.inference_mode(),
@@ -329,15 +323,18 @@ class Executor:
): ):
outputs = self.model( outputs = self.model(
input_ids, input_ids,
input_mask=input_mask,
position_ids=position_ids, position_ids=position_ids,
kv_cache=self.task_cache.bind( kv_cache=self.task_cache.bind(
task_ids, task_ids,
self._workspace, self._workspace,
start_pos=start_pos, start_pos=start_pos,
), ),
fwd="prefill",
) )
logits = outputs["logits"][:, -1, :] q_len = prompt_len - start_pos
logits = outputs["logits"][
torch.arange(1, batch_sz + 1, device=self.device) * q_len - 1
]
return tasks, self._sample_logits(logits, tasks, return_logprobs) return tasks, self._sample_logits(logits, tasks, return_logprobs)
@@ -391,9 +388,6 @@ class Executor:
) )
self._decode_cache = DecodeSteadyState(task_sig, cur_positions, info) self._decode_cache = DecodeSteadyState(task_sig, cur_positions, info)
total_len = max(cur_positions) + 1
input_mask = ws.decode_mask(ws.position_ids[:b], total_len)
# ---- forward (graph replay or live run + capture) ---- # ---- forward (graph replay or live run + capture) ----
use_graph = ( use_graph = (
@@ -402,9 +396,6 @@ class Executor:
and get_backend().supports_graph() and get_backend().supports_graph()
) )
key = (b,) key = (b,)
if use_graph:
input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
with ( with (
torch.inference_mode(), torch.inference_mode(),
timed(f"execute_decode forward b={b}", logger), timed(f"execute_decode forward b={b}", logger),
@@ -413,18 +404,18 @@ class Executor:
outputs = self._graph_ctx.forward( outputs = self._graph_ctx.forward(
self.model, self.model,
key=key, key=key,
input_ids=input_ids.unsqueeze(1), input_ids=input_ids,
input_mask=input_mask,
kv_cache=kv_cache, kv_cache=kv_cache,
position_ids=ws.position_ids[:b].unsqueeze(1), position_ids=ws.position_ids[:b],
fwd="decode",
) )
else: else:
outputs = self.model( outputs = self.model(
input_ids.unsqueeze(1), input_ids,
input_mask=input_mask,
kv_cache=kv_cache, kv_cache=kv_cache,
position_ids=ws.position_ids[:b].unsqueeze(1), position_ids=ws.position_ids[:b],
fwd="decode",
) )
logits = outputs["logits"][:, -1, :] logits = outputs["logits"]
return self._sample_logits(logits, tasks, return_logprobs, info=info) return self._sample_logits(logits, tasks, return_logprobs, info=info)
astrai/model/components/attention.py
+11 -9
View File
@@ -56,9 +56,7 @@ class GQA(nn.Module):
self.gate = Linear(dim, dim) self.gate = Linear(dim, dim)
def _split_heads(self, x: Tensor, n_heads) -> Tensor: def _split_heads(self, x: Tensor, n_heads) -> Tensor:
batch_size, seq_len, _ = x.shape return x.reshape(*x.shape[:-1], n_heads, self.head_dim)
x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
return x
def forward( def forward(
self, self,
@@ -67,6 +65,7 @@ class GQA(nn.Module):
attn_mask: Tensor = None, attn_mask: Tensor = None,
kv_cache: Optional[KVCache] = None, kv_cache: Optional[KVCache] = None,
is_causal: bool = False, is_causal: bool = False,
fwd: Optional[str] = None,
) -> Tensor: ) -> Tensor:
q = self._split_heads(self.q_proj(x), self.n_heads) q = self._split_heads(self.q_proj(x), self.n_heads)
k = self._split_heads(self.k_proj(x), self.n_kv_heads) k = self._split_heads(self.k_proj(x), self.n_kv_heads)
@@ -76,7 +75,9 @@ class GQA(nn.Module):
if self.use_qk_norm: if self.use_qk_norm:
q, k = self.q_norm(q), self.k_norm(k) q, k = self.q_norm(q), self.k_norm(k)
sdqa_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal) sdqa_out = attention(
q, k, v, kv_cache, self.layer_id, attn_mask, is_causal, fwd
).reshape(*x.shape[:-1], self.dim)
if self.use_gated_attention: if self.use_gated_attention:
sdqa_out = sdqa_out * F.sigmoid(self.gate(x)) sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
@@ -141,17 +142,16 @@ class MLA(nn.Module):
attn_mask: Tensor = None, attn_mask: Tensor = None,
kv_cache: Optional[KVCache] = None, kv_cache: Optional[KVCache] = None,
is_causal: bool = False, is_causal: bool = False,
fwd: Optional[str] = None,
) -> Tensor: ) -> Tensor:
bsz, seq_len, _ = x.size()
q = self.q_proj(x) q = self.q_proj(x)
q = q.view(bsz, seq_len, self.n_heads, self.head_dim) q = q.reshape(*x.shape[:-1], self.n_heads, self.head_dim)
kv_compressed = self.kv_a_proj(x) kv_compressed = self.kv_a_proj(x)
kv_compressed = self.kv_norm(kv_compressed) kv_compressed = self.kv_norm(kv_compressed)
kv = self.kv_b_proj(kv_compressed) kv = self.kv_b_proj(kv_compressed)
kv = kv.view(bsz, seq_len, self.n_kv_heads, -1) kv = kv.reshape(*x.shape[:-1], self.n_kv_heads, -1)
k_nope, k_rope, v = torch.split( k_nope, k_rope, v = torch.split(
kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1 kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1
@@ -171,7 +171,9 @@ class MLA(nn.Module):
q = self.q_norm(q) q = self.q_norm(q)
k = self.k_norm(k) k = self.k_norm(k)
attn_out = attention(q, k, v, kv_cache, self.layer_id, attn_mask, is_causal) attn_out = attention(
q, k, v, kv_cache, self.layer_id, attn_mask, is_causal, fwd
).reshape(*x.shape[:-1], self.dim)
if self.use_gated_attention: if self.use_gated_attention:
attn_out = attn_out * F.sigmoid(self.gate(x)) attn_out = attn_out * F.sigmoid(self.gate(x))
astrai/model/components/decoder_block.py
+2
View File
@@ -54,6 +54,7 @@ class DecoderBlock(nn.Module):
attention_mask: Optional[Tensor] = None, attention_mask: Optional[Tensor] = None,
kv_cache: Optional[KVCache] = None, kv_cache: Optional[KVCache] = None,
is_causal: bool = False, is_causal: bool = False,
fwd: Optional[str] = None,
) -> DecoderOutput: ) -> DecoderOutput:
attn_output = self.attention( attn_output = self.attention(
self.input_norm(x), self.input_norm(x),
@@ -61,6 +62,7 @@ class DecoderBlock(nn.Module):
attention_mask, attention_mask,
kv_cache, kv_cache,
is_causal, is_causal,
fwd,
) )
x = attn_output + x x = attn_output + x
normalized = self.post_attention_norm(x) normalized = self.post_attention_norm(x)
astrai/model/components/mlp.py
+3 -2
View File
@@ -100,13 +100,14 @@ class DeepSeekMoE(nn.Module):
def forward(self, x: Tensor) -> FFNOutput: def forward(self, x: Tensor) -> FFNOutput:
include_aux_loss = self.training and torch.is_grad_enabled() include_aux_loss = self.training and torch.is_grad_enabled()
bsz, seq_len, dim = x.shape shape = x.shape
dim = shape[-1]
x_flat = x.view(-1, dim) x_flat = x.view(-1, dim)
shared_out = self._shared_forward(x_flat) shared_out = self._shared_forward(x_flat)
routed_output = self._routed_forward(x_flat, include_aux_loss) routed_output = self._routed_forward(x_flat, include_aux_loss)
out = (shared_out + routed_output["hidden_states"]).view(bsz, seq_len, dim) out = (shared_out + routed_output["hidden_states"]).view(shape)
return { return {
"hidden_states": out, "hidden_states": out,
"aux_loss": routed_output["aux_loss"], "aux_loss": routed_output["aux_loss"],
astrai/model/components/rope.py
+8 -5
View File
@@ -65,9 +65,12 @@ class RotaryEmbedding(nn.Module):
[batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs. [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
""" """
if position_ids is None: if position_ids is None:
position_ids = ( if x.ndim == 2:
torch.arange(x.size(1), device=x.device) position_ids = torch.arange(x.size(0), device=x.device)
.unsqueeze(0) else:
.expand(x.size(0), -1) position_ids = (
) torch.arange(x.size(1), device=x.device)
.unsqueeze(0)
.expand(x.size(0), -1)
)
return self.freqs_cis[position_ids].float() return self.freqs_cis[position_ids].float()
astrai/model/transformer.py
+14 -1
View File
@@ -105,8 +105,20 @@ class AutoRegressiveLM(AutoModel):
input_mask: Optional[Tensor] = None, input_mask: Optional[Tensor] = None,
kv_cache: Optional[KVCache] = None, kv_cache: Optional[KVCache] = None,
position_ids: Optional[Tensor] = None, position_ids: Optional[Tensor] = None,
fwd: Optional[str] = None,
) -> Dict[str, Tensor]: ) -> Dict[str, Tensor]:
assert input_ids.ndim == 2 if fwd is None:
if input_ids.ndim != 2:
raise ValueError("training input_ids must be [batch, seq_len]")
if kv_cache is not None:
raise ValueError("training forward does not accept a KV cache")
elif fwd in ("prefill", "decode"):
if input_ids.ndim != 1:
raise ValueError("inference input_ids must be packed [tokens]")
if kv_cache is None:
raise ValueError("inference forward requires a KV cache")
else:
raise ValueError(f"unsupported forward mode: {fwd}")
x = self.embed_tokens(input_ids) x = self.embed_tokens(input_ids)
rotary_emb = self.rotary_embedding(x, position_ids) rotary_emb = self.rotary_embedding(x, position_ids)
@@ -122,6 +134,7 @@ class AutoRegressiveLM(AutoModel):
attn_mask, attn_mask,
kv_cache, kv_cache,
use_sdpa_causal_mask, use_sdpa_causal_mask,
fwd,
) )
x = layer_output["hidden_states"] x = layer_output["hidden_states"]
stats = layer_output.get("router_stats") stats = layer_output.get("router_stats")
csrc/kernels/rotary_emb.cu
+19 -20
View File
@@ -7,13 +7,12 @@ __global__ void rotary_emb_kernel(
const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ x,
const float* __restrict__ freqs_cis, const float* __restrict__ freqs_cis,
__nv_bfloat16* __restrict__ out, __nv_bfloat16* __restrict__ out,
int batch, int n_tokens,
int seq_len,
int n_heads, int n_heads,
int head_dim int head_dim
) { ) {
const int half_dim = head_dim >> 1; const int half_dim = head_dim >> 1;
const int total = batch * seq_len * n_heads * half_dim; const int total = n_tokens * n_heads * half_dim;
for (int idx = blockIdx.x * blockDim.x + threadIdx.x; for (int idx = blockIdx.x * blockDim.x + threadIdx.x;
idx < total; idx < total;
@@ -23,11 +22,10 @@ __global__ void rotary_emb_kernel(
int tmp = idx / half_dim; int tmp = idx / half_dim;
int head = tmp % n_heads; int head = tmp % n_heads;
tmp /= n_heads; tmp /= n_heads;
int seq = tmp % seq_len; int token = tmp;
int b = tmp / seq_len;
int x_offset = ((b * seq_len + seq) * n_heads + head) * head_dim + (pair << 1); int x_offset = (token * n_heads + head) * head_dim + (pair << 1);
int cs_offset = ((b * seq_len + seq) * half_dim + pair) * 2; int cs_offset = (token * half_dim + pair) * 2;
__nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset); __nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
float x_even = __bfloat162float(__low2bfloat16(x_pair)); float x_even = __bfloat162float(__low2bfloat16(x_pair));
@@ -54,27 +52,28 @@ torch::Tensor rotary_emb(
TORCH_CHECK(x.is_cuda(), "x must be on CUDA"); TORCH_CHECK(x.is_cuda(), "x must be on CUDA");
TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis must be on CUDA"); TORCH_CHECK(freqs_cis.is_cuda(), "freqs_cis must be on CUDA");
TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16"); TORCH_CHECK(x.scalar_type() == torch::kBFloat16, "x must be bf16");
TORCH_CHECK(x.dim() == 4, "x must be 4D [batch, seq_len, n_heads, head_dim]"); TORCH_CHECK(x.dim() == 3 || x.dim() == 4,
"x must be [tokens, n_heads, head_dim] or "
"[batch, seq_len, n_heads, head_dim]");
TORCH_CHECK(x.is_contiguous(), "x must be contiguous"); TORCH_CHECK(x.is_contiguous(), "x must be contiguous");
TORCH_CHECK(freqs_cis.dim() == 4, "freqs_cis must be 4D [batch, seq_len, dim/2, 2]"); TORCH_CHECK(freqs_cis.dim() == x.dim(), "freqs_cis rank must match x rank");
TORCH_CHECK(freqs_cis.is_contiguous(), "freqs_cis must be contiguous"); TORCH_CHECK(freqs_cis.is_contiguous(), "freqs_cis must be contiguous");
TORCH_CHECK(freqs_cis.scalar_type() == torch::kFloat32, "freqs_cis must be f32"); TORCH_CHECK(freqs_cis.scalar_type() == torch::kFloat32, "freqs_cis must be f32");
int batch = x.size(0); int n_tokens = x.dim() == 3 ? x.size(0) : x.size(0) * x.size(1);
int seq_len = x.size(1); int n_heads = x.size(x.dim() - 2);
int n_heads = x.size(2); int head_dim = x.size(x.dim() - 1);
int head_dim = x.size(3);
TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even"); TORCH_CHECK(head_dim % 2 == 0, "head_dim must be even");
TORCH_CHECK(freqs_cis.size(0) == batch, "freqs_cis batch mismatch"); TORCH_CHECK(freqs_cis.numel() == (int64_t)n_tokens * head_dim,
TORCH_CHECK(freqs_cis.size(1) == seq_len, "freqs_cis seq_len mismatch"); "freqs_cis token or rotary dimension mismatch");
TORCH_CHECK(freqs_cis.size(2) == head_dim / 2, "freqs_cis dim/2 mismatch"); TORCH_CHECK(freqs_cis.size(-2) == head_dim / 2, "freqs_cis dim/2 mismatch");
TORCH_CHECK(freqs_cis.size(3) == 2, "freqs_cis last dim must be 2 [cos, sin]"); TORCH_CHECK(freqs_cis.size(-1) == 2, "freqs_cis last dim must be 2 [cos, sin]");
auto out = torch::empty_like(x); auto out = torch::empty_like(x);
int half_dim = head_dim / 2; int half_dim = head_dim / 2;
int total = batch * seq_len * n_heads * half_dim; int total = n_tokens * n_heads * half_dim;
int block = 256; int block = 256;
int grid = std::min((total + block - 1) / block, 1024); int grid = std::min((total + block - 1) / block, 1024);
@@ -82,7 +81,7 @@ torch::Tensor rotary_emb(
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()), reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
freqs_cis.data_ptr<float>(), freqs_cis.data_ptr<float>(),
reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
batch, seq_len, n_heads, head_dim n_tokens, n_heads, head_dim
); );
C10_CUDA_CHECK(cudaGetLastError()); C10_CUDA_CHECK(cudaGetLastError());
@@ -93,6 +92,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("rotary_emb", &rotary_emb, m.def("rotary_emb", &rotary_emb,
py::arg("x"), py::arg("x"),
py::arg("freqs_cis"), py::arg("freqs_cis"),
"Fused rotary embedding (bf16 x, f32 freqs_cis [b,s,d/2,2], bf16 out)" "Fused rotary embedding for packed 3D or dense 4D tensors"
); );
} }
scripts/tools/benchmark.py
+21 -40
View File
@@ -118,16 +118,11 @@ class GenerationBenchmark:
workspace: InferenceWorkspace, workspace: InferenceWorkspace,
) -> list: ) -> list:
input_ids = torch.randint( input_ids = torch.randint(
0, self.config.vocab_size, (batch_size, prompt_len), device=self.device 0, self.config.vocab_size, (batch_size * prompt_len,), device=self.device
)
position_ids = (
torch.arange(0, prompt_len, dtype=torch.long, device=self.device)
.unsqueeze(0)
.expand(batch_size, -1)
)
input_mask = position_ids.unsqueeze(-1) >= torch.arange(
prompt_len, device=self.device
) )
position_ids = torch.arange(
prompt_len, dtype=torch.long, device=self.device
).repeat(batch_size)
task_ids = [f"bench_{i}" for i in range(batch_size)] task_ids = [f"bench_{i}" for i in range(batch_size)]
for tid in task_ids: for tid in task_ids:
@@ -137,9 +132,9 @@ class GenerationBenchmark:
with torch.inference_mode(), attn_backend(self.backend): with torch.inference_mode(), attn_backend(self.backend):
self.model( self.model(
input_ids, input_ids,
input_mask=input_mask,
kv_cache=kv_cache, kv_cache=kv_cache,
position_ids=position_ids, position_ids=position_ids,
fwd="prefill",
) )
torch.cuda.synchronize() torch.cuda.synchronize()
return task_ids return task_ids
@@ -154,24 +149,20 @@ class GenerationBenchmark:
): ):
batch_size = len(task_ids) batch_size = len(task_ids)
input_ids = torch.randint( input_ids = torch.randint(
0, self.config.vocab_size, (batch_size, 1), device=self.device 0, self.config.vocab_size, (batch_size,), device=self.device
) )
position_ids = torch.tensor( position_ids = torch.tensor(
[[seq_len] for _ in range(batch_size)], dtype=torch.long, device=self.device [seq_len] * batch_size, dtype=torch.long, device=self.device
) )
total_len = seq_len + 1
for tid in task_ids: for tid in task_ids:
task_cache.task_extend(tid, seq_len) task_cache.task_extend(tid, seq_len)
input_mask = position_ids[:, :, None] >= torch.arange(
total_len, device=self.device
)
kv_cache = task_cache.bind(task_ids, workspace, self.device) kv_cache = task_cache.bind(task_ids, workspace, self.device)
with torch.inference_mode(), attn_backend(self.backend): with torch.inference_mode(), attn_backend(self.backend):
self.model( self.model(
input_ids, input_ids,
input_mask=input_mask,
kv_cache=kv_cache, kv_cache=kv_cache,
position_ids=position_ids, position_ids=position_ids,
fwd="decode",
) )
def run_prefill_benchmark( def run_prefill_benchmark(
@@ -188,25 +179,23 @@ class GenerationBenchmark:
task_cache.task_alloc(tid, list(range(prompt_length))) task_cache.task_alloc(tid, list(range(prompt_length)))
input_ids = torch.randint( input_ids = torch.randint(
0, self.config.vocab_size, (batch_size, prompt_length), device=self.device 0,
) self.config.vocab_size,
position_ids = ( (batch_size * prompt_length,),
torch.arange(0, prompt_length, dtype=torch.long, device=self.device) device=self.device,
.unsqueeze(0)
.expand(batch_size, -1)
)
input_mask = position_ids.unsqueeze(-1) >= torch.arange(
prompt_length, device=self.device
) )
position_ids = torch.arange(
prompt_length, dtype=torch.long, device=self.device
).repeat(batch_size)
kv_cache = task_cache.bind(task_ids, workspace, self.device, start_pos=0) kv_cache = task_cache.bind(task_ids, workspace, self.device, start_pos=0)
for _ in range(3): for _ in range(3):
with torch.inference_mode(), attn_backend(self.backend): with torch.inference_mode(), attn_backend(self.backend):
self.model( self.model(
input_ids, input_ids,
input_mask=input_mask,
kv_cache=kv_cache, kv_cache=kv_cache,
position_ids=position_ids, position_ids=position_ids,
fwd="prefill",
) )
torch.cuda.synchronize() torch.cuda.synchronize()
@@ -215,9 +204,9 @@ class GenerationBenchmark:
with torch.inference_mode(), attn_backend(self.backend): with torch.inference_mode(), attn_backend(self.backend):
self.model( self.model(
input_ids, input_ids,
input_mask=input_mask,
kv_cache=kv_cache, kv_cache=kv_cache,
position_ids=position_ids, position_ids=position_ids,
fwd="prefill",
) )
torch.cuda.synchronize() torch.cuda.synchronize()
elapsed = time.perf_counter() - t0 elapsed = time.perf_counter() - t0
@@ -311,37 +300,29 @@ class GenerationBenchmark:
) )
b = batch_size b = batch_size
input_ids_buf = torch.zeros(b, 1, dtype=torch.long, device=self.device) input_ids_buf = torch.zeros(b, dtype=torch.long, device=self.device)
position_ids_buf = torch.zeros(b, dtype=torch.long, device=self.device) position_ids_buf = torch.zeros(b, dtype=torch.long, device=self.device)
arange = torch.arange(max_seq_len, device=self.device)
gctx = CudaGraphContext(enabled=True) gctx = CudaGraphContext(enabled=True)
graph_key = (b,) graph_key = (b,)
def _decode_graph_step(seq_len): def _decode_graph_step(seq_len):
input_ids_buf.copy_( input_ids_buf.copy_(
torch.randint(0, self.config.vocab_size, (b, 1), device=self.device) torch.randint(0, self.config.vocab_size, (b,), device=self.device)
) )
position_ids_buf[:] = seq_len position_ids_buf[:] = seq_len
for tid in task_ids: for tid in task_ids:
task_cache.task_extend(tid, seq_len) task_cache.task_extend(tid, seq_len)
kv_cache = task_cache.bind(task_ids, workspace, self.device) kv_cache = task_cache.bind(task_ids, workspace, self.device)
input_mask = torch.ge(
position_ids_buf[:, None],
arange,
out=workspace.input_mask[:b, 0, :max_seq_len],
)
input_mask = input_mask.unsqueeze(1)
with torch.inference_mode(), attn_backend(self.backend): with torch.inference_mode(), attn_backend(self.backend):
return gctx.forward( return gctx.forward(
self.model, self.model,
key=graph_key, key=graph_key,
input_ids=input_ids_buf, input_ids=input_ids_buf,
input_mask=input_mask,
kv_cache=kv_cache, kv_cache=kv_cache,
position_ids=position_ids_buf.unsqueeze(1), position_ids=position_ids_buf,
fwd="decode",
) )
for i in range(5): for i in range(5):
tests/extension/test_backend_equivalence.py
+22 -42
View File
@@ -59,17 +59,9 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
"""Inference prefill with KV cache should match torch backend.""" """Inference prefill with KV cache should match torch backend."""
model, _ = cuda_model model, _ = cuda_model
prompt_ids = [[1, 2, 3, 4, 5, 6, 7, 8], [10, 11, 12, 13, 14, 15]] 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" device = "cuda"
input_ids = torch.zeros(batch, max_len, dtype=torch.long, device=device) input_ids = torch.tensor(sum(prompt_ids, []), dtype=torch.long, device=device)
input_mask = torch.zeros(batch, max_len, dtype=torch.bool, device=device) position_ids = torch.cat([torch.arange(len(p), device=device) for p in prompt_ids])
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( cache = PagePool(
n_layers=2, n_layers=2,
@@ -88,7 +80,7 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
kv1 = task_cache.bind(["t1", "t2"], ws, start_pos=0) kv1 = task_cache.bind(["t1", "t2"], ws, start_pos=0)
with torch.inference_mode(): with torch.inference_mode():
out_torch = model( out_torch = model(
input_ids, input_mask=input_mask, kv_cache=kv1, position_ids=position_ids input_ids, kv_cache=kv1, position_ids=position_ids, fwd="prefill"
) )
task_cache.task_free("t1") task_cache.task_free("t1")
@@ -100,22 +92,24 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
with torch.inference_mode(): with torch.inference_mode():
out_cuda = model( out_cuda = model(
input_ids, input_ids,
input_mask=input_mask,
kv_cache=kv2, kv_cache=kv2,
position_ids=position_ids, position_ids=position_ids,
fwd="prefill",
) )
offset = 0
for i, p in enumerate(prompt_ids): for i, p in enumerate(prompt_ids):
d = ( d = (
( (
out_torch["logits"][i, : len(p)].float() out_torch["logits"][offset : offset + len(p)].float()
- out_cuda["logits"][i, : len(p)].float() - out_cuda["logits"][offset : offset + len(p)].float()
) )
.abs() .abs()
.max() .max()
.item() .item()
) )
assert d == 0.0, f"Prefill diff for sample {i}: {d}" assert d == 0.0, f"Prefill diff for sample {i}: {d}"
offset += len(p)
@skip_no_kernel @skip_no_kernel
@@ -136,15 +130,8 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
) )
# Prefill to populate cache # Prefill to populate cache
max_len = max(len(p) for p in prompt_ids) input_ids = torch.tensor(sum(prompt_ids, []), dtype=torch.long, device=device)
batch = len(prompt_ids) position_ids = torch.cat([torch.arange(len(p), device=device) for p in 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)
task_cache = _mk_task_cache(cache) task_cache = _mk_task_cache(cache)
ws = _ws(cache) ws = _ws(cache)
@@ -152,28 +139,22 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
task_cache.task_alloc("t2", prompt_ids[1]) task_cache.task_alloc("t2", prompt_ids[1])
kv = task_cache.bind(["t1", "t2"], ws, start_pos=0) kv = task_cache.bind(["t1", "t2"], ws, start_pos=0)
with torch.inference_mode(): with torch.inference_mode():
model(input_ids, input_mask=input_mask, kv_cache=kv, position_ids=position_ids) model(input_ids, kv_cache=kv, position_ids=position_ids, fwd="prefill")
# Decode step — seq_lens are 9 and 7 (after extending) # Decode step — seq_lens are 9 and 7 (after extending)
dec_ids = torch.tensor([[99], [98]], dtype=torch.long, device=device) dec_ids = torch.tensor([99, 98], dtype=torch.long, device=device)
dec_pos = torch.tensor([[8], [6]], 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)
task_cache.task_extend("t1", 8) task_cache.task_extend("t1", 8)
task_cache.task_extend("t2", 6) task_cache.task_extend("t2", 6)
kv_t = task_cache.bind(["t1", "t2"], ws) kv_t = task_cache.bind(["t1", "t2"], ws)
with torch.inference_mode(): with torch.inference_mode():
out_torch = model( out_torch = model(dec_ids, kv_cache=kv_t, position_ids=dec_pos, fwd="decode")
dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
)
kv_c = task_cache.bind(["t1", "t2"], ws) kv_c = task_cache.bind(["t1", "t2"], ws)
with attn_backend(ATTN_BACKEND.CUDA): with attn_backend(ATTN_BACKEND.CUDA):
with torch.inference_mode(): with torch.inference_mode():
out_cuda = model( out_cuda = model(dec_ids, kv_cache=kv_c, position_ids=dec_pos, fwd="decode")
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() diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}" assert diff < 0.05, f"Decode diff (mixed seq_lens): {diff}"
@@ -198,16 +179,15 @@ def test_decode_cuda_graph_replay_is_exact(cuda_model):
ws = _ws(cache) ws = _ws(cache)
task_cache.task_alloc("t1", prompt_ids) task_cache.task_alloc("t1", prompt_ids)
input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device) input_ids = torch.tensor(prompt_ids, dtype=torch.long, device=device)
position_ids = torch.arange(len(prompt_ids), device=device).unsqueeze(0) position_ids = torch.arange(len(prompt_ids), device=device)
input_mask = torch.ones(1, len(prompt_ids), dtype=torch.bool, device=device)
with attn_backend(ATTN_BACKEND.CUDA), torch.inference_mode(): with attn_backend(ATTN_BACKEND.CUDA), torch.inference_mode():
model( model(
input_ids, input_ids,
input_mask=input_mask,
position_ids=position_ids, position_ids=position_ids,
kv_cache=task_cache.bind(["t1"], ws, start_pos=0), kv_cache=task_cache.bind(["t1"], ws, start_pos=0),
fwd="prefill",
) )
task_cache.task_extend("t1", len(prompt_ids)) task_cache.task_extend("t1", len(prompt_ids))
@@ -217,16 +197,16 @@ def test_decode_cuda_graph_replay_is_exact(cuda_model):
assert kv_cache.out_cache_loc.dtype == torch.int32 assert kv_cache.out_cache_loc.dtype == torch.int32
decode_args = { decode_args = {
"input_ids": torch.tensor([[9]], dtype=torch.long, device=device), "input_ids": torch.tensor([9], dtype=torch.long, device=device),
"input_mask": torch.ones(1, 1, 64, dtype=torch.bool, device=device), "position_ids": torch.tensor([len(prompt_ids)], device=device),
"position_ids": torch.tensor([[len(prompt_ids)]], device=device),
"kv_cache": kv_cache, "kv_cache": kv_cache,
"fwd": "decode",
} }
graph = CudaGraphContext(enabled=True) graph = CudaGraphContext(enabled=True)
graph.forward(model, key=(1,), **decode_args) graph.forward(model, key=(1,), **decode_args)
graph.forward(model, key=(1,), **decode_args) graph.forward(model, key=(1,), **decode_args)
first = graph.forward(model, key=(1,), **decode_args)["logits"].clone() first = graph.forward(model, key=(1,), **decode_args)["logits"].clone()
slot = kv_cache.out_cache_loc[0, 0] slot = kv_cache.out_cache_loc[0]
first_k = kv_cache.k_buffer[:, slot].clone() first_k = kv_cache.k_buffer[:, slot].clone()
first_v = kv_cache.v_buffer[:, slot].clone() first_v = kv_cache.v_buffer[:, slot].clone()
tests/inference/test_cache.py
+2 -2
View File
@@ -279,7 +279,7 @@ def test_page_pool_contiguous_bind_tasks_prefill():
task_cache.task_alloc("t1", list(range(10))) task_cache.task_alloc("t1", list(range(10)))
task_cache.task_alloc("t2", list(range(10))) task_cache.task_alloc("t2", list(range(10)))
kv = task_cache.bind(["t1", "t2"], _ws(pool), start_pos=0) kv = task_cache.bind(["t1", "t2"], _ws(pool), start_pos=0)
assert kv.out_cache_loc.shape == (2, 10) assert kv.out_cache_loc.shape == (20,)
assert kv.out_cache_loc.dtype == torch.int32 assert kv.out_cache_loc.dtype == torch.int32
assert kv.seq_lens.tolist() == [10, 10] assert kv.seq_lens.tolist() == [10, 10]
assert kv.req_pool_indices.shape == (2,) assert kv.req_pool_indices.shape == (2,)
@@ -295,7 +295,7 @@ def test_page_pool_contiguous_bind_tasks_decode():
assert task_cache.task_extend("t1", 10) assert task_cache.task_extend("t1", 10)
assert task_cache.task_extend("t2", 8) assert task_cache.task_extend("t2", 8)
kv = task_cache.bind(["t1", "t2"], _ws(pool)) kv = task_cache.bind(["t1", "t2"], _ws(pool))
assert kv.out_cache_loc.shape == (2, 1) assert kv.out_cache_loc.shape == (2,)
assert kv.seq_lens.tolist() == [11, 9] assert kv.seq_lens.tolist() == [11, 9]
tests/module/test_model_forward.py
+46
View File
@@ -39,6 +39,52 @@ def _make_model(config=None) -> AutoRegressiveLM:
return AutoRegressiveLM(config) return AutoRegressiveLM(config)
def test_model_forward_contract_uses_dense_training_and_packed_inference():
from astrai.inference.cache import PagePool, TaskCacheManager
from astrai.inference.workspace import InferenceWorkspace
config = AutoRegressiveLMConfig(**TINY_CONFIG)
model = AutoRegressiveLM(config).eval()
dense = model(torch.tensor([[1, 2, 3]]))
assert dense["logits"].shape == (1, 3, config.vocab_size)
pool = PagePool(
n_layers=config.num_hidden_layers,
n_kv_heads=config.num_key_value_heads,
head_dim=config.hidden_size // config.num_attention_heads,
max_batch_size=1,
max_seq_len=config.max_position_embeddings,
device="cpu",
dtype=torch.float32,
)
cache = TaskCacheManager(pool)
workspace = InferenceWorkspace(
1,
config.max_position_embeddings,
config.num_attention_heads,
config.hidden_size // config.num_attention_heads,
torch.device("cpu"),
torch.float32,
)
assert cache.task_alloc("t", [1, 2, 3])
packed = model(
torch.tensor([1, 2, 3]),
position_ids=torch.arange(3),
kv_cache=cache.bind(["t"], workspace, start_pos=0),
fwd="prefill",
)
assert packed["logits"].shape == (3, config.vocab_size)
with pytest.raises(ValueError, match="training input_ids"):
model(torch.tensor([1, 2, 3]))
with pytest.raises(ValueError, match="inference input_ids"):
model(
torch.tensor([[1, 2, 3]]),
kv_cache=cache.bind(["t"], workspace, start_pos=0),
fwd="prefill",
)
def _router_stats(probs, topk_indices): def _router_stats(probs, topk_indices):
return {"probs": probs, "topk_indices": topk_indices} return {"probs": probs, "topk_indices": topk_indices}