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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user