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/extension/attention_backend.py

@@ -511,6 +511,7 @@ class CudaBackend(AttentionBackend):
             is_causal=True,
             o_part_buf=kv_cache.decode_o_part,
             ml_part_buf=kv_cache.decode_ml_part,
+            out_buf=kv_cache.decode_out,
         )
         return out.unsqueeze(1).flatten(2)
 

astrai/extension/attention_ops.py

@@ -102,6 +102,7 @@ def attn_paged_decode(
     is_causal: bool = False,
     o_part_buf: Optional[torch.Tensor] = None,
     ml_part_buf: Optional[torch.Tensor] = None,
+    out_buf: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     """SGLang-style paged decode (q_len == 1, flat KV pool).
 
@@ -121,6 +122,7 @@ def attn_paged_decode(
         is_causal: apply causal mask
         o_part_buf: pre-allocated split-KV o partial buffer (workflow bypass)
         ml_part_buf: pre-allocated split-KV m/l buffer (workflow bypass)
+        out_buf: pre-allocated output buffer [batch, n_heads, head_dim] (graph-safe)
 
     Returns:
         [batch, n_heads, head_dim] (bf16, 3D)
@@ -139,6 +141,7 @@ def attn_paged_decode(
         causal_offset=causal_offset,
         o_part_buf=o_part_buf,
         ml_part_buf=ml_part_buf,
+        out_buf=out_buf,
     )