[Bugfix][MM][CG] Enable dual-path ViT CUDA graph for Step3-VL (#46034)
Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Shanshan Shen committed
Jun 28, 2026 at 15:06 UTC
35e6c86caaaced4fd1398739fb04140b65d6ca89
2 files changed
+87
-125
docs/design/cuda_graphs_multimodal.md
+1
-1
@@ -136,7 +136,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
136
| `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ | ❌︎ |
137
| `Qwen3_5ForConditionalGeneration` | `Qwen3.5`, `Qwen3.6` | ✅︎ | ✅︎ | ❌︎ |
138
| `Qwen3_5MoeForConditionalGeneration` | `Qwen3.5-MoE`, `Qwen3.6-MoE` | ✅︎ | ✅︎ | ❌︎ |
139
-| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ❌︎ |
139
+| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ | ✅︎ |
140
141
!!! note
142
Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
vllm/model_executor/models/step3_vl.py
+86
-124
@@ -589,6 +589,31 @@ class Step3VLForConditionalGeneration(
589
h2 = (h1 - 1) // 2 + 1
590
return h2 * h2
591
592
+ @property
593
+ def img_output_tokens(self) -> int:
594
+ return self._compute_spatial_tokens(
595
+ self.config.vision_config.image_size,
596
+ self.config.vision_config.patch_size,
597
+ self.config.understand_projector_stride,
598
+ )
599
+
600
+ @property
601
+ def patch_output_tokens(self) -> int:
602
+ return self._compute_spatial_tokens(
603
+ 504,
604
+ self.config.vision_config.patch_size,
605
+ self.config.understand_projector_stride,
606
+ )
607
+
608
+ def _batched_encoder_forward(
609
+ self,
610
+ pixel_values: torch.Tensor,
611
+ ) -> torch.Tensor:
612
+ image_features = self._process_image_features(
613
+ self._get_vision_model_output(pixel_values)
614
+ )
615
+ return image_features.reshape(-1, image_features.shape[-1])
616
+
617
def _parse_and_validate_image_input(
618
self, **kwargs: object
619
) -> Step3VLImageInputs | None:
@@ -695,6 +720,8 @@ class Step3VLForConditionalGeneration(
720
is_multimodal=is_multimodal,
721
)
722
723
+ # -- SupportsEncoderCudaGraph protocol methods --
724
+
725
def get_encoder_cudagraph_config(self):
726
from vllm.v1.worker.encoder_cudagraph_defs import (
727
EncoderCudaGraphConfig,
@@ -707,18 +734,16 @@ class Step3VLForConditionalGeneration(
734
"patch_pixel_values",
735
],
736
out_hidden_size=self.config.hidden_size,
737
+ enable_dual_path_graph=True,
738
+ global_token_per_image=self.img_output_tokens,
739
+ local_token_per_patch=self.patch_output_tokens,
740
)
741
742
def get_encoder_cudagraph_budget_range(
743
self,
744
vllm_config: "VllmConfig",
745
) -> tuple[int, int]:
716
- # An image without patches
717
- min_budget = self._compute_spatial_tokens(
718
- self.config.vision_config.image_size,
719
- self.config.vision_config.patch_size,
720
- self.config.understand_projector_stride,
721
- )
746
+ min_budget = self.img_output_tokens
747
max_budget = min(
748
vllm_config.scheduler_config.max_num_batched_tokens,
749
self.model_config.max_model_len,
@@ -732,22 +757,6 @@ class Step3VLForConditionalGeneration(
757
from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
758
759
num_patches = mm_kwargs.get("num_patches")
735
- img_output_tokens = self._compute_spatial_tokens(
736
- self.config.vision_config.image_size,
737
- self.config.vision_config.patch_size,
738
- self.config.understand_projector_stride,
739
- )
740
-
741
- # NOTE: 504 is the hard coded size for each patch after processing
742
- # by the vision model, which is determined by the current architecture
743
- # of the vision model and may need to be updated if the architecture changes.
744
- # The number of tokens for each patch is calculated based on this
745
- # size and the patch size.
746
- patch_output_tokens = self._compute_spatial_tokens(
747
- 504,
748
- self.config.vision_config.patch_size,
749
- self.config.understand_projector_stride,
750
- )
760
761
img_grid = (
762
self.config.vision_config.image_size // self.config.vision_config.patch_size
@@ -759,7 +768,11 @@ class Step3VLForConditionalGeneration(
768
return [
769
EncoderItemSpec(
770
input_size=(total_image_pixel + num_patch * total_patch_pixel),
762
- output_tokens=(img_output_tokens + num_patch * patch_output_tokens),
771
+ output_tokens=(
772
+ self.img_output_tokens + num_patch * self.patch_output_tokens
773
+ ),
774
+ global_output_tokens=self.img_output_tokens,
775
+ local_output_tokens=num_patch * self.patch_output_tokens,
776
)
777
for num_patch in num_patches
778
]
@@ -810,46 +823,30 @@ class Step3VLForConditionalGeneration(
823
EncoderCudaGraphCaptureInputs,
824
)
825
813
- # For pixel_value, the max input size is max_batch_size
814
- img_output_tokens = self._compute_spatial_tokens(
815
- self.config.vision_config.image_size,
816
- self.config.vision_config.patch_size,
817
- self.config.understand_projector_stride,
818
- )
819
- patch_output_tokens = self._compute_spatial_tokens(
820
- 504,
821
- self.config.vision_config.patch_size,
822
- self.config.understand_projector_stride,
823
- )
824
- dummy_pixel_values = torch.randn(
825
- max_batch_size,
826
- 3,
827
- self.config.vision_config.image_size,
828
- self.config.vision_config.image_size,
829
- device=device,
830
- dtype=dtype,
831
- )
832
- # max_num_patches is the max total patches across the whole batch.
833
- # token_budget = max_batch_size * img_out + max_num_patches * patch_out
834
- max_num_patches = max(
835
- 0,
836
- (token_budget - max_batch_size * img_output_tokens) // patch_output_tokens,
837
- )
838
- dummy_patch_pixel_values = torch.randn(
839
- max_num_patches,
840
- 3,
841
- 504,
842
- 504,
843
- device=device,
844
- dtype=dtype,
845
- )
846
- # num_patches is NOT in values -- the per-item merge is done
847
- # CPU-side by finalize_encoder_cudagraph_output using the actual
848
- # batch's num_patches from mm_kwargs.
849
- values = {
850
- "pixel_values": dummy_pixel_values,
851
- "patch_pixel_values": dummy_patch_pixel_values,
852
- }
826
+ assert path in ("global", "local")
827
+ if path == "global":
828
+ max_num_images = token_budget // self.img_output_tokens
829
+ max_batch_size = min(max_batch_size, max_num_images)
830
+ dummy_pixel_values = torch.randn(
831
+ max_batch_size,
832
+ 3,
833
+ self.config.vision_config.image_size,
834
+ self.config.vision_config.image_size,
835
+ device=device,
836
+ dtype=dtype,
837
+ )
838
+ values = {"pixel_values": dummy_pixel_values}
839
+ else:
840
+ max_num_patches = token_budget // self.patch_output_tokens
841
+ dummy_patch_pixel_values = torch.randn(
842
+ max_num_patches,
843
+ 3,
844
+ 504,
845
+ 504,
846
+ device=device,
847
+ dtype=dtype,
848
+ )
849
+ values = {"patch_pixel_values": dummy_patch_pixel_values}
850
851
return EncoderCudaGraphCaptureInputs(
852
values=values,
@@ -860,42 +857,22 @@ class Step3VLForConditionalGeneration(
857
values: dict[str, torch.Tensor],
858
path: str = "default",
859
) -> torch.Tensor:
863
- # Graph captures only the compute (vision model + conv projector).
864
- # Per-item merge happens CPU-side in finalize_encoder_cudagraph_output
865
- # using actual num_patches from the batch data.
866
- pixel_values = values["pixel_values"]
867
- patch_pixel_values = values["patch_pixel_values"]
868
-
869
- image_features = self._process_image_features(
870
- self._get_vision_model_output(pixel_values)
871
- )
872
-
873
- has_patches = len(patch_pixel_values) > 0
874
- if has_patches:
875
- patch_features = self._process_image_features(
876
- self._get_vision_model_output(patch_pixel_values)
877
- )
878
-
879
- # Deterministic single cat: [all_img_flat, all_patch_flat]
880
- img_flat = image_features.reshape(-1, image_features.shape[-1])
881
- if has_patches:
882
- patch_flat = patch_features.reshape(-1, patch_features.shape[-1])
883
- return torch.cat([img_flat, patch_flat], dim=0)
884
- return img_flat
860
+ assert path in ("global", "local")
861
+ if path == "global":
862
+ return self._batched_encoder_forward(values["pixel_values"])
863
+ else:
864
+ return self._batched_encoder_forward(values["patch_pixel_values"])
865
866
def encoder_eager_forward(
867
self,
868
mm_kwargs: dict[str, Any],
869
path: str = "default",
870
) -> torch.Tensor:
891
- image_input = Step3VLImagePixelInputs(
892
- type="pixel_values",
893
- pixel_values=mm_kwargs["pixel_values"],
894
- patch_pixel_values=mm_kwargs["patch_pixel_values"],
895
- num_patches=mm_kwargs["num_patches"],
896
- )
897
- vision_embeddings = self._process_image_input(image_input)
898
- return torch.cat(vision_embeddings, dim=0)
871
+ assert path in ("global", "local")
872
+ if path == "global":
873
+ return self._batched_encoder_forward(mm_kwargs["pixel_values"])
874
+ else:
875
+ return self._batched_encoder_forward(mm_kwargs["patch_pixel_values"])
876
877
def postprocess_encoder_output(
878
self,
@@ -907,38 +884,24 @@ class Step3VLForConditionalGeneration(
884
batch_mm_kwargs: dict[str, Any] | None = None,
885
local_output: torch.Tensor | None = None,
886
):
910
- """CPU-side per-item merge after graph replay.
887
+ """CPU-side per-item merge after dual-path graph replay.
888
912
- The graph output is ``[all_img_flat, all_patch_flat]``.
913
- This method splits the flat output into image and patch features,
914
- then reassembles per-item embeddings using the *actual* batch
915
- ``num_patches`` from ``batch_mm_kwargs`` (not the capture-time values).
889
+ ``output`` contains global-image features and ``local_output``
890
+ contains local-patch features (or ``None`` when there are no patches).
891
"""
892
num_patches = batch_mm_kwargs["num_patches"]
893
hidden = output.shape[-1]
894
bsz = len(indices)
895
921
- img_out = self._compute_spatial_tokens(
922
- self.config.vision_config.image_size,
923
- self.config.vision_config.patch_size,
924
- self.config.understand_projector_stride,
925
- )
926
- patch_out = self._compute_spatial_tokens(
927
- 504,
928
- self.config.vision_config.patch_size,
929
- self.config.understand_projector_stride,
930
- )
931
-
932
- # Valid portion: bsz images, actual_total_patches patches
896
actual_np = [int(np) for np in num_patches]
897
total_patches = sum(actual_np)
935
- img_tokens = bsz * img_out
936
- patch_tokens = total_patches * patch_out
898
+ img_tokens = bsz * self.img_output_tokens
899
+ patch_tokens = total_patches * self.patch_output_tokens
900
938
- img_part = output[:img_tokens].reshape(bsz, img_out, hidden)
901
+ global_part = output[:img_tokens].reshape(bsz, self.img_output_tokens, hidden)
902
if total_patches > 0:
940
- patch_part = output[img_tokens : img_tokens + patch_tokens].reshape(
941
- -1, patch_out, hidden
903
+ patch_part = local_output[:patch_tokens].reshape(
904
+ -1, self.patch_output_tokens, hidden
905
)
906
else:
907
patch_part = None
@@ -951,7 +914,7 @@ class Step3VLForConditionalGeneration(
914
if patch_part is not None and np > 0:
915
parts.append(patch_part[cur_patch : cur_patch + np].reshape(-1, hidden))
916
cur_patch += np
954
- parts.append(img_part[i].reshape(-1, hidden))
917
+ parts.append(global_part[i].reshape(-1, hidden))
918
merged[idx] = torch.cat(parts, dim=0) if len(parts) > 1 else parts[0]
919
920
out = [merged[i] for i in indices]
@@ -969,14 +932,13 @@ class Step3VLForConditionalGeneration(
932
EncoderCudaGraphReplayBuffers,
933
)
934
972
- # Only patch_pixel_values lives in the values dict; num_patches is
973
- # processed CPU-side by finalize_encoder_cudagraph_output.
974
- return EncoderCudaGraphReplayBuffers(
975
- values={
976
- "pixel_values": mm_kwargs["pixel_values"],
977
- "patch_pixel_values": mm_kwargs["patch_pixel_values"],
978
- },
979
- )
935
+ assert path in ("global", "local")
936
+ if path == "global":
937
+ values = {"pixel_values": mm_kwargs["pixel_values"]}
938
+ else:
939
+ values = {"patch_pixel_values": mm_kwargs["patch_pixel_values"]}
940
+
941
+ return EncoderCudaGraphReplayBuffers(values=values)
942
943
def forward(
944
self,