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 10, 16:30 UTC
[Models]: Use `MMEncoderAttention` for MoonViT (#31738)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: h100 <h100@inferact.ai> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: h100 <h100@inferact.ai>
Isotr0py committed
Jan 6, 2026 at 16:00 UTC
7101e0851f73b8157145cf04835a87dc0837d450
2 files changed
+72
-158
vllm/model_executor/models/kimi_vl.py
+1
-1
@@ -325,7 +325,7 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):
325
self.hidden_size = config.text_config.hidden_size
326
self.vision_tower = MoonVitPretrainedModel(
327
config.vision_config,
328
- self.use_data_parallel,
328
+ multimodal_config=model_config.multimodal_config,
329
prefix=maybe_prefix(prefix, "vision_tower"),
330
)
331
vllm/model_executor/models/moonvit.py
+71
-157
@@ -51,118 +51,20 @@ import torch.nn as nn
51
import torch.nn.functional as F
52
from transformers.activations import ACT2FN
53
from transformers.modeling_utils import PreTrainedModel
54
-from transformers.utils import is_flash_attn_2_available
54
55
+from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
56
+from vllm.config import MultiModalConfig
57
+from vllm.distributed import divide, get_tensor_model_parallel_world_size
58
from vllm.model_executor.layers.conv import Conv2dLayer
57
-from vllm.model_executor.layers.linear import ReplicatedLinear
59
+from vllm.model_executor.layers.linear import (
60
+ ColumnParallelLinear,
61
+ QKVParallelLinear,
62
+ RowParallelLinear,
63
+)
64
from vllm.model_executor.models.utils import maybe_prefix
65
from vllm.platforms import current_platform
66
from vllm.transformers_utils.configs.moonvit import MoonViTConfig
67
62
-if is_flash_attn_2_available():
63
- from flash_attn import flash_attn_varlen_func
64
-elif current_platform.is_xpu():
65
- from vllm.attention.utils.fa_utils import flash_attn_varlen_func
66
-else:
67
- flash_attn_varlen_func = None
68
-
69
-
70
-def multihead_attention(
71
- q: torch.Tensor,
72
- k: torch.Tensor,
73
- v: torch.Tensor,
74
- q_cu_seqlens: torch.Tensor | None = None,
75
- k_cu_seqlens: torch.Tensor | None = None,
76
-) -> torch.Tensor:
77
- """Multi-head attention using flash attention 2.
78
-
79
- Args:
80
- q: Query tensor of shape (batch_size, seqlen, num_heads, head_dim),
81
- or (tot_seqlens, num_heads, head_dim) if packing.
82
- k: Key tensor of shape (batch_size, seqlen, num_heads, head_dim),
83
- or (tot_seqlens, num_heads, head_dim) if packing.
84
- v: Value tensor of shape (batch_size, seqlen, num_heads, head_dim),
85
- or (tot_seqlens, num_heads, head_dim) if packing.
86
- q_cu_seqlens (torch.Tensor): cumulative sequence lengths of q.
87
- The first element should be 0 and the last element should be q.shape[0].
88
- k_cu_seqlens (torch.Tensor): cumulative sequence lengths of k.
89
- The first element should be 0 and the last element should be k.shape[0].
90
-
91
- Returns:
92
- output: shape (batch_size, seqlen, dim) or (tot_seqlens, dim) if packing,
93
- where dim = num_heads * head_dim
94
- """
95
- # Unified format legal check
96
- assert q.dim() == k.dim() == v.dim() == 3, "q, k, v must have 3 dims"
97
- assert q_cu_seqlens[-1] == q.shape[0], "q_cu_seqlens must sum to q.shape[0]"
98
- assert k_cu_seqlens[-1] == k.shape[0] == v.shape[0], (
99
- "k_cu_seqlens must sum to k.shape[0]"
100
- )
101
- assert q.dtype in [
102
- torch.bfloat16,
103
- torch.float16,
104
- ], f"unsupported dtype {q.dtype} for multihead attn"
105
-
106
- max_seqlen_q = (q_cu_seqlens[1:] - q_cu_seqlens[:-1]).max().item()
107
- max_seqlen_k = (k_cu_seqlens[1:] - k_cu_seqlens[:-1]).max().item()
108
- attn_out = flash_attn_varlen_func(
109
- q,
110
- k,
111
- v,
112
- cu_seqlens_q=q_cu_seqlens,
113
- cu_seqlens_k=k_cu_seqlens,
114
- max_seqlen_q=max_seqlen_q,
115
- max_seqlen_k=max_seqlen_k,
116
- causal=False,
117
- )
118
- attn_out = attn_out.flatten(start_dim=-2)
119
-
120
- return attn_out
121
-
122
-
123
-def sdpa_attention(
124
- q: torch.Tensor,
125
- k: torch.Tensor,
126
- v: torch.Tensor,
127
- q_cu_seqlens: torch.Tensor | None = None,
128
- k_cu_seqlens: torch.Tensor | None = None,
129
-) -> torch.Tensor:
130
- """SDPA attention.
131
-
132
- Args:
133
- q: Query tensor of shape (batch_size, seqlen, num_heads, head_dim),
134
- or (tot_seqlens, num_heads, head_dim) if packing.
135
- k: Key tensor of shape (batch_size, seqlen, num_heads, head_dim),
136
- or (tot_seqlens, num_heads, head_dim) if packing.
137
- v: Value tensor of shape (batch_size, seqlen, num_heads, head_dim),
138
- or (tot_seqlens, num_heads, head_dim) if packing.
139
- q_cu_seqlens: Optional cumulative sequence lengths of q.
140
- k_cu_seqlens: Optional cumulative sequence lengths of k.
141
- """
142
- seq_length = q.shape[0]
143
- attention_mask = torch.zeros(
144
- [1, seq_length, seq_length], device=q.device, dtype=torch.bool
145
- )
146
- for i in range(1, len(q_cu_seqlens)):
147
- attention_mask[
148
- ...,
149
- q_cu_seqlens[i - 1] : q_cu_seqlens[i],
150
- q_cu_seqlens[i - 1] : q_cu_seqlens[i],
151
- ] = True
152
- q = q.transpose(0, 1)
153
- k = k.transpose(0, 1)
154
- v = v.transpose(0, 1)
155
- attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
156
- attn_output = attn_output.transpose(0, 1)
157
- attn_output = attn_output.reshape(seq_length, -1)
158
- return attn_output
159
-
160
-
161
-VL_VISION_ATTENTION_FUNCTIONS = {
162
- "flash_attention_2": multihead_attention,
163
- "sdpa": sdpa_attention,
164
-}
165
-
68
69
def _apply_rope_input_validation(x, freqs_cis):
70
assert x.ndim == freqs_cis.ndim + 1, (x.shape, freqs_cis.shape)
@@ -411,11 +313,19 @@ class MLP2(nn.Module):
313
super().__init__()
314
assert len(dims) == 3
315
self.use_data_parallel = use_data_parallel
414
- self.fc0 = ReplicatedLinear(
415
- dims[0], dims[1], bias=bias, prefix=maybe_prefix(prefix, "fc0")
316
+ self.fc0 = ColumnParallelLinear(
317
+ dims[0],
318
+ dims[1],
319
+ bias=bias,
320
+ prefix=maybe_prefix(prefix, "fc0"),
321
+ disable_tp=self.use_data_parallel,
322
)
417
- self.fc1 = ReplicatedLinear(
418
- dims[1], dims[2], bias=bias, prefix=maybe_prefix(prefix, "fc1")
323
+ self.fc1 = RowParallelLinear(
324
+ dims[1],
325
+ dims[2],
326
+ bias=bias,
327
+ prefix=maybe_prefix(prefix, "fc1"),
328
+ disable_tp=self.use_data_parallel,
329
)
330
self.activation = activation
331
@@ -433,35 +343,55 @@ class MoonVitEncoderLayer(nn.Module):
343
hidden_dim: int,
344
mlp_dim: int,
345
prefix: str = "",
436
- use_data_parallel: bool = False,
346
+ multimodal_config: MultiModalConfig | None = None,
347
*,
438
- attn_implementation: str = "sdpa",
348
activation=F.gelu,
349
attn_bias: bool = False,
350
):
351
super().__init__()
352
+ self.use_data_parallel = (
353
+ multimodal_config.mm_encoder_tp_mode == "data"
354
+ if multimodal_config
355
+ else False
356
+ )
357
+
358
self.num_heads = num_heads
359
self.hidden_dim = hidden_dim
360
self.hidden_size_per_attention_head = self.hidden_dim // self.num_heads
446
- self.attn_implementation = attn_implementation
447
- # use fa2 in vllm by default
448
- if is_flash_attn_2_available() or current_platform.is_xpu():
449
- self.attn_implementation = "flash_attention_2"
361
+ self.tp_size = (
362
+ 1 if self.use_data_parallel else get_tensor_model_parallel_world_size()
363
+ )
364
+ self.num_attention_heads_per_partition = divide(num_heads, self.tp_size)
365
366
self.norm0 = nn.LayerNorm(hidden_dim)
367
self.norm1 = nn.LayerNorm(hidden_dim)
453
- self.use_data_parallel = use_data_parallel
368
self.mlp = MLP2(
369
[hidden_dim, mlp_dim, hidden_dim],
370
activation,
371
prefix=f"{prefix}.mlp",
458
- use_data_parallel=use_data_parallel,
372
+ use_data_parallel=self.use_data_parallel,
373
)
460
- self.wqkv = ReplicatedLinear(
461
- hidden_dim, hidden_dim * 3, bias=attn_bias, prefix=f"{prefix}.wqkv"
374
+ self.wqkv = QKVParallelLinear(
375
+ hidden_size=hidden_dim,
376
+ head_size=self.hidden_size_per_attention_head,
377
+ total_num_heads=num_heads,
378
+ total_num_kv_heads=num_heads,
379
+ bias=attn_bias,
380
+ prefix=f"{prefix}.wqkv",
381
+ disable_tp=self.use_data_parallel,
382
)
463
- self.wo = ReplicatedLinear(
464
- hidden_dim, hidden_dim, bias=attn_bias, prefix=f"{prefix}.wo"
383
+ self.wo = RowParallelLinear(
384
+ hidden_dim,
385
+ hidden_dim,
386
+ bias=attn_bias,
387
+ prefix=f"{prefix}.wo",
388
+ disable_tp=self.use_data_parallel,
389
+ )
390
+ self.attn = MMEncoderAttention(
391
+ num_heads=self.num_attention_heads_per_partition,
392
+ head_size=self.hidden_size_per_attention_head,
393
+ multimodal_config=multimodal_config,
394
+ prefix=f"{prefix}.attn",
395
)
396
397
def attention_qkvpacked(
@@ -472,14 +402,15 @@ class MoonVitEncoderLayer(nn.Module):
402
):
403
"""
404
Args:
475
- x (torch.Tensor): (batch_size, seqlen, hidden_dim)
405
+ x (torch.Tensor): (seqlen, hidden_dim)
406
cu_seqlens (torch.Tensor):
407
"""
408
+ seq_length = x.size(0)
409
xqkv, _ = self.wqkv(x)
410
411
qkv_shape = xqkv.size()[:-1] + (
412
3,
482
- self.num_heads,
413
+ self.num_attention_heads_per_partition,
414
self.hidden_size_per_attention_head,
415
)
416
# xqkv: (batch_size, seqlen, 3, nheads, headdim)
@@ -488,9 +419,18 @@ class MoonVitEncoderLayer(nn.Module):
419
420
xq, xk = apply_rope(xq, xk, rope_freqs_cis)
421
491
- attn_func = VL_VISION_ATTENTION_FUNCTIONS[self.attn_implementation]
492
- attn_out = attn_func(
493
- xq, xk, xv, q_cu_seqlens=cu_seqlens, k_cu_seqlens=cu_seqlens
422
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
423
+ attn_out = self.attn(
424
+ xq.unsqueeze(0),
425
+ xk.unsqueeze(0),
426
+ xv.unsqueeze(0),
427
+ cu_seqlens=cu_seqlens,
428
+ max_seqlen=max_seqlen,
429
+ )
430
+ attn_out = attn_out.reshape(
431
+ seq_length,
432
+ self.num_attention_heads_per_partition
433
+ * self.hidden_size_per_attention_head,
434
)
435
attn_out, _ = self.wo(attn_out)
436
return attn_out
@@ -528,7 +468,7 @@ class MoonVitEncoder(nn.Module):
468
num_layers: int,
469
block_cfg: dict,
470
prefix: str = "",
531
- use_data_parallel: bool = False,
471
+ multimodal_config: MultiModalConfig | None = None,
472
) -> None:
473
super().__init__()
474
@@ -538,7 +478,7 @@ class MoonVitEncoder(nn.Module):
478
self.blocks = nn.ModuleList(
479
[
480
MoonVitEncoderLayer(
541
- use_data_parallel=use_data_parallel,
481
+ multimodal_config=multimodal_config,
482
prefix=f"{prefix}.blocks.{layer_idx}",
483
**block_cfg,
484
)
@@ -599,31 +539,6 @@ def patch_merger(
539
return outputs
540
541
602
-class MoonVitVLProjector(nn.Module):
603
- def __init__(
604
- self,
605
- in_channels: int,
606
- merge_kernel_size: list[int, int],
607
- hidden_act: str = "gelu",
608
- ln_eps: float = 1e-5,
609
- out_dim: int = 4096,
610
- ):
611
- super().__init__()
612
- self.hidden_size = in_channels * merge_kernel_size[0] * merge_kernel_size[1]
613
-
614
- self.pre_norm = nn.nn.LayerNorm(in_channels, eps=ln_eps)
615
- self.linear_1 = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
616
- self.act = ACT2FN[hidden_act]
617
- self.linear_2 = nn.Linear(self.hidden_size, out_dim, bias=True)
618
-
619
- def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
620
- hidden_states = self.pre_norm(hidden_states).view(-1, self.hidden_size)
621
- hidden_states = self.linear_1(hidden_states)
622
- hidden_states = self.act(hidden_states)
623
- hidden_states = self.linear_2(hidden_states)
624
- return hidden_states
625
-
626
-
542
class MoonVitPretrainedModel(PreTrainedModel):
543
config_class = MoonViTConfig
544
model_type = "moonvit"
@@ -634,14 +549,13 @@ class MoonVitPretrainedModel(PreTrainedModel):
549
def __init__(
550
self,
551
config: MoonViTConfig,
637
- use_data_parallel: bool = False,
552
+ multimodal_config: MultiModalConfig | None = None,
553
prefix: str = "",
554
*inputs,
555
**kwargs,
556
):
557
super().__init__(config, *inputs, **kwargs)
558
config = deepcopy(config)
644
- self.use_data_parallel = use_data_parallel
559
self.merge_kernel_size = config.merge_kernel_size
560
self.hidden_size = config.hidden_size
561
self.patch_size = config.patch_size
@@ -662,9 +576,9 @@ class MoonVitPretrainedModel(PreTrainedModel):
576
"mlp_dim": config.intermediate_size,
577
"activation": ACT2FN["gelu_pytorch_tanh"],
578
"attn_bias": True,
665
- "attn_implementation": config._attn_implementation,
579
},
580
prefix=f"{prefix}.encoder",
581
+ multimodal_config=multimodal_config,
582
)
583
584
def forward(