@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9c3fe9936b

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

Flashinfer cuDNN backend for Qwen3 VL ViT attention (#34580)

Signed-off-by: Max Hu <maxhu@nvidia.com> Signed-off-by: Max Hu <hyoung2991@gmail.com> Co-authored-by: Max Hu <maxhu@nvidia.com> Co-authored-by: Shang Wang <shangw@nvidia.com>

Max Hu committed Feb 27, 2026 at 20:20 UTC 9c3fe9936b929b5503d780bd4e8e3cd524de1c4e

6 files changed +405 -21

tests/kernels/attention/test_mha_attn.py

+110

@@ -9,9 +9,12 @@ Test:

9 import itertools

10 from unittest.mock import patch

12 +import numpy as np

13 import pytest

14 import torch

16 +from vllm.config import get_current_vllm_config

17 +from vllm.config.multimodal import MultiModalConfig

18 from vllm.model_executor.layers.attention import MMEncoderAttention

19 from vllm.platforms import current_platform

20 from vllm.platforms.cpu import CpuPlatform

@@ -224,3 +227,110 @@ def test_mha_attn_varlen_forward(

227 ref_output.append(output_i)

228 ref_output = torch.cat(ref_output, dim=1)

229 torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)

230 +

231 +

232 +@pytest.mark.parametrize("var_seq_len", VAR_SEQ_LENS)

233 +@pytest.mark.parametrize(

234 + "dtype",

235 + [torch.bfloat16, torch.half],

236 +)

237 +@pytest.mark.parametrize("device", CUDA_DEVICES)

238 +def test_mha_attn_varlen_forward_flashinfer(

239 + default_vllm_config,

240 + var_seq_len: list[int],

241 + dtype: torch.dtype,

242 + device: str,

243 +):

244 + """Test MMEncoderAttention varlen forward with FLASHINFER backend (head_size=72).

245 +

246 + Exercises the path that uses --mm-encoder-attn-backend=FLASHINFER with

247 + recomputed cu_seqlens, max_seqlen, and sequence_lengths as in qwen3_vl

248 + vision encoder.

249 + """

250 + pytest.importorskip("flashinfer")

251 +

252 + num_heads = 16

253 + head_size = 72

254 + set_random_seed(0)

255 + torch.set_default_device(device)

256 + torch.set_default_dtype(dtype)

257 +

258 + # Override vllm config so get_vit_attn_backend returns FLASHINFER (simulates

259 + # --mm-encoder-attn-backend=FLASHINFER).

260 + vllm_config = get_current_vllm_config()

261 + old_model_config = getattr(vllm_config, "model_config", None)

262 + minimal_model_config = type(

263 + "MinimalModelConfig",

264 + (),

265 + {

266 + "multimodal_config": MultiModalConfig(

267 + mm_encoder_attn_backend=AttentionBackendEnum.FLASHINFER

268 + ),

269 + },

270 + )()

271 + vllm_config.model_config = minimal_model_config

272 + try:

273 + total_len = sum(var_seq_len)

274 + # Stride of second dim = 3 * num_heads * head_size (same as qwen2_5_vl

275 + # after qkv rearrange and unbind: qkv shape (b, s, 3, head, head_dim)).

276 + qkv = torch.randn(1, total_len, 3, num_heads, head_size)

277 + q, k, v = qkv.unbind(dim=2)

278 +

279 + cu_seqlens_np = np.array(

280 + [0] + list(itertools.accumulate(var_seq_len)), dtype=np.int32

281 + )

282 + hidden_size = num_heads * head_size

283 + tp_size = 1

284 +

285 + sequence_lengths_np = MMEncoderAttention.maybe_compute_sequence_lengths(

286 + AttentionBackendEnum.FLASHINFER, cu_seqlens_np

287 + )

288 + sequence_lengths = torch.from_numpy(sequence_lengths_np).to(

289 + device, dtype=torch.int32, non_blocking=True

290 + )

291 +

292 + max_seqlen_val = MMEncoderAttention.compute_max_seqlen(

293 + AttentionBackendEnum.FLASHINFER, cu_seqlens_np

294 + )

295 + max_seqlen = torch.tensor(max_seqlen_val, device=device, dtype=torch.int32)

296 +

297 + cu_seqlens_np = MMEncoderAttention.maybe_recompute_cu_seqlens(

298 + AttentionBackendEnum.FLASHINFER,

299 + cu_seqlens_np,

300 + hidden_size,

301 + tp_size,

302 + )

303 + cu_seqlens = torch.from_numpy(cu_seqlens_np).to(

304 + device, dtype=torch.int32, non_blocking=True

305 + )

306 +

307 + scale = 1.0 / head_size**0.5

308 + attn = MMEncoderAttention(

309 + num_heads,

310 + head_size,

311 + scale=scale,

312 + num_kv_heads=num_heads,

313 + )

314 + assert attn.attn_backend == AttentionBackendEnum.FLASHINFER

315 +

316 + output = attn(

317 + q,

318 + k,

319 + v,

320 + cu_seqlens=cu_seqlens,

321 + max_seqlen=max_seqlen,

322 + sequence_lengths=sequence_lengths,

323 + )

324 +

325 + ref_output = []

326 + for q_i, k_i, v_i in zip(

327 + torch.split(q, var_seq_len, dim=1),

328 + torch.split(k, var_seq_len, dim=1),

329 + torch.split(v, var_seq_len, dim=1),

330 + ):

331 + output_i = ref_attention(q_i, k_i, v_i, scale=scale)

332 + ref_output.append(output_i)

333 + ref_output = torch.cat(ref_output, dim=1)

334 + torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=1e-2)

