@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f0feb15e7f

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 2, 23:15 UTC

[Multimodal] Simplify ViT CUDA graph interfaces (#41234)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Isotr0py committed May 22, 2026 at 22:31 UTC f0feb15e7fc521544d23c2d23de0e327a509876b

9 files changed +148 -146

tests/models/multimodal/generation/test_vit_cudagraph.py

-3

@@ -10,7 +10,6 @@ from vllm.multimodal.video import sample_frames_from_video

10 from vllm.platforms import current_platform

12 from ....conftest import IMAGE_ASSETS, VIDEO_ASSETS

13 -from ....utils import create_new_process_for_each_test

13 from ...utils import dummy_hf_overrides

14 from .vlm_utils.builders import sample_frames_with_video_metadata

@@ -142,7 +141,6 @@ def get_compilation_config(config: VitCudagraphTestConfig):

141

142 @pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS))

143 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")

145 -@create_new_process_for_each_test()

144 def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):

145 config = MODEL_CONFIGS[model_id]

146

@@ -184,7 +182,6 @@ def test_vit_cudagraph_image(model_id, vllm_runner, image_assets):

182

183 @pytest.mark.parametrize("model_id", params_with_marks(MODEL_CONFIGS))

184 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")

187 -@create_new_process_for_each_test()

185 def test_vit_cudagraph_video(model_id, vllm_runner, video_assets):

186 config = MODEL_CONFIGS[model_id]

187

tests/v1/cudagraph/test_encoder_cudagraph.py

+28 -29

@@ -26,6 +26,7 @@ from vllm.v1.worker.encoder_cudagraph_defs import (

26 EncoderCudaGraphCaptureInputs,

27 EncoderCudaGraphConfig,

28 EncoderCudaGraphReplayBuffers,

29 + EncoderItemSpec,

30 )

32 # ---------------------------------------------------------------------------

@@ -288,24 +289,18 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):

289 # For tests: min=4, max=128 (small values for fast capture)

290 return (4, 128)

291

291 - def get_encoder_cudagraph_num_items(

292 + def get_encoder_cudagraph_item_specs(

293 self,

294 mm_kwargs: dict[str, Any],

294 - ) -> int:

295 - return len(mm_kwargs["image_grid_thw"])

296 -

297 - def get_encoder_cudagraph_per_item_output_tokens(

298 - self,

299 - mm_kwargs: dict[str, Any],

300 - ) -> list[int]:

295 + ) -> list[EncoderItemSpec]:

296 m = _SPATIAL_MERGE

302 - return [t * (h // m) * (w // m) for t, h, w in mm_kwargs["image_grid_thw"]]

303 -

304 - def get_encoder_cudagraph_per_item_input_sizes(

305 - self,

306 - mm_kwargs: dict[str, Any],

307 - ) -> list[int]:

308 - return [t * h * w for t, h, w in mm_kwargs["image_grid_thw"]]

297 + return [

298 + EncoderItemSpec(

299 + input_size=t * h * w,

300 + output_tokens=t * (h // m) * (w // m),

301 + )

302 + for t, h, w in mm_kwargs["image_grid_thw"]

303 + ]

304

305 def select_encoder_cudagraph_items(

306 self,

@@ -591,19 +586,18 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):

586 # Protocol overrides that depend on modality keys

587 # ------------------------------------------------------------------

588

594 - def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int:

595 - return len(self._get_grid_thw(mm_kwargs))

596 -

597 - def get_encoder_cudagraph_per_item_output_tokens(

598 - self, mm_kwargs: dict[str, Any]

599 - ) -> list[int]:

589 + def get_encoder_cudagraph_item_specs(

590 + self,

591 + mm_kwargs: dict[str, Any],

592 + ) -> list[EncoderItemSpec]:

593 m = _SPATIAL_MERGE

601 - return [t * (h // m) * (w // m) for t, h, w in self._get_grid_thw(mm_kwargs)]

602 -

603 - def get_encoder_cudagraph_per_item_input_sizes(

604 - self, mm_kwargs: dict[str, Any]

605 - ) -> list[int]:

606 - return [t * h * w for t, h, w in self._get_grid_thw(mm_kwargs)]

594 + return [

595 + EncoderItemSpec(

596 + input_size=t * h * w,

597 + output_tokens=t * (h // m) * (w // m),

598 + )

599 + for t, h, w in self._get_grid_thw(mm_kwargs)

600 + ]

601

602 def select_encoder_cudagraph_items(

603 self, mm_kwargs: dict[str, Any], indices: list[int]

@@ -679,11 +673,16 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):

673 return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})

674

675 def encoder_cudagraph_forward(

682 - self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor]

676 + self,

677 + mm_kwargs: dict[str, Any],

678 + buffers: dict[str, torch.Tensor],

679 ) -> torch.Tensor:

680 return self._forward(self._get_pixel_values(mm_kwargs))

681

686 - def encoder_eager_forward(self, mm_kwargs: dict[str, Any]) -> torch.Tensor:

682 + def encoder_eager_forward(

683 + self,

684 + mm_kwargs: dict[str, Any],

685 + ) -> torch.Tensor:

686 return self._forward(self._get_pixel_values(mm_kwargs))

687

688

vllm/model_executor/models/interfaces.py

+7 -21

@@ -50,6 +50,7 @@ if TYPE_CHECKING:

50 EncoderCudaGraphCaptureInputs,

51 EncoderCudaGraphConfig,

52 EncoderCudaGraphReplayBuffers,

53 + EncoderItemSpec,

54 )

55 else:

56 VllmConfig = object

@@ -1549,30 +1550,15 @@ class SupportsEncoderCudaGraph(Protocol):

1550 """

1551 ...

1552

1552 - def get_encoder_cudagraph_num_items(

1553 + def get_encoder_cudagraph_item_specs(

1554 self,

1555 mm_kwargs: dict[str, Any],

1555 - ) -> int:

1556 - """Return the number of items (e.g. images) in the batch."""

1557 - ...

1558 -

1559 - def get_encoder_cudagraph_per_item_output_tokens(

1560 - self,

1561 - mm_kwargs: dict[str, Any],

1562 - ) -> list[int]:

1563 - """Return output token count for each item.

1564 -

1565 - Used for greedy packing and DP load balancing.

1566 - """

1567 - ...

1568 -

1569 - def get_encoder_cudagraph_per_item_input_sizes(

1570 - self,

1571 - mm_kwargs: dict[str, Any],

1572 - ) -> list[int]:

1573 - """Return input size (e.g. patch count) for each item.

1556 + ) -> list["EncoderItemSpec"]:

1557 + """Return specs describing each item in the batch.

1558

1575 - Used for input tensor slicing offsets.

1559 + Replaces the former separate methods for num_items,

1560 + per_item_output_tokens, and per_item_input_sizes.

1561 + The manager derives all three from this single return value.

1562 """

1563 ...

1564

vllm/model_executor/models/qwen2_5_vl.py

+24 -19

@@ -1578,6 +1578,7 @@ class Qwen2_5_VLForConditionalGeneration(

1578 # embedding formats between eager and cudagraph paths.

1579 modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"]

1580

1581 + max_frames = self.get_max_frames_per_video() if "video" in modalities else 1

1582 return EncoderCudaGraphConfig(

1583 modalities=modalities,

1584 input_key_by_modality={

@@ -1595,6 +1596,7 @@ class Qwen2_5_VLForConditionalGeneration(

1596 "max_seqlen_window",

1597 ],

1598 out_hidden_size=self.visual.out_hidden_size,

1599 + max_frames_per_video=max_frames,

1600 )

1601

1602 def get_input_modality(

@@ -1603,7 +1605,9 @@ class Qwen2_5_VLForConditionalGeneration(

1605 ) -> str:

1606 if "image_grid_thw" in mm_kwargs:

1607 return "image"

1606 - return "video"

1608 + elif "video_grid_thw" in mm_kwargs:

1609 + return "video"

1610 + raise AssertionError("This line should be unreachable.")

1611

1612 def get_max_frames_per_video(self) -> int:

1613 mm_registry = MULTIMODAL_REGISTRY

@@ -1649,26 +1653,21 @@ class Qwen2_5_VLForConditionalGeneration(

1653 grid_thw = grid_thw.tolist()

1654 return grid_thw

1655

1652 - def get_encoder_cudagraph_num_items(

1656 + def get_encoder_cudagraph_item_specs(

1657 self,

1658 mm_kwargs: dict[str, Any],

1655 - ) -> int:

1656 - return len(self._get_grid_thw_by_modality(mm_kwargs))

1659 + ):

1660 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

1661

1658 - def get_encoder_cudagraph_per_item_output_tokens(

1659 - self,

1660 - mm_kwargs: dict[str, Any],

1661 - ) -> list[int]:

1662 m = self.visual.spatial_merge_size

1663 grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1664 - return [t * (h // m) * (w // m) for t, h, w in grid_thw]

1665 -

1666 - def get_encoder_cudagraph_per_item_input_sizes(

1667 - self,

1668 - mm_kwargs: dict[str, Any],

1669 - ) -> list[int]:

1670 - grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1671 - return [t * h * w for t, h, w in grid_thw]

1664 + return [

1665 + EncoderItemSpec(

1666 + input_size=t * h * w,

1667 + output_tokens=t * (h // m) * (w // m),

1668 + )

1669 + for t, h, w in grid_thw

1670 + ]

1671

1672 def select_encoder_cudagraph_items(

1673 self,

@@ -1684,11 +1683,13 @@ class Qwen2_5_VLForConditionalGeneration(

1683 "pixel_values": pixel_values[:0],

1684 "image_grid_thw": [],

1685 }

1687 - else:

1686 + elif self.get_input_modality(mm_kwargs) == "video":

1687 return {

1688 "pixel_values_videos": pixel_values[:0],

1689 "video_grid_thw": [],

1690 }

1691 + else:

1692 + raise AssertionError("This line should be unreachable.")

1693

1694 # Compute cumulative patch offsets for slicing pixel_values

1695 patches_per_item = [t * h * w for t, h, w in grid_thw]

@@ -1706,11 +1707,13 @@ class Qwen2_5_VLForConditionalGeneration(

1707 "pixel_values": selected_pv,

1708 "image_grid_thw": selected_grid,

1709 }

1709 - else:

1710 + elif self.get_input_modality(mm_kwargs) == "video":

1711 return {

1712 "pixel_values_videos": selected_pv,

1713 "video_grid_thw": selected_grid,

1714 }

1715 + else:

1716 + raise AssertionError("This line should be unreachable.")

1717

1718 def prepare_encoder_cudagraph_capture_inputs(

1719 self,

@@ -1831,7 +1834,7 @@ class Qwen2_5_VLForConditionalGeneration(

1834 self.model_config.max_model_len,

1835 ),

1836 )

1834 - else:

1837 + elif modality == "video":

1838 buffers = self.visual.prepare_encoder_metadata(

1839 grid_thw_list,

1840 max_frames_per_batch=max_frames_per_batch,

@@ -1840,6 +1843,8 @@ class Qwen2_5_VLForConditionalGeneration(

1843 self.model_config.max_model_len,

1844 ),

1845 )

1846 + else:

1847 + raise AssertionError("This line should be unreachable.")

1848

1849 return EncoderCudaGraphReplayBuffers(buffers=buffers)

1850

vllm/model_executor/models/qwen2_vl.py

+14 -12

@@ -1457,6 +1457,7 @@ class Qwen2VLForConditionalGeneration(

1457 EncoderCudaGraphConfig,

1458 )

1459

1460 + max_frames = self.get_max_frames_per_video()

1461 return EncoderCudaGraphConfig(

1462 modalities=["image", "video"],

1463 input_key_by_modality={

@@ -1470,6 +1471,7 @@ class Qwen2VLForConditionalGeneration(

1471 "max_seqlen",

1472 ],

1473 out_hidden_size=self.visual.out_hidden_size,

1474 + max_frames_per_video=max_frames,

1475 )

1476

1477 def get_input_modality(self, mm_kwargs: dict[str, Any]) -> str:

@@ -1515,21 +1517,21 @@ class Qwen2VLForConditionalGeneration(

1517 grid_thw = grid_thw.tolist()

1518 return grid_thw

1519

1518 - def get_encoder_cudagraph_num_items(self, mm_kwargs: dict[str, Any]) -> int:

1519 - return len(self._get_grid_thw_by_modality(mm_kwargs))

1520 + def get_encoder_cudagraph_item_specs(

1521 + self,

1522 + mm_kwargs: dict[str, Any],

1523 + ):

1524 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

1525

1521 - def get_encoder_cudagraph_per_item_output_tokens(

1522 - self, mm_kwargs: dict[str, Any]

1523 - ) -> list[int]:

1526 m = self.visual.spatial_merge_size

1527 grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1526 - return [t * (h // m) * (w // m) for t, h, w in grid_thw]

1527 -

1528 - def get_encoder_cudagraph_per_item_input_sizes(

1529 - self, mm_kwargs: dict[str, Any]

1530 - ) -> list[int]:

1531 - grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1532 - return [t * h * w for t, h, w in grid_thw]

1528 + return [

1529 + EncoderItemSpec(

1530 + input_size=t * h * w,

1531 + output_tokens=t * (h // m) * (w // m),

1532 + )

1533 + for t, h, w in grid_thw

1534 + ]

1535

1536 def select_encoder_cudagraph_items(

1537 self, mm_kwargs: dict[str, Any], indices: list[int]

vllm/model_executor/models/qwen3_vl.py

+26 -22

@@ -1776,6 +1776,9 @@ class Qwen3VLForConditionalGeneration(

1776 # disable CUDA graph for all modalities when pruning is on.

1777 modalities = [] if self.is_multimodal_pruning_enabled else ["image", "video"]

1778

1779 + # Compute max_frames_per_video for budget sizing.

1780 + max_frames = self.get_max_frames_per_video() if "video" in modalities else 1

1781 +

1782 return EncoderCudaGraphConfig(

1783 modalities=modalities,

1784 input_key_by_modality={

@@ -1791,6 +1794,7 @@ class Qwen3VLForConditionalGeneration(

1794 "sequence_lengths",

1795 ],

1796 out_hidden_size=self.visual.out_hidden_size,

1797 + max_frames_per_video=max_frames,

1798 )

1799

1800 def get_input_modality(

@@ -1799,7 +1803,9 @@ class Qwen3VLForConditionalGeneration(

1803 ) -> str:

1804 if "image_grid_thw" in mm_kwargs:

1805 return "image"

1802 - return "video"

1806 + elif "video_grid_thw" in mm_kwargs:

1807 + return "video"

1808 + raise AssertionError("This line should be unreachable.")

1809

1810 def get_max_frames_per_video(self) -> int:

1811 mm_registry = MULTIMODAL_REGISTRY

@@ -1830,11 +1836,12 @@ class Qwen3VLForConditionalGeneration(

1836 self,

1837 mm_kwargs: dict[str, Any],

1838 ) -> torch.Tensor:

1833 - if self.get_input_modality(mm_kwargs) == "image":

1834 - pixel_values = mm_kwargs["pixel_values"]

1835 - else:

1836 - pixel_values = mm_kwargs["pixel_values_videos"]

1837 - return pixel_values

1839 + modality = self.get_input_modality(mm_kwargs)

1840 + if modality == "image":

1841 + return mm_kwargs["pixel_values"]

1842 + elif modality == "video":

1843 + return mm_kwargs["pixel_values_videos"]

1844 + raise AssertionError("This line should be unreachable.")

1845

1846 def _get_grid_thw_by_modality(

1847 self,

@@ -1846,26 +1853,21 @@ class Qwen3VLForConditionalGeneration(

1853 grid_thw = grid_thw.tolist()

1854 return grid_thw

1855

1849 - def get_encoder_cudagraph_num_items(

1856 + def get_encoder_cudagraph_item_specs(

1857 self,

1858 mm_kwargs: dict[str, Any],

1852 - ) -> int:

1853 - return len(self._get_grid_thw_by_modality(mm_kwargs))

1859 + ):

1860 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

1861

1855 - def get_encoder_cudagraph_per_item_output_tokens(

1856 - self,

1857 - mm_kwargs: dict[str, Any],

1858 - ) -> list[int]:

1862 m = self.visual.spatial_merge_size

1863 grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1861 - return [t * (h // m) * (w // m) for t, h, w in grid_thw]

1862 -

1863 - def get_encoder_cudagraph_per_item_input_sizes(

1864 - self,

1865 - mm_kwargs: dict[str, Any],

1866 - ) -> list[int]:

1867 - grid_thw = self._get_grid_thw_by_modality(mm_kwargs)

1868 - return [t * h * w for t, h, w in grid_thw]

1864 + return [

1865 + EncoderItemSpec(

1866 + input_size=t * h * w,

1867 + output_tokens=t * (h // m) * (w // m),

1868 + )

1869 + for t, h, w in grid_thw

1870 + ]

1871

1872 def select_encoder_cudagraph_items(

1873 self,

@@ -2007,11 +2009,13 @@ class Qwen3VLForConditionalGeneration(

2009 grid_thw_list,

2010 max_batch_size=max_batch_size,

2011 )

2010 - else:

2012 + elif modality == "video":

2013 buffers = self.visual.prepare_encoder_metadata(

2014 grid_thw_list,

2015 max_frames_per_batch=max_frames_per_batch,

2016 )

2017 + else:

2018 + raise AssertionError("This line should be unreachable.")

2019

2020 return EncoderCudaGraphReplayBuffers(buffers=buffers)

2021

vllm/model_executor/models/step3_vl.py

+13 -28

@@ -713,11 +713,6 @@ class Step3VLForConditionalGeneration(

713 ) -> str:

714 return "image"

715

716 - def get_max_frames_per_video(

717 - self,

718 - ) -> int:

719 - return 0

720 -

716 def get_encoder_cudagraph_budget_range(

717 self,

718 vllm_config: "VllmConfig",

@@ -734,52 +729,42 @@ class Step3VLForConditionalGeneration(

729 )

730 return min_budget, max_budget

731

737 - def get_encoder_cudagraph_num_items(

732 + def get_encoder_cudagraph_item_specs(

733 self,

734 mm_kwargs: dict[str, Any],

740 - ) -> int:

741 - return len(mm_kwargs.get("pixel_values", []))

735 + ):

736 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

737

743 - def get_encoder_cudagraph_per_item_output_tokens(

744 - self,

745 - mm_kwargs: dict[str, Any],

746 - ) -> list[int]:

738 num_patches = mm_kwargs.get("num_patches")

739 img_output_tokens = self._compute_spatial_tokens(

740 self.config.vision_config.image_size,

741 self.config.vision_config.patch_size,

742 self.config.understand_projector_stride,

743 )

744 +

745 + # NOTE: 504 is the hard coded size for each patch after processing

746 + # by the vision model, which is determined by the current architecture

747 + # of the vision model and may need to be updated if the architecture changes.

748 + # The number of tokens for each patch is calculated based on this

749 + # size and the patch size.

750 patch_output_tokens = self._compute_spatial_tokens(

751 504,

752 self.config.vision_config.patch_size,

753 self.config.understand_projector_stride,

754 )

758 - return [

759 - img_output_tokens + num_patch * patch_output_tokens

760 - for num_patch in num_patches

761 - ]

755

763 - def get_encoder_cudagraph_per_item_input_sizes(

764 - self,

765 - mm_kwargs: dict[str, Any],

766 - ) -> list[int]:

756 img_grid = (

757 self.config.vision_config.image_size // self.config.vision_config.patch_size

758 )

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.

759 patch_grid = 504 // self.config.vision_config.patch_size

760 total_image_pixel = img_grid * img_grid

761 total_patch_pixel = patch_grid * patch_grid

762

780 - num_patches = mm_kwargs.get("num_patches")

763 return [

782 - total_image_pixel + num_patch * total_patch_pixel

764 + EncoderItemSpec(

765 + input_size=(total_image_pixel + num_patch * total_patch_pixel),

766 + output_tokens=(img_output_tokens + num_patch * patch_output_tokens),

767 + )

768 for num_patch in num_patches

769 ]

770

vllm/v1/worker/encoder_cudagraph.py

+14 -12

@@ -21,6 +21,7 @@ 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,

24 + EncoderItemSpec,

25 )

27 logger = init_logger(__name__)

@@ -141,8 +142,8 @@ class EncoderCudaGraphManager:

142 elif user_max_frames is not None:

143 self.max_frames_per_batch = user_max_frames

144 else:

144 - # Set it to the model-specific value according to its `processing_info`.

145 - max_frames_per_video = self.model.get_max_frames_per_video()

145 + # Set it to the model-specific value from config.

146 + max_frames_per_video = self.config.max_frames_per_video

147 self.max_frames_per_batch = self.max_batch_size * max_frames_per_video

148

149 mm_config = vllm_config.model_config.multimodal_config

@@ -250,12 +251,13 @@ class EncoderCudaGraphManager:

251 return budget

252 return None

253

254 + def _get_item_specs(self, mm_kwargs: dict[str, Any]) -> list[EncoderItemSpec]:

255 + """Get item specs from the model."""

256 + return self.model.get_encoder_cudagraph_item_specs(mm_kwargs)

257 +

258 def _get_per_item_out_tokens(self, mm_kwargs: dict[str, Any]) -> list[int]:

259 """Get per-item output token counts as plain ints."""

255 - return [

256 - int(t)

257 - for t in self.model.get_encoder_cudagraph_per_item_output_tokens(mm_kwargs)

258 - ]

260 + return [spec.output_tokens for spec in self._get_item_specs(mm_kwargs)]

261

262 def _run_budget_graph(

263 self,

@@ -274,7 +276,7 @@ class EncoderCudaGraphManager:

276 Returns:

277 Encoder outputs, or None if graph not captured.

278 """

277 - num_items = self.model.get_encoder_cudagraph_num_items(mm_kwargs)

279 + num_items = len(self._get_item_specs(mm_kwargs))

280 if token_budget not in self.budget_graphs:

281 self.graph_misses += num_items

282 return None

@@ -332,10 +334,11 @@ class EncoderCudaGraphManager:

334 always satisfy total_tokens <= max_budget and therefore

335 always find a valid budget (no miss).

336 """

335 - num_items = self.model.get_encoder_cudagraph_num_items(mm_kwargs)

337 + item_specs = self._get_item_specs(mm_kwargs)

338 + num_items = len(item_specs)

339 max_budget = self.token_budgets[-1]

340

338 - per_item_out_tokens = self._get_per_item_out_tokens(mm_kwargs)

341 + per_item_out_tokens = [spec.output_tokens for spec in item_specs]

342

343 # Sort ascending by output token count (smallest first)

344 sorted_indices = sorted(range(num_items), key=lambda i: per_item_out_tokens[i])

@@ -459,9 +462,8 @@ class EncoderCudaGraphManager:

462 tp_size = get_tensor_model_parallel_world_size()

463 current_rank = get_tensor_model_parallel_rank()

464

462 - per_item_input_sizes = self.model.get_encoder_cudagraph_per_item_input_sizes(

463 - mm_kwargs

464 - )

465 + item_specs = self._get_item_specs(mm_kwargs)

466 + per_item_input_sizes = [spec.input_size for spec in item_specs]

467

468 (image_rank_assignment, images_per_rank, input_patches_per_rank) = (

469 get_load_balance_assignment(per_item_input_sizes, tp_size)

vllm/v1/worker/encoder_cudagraph_defs.py

+22

@@ -8,6 +8,23 @@ from typing import Any

8 import torch

11 +@dataclass

12 +class EncoderItemSpec:

13 + """Description of a single encoder input item.

14 +

15 + Returned by ``get_encoder_cudagraph_item_specs()`` to describe each

16 + image or video in a batch without the manager needing to understand

17 + model-specific input formats.

18 + """

19 +

20 + input_size: int

21 + """Number of input patches/rows for this item."""

22 +

23 + output_tokens: int

24 + """Number of output tokens after encoder processing (e.g. after

25 + spatial merge)."""

26 +

27 +

28 @dataclass

29 class EncoderCudaGraphConfig:

30 """Configuration for encoder CUDA graph management.

@@ -34,6 +51,11 @@ class EncoderCudaGraphConfig:

51 """Output hidden dim of the vision encoder.

52 Used for DP gather buffer allocation."""

54 + max_frames_per_video: int = 1

55 + """Maximum number of frames per video.

56 + Only relevant when "video" is in ``modalities``.

57 + Image-only models can use the default of 1."""

58 +

60 @dataclass

61 class EncoderCudaGraphCaptureInputs: