@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 71cd89264f

Fork 0

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 8, 06:15 UTC

[MM Encoder] Add Triton ViT attention backend (#32183)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Isotr0py committed Feb 15, 2026 at 22:32 UTC 71cd89264f6c88ee4cca7c3cc556885e8844fc92

14 files changed +178 -51

tests/kernels/attention/test_mha_attn.py

+16 -2

@@ -17,7 +17,7 @@ from vllm.platforms import current_platform

17 from vllm.platforms.cpu import CpuPlatform

18 from vllm.platforms.cuda import CudaPlatform

19 from vllm.platforms.rocm import RocmPlatform

20 -from vllm.utils.torch_utils import set_random_seed

20 +from vllm.utils.torch_utils import set_default_torch_dtype, set_random_seed

21 from vllm.v1.attention.backends.registry import AttentionBackendEnum

22 from vllm.v1.attention.selector import _cached_get_attn_backend

@@ -71,6 +71,15 @@ def test_mha_attn_platform(default_vllm_config, device: str):

71 attn = MMEncoderAttention(16, 72, scale=1)

72 assert attn.attn_backend == AttentionBackendEnum.FLASH_ATTN

74 + # Test CUDA with head_size=72 (not divisible by 32)

75 + # - should use vLLM's FlashAttention

76 + with (

77 + patch("vllm.model_executor.models.vision.current_platform", CudaPlatform()),

78 + set_default_torch_dtype(torch.float32),

79 + ):

80 + attn = MMEncoderAttention(16, 72, scale=1)

81 + assert attn.attn_backend == AttentionBackendEnum.TRITON_ATTN

82 +

84 def ref_attention(

85 query: torch.Tensor,

@@ -153,7 +162,12 @@ def test_mha_attn_forward(

162 v,

163 scale=scale,

164 ).reshape(batch_size, seq_len, num_heads * head_size)

156 - torch.testing.assert_close(output, ref_output)

165 + tol_kwargs = (

166 + dict(rtol=1e-3, atol=1e-3)

167 + if attn.attn_backend == AttentionBackendEnum.TRITON_ATTN

168 + else {}

169 + )

170 + torch.testing.assert_close(output, ref_output, **tol_kwargs)

171

172

173 @pytest.mark.parametrize("var_seq_len", VAR_SEQ_LENS)

vllm/model_executor/layers/attention/mm_encoder_attention.py

+38

@@ -12,6 +12,7 @@ from vllm.v1.attention.backends.registry import AttentionBackendEnum

12 from vllm.v1.attention.ops.vit_attn_wrappers import (

13 vit_flash_attn_wrapper,

14 vit_torch_sdpa_wrapper,

15 + vit_triton_attn_wrapper,

16 )

18 logger = init_logger(__name__)

@@ -165,6 +166,41 @@ class MMEncoderAttention(CustomOp):

166 output = output.reshape(bsz, q_len, -1)

167 return output

168

169 + def _forward_triton(

170 + self,

171 + query: torch.Tensor,

172 + key: torch.Tensor,

173 + value: torch.Tensor,

174 + cu_seqlens: torch.Tensor | None = None,

175 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention

176 + ) -> torch.Tensor:

177 + """Input shape:

178 + (batch_size x seq_len x hidden_size) or

179 + (batch_size x seq_len x num_heads x head_size)

180 + """

181 + assert (cu_seqlens is not None and max_seqlen is not None) or (

182 + cu_seqlens is None and max_seqlen is None

183 + ), "cu_seqlens and max_seqlen should be both set or both None."

184 +

185 + bsz, q_len = query.size()[:2]

186 + kv_len = key.size(1)

187 + is_reshaped = query.dim() != 4

188 +

189 + query, key, value = self.view_qkv_to_4d(query, key, value, bsz, q_len, kv_len)

190 +

191 + output = vit_triton_attn_wrapper(

192 + q=query,

193 + k=key,

194 + v=value,

195 + batch_size=bsz,

196 + scale=self.scale,

197 + cu_seqlens=cu_seqlens,

198 + max_seqlen=max_seqlen,

199 + )

200 + if is_reshaped:

201 + output = output.reshape(bsz, q_len, -1)

202 + return output

203 +

204 def forward_native(

205 self,

206 query: torch.Tensor,

@@ -185,6 +221,8 @@ class MMEncoderAttention(CustomOp):

221 ) -> torch.Tensor:

222 if self.is_flash_attn_backend:

223 return self._forward_fa(query, key, value, cu_seqlens, max_seqlen)

224 + elif self.attn_backend == AttentionBackendEnum.TRITON_ATTN:

225 + return self._forward_triton(query, key, value, cu_seqlens, max_seqlen)

226 elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:

227 return self._forward_sdpa(query, key, value, cu_seqlens)

228 else:

vllm/model_executor/models/dots_ocr.py

+5 -4

@@ -573,10 +573,11 @@ class DotsVisionTransformer(nn.Module):

573

574 def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> int | None:

575 max_seqlen = None

576 - if (

577 - self.attn_backend == AttentionBackendEnum.FLASH_ATTN

578 - or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA

579 - ):

576 + if self.attn_backend in {

577 + AttentionBackendEnum.FLASH_ATTN,

578 + AttentionBackendEnum.ROCM_AITER_FA,

579 + AttentionBackendEnum.TRITON_ATTN,

580 + }:

581 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

582 return max_seqlen

583

vllm/model_executor/models/ernie45_vl.py

+5 -4

@@ -446,10 +446,11 @@ class Ernie4_5_VisionTransformer(nn.Module):

446

447 def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> torch.Tensor | None:

448 max_seqlen = None

449 - if (

450 - self.attn_backend == AttentionBackendEnum.FLASH_ATTN

451 - or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA

452 - ):

449 + if self.attn_backend in {

450 + AttentionBackendEnum.FLASH_ATTN,

451 + AttentionBackendEnum.ROCM_AITER_FA,

452 + AttentionBackendEnum.TRITON_ATTN,

453 + }:

454 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

455 return max_seqlen

456

vllm/model_executor/models/glm4_1v.py

+5 -4

@@ -723,10 +723,11 @@ class Glm4vVisionTransformer(nn.Module):

723 cu_seqlens: torch.Tensor,

724 ) -> torch.Tensor | None:

725 max_seqlen = None

726 - if (

727 - self.attn_backend == AttentionBackendEnum.FLASH_ATTN

728 - or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA

729 - ):

726 + if self.attn_backend in {

727 + AttentionBackendEnum.FLASH_ATTN,

728 + AttentionBackendEnum.ROCM_AITER_FA,

729 + AttentionBackendEnum.TRITON_ATTN,

730 + }:

731 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

732 return max_seqlen

733

vllm/model_executor/models/paddleocr_vl.py

+2 -8

@@ -730,14 +730,7 @@ class SiglipEncoder(nn.Module):

730 head_size=head_dim,

731 dtype=torch.get_default_dtype(),

732 )