335 + finally:

336 + vllm_config.model_config = old_model_config

vllm/model_executor/layers/attention/mm_encoder_attention.py

+170 -2

@@ -2,21 +2,93 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

5 +import numpy as np

6 import torch

8 from vllm.logger import init_logger

9 from vllm.model_executor.custom_op import CustomOp

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

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

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

15 vit_flash_attn_wrapper,

16 + vit_flashinfer_wrapper,

17 vit_torch_sdpa_wrapper,

18 vit_triton_attn_wrapper,

19 )

21 logger = init_logger(__name__)

23 +# Batch buckets for cuDNN graph caching.

24 +# Graphs use batch size and max sequence length as cache key.

25 +# This avoids creating a new graph for each unique set of

26 +# batch size and max sequence length at runtime.

27 +# From the cuDNN team's performance measurements, there

28 +# is no significant kernel performance difference between padding

29 +# to a smaller batch size/seq length and padding to larger

30 +# ones. The bucketing here is solely used to avoid memory

31 +# operation overhead, which won't be needed if we have CUDA

32 +# graph support in the future.

33 +# TODO: Remove buckets after issue #34763

34 +# (cuda graph support) is addressed.

35 +FLASHINFER_BATCH_BUCKETS = [8, 16, 32, 64]

36 +FLASHINFER_MAX_SEQLEN_BUCKETS = [

37 + 1 * 1024,

38 + 2 * 1024,

39 + 4 * 1024,

40 + 8 * 1024,

41 + 16 * 1024,

42 + 32 * 1024,

43 + 64 * 1024,

44 + 128 * 1024,

45 +]

46 +

47 +# Workspace buffer for FlashInfer CuDNN backend

48 +FLASHINFER_CUDNN_WORKSPACE_SIZE_BYTES = 128 * 1024 * 1024

49 +_flashinfer_workspace_buffer: torch.Tensor | None = None

50 +

51 +

52 +def _get_flashinfer_workspace_buffer() -> torch.Tensor:

53 + global _flashinfer_workspace_buffer

54 + if _flashinfer_workspace_buffer is None:

55 + _flashinfer_workspace_buffer = torch.zeros(

56 + FLASHINFER_CUDNN_WORKSPACE_SIZE_BYTES,

57 + dtype=torch.uint8,

58 + device="cuda",

59 + )

60 + return _flashinfer_workspace_buffer

61 +

62 +

63 +def add_padding_to_seqlens(

64 + seq: np.ndarray,

65 + batch_size: int,

66 + padding_value: int,

67 +) -> np.ndarray:

68 + batch_size_padded = next(

69 + (b for b in FLASHINFER_BATCH_BUCKETS if b >= batch_size),

70 + round_up(batch_size, FLASHINFER_BATCH_BUCKETS[0]),

71 + )

72 + if batch_size_padded == batch_size:

73 + return seq

74 + return np.concatenate(

75 + [

76 + seq,

77 + np.full((batch_size_padded - batch_size,), padding_value, dtype=seq.dtype),

78 + ]

79 + )

