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:
2026-08-07 19:45:59 +08:00
parent 6572be4f98
commit af25833fab
7 changed files with 132 additions and 4 deletions
+1
View File
@@ -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)