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 6, 20:30 UTC
[MM][OOT] Support CPU `seq_lens` for OOT MMEncoderAttention kernels (#36605)
Signed-off-by: shen-shanshan <467638484@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Shanshan Shen committed
Mar 12, 2026 at 18:28 UTC
f0d3658c0f10700e7b8f7b4c7546059a3b7c027b
5 files changed
+52
-40
tests/kernels/attention/test_mha_attn.py
+6
-9
@@ -297,11 +297,10 @@ def test_mha_attn_varlen_forward_flashinfer(
297
hidden_size = num_heads * head_size
298
tp_size = 1
299
300
- sequence_lengths_np = MMEncoderAttention.maybe_compute_sequence_lengths(
301
- AttentionBackendEnum.FLASHINFER, cu_seqlens_np
302
- )
303
- sequence_lengths = torch.from_numpy(sequence_lengths_np).to(
304
- device, dtype=torch.int32, non_blocking=True
300
+ sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
301
+ AttentionBackendEnum.FLASHINFER,
302
+ cu_seqlens_np,
303
+ device,
304
)
305
306
max_seqlen_val = MMEncoderAttention.compute_max_seqlen(
@@ -309,14 +308,12 @@ def test_mha_attn_varlen_forward_flashinfer(
308
)
309
max_seqlen = torch.tensor(max_seqlen_val, device=device, dtype=torch.int32)
310
312
- cu_seqlens_np = MMEncoderAttention.maybe_recompute_cu_seqlens(
311
+ cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
312
AttentionBackendEnum.FLASHINFER,
313
cu_seqlens_np,
314
hidden_size,
315
tp_size,
317
- )
318
- cu_seqlens = torch.from_numpy(cu_seqlens_np).to(
319
- device, dtype=torch.int32, non_blocking=True
316
+ device,
317
)
318
319
scale = 1.0 / head_size**0.5
vllm/model_executor/custom_op.py
+6
@@ -22,6 +22,12 @@ op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
22
op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
23
24
25
+def get_oot_class_by_name(class_name: str) -> type | None:
26
+ if class_name in op_registry_oot:
27
+ return op_registry_oot[class_name]
28
+ return None
29
+
30
+
31
class PluggableLayer(nn.Module):
32
"""
33
Base class for pluggable layers.
vllm/model_executor/layers/attention/mm_encoder_attention.py
+33
-18
@@ -6,7 +6,7 @@ import numpy as np
6
import torch
7
8
from vllm.logger import init_logger
9
-from vllm.model_executor.custom_op import CustomOp
9
+from vllm.model_executor.custom_op import CustomOp, get_oot_class_by_name
10
from vllm.model_executor.models.vision import get_vit_attn_backend
11
from vllm.utils.math_utils import round_up
12
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
@@ -119,17 +119,25 @@ class MMEncoderAttention(CustomOp):
119
return max_seqlen
120
121
@classmethod
122
- def maybe_compute_sequence_lengths(
122
+ def maybe_compute_seq_lens(
123
cls,
124
attn_backend: AttentionBackendEnum,
125
cu_seqlens: np.ndarray,
126
- ) -> np.ndarray | None:
126
+ device: torch.device,
127
+ ) -> torch.Tensor | None:
128
+ if (oot_class := get_oot_class_by_name(cls.__name__)) is not None:
129
+ return oot_class.maybe_compute_seq_lens(attn_backend, cu_seqlens, device) # type: ignore[attr-defined]
130
+
131
if attn_backend != AttentionBackendEnum.FLASHINFER:
132
return None
133
+
134
sequence_lengths = cu_seqlens[1:] - cu_seqlens[:-1]
135
sequence_lengths = add_padding_to_seqlens(
136
sequence_lengths, len(sequence_lengths), 0
137
)
138
+ sequence_lengths = torch.from_numpy(sequence_lengths).to(
139
+ device, non_blocking=True
140
+ )
141
return sequence_lengths
142
143
@classmethod
@@ -139,24 +147,31 @@ class MMEncoderAttention(CustomOp):
147
cu_seqlens: np.ndarray,
148
hidden_size: int,
149
tp_size: int,
142
- ) -> np.ndarray:
143
- if attn_backend != AttentionBackendEnum.FLASHINFER:
144
- return cu_seqlens
150
+ device: torch.device,
151
+ ) -> torch.Tensor:
152
+ if (oot_class := get_oot_class_by_name(cls.__name__)) is not None:
153
+ return oot_class.maybe_recompute_cu_seqlens( # type: ignore[attr-defined]
154
+ attn_backend, cu_seqlens, hidden_size, tp_size, device
155
+ )
156
146
- batch_size = len(cu_seqlens) - 1
147
- scale = hidden_size // tp_size
148
- cu_seqlens = cu_seqlens * scale
157
+ if attn_backend == AttentionBackendEnum.FLASHINFER:
158
+ batch_size = len(cu_seqlens) - 1
159
+ scale = hidden_size // tp_size
160
+ cu_seqlens = cu_seqlens * scale
161
150
- cu_seqlens_qko = cu_seqlens
151
- cu_seqlens_v = cu_seqlens * 3
162
+ cu_seqlens_qko = cu_seqlens
163
+ cu_seqlens_v = cu_seqlens * 3
164
153
- cu_seqlens_qko = add_padding_to_seqlens(
154
- cu_seqlens_qko, batch_size, cu_seqlens_qko[-1]
155
- )
156
- cu_seqlens_v = add_padding_to_seqlens(
157
- cu_seqlens_v, batch_size, cu_seqlens_v[-1]
158
- )
159
- return np.concatenate([cu_seqlens_qko, cu_seqlens_v])
165
+ cu_seqlens_qko = add_padding_to_seqlens(
166
+ cu_seqlens_qko, batch_size, cu_seqlens_qko[-1]
167
+ )
168
+ cu_seqlens_v = add_padding_to_seqlens(
169
+ cu_seqlens_v, batch_size, cu_seqlens_v[-1]
170
+ )
171
+ cu_seqlens = np.concatenate([cu_seqlens_qko, cu_seqlens_v])
172
+
173
+ cu_seqlens = torch.from_numpy(cu_seqlens).to(device, non_blocking=True)
174
+ return cu_seqlens
175
176
def __init__(
177
self,
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+4
-6
@@ -983,13 +983,11 @@ class Qwen3Omni_VisionTransformer(nn.Module):
983
grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0]
984
).cumsum(axis=0, dtype=np.int32)
985
cu_seqlens_np = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens_np])
986
- sequence_lengths = MMEncoderAttention.maybe_compute_sequence_lengths(
987
- self.attn_backend, cu_seqlens_np
986
+ sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
987
+ self.attn_backend,
988
+ cu_seqlens_np,
989
+ self.device,
990
)
989
- if sequence_lengths is not None:
990
- sequence_lengths = torch.from_numpy(sequence_lengths).to(
991
- self.device, non_blocking=True
992
- )
991
992
hidden_states_list = []
993
deepstack_visual_indexes = self.deepstack_visual_indexes
vllm/model_executor/models/qwen3_vl.py
+3
-7
@@ -550,13 +550,9 @@ class Qwen3_VisionTransformer(nn.Module):
550
axis=0, dtype=np.int32
551
)
552
cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])
553
- sequence_lengths = MMEncoderAttention.maybe_compute_sequence_lengths(
554
- self.attn_backend, cu_seqlens
553
+ sequence_lengths = MMEncoderAttention.maybe_compute_seq_lens(
554
+ self.attn_backend, cu_seqlens, self.device
555
)
556
- if sequence_lengths is not None:
557
- sequence_lengths = torch.from_numpy(sequence_lengths).to(
558
- self.device, non_blocking=True
559
- )
556
max_seqlen = torch.tensor(
557
MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens),
558
dtype=torch.int32,
@@ -567,8 +563,8 @@ class Qwen3_VisionTransformer(nn.Module):
563
cu_seqlens,
564
self.hidden_size,
565
self.tp_size,
566
+ self.device,
567
)
571
- cu_seqlens = torch.from_numpy(cu_seqlens).to(self.device, non_blocking=True)
568
hidden_states = hidden_states.unsqueeze(1)
569
570
deepstack_feature_lists = []