80 +

81 +

82 +def bucket_flashinfer_max_seqlen(

83 + real_max_seqlen: int,

84 +) -> int:

85 + if real_max_seqlen <= 0:

86 + return FLASHINFER_MAX_SEQLEN_BUCKETS[0]

87 + return next(

88 + (s for s in FLASHINFER_MAX_SEQLEN_BUCKETS if s >= real_max_seqlen),

89 + round_up(real_max_seqlen, FLASHINFER_MAX_SEQLEN_BUCKETS[-1]),

90 + )

91 +

93 # --8<-- [start:mm_encoder_attn]

94 @CustomOp.register("mm_encoder_attn")

@@ -24,6 +96,67 @@ class MMEncoderAttention(CustomOp):

96 """Multi-headed attention without any cache, used for multimodal encoder."""

98 # --8<-- [end:mm_encoder_attn]

99 + @classmethod

100 + def compute_max_seqlen(

101 + cls,

102 + attn_backend: AttentionBackendEnum,

103 + cu_seqlens: np.ndarray,

104 + ) -> int:

105 + max_seqlen = 0

106 + if (

107 + attn_backend

108 + in (

109 + AttentionBackendEnum.FLASH_ATTN,

110 + AttentionBackendEnum.ROCM_AITER_FA,

111 + AttentionBackendEnum.TRITON_ATTN,

112 + AttentionBackendEnum.FLASHINFER,

113 + )

114 + and len(cu_seqlens) >= 2

115 + ):

116 + max_seqlen = int((cu_seqlens[1:] - cu_seqlens[:-1]).max())

117 + if attn_backend == AttentionBackendEnum.FLASHINFER:

118 + max_seqlen = bucket_flashinfer_max_seqlen(max_seqlen)

119 + return max_seqlen

120 +

121 + @classmethod

122 + def maybe_compute_sequence_lengths(

123 + cls,

124 + attn_backend: AttentionBackendEnum,

125 + cu_seqlens: np.ndarray,

126 + ) -> np.ndarray | None:

127 + if attn_backend != AttentionBackendEnum.FLASHINFER:

128 + return None

129 + sequence_lengths = cu_seqlens[1:] - cu_seqlens[:-1]

130 + sequence_lengths = add_padding_to_seqlens(

131 + sequence_lengths, len(sequence_lengths), 0

132 + )

133 + return sequence_lengths

134 +

135 + @classmethod

136 + def maybe_recompute_cu_seqlens(

137 + cls,

138 + attn_backend: AttentionBackendEnum,

139 + cu_seqlens: np.ndarray,

140 + hidden_size: int,

141 + tp_size: int,

142 + ) -> np.ndarray:

143 + if attn_backend != AttentionBackendEnum.FLASHINFER:

144 + return cu_seqlens

145 +

146 + batch_size = len(cu_seqlens) - 1

147 + scale = hidden_size // tp_size

148 + cu_seqlens = cu_seqlens * scale

149 +

150 + cu_seqlens_qko = cu_seqlens

151 + cu_seqlens_v = cu_seqlens * 3

152 +

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])

160

