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 1, 23:45 UTC
Fix Kimi-K2.5 FlashInfer ViT metadata (#44493)
Signed-off-by: Kevin-XiongC <kevin_xiong1997@outlook.com>
Kevin_Xiong committed
Jun 4, 2026 at 16:14 UTC
1bdc60ed53adcb2f13e09f6b6a13a9b93c68c5fc
2 files changed
+109
-28
vllm/model_executor/models/kimi_k25.py
+1
-1
@@ -235,7 +235,7 @@ class KimiK25MultiModalProcessor(BaseMultiModalProcessor[KimiK25ProcessingInfo])
235
pixel_values=MultiModalFieldConfig.flat_from_sizes(
236
"vision_chunk", grid_sizes
237
),
238
- grid_thws=MultiModalFieldConfig.batched("vision_chunk"),
238
+ grid_thws=MultiModalFieldConfig.batched("vision_chunk", keep_on_cpu=True),
239
)
240
241
def _call_hf_processor(
vllm/model_executor/models/kimi_k25_vit.py
+108
-27
@@ -154,9 +154,12 @@ class Learnable2DInterpPosEmbDivided_fixed(nn.Module):
154
def reset_parameters(self):
155
nn.init.normal_(self.weight)
156
157
- def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor:
157
+ def forward(
158
+ self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]]
159
+ ) -> torch.Tensor:
160
pos_embs = []
159
- for t, h, w in grid_thws.tolist():
161
+ grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
162
+ for t, h, w in grid_thw_list:
163
assert t <= self.num_frames, f"t:{t} > self.num_frames:{self.num_frames}"
164
if (h, w) == self.weight.shape[:-1]:
165
pos_emb_2d = self.weight.flatten(end_dim=1)
@@ -218,7 +221,9 @@ class MoonVision3dPatchEmbed(nn.Module):
221
else:
222
raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}")
223
221
- def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor:
224
+ def forward(
225
+ self, x: torch.Tensor, grid_thws: torch.Tensor | list[list[int]]
226
+ ) -> torch.Tensor:
227
x = self.proj(x).view(x.size(0), -1)
228
# apply positional embedding
229
x = self.pos_emb(x, grid_thws)
@@ -265,7 +270,7 @@ class Rope2DPosEmbRepeated(nn.Module):
270
return freqs_cis
271
272
def get_freqs_cis(
268
- self, grid_thws: torch.Tensor, device: torch.device
273
+ self, grid_thws: torch.Tensor | list[list[int]], device: torch.device
274
) -> torch.Tensor:
275
"""
276
Args:
@@ -279,7 +284,7 @@ class Rope2DPosEmbRepeated(nn.Module):
284
"freqs_cis", self._precompute_freqs_cis(device), persistent=False
285
)
286
282
- shapes = grid_thws.tolist()
287
+ shapes = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
288
assert all(
289
1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes
290
), (
@@ -401,6 +406,8 @@ class MoonViTEncoderLayer(nn.Module):
406
x: torch.Tensor,
407
cu_seqlens: torch.Tensor,
408
rope_freqs_cis: torch.Tensor | None = None,
409
+ max_seqlen: torch.Tensor | None = None,
410
+ sequence_lengths: torch.Tensor | None = None,
411
):
412
"""Compute self-attention with packed QKV.
413
@@ -422,13 +429,15 @@ class MoonViTEncoderLayer(nn.Module):
429
430
xq, xk = apply_rope(xq, xk, rope_freqs_cis)
431
425
- max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
432
+ if max_seqlen is None:
433
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
434
attn_out = self.attn(
435
xq.unsqueeze(0),
436
xk.unsqueeze(0),
437
xv.unsqueeze(0),
438
cu_seqlens=cu_seqlens,
439
max_seqlen=max_seqlen,
440
+ sequence_lengths=sequence_lengths,
441
)
442
attn_out = attn_out.reshape(
443
seq_length,
@@ -443,12 +452,18 @@ class MoonViTEncoderLayer(nn.Module):
452
hidden_states: torch.Tensor,
453
cu_seqlens: torch.Tensor,
454
rope_freqs_cis: torch.Tensor | None = None,
455
+ max_seqlen: torch.Tensor | None = None,
456
+ sequence_lengths: torch.Tensor | None = None,
457
):
458
residual = hidden_states
459
hidden_states = self.norm0(hidden_states)
460
461
hidden_states = self.attention_qkvpacked(
451
- hidden_states, cu_seqlens, rope_freqs_cis
462
+ hidden_states,
463
+ cu_seqlens,
464
+ rope_freqs_cis,
465
+ max_seqlen=max_seqlen,
466
+ sequence_lengths=sequence_lengths,
467
)
468
hidden_states = residual + hidden_states
469
@@ -493,27 +508,70 @@ class MoonViT3dEncoder(nn.Module):
508
)
509
self.final_layernorm = nn.LayerNorm(hidden_dim)
510
511
+ def prepare_encoder_metadata(
512
+ self,
513
+ grid_thw_list: list[list[int]],
514
+ *,
515
+ device: torch.device,
516
+ ) -> dict[str, torch.Tensor | None]:
517
+ metadata: dict[str, torch.Tensor | None] = {}
518
+ metadata["rope_freqs_cis"] = self.rope_2d.get_freqs_cis(
519
+ grid_thw_list, device=device
520
+ )
521
+
522
+ grid_thw_np = np.array(grid_thw_list, dtype=np.int32)
523
+ lengths = grid_thw_np[:, 0] * grid_thw_np[:, 1] * grid_thw_np[:, 2]
524
+ cu_seqlens = np.concatenate(
525
+ [np.zeros(1, dtype=np.int32), lengths.cumsum(dtype=np.int32)]
526
+ )
527
+
528
+ attn_backend = self.blocks[0].attn.attn_backend
529
+ metadata["sequence_lengths"] = MMEncoderAttention.maybe_compute_seq_lens(
530
+ attn_backend, cu_seqlens, device
531
+ )
532
+ metadata["max_seqlen"] = torch.tensor(
533
+ MMEncoderAttention.compute_max_seqlen(attn_backend, cu_seqlens),
534
+ dtype=torch.int32,
535
+ )
536
+ metadata["cu_seqlens"] = MMEncoderAttention.maybe_recompute_cu_seqlens(
537
+ attn_backend,
538
+ cu_seqlens,
539
+ self.blocks[0].hidden_dim,
540
+ self.blocks[0].tp_size,
541
+ device,
542
+ )
543
+ return metadata
544
+
545
def forward(
546
self,
547
hidden_states: torch.Tensor,
499
- grid_thws: torch.Tensor,
548
+ grid_thws: torch.Tensor | list[list[int]],
549
+ *,
550
+ encoder_metadata: dict[str, torch.Tensor | None] | None = None,
551
) -> torch.Tensor:
501
- rope_freqs_cis = self.rope_2d.get_freqs_cis(
502
- grid_thws=grid_thws, device=hidden_states.device
503
- )
504
-
505
- lengths = torch.cat(
506
- (
507
- torch.zeros(1, dtype=grid_thws.dtype, device=grid_thws.device),
508
- grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2],
552
+ if encoder_metadata is None:
553
+ grid_thw_list = (
554
+ grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
555
+ )
556
+ encoder_metadata = self.prepare_encoder_metadata(
557
+ grid_thw_list, device=hidden_states.device
558
)
510
- )
559
512
- cu_seqlens = lengths.to(hidden_states.device).cumsum(dim=0, dtype=torch.int32)
560
+ rope_freqs_cis = encoder_metadata["rope_freqs_cis"]
561
+ cu_seqlens = encoder_metadata["cu_seqlens"]
562
+ max_seqlen = encoder_metadata["max_seqlen"]
563
+ sequence_lengths = encoder_metadata.get("sequence_lengths")
564
+ assert rope_freqs_cis is not None
565
+ assert cu_seqlens is not None
566
+ assert max_seqlen is not None
567
568
for block in self.blocks:
569
hidden_states = block(
516
- hidden_states, cu_seqlens, rope_freqs_cis=rope_freqs_cis
570
+ hidden_states,
571
+ cu_seqlens,
572
+ rope_freqs_cis=rope_freqs_cis,
573
+ max_seqlen=max_seqlen,
574
+ sequence_lengths=sequence_lengths,
575
)
576
577
hidden_states = self.final_layernorm(hidden_states)
@@ -523,16 +581,17 @@ class MoonViT3dEncoder(nn.Module):
581
582
def tpool_patch_merger(
583
x: torch.Tensor,
526
- grid_thws: torch.Tensor,
584
+ grid_thws: torch.Tensor | list[list[int]],
585
merge_kernel_size: tuple[int, int] = (2, 2),
586
) -> list[torch.Tensor]:
587
"""Temporal pooling patch merger."""
588
kh, kw = merge_kernel_size
531
- lengths = (grid_thws[:, 0] * grid_thws[:, 1] * grid_thws[:, 2]).tolist()
589
+ grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
590
+ lengths = [t * h * w for t, h, w in grid_thw_list]
591
seqs = x.split(lengths, dim=0)
592
593
outputs = []
535
- for seq, (t, h, w) in zip(seqs, grid_thws.tolist()):
594
+ for seq, (t, h, w) in zip(seqs, grid_thw_list):
595
nh, nw = h // kh, w // kw
596
# Reshape: (t*h*w, d) -> (t, nh, kh, nw, kw, d)
597
v = seq.view(t, nh, kh, nw, kw, -1)
@@ -589,7 +648,11 @@ class MoonViT3dPretrainedModel(nn.Module):
648
)
649
650
def forward(
592
- self, pixel_values: torch.Tensor, grid_thws: torch.Tensor
651
+ self,
652
+ pixel_values: torch.Tensor,
653
+ grid_thws: torch.Tensor | list[list[int]],
654
+ *,
655
+ encoder_metadata: dict[str, torch.Tensor | None] | None = None,
656
) -> torch.Tensor:
657
"""
658
Args:
@@ -599,13 +662,23 @@ class MoonViT3dPretrainedModel(nn.Module):
662
Returns:
663
torch.Tensor: The output tokens.
664
"""
602
- hidden_states = self.patch_embed(pixel_values, grid_thws)
603
- hidden_states = self.encoder(hidden_states, grid_thws)
665
+ grid_thw_list = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist()
666
+ if encoder_metadata is None:
667
+ encoder_metadata = self.encoder.prepare_encoder_metadata(
668
+ grid_thw_list, device=pixel_values.device
669
+ )
670
+
671
+ hidden_states = self.patch_embed(pixel_values, grid_thw_list)
672
+ hidden_states = self.encoder(
673
+ hidden_states,
674
+ grid_thw_list,
675
+ encoder_metadata=encoder_metadata,
676
+ )
677
if (
678
self.merge_type == "sd2_tpool"
679
): # spatial downsampling 2x with temporal pooling all
680
hidden_states = tpool_patch_merger(
608
- hidden_states, grid_thws, merge_kernel_size=self.merge_kernel_size
681
+ hidden_states, grid_thw_list, merge_kernel_size=self.merge_kernel_size
682
)
683
else:
684
raise NotImplementedError(f"Not support {self.merge_type}")
@@ -649,7 +722,15 @@ def vision_tower_forward(
722
rope_type="rope_2d",
723
)
724
else:
652
- vt_outputs = vision_tower(pixel_values, grid_thw)
725
+ grid_thw_list = grid_thw.tolist()
726
+ encoder_metadata = vision_tower.encoder.prepare_encoder_metadata(
727
+ grid_thw_list, device=pixel_values.device
728
+ )
729
+ vt_outputs = vision_tower(
730
+ pixel_values,
731
+ grid_thw_list,
732
+ encoder_metadata=encoder_metadata,
733
+ )
734
tensors = mm_projector_forward(mm_projector, list(vt_outputs))
735
return list(tensors)
736