fix: add out_buf to attn_paged_decode for CUDA graph capture compatibility

- Pre-allocate decode_out in InferenceWorkspace so attn_paged_decode does not call torch::empty inside graph capture
- Wire decode_out through KVCache, PagePool.bind_tasks, and CudaBackend.fwd_decode
- Run live forward before graph capture to get valid output (graph pool memory is zeroed after capture block exits)
- Greedy generation with graph replay is bit-exact across all batch sizes
- Decode speedups vs no-graph: B=1 2.09x, B=4 1.80x, B=8 1.94x, B=16 1.76x

astrai/inference/core/cache.py

@@ -263,6 +263,7 @@ class KVCache:
         qo_indptr: [batch+1] int32 — prefill qo prefix-sum (None in decode)
         decode_o_part: split-KV o partial workspace (mirrors FlashInfer)
         decode_ml_part: split-KV m/l partial workspace (mirrors FlashInfer)
+        decode_out: pre-allocated decode output buffer (graph-safe)
     """
 
     k_buffer: Tensor
@@ -276,6 +277,7 @@ class KVCache:
     qo_indptr: Optional[Tensor] = None
     decode_o_part: Optional[Tensor] = None
     decode_ml_part: Optional[Tensor] = None
+    decode_out: Optional[Tensor] = None
 
 
 class PagePool:
@@ -571,6 +573,7 @@ class PagePool:
             )
             qo_indptr = workspace.qo_indptr[: b + 1]
             decode_o_part, decode_ml_part = None, None
+            decode_out = None
         else:
             write_pos = seq_lens_t - 1
             loc = self._req_pool.req_to_token[req_pool_indices, write_pos].unsqueeze(-1)
@@ -579,6 +582,7 @@ class PagePool:
             qo_indptr = None
             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)
 
         return KVCache(
             k_buffer=self._storage.k_buffer,
@@ -592,6 +596,7 @@ class PagePool:
             qo_indptr=qo_indptr,
             decode_o_part=decode_o_part,
             decode_ml_part=decode_ml_part,
+            decode_out=decode_out,
         )
 
     # ---- internals ----

astrai/inference/core/graph.py

@@ -86,12 +86,14 @@ class CudaGraphContext:
         if key in self._graphs:
             self._graphs[key].replay()
         elif key in self._warmed:
+            cap_output = model(**kwargs)
             torch.cuda.synchronize()
             graph = torch.cuda.CUDAGraph()
             with torch.cuda.graph(graph):
                 self._outputs[key] = model(**kwargs)
             self._graphs[key] = graph
             self._warmed.discard(key)
+            return cap_output
         else:
             self._warmed.add(key)
             self._outputs[key] = model(**kwargs)

astrai/inference/core/workspace.py

@@ -108,6 +108,14 @@ class InferenceWorkspace:
             device=device,
         )
 
+        # Decode output buffer (graph-safe pre-alloc).  Shape matches the
+        # decode kernel's output: [batch, q_head, head_dim].
+        self.decode_out = torch.empty(
+            (max_batch_size, max_q_heads, head_dim),
+            dtype=dtype,
+            device=device,
+        )
+
     def decode_buffers(self, batch: int, q_heads: int):
         """Return ``(o_part, ml_part)`` view sliced to live dimensions."""
         return (