161 def __init__(

162 self,

@@ -46,10 +179,9 @@ class MMEncoderAttention(CustomOp):

179

180 self.num_heads = num_heads

181 self.head_size = head_size

49 - self.scale = scale

182 + self.scale = 1.0 / (head_size**0.5) if scale is None else scale

183 self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads

184 self.layer_name = prefix

52 -

185 assert self.num_heads % self.num_kv_heads == 0, (

186 f"num_heads ({self.num_heads}) is not "

187 f"divisible by num_kv_heads ({self.num_kv_heads})"

@@ -75,6 +207,9 @@ class MMEncoderAttention(CustomOp):

207 get_flash_attn_version() if self.is_flash_attn_backend else None

208 )

209

210 + if self.attn_backend == AttentionBackendEnum.FLASHINFER:

211 + _get_flashinfer_workspace_buffer()

212 +

213 logger.info_once(f"Using {self.attn_backend} for MMEncoderAttention.")

214

215 @classmethod

@@ -201,6 +336,27 @@ class MMEncoderAttention(CustomOp):

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

337 return output

338

339 + def _forward_flashinfer(

340 + self,

341 + query: torch.Tensor,

342 + key: torch.Tensor,

343 + value: torch.Tensor,

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

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

346 + sequence_lengths: torch.Tensor

347 + | None = None, # Only used for FlashInfer CuDNN backend

348 + ) -> torch.Tensor:

349 + return vit_flashinfer_wrapper(

350 + q=query,

351 + k=key,

352 + v=value,

353 + scale=self.scale,

354 + workspace_buffer=_get_flashinfer_workspace_buffer(),

355 + cu_seqlens=cu_seqlens,

356 + max_seqlen=max_seqlen,

357 + sequence_lengths=sequence_lengths,

358 + )

359 +

360 def forward_native(

361 self,

362 query: torch.Tensor,

@@ -208,6 +364,8 @@ class MMEncoderAttention(CustomOp):

364 value: torch.Tensor,

365 cu_seqlens: torch.Tensor | None = None,

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

367 + sequence_lengths: torch.Tensor

368 + | None = None, # Only used for FlashInfer CuDNN backend

369 ) -> torch.Tensor:

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

371

@@ -218,11 +376,17 @@ class MMEncoderAttention(CustomOp):

376 value: torch.Tensor,

377 cu_seqlens: torch.Tensor | None = None,

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

379 + sequence_lengths: torch.Tensor

380 + | None = None, # Only used for FlashInfer CuDNN backend

381 ) -> torch.Tensor:

382 if self.is_flash_attn_backend:

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

384 elif self.attn_backend == AttentionBackendEnum.TRITON_ATTN:

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

386 + elif self.attn_backend == AttentionBackendEnum.FLASHINFER:

387 + return self._forward_flashinfer(

388 + query, key, value, cu_seqlens, max_seqlen, sequence_lengths

389 + )

390 elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:

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

392 else:

@@ -238,6 +402,8 @@ class MMEncoderAttention(CustomOp):

402 value: torch.Tensor,

403 cu_seqlens: torch.Tensor | None = None,

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

405 + sequence_lengths: torch.Tensor

406 + | None = None, # Only used for FlashInfer CuDNN backend

407 ) -> torch.Tensor:

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

409

@@ -248,6 +414,8 @@ class MMEncoderAttention(CustomOp):

414 value: torch.Tensor,

415 cu_seqlens: torch.Tensor | None = None,

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

417 + sequence_lengths: torch.Tensor

418 + | None = None, # Only used for FlashInfer CuDNN backend

419 ) -> torch.Tensor:

420 if self.attn_backend == AttentionBackendEnum.FLASH_ATTN:

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

vllm/model_executor/models/qwen2_5_vl.py

@@ -357,6 +357,7 @@ class Qwen2_5_VisionAttention(nn.Module):

357 rotary_pos_emb_cos: torch.Tensor,

358 rotary_pos_emb_sin: torch.Tensor,

359 max_seqlen: torch.Tensor, # Only used for Flash Attention

360 + sequence_lengths: torch.Tensor, # Only used for FlashInfer CuDNN backend

361 ) -> torch.Tensor:

362 # [s, b, c] --> [s, b, head * 3 * head_dim]

363 x, _ = self.qkv(x)

@@ -398,6 +399,7 @@ class Qwen2_5_VisionAttention(nn.Module):

399 value=v,

400 cu_seqlens=cu_seqlens,

401 max_seqlen=max_seqlen,

402 + sequence_lengths=sequence_lengths,

403 )

404

405 context_layer = einops.rearrange(

@@ -463,6 +465,7 @@ class Qwen2_5_VisionBlock(nn.Module):

465 rotary_pos_emb_cos=rotary_pos_emb_cos,

466 rotary_pos_emb_sin=rotary_pos_emb_sin,

467 max_seqlen=max_seqlen,

468 + sequence_lengths=None,

469 )

470 x_fused_norm, residual = self.norm2(x, residual=x_attn)

471 x = residual + self.mlp(x_fused_norm)

vllm/model_executor/models/qwen3_vl.py

+33 -19

@@ -51,9 +51,12 @@ from transformers.video_utils import VideoMetadata

51 from vllm.compilation.decorators import support_torch_compile

52 from vllm.config import VllmConfig

53 from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions

54 -from vllm.distributed import get_pp_group

54 +from vllm.distributed import get_pp_group, parallel_state

