From daf02bde3306863ca4478a905b9ffa90e9332dac Mon Sep 17 00:00:00 2001 From: Jianying <53503712+jianyingzhu@users.noreply.github.com> Date: Sat, 28 Mar 2026 12:56:21 +0800 Subject: [PATCH] Fix Piecewise CUDA Graph crash with `-enable-mixed-chunk` (#20441) Co-authored-by: jianyingzhu --- .../piecewise_cuda_graph_runner.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 18fc12785..3fd7c3adc 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -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,