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
|
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(
|
static_forward_batch = ForwardBatch(
|
||||||
forward_mode=forward_batch.forward_mode,
|
forward_mode=pcg_forward_mode,
|
||||||
batch_size=bs,
|
batch_size=bs,
|
||||||
input_ids=input_ids,
|
input_ids=input_ids,
|
||||||
input_embeds=input_embeds,
|
input_embeds=input_embeds,
|
||||||
@@ -721,7 +733,7 @@ class PiecewiseCudaGraphRunner:
|
|||||||
spec_info=forward_batch.spec_info,
|
spec_info=forward_batch.spec_info,
|
||||||
capture_hidden_mode=forward_batch.capture_hidden_mode,
|
capture_hidden_mode=forward_batch.capture_hidden_mode,
|
||||||
num_token_non_padded=forward_batch.num_token_non_padded,
|
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,
|
lora_ids=forward_batch.lora_ids,
|
||||||
sampling_info=forward_batch.sampling_info,
|
sampling_info=forward_batch.sampling_info,
|
||||||
mm_inputs=forward_batch.mm_inputs,
|
mm_inputs=forward_batch.mm_inputs,
|
||||||
|
|||||||
Reference in New Issue
Block a user