733 - if self.attn_backend not in {

734 - AttentionBackendEnum.FLASH_ATTN,

735 - AttentionBackendEnum.TORCH_SDPA,

736 - AttentionBackendEnum.ROCM_AITER_FA,

737 - }:

738 - raise RuntimeError(

739 - f"PaddleOCR-VL does not support {self.attn_backend} backend now."

740 - )

733 +

734 self.layers = nn.ModuleList(

735 [

736 SiglipEncoderLayer(

@@ -805,6 +798,7 @@ class SiglipEncoder(nn.Module):

798 if self.attn_backend in {

799 AttentionBackendEnum.FLASH_ATTN,

800 AttentionBackendEnum.ROCM_AITER_FA,

801 + AttentionBackendEnum.TRITON_ATTN,

802 }:

803 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

804

vllm/model_executor/models/qwen2_5_vl.py

+1 -9

@@ -607,15 +607,6 @@ class Qwen2_5_VisionTransformer(nn.Module):

607 dtype=torch.get_default_dtype(),

608 )

609

610 - if self.attn_backend not in {

611 - AttentionBackendEnum.FLASH_ATTN,

612 - AttentionBackendEnum.TORCH_SDPA,

613 - AttentionBackendEnum.ROCM_AITER_FA,

614 - }:

615 - raise RuntimeError(

616 - f"Qwen2.5-VL does not support {self.attn_backend} backend now."

617 - )

618 -

610 with set_model_tag("Qwen2_5_VisionBlock", is_encoder=True):

611 self.blocks = nn.ModuleList(

612 [

@@ -761,6 +752,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

752 if self.attn_backend in {

753 AttentionBackendEnum.FLASH_ATTN,

754 AttentionBackendEnum.ROCM_AITER_FA,

755 + AttentionBackendEnum.TRITON_ATTN,

756 }:

757 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

758 return max_seqlen

vllm/model_executor/models/qwen2_vl.py

@@ -642,6 +642,7 @@ class Qwen2VisionTransformer(nn.Module):

642 if self.attn_backend in {

643 AttentionBackendEnum.FLASH_ATTN,

644 AttentionBackendEnum.ROCM_AITER_FA,

645 + AttentionBackendEnum.TRITON_ATTN,

646 }:

647 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

648 return max_seqlen

vllm/model_executor/models/qwen3_omni_moe_thinker.py

@@ -391,6 +391,7 @@ class Qwen3OmniMoeAudioEncoder(nn.Module):

391 if self.attn_backend in {

392 AttentionBackendEnum.FLASH_ATTN,

393 AttentionBackendEnum.ROCM_AITER_FA,

394 + AttentionBackendEnum.TRITON_ATTN,

395 }:

396 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

397 return max_seqlen

@@ -919,6 +920,7 @@ class Qwen3Omni_VisionTransformer(nn.Module):

920 if self.attn_backend in {

921 AttentionBackendEnum.FLASH_ATTN,

922 AttentionBackendEnum.ROCM_AITER_FA,

923 + AttentionBackendEnum.TRITON_ATTN,

924 }:

925 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

926 return max_seqlen

vllm/model_executor/models/qwen3_vl.py

+4 -11

@@ -385,14 +385,6 @@ class Qwen3_VisionTransformer(nn.Module):

385 dtype=torch.get_default_dtype(),

386 )

387

388 - if self.attn_backend not in {

389 - AttentionBackendEnum.FLASH_ATTN,

390 - AttentionBackendEnum.TORCH_SDPA,

391 - AttentionBackendEnum.ROCM_AITER_FA,

392 - }:

393 - raise RuntimeError(

394 - f"Qwen3-VL does not support {self.attn_backend} backend now."

395 - )

388 self.blocks = nn.ModuleList(

389 [

390 Qwen3_VisionBlock(

@@ -526,9 +518,10 @@ class Qwen3_VisionTransformer(nn.Module):

518 cu_seqlens: torch.Tensor,

519 ) -> torch.Tensor:

520 max_seqlen = torch.zeros([], device=cu_seqlens.device)

529 - if (

530 - self.attn_backend == AttentionBackendEnum.FLASH_ATTN

531 - or self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA

521 + if self.attn_backend in (

522 + AttentionBackendEnum.FLASH_ATTN,

523 + AttentionBackendEnum.ROCM_AITER_FA,

524 + AttentionBackendEnum.TRITON_ATTN,

525 ):

526 max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

527 return max_seqlen

vllm/model_executor/models/vision.py

+2 -2

@@ -108,7 +108,7 @@ def get_vit_attn_backend(

108 multimodal_config: MultiModalConfig | None = (

109 model_config.multimodal_config if model_config is not None else None

110 )

111 - except AssertionError:

111 + except (AssertionError, AttributeError):

112 multimodal_config = None

113

114 attn_backend_override = (

@@ -134,7 +134,7 @@ def is_vit_use_data_parallel():

134 multimodal_config: MultiModalConfig | None = (

135 model_config.multimodal_config if model_config is not None else None

136 )

137 - except AssertionError:

137 + except (AssertionError, AttributeError):

138 multimodal_config = None

139

140 mm_encoder_tp_mode = (

vllm/platforms/cuda.py

+19 -7

@@ -411,8 +411,9 @@ class CudaPlatformBase(Platform):

411 @classmethod

412 def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:

413 return [

414 - AttentionBackendEnum.TORCH_SDPA,

414 AttentionBackendEnum.FLASH_ATTN,

415 + AttentionBackendEnum.TRITON_ATTN,

416 + AttentionBackendEnum.TORCH_SDPA,

417 ]

418

419 @classmethod

@@ -430,14 +431,25 @@ class CudaPlatformBase(Platform):

431 logger.info_once(f"Using backend {backend} for vit attention")

432 return backend

433

433 - # Try FlashAttention first

434 - if (cc := cls.get_device_capability()) and cc.major >= 8:

434 + cc = cls.get_device_capability()

435 + for vit_attn_backend in cls.get_supported_vit_attn_backends():

436 + if vit_attn_backend == AttentionBackendEnum.TORCH_SDPA:

437 + continue

438 try:

436 - backend_class = AttentionBackendEnum.FLASH_ATTN.get_class()

437 - if backend_class.supports_head_size(

439 + backend_class = vit_attn_backend.get_class()

440 + is_backend_supported = backend_class.supports_head_size(

441 head_size

439 - ) and backend_class.supports_dtype(dtype):

440 - return AttentionBackendEnum.FLASH_ATTN

442 + ) and backend_class.supports_dtype(dtype)

443 + if cc is not None:

444 + is_backend_supported = (

445 + is_backend_supported

446 + and backend_class.supports_compute_capability(cc)

447 + )

448 + if is_backend_supported:

449 + logger.info_once(

450 + f"Using backend {vit_attn_backend} for vit attention"

451 + )

452 + return vit_attn_backend

453 except ImportError:

454 pass

455

vllm/platforms/rocm.py

@@ -384,6 +384,7 @@ class RocmPlatform(Platform):

384 return [

385 AttentionBackendEnum.FLASH_ATTN,

386 AttentionBackendEnum.ROCM_AITER_FA,

387 + AttentionBackendEnum.TRITON_ATTN,

388 AttentionBackendEnum.TORCH_SDPA,

389 ]

390

vllm/v1/attention/ops/vit_attn_wrappers.py

+77

@@ -110,6 +110,83 @@ def vit_flash_attn_wrapper(

110 )

111

112

113 +def triton_attn_wrapper(

114 + q: torch.Tensor,

115 + k: torch.Tensor,

116 + v: torch.Tensor,

117 + batch_size: int,

118 + scale: float | None = None,

119 + cu_seqlens: torch.Tensor | None = None,

120 + max_seqlen: torch.Tensor | None = None,

121 +) -> torch.Tensor:

122 + from vllm.v1.attention.ops.triton_prefill_attention import context_attention_fwd

123 +

124 + q_len = q.size(1)

125 + if cu_seqlens is None:

126 + cu_seqlens = torch.arange(

127 + 0, (batch_size + 1) * q_len, step=q_len, dtype=torch.int32, device=q.device

128 + )

129 + max_seqlen = q_len if max_seqlen is None else max_seqlen.item()

130 +

131 + q, k, v = (einops.rearrange(x, "b s ... -> (b s) ...") for x in [q, k, v])

132 + output = torch.empty_like(q)

133 + context_attention_fwd(

134 + q,

135 + k,

136 + v,

137 + output,

138 + b_start_loc=cu_seqlens[:-1],

139 + b_seq_len=cu_seqlens[1:] - cu_seqlens[:-1],

140 + max_input_len=max_seqlen,

141 + is_causal=False,

142 + sliding_window_q=None,

143 + sliding_window_k=None,

144 + softmax_scale=scale,

145 + )

146 +

147 + context_layer = einops.rearrange(output, "(b s) h d -> b s h d", b=batch_size)

148 + return context_layer

149 +

150 +

151 +def triton_attn_wrapper_fake(

152 + q: torch.Tensor,

153 + k: torch.Tensor,

154 + v: torch.Tensor,

155 + batch_size: int,

156 + scale: float | None = None,

157 + cu_seqlens: torch.Tensor | None = None,

158 + max_seqlen: torch.Tensor | None = None,

159 +) -> torch.Tensor:

160 + return torch.empty_like(q)

161 +

162 +

163 +direct_register_custom_op(

164 + op_name="triton_attn_wrapper",

165 + op_func=triton_attn_wrapper,

166 + fake_impl=triton_attn_wrapper_fake,

167 +)

168 +

169 +

170 +def vit_triton_attn_wrapper(

171 + q: torch.Tensor,

172 + k: torch.Tensor,

173 + v: torch.Tensor,

174 + batch_size: int,

175 + scale: float | None = None,

176 + cu_seqlens: torch.Tensor | None = None,

177 + max_seqlen: torch.Tensor | None = None,

178 +) -> torch.Tensor:

179 + return torch.ops.vllm.triton_attn_wrapper(

180 + q,

181 + k,

182 + v,

183 + batch_size,

184 + scale,

185 + cu_seqlens,

186 + max_seqlen,

187 + )

188 +

189 +

190 def apply_sdpa(

191 q: torch.Tensor,

192 k: torch.Tensor,