55 from vllm.logger import init_logger

56 from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY

57 +from vllm.model_executor.layers.attention.mm_encoder_attention import (

58 + MMEncoderAttention,

59 +)

60 from vllm.model_executor.layers.conv import Conv3dLayer

61 from vllm.model_executor.layers.linear import (

62 ColumnParallelLinear,

@@ -92,7 +95,6 @@ from vllm.multimodal.processing import (

95 from vllm.sequence import IntermediateTensors

96 from vllm.utils.collection_utils import is_list_of

97 from vllm.utils.math_utils import round_up

95 -from vllm.v1.attention.backends.registry import AttentionBackendEnum

99 from .interfaces import (

100 MultiModalEmbeddings,

@@ -244,6 +246,7 @@ class Qwen3_VisionBlock(nn.Module):

246 rotary_pos_emb_cos: torch.Tensor,

247 rotary_pos_emb_sin: torch.Tensor,

248 max_seqlen: torch.Tensor, # Only used for Flash Attention

249 + sequence_lengths: torch.Tensor, # Only used for FlashInfer CuDNN backend

250 ) -> torch.Tensor:

251 x = x + self.attn(

252 self.norm1(x),

@@ -251,6 +254,7 @@ class Qwen3_VisionBlock(nn.Module):

254 rotary_pos_emb_cos=rotary_pos_emb_cos,

255 rotary_pos_emb_sin=rotary_pos_emb_sin,

256 max_seqlen=max_seqlen,

257 + sequence_lengths=sequence_lengths,

258 )

259

260 x = x + self.mlp(self.norm2(x))

@@ -332,6 +336,13 @@ class Qwen3_VisionTransformer(nn.Module):

336 )

337 self.num_grid_per_side = int(self.num_position_embeddings**0.5)

338

339 + use_data_parallel = is_vit_use_data_parallel()

340 + self.tp_size = (

341 + 1

342 + if use_data_parallel

343 + else parallel_state.get_tensor_model_parallel_world_size()

344 + )

345 +

346 # NOTE: This is used for creating empty tensor for all_gather for

347 # DP ViT. Here out_hidden_size is enlarged due to deepstack

348 self.out_hidden_size = vision_config.out_hidden_size * (

@@ -513,19 +524,6 @@ class Qwen3_VisionTransformer(nn.Module):

524

525 return torch.cat(outputs, dim=0)

526

516 - def compute_attn_mask_seqlen(

517 - self,

518 - cu_seqlens: torch.Tensor,

519 - ) -> torch.Tensor:

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

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

528 -

527 def forward(

528 self,

529 x: torch.Tensor,

@@ -549,11 +547,26 @@ class Qwen3_VisionTransformer(nn.Module):

547 axis=0, dtype=np.int32

548 )

549 cu_seqlens = np.concatenate([np.zeros(1, dtype=np.int32), cu_seqlens])

552 - cu_seqlens = torch.from_numpy(cu_seqlens)

553 -

550 + sequence_lengths = MMEncoderAttention.maybe_compute_sequence_lengths(

551 + self.attn_backend, cu_seqlens

552 + )

553 + if sequence_lengths is not None:

554 + sequence_lengths = torch.from_numpy(sequence_lengths).to(

555 + self.device, non_blocking=True

556 + )

557 + max_seqlen = torch.tensor(

558 + MMEncoderAttention.compute_max_seqlen(self.attn_backend, cu_seqlens),

559 + dtype=torch.int32,

560 + device=self.device,

561 + )

562 + cu_seqlens = MMEncoderAttention.maybe_recompute_cu_seqlens(

563 + self.attn_backend,

564 + cu_seqlens,

565 + self.hidden_size,

566 + self.tp_size,

567 + )

568 + cu_seqlens = torch.from_numpy(cu_seqlens).to(self.device, non_blocking=True)

569 hidden_states = hidden_states.unsqueeze(1)

555 - max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

556 - cu_seqlens = cu_seqlens.to(self.device, non_blocking=True)

570

571 deepstack_feature_lists = []

572 for layer_num, blk in enumerate(self.blocks):

@@ -563,6 +576,7 @@ class Qwen3_VisionTransformer(nn.Module):

576 rotary_pos_emb_cos=rotary_pos_emb_cos,

577 rotary_pos_emb_sin=rotary_pos_emb_sin,

578 max_seqlen=max_seqlen,

579 + sequence_lengths=sequence_lengths,

580 )

581 if layer_num in self.deepstack_visual_indexes:

582 deepstack_merger_idx = self.deepstack_visual_indexes.index(layer_num)

vllm/platforms/cuda.py

@@ -414,6 +414,7 @@ class CudaPlatformBase(Platform):

414 AttentionBackendEnum.FLASH_ATTN,

415 AttentionBackendEnum.TRITON_ATTN,

416 AttentionBackendEnum.TORCH_SDPA,

417 + AttentionBackendEnum.FLASHINFER,

418 ]

419

420 @classmethod

vllm/v1/attention/ops/vit_attn_wrappers.py

+88

@@ -268,3 +268,91 @@ def vit_torch_sdpa_wrapper(

268 return torch.ops.vllm.torch_sdpa_wrapper(

269 q, k, v, scale, cu_seqlens, enable_gqa=enable_gqa

270 )

271 +

272 +

273 +def flashinfer_wrapper(

274 + q: torch.Tensor,

275 + k: torch.Tensor,

276 + v: torch.Tensor,

277 + scale: float,

278 + workspace_buffer: torch.Tensor,

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

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

281 + sequence_lengths: torch.Tensor | None = None,

282 +) -> torch.Tensor:

283 + from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache

284 +

285 + is_reshaped = q.dim() == 4

286 +

287 + if is_reshaped:

288 + reshape_batch_size = q.shape[0]

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

290 + # cuDNN <= 9.10.2.21 requires q, k to be contiguous

291 + # this comes with no cost for ViTs with RoPE because

292 + # RoPE has already made q and k contiguous.

293 + q, k = q.contiguous(), k.contiguous()

294 +

295 + assert len(cu_seqlens) % 2 == 0, "cu_seqlens must be divisible by 2"

296 + cu_seqlength = len(cu_seqlens) // 2

297 + batch_offsets_qko = cu_seqlens[:cu_seqlength].view(-1, 1, 1, 1)

298 + batch_offsets_v = cu_seqlens[cu_seqlength:].view(-1, 1, 1, 1)

299 + sequence_lengths = sequence_lengths.view(-1, 1, 1, 1)

300 + max_seqlen = max_seqlen.item()

301 +

302 + output, _ = cudnn_batch_prefill_with_kv_cache(

303 + q,

304 + k,

305 + v,

306 + scale,

307 + workspace_buffer,

308 + max_token_per_sequence=max_seqlen,

309 + max_sequence_kv=max_seqlen,

310 + actual_seq_lens_q=sequence_lengths,

311 + actual_seq_lens_kv=sequence_lengths,

312 + causal=False,

313 + return_lse=False,

314 + batch_offsets_q=batch_offsets_qko,

315 + batch_offsets_k=batch_offsets_qko,

316 + batch_offsets_v=batch_offsets_v,

317 + batch_offsets_o=batch_offsets_qko,

318 + )

319 +

320 + if is_reshaped:

321 + output = einops.rearrange(output, "(b s) h d -> b s h d", b=reshape_batch_size)

322 +

323 + return output

324 +

325 +

326 +def vit_flashinfer_wrapper_fake(

327 + q: torch.Tensor,

328 + k: torch.Tensor,

329 + v: torch.Tensor,

330 + scale: float,

331 + workspace_buffer: torch.Tensor,

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

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

334 + sequence_lengths: torch.Tensor | None = None,

335 +) -> torch.Tensor:

336 + return torch.empty_like(q)

337 +

338 +

339 +direct_register_custom_op(

340 + op_name="flashinfer_wrapper",

341 + op_func=flashinfer_wrapper,

342 + fake_impl=vit_flashinfer_wrapper_fake,

343 +)

344 +

345 +

346 +def vit_flashinfer_wrapper(

347 + q: torch.Tensor,

348 + k: torch.Tensor,

349 + v: torch.Tensor,

350 + scale: float,

351 + workspace_buffer: torch.Tensor,

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

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

354 + sequence_lengths: torch.Tensor | None = None,

355 +) -> torch.Tensor:

356 + return torch.ops.vllm.flashinfer_wrapper(

357 + q, k, v, scale, workspace_buffer, cu_seqlens, max_seqlen, sequence_lengths

358 + )