Fix Piecewise CUDA Graph crash with -enable-mixed-chunk (#20441)

Co-authored-by: jianyingzhu <joeyzhu@nvidia.com>
This commit is contained in:
Jianying
2026-03-27 21:56:21 -07:00
committed by GitHub
co-authored by jianyingzhu
parent 19b1f75186
commit daf02bde33
@@ -682,8 +682,20 @@ class PiecewiseCudaGraphRunner:
next_token_logits_buffer = None
# Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1) doesn't fail on MIXED=3.
pcg_forward_mode = (
ForwardMode.EXTEND
if forward_batch.forward_mode == ForwardMode.MIXED
else forward_batch.forward_mode
)
pcg_global_forward_mode = (
ForwardMode.EXTEND
if forward_batch.global_forward_mode == ForwardMode.MIXED
else forward_batch.global_forward_mode
)
static_forward_batch = ForwardBatch(
forward_mode=forward_batch.forward_mode,
forward_mode=pcg_forward_mode,
batch_size=bs,
input_ids=input_ids,
input_embeds=input_embeds,
@@ -721,7 +733,7 @@ class PiecewiseCudaGraphRunner:
spec_info=forward_batch.spec_info,
capture_hidden_mode=forward_batch.capture_hidden_mode,
num_token_non_padded=forward_batch.num_token_non_padded,
global_forward_mode=forward_batch.global_forward_mode,
global_forward_mode=pcg_global_forward_mode,
lora_ids=forward_batch.lora_ids,
sampling_info=forward_batch.sampling_info,
mm_inputs=forward_batch.mm_inputs,