@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 990f49bdcb
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 11:45 UTC

[MM][CG] Enable encoder Cudagraph for Step3VL (#42224)

Signed-off-by: JisoLya <523420504@qq.com> Signed-off-by: Soyaazz <523420504@qq.com>

Soyaazz committed May 18, 2026 at 11:19 UTC 990f49bdcb8ff51c0ceb1d784c3ca16e6c276927
8 files changed +384 -22
docs/design/cuda_graphs_multimodal.md
+2
@@ -77,6 +77,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
77 * `encoder_eager_forward(...)` — fallback eager forward when no graph fits.
78 * `get_input_modality(...)` - return the modality of the inputs.
79 * `get_max_frames_per_video()` - return model-specific max frames per video.
80 +* `postprocess_encoder_output(...)` - post process encoder output, directly call scatter_output_slices by default
81
82 !!! note
83 The `SupportsEncoderCudaGraph` protocol is designed to be model-agnostic. New vision encoder models can opt-in by implementing the protocol methods without modifying the manager.
@@ -89,6 +90,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
90 | `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ |
91 | `Qwen3VLForConditionalGeneration` | `Qwen3-VL` | ✅︎ | ✅︎ |
92 | `Qwen3_5ForConditionalGeneration` | `Qwen3.5` | ✅︎ | ✅︎ |
93 +| `Step3VLForConditionalGeneration` | `Step3-VL` | ✅︎ | ❌︎ |
94
95 !!! note
96 Encoder CUDA Graphs have currently been tested with `--mm-encoder-attn-backend=FLASH_ATTN` and `--mm-encoder-attn-backend=FLASHINFER` on Blackwell GPUs.
examples/generate/multimodal/vision_language_offline.py
+1
@@ -2560,6 +2560,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2560 "qwen2_vl",
2561 "qwen3_5",
2562 "qwen3_5_moe",
2563 + "stepvl",
2564 ]
2565
2566
tests/models/multimodal/generation/test_vit_cudagraph.py
+12
@@ -41,6 +41,13 @@ def qwen_vl_chat_template(content: str) -> str:
41 return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
42
43
44 +def step3_vl_chat_template(content: str) -> str:
45 + return (
46 + "<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n "
47 + f"<im_patch>{content} <|EOT|><|BOT|>assistant\n"
48 + )
49 +
50 +
51 MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
52 "qwen2_5_vl": VitCudagraphTestConfig(
53 model="Qwen/Qwen2.5-VL-3B-Instruct",
@@ -90,6 +97,11 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
97 needs_video_metadata=False,
98 marks=[pytest.mark.core_model],
99 ),
100 + "step3_vl": VitCudagraphTestConfig(
101 + model="stepfun-ai/Step3-VL-10B",
102 + image_prompt=step3_vl_chat_template("What is in this image?"),
103 + video_prompt=None,
104 + ),
105 }
106
107
vllm/model_executor/models/interfaces.py
+21
@@ -1594,6 +1594,27 @@ class SupportsEncoderCudaGraph(Protocol):
1594 """
1595 ...
1596
1597 + def postprocess_encoder_output(
1598 + self,
1599 + output: torch.Tensor,
1600 + indices: list[int],
1601 + per_item_out_tokens: list[int],
1602 + dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
1603 + clone: bool = False,
1604 + batch_mm_kwargs: dict[str, Any] | None = None,
1605 + ) -> None:
1606 + """
1607 + Post-process encoder output, directly call scatter_output_slices by default.
1608 +
1609 + By default, delegates directly to scatter_output_slices.
1610 + Override this for models that require additional processing on the raw
1611 + encoder output prior to scattering, e.g. Step3-VL, which merges features
1612 + according to dynamic patch counts before scattering.
1613 + """
1614 + from vllm.model_executor.models.utils import scatter_output_slices
1615 +
1616 + scatter_output_slices(output, indices, per_item_out_tokens, dest, clone)
1617 +
1618 def prepare_encoder_cudagraph_capture_inputs(
1619 self,
1620 token_budget: int,
vllm/model_executor/models/step3_vl.py
+323 -2
@@ -46,7 +46,12 @@ from vllm.transformers_utils.processors.step3_vl import (
46 )
47 from vllm.utils.tensor_schema import TensorSchema, TensorShape
48
49 -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
49 +from .interfaces import (
50 + MultiModalEmbeddings,
51 + SupportsEncoderCudaGraph,
52 + SupportsMultiModal,
53 + SupportsPP,
54 +)
55 from .utils import (
56 AutoWeightsLoader,
57 WeightsMapper,
@@ -487,7 +492,9 @@ class Step3VisionTransformer(nn.Module):
492 info=Step3VLProcessingInfo,
493 dummy_inputs=Step3VLDummyInputsBuilder,
494 )
490 -class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
495 +class Step3VLForConditionalGeneration(
496 + nn.Module, SupportsMultiModal, SupportsPP, SupportsEncoderCudaGraph
497 +):
498 hf_to_vllm_mapper = WeightsMapper(
499 orig_to_new_prefix={
500 "model.": "language_model.model.",
@@ -510,6 +517,7 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
517 multimodal_config = vllm_config.model_config.multimodal_config
518
519 self.config = config
520 + self.model_config = vllm_config.model_config
521 self.multimodal_config = multimodal_config
522 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
523
@@ -570,6 +578,17 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
578 def dtype(self):
579 return next(self.parameters()).dtype
580
581 + @staticmethod
582 + def _compute_spatial_tokens(size, patch_size, stride):
583 + # Compute the number of spatial tokens after two rounds of
584 + # downsampling with given patch size and stride.
585 + grid = size // patch_size
586 + vit_tokens = grid * grid
587 + spatial = int(math.sqrt(vit_tokens))
588 + h1 = (spatial - 2) // stride + 1
589 + h2 = (h1 - 1) // 2 + 1
590 + return h2 * h2
591 +
592 def _parse_and_validate_image_input(
593 self, **kwargs: object
594 ) -> Step3VLImageInputs | None:
@@ -676,6 +695,308 @@ class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP)
695 is_multimodal=is_multimodal,
696 )
697
698 + def get_encoder_cudagraph_config(self):
699 + from vllm.v1.worker.encoder_cudagraph_defs import (
700 + EncoderCudaGraphConfig,
701 + )
702 +
703 + return EncoderCudaGraphConfig(
704 + modalities=["image"],
705 + input_key_by_modality={"image": "pixel_values"},
706 + buffer_keys=["patch_pixel_values"],
707 + out_hidden_size=self.config.hidden_size,
708 + )
709 +
710 + def get_input_modality(
711 + self,
712 + mm_kwargs: dict[str, Any],
713 + ) -> str:
714 + return "image"
715 +
716 + def get_max_frames_per_video(
717 + self,
718 + ) -> int:
719 + return 0
720 +
721 + def get_encoder_cudagraph_budget_range(
722 + self,
723 + vllm_config: "VllmConfig",
724 + ) -> tuple[int, int]:
725 + # An image without patches
726 + min_budget = self._compute_spatial_tokens(
727 + self.config.vision_config.image_size,
728 + self.config.vision_config.patch_size,
729 + self.config.understand_projector_stride,
730 + )
731 + max_budget = min(
732 + vllm_config.scheduler_config.max_num_batched_tokens,
733 + self.model_config.max_model_len,
734 + )
735 + return min_budget, max_budget
736 +
737 + def get_encoder_cudagraph_num_items(
738 + self,
739 + mm_kwargs: dict[str, Any],
740 + ) -> int:
741 + return len(mm_kwargs.get("pixel_values", []))
742 +
743 + def get_encoder_cudagraph_per_item_output_tokens(
744 + self,
745 + mm_kwargs: dict[str, Any],
746 + ) -> list[int]:
747 + num_patches = mm_kwargs.get("num_patches")
748 + img_output_tokens = self._compute_spatial_tokens(
749 + self.config.vision_config.image_size,
750 + self.config.vision_config.patch_size,
751 + self.config.understand_projector_stride,
752 + )
753 + patch_output_tokens = self._compute_spatial_tokens(
754 + 504,
755 + self.config.vision_config.patch_size,
756 + self.config.understand_projector_stride,
757 + )
758 + return [
759 + img_output_tokens + num_patch * patch_output_tokens
760 + for num_patch in num_patches
761 + ]
762 +
763 + def get_encoder_cudagraph_per_item_input_sizes(
764 + self,
765 + mm_kwargs: dict[str, Any],
766 + ) -> list[int]:
767 + img_grid = (
768 + self.config.vision_config.image_size // self.config.vision_config.patch_size
769 + )
770 +
771 + # NOTE: 504 is the hard coded size for each patch after processing
772 + # by the vision model, which is determined by the current architecture
773 + # of the vision model and may need to be updated if the architecture changes.
774 + # The number of tokens for each patch is calculated based on this
775 + # size and the patch size.
776 + patch_grid = 504 // self.config.vision_config.patch_size
777 + total_image_pixel = img_grid * img_grid
778 + total_patch_pixel = patch_grid * patch_grid
779 +
780 + num_patches = mm_kwargs.get("num_patches")
781 + return [
782 + total_image_pixel + num_patch * total_patch_pixel
783 + for num_patch in num_patches
784 + ]
785 +
786 + def select_encoder_cudagraph_items(
787 + self,
788 + mm_kwargs: dict[str, Any],
789 + indices: list[int],
790 + ) -> dict[str, Any]:
791 + pixel_values = mm_kwargs["pixel_values"]
792 + patch_pixel_values = mm_kwargs["patch_pixel_values"]
793 + num_patches = mm_kwargs["num_patches"]
794 +
795 + # calcute the accumulated patch counts
796 + cum_patches = [0]
797 + for p in num_patches:
798 + cum_patches.append(cum_patches[-1] + p)
799 +
800 + if len(indices) == 0:
801 + return {
802 + "pixel_values": pixel_values[:0],
803 + "patch_pixel_values": patch_pixel_values[:0],
804 + "num_patches": num_patches[:0],
805 + }
806 +
807 + selected_pv = pixel_values[indices]
808 + selected_np = num_patches[indices]
809 + selected_ppv = torch.cat(
810 + [patch_pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
811 + )
812 +
813 + return {
814 + "pixel_values": selected_pv,
815 + "patch_pixel_values": selected_ppv,
816 + "num_patches": selected_np,
817 + }
818 +
819 + def prepare_encoder_cudagraph_capture_inputs(
820 + self,
821 + token_budget: int,
822 + max_batch_size: int,
823 + max_frames_per_batch: int,
824 + device: torch.device,
825 + dtype: torch.dtype,
826 + ):
827 + from vllm.v1.worker.encoder_cudagraph_defs import (
828 + EncoderCudaGraphCaptureInputs,
829 + )
830 +
831 + # For pixel_value, the max input size is max_batch_size
832 + img_output_tokens = self._compute_spatial_tokens(
833 + self.config.vision_config.image_size,
834 + self.config.vision_config.patch_size,
835 + self.config.understand_projector_stride,
836 + )
837 + patch_output_tokens = self._compute_spatial_tokens(
838 + 504,
839 + self.config.vision_config.patch_size,
840 + self.config.understand_projector_stride,
841 + )
842 + dummy_pixel_values = torch.randn(
843 + max_batch_size,
844 + 3,
845 + self.config.vision_config.image_size,
846 + self.config.vision_config.image_size,
847 + device=device,
848 + dtype=dtype,
849 + )
850 + # max_num_patches is the max total patches across the whole batch.
851 + # token_budget = max_batch_size * img_out + max_num_patches * patch_out
852 + max_num_patches = max(
853 + 0,
854 + (token_budget - max_batch_size * img_output_tokens) // patch_output_tokens,
855 + )
856 + dummy_patch_pixel_values = torch.randn(
857 + max_num_patches,
858 + 3,
859 + 504,
860 + 504,
861 + device=device,
862 + dtype=dtype,
863 + )
864 + # num_patches is NOT in buffers -- the per-item merge is done
865 + # CPU-side by finalize_encoder_cudagraph_output using the actual
866 + # batch's num_patches from mm_kwargs.
867 + mm_kwargs = {
868 + "pixel_values": dummy_pixel_values,
869 + "patch_pixel_values": dummy_patch_pixel_values,
870 + }
871 +
872 + buffers = {
873 + "patch_pixel_values": dummy_patch_pixel_values,
874 + }
875 +
876 + return EncoderCudaGraphCaptureInputs(
877 + mm_kwargs=mm_kwargs,
878 + buffers=buffers,
879 + )
880 +
881 + def encoder_cudagraph_forward(
882 + self,
883 + mm_kwargs: dict[str, Any],
884 + buffers: dict[str, torch.Tensor],
885 + ) -> torch.Tensor:
886 + # Graph captures only the compute (vision model + conv projector).
887 + # Per-item merge happens CPU-side in finalize_encoder_cudagraph_output
888 + # using actual num_patches from the batch data.
889 + pixel_values = mm_kwargs["pixel_values"]
890 + patch_pixel_values = buffers["patch_pixel_values"]
891 +
892 + image_features = self._process_image_features(
893 + self._get_vision_model_output(pixel_values)
894 + )
895 +
896 + has_patches = len(patch_pixel_values) > 0
897 + if has_patches:
898 + patch_features = self._process_image_features(
899 + self._get_vision_model_output(patch_pixel_values)
900 + )
901 +
902 + # Deterministic single cat: [all_img_flat, all_patch_flat]
903 + img_flat = image_features.reshape(-1, image_features.shape[-1])
904 + if has_patches:
905 + patch_flat = patch_features.reshape(-1, patch_features.shape[-1])
906 + return torch.cat([img_flat, patch_flat], dim=0)
907 + return img_flat
908 +
909 + def encoder_eager_forward(
910 + self,
911 + mm_kwargs: dict[str, Any],
912 + ) -> torch.Tensor:
913 + image_input = Step3VLImagePixelInputs(
914 + type="pixel_values",
915 + pixel_values=mm_kwargs["pixel_values"],
916 + patch_pixel_values=mm_kwargs["patch_pixel_values"],
917 + num_patches=mm_kwargs["num_patches"],
918 + )
919 + vision_embeddings = self._process_image_input(image_input)
920 + return torch.cat(vision_embeddings, dim=0)
921 +
922 + def postprocess_encoder_output(
923 + self,
924 + output: torch.Tensor,
925 + indices: list[int],
926 + per_item_out_tokens: list[int],
927 + dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
928 + clone: bool = False,
929 + batch_mm_kwargs: dict[str, Any] | None = None,
930 + ):
931 + """CPU-side per-item merge after graph replay.
932 +
933 + The graph output is ``[all_img_flat, all_patch_flat]``.
934 + This method splits the flat output into image and patch features,
935 + then reassembles per-item embeddings using the *actual* batch
936 + ``num_patches`` from ``batch_mm_kwargs`` (not the capture-time values).
937 + """
938 + num_patches = batch_mm_kwargs["num_patches"]
939 + hidden = output.shape[-1]
940 + bsz = len(indices)
941 +
942 + img_out = self._compute_spatial_tokens(
943 + self.config.vision_config.image_size,
944 + self.config.vision_config.patch_size,
945 + self.config.understand_projector_stride,
946 + )
947 + patch_out = self._compute_spatial_tokens(
948 + 504,
949 + self.config.vision_config.patch_size,
950 + self.config.understand_projector_stride,
951 + )
952 +
953 + # Valid portion: bsz images, actual_total_patches patches
954 + actual_np = [int(np) for np in num_patches]
955 + total_patches = sum(actual_np)
956 + img_tokens = bsz * img_out
957 + patch_tokens = total_patches * patch_out
958 +
959 + img_part = output[:img_tokens].reshape(bsz, img_out, hidden)
960 + if total_patches > 0:
961 + patch_part = output[img_tokens : img_tokens + patch_tokens].reshape(
962 + -1, patch_out, hidden
963 + )
964 + else:
965 + patch_part = None
966 +
967 + merged: dict[int, torch.Tensor] = {}
968 + cur_patch = 0
969 + for i, idx in enumerate(indices):
970 + np = actual_np[i]
971 + parts: list[torch.Tensor] = []
972 + if patch_part is not None and np > 0:
973 + parts.append(patch_part[cur_patch : cur_patch + np].reshape(-1, hidden))
974 + cur_patch += np
975 + parts.append(img_part[i].reshape(-1, hidden))
976 + merged[idx] = torch.cat(parts, dim=0) if len(parts) > 1 else parts[0]
977 +
978 + out = [merged[i] for i in indices]
979 + for i, idx in enumerate(indices):
980 + dest[idx] = out[i]
981 +
982 + def prepare_encoder_cudagraph_replay_buffers(
983 + self,
984 + mm_kwargs: dict[str, Any],
985 + max_batch_size: int,
986 + max_frames_per_batch: int,
987 + ):
988 + from vllm.v1.worker.encoder_cudagraph_defs import (
989 + EncoderCudaGraphReplayBuffers,
990 + )
991 +
992 + # Only patch_pixel_values lives in the buffers dict; num_patches is
993 + # processed CPU-side by finalize_encoder_cudagraph_output.
994 + return EncoderCudaGraphReplayBuffers(
995 + buffers={
996 + "patch_pixel_values": mm_kwargs["patch_pixel_values"],
997 + },
998 + )
999 +
1000 def forward(
1001 self,
1002 input_ids: torch.Tensor | None,
vllm/model_executor/models/step_vl.py
+1
@@ -500,6 +500,7 @@ class StepVLForConditionalGeneration(Step3VLForConditionalGeneration):
500 quant_config = vllm_config.quant_config
501
502 self.config = config
503 + self.model_config = vllm_config.model_config
504 self.multimodal_config = multimodal_config
505 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
506
vllm/model_executor/models/utils.py
+16
@@ -884,3 +884,19 @@ def get_layer_index(feature_layer_index: int, num_hidden_layers: int) -> int:
884 if feature_layer_index < 0:
885 return num_hidden_layers + feature_layer_index + 1
886 return feature_layer_index
887 +
888 +
889 +def scatter_output_slices(
890 + output: torch.Tensor,
891 + indices: list[int],
892 + per_item_out_tokens: list[int],
893 + dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
894 + clone: bool = False,
895 +) -> None:
896 + """Slice a concatenated output tensor and scatter into dest by index."""
897 + offset = 0
898 + for idx in indices:
899 + n_tok = per_item_out_tokens[idx]
900 + sliced = output[offset : offset + n_tok]
901 + dest[idx] = sliced.clone() if clone else sliced
902 + offset += n_tok
vllm/v1/worker/encoder_cudagraph.py
+8 -20
@@ -14,7 +14,10 @@ from vllm.distributed import (
14 tensor_model_parallel_all_gather,
15 )
16 from vllm.logger import init_logger
17 -from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
17 +from vllm.model_executor.models.interfaces import (
18 + SupportsEncoderCudaGraph,
19 +)
20 +from vllm.model_executor.models.utils import scatter_output_slices
21 from vllm.model_executor.models.vision import get_load_balance_assignment
22 from vllm.v1.worker.encoder_cudagraph_defs import (
23 EncoderCudaGraphConfig,
@@ -254,22 +257,6 @@ class EncoderCudaGraphManager:
257 for t in self.model.get_encoder_cudagraph_per_item_output_tokens(mm_kwargs)
258 ]
259
257 - @staticmethod
258 - def _scatter_output_slices(
259 - output: torch.Tensor,
260 - indices: list[int],
261 - per_item_out_tokens: list[int],
262 - dest: dict[int, torch.Tensor] | list[torch.Tensor | None],
263 - clone: bool = False,
264 - ) -> None:
265 - """Slice a concatenated output tensor and scatter into dest by index."""
266 - offset = 0
267 - for idx in indices:
268 - n_tok = per_item_out_tokens[idx]
269 - sliced = output[offset : offset + n_tok]
270 - dest[idx] = sliced.clone() if clone else sliced
271 - offset += n_tok
272 -
260 def _run_budget_graph(
261 self,
262 mm_kwargs: dict[str, Any],
@@ -414,7 +401,7 @@ class EncoderCudaGraphManager:
401 self.graph_misses += len(batch_orig_indices)
402 with torch.inference_mode():
403 raw = self.model.encoder_eager_forward(batch_mm_kwargs)
417 - self._scatter_output_slices(
404 + scatter_output_slices(
405 raw,
406 batch_orig_indices,
407 per_item_out_tokens,
@@ -440,12 +427,13 @@ class EncoderCudaGraphManager:
427 batch_mm_kwargs, token_budget, replay.buffers
428 )
429 assert output is not None
443 - self._scatter_output_slices(
430 + self.model.postprocess_encoder_output(
431 output,
432 batch_orig_indices,
433 per_item_out_tokens,
434 outputs_by_orig_idx,
435 clone=True,
436 + batch_mm_kwargs=batch_mm_kwargs,
437 )
438
439 # Return in original batch order (caller maps outputs to token positions)
@@ -573,7 +561,7 @@ class EncoderCudaGraphManager:
561 count = images_per_rank[rank]
562 if count > 0:
563 rank_items = image_rank_assignment[current_idx : current_idx + count]
576 - self._scatter_output_slices(
564 + scatter_output_slices(
565 rank_outputs[rank],
566 rank_items,
567 per_item_out_tokens,