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
This commit is contained in:
@@ -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 (
|
||||
|
||||
Reference in New Issue
Block a user