[cuda_graph] Gate breakable-CG capture_inputs retention to DP-gather paths (#32678)

Signed-off-by: xingyuliu <charlotteliu12x@gmail.com>
2026-08-04 02:23:34 -07:00
parent b57721ccf7
commit 5e6c37f2b4
1 changed files with 5 additions and 1 deletions
@@ -1392,7 +1392,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
            run_once,
            # DP padding can install capture-only tensors on this dummy batch;
            # 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,
        )