refactor: standardize packed 3d inference

- 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

astrai/extension/attention_backend.py

@@ -123,6 +123,7 @@ def _backend_supports(
     kv_cache: Optional["KVCache"],
     attn_mask: Optional[Tensor],
     is_causal: bool,
+    fwd: Optional[str],
 ) -> bool:
     """Whether ``backend`` can run this attention call.
 
@@ -131,17 +132,20 @@ def _backend_supports(
     """
     if isinstance(backend, CudaBackend):
         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.size(-1) in (32, 64, 128, 256)
+            and is_available(f"attn_paged_{fwd}")
         )
     if isinstance(backend, FlashAttnBackend):
         if not flash_attn_available():
             return False
         if q.dtype not in (torch.float16, torch.bfloat16):
             return False
-        if q.size(1) == 1 and kv_cache is not None:
-            return True
+        if fwd is not None:
+            return q.ndim == 3 and hasattr(_get_flash_attn(), "flash_attn_varlen_func")
         if attn_mask is None or is_causal:
             return True
         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:
     """Expand KV heads to match Q heads for GQA."""
-    bs, slen, n_heads, head_dim = x.shape
     if n_rep == 1:
         return x
+    n_heads, head_dim = x.shape[-2:]
     return (
-        x[:, :, :, None, :]
-        .expand(bs, slen, n_heads, n_rep, head_dim)
-        .reshape(bs, slen, n_heads * n_rep, head_dim)
+        x.unsqueeze(-2)
+        .expand(*x.shape[:-2], 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,
     attn_mask: Optional[Tensor] = None,
     is_causal: bool = False,
+    fwd: Optional[str] = None,
 ) -> Tensor:
     """Functional attention entry point — mirrors ``F.scaled_dot_product_attention``.
 
@@ -302,9 +307,11 @@ def attention(
     Returns:
         [batch, q_len, n_heads * head_dim]
     """
+    explicit = get_backend(use_default=False)
     backend = get_backend()
-    if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal):
-        explicit = get_backend(use_default=False)
+    if fwd is None and explicit is None:
+        backend = TorchNativeBackend()
+    if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal, fwd):
         if explicit is not None:
             raise RuntimeError(
                 f"Explicitly-set backend {type(backend).__name__} cannot "
@@ -316,10 +323,10 @@ def attention(
         for candidate in _priority_backends():
             if isinstance(candidate, type(backend)):
                 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
                 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):
@@ -355,6 +362,7 @@ class AttentionBackend(ABC):
         layer_id: int,
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> Tensor:
         """Dispatch to decode or extend based on q_len.
 
@@ -370,9 +378,11 @@ class AttentionBackend(ABC):
         Returns:
             [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_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
     def fwd_decode(
@@ -466,23 +476,52 @@ class TorchNativeBackend(AttentionBackend):
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
     ) -> Tensor:
-        if kv_cache is not None:
-            k, v = _write_and_gather_kv(kv_cache, k, v, layer_id, q, attn_mask)
+        if q.ndim == 4:
+            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 n_rep > 1:
-            k = repeat_kv(k, n_rep)
-            v = repeat_kv(v, n_rep)
-
-        out = 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,
-        )
-        out = out.permute(0, 2, 1, 3).contiguous().flatten(2)
-        return out
+        if kv_cache is None or kv_cache.qo_indptr is None:
+            raise ValueError("packed attention requires KV cache metadata")
+        kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
+        kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
+        outputs = []
+        n_rep = q.size(1) // k.size(1)
+        for i in range(kv_cache.req_pool_indices.numel()):
+            q_start = int(kv_cache.qo_indptr[i])
+            q_end = int(kv_cache.qo_indptr[i + 1])
+            indices = kv_cache.req_to_token[
+                kv_cache.req_pool_indices[i], : kv_cache.seq_lens[i]
+            ]
+            k_i = kv_cache.k_buffer[layer_id, indices]
+            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)
@@ -530,16 +569,14 @@ class CudaBackend(AttentionBackend):
         if kv_cache is None:
             raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
 
-        loc = kv_cache.out_cache_loc[:, 0]
-        kv_cache.k_buffer[layer_id, loc] = k[:, 0]
-        kv_cache.v_buffer[layer_id, loc] = v[:, 0]
-
-        q_3d = q.squeeze(1)
+        loc = kv_cache.out_cache_loc
+        kv_cache.k_buffer[layer_id, loc] = k
+        kv_cache.v_buffer[layer_id, loc] = v
 
         kv_indptr = kv_cache.kv_indptr
 
         out = attn_paged_decode(
-            q_3d,
+            q,
             kv_cache.k_buffer[layer_id],
             kv_cache.v_buffer[layer_id],
             kv_cache.req_to_token,
@@ -550,7 +587,7 @@ class CudaBackend(AttentionBackend):
             ml_part_buf=kv_cache.decode_ml_part,
             out_buf=kv_cache.decode_out,
         )
-        return out.unsqueeze(1).flatten(2)
+        return out
 
     def fwd_prefill(
         self,
@@ -565,30 +602,22 @@ class CudaBackend(AttentionBackend):
         if kv_cache is None:
             raise RuntimeError("CudaBackend does not support training (kv_cache=None)")
 
-        loc = kv_cache.out_cache_loc.reshape(-1)
-        kv_cache.k_buffer[layer_id, loc] = k.reshape(-1, k.size(2), k.size(3))
-        kv_cache.v_buffer[layer_id, loc] = v.reshape(-1, v.size(2), v.size(3))
-
-        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))
+        loc = kv_cache.out_cache_loc
+        kv_cache.k_buffer[layer_id, loc] = k
+        kv_cache.v_buffer[layer_id, loc] = v
 
         out = attn_paged_prefill(
-            q_flat,
+            q,
             kv_cache.k_buffer[layer_id],
             kv_cache.v_buffer[layer_id],
             kv_cache.req_to_token,
             kv_cache.req_pool_indices,
-            kv_indptr,
-            qo_indptr,
+            kv_cache.kv_indptr,
+            kv_cache.qo_indptr,
             attn_mask,
             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)
@@ -617,7 +646,7 @@ class FlashAttnBackend(AttentionBackend):
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
     ) -> 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(
         self,
@@ -629,25 +658,18 @@ class FlashAttnBackend(AttentionBackend):
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
     ) -> 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,
         q: Tensor,
         k: Tensor,
         v: Tensor,
-        kv_cache: Optional["KVCache"],
-        layer_id: int,
         attn_mask: Optional[Tensor] = None,
         is_causal: bool = False,
     ) -> 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)
         if n_rep > 1:
             k = repeat_kv(k, n_rep)
@@ -670,9 +692,9 @@ class FlashAttnBackend(AttentionBackend):
             v.contiguous(),
             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,
         q: Tensor,
         k: Tensor,
@@ -680,22 +702,27 @@ class FlashAttnBackend(AttentionBackend):
         kv_cache: "KVCache",
         layer_id: int,
     ) -> 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()
-        out = fa.flash_attn_with_kvcache(
-            q=q,
-            k_cache=k_cache,
-            v_cache=v_cache,
-            k=k,
-            v=v,
-            cache_seqlens=(kv_cache.seq_lens - 1).to(torch.int32),
-            cache_batch_idx=kv_cache.req_pool_indices.to(torch.int32),
+        if fa is None or not hasattr(fa, "flash_attn_varlen_func"):
+            raise RuntimeError("packed inference requires flash_attn_varlen_func")
+        kv_cache.k_buffer[layer_id, kv_cache.out_cache_loc] = k
+        kv_cache.v_buffer[layer_id, kv_cache.out_cache_loc] = v
+        page_table = kv_cache.req_to_token[
+            kv_cache.req_pool_indices, : kv_cache.max_len
+        ]
+        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,
         )
-        return out.flatten(2)
+        return out

astrai/extension/rotary_backend.py

@@ -26,7 +26,7 @@ def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
     dtype = x.dtype
     x_ = x.float().reshape(*x.shape[:-1], -1, 2)
     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_out = torch.view_as_real(x_rotated).flatten(-2)
     return x_out.to(dtype)

astrai/extension/rotary_ops.py

@@ -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
 responsibility of ``astrai.extension.rotary_backend.apply_rotary_emb``.
 
-Layout: x is [batch, seq_len, n_heads, head_dim] (bf16, contiguous).
-freqs_cis is [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs.
+Layout: x is packed [tokens, n_heads, head_dim] or dense
+[batch, seq_len, n_heads, head_dim]. ``freqs_cis`` has matching token axes.
 """
 
 import torch
@@ -25,11 +25,11 @@ def rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
     """Fused rotary embedding kernel.
 
     Args:
-        x: [batch, seq_len, n_heads, head_dim] (bf16, contiguous)
-        freqs_cis: [batch, seq_len, head_dim/2, 2] (f32, contiguous) — [cos, sin] pairs
+        x: packed 3D or dense 4D bf16 tensor.
+        freqs_cis: matching token axes followed by [head_dim/2, 2].
 
     Returns:
-        [batch, seq_len, n_heads, head_dim] (bf16)
+        Tensor with the same shape as ``x``.
     """
     _check_available()
     if not x.is_contiguous():

astrai/inference/cache/pool.py

@@ -200,14 +200,21 @@ class PagePool:
         kv_indptr = kvp_buf[: b + 1]
 
         if start_pos is not None:
-            # ---- prefill: out_cache_loc covers prefix range [start_pos:seq_len] ----
-            seq_len = seq_lens[0]
-            out_cache_loc = self._req_pool.req_to_token[
-                req_pool_indices, start_pos:seq_len
-            ]
-            q_len = seq_len - start_pos
-            workspace.qo_indptr[: b + 1].copy_(
-                torch.arange(b + 1, dtype=torch.int32, device=device) * q_len
+            # Packed prefill concatenates each request's query tokens.
+            q_lens = [seq_len - start_pos for seq_len in seq_lens]
+            if any(q_len <= 0 for q_len in q_lens):
+                raise ValueError("prefill sequence lengths must exceed start_pos")
+            out_cache_loc = torch.cat(
+                [
+                    self._req_pool.req_to_token[
+                        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]
             decode_o_part = decode_ml_part = decode_out = None
@@ -216,8 +223,9 @@ class PagePool:
             write_pos = seq_lens_t - 1
             loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
             ocl_buf[:b].copy_(loc)
-            out_cache_loc = ocl_buf[:b]
-            qo_indptr = None
+            out_cache_loc = ocl_buf[:b].reshape(-1)
+            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_ml_part = getattr(workspace, "decode_ml_part", None)
             decode_out = getattr(workspace, "decode_out", None)

astrai/inference/runtime/executor.py

@@ -118,13 +118,13 @@ def _warmup_cuda_graphs(
             timed("warmup prefill", logger),
         ):
             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
             model(
                 ids_in,
-                input_mask=pos_in.unsqueeze(-1) >= torch.arange(warmup_len, device=dev),
                 kv_cache=kv,
                 position_ids=pos_in,
+                fwd="prefill",
             )
         task_cache.task_free(tid)
 
@@ -159,15 +159,14 @@ def _warmup_cuda_graphs(
                 for tid in task_ids:
                     task_cache.task_extend(tid, seq_pos)
                 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)
                 gctx.forward(
                     model,
                     key=(b,),
-                    input_ids=ids_buf.unsqueeze(1),
-                    input_mask=input_mask,
+                    input_ids=ids_buf,
                     kv_cache=kv,
-                    position_ids=ws.position_ids[:b].unsqueeze(1),
+                    position_ids=ws.position_ids[:b],
+                    fwd="decode",
                 )
 
         for tid in task_ids:
@@ -308,20 +307,15 @@ class Executor:
         batch_sz = len(tasks)
 
         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,
             device=self.device,
         )
 
         task_ids = [t.task_id for t in tasks]
-        position_ids = (
-            torch.arange(start_pos, prompt_len, dtype=torch.long, device=self.device)
-            .unsqueeze(0)
-            .expand(batch_sz, -1)
-        )
-        input_mask = position_ids.unsqueeze(-1) >= torch.arange(
-            prompt_len, device=self.device
-        )
+        position_ids = torch.arange(
+            start_pos, prompt_len, dtype=torch.long, device=self.device
+        ).repeat(batch_sz)
 
         with (
             torch.inference_mode(),
@@ -329,15 +323,18 @@ class Executor:
         ):
             outputs = self.model(
                 input_ids,
-                input_mask=input_mask,
                 position_ids=position_ids,
                 kv_cache=self.task_cache.bind(
                     task_ids,
                     self._workspace,
                     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)
 
@@ -391,9 +388,6 @@ class Executor:
             )
         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) ----
 
         use_graph = (
@@ -402,9 +396,6 @@ class Executor:
             and get_backend().supports_graph()
         )
         key = (b,)
-        if use_graph:
-            input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
-
         with (
             torch.inference_mode(),
             timed(f"execute_decode forward b={b}", logger),
@@ -413,18 +404,18 @@ class Executor:
                 outputs = self._graph_ctx.forward(
                     self.model,
                     key=key,
-                    input_ids=input_ids.unsqueeze(1),
-                    input_mask=input_mask,
+                    input_ids=input_ids,
                     kv_cache=kv_cache,
-                    position_ids=ws.position_ids[:b].unsqueeze(1),
+                    position_ids=ws.position_ids[:b],
+                    fwd="decode",
                 )
             else:
                 outputs = self.model(
-                    input_ids.unsqueeze(1),
-                    input_mask=input_mask,
+                    input_ids,
                     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)

astrai/model/components/attention.py

@@ -56,9 +56,7 @@ class GQA(nn.Module):
             self.gate = Linear(dim, dim)
 
     def _split_heads(self, x: Tensor, n_heads) -> Tensor:
-        batch_size, seq_len, _ = x.shape
-        x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
-        return x
+        return x.reshape(*x.shape[:-1], n_heads, self.head_dim)
 
     def forward(
         self,
@@ -67,6 +65,7 @@ class GQA(nn.Module):
         attn_mask: Tensor = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> Tensor:
         q = self._split_heads(self.q_proj(x), self.n_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:
             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:
             sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
@@ -141,17 +142,16 @@ class MLA(nn.Module):
         attn_mask: Tensor = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> Tensor:
-        bsz, seq_len, _ = x.size()
-
         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_norm(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(
             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)
             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:
             attn_out = attn_out * F.sigmoid(self.gate(x))

astrai/model/components/decoder_block.py

@@ -54,6 +54,7 @@ class DecoderBlock(nn.Module):
         attention_mask: Optional[Tensor] = None,
         kv_cache: Optional[KVCache] = None,
         is_causal: bool = False,
+        fwd: Optional[str] = None,
     ) -> DecoderOutput:
         attn_output = self.attention(
             self.input_norm(x),
@@ -61,6 +62,7 @@ class DecoderBlock(nn.Module):
             attention_mask,
             kv_cache,
             is_causal,
+            fwd,
         )
         x = attn_output + x
         normalized = self.post_attention_norm(x)

astrai/model/components/mlp.py

@@ -100,13 +100,14 @@ class DeepSeekMoE(nn.Module):
 
     def forward(self, x: Tensor) -> FFNOutput:
         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)
 
         shared_out = self._shared_forward(x_flat)
         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 {
             "hidden_states": out,
             "aux_loss": routed_output["aux_loss"],

astrai/model/components/rope.py

@@ -65,9 +65,12 @@ class RotaryEmbedding(nn.Module):
             [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
         """
         if position_ids is None:
-            position_ids = (
-                torch.arange(x.size(1), device=x.device)
-                .unsqueeze(0)
-                .expand(x.size(0), -1)
-            )
+            if x.ndim == 2:
+                position_ids = torch.arange(x.size(0), device=x.device)
+            else:
+                position_ids = (
+                    torch.arange(x.size(1), device=x.device)
+                    .unsqueeze(0)
+                    .expand(x.size(0), -1)
+                )
         return self.freqs_cis[position_ids].float()

astrai/model/transformer.py

@@ -105,8 +105,20 @@ class AutoRegressiveLM(AutoModel):
         input_mask: Optional[Tensor] = None,
         kv_cache: Optional[KVCache] = None,
         position_ids: Optional[Tensor] = None,
+        fwd: Optional[str] = None,
     ) -> 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)
         rotary_emb = self.rotary_embedding(x, position_ids)
@@ -122,6 +134,7 @@ class AutoRegressiveLM(AutoModel):
                 attn_mask,
                 kv_cache,
                 use_sdpa_causal_mask,
+                fwd,
             )
             x = layer_output["hidden_states"]
             stats = layer_output.get("router_stats")