diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py
index b712fa2..fa5c0c0 100644
--- a/astrai/extension/attention_backend.py
+++ b/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
diff --git a/astrai/extension/rotary_backend.py b/astrai/extension/rotary_backend.py
index 4b3d3d3..91bd2ef 100644
--- a/astrai/extension/rotary_backend.py
+++ b/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)
diff --git a/astrai/extension/rotary_ops.py b/astrai/extension/rotary_ops.py
index 7f94efd..37d2a90 100644
--- a/astrai/extension/rotary_ops.py
+++ b/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():
diff --git a/astrai/inference/cache/pool.py b/astrai/inference/cache/pool.py
index eebf4f4..72ec88b 100644
--- a/astrai/inference/cache/pool.py
+++ b/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)
diff --git a/astrai/inference/runtime/executor.py b/astrai/inference/runtime/executor.py
index 810b965..0b55cf9 100644
--- a/astrai/inference/runtime/executor.py
+++ b/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)
diff --git a/astrai/model/components/attention.py b/astrai/model/components/attention.py
index 8e06798..831323f 100644
--- a/astrai/model/components/attention.py
+++ b/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))
diff --git a/astrai/model/components/decoder_block.py b/astrai/model/components/decoder_block.py
index 65e1497..2fa8a39 100644
--- a/astrai/model/components/decoder_block.py
+++ b/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)
diff --git a/astrai/model/components/mlp.py b/astrai/model/components/mlp.py
index 2a60dea..58c5d99 100644
--- a/astrai/model/components/mlp.py
+++ b/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"],
diff --git a/astrai/model/components/rope.py b/astrai/model/components/rope.py
index c051fe1..1384b21 100644
--- a/astrai/model/components/rope.py
+++ b/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()
diff --git a/astrai/model/transformer.py b/astrai/model/transformer.py
index a2f2167..267399b 100644
--- a/astrai/model/transformer.py
+++ b/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")
diff --git a/csrc/kernels/rotary_emb.cu b/csrc/kernels/rotary_emb.cu
index 8e9c76e..db69e76 100644
--- a/csrc/kernels/rotary_emb.cu
+++ b/csrc/kernels/rotary_emb.cu
@@ -7,13 +7,12 @@ __global__ void rotary_emb_kernel(
     const __nv_bfloat16* __restrict__ x,
     const float* __restrict__ freqs_cis,
     __nv_bfloat16* __restrict__ out,
-    int batch,
-    int seq_len,
+    int n_tokens,
     int n_heads,
     int head_dim
 ) {
     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;
          idx < total;
@@ -23,11 +22,10 @@ __global__ void rotary_emb_kernel(
         int tmp = idx / half_dim;
         int head = tmp % n_heads;
         tmp /= n_heads;
-        int seq = tmp % seq_len;
-        int b = tmp / seq_len;
+        int token = tmp;
 
-        int x_offset = ((b * seq_len + seq) * n_heads + head) * head_dim + (pair << 1);
-        int cs_offset = ((b * seq_len + seq) * half_dim + pair) * 2;
+        int x_offset = (token * n_heads + head) * head_dim + (pair << 1);
+        int cs_offset = (token * half_dim + pair) * 2;
 
         __nv_bfloat162 x_pair = *reinterpret_cast<const __nv_bfloat162*>(x + x_offset);
         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(freqs_cis.is_cuda(), "freqs_cis must be on CUDA");
     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(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.scalar_type() == torch::kFloat32, "freqs_cis must be f32");
 
-    int batch = x.size(0);
-    int seq_len = x.size(1);
-    int n_heads = x.size(2);
-    int head_dim = x.size(3);
+    int n_tokens = x.dim() == 3 ? x.size(0) : x.size(0) * x.size(1);
+    int n_heads = x.size(x.dim() - 2);
+    int head_dim = x.size(x.dim() - 1);
 
     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.size(1) == seq_len, "freqs_cis seq_len 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.numel() == (int64_t)n_tokens * head_dim,
+                "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(-1) == 2, "freqs_cis last dim must be 2 [cos, sin]");
 
     auto out = torch::empty_like(x);
 
     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 grid = std::min((total + block - 1) / block, 1024);
 
@@ -82,7 +81,7 @@ torch::Tensor rotary_emb(
         reinterpret_cast<const __nv_bfloat16*>(x.data_ptr()),
         freqs_cis.data_ptr<float>(),
         reinterpret_cast<__nv_bfloat16*>(out.data_ptr()),
-        batch, seq_len, n_heads, head_dim
+        n_tokens, n_heads, head_dim
     );
     C10_CUDA_CHECK(cudaGetLastError());
 
@@ -93,6 +92,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("rotary_emb", &rotary_emb,
         py::arg("x"),
         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"
     );
 }
diff --git a/scripts/tools/benchmark.py b/scripts/tools/benchmark.py
index 2fc50c5..45a4ad6 100644
--- a/scripts/tools/benchmark.py
+++ b/scripts/tools/benchmark.py
@@ -118,16 +118,11 @@ class GenerationBenchmark:
         workspace: InferenceWorkspace,
     ) -> list:
         input_ids = torch.randint(
-            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
+            0, self.config.vocab_size, (batch_size * 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)]
         for tid in task_ids:
@@ -137,9 +132,9 @@ class GenerationBenchmark:
         with torch.inference_mode(), attn_backend(self.backend):
             self.model(
                 input_ids,
-                input_mask=input_mask,
                 kv_cache=kv_cache,
                 position_ids=position_ids,
+                fwd="prefill",
             )
         torch.cuda.synchronize()
         return task_ids
@@ -154,24 +149,20 @@ class GenerationBenchmark:
     ):
         batch_size = len(task_ids)
         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(
-            [[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:
             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)
         with torch.inference_mode(), attn_backend(self.backend):
             self.model(
                 input_ids,
-                input_mask=input_mask,
                 kv_cache=kv_cache,
                 position_ids=position_ids,
+                fwd="decode",
             )
 
     def run_prefill_benchmark(
@@ -188,25 +179,23 @@ class GenerationBenchmark:
             task_cache.task_alloc(tid, list(range(prompt_length)))
 
         input_ids = torch.randint(
-            0, self.config.vocab_size, (batch_size, prompt_length), device=self.device
-        )
-        position_ids = (
-            torch.arange(0, prompt_length, dtype=torch.long, device=self.device)
-            .unsqueeze(0)
-            .expand(batch_size, -1)
-        )
-        input_mask = position_ids.unsqueeze(-1) >= torch.arange(
-            prompt_length, device=self.device
+            0,
+            self.config.vocab_size,
+            (batch_size * 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)
 
         for _ in range(3):
             with torch.inference_mode(), attn_backend(self.backend):
                 self.model(
                     input_ids,
-                    input_mask=input_mask,
                     kv_cache=kv_cache,
                     position_ids=position_ids,
+                    fwd="prefill",
                 )
 
         torch.cuda.synchronize()
@@ -215,9 +204,9 @@ class GenerationBenchmark:
             with torch.inference_mode(), attn_backend(self.backend):
                 self.model(
                     input_ids,
-                    input_mask=input_mask,
                     kv_cache=kv_cache,
                     position_ids=position_ids,
+                    fwd="prefill",
                 )
         torch.cuda.synchronize()
         elapsed = time.perf_counter() - t0
@@ -311,37 +300,29 @@ class GenerationBenchmark:
         )
 
         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)
-        arange = torch.arange(max_seq_len, device=self.device)
 
         gctx = CudaGraphContext(enabled=True)
         graph_key = (b,)
 
         def _decode_graph_step(seq_len):
             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
             for tid in task_ids:
                 task_cache.task_extend(tid, seq_len)
             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):
                 return gctx.forward(
                     self.model,
                     key=graph_key,
                     input_ids=input_ids_buf,
-                    input_mask=input_mask,
                     kv_cache=kv_cache,
-                    position_ids=position_ids_buf.unsqueeze(1),
+                    position_ids=position_ids_buf,
+                    fwd="decode",
                 )
 
         for i in range(5):
diff --git a/tests/extension/test_backend_equivalence.py b/tests/extension/test_backend_equivalence.py
index a775c05..2ec240e 100644
--- a/tests/extension/test_backend_equivalence.py
+++ b/tests/extension/test_backend_equivalence.py
@@ -59,17 +59,9 @@ 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)
+    input_ids = torch.tensor(sum(prompt_ids, []), dtype=torch.long, device=device)
+    position_ids = torch.cat([torch.arange(len(p), device=device) for p in prompt_ids])
 
     cache = PagePool(
         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)
     with torch.inference_mode():
         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")
@@ -100,22 +92,24 @@ def test_prefill_with_kv_cache_matches_torch(cuda_model):
         with torch.inference_mode():
             out_cuda = model(
                 input_ids,
-                input_mask=input_mask,
                 kv_cache=kv2,
                 position_ids=position_ids,
+                fwd="prefill",
             )
 
+    offset = 0
     for i, p in enumerate(prompt_ids):
         d = (
             (
-                out_torch["logits"][i, : len(p)].float()
-                - out_cuda["logits"][i, : len(p)].float()
+                out_torch["logits"][offset : offset + len(p)].float()
+                - out_cuda["logits"][offset : offset + len(p)].float()
             )
             .abs()
             .max()
             .item()
         )
         assert d == 0.0, f"Prefill diff for sample {i}: {d}"
+        offset += len(p)
 
 
 @skip_no_kernel
@@ -136,15 +130,8 @@ def test_decode_mixed_seq_lens_matches_torch(cuda_model):
     )
 
     # 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)
+    input_ids = torch.tensor(sum(prompt_ids, []), dtype=torch.long, device=device)
+    position_ids = torch.cat([torch.arange(len(p), device=device) for p in prompt_ids])
 
     task_cache = _mk_task_cache(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])
     kv = task_cache.bind(["t1", "t2"], ws, start_pos=0)
     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)
-    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)
+    dec_ids = torch.tensor([99, 98], dtype=torch.long, device=device)
+    dec_pos = torch.tensor([8, 6], dtype=torch.long, device=device)
 
     task_cache.task_extend("t1", 8)
     task_cache.task_extend("t2", 6)
     kv_t = task_cache.bind(["t1", "t2"], ws)
     with torch.inference_mode():
-        out_torch = model(
-            dec_ids, input_mask=dec_mask, kv_cache=kv_t, position_ids=dec_pos
-        )
+        out_torch = model(dec_ids, kv_cache=kv_t, position_ids=dec_pos, fwd="decode")
 
     kv_c = task_cache.bind(["t1", "t2"], ws)
     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
-            )
+            out_cuda = model(dec_ids, kv_cache=kv_c, position_ids=dec_pos, fwd="decode")
 
     diff = (out_torch["logits"].float() - out_cuda["logits"].float()).abs().max().item()
     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)
     task_cache.task_alloc("t1", prompt_ids)
 
-    input_ids = torch.tensor([prompt_ids], dtype=torch.long, device=device)
-    position_ids = torch.arange(len(prompt_ids), device=device).unsqueeze(0)
-    input_mask = torch.ones(1, len(prompt_ids), dtype=torch.bool, device=device)
+    input_ids = torch.tensor(prompt_ids, dtype=torch.long, device=device)
+    position_ids = torch.arange(len(prompt_ids), device=device)
 
     with attn_backend(ATTN_BACKEND.CUDA), torch.inference_mode():
         model(
             input_ids,
-            input_mask=input_mask,
             position_ids=position_ids,
             kv_cache=task_cache.bind(["t1"], ws, start_pos=0),
+            fwd="prefill",
         )
 
         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
 
         decode_args = {
-            "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),
+            "input_ids": torch.tensor([9], dtype=torch.long, device=device),
+            "position_ids": torch.tensor([len(prompt_ids)], device=device),
             "kv_cache": kv_cache,
+            "fwd": "decode",
         }
         graph = CudaGraphContext(enabled=True)
         graph.forward(model, key=(1,), **decode_args)
         graph.forward(model, key=(1,), **decode_args)
         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_v = kv_cache.v_buffer[:, slot].clone()
 
diff --git a/tests/inference/test_cache.py b/tests/inference/test_cache.py
index 0d4a559..ed277aa 100644
--- a/tests/inference/test_cache.py
+++ b/tests/inference/test_cache.py
@@ -279,7 +279,7 @@ def test_page_pool_contiguous_bind_tasks_prefill():
     task_cache.task_alloc("t1", list(range(10)))
     task_cache.task_alloc("t2", list(range(10)))
     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.seq_lens.tolist() == [10, 10]
     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("t2", 8)
     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]
 
 
diff --git a/tests/module/test_model_forward.py b/tests/module/test_model_forward.py
index fa47bc4..a88726d 100644
--- a/tests/module/test_model_forward.py
+++ b/tests/module/test_model_forward.py
@@ -39,6 +39,52 @@ def _make_model(config=None) -> AutoRegressiveLM:
     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):
     return {"probs": probs, "topk_indices": topk_indices}