A high-throughput and memory-efficient inference and serving engine for LLMs
[MM][Perf][CG] Support ViT full CUDA graph for Kimi-VL (#41992)
Signed-off-by: oguz <oguzhankir17@gmail.com>
Oğuzhan KIR committed
Jun 17, 2026 at 15:14 UTC
fa85ead2f378621fc412ac3679eae88a5604ebe8
5 files changed
+498
-39
docs/design/cuda_graphs_multimodal.md
+1
@@ -129,6 +129,7 @@ Models opt-in to encoder CUDA Graphs by implementing the [SupportsEncoderCudaGra
129
| `DeepseekOCRForCausalLM` | `DeepSeek-OCR` | ✅︎ | ❌︎ | ✅︎ |
130
| `Glm4vForConditionalGeneration` | `GLM-4.1V, GLM-4.6V-Flash` | ✅︎ | ✅︎ | ❌︎ |
131
| `InternVLChatModel` | `InternVL3.5`, `InternVL3`, `InternVL2.5`, `InternVL2` | ✅︎ | ✅︎ | ❌︎ |
132
+| `KimiVLForConditionalGeneration` | `Kimi-VL` | ✅︎ | ❌︎ | ❌︎ |
133
| `Llama4ForConditionalGeneration` | `Llama 4` | ✅︎ | ❌︎ | ❌︎ |
134
| `Qwen2VLForConditionalGeneration` | `Qwen2-VL` | ✅︎ | ✅︎ | ❌︎ |
135
| `Qwen2_5_VLForConditionalGeneration` | `Qwen2.5-VL` | ✅︎ | ✅︎ | ❌︎ |
examples/generate/multimodal/vision_language_offline.py
+1
@@ -2537,6 +2537,7 @@ MODELS_SUPPORT_VIT_CUDA_GRAPH = [
2537
"qwen2_5_vl",
2538
"qwen3_vl",
2539
"qwen3_vl_moe",
2540
+ "kimi_vl",
2541
"qwen3_5",
2542
"qwen3_5_moe",
2543
"internvl_chat",
tests/models/multimodal/generation/test_vit_cudagraph.py
+35
@@ -48,6 +48,13 @@ def internvl_chat_template(content: str) -> str:
48
return f"<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n"
49
50
51
+def kimi_vl_chat_template(content: str) -> str:
52
+ return (
53
+ f"<|im_user|>user<|im_middle|>{content}<|im_end|>"
54
+ "<|im_assistant|>assistant<|im_middle|>"
55
+ )
56
+
57
+
58
def step3_vl_chat_template(content: str) -> str:
59
return (
60
"<｜begin▁of▁sentence｜> You are a helpful assistant.<|BOT|>user\n "
@@ -100,6 +107,34 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = {
107
needs_video_metadata=False,
108
marks=[pytest.mark.core_model],
109
),
110
+ "kimi_vl": VitCudagraphTestConfig(
111
+ model="moonshotai/Kimi-VL-A3B-Instruct",
112
+ modalities=["image"],
113
+ image_prompt=kimi_vl_chat_template(
114
+ "<|media_start|>image<|media_content|><|media_pad|><|media_end|>"
115
+ "What is in this image?"
116
+ ),
117
+ needs_video_metadata=False,
118
+ # Single bucket sized to cover the test images' output tokens.
119
+ # The default auto-inferred range fans out into multiple power-of-2
120
+ # buckets, each holding a full ViT capture pool.
121
+ compilation_config_overrides={
122
+ "encoder_cudagraph_token_budgets": [1024],
123
+ },
124
+ # Shrink to 1 text + 1 vision layer with random weights so the
125
+ # test runs on any CI GPU (incl. L4) and skips the multi-GiB
126
+ # weight download. The test only validates that encoder CG
127
+ # capture/replay functions correctly, not output quality.
128
+ vllm_runner_kwargs={
129
+ "trust_remote_code": True,
130
+ "load_format": "dummy",
131
+ "hf_overrides": partial(
132
+ dummy_hf_overrides,
133
+ model_arch="KimiVLForConditionalGeneration",
134
+ ),
135
+ },
136
+ marks=[pytest.mark.core_model],
137
+ ),
138
"qwen3_vl": VitCudagraphTestConfig(
139
model="Qwen/Qwen3-VL-2B-Instruct",
140
image_prompt=qwen_vl_chat_template(
vllm/model_executor/models/kimi_vl.py
+195
-2
@@ -56,7 +56,11 @@ from vllm.config import VllmConfig
56
from vllm.config.multimodal import BaseDummyOptions
57
from vllm.inputs import MultiModalDataDict
58
from vllm.model_executor.layers.linear import ReplicatedLinear
59
-from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
59
+from vllm.model_executor.models.interfaces import (
60
+ SupportsEncoderCudaGraph,
61
+ SupportsMultiModal,
62
+ SupportsPP,
63
+)
64
from vllm.model_executor.models.moonvit import MoonVitPretrainedModel
65
from vllm.multimodal import MULTIMODAL_REGISTRY
66
from vllm.multimodal.inputs import (
@@ -79,6 +83,7 @@ from vllm.multimodal.processing import (
83
from vllm.sequence import IntermediateTensors
84
from vllm.transformers_utils.configs.kimi_vl import KimiVLConfig, MoonViTConfig
85
from vllm.utils.tensor_schema import TensorSchema, TensorShape
86
+from vllm.v1.worker.encoder_cudagraph_defs import EncoderCudaGraphReplayBuffers
87
88
from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
89
from .vision import is_vit_use_data_parallel, run_dp_sharded_mrope_vision_model
@@ -287,7 +292,9 @@ class KimiVLMultiModalProcessor(BaseMultiModalProcessor[KimiVLProcessingInfo]):
292
info=KimiVLProcessingInfo,
293
dummy_inputs=KimiVLDummyInputsBuilder,
294
)
290
-class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
295
+class KimiVLForConditionalGeneration(
296
+ nn.Module, SupportsMultiModal, SupportsEncoderCudaGraph, SupportsPP
297
+):
298
supports_encoder_tp_data = True
299
300
@classmethod
@@ -340,6 +347,192 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
347
348
self.media_placeholder: int = self.config.media_placeholder_token_id
349
350
+ self.model_config = model_config
351
+
352
+ # -- SupportsEncoderCudaGraph protocol methods --
353
+
354
+ def get_encoder_cudagraph_config(self):
355
+ from vllm.v1.worker.encoder_cudagraph_defs import (
356
+ EncoderCudaGraphConfig,
357
+ )
358
+
359
+ return EncoderCudaGraphConfig(
360
+ modalities=["image"],
361
+ buffer_keys=[
362
+ "pixel_values",
363
+ "pos_embeds",
364
+ "rope_freqs_cis",
365
+ "cu_seqlens",
366
+ "max_seqlen",
367
+ "merge_gather_idx",
368
+ ],
369
+ out_hidden_size=self.hidden_size,
370
+ )
371
+
372
+ def get_encoder_cudagraph_budget_range(
373
+ self,
374
+ vllm_config,
375
+ ) -> tuple[int, int]:
376
+ # Min: estimated smallest possible encoder input.
377
+ # 224x224 image with patch_size=14 -> 16x16 patches, then merge
378
+ # kernel (2,2) -> 8x8 = 64 output tokens.
379
+ min_budget = 64
380
+ max_budget = min(
381
+ vllm_config.scheduler_config.max_num_batched_tokens,
382
+ self.model_config.max_model_len,
383
+ )
384
+ return (min_budget, max_budget)
385
+
386
+ def _get_grid_hws(
387
+ self,
388
+ mm_kwargs: dict[str, Any],
389
+ ) -> list[tuple[int, int]]:
390
+ grid_hws = mm_kwargs["image_grid_hws"]
391
+ if not isinstance(grid_hws, list):
392
+ grid_hws = grid_hws.tolist()
393
+ return grid_hws
394
+
395
+ def get_encoder_cudagraph_item_specs(
396
+ self,
397
+ mm_kwargs: dict[str, Any],
398
+ ):
399
+ from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec
400
+
401
+ kh, kw = self.config.vision_config.merge_kernel_size
402
+ return [
403
+ EncoderItemSpec(
404
+ input_size=h * w,
405
+ output_tokens=(h // kh) * (w // kw),
406
+ )
407
+ for h, w in self._get_grid_hws(mm_kwargs)
408
+ ]
409
+
410
+ def select_encoder_cudagraph_items(
411
+ self,
412
+ mm_kwargs: dict[str, Any],
413
+ indices: list[int],
414
+ ) -> dict[str, Any]:
415
+ grid_hws = self._get_grid_hws(mm_kwargs)
416
+ pixel_values = mm_kwargs["pixel_values"]
417
+
418
+ if len(indices) == 0:
419
+ return {
420
+ "pixel_values": pixel_values[:0],
421
+ "image_grid_hws": pixel_values.new_zeros((0, 2), dtype=torch.long),
422
+ }
423
+
424
+ patches_per_item = [h * w for h, w in grid_hws]
425
+ cum_patches = [0]
426
+ for p in patches_per_item:
427
+ cum_patches.append(cum_patches[-1] + p)
428
+
429
+ selected_pv = torch.cat(
430
+ [pixel_values[cum_patches[i] : cum_patches[i + 1]] for i in indices]
431
+ )
432
+ selected_grid = torch.tensor(
433
+ [grid_hws[i] for i in indices],
434
+ dtype=torch.long,
435
+ device=pixel_values.device,
436
+ )
437
+ return {
438
+ "pixel_values": selected_pv,
439
+ "image_grid_hws": selected_grid,
440
+ }
441
+
442
+ def prepare_encoder_cudagraph_capture_inputs(
443
+ self,
444
+ token_budget: int,
445
+ max_batch_size: int,
446
+ max_frames_per_batch: int,
447
+ device: torch.device,
448
+ dtype: torch.dtype,
449
+ path: str = "default",
450
+ ):
451
+ from vllm.v1.worker.encoder_cudagraph_defs import (
452
+ EncoderCudaGraphCaptureInputs,
453
+ )
454
+
455
+ kh, kw = self.config.vision_config.merge_kernel_size
456
+ # Ceil so the buffer fits the worst case of one item using the full
457
+ # budget. Floor under-allocates when budget is not a multiple of
458
+ # max_batch_size.
459
+ per_mm_item_output = (token_budget + max_batch_size - 1) // max_batch_size
460
+
461
+ # Shape the synthetic grid so neither dimension exceeds Rope2DPosEmb's
462
+ # precomputed range. Pack as wide a row as fits, then add rows.
463
+ rope = self.vision_tower.encoder.rope_2d
464
+ max_wo = rope.max_width // kw
465
+ wo = min(per_mm_item_output, max_wo)
466
+ ho = (per_mm_item_output + wo - 1) // wo
467
+ assert ho * kh <= rope.max_height, (
468
+ f"per_mm_item_output={per_mm_item_output} exceeds RoPE grid capacity "
469
+ f"(max {(rope.max_height // kh) * (rope.max_width // kw)} tokens)"
470
+ )
471
+ grid_hws_list = [(ho * kh, wo * kw) for _ in range(max_batch_size)]
472
+
473
+ patch_size = self.config.vision_config.patch_size
474
+ if isinstance(patch_size, int):
475
+ patch_size = (patch_size, patch_size)
476
+
477
+ total_patches = sum(h * w for h, w in grid_hws_list)
478
+ in_channels = 3
479
+ dummy_pixel_values = torch.randn(
480
+ total_patches,
481
+ in_channels,
482
+ patch_size[0],
483
+ patch_size[1],
484
+ device=device,
485
+ dtype=dtype,
486
+ )
487
+
488
+ buffers = self.vision_tower.prepare_encoder_metadata(
489
+ grid_hws_list,
490
+ max_batch_size=max_batch_size,
491
+ max_seqlen_override=token_budget,
492
+ device=device,
493
+ )
494
+ values = buffers | {"pixel_values": dummy_pixel_values}
495
+
496
+ return EncoderCudaGraphCaptureInputs(values=values)
497
+
498
+ def prepare_encoder_cudagraph_replay_buffers(
499
+ self,
500
+ mm_kwargs: dict[str, Any],
501
+ max_batch_size: int,
502
+ max_frames_per_batch: int,
503
+ path: str = "default",
504
+ ):
505
+ grid_hws_list = self._get_grid_hws(mm_kwargs)
506
+ buffers = self.vision_tower.prepare_encoder_metadata(
507
+ grid_hws_list,
508
+ max_batch_size=max_batch_size,
509
+ device=mm_kwargs["pixel_values"].device,
510
+ )
511
+ values = buffers | {"pixel_values": mm_kwargs["pixel_values"]}
512
+ return EncoderCudaGraphReplayBuffers(values=values)
513
+
514
+ def encoder_cudagraph_forward(
515
+ self,
516
+ values: dict[str, torch.Tensor],
517
+ path: str = "default",
518
+ ) -> torch.Tensor:
519
+ pixel_values = values.pop("pixel_values")
520
+ metadata = values
521
+ image_features = self.vision_tower(
522
+ pixel_values, grid_hw=None, encoder_metadata=metadata
523
+ )
524
+ return self.multi_modal_projector(image_features)
525
+
526
+ def encoder_eager_forward(
527
+ self,
528
+ mm_kwargs: dict[str, Any],
529
+ path: str = "default",
530
+ ) -> torch.Tensor:
531
+ pixel_values = mm_kwargs["pixel_values"]
532
+ image_grid_hws = mm_kwargs["image_grid_hws"]
533
+ image_features = self.vision_tower(pixel_values, image_grid_hws)
534
+ return self.multi_modal_projector(torch.cat(image_features))
535
+
536
def _parse_and_validate_image_input(
537
self, **kwargs: object
538
) -> KimiVLImageInputs | None:
vllm/model_executor/models/moonvit.py
+266
-37
@@ -45,7 +45,9 @@
45
from collections.abc import Sequence
46
from copy import deepcopy
47
from functools import cached_property
48
+from typing import Any
49
50
+import numpy as np
51
import torch
52
import torch.nn as nn
53
import torch.nn.functional as F
@@ -110,23 +112,42 @@ class Learnable2DInterpPosEmb(nn.Module):
112
def reset_parameters(self):
113
nn.init.normal_(self.weight)
114
113
- def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
114
- pos_embs = []
115
- for shape in grid_hws.tolist():
116
- if shape == self.weight.shape[:-1]:
115
+ def get_pos_embeds(
116
+ self,
117
+ grid_hws_list: list[list[int]] | list[tuple[int, int]],
118
+ ) -> torch.Tensor:
119
+ """Build packed per-token positional embeddings for a list of grids.
120
+
121
+ Returns a tensor of shape ``(sum(h * w), dim)`` formed by interpolating
122
+ the learned ``(height, width, dim)`` weight to each ``(h, w)`` grid and
123
+ concatenating the flattened results in the same order as
124
+ ``grid_hws_list``. Lives outside the captured CUDA graph so the
125
+ per-grid Python iteration is safe.
126
+ """
127
+ weight_shape = list(self.weight.shape[:-1])
128
+ pos_embs: list[torch.Tensor] = []
129
+ for shape in grid_hws_list:
130
+ shape_list = [int(shape[0]), int(shape[1])]
131
+ if shape_list == weight_shape:
132
pos_embs.append(self.weight.flatten(end_dim=1))
133
else:
134
pos_embs.append(
135
F.interpolate(
136
self.weight.permute((2, 0, 1)).unsqueeze(0),
122
- size=shape,
137
+ size=tuple(shape_list),
138
mode=self.interpolation_mode,
139
)
140
.squeeze(0)
141
.permute((1, 2, 0))
142
.flatten(end_dim=1)
143
)
129
- out = x + torch.cat(pos_embs)
144
+ if not pos_embs:
145
+ return self.weight.new_zeros((0, self.weight.shape[-1]))
146
+ return torch.cat(pos_embs)
147
+
148
+ def forward(self, x: torch.Tensor, grid_hws: torch.Tensor) -> torch.Tensor:
149
+ pos_embs = self.get_pos_embeds(grid_hws.tolist())
150
+ out = x + pos_embs
151
return out
152
153
@@ -158,19 +179,29 @@ class MoonVisionPatchEmbed(nn.Module):
179
height=pos_emb_height, width=pos_emb_width, dim=out_dim
180
)
181
161
- def forward(self, x: torch.Tensor, grid_hw: torch.Tensor) -> torch.Tensor:
182
+ def forward(
183
+ self,
184
+ x: torch.Tensor,
185
+ grid_hw: torch.Tensor | None = None,
186
+ *,
187
+ pos_embeds: torch.Tensor | None = None,
188
+ ) -> torch.Tensor:
189
"""
190
Args:
191
x (L, Channels): input tensor
192
grid_hw (N, 2): grid height and width
193
+ pos_embeds: precomputed positional embeddings of shape
194
+ ``(L, Cout)``. When provided, ``grid_hw`` is unused and the
195
+ CUDA-graph-incompatible interpolation in ``self.pos_emb`` is
196
+ skipped.
197
198
Returns:
199
(L, Cout) tensor
200
"""
201
x = self.proj(x).view(x.size(0), -1)
171
- # apply positional embedding
172
- x = self.pos_emb(x, grid_hw)
173
- return x
202
+ if pos_embeds is not None:
203
+ return x + pos_embeds
204
+ return self.pos_emb(x, grid_hw)
205
206
207
class Rope2DPosEmb(nn.Module):
@@ -243,30 +274,44 @@ class Rope2DPosEmb(nn.Module):
274
freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1)
275
return freqs_cis
276
246
- def get_freqs_cis_by_seqlens(self, grid_hws: torch.Tensor) -> torch.Tensor:
247
- """
248
- Args:
249
- grid_hws (torch.Tensor): containing list of (height, width) or (t, height, width) tuples.
250
- Returns:
251
- freqs_cis: tensor of shape (sum(t * height * width), dim//2)
277
+ def get_freqs_cis_by_seqlens_list(
278
+ self,
279
+ grid_hws_list: list[list[int]] | list[tuple[int, int]],
280
+ ) -> torch.Tensor:
281
+ """List-based variant of :meth:`get_freqs_cis_by_seqlens`.
282
+
283
+ Accepts a Python list of ``(h, w)`` pairs so callers that already
284
+ operate outside the captured CUDA graph can avoid materializing a
285
+ tensor + ``.tolist()`` round-trip.
286
"""
253
- shapes = grid_hws.tolist()
287
assert all(
255
- 1 <= h <= self.max_height and 1 <= w <= self.max_width for h, w in shapes
288
+ 1 <= h <= self.max_height and 1 <= w <= self.max_width
289
+ for h, w in grid_hws_list
290
), (
257
- shapes,
291
+ grid_hws_list,
292
self.max_height,
293
self.max_width,
294
)
295
+ if not grid_hws_list:
296
+ return self.precomputed_freqs_cis.new_zeros((0, self.dim // 2))
297
freqs_cis = torch.cat(
298
[
299
self.precomputed_freqs_cis[:h, :w].reshape(-1, self.dim // 2)
264
- for h, w in shapes
300
+ for h, w in grid_hws_list
301
],
302
dim=0,
303
)
304
return freqs_cis
305
306
+ def get_freqs_cis_by_seqlens(self, grid_hws: torch.Tensor) -> torch.Tensor:
307
+ """
308
+ Args:
309
+ grid_hws (torch.Tensor): containing list of (height, width) or (t, height, width) tuples.
310
+ Returns:
311
+ freqs_cis: tensor of shape (sum(t * height * width), dim//2)
312
+ """
313
+ return self.get_freqs_cis_by_seqlens_list(grid_hws.tolist())
314
+
315
def get_freqs_cis_by_idx(
316
self, pos_idx: torch.Tensor, pos_idx_mask: torch.Tensor
317
) -> torch.Tensor:
@@ -392,11 +437,15 @@ class MoonVitEncoderLayer(nn.Module):
437
x: torch.Tensor,
438
cu_seqlens: torch.Tensor,
439
rope_freqs_cis: torch.Tensor | None = None,
440
+ max_seqlen: torch.Tensor | None = None,
441
):
442
"""
443
Args:
444
x (torch.Tensor): (seqlen, hidden_dim)
445
cu_seqlens (torch.Tensor):
446
+ max_seqlen: Optional precomputed scalar tensor. When omitted it
447
+ is derived from ``cu_seqlens``, which produces a GPU scalar
448
+ that breaks CUDA graph capture.
449
"""
450
seq_length = x.size(0)
451
xqkv, _ = self.wqkv(x)
@@ -412,7 +461,8 @@ class MoonVitEncoderLayer(nn.Module):
461
462
xq, xk = apply_rope(xq, xk, rope_freqs_cis)
463
415
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
464
+ if max_seqlen is None:
465
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
466
attn_out = self.attn(
467
xq.unsqueeze(0),
468
xk.unsqueeze(0),
@@ -433,10 +483,12 @@ class MoonVitEncoderLayer(nn.Module):
483
hidden_states: torch.Tensor,
484
cu_seqlens: torch.Tensor,
485
rope_freqs_cis: torch.Tensor | None = None,
486
+ max_seqlen: torch.Tensor | None = None,
487
) -> torch.Tensor:
488
"""
489
Args:
490
hidden_states: non-packed (B, N, D) or packed (L, D). if non-packed, seqlens should be None, if packed, seqlens should be set
491
+ max_seqlen: optional precomputed max-sequence-length scalar.
492
493
Returns:
494
output: same shape of input, non-packed (B, N, D) for non-packed input, (L, D) for packed input
@@ -444,7 +496,10 @@ class MoonVitEncoderLayer(nn.Module):
496
residual = hidden_states
497
hidden_states = self.norm0(hidden_states)
498
attn_out = self.attention_qkvpacked(
447
- hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
499
+ hidden_states,
500
+ cu_seqlens,
501
+ rope_freqs_cis=rope_freqs_cis,
502
+ max_seqlen=max_seqlen,
503
)
504
hidden_states = residual + attn_out
505
@@ -478,22 +533,39 @@ class MoonVitEncoder(nn.Module):
533
)
534
self.final_layernorm = nn.LayerNorm(hidden_dim)
535
481
- def forward(
482
- self, hidden_states: torch.Tensor, grid_hw: torch.Tensor
536
+ def get_rope_freqs_cis(
537
+ self,
538
+ grid_hws_list: list[list[int]] | list[tuple[int, int]],
539
) -> torch.Tensor:
484
- rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(grid_hws=grid_hw)
540
+ return self.rope_2d.get_freqs_cis_by_seqlens_list(grid_hws_list)
541
486
- lengths = torch.cat(
487
- (
488
- torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype),
489
- (grid_hw[:, 0] * grid_hw[:, 1]).to(hidden_states.device),
542
+ def forward(
543
+ self,
544
+ hidden_states: torch.Tensor,
545
+ grid_hw: torch.Tensor | None = None,
546
+ *,
547
+ cu_seqlens: torch.Tensor | None = None,
548
+ rope_freqs_cis: torch.Tensor | None = None,
549
+ max_seqlen: torch.Tensor | None = None,
550
+ ) -> torch.Tensor:
551
+ if rope_freqs_cis is None:
552
+ rope_freqs_cis = self.rope_2d.get_freqs_cis_by_seqlens(grid_hws=grid_hw)
553
+
554
+ if cu_seqlens is None:
555
+ lengths = torch.cat(
556
+ (
557
+ torch.zeros(1, device=hidden_states.device, dtype=grid_hw.dtype),
558
+ (grid_hw[:, 0] * grid_hw[:, 1]).to(hidden_states.device),
559
+ )
560
)
491
- )
492
- cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32)
561
+ cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32)
562
563
for _, block in enumerate(self.blocks):
564
hidden_states = block(
496
- hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
565
+ hidden_states,
566
+ cu_seqlens,
567
+ rope_freqs_cis=rope_freqs_cis,
568
+ max_seqlen=max_seqlen,
569
)
570
571
hidden_states = self.final_layernorm(hidden_states)
@@ -530,6 +602,54 @@ def patch_merger(
602
return outputs
603
604
605
+def patch_merger_packed(
606
+ x: torch.Tensor,
607
+ gather_idx: torch.Tensor,
608
+ merge_kernel_size: tuple[int, int],
609
+) -> torch.Tensor:
610
+ """CUDA-graph-safe equivalent of :func:`patch_merger`.
611
+
612
+ Uses a precomputed index tensor to gather the per-token reshape +
613
+ permute that ``patch_merger`` does inside a Python loop. The output is
614
+ the concatenated 3D tensor ``(sum(new_h * new_w), kh * kw, d_model)``,
615
+ matching what ``torch.cat(patch_merger(...))`` would produce.
616
+ """
617
+ kh, kw = merge_kernel_size
618
+ d_model = x.size(-1)
619
+ return x.index_select(0, gather_idx).view(-1, kh * kw, d_model)
620
+
621
+
622
+def _build_merge_gather_idx(
623
+ grid_hws_list: list[list[int]] | list[tuple[int, int]],
624
+ merge_kernel_size: tuple[int, int],
625
+) -> np.ndarray:
626
+ """Build the per-token gather indices used by :func:`patch_merger_packed`.
627
+
628
+ For each item with grid (h, w) and merge kernel (kh, kw), the output
629
+ block at position (nh, nw) gathers the kh*kw input tokens at rows
630
+ (nh*kh + ih, nw*kw + iw) of that item, in (ih, iw) row-major order.
631
+ """
632
+ kh, kw = merge_kernel_size
633
+ parts: list[np.ndarray] = []
634
+ pre_sum = 0
635
+ for h, w in grid_hws_list:
636
+ new_h, new_w = h // kh, w // kw
637
+ nh = np.arange(new_h, dtype=np.int64).reshape(new_h, 1, 1, 1)
638
+ nw = np.arange(new_w, dtype=np.int64).reshape(1, new_w, 1, 1)
639
+ ih = np.arange(kh, dtype=np.int64).reshape(1, 1, kh, 1)
640
+ iw = np.arange(kw, dtype=np.int64).reshape(1, 1, 1, kw)
641
+ # Linearized input row = (nh*kh + ih) * w + (nw*kw + iw), offset by
642
+ # the per-item base ``pre_sum``. Output is laid out as
643
+ # (new_h, new_w, kh, kw) which patch_merger flattens to
644
+ # (new_h*new_w, kh*kw).
645
+ idx = pre_sum + (nh * kh + ih) * w + (nw * kw + iw)
646
+ parts.append(idx.reshape(-1))
647
+ pre_sum += h * w
648
+ if not parts:
649
+ return np.zeros(0, dtype=np.int64)
650
+ return np.concatenate(parts)
651
+
652
+
653
class MoonVitPretrainedModel(PreTrainedModel):
654
config_class = MoonViTConfig
655
model_type = "moonvit"
@@ -570,17 +690,126 @@ class MoonVitPretrainedModel(PreTrainedModel):
690
prefix=f"{prefix}.encoder",
691
)
692
693
+ def prepare_encoder_metadata(
694
+ self,
695
+ grid_hws_list: list[list[int]] | list[tuple[int, int]],
696
+ *,
697
+ max_batch_size: int | None = None,
698
+ max_seqlen_override: int | None = None,
699
+ device: torch.device | None = None,
700
+ ) -> dict[str, Any]:
701
+ """Precompute every grid-dependent input the encoder needs.
702
+
703
+ Used by the CUDA graph capture and replay paths to precompute
704
+ every grid-dependent input outside the captured graph, so per-grid
705
+ Python iteration and ``.tolist()`` round-trips are fine; the
706
+ values are then copied into fixed-shape buffers for replay.
707
+
708
+ Args:
709
+ grid_hws_list: List of ``(h, w)`` patch-grid sizes per image.
710
+ max_batch_size: When set, ``cu_seqlens`` is right-padded with
711
+ its last value so the buffer covers up to this many
712
+ sequences. Required at CUDA graph capture/replay so the
713
+ buffer shape matches what was recorded; padding entries
714
+ are zero-length sequences and are ignored by varlen
715
+ attention.
716
+ max_seqlen_override: Override the per-replay max sequence
717
+ length scalar. At capture this must be a safe upper bound
718
+ (worst case: a single image consuming the full token
719
+ budget) because the value is baked into the captured
720
+ graph.
721
+ device: Device for the metadata tensors. Defaults to the
722
+ model's parameter device.
723
+ """
724
+ if device is None:
725
+ device = next(self.parameters()).device
726
+
727
+ # Normalize to a list of plain Python int pairs so the helpers
728
+ # below never need ``.tolist()`` on a tensor.
729
+ grid_pairs: list[tuple[int, int]] = [(int(h), int(w)) for h, w in grid_hws_list]
730
+
731
+ metadata: dict[str, Any] = {}
732
+
733
+ pos_embeds = self.patch_embed.pos_emb.get_pos_embeds(grid_pairs)
734
+ metadata["pos_embeds"] = pos_embeds.to(device=device)
735
+
736
+ rope_freqs_cis = self.encoder.get_rope_freqs_cis(grid_pairs)
737
+ metadata["rope_freqs_cis"] = rope_freqs_cis.to(device=device)
738
+
739
+ grid_arr = np.array(grid_pairs, dtype=np.int64)
740
+ seq_lens = (grid_arr[:, 0] * grid_arr[:, 1]).astype(np.int32)
741
+ cu_seqlens_np = np.concatenate(
742
+ [
743
+ np.zeros(1, dtype=np.int32),
744
+ seq_lens.cumsum(dtype=np.int32),
745
+ ]
746
+ )
747
+
748
+ if max_batch_size is not None:
749
+ num_seqs = len(cu_seqlens_np) - 1
750
+ if num_seqs < max_batch_size:
751
+ cu_seqlens_np = np.concatenate(
752
+ [
753
+ cu_seqlens_np,
754
+ np.full(
755
+ max_batch_size - num_seqs,
756
+ cu_seqlens_np[-1],
757
+ dtype=np.int32,
758
+ ),
759
+ ]
760
+ )
761
+ metadata["cu_seqlens"] = torch.from_numpy(cu_seqlens_np).to(device)
762
+
763
+ if max_seqlen_override is not None:
764
+ max_seqlen_val = int(max_seqlen_override)
765
+ else:
766
+ max_seqlen_val = int(seq_lens.max()) if len(seq_lens) > 0 else 0
767
+ # Keep on CPU: attention wrappers may call .item() on this scalar
768
+ # and we want that materialization to happen outside the captured
769
+ # graph (the value is constant per capture anyway).
770
+ metadata["max_seqlen"] = torch.tensor(max_seqlen_val, dtype=torch.int32)
771
+
772
+ gather_idx_np = _build_merge_gather_idx(grid_pairs, self.merge_kernel_size)
773
+ metadata["merge_gather_idx"] = torch.from_numpy(gather_idx_np).to(device)
774
+
775
+ return metadata
776
+
777
def forward(
574
- self, pixel_values: torch.Tensor, grid_hw: torch.Tensor
575
- ) -> torch.Tensor:
778
+ self,
779
+ pixel_values: torch.Tensor,
780
+ grid_hw: torch.Tensor,
781
+ *,
782
+ encoder_metadata: dict[str, Any] | None = None,
783
+ ) -> torch.Tensor | list[torch.Tensor]:
784
"""
785
Args:
786
pixel_values (torch.Tensor): The input pixel values.
787
grid_hw (torch.Tensor): The grid height and width.
580
-
581
- Returns:
582
- torch.Tensor: The output tokens.
788
+ encoder_metadata: Optional precomputed metadata produced by
789
+ :meth:`prepare_encoder_metadata`. When provided every
790
+ ``.tolist()`` call in the forward path is skipped, the
791
+ returned tensor is the packed
792
+ ``(sum(new_h*new_w), kh*kw, hidden_size)`` form (suitable
793
+ for CUDA graph capture/replay), and ``grid_hw`` is unused.
794
+ When ``None`` the legacy path runs and returns a list of
795
+ per-image tensors.
796
"""
797
+ if encoder_metadata is not None:
798
+ hidden_states = self.patch_embed(
799
+ pixel_values, pos_embeds=encoder_metadata["pos_embeds"]
800
+ )
801
+ hidden_states = self.encoder(
802
+ hidden_states,
803
+ cu_seqlens=encoder_metadata["cu_seqlens"],
804
+ rope_freqs_cis=encoder_metadata["rope_freqs_cis"],
805
+ max_seqlen=encoder_metadata["max_seqlen"],
806
+ )
807
+ return patch_merger_packed(
808
+ hidden_states,
809
+ encoder_metadata["merge_gather_idx"],
810
+ merge_kernel_size=self.merge_kernel_size,
811
+ )
812
+
813
hidden_states = self.patch_embed(pixel_values, grid_hw)
814
hidden_states = self.encoder(hidden_states, grid_hw)
815
hidden_states = patch_merger(