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, 22:00 UTC
[MM] Enable FlashInfer metadata support for Qwen2.5-VL vision attention (#42787)
Signed-off-by: Hua Huang <huah@nvidia.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Hua Huang committed
May 23, 2026 at 09:08 UTC
a0be71ee47d3e0bed797ca08dff2ec863b2e28d6
1 file changed
+72
-13
vllm/model_executor/models/qwen2_5_vl.py
+72
-13
@@ -113,6 +113,7 @@ from .utils import (
113
maybe_prefix,
114
)
115
from .vision import (
116
+ get_fp8_padded_hidden_size,
117
get_vit_attn_backend,
118
is_vit_use_data_parallel,
119
run_dp_sharded_mrope_vision_model,
@@ -369,7 +370,8 @@ class Qwen2_5_VisionAttention(nn.Module):
370
rotary_pos_emb_cos: torch.Tensor,
371
rotary_pos_emb_sin: torch.Tensor,
372
max_seqlen: torch.Tensor, # Only used for Flash Attention
372
- sequence_lengths: torch.Tensor, # Only used for FlashInfer CuDNN backend
373
+ # Only used for FlashInfer CuDNN backend.
374
+ sequence_lengths: torch.Tensor | None,
375
) -> torch.Tensor:
376
# [s, b, c] --> [s, b, head * 3 * head_dim]
377
x, _ = self.qkv(x)
@@ -426,6 +428,7 @@ class Qwen2_5_VisionAttention(nn.Module):
428
dynamic_arg_dims={
429
"x": 0,
430
"cu_seqlens": 0,
431
+ "sequence_lengths": 0,
432
"rotary_pos_emb_cos": 0,
433
"rotary_pos_emb_sin": 0,
434
},
@@ -471,6 +474,8 @@ class Qwen2_5_VisionBlock(nn.Module):
474
rotary_pos_emb_cos: torch.Tensor,
475
rotary_pos_emb_sin: torch.Tensor,
476
max_seqlen: torch.Tensor, # Only used for Flash Attention
477
+ # Only used for FlashInfer CuDNN backend.
478
+ sequence_lengths: torch.Tensor | None = None,
479
) -> torch.Tensor:
480
x_attn = self.attn(
481
self.norm1(x),
@@ -478,7 +483,7 @@ class Qwen2_5_VisionBlock(nn.Module):
483
rotary_pos_emb_cos=rotary_pos_emb_cos,
484
rotary_pos_emb_sin=rotary_pos_emb_sin,
485
max_seqlen=max_seqlen,
481
- sequence_lengths=None,
486
+ sequence_lengths=sequence_lengths,
487
)
488
x_fused_norm, residual = self.norm2(x, residual=x_attn)
489
x = residual + self.mlp(x_fused_norm)
@@ -598,6 +603,12 @@ class Qwen2_5_VisionTransformer(nn.Module):
603
self.spatial_merge_size = vision_config.spatial_merge_size
604
self.fullatt_block_indexes = vision_config.fullatt_block_indexes
605
self.spatial_merge_unit = self.spatial_merge_size**2
606
+ use_data_parallel = is_vit_use_data_parallel()
607
+ self.tp_size = (
608
+ 1
609
+ if use_data_parallel
610
+ else parallel_state.get_tensor_model_parallel_world_size()
611
+ )
612
self.patch_embed = Qwen2_5_VisionPatchEmbed(
613
patch_size=patch_size,
614
temporal_patch_size=temporal_patch_size,
@@ -607,6 +618,11 @@ class Qwen2_5_VisionTransformer(nn.Module):
618
619
norm_layer = partial(RMSNorm, eps=norm_eps)
620
head_dim = self.hidden_size // self.num_heads
621
+ # FP8 attention: Q/K/V become independent contiguous tensors after
622
+ # quantization, so FlashInfer cu_seqlens uses uniform stride.
623
+ self.fp8_padded_hidden_size = get_fp8_padded_hidden_size(
624
+ self.num_heads, head_dim
625
+ )
626
self.rotary_pos_emb = get_rope(
627
head_size=head_dim,
628
max_position=8192,
@@ -902,21 +918,55 @@ class Qwen2_5_VisionTransformer(nn.Module):
918
)
919
)
920
905
- # transformers
906
- # pre-compute seqlens for window/full attn to reduce cuMemcpy operations
921
+ cu_seqlens_np = cu_seqlens.cpu().numpy()
922
+ cu_window_seqlens_np = cu_window_seqlens.cpu().numpy()
923
+
924
+ # FlashInfer needs the real per-sequence lengths in addition to
925
+ # cu_seqlens. For other backends this returns None and is ignored.
926
+ sequence_lengths_full = MMEncoderAttention.maybe_compute_seq_lens(
927
+ self.attn_backend, cu_seqlens_np, device
928
+ )
929
+ sequence_lengths_window = MMEncoderAttention.maybe_compute_seq_lens(
930
+ self.attn_backend, cu_window_seqlens_np, device
931
+ )
932
+
933
+ # Pre-compute max sequence lengths for window/full attention. FlashInfer
934
+ # buckets this value for cuDNN graph reuse; other backends keep the exact
935
+ # maximum. Keep the scalar on CPU because attention wrappers call .item().
936
if max_seqlen_override is None:
908
- max_seqlen_full = self.compute_attn_mask_seqlen(cu_seqlens)
937
+ max_seqlen_full_val = MMEncoderAttention.compute_max_seqlen(
938
+ self.attn_backend, cu_seqlens_np
939
+ )
940
else:
910
- max_seqlen_full = torch.tensor(max_seqlen_override, dtype=torch.int32)
941
+ max_seqlen_full_val = max_seqlen_override
942
+ max_seqlen_full = torch.tensor(max_seqlen_full_val, dtype=torch.int32)
943
if max_seqlen_window_override is None:
912
- max_seqlen_window = self.compute_attn_mask_seqlen(cu_window_seqlens)
913
- else:
914
- max_seqlen_window = torch.tensor(
915
- max_seqlen_window_override, dtype=torch.int32
944
+ max_seqlen_window_val = MMEncoderAttention.compute_max_seqlen(
945
+ self.attn_backend, cu_window_seqlens_np
946
)
917
-
918
- cu_seqlens = cu_seqlens.to(device=device, non_blocking=True)
919
- cu_window_seqlens = cu_window_seqlens.to(device=device, non_blocking=True)
947
+ else:
948
+ max_seqlen_window_val = max_seqlen_window_override
949
+ max_seqlen_window = torch.tensor(max_seqlen_window_val, dtype=torch.int32)
950
+
951
+ # FlashInfer uses backend-specific cu_seqlens offsets into the flattened
952
+ # Q/K/O and V buffers. Other backends receive the original cumulative
953
+ # token offsets unchanged.
954
+ cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
955
+ self.attn_backend,
956
+ cu_seqlens_np,
957
+ self.hidden_size,
958
+ self.tp_size,
959
+ device,
960
+ fp8_padded_hidden_size=self.fp8_padded_hidden_size,
961
+ )
962
+ cu_window_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(
963
+ self.attn_backend,
964
+ cu_window_seqlens_np,
965
+ self.hidden_size,
966
+ self.tp_size,
967
+ device,
968
+ fp8_padded_hidden_size=self.fp8_padded_hidden_size,
969
+ )
970
rotary_pos_emb_cos = rotary_pos_emb_cos.to(device=device, non_blocking=True)
971
rotary_pos_emb_sin = rotary_pos_emb_sin.to(device=device, non_blocking=True)
972
window_index = window_index.to(device=device, non_blocking=True)
@@ -930,6 +980,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
980
metadata["cu_window_seqlens"] = cu_window_seqlens
981
metadata["max_seqlen_full"] = max_seqlen_full
982
metadata["max_seqlen_window"] = max_seqlen_window
983
+ metadata["sequence_lengths_full"] = sequence_lengths_full
984
+ metadata["sequence_lengths_window"] = sequence_lengths_window
985
986
return metadata
987
@@ -955,6 +1007,8 @@ class Qwen2_5_VisionTransformer(nn.Module):
1007
cu_window_seqlens = encoder_metadata["cu_window_seqlens"]
1008
max_seqlen_full = encoder_metadata["max_seqlen_full"]
1009
max_seqlen_window = encoder_metadata["max_seqlen_window"]
1010
+ sequence_lengths_full = encoder_metadata.get("sequence_lengths_full")
1011
+ sequence_lengths_window = encoder_metadata.get("sequence_lengths_window")
1012
1013
hidden_states = hidden_states.reshape(
1014
seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1
@@ -968,9 +1022,11 @@ class Qwen2_5_VisionTransformer(nn.Module):
1022
if layer_num in self.fullatt_block_indexes:
1023
cu_seqlens_now = cu_seqlens
1024
max_seqlen_now = max_seqlen_full
1025
+ sequence_lengths_now = sequence_lengths_full
1026
else:
1027
cu_seqlens_now = cu_window_seqlens
1028
max_seqlen_now = max_seqlen_window
1029
+ sequence_lengths_now = sequence_lengths_window
1030
1031
hidden_states = blk(
1032
hidden_states,
@@ -978,6 +1034,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
1034
rotary_pos_emb_cos=rotary_pos_emb_cos,
1035
rotary_pos_emb_sin=rotary_pos_emb_sin,
1036
max_seqlen=max_seqlen_now,
1037
+ sequence_lengths=sequence_lengths_now,
1038
)
1039
1040
# For Qwen2.5-VL-3B, float16 will overflow at last block
@@ -1594,6 +1651,8 @@ class Qwen2_5_VLForConditionalGeneration(
1651
"cu_window_seqlens",
1652
"max_seqlen_full",
1653
"max_seqlen_window",
1654
+ "sequence_lengths_full",
1655
+ "sequence_lengths_window",
1656
],
1657
out_hidden_size=self.visual.out_hidden_size,
1658
max_frames_per_video=max_frames,