[cuda_graph] Gate breakable-CG capture_inputs retention to DP-gather paths (#32678)
Signed-off-by: xingyuliu <charlotteliu12x@gmail.com>
This commit is contained in:
@@ -1392,7 +1392,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
run_once,
|
run_once,
|
||||||
# DP padding can install capture-only tensors on this dummy batch;
|
# DP padding can install capture-only tensors on this dummy batch;
|
||||||
# BCG retains it so their recorded addresses remain valid.
|
# BCG retains it so their recorded addresses remain valid.
|
||||||
capture_inputs=forward_batch,
|
capture_inputs=(
|
||||||
|
forward_batch
|
||||||
|
if forward_batch.global_num_tokens_gpu is not None
|
||||||
|
else None
|
||||||
|
),
|
||||||
post_warmup_hook=post_warmup_hook,
|
post_warmup_hook=post_warmup_hook,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user