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 14, 05:30 UTC
[Qwen][ROCm] Flash Attention Rotary Embeddings (#24642)
Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>
vllmellm committed
Oct 2, 2025 at 23:26 UTC
5e4a8223c6445fb2dbfe10e8c474e6d9ae6f914e
2 files changed
+28
-5
vllm/model_executor/layers/rotary_embedding/common.py
+23
@@ -2,15 +2,21 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import math
5
+from functools import cache
6
+from importlib.util import find_spec
7
+from typing import Callable
8
9
import torch
10
11
+from vllm.logger import init_logger
12
from vllm.platforms import current_platform
13
from vllm.utils import direct_register_custom_op
14
15
if current_platform.is_cuda():
16
from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
17
18
+logger = init_logger(__name__)
19
+
20
21
# common functions
22
def rotate_neox(x: torch.Tensor) -> torch.Tensor:
@@ -65,6 +71,23 @@ def apply_rotary_emb_dispatch(x: torch.Tensor, cos: torch.Tensor,
71
return apply_rotary_emb_torch(x, cos, sin, is_neox_style)
72
73
74
+@cache
75
+def dispatch_rotary_emb_function() -> Callable[..., torch.Tensor]:
76
+ if current_platform.is_cuda():
77
+ return apply_rotary_emb
78
+
79
+ if current_platform.is_rocm():
80
+ if find_spec("flash_attn") is not None:
81
+ from flash_attn.ops.triton.rotary import apply_rotary
82
+ return apply_rotary
83
+ else:
84
+ logger.warning(
85
+ "flash_attn is not installed. Falling back to PyTorch "
86
+ "implementation for rotary embeddings.")
87
+
88
+ return apply_rotary_emb_torch
89
+
90
+
91
# yarn functions
92
# Inverse dim formula to find dim based on number of rotations
93
def yarn_find_correction_dim(num_rotations: int,
vllm/model_executor/models/qwen2_vl.py
+5
-5
@@ -50,6 +50,8 @@ from vllm.model_executor.layers.activation import QuickGELU
50
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
51
RowParallelLinear)
52
from vllm.model_executor.layers.quantization import QuantizationConfig
53
+from vllm.model_executor.layers.rotary_embedding.common import (
54
+ dispatch_rotary_emb_function)
55
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
56
from vllm.model_executor.models.module_mapping import MultiModelKeys
57
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -63,7 +65,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
65
BaseProcessingInfo, PromptReplacement,
66
PromptUpdate)
67
from vllm.multimodal.profiling import BaseDummyInputsBuilder
66
-from vllm.platforms import _Backend, current_platform
68
+from vllm.platforms import _Backend
69
from vllm.sequence import IntermediateTensors
70
from vllm.transformers_utils.tokenizer import AnyTokenizer
71
from vllm.utils.tensor_schema import TensorSchema, TensorShape
@@ -272,13 +274,11 @@ def apply_rotary_emb_torch(x: torch.Tensor,
274
275
def apply_rotary_pos_emb_vision(t: torch.Tensor,
276
freqs: torch.Tensor) -> torch.Tensor:
277
+ rotary_emb_function = dispatch_rotary_emb_function()
278
t_ = t.float()
279
cos = freqs.cos()
280
sin = freqs.sin()
278
- apply_rotary_emb = apply_rotary_emb_torch
279
- if current_platform.is_cuda():
280
- from vllm.vllm_flash_attn.layers.rotary import apply_rotary_emb
281
- output = apply_rotary_emb(t_, cos, sin).type_as(t)
281
+ output = rotary_emb_function(t_, cos, sin).type_as(t)
282
return output
283
284