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 7, 06:30 UTC
Use MMEncoderAttention (=use FlashAttention) instead of torch.sdpa in radio.py (#35653)
Netanel Haber committed
Mar 4, 2026 at 18:43 UTC
289fc48ab73fb1eb610a72b4ddde9694e529bfba
1 file changed
+36
-44
vllm/model_executor/models/radio.py
+36
-44
@@ -10,7 +10,8 @@
10
11
import math
12
from collections.abc import Iterable
13
-from itertools import repeat
13
+from dataclasses import dataclass
14
+from itertools import accumulate, repeat
15
from typing import TypeAlias
16
17
import torch
@@ -477,28 +478,27 @@ class ViTPatchLinear(nn.Linear):
478
self.patch_size = patch_size
479
480
481
+@dataclass(frozen=True, kw_only=True)
482
+class MaskMetadata:
483
+ cu_seqlens: torch.Tensor
484
+ max_seqlen: torch.Tensor
485
+
486
+
487
class RadioParallelAttention(InternParallelAttention):
488
def forward(
482
- self, x: torch.Tensor, attn_mask: torch.Tensor | None = None
489
+ self, x: torch.Tensor, mask_meta: MaskMetadata | None = None
490
) -> torch.Tensor:
484
- if attn_mask is None:
485
- return super().forward(x)
486
-
487
- B, N, _ = x.shape
491
qkv, _ = self.qkv(x)
492
q, k, v = qkv.chunk(3, dim=-1)
493
494
if self.qk_normalization:
495
q, k = self._apply_qk_norm(q, k)
496
494
- q = q.view(B, N, self.num_heads_per_partition, self.head_dim)
495
- k = k.view(B, N, self.num_heads_per_partition, self.head_dim)
496
- v = v.view(B, N, self.num_heads_per_partition, self.head_dim)
497
- q, k, v = (t.transpose(1, 2) for t in (q, k, v))
498
- out = F.scaled_dot_product_attention(
499
- q, k, v, attn_mask=attn_mask, scale=self.scale
500
- )
501
- out = out.transpose(1, 2).reshape(B, N, -1)
497
+ cu_seqlens, max_seqlen = None, None
498
+ if mask_meta is not None:
499
+ cu_seqlens = mask_meta.cu_seqlens
500
+ max_seqlen = mask_meta.max_seqlen
501
+ out = self.attn(q, k, v, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
502
out, _ = self.proj(out)
503
return out
504
@@ -510,11 +510,11 @@ class RadioVisionEncoderLayer(InternVisionEncoderLayer):
510
def forward(
511
self,
512
hidden_states: torch.Tensor,
513
- attn_mask: torch.Tensor | None = None,
513
+ mask_meta: MaskMetadata | None = None,
514
):
515
hidden_states = (
516
hidden_states
517
- + self.attn(self.norm1(hidden_states), attn_mask=attn_mask) * self.ls1
517
+ + self.attn(self.norm1(hidden_states), mask_meta=mask_meta) * self.ls1
518
)
519
520
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states)) * self.ls2
@@ -529,11 +529,11 @@ class RadioVisionEncoder(InternVisionEncoder):
529
def forward(
530
self,
531
inputs_embeds: torch.Tensor,
532
- attn_mask: torch.Tensor | None = None,
532
+ mask_meta: MaskMetadata | None = None,
533
):
534
hidden_states = inputs_embeds
535
for encoder_layer in self.layers:
536
- hidden_states = encoder_layer(hidden_states, attn_mask=attn_mask)
536
+ hidden_states = encoder_layer(hidden_states, mask_meta=mask_meta)
537
return hidden_states
538
539
@@ -590,44 +590,36 @@ class RadioInternVisionModel(nn.Module):
590
def get_input_embeddings(self):
591
return self.embeddings
592
593
- def create_inter_image_attention_mask(
593
+ def inter_image_mask_metadata(
594
self, imgs_sizes: list[tuple[int, int]], device: torch.device
595
- ) -> torch.Tensor:
595
+ ) -> MaskMetadata:
596
patch_size = self.patch_generator.patch_size
597
num_skip = self.patch_generator.num_skip
598
599
seq_lens = calc_seq_lens(imgs_sizes, patch_size)
600
- patch_counts = [seq_len + num_skip for seq_len in seq_lens]
601
- total_patches = sum(patch_counts)
602
-
603
- # Create attention mask - default to False (mask out)
604
- mask = torch.zeros(
605
- total_patches, total_patches, dtype=torch.bool, device=device
600
+ adjusted = [s + num_skip for s in seq_lens]
601
+ cu_seqlens = torch.tensor(
602
+ list(accumulate(adjusted, initial=0)), dtype=torch.int32, device=device
603
)
607
-
608
- # Each image's patches can only attend to patches from the same image
609
- start_idx = 0
610
- for patch_count in patch_counts:
611
- end_idx = start_idx + patch_count
612
- # Allow attention within this image's patches
613
- mask[start_idx:end_idx, start_idx:end_idx] = True
614
- start_idx = end_idx
615
-
616
- return mask
604
+ # Keep max_seqlen on CPU to avoid .item() sync
605
+ # See: https://github.com/vllm-project/vllm/blob/20b6b01/vllm/v1/attention/ops/vit_attn_wrappers.py#L48
606
+ max_seqlen = torch.tensor(max(adjusted), dtype=torch.int32)
607
+ return MaskMetadata(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)
608
609
def forward(
610
self,
611
x: torch.Tensor,
621
- imgs_sizes: torch.Tensor | None = None,
612
+ imgs_sizes: list[tuple[int, int]] | None = None,
613
) -> torch.FloatTensor:
614
hidden_states = self.patch_generator(x, imgs_sizes=imgs_sizes)
624
- attn_mask = None
625
- if imgs_sizes is not None and len(imgs_sizes) > 1:
626
- # Dynamic Resolution
627
- attn_mask = self.create_inter_image_attention_mask(
628
- imgs_sizes, device=x.device
615
+ mask_meta = None
616
+ if imgs_sizes is not None:
617
+ assert len(imgs_sizes) > 0
618
+ # Dynamic resolution: process each image as an independent sequence.
619
+ mask_meta = self.inter_image_mask_metadata(
620
+ imgs_sizes, device=hidden_states.device
621
)
630
- encoder_outputs = self.encoder(inputs_embeds=hidden_states, attn_mask=attn_mask)
622
+ encoder_outputs = self.encoder(inputs_embeds=hidden_states, mask_meta=mask_meta)
623
return encoder_outputs
624
625
@@ -670,7 +662,7 @@ class RadioModel(nn.Module):
662
pixel_values: torch.Tensor | None = None,
663
pixel_embeds: torch.Tensor | None = None,
664
*,
673
- imgs_sizes: torch.Tensor | None = None,
665
+ imgs_sizes: list[tuple[int, int]] | None = None,
666
) -> tuple[torch.FloatTensor, torch.FloatTensor]:
667
y = self.model(pixel_values, imgs_sizes=imgs_sizes)
668
return self._extract_final(y, imgs_sizes=imgs_sizes)