Fix Piecewise CUDA Graph crash with -enable-mixed-chunk (#20441)
Co-authored-by: jianyingzhu <joeyzhu@nvidia.com>
This commit is contained in:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user