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, 12:15 UTC
Adjust design around encoder_cudagraph_forward (#42288)
Signed-off-by: Weida Hong <wdhongtw@google.com>
Weida Hong committed
May 29, 2026 at 11:02 UTC
dfe8ba7c80606d0d4230d752d0235c487fe0e8b5
8 files changed
+105
-154
tests/v1/cudagraph/test_encoder_cudagraph.py
+24
-31
@@ -87,8 +87,10 @@ class _MockModel(SupportsEncoderCudaGraph):
87
def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
88
return EncoderCudaGraphConfig(
89
modalities=["image"],
90
- input_key_by_modality={"image": "pixel_values"},
91
- buffer_keys=["dummy_buf"],
90
+ buffer_keys=[
91
+ "pixel_values",
92
+ "dummy_buf",
93
+ ],
94
out_hidden_size=32,
95
)
96
@@ -269,9 +271,6 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
271
def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
272
return EncoderCudaGraphConfig(
273
modalities=["image"],
272
- input_key_by_modality={
273
- "image": "pixel_values",
274
- },
274
buffer_keys=["dummy_buf"],
275
out_hidden_size=_HIDDEN,
276
)
@@ -350,11 +349,10 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
349
n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)
350
dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)
351
return EncoderCudaGraphCaptureInputs(
353
- mm_kwargs={
352
+ values={
353
"pixel_values": dummy_pixel_values,
355
- "image_grid_thw": grid_config,
354
+ "dummy_buf": dummy_buf,
355
},
357
- buffers={"dummy_buf": dummy_buf},
356
)
357
358
def prepare_encoder_cudagraph_replay_buffers(
@@ -367,14 +365,18 @@ class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
365
n_out = _count_output_tokens(grid_thw, _SPATIAL_MERGE)
366
p = next(self.parameters())
367
dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)
370
- return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
368
+ return EncoderCudaGraphReplayBuffers(
369
+ values={
370
+ "pixel_values": mm_kwargs["pixel_values"],
371
+ "dummy_buf": dummy_buf,
372
+ }
373
+ )
374
375
def encoder_cudagraph_forward(
376
self,
374
- mm_kwargs: dict[str, Any],
375
- buffers: dict[str, torch.Tensor],
377
+ values: dict[str, torch.Tensor],
378
) -> torch.Tensor:
377
- return self._forward(mm_kwargs["pixel_values"])
379
+ return self._forward(values["pixel_values"])
380
381
def encoder_eager_forward(
382
self,
@@ -551,10 +553,6 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
553
def get_encoder_cudagraph_config(self) -> EncoderCudaGraphConfig:
554
return EncoderCudaGraphConfig(
555
modalities=["image", "video"],
554
- input_key_by_modality={
555
- "image": "pixel_values",
556
- "video": "pixel_values_videos",
557
- },
556
buffer_keys=["dummy_buf"],
557
out_hidden_size=_HIDDEN,
558
)
@@ -654,11 +652,10 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
652
n_out = _count_output_tokens(grid_config, _SPATIAL_MERGE)
653
dummy_buf = torch.zeros(n_out, _HIDDEN, device=device, dtype=dtype)
654
return EncoderCudaGraphCaptureInputs(
657
- mm_kwargs={
655
+ values={
656
"pixel_values": dummy_pixel_values,
659
- "image_grid_thw": grid_config,
657
+ "dummy_buf": dummy_buf,
658
},
661
- buffers={"dummy_buf": dummy_buf},
659
)
660
661
def prepare_encoder_cudagraph_replay_buffers(
@@ -670,14 +667,18 @@ class SimpleMockViTVideoModel(SimpleMockViTModel):
667
n_out = _count_output_tokens(self._get_grid_thw(mm_kwargs), _SPATIAL_MERGE)
668
p = next(self.parameters())
669
dummy_buf = torch.zeros(n_out, _HIDDEN, device=p.device, dtype=p.dtype)
673
- return EncoderCudaGraphReplayBuffers(buffers={"dummy_buf": dummy_buf})
670
+ return EncoderCudaGraphReplayBuffers(
671
+ values={
672
+ "pixel_values": self._get_pixel_values(mm_kwargs),
673
+ "dummy_buf": dummy_buf,
674
+ }
675
+ )
676
677
def encoder_cudagraph_forward(
678
self,
677
- mm_kwargs: dict[str, Any],
678
- buffers: dict[str, torch.Tensor],
679
+ values: dict[str, torch.Tensor],
680
) -> torch.Tensor:
680
- return self._forward(self._get_pixel_values(mm_kwargs))
681
+ return self._forward(values["pixel_values"])
682
683
def encoder_eager_forward(
684
self,
@@ -718,14 +719,6 @@ class TestGetInputModality:
719
}
720
assert model.get_input_modality(mm_kwargs) == "video"
721
721
- def test_video_model_config_has_both_modalities(self):
722
- model = SimpleMockViTVideoModel()
723
- cfg = model.get_encoder_cudagraph_config()
724
- assert "image" in cfg.modalities
725
- assert "video" in cfg.modalities
726
- assert cfg.input_key_by_modality["image"] == "pixel_values"
727
- assert cfg.input_key_by_modality["video"] == "pixel_values_videos"
728
-
722
723
# ---------------------------------------------------------------------------
724
# GPU tests — video capture, replay, fallback, and mixed image+video
vllm/model_executor/models/interfaces.py
+1
-2
@@ -1623,8 +1623,7 @@ class SupportsEncoderCudaGraph(Protocol):
1623
1624
def encoder_cudagraph_forward(
1625
self,
1626
- mm_kwargs: dict[str, Any],
1627
- buffers: dict[str, torch.Tensor],
1626
+ inputs: dict[str, torch.Tensor],
1627
) -> torch.Tensor:
1628
"""Run the encoder forward pass with precomputed buffers.
1629
vllm/model_executor/models/qwen2_5_vl.py
+15
-18
@@ -1057,7 +1057,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
1057
def forward(
1058
self,
1059
x: torch.Tensor,
1060
- grid_thw: list[list[int]],
1060
+ grid_thw: list[list[int]] | None,
1061
*,
1062
encoder_metadata: dict[str, torch.Tensor] | None = None,
1063
) -> torch.Tensor:
@@ -1712,11 +1712,8 @@ class Qwen2_5_VLForConditionalGeneration(
1712
)
1713
return EncoderCudaGraphConfig(
1714
modalities=modalities,
1715
- input_key_by_modality={
1716
- "image": "pixel_values",
1717
- "video": "pixel_values_videos",
1718
- },
1715
buffer_keys=[
1716
+ "pixel_values",
1717
"rotary_pos_emb_cos",
1718
"rotary_pos_emb_sin",
1719
"window_index",
@@ -1930,7 +1927,7 @@ class Qwen2_5_VLForConditionalGeneration(
1927
// self.visual.patch_size
1928
)
1929
max_seqlen_window_override = vit_merger_window_size**2 * (spatial_merge_size**2)
1933
- buffers = self.visual.prepare_encoder_metadata(
1930
+ metadata = self.visual.prepare_encoder_metadata(
1931
grid_config,
1932
max_batch_size=max_batch_size,
1933
max_frames_per_batch=max_frames_per_batch,
@@ -1942,14 +1939,12 @@ class Qwen2_5_VLForConditionalGeneration(
1939
1940
# Just use image-modality dummy input_buffer for capturing, since it's also
1941
# compatible for video inputs (has the same shape: [num_patches, C*T*P*P]).
1945
- mm_kwargs = {
1942
+ values = metadata | {
1943
"pixel_values": dummy_pixel_values,
1947
- "image_grid_thw": grid_config,
1944
}
1945
1946
return EncoderCudaGraphCaptureInputs(
1951
- mm_kwargs=mm_kwargs,
1952
- buffers=buffers,
1947
+ values=values,
1948
)
1949
1950
def prepare_encoder_cudagraph_replay_buffers(
@@ -1968,28 +1963,30 @@ class Qwen2_5_VLForConditionalGeneration(
1963
# bound and can over-pad window attention into many empty FlashAttention
1964
# CTAs.
1965
if modality == "image":
1971
- buffers = self.visual.prepare_encoder_metadata(
1966
+ metadata = self.visual.prepare_encoder_metadata(
1967
grid_thw_list,
1968
max_batch_size=max_batch_size,
1969
)
1970
elif modality == "video":
1976
- buffers = self.visual.prepare_encoder_metadata(
1971
+ metadata = self.visual.prepare_encoder_metadata(
1972
grid_thw_list,
1973
max_frames_per_batch=max_frames_per_batch,
1974
)
1975
else:
1976
raise AssertionError("This line should be unreachable.")
1977
1983
- return EncoderCudaGraphReplayBuffers(buffers=buffers)
1978
+ values = metadata | {
1979
+ "pixel_values": self._get_pixel_values_by_modality(mm_kwargs),
1980
+ }
1981
+ return EncoderCudaGraphReplayBuffers(values=values)
1982
1983
def encoder_cudagraph_forward(
1984
self,
1987
- mm_kwargs: dict[str, Any],
1988
- buffers: dict[str, torch.Tensor],
1985
+ values: dict[str, torch.Tensor],
1986
) -> torch.Tensor:
1990
- pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
1991
- grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
1992
- return self.visual(pixel_values, grid_thw, encoder_metadata=buffers)
1987
+ pixel_values = values.pop("pixel_values")
1988
+ metadata = values
1989
+ return self.visual(pixel_values, None, encoder_metadata=metadata)
1990
1991
def encoder_eager_forward(
1992
self,
vllm/model_executor/models/qwen2_vl.py
+20
-19
@@ -707,7 +707,7 @@ class Qwen2VisionTransformer(nn.Module):
707
def forward(
708
self,
709
x: torch.Tensor,
710
- grid_thw: torch.Tensor | list[list[int]],
710
+ grid_thw: torch.Tensor | list[list[int]] | None,
711
*,
712
encoder_metadata: dict[str, torch.Tensor] | None = None,
713
) -> torch.Tensor:
@@ -715,9 +715,11 @@ class Qwen2VisionTransformer(nn.Module):
715
x = x.to(device=self.device, dtype=self.dtype)
716
x = self.patch_embed(x)
717
718
- grid_thw_list = grid_thw if isinstance(grid_thw, list) else grid_thw.tolist()
719
-
718
if encoder_metadata is None:
719
+ assert grid_thw is not None
720
+ grid_thw_list = (
721
+ grid_thw if isinstance(grid_thw, list) else grid_thw.tolist()
722
+ )
723
encoder_metadata = self.prepare_encoder_metadata(grid_thw_list)
724
725
rotary_pos_emb_cos = encoder_metadata["rotary_pos_emb_cos"]
@@ -1460,11 +1462,8 @@ class Qwen2VLForConditionalGeneration(
1462
max_frames = self.get_max_frames_per_video()
1463
return EncoderCudaGraphConfig(
1464
modalities=["image", "video"],
1463
- input_key_by_modality={
1464
- "image": "pixel_values",
1465
- "video": "pixel_values_videos",
1466
- },
1465
buffer_keys=[
1466
+ "pixel_values",
1467
"rotary_pos_emb_cos",
1468
"rotary_pos_emb_sin",
1469
"cu_seqlens",
@@ -1622,7 +1621,7 @@ class Qwen2VLForConditionalGeneration(
1621
)
1622
1623
# max_seqlen.item() gets baked into the CUDA graph at capture time.
1625
- buffers = self.visual.prepare_encoder_metadata(
1624
+ metadata = self.visual.prepare_encoder_metadata(
1625
grid_config,
1626
max_batch_size=max_batch_size,
1627
max_frames_per_batch=max_frames_per_batch,
@@ -1632,14 +1631,12 @@ class Qwen2VLForConditionalGeneration(
1631
1632
# Capture with image-format kwargs; pixel_values shape is compatible with
1633
# both image and video replay paths.
1635
- mm_kwargs = {
1634
+ values = metadata | {
1635
"pixel_values": dummy_pixel_values,
1637
- "image_grid_thw": grid_config,
1636
}
1637
1638
return EncoderCudaGraphCaptureInputs(
1641
- mm_kwargs=mm_kwargs,
1642
- buffers=buffers,
1639
+ values=values,
1640
)
1641
1642
def prepare_encoder_cudagraph_replay_buffers(
@@ -1652,24 +1649,28 @@ class Qwen2VLForConditionalGeneration(
1649
grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
1650
1651
if modality == "image":
1655
- buffers = self.visual.prepare_encoder_metadata(
1652
+ metadata = self.visual.prepare_encoder_metadata(
1653
grid_thw_list,
1654
max_batch_size=max_batch_size,
1655
)
1656
else:
1660
- buffers = self.visual.prepare_encoder_metadata(
1657
+ metadata = self.visual.prepare_encoder_metadata(
1658
grid_thw_list,
1659
max_frames_per_batch=max_frames_per_batch,
1660
)
1661
1665
- return EncoderCudaGraphReplayBuffers(buffers=buffers)
1662
+ values = metadata | {
1663
+ "pixel_values": self._get_pixel_values_by_modality(mm_kwargs),
1664
+ }
1665
+ return EncoderCudaGraphReplayBuffers(values=values)
1666
1667
def encoder_cudagraph_forward(
1668
- self, mm_kwargs: dict[str, Any], buffers: dict[str, torch.Tensor]
1668
+ self,
1669
+ values: dict[str, torch.Tensor],
1670
) -> torch.Tensor:
1670
- pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
1671
- grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
1672
- return self.visual(pixel_values, grid_thw, encoder_metadata=buffers)
1671
+ pixel_values = values.pop("pixel_values")
1672
+ metadata = values
1673
+ return self.visual(pixel_values, None, encoder_metadata=metadata)
1674
1675
def encoder_eager_forward(
1676
self,
vllm/model_executor/models/qwen3_vl.py
+14
-17
@@ -1784,11 +1784,8 @@ class Qwen3VLForConditionalGeneration(
1784
1785
return EncoderCudaGraphConfig(
1786
modalities=modalities,
1787
- input_key_by_modality={
1788
- "image": "pixel_values",
1789
- "video": "pixel_values_videos",
1790
- },
1787
buffer_keys=[
1788
+ "pixel_values",
1789
"pos_embeds",
1790
"rotary_pos_emb_cos",
1791
"rotary_pos_emb_sin",
@@ -1978,7 +1975,7 @@ class Qwen3VLForConditionalGeneration(
1975
# so the capture value must cover any replay scenario.
1976
# Worst case: 1 item consuming the full budget ->
1977
# seq_len = token_budget * spatial_merge_size^2.
1981
- buffers = self.visual.prepare_encoder_metadata(
1978
+ metadata = self.visual.prepare_encoder_metadata(
1979
grid_config,
1980
max_batch_size=max_batch_size,
1981
max_frames_per_batch=max_frames_per_batch,
@@ -1988,14 +1985,12 @@ class Qwen3VLForConditionalGeneration(
1985
1986
# Just use image-modality dummy input_buffer for capturing, since it's also
1987
# compatible for video inputs (has the same shape: [num_patches, C*T*P*P]).
1991
- mm_kwargs = {
1988
+ values = metadata | {
1989
"pixel_values": dummy_pixel_values,
1993
- "image_grid_thw": grid_config,
1990
}
1991
1992
return EncoderCudaGraphCaptureInputs(
1997
- mm_kwargs=mm_kwargs,
1998
- buffers=buffers,
1993
+ values=values,
1994
)
1995
1996
def prepare_encoder_cudagraph_replay_buffers(
@@ -2008,28 +2003,30 @@ class Qwen3VLForConditionalGeneration(
2003
grid_thw_list = self._get_grid_thw_by_modality(mm_kwargs)
2004
2005
if modality == "image":
2011
- buffers = self.visual.prepare_encoder_metadata(
2006
+ metadata = self.visual.prepare_encoder_metadata(
2007
grid_thw_list,
2008
max_batch_size=max_batch_size,
2009
)
2010
elif modality == "video":
2016
- buffers = self.visual.prepare_encoder_metadata(
2011
+ metadata = self.visual.prepare_encoder_metadata(
2012
grid_thw_list,
2013
max_frames_per_batch=max_frames_per_batch,
2014
)
2015
else:
2016
raise AssertionError("This line should be unreachable.")
2017
2023
- return EncoderCudaGraphReplayBuffers(buffers=buffers)
2018
+ values = metadata | {
2019
+ "pixel_values": self._get_pixel_values_by_modality(mm_kwargs),
2020
+ }
2021
+ return EncoderCudaGraphReplayBuffers(values=values)
2022
2023
def encoder_cudagraph_forward(
2024
self,
2027
- mm_kwargs: dict[str, Any],
2028
- buffers: dict[str, torch.Tensor],
2025
+ values: dict[str, torch.Tensor],
2026
) -> torch.Tensor:
2030
- pixel_values = self._get_pixel_values_by_modality(mm_kwargs)
2031
- grid_thw = self._get_grid_thw_by_modality(mm_kwargs)
2032
- return self.visual(pixel_values, grid_thw, encoder_metadata=buffers)
2027
+ pixel_values = values.pop("pixel_values")
2028
+ metadata = values
2029
+ return self.visual(pixel_values, None, encoder_metadata=metadata)
2030
2031
def encoder_eager_forward(
2032
self,
vllm/model_executor/models/step3_vl.py
+13
-16
@@ -702,8 +702,10 @@ class Step3VLForConditionalGeneration(
702
703
return EncoderCudaGraphConfig(
704
modalities=["image"],
705
- input_key_by_modality={"image": "pixel_values"},
706
- buffer_keys=["patch_pixel_values"],
705
+ buffer_keys=[
706
+ "pixel_values",
707
+ "patch_pixel_values",
708
+ ],
709
out_hidden_size=self.config.hidden_size,
710
)
711
@@ -846,33 +848,27 @@ class Step3VLForConditionalGeneration(
848
device=device,
849
dtype=dtype,
850
)
849
- # num_patches is NOT in buffers -- the per-item merge is done
851
+ # num_patches is NOT in values -- the per-item merge is done
852
# CPU-side by finalize_encoder_cudagraph_output using the actual
853
# batch's num_patches from mm_kwargs.
852
- mm_kwargs = {
854
+ values = {
855
"pixel_values": dummy_pixel_values,
856
"patch_pixel_values": dummy_patch_pixel_values,
857
}
858
857
- buffers = {
858
- "patch_pixel_values": dummy_patch_pixel_values,
859
- }
860
-
859
return EncoderCudaGraphCaptureInputs(
862
- mm_kwargs=mm_kwargs,
863
- buffers=buffers,
860
+ values=values,
861
)
862
863
def encoder_cudagraph_forward(
864
self,
868
- mm_kwargs: dict[str, Any],
869
- buffers: dict[str, torch.Tensor],
865
+ values: dict[str, torch.Tensor],
866
) -> torch.Tensor:
867
# Graph captures only the compute (vision model + conv projector).
868
# Per-item merge happens CPU-side in finalize_encoder_cudagraph_output
869
# using actual num_patches from the batch data.
874
- pixel_values = mm_kwargs["pixel_values"]
875
- patch_pixel_values = buffers["patch_pixel_values"]
870
+ pixel_values = values["pixel_values"]
871
+ patch_pixel_values = values["patch_pixel_values"]
872
873
image_features = self._process_image_features(
874
self._get_vision_model_output(pixel_values)
@@ -974,10 +970,11 @@ class Step3VLForConditionalGeneration(
970
EncoderCudaGraphReplayBuffers,
971
)
972
977
- # Only patch_pixel_values lives in the buffers dict; num_patches is
973
+ # Only patch_pixel_values lives in the values dict; num_patches is
974
# processed CPU-side by finalize_encoder_cudagraph_output.
975
return EncoderCudaGraphReplayBuffers(
980
- buffers={
976
+ values={
977
+ "pixel_values": mm_kwargs["pixel_values"],
978
"patch_pixel_values": mm_kwargs["patch_pixel_values"],
979
},
980
)
vllm/v1/worker/encoder_cudagraph.py
+16
-39
@@ -32,23 +32,20 @@ class BudgetGraphMetadata:
32
"""Metadata for a single budget graph.
33
34
CUDA graph replay pattern:
35
- 1. Copy new batch data into input_buffer (e.g. pixel_values)
36
- 2. Copy precomputed values into metadata_buffers
37
- 3. Replay graph
38
- 4. Read encoder outputs from output_buffer
35
+ * Copy precomputed values into input_buffers
36
+ * Replay graph
37
+ * Read encoder outputs from output_buffer
38
"""
39
40
token_budget: int
41
max_batch_size: int # Max number of images/videos per batch
42
max_frames_per_batch: int # Max total frames per batch (for video)
43
graph: torch.cuda.CUDAGraph
45
- # The input tensor updated before replay (e.g. pixel_values)
46
- input_buffer: torch.Tensor
44
# Buffers recorded into the CUDA graph (e.g. embeddings, sequence metadata).
45
# Before replay the manager updates these in-place. By default buffers are
46
# zeroed before slice-copying the actual values; model-specific padding
47
# behavior is provided by EncoderCudaGraphConfig.padding_logics.
51
- metadata_buffers: dict[str, torch.Tensor]
48
+ input_buffers: dict[str, torch.Tensor]
49
# Output written by graph, read after replay
50
output_buffer: torch.Tensor
51
@@ -214,29 +211,23 @@ class EncoderCudaGraphManager:
211
self.dtype,
212
)
213
217
- mm_kwargs = capture_inputs.mm_kwargs
218
- buffers = capture_inputs.buffers
214
+ values = capture_inputs.values
215
216
with torch.inference_mode():
221
- output = self.model.encoder_cudagraph_forward(mm_kwargs, buffers)
217
+ output = self.model.encoder_cudagraph_forward({**values})
218
output_buffer = torch.empty_like(output)
219
220
graph = torch.cuda.CUDAGraph()
221
with torch.inference_mode(), torch.cuda.graph(graph):
226
- output = self.model.encoder_cudagraph_forward(mm_kwargs, buffers)
222
+ output = self.model.encoder_cudagraph_forward({**values})
223
output_buffer.copy_(output)
224
229
- # Since the image and video modalities share the same per-patch shape,
230
- # so we can use the image dummy inputs to capture CUDA graph for both
231
- # image and video.
232
- input_key = self.config.input_key_by_modality["image"]
225
self.budget_graphs[token_budget] = BudgetGraphMetadata(
226
token_budget=token_budget,
227
max_batch_size=self.max_batch_size,
228
max_frames_per_batch=self.max_frames_per_batch,
229
graph=graph,
238
- input_buffer=mm_kwargs[input_key],
239
- metadata_buffers=buffers,
230
+ input_buffers=values,
231
output_buffer=output_buffer,
232
)
233
@@ -273,15 +264,12 @@ class EncoderCudaGraphManager:
264
self,
265
mm_kwargs: dict[str, Any],
266
token_budget: int,
276
- replay_buffers: dict[str, torch.Tensor | None],
267
) -> torch.Tensor | None:
268
"""Execute budget graph.
269
270
Args:
271
mm_kwargs: Multimodal inputs for the batch.
272
token_budget: Token budget to use.
283
- replay_buffers: Buffer values to copy into captured buffers.
284
- None values leave the corresponding buffer unchanged.
273
274
Returns:
275
Encoder outputs, or None if graph not captured.
@@ -293,22 +281,18 @@ class EncoderCudaGraphManager:
281
282
graph_meta = self.budget_graphs[token_budget]
283
296
- # Copy the input tensor. Buffers are sized for the full budget;
297
- # actual inputs may be smaller. Zero then slice-copy so padded
298
- # positions are invisible to attention (cu_seqlens masks them out).
299
- input_key = self.config.input_key_by_modality[
300
- self.model.get_input_modality(mm_kwargs)
301
- ]
302
- src = mm_kwargs[input_key]
303
- n = src.shape[0]
304
- graph_meta.input_buffer[:n].copy_(src)
284
+ replay = self.model.prepare_encoder_cudagraph_replay_buffers(
285
+ mm_kwargs,
286
+ self.max_batch_size,
287
+ self.max_frames_per_batch,
288
+ )
289
290
# Copy metadata buffers using keys from config.buffer_keys.
291
for key in self.config.buffer_keys:
308
- src = replay_buffers.get(key)
292
+ src = replay.values.get(key)
293
if src is None:
294
continue
311
- buf = graph_meta.metadata_buffers[key]
295
+ buf = graph_meta.input_buffers[key]
296
if src.ndim == 0:
297
buf.copy_(src)
298
else:
@@ -430,16 +414,9 @@ class EncoderCudaGraphManager:
414
token_budget,
415
(token_budget - batch_out_tokens) / token_budget * 100,
416
)
433
- replay = self.model.prepare_encoder_cudagraph_replay_buffers(
434
- batch_mm_kwargs,
435
- self.max_batch_size,
436
- self.max_frames_per_batch,
437
- )
417
418
# graph_hits counted inside _run_budget_graph after replay.
440
- output = self._run_budget_graph(
441
- batch_mm_kwargs, token_budget, replay.buffers
442
- )
419
+ output = self._run_budget_graph(batch_mm_kwargs, token_budget)
420
assert output is not None
421
self.model.postprocess_encoder_output(
422
output,
vllm/v1/worker/encoder_cudagraph_defs.py
+2
-12
@@ -4,7 +4,6 @@
4
5
from collections.abc import Callable
6
from dataclasses import dataclass, field
7
-from typing import Any
7
8
import torch
9
@@ -40,11 +39,6 @@ class EncoderCudaGraphConfig:
39
modalities: list[str]
40
"""Supported modalities (e.g. ["image"])."""
41
43
- input_key_by_modality: dict[str, str]
44
- """Per-modality input tensor key mapping, e.g.
45
- {"image": "pixel_values", "video": "pixel_values_videos"}.
46
- """
47
-
42
buffer_keys: list[str]
43
"""Keys for the tensor buffers recorded into the CUDA graph.
44
Before replay the manager zeros then slice-copies new data
@@ -74,11 +68,7 @@ class EncoderCudaGraphCaptureInputs:
68
Returned by ``prepare_encoder_cudagraph_capture_inputs()``.
69
"""
70
77
- mm_kwargs: dict[str, Any]
78
- """Dummy forward inputs (model-specific keys).
79
- For Qwen3-VL this contains pixel_values and grid_thw."""
80
-
81
- buffers: dict[str, torch.Tensor]
71
+ values: dict[str, torch.Tensor]
72
"""Precomputed tensor buffers that will be recorded into the
73
CUDA graph. The manager stores references to these exact
74
tensor objects and copies new data into them before each
@@ -94,7 +84,7 @@ class EncoderCudaGraphReplayBuffers:
84
Keys match ``EncoderCudaGraphConfig.buffer_keys``.
85
"""
86
97
- buffers: dict[str, torch.Tensor | None]
87
+ values: dict[str, torch.Tensor | None]
88
"""Data to copy into the captured buffers before replay.
89
``None`` values leave the corresponding captured buffer
90
unchanged."""