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
[ROCm] Split AITER unified attention into its own backend (#25507)
Signed-off-by: Gregory Shtrasberg <Gregory.Shtrasberg@amd.com>
Gregory Shtrasberg committed
Oct 6, 2025 at 18:49 UTC
f231e5bc21d5c972bffe3aae286c0102c5623e48
8 files changed
+325
-301
tests/compile/test_fusion_attn.py
+54
-164
@@ -7,9 +7,7 @@ import pytest
7
import torch._dynamo
8
9
from tests.compile.backend import LazyInitPass, TestBackend
10
-from tests.models.utils import check_outputs_equal
10
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
12
-from vllm import LLM, SamplingParams
11
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12
from vllm.attention import Attention, AttentionMetadata
13
from vllm.attention.backends.registry import _Backend
@@ -31,7 +29,6 @@ from vllm.config import (
29
)
30
from vllm.forward_context import get_forward_context, set_forward_context
31
from vllm.model_executor.layers.quantization.utils.quant_utils import (
34
- QuantKey,
32
kFp8StaticTensorSym,
33
kNvfp4Quant,
34
)
@@ -48,132 +45,6 @@ backend: Optional[TestBackend] = None
45
backend_unfused: Optional[TestBackend] = None
46
47
51
-@pytest.mark.parametrize(
52
- "model, quant_key", [("amd/Llama-3.1-8B-Instruct-FP8-KV", kFp8StaticTensorSym)]
53
-)
54
-@pytest.mark.parametrize("use_triton_fa", [True, False])
55
-@pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")
56
-@pytest.mark.skipif(
57
- not current_platform.is_rocm(), reason="V0 attn quant fusion only on ROCm"
58
-)
59
-def test_attention_fusion_v0(
60
- example_prompts, monkeypatch, model: str, quant_key: QuantKey, use_triton_fa: bool
61
-):
62
- # Clean Dynamo cache to avoid reusing other test cases
63
- # (for some reason the reset at the end is not enough)
64
- torch._dynamo.reset()
65
-
66
- # Use global backends
67
- global backend, backend_unfused
68
-
69
- monkeypatch.setenv("VLLM_USE_V1", "1")
70
- monkeypatch.setenv("VLLM_USE_TRITON_FLASH_ATTN", str(int(use_triton_fa)))
71
-
72
- # Prompt 4 seems too open-ended, differs between fused and unfused
73
- # (both outputs look reasonable though)
74
- prompts = example_prompts[:4] + example_prompts[5:]
75
-
76
- compile_config = CompilationConfig(
77
- # DYNAMO_AS_IS triggers custom backend & does full Dynamo compilation
78
- # DYNAMO_ONCE does not properly propagate shapes.
79
- level=CompilationLevel.DYNAMO_AS_IS,
80
- backend="tests.compile.test_fusion_attn.backend_unfused",
81
- custom_ops=["+quant_fp8"],
82
- )
83
- vllm_config = VllmConfig(
84
- compilation_config=compile_config,
85
- model_config=ModelConfig(
86
- model=model,
87
- dtype=torch.bfloat16,
88
- ),
89
- )
90
- backend_unfused = TestBackend(NoOpEliminationPass(vllm_config))
91
-
92
- llm = LLM(
93
- model,
94
- enforce_eager=True,
95
- compilation_config=compile_config,
96
- gpu_memory_utilization=0.5,
97
- max_model_len=2048,
98
- )
99
-
100
- sampling_params = SamplingParams(temperature=0.0, max_tokens=10, top_p=0.95)
101
-
102
- unfused_output = llm.generate(prompts, sampling_params)
103
- backend_unfused = None # Reset backend to make sure llm gets released
104
- del llm
105
-
106
- compile_config = CompilationConfig(
107
- # DYNAMO_AS_IS triggers custom backend & does full Dynamo compilation
108
- # DYNAMO_ONCE does not properly propagate shapes.
109
- level=CompilationLevel.DYNAMO_AS_IS,
110
- backend="tests.compile.test_fusion_attn.backend",
111
- custom_ops=["+quant_fp8"],
112
- )
113
- vllm_config = VllmConfig(
114
- compilation_config=compile_config,
115
- model_config=ModelConfig(
116
- model=model,
117
- dtype=torch.bfloat16,
118
- ),
119
- )
120
-
121
- # AttnFusionPass needs attention layers to be registered in config upon init
122
- # so we initialize it during compilation.
123
- attn_pass = LazyInitPass(AttnFusionPass, vllm_config)
124
- backend = TestBackend(NoOpEliminationPass(vllm_config), attn_pass)
125
- llm2 = LLM(
126
- model,
127
- enforce_eager=True,
128
- compilation_config=compile_config,
129
- gpu_memory_utilization=0.5,
130
- max_model_len=2048,
131
- )
132
-
133
- # check support
134
- attn_fusion_supported = [
135
- layer.impl.fused_output_quant_supported(quant_key)
136
- for key, layer in compile_config.static_forward_context.items()
137
- ]
138
-
139
- print(f"{attn_fusion_supported=}")
140
- if any(attn_fusion_supported):
141
- # Check quant ops
142
- backend.check_before_ops([QUANT_OPS[quant_key]], fully_replaced=False)
143
-
144
- # attention ops present in both, just output_scale param changes
145
- attn_nodes_pre = list(find_op_nodes(ATTN_OP, backend.graph_pre_pass))
146
- attn_nodes_post = list(find_op_nodes(ATTN_OP, backend.graph_post_pass))
147
- assert len(attn_nodes_pre) == len(attn_nodes_post)
148
-
149
- for i in range(len(attn_nodes_pre)):
150
- assert attn_nodes_pre[i].kwargs["output_scale"] is None
151
- fused = attn_nodes_post[i].kwargs["output_scale"] is not None
152
- assert fused == attn_fusion_supported[i], (
153
- f"Node {i} {'' if fused else 'not '} expected to have fused output quant"
154
- )
155
-
156
- # check outputs
157
- fused_output = llm2.generate(prompts, sampling_params)
158
-
159
- # transform outputs to format expected by check_outputs_equal
160
- sample_outs = lambda s: (list(s.token_ids), s.text)
161
- outs_lst = lambda ros: [sample_outs(ro.outputs[0]) for ro in ros]
162
-
163
- check_outputs_equal(
164
- outputs_0_lst=outs_lst(unfused_output),
165
- outputs_1_lst=outs_lst(fused_output),
166
- name_0="unfused",
167
- name_1="fused",
168
- )
169
-
170
- # Clean Dynamo cache to avoid polluting other case(s)
171
- torch._dynamo.reset()
172
-
173
- # Reset backend to make sure llm2 gets released
174
- backend = None
175
-
176
-
48
class AttentionQuantPatternModel(torch.nn.Module):
49
"""Base model for AttentionQuantPattern fusion."""
50
@@ -221,7 +92,7 @@ class AttentionQuantPatternModel(torch.nn.Module):
92
device=self.device,
93
)
94
224
- def build_attn_metadata(self, batch_size: int, use_hnd: bool) -> AttentionMetadata:
95
+ def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:
96
"""Initialize attention metadata."""
97
98
# Create common attn metadata
@@ -232,30 +103,57 @@ class AttentionQuantPatternModel(torch.nn.Module):
103
104
max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size
105
num_blocks = batch_size * max_blocks
106
+ backend = self.attn.backend
107
236
- # Create dummy KV cache for FlashInfer TRTLLM
237
- # - NHD: [num_blocks, block_size, num_kv_heads, head_size]
238
- # - HND: [num_blocks, num_kv_heads, block_size, head_size]
239
- kv_cache = torch.zeros(
240
- num_blocks,
241
- 2,
242
- self.num_kv_heads,
243
- self.block_size,
244
- self.head_size,
245
- dtype=self.kv_cache_dtype,
246
- device=self.device,
247
- )
248
- if current_platform.is_rocm():
108
+ # Create dummy KV cache for the selected backend
109
+ if backend == _Backend.ROCM_ATTN:
110
# k/v as 1st dimention
250
- if use_hnd:
251
- kv_cache = kv_cache.permute(1, 0, 2, 3, 4)
252
- else:
253
- kv_cache = kv_cache.permute(1, 0, 3, 2, 4)
254
- else:
111
+ # HND: [num_blocks, num_kv_heads, block_size, head_size]
112
+ kv_cache = torch.zeros(
113
+ 2,
114
+ num_blocks,
115
+ self.num_kv_heads,
116
+ self.block_size,
117
+ self.head_size,
118
+ dtype=self.kv_cache_dtype,
119
+ device=self.device,
120
+ )
121
+ elif backend == _Backend.ROCM_AITER_UNIFIED_ATTN:
122
+ # k/v as 1st dimention
123
+ # NHD: [num_blocks, block_size, num_kv_heads, head_size]
124
+ kv_cache = torch.zeros(
125
+ 2,
126
+ num_blocks,
127
+ self.block_size,
128
+ self.num_kv_heads,
129
+ self.head_size,
130
+ dtype=self.kv_cache_dtype,
131
+ device=self.device,
132
+ )
133
+ elif backend == _Backend.TRITON_ATTN:
134
# k/v as 2nd dimention
256
- # Create kv_cache in HND layout and permute to NHD layout
257
- # (later will be permuted back to HND layout in forward pass)
258
- kv_cache = kv_cache.permute(0, 1, 3, 2, 4)
135
+ # NHD: [num_blocks, block_size, num_kv_heads, head_size]
136
+ kv_cache = torch.zeros(
137
+ num_blocks,
138
+ 2,
139
+ self.num_kv_heads,
140
+ self.block_size,
141
+ self.head_size,
142
+ dtype=self.kv_cache_dtype,
143
+ device=self.device,
144
+ )
145
+ elif backend == _Backend.FLASHINFER:
146
+ kv_cache = torch.zeros(
147
+ num_blocks,
148
+ 2,
149
+ self.num_kv_heads,
150
+ self.block_size,
151
+ self.head_size,
152
+ dtype=self.kv_cache_dtype,
153
+ device=self.device,
154
+ ).permute(0, 1, 3, 2, 4)
155
+ else:
156
+ raise ValueError(f"Unsupported backend: {backend}")
157
self.attn.kv_cache = [kv_cache]
158
159
# Build attn metadata
@@ -375,10 +273,9 @@ else:
273
@pytest.mark.parametrize("model_name, model_class", MODELS)
274
@pytest.mark.parametrize(
275
"backend",
378
- [_Backend.FLASHINFER] if current_platform.is_cuda() else [_Backend.TRITON_ATTN],
379
-)
380
-@pytest.mark.parametrize(
381
- "split_attention", [False, True] if current_platform.is_rocm() else [False]
276
+ [_Backend.FLASHINFER]
277
+ if current_platform.is_cuda()
278
+ else [_Backend.ROCM_AITER_UNIFIED_ATTN, _Backend.ROCM_ATTN, _Backend.TRITON_ATTN],
279
)
280
# TODO(boyuan): test inductor graph partition on rocm
281
@pytest.mark.parametrize(
@@ -405,7 +302,6 @@ def test_attention_quant_pattern(
302
model_name: str,
303
model_class: type[AttentionQuantPatternModel],
304
backend: _Backend,
408
- split_attention: bool,
305
use_inductor_graph_partition: bool,
306
monkeypatch,
307
dist_init,
@@ -417,8 +313,6 @@ def test_attention_quant_pattern(
313
pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
314
315
monkeypatch.setenv("VLLM_USE_V1", "1")
420
- if split_attention:
421
- monkeypatch.setenv("VLLM_V1_USE_PREFILL_DECODE_ATTENTION", "1")
316
317
device = torch.device("cuda:0")
318
torch.manual_seed(42)
@@ -466,9 +360,7 @@ def test_attention_quant_pattern(
360
model_unfused = model_unfused.to(device)
361
362
forward_ctx = get_forward_context()
469
- forward_ctx.attn_metadata = model_unfused.build_attn_metadata(
470
- batch_size, use_hnd=split_attention
471
- )
363
+ forward_ctx.attn_metadata = model_unfused.build_attn_metadata(batch_size)
364
365
# Run model directly without compilation and fusion
366
result_unfused = model_unfused(q, k, v)
@@ -494,9 +386,7 @@ def test_attention_quant_pattern(
386
model_fused = model_fused.to(device)
387
388
forward_ctx = get_forward_context()
497
- forward_ctx.attn_metadata = model_fused.build_attn_metadata(
498
- batch_size, use_hnd=split_attention
499
- )
389
+ forward_ctx.attn_metadata = model_fused.build_attn_metadata(batch_size)
390
391
# Create test backend with fusion passes enabled
392
noop_pass = NoOpEliminationPass(vllm_config)
vllm/attention/backends/registry.py
+1
@@ -25,3 +25,4 @@ class _Backend(enum.Enum):
25
FLEX_ATTENTION = enum.auto()
26
TREE_ATTN = enum.auto()
27
ROCM_ATTN = enum.auto()
28
+ ROCM_AITER_UNIFIED_ATTN = enum.auto()
vllm/attention/selector.py
+1
@@ -254,3 +254,4 @@ def global_force_attn_backend_context_manager(
254
finally:
255
# Revert the original global backend override, if any
256
global_force_attn_backend(original_value)
257
+ _cached_get_attn_backend.cache_clear()
vllm/engine/arg_utils.py
+1
@@ -1623,6 +1623,7 @@ class EngineArgs:
1623
"TREE_ATTN",
1624
"XFORMERS",
1625
"ROCM_ATTN",
1626
+ "ROCM_AITER_UNIFIED_ATTN",
1627
]
1628
if (
1629
envs.is_set("VLLM_ATTENTION_BACKEND")
vllm/envs.py
+7
-6
@@ -18,7 +18,6 @@ if TYPE_CHECKING:
18
LD_LIBRARY_PATH: Optional[str] = None
19
VLLM_USE_TRITON_FLASH_ATTN: bool = True
20
VLLM_V1_USE_PREFILL_DECODE_ATTENTION: bool = False
21
- VLLM_USE_AITER_UNIFIED_ATTENTION: bool = False
21
VLLM_FLASH_ATTN_VERSION: Optional[int] = None
22
LOCAL_RANK: int = 0
23
CUDA_VISIBLE_DEVICES: Optional[str] = None
@@ -109,6 +108,7 @@ if TYPE_CHECKING:
108
VLLM_ROCM_USE_AITER_FP4_ASM_GEMM: bool = False
109
VLLM_ROCM_USE_TRITON_ROPE: bool = False
110
VLLM_ROCM_USE_AITER_FP8BMM: bool = True
111
+ VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False
112
VLLM_ROCM_USE_SKINNY_GEMM: bool = True
113
VLLM_ROCM_FP8_PADDING: bool = True
114
VLLM_ROCM_MOE_PADDING: bool = True
@@ -475,10 +475,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
475
os.getenv("VLLM_V1_USE_PREFILL_DECODE_ATTENTION", "False").lower()
476
in ("true", "1")
477
),
478
- # Use AITER triton unified attention for V1 attention
479
- "VLLM_USE_AITER_UNIFIED_ATTENTION": lambda: (
480
- os.getenv("VLLM_USE_AITER_UNIFIED_ATTENTION", "False").lower() in ("true", "1")
481
- ),
478
# Force vllm to use a specific flash-attention version (2 or 3), only valid
479
# when using the flash-attention backend.
480
"VLLM_FLASH_ATTN_VERSION": lambda: maybe_convert_int(
@@ -896,6 +892,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
892
"VLLM_ROCM_USE_AITER_FP8BMM": lambda: (
893
os.getenv("VLLM_ROCM_USE_AITER_FP8BMM", "True").lower() in ("true", "1")
894
),
895
+ # Use AITER triton unified attention for V1 attention
896
+ "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION": lambda: (
897
+ os.getenv("VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION", "False").lower()
898
+ in ("true", "1")
899
+ ),
900
# use rocm skinny gemms
901
"VLLM_ROCM_USE_SKINNY_GEMM": lambda: (
902
os.getenv("VLLM_ROCM_USE_SKINNY_GEMM", "True").lower() in ("true", "1")
@@ -1434,7 +1435,6 @@ def compute_hash() -> str:
1435
"VLLM_FUSED_MOE_CHUNK_SIZE",
1436
"VLLM_FLASHINFER_MOE_BACKEND",
1437
"VLLM_V1_USE_PREFILL_DECODE_ATTENTION",
1437
- "VLLM_USE_AITER_UNIFIED_ATTENTION",
1438
"VLLM_ATTENTION_BACKEND",
1439
"VLLM_USE_FLASHINFER_SAMPLER",
1440
"VLLM_DISABLED_KERNELS",
@@ -1462,6 +1462,7 @@ def compute_hash() -> str:
1462
"VLLM_ROCM_USE_AITER_FP4_ASM_GEMM",
1463
"VLLM_ROCM_USE_TRITON_ROPE",
1464
"VLLM_ROCM_USE_AITER_FP8BMM",
1465
+ "VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION",
1466
"VLLM_ROCM_USE_SKINNY_GEMM",
1467
"VLLM_ROCM_FP8_PADDING",
1468
"VLLM_ROCM_MOE_PADDING",
vllm/platforms/rocm.py
+18
-10
@@ -276,25 +276,33 @@ class RocmPlatform(Platform):
276
)
277
278
if envs.VLLM_USE_V1:
279
- if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
280
- logger.info("Using Flash Attention backend on V1 engine.")
279
+ if (
280
+ envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9()
281
+ ) or selected_backend == _Backend.ROCM_AITER_FA:
282
+ logger.info("Using Aiter Flash Attention backend on V1 engine.")
283
return (
284
"vllm.v1.attention.backends."
285
"rocm_aiter_fa.AiterFlashAttentionBackend"
286
)
285
- elif (
286
- (envs.VLLM_ROCM_USE_AITER and envs.VLLM_USE_AITER_UNIFIED_ATTENTION)
287
- or envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
287
+ if (
288
+ envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION
289
+ ) or selected_backend == _Backend.ROCM_AITER_UNIFIED_ATTN:
290
+ logger.info("Using Aiter Unified Attention backend on V1 engine.")
291
+ return (
292
+ "vllm.v1.attention.backends."
293
+ "rocm_aiter_unified_attn.RocmAiterUnifiedAttentionBackend"
294
+ )
295
+ if (
296
+ envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
297
or selected_backend == _Backend.ROCM_ATTN
298
):
299
# rocm specific backend, with aiter and/or
300
# triton prefix-prefill
292
- logger.info("Using Rocm/Aiter Attention backend on V1 engine.")
301
+ logger.info("Using Rocm Attention backend on V1 engine.")
302
return "vllm.v1.attention.backends.rocm_attn.RocmAttentionBackend"
294
- else:
295
- # default case, using triton unified attention
296
- logger.info("Using Triton Attention backend on V1 engine.")
297
- return "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"
303
+ # default case, using triton unified attention
304
+ logger.info("Using Triton Attention backend on V1 engine.")
305
+ return "vllm.v1.attention.backends.triton_attn.TritonAttentionBackend"
306
raise RuntimeError(
307
"V0 attention backends have been removed. Set VLLM_USE_V1=1 "
308
"to select a supported backend."
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
new
+203
@@ -0,0 +1,203 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Attention layer with PagedAttention and Triton prefix prefill."""
4
+
5
+from typing import Optional
6
+
7
+import torch
8
+
9
+from vllm import _custom_ops as ops
10
+from vllm.attention.backends.abstract import AttentionMetadata, AttentionType
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ QuantKey,
14
+ kFp8StaticTensorSym,
15
+)
16
+from vllm.v1.attention.backends.flash_attn import FlashAttentionMetadata
17
+from vllm.v1.attention.backends.rocm_attn import (
18
+ RocmAttentionBackend,
19
+ RocmAttentionImpl,
20
+ RocmAttentionMetadata,
21
+ RocmAttentionMetadataBuilder,
22
+)
23
+
24
+logger = init_logger(__name__)
25
+
26
+
27
+class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
28
+ accept_output_buffer: bool = True
29
+
30
+ @staticmethod
31
+ def get_name() -> str:
32
+ return "ROCM_AITER_UNIFIED_ATTN"
33
+
34
+ @staticmethod
35
+ def get_impl_cls() -> type["RocmAiterUnifiedAttentionImpl"]:
36
+ return RocmAiterUnifiedAttentionImpl
37
+
38
+ @staticmethod
39
+ def get_metadata_cls() -> type["AttentionMetadata"]:
40
+ return RocmAttentionMetadata
41
+
42
+ @staticmethod
43
+ def get_kv_cache_shape(
44
+ num_blocks: int,
45
+ block_size: int,
46
+ num_kv_heads: int,
47
+ head_size: int,
48
+ cache_dtype_str: str = "auto",
49
+ ) -> tuple[int, ...]:
50
+ if block_size % 16 != 0:
51
+ raise ValueError("Block size must be a multiple of 16.")
52
+ return (2, num_blocks, block_size, num_kv_heads, head_size)
53
+
54
+ @staticmethod
55
+ def use_cascade_attention(*args, **kwargs) -> bool:
56
+ return False
57
+
58
+ @staticmethod
59
+ def get_builder_cls() -> type["RocmAttentionMetadataBuilder"]:
60
+ return RocmAttentionMetadataBuilder
61
+
62
+
63
+class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
64
+ def fused_output_quant_supported(self, quant_key: QuantKey):
65
+ return quant_key == kFp8StaticTensorSym
66
+
67
+ def __init__(
68
+ self,
69
+ num_heads: int,
70
+ head_size: int,
71
+ scale: float,
72
+ num_kv_heads: int,
73
+ alibi_slopes: Optional[list[float]],
74
+ sliding_window: Optional[int],
75
+ kv_cache_dtype: str,
76
+ logits_soft_cap: Optional[float] = None,
77
+ attn_type: AttentionType = AttentionType.DECODER,
78
+ kv_sharing_target_layer_name: Optional[int] = None,
79
+ sinks: Optional[torch.Tensor] = None,
80
+ ) -> None:
81
+ super().__init__(
82
+ num_heads,
83
+ head_size,
84
+ scale,
85
+ num_kv_heads,
86
+ alibi_slopes,
87
+ sliding_window,
88
+ kv_cache_dtype,
89
+ logits_soft_cap,
90
+ attn_type,
91
+ kv_sharing_target_layer_name,
92
+ sinks,
93
+ )
94
+ logger.info_once(
95
+ "Using aiter unified attention for RocmAiterUnifiedAttentionImpl"
96
+ )
97
+ from aiter.ops.triton.unified_attention import unified_attention
98
+
99
+ self.unified_attention = unified_attention
100
+
101
+ def forward(
102
+ self,
103
+ layer: torch.nn.Module,
104
+ query: torch.Tensor,
105
+ key: torch.Tensor,
106
+ value: torch.Tensor,
107
+ kv_cache: torch.Tensor,
108
+ attn_metadata: FlashAttentionMetadata,
109
+ output: Optional[torch.Tensor] = None,
110
+ output_scale: Optional[torch.Tensor] = None,
111
+ output_block_scale: Optional[torch.Tensor] = None,
112
+ ) -> torch.Tensor:
113
+ """Forward pass with FlashAttention.
114
+
115
+ Args:
116
+ query: shape = [num_tokens, num_heads, head_size]
117
+ key: shape = [num_tokens, num_kv_heads, head_size]
118
+ value: shape = [num_tokens, num_kv_heads, head_size]
119
+ kv_cache: shape =
120
+ [2, num_blocks, block_size, num_kv_heads, head_size]
121
+ attn_metadata: Metadata for attention.
122
+ Returns:
123
+ shape = [num_tokens, num_heads * head_size]
124
+ """
125
+ assert output is not None, "Output tensor must be provided."
126
+
127
+ if output_block_scale is not None:
128
+ raise NotImplementedError(
129
+ "fused block_scale output quantization is not yet supported"
130
+ " for RocmAttentionImpl"
131
+ )
132
+
133
+ if attn_metadata is None:
134
+ # Profiling run.
135
+ return output
136
+
137
+ assert attn_metadata.use_cascade is False
138
+
139
+ # IMPORTANT!
140
+ # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
141
+ # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
142
+ # in this method. For example, `view` and `slice` (or `[:n]`) operations
143
+ # are surprisingly slow even in the case they do not invoke any GPU ops.
144
+ # Minimize the PyTorch ops in this method as much as possible.
145
+ # Whenever making a change in this method, please benchmark the
146
+ # performance to make sure it does not introduce any overhead.
147
+
148
+ num_actual_tokens = attn_metadata.num_actual_tokens
149
+
150
+ key_cache, value_cache = kv_cache.unbind(0)
151
+
152
+ if self.kv_sharing_target_layer_name is None:
153
+ # Reshape the input keys and values and store them in the cache.
154
+ # Skip this if sharing KV cache with an earlier attention layer.
155
+ ops.reshape_and_cache_flash(
156
+ key,
157
+ value,
158
+ key_cache,
159
+ value_cache,
160
+ attn_metadata.slot_mapping,
161
+ self.kv_cache_dtype,
162
+ layer._k_scale,
163
+ layer._v_scale,
164
+ )
165
+
166
+ if self.kv_cache_dtype.startswith("fp8"):
167
+ key_cache = key_cache.view(self.fp8_dtype)
168
+ value_cache = value_cache.view(self.fp8_dtype)
169
+ assert layer._q_scale_float == 1.0, (
170
+ "A non 1.0 q_scale is not currently supported."
171
+ )
172
+
173
+ cu_seqlens_q = attn_metadata.query_start_loc
174
+ seqused_k = attn_metadata.seq_lens
175
+ max_seqlen_q = attn_metadata.max_query_len
176
+ max_seqlen_k = attn_metadata.max_seq_len
177
+ block_table = attn_metadata.block_table
178
+
179
+ descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
180
+
181
+ self.unified_attention(
182
+ q=query[:num_actual_tokens],
183
+ k=key_cache,
184
+ v=value_cache,
185
+ out=output[:num_actual_tokens],
186
+ cu_seqlens_q=cu_seqlens_q,
187
+ max_seqlen_q=max_seqlen_q,
188
+ seqused_k=seqused_k,
189
+ max_seqlen_k=max_seqlen_k,
190
+ softmax_scale=self.scale,
191
+ causal=True,
192
+ alibi_slopes=self.alibi_slopes,
193
+ window_size=self.sliding_window,
194
+ block_table=block_table,
195
+ softcap=self.logits_soft_cap,
196
+ q_descale=None, # Not supported
197
+ k_descale=layer._k_scale.expand(descale_shape),
198
+ v_descale=layer._v_scale.expand(descale_shape),
199
+ sinks=self.sinks,
200
+ output_scale=output_scale,
201
+ )
202
+
203
+ return output
vllm/v1/attention/backends/rocm_attn.py
+40
-121
@@ -3,13 +3,10 @@
3
"""Attention layer with PagedAttention and Triton prefix prefill."""
4
5
from dataclasses import dataclass
6
-from functools import cache
6
from typing import ClassVar, Optional
7
8
import torch
9
11
-from vllm import _custom_ops as ops
12
-from vllm import envs
10
from vllm.attention.backends.abstract import (
11
AttentionBackend,
12
AttentionImpl,
@@ -96,12 +93,11 @@ class RocmAttentionMetadataBuilder(AttentionMetadataBuilder[RocmAttentionMetadat
93
# slow, so here we set it to 1.
94
attn_metadata.seq_lens.fill_(1)
95
99
- if envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION:
100
- # Here we set the query start locs to 0. This is to
101
- # cover up an invalid memory access in the prefix_prefil kernel
102
- # that we run into during graph capture (#25985)
103
- common_attn_metadata.query_start_loc.zero_()
104
- common_attn_metadata.query_start_loc_cpu.zero_()
96
+ # Here we set the query start locs to 0. This is to
97
+ # cover up an invalid memory access in the prefix_prefil kernel
98
+ # that we run into during graph capture (#25985)
99
+ common_attn_metadata.query_start_loc.zero_()
100
+ common_attn_metadata.query_start_loc_cpu.zero_()
101
102
return attn_metadata
103
@@ -211,14 +207,6 @@ class RocmAttentionBackend(AttentionBackend):
207
return RocmAttentionMetadataBuilder
208
209
214
-@cache
215
-def use_aiter_unified_attention() -> bool:
216
- """Check if aiter unified attention should be used."""
217
- # VLLM_ROCM_USE_AITER_MHA needs to set to 0 as well as it is set
218
- # to 1 as default
219
- return envs.VLLM_ROCM_USE_AITER and envs.VLLM_USE_AITER_UNIFIED_ATTENTION
220
-
221
-
210
class RocmAttentionImpl(AttentionImpl):
211
def fused_output_quant_supported(self, quant_key: QuantKey):
212
return quant_key == kFp8StaticTensorSym
@@ -268,23 +256,6 @@ class RocmAttentionImpl(AttentionImpl):
256
)
257
258
self.fp8_dtype = current_platform.fp8_dtype()
271
- self.force_prefill_decode_attn = envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION
272
-
273
- if not self.force_prefill_decode_attn:
274
- # If not using prefill decode attention, we use the Triton
275
- # unified attention implementation.
276
- if use_aiter_unified_attention():
277
- logger.info_once("Using aiter unified attention for RocmAttentionImpl")
278
- from aiter.ops.triton.unified_attention import unified_attention
279
-
280
- self.unified_attention = unified_attention
281
- else:
282
- logger.info_once("Using vllm unified attention for RocmAttentionImpl")
283
- from vllm.attention.ops.triton_unified_attention import (
284
- unified_attention,
285
- )
286
-
287
- self.unified_attention = unified_attention
259
260
self.sinks = sinks
261
if sinks is not None:
@@ -341,58 +312,32 @@ class RocmAttentionImpl(AttentionImpl):
312
# Whenever making a change in this method, please benchmark the
313
# performance to make sure it does not introduce any overhead.
314
344
- use_prefill_decode_attn = self.force_prefill_decode_attn
315
num_actual_tokens = attn_metadata.num_actual_tokens
316
347
- if use_prefill_decode_attn:
348
- key_cache, value_cache = PagedAttention.split_kv_cache(
349
- kv_cache, self.num_kv_heads, self.head_size
350
- )
351
- else:
352
- key_cache, value_cache = kv_cache.unbind(0)
317
+ key_cache, value_cache = PagedAttention.split_kv_cache(
318
+ kv_cache, self.num_kv_heads, self.head_size
319
+ )
320
321
if self.kv_sharing_target_layer_name is None:
322
# Reshape the input keys and values and store them in the cache.
323
# Skip this if sharing KV cache with an earlier attention layer.
357
- if use_prefill_decode_attn:
358
- PagedAttention.write_to_paged_cache(
359
- key,
360
- value,
361
- key_cache,
362
- value_cache,
363
- attn_metadata.slot_mapping,
364
- self.kv_cache_dtype,
365
- layer._k_scale,
366
- layer._v_scale,
367
- )
368
- else:
369
- ops.reshape_and_cache_flash(
370
- key,
371
- value,
372
- key_cache,
373
- value_cache,
374
- attn_metadata.slot_mapping,
375
- self.kv_cache_dtype,
376
- layer._k_scale,
377
- layer._v_scale,
378
- )
324
+ PagedAttention.write_to_paged_cache(
325
+ key,
326
+ value,
327
+ key_cache,
328
+ value_cache,
329
+ attn_metadata.slot_mapping,
330
+ self.kv_cache_dtype,
331
+ layer._k_scale,
332
+ layer._v_scale,
333
+ )
334
335
if self.kv_cache_dtype.startswith("fp8"):
336
key_cache = key_cache.view(self.fp8_dtype)
337
value_cache = value_cache.view(self.fp8_dtype)
383
- num_tokens, num_heads, head_size = query.shape
338
assert layer._q_scale_float == 1.0, (
339
"A non 1.0 q_scale is not currently supported."
340
)
387
- if current_platform.is_cuda():
388
- # Skip Q quantization on ROCm and XPU, enable this on cuda
389
- # only, since dequantizing back to f32 in the attention kernel
390
- # is not supported.
391
- query, _ = ops.scaled_fp8_quant(
392
- query.reshape((num_tokens, num_heads * head_size)).contiguous(),
393
- layer._q_scale,
394
- )
395
- query = query.reshape((num_tokens, num_heads, head_size))
341
342
cu_seqlens_q = attn_metadata.query_start_loc
343
seqused_k = attn_metadata.seq_lens
@@ -400,53 +345,27 @@ class RocmAttentionImpl(AttentionImpl):
345
max_seqlen_k = attn_metadata.max_seq_len
346
block_table = attn_metadata.block_table
347
403
- if use_prefill_decode_attn:
404
- # Compute attention and update output up to `num_actual_tokens`.
405
- chunked_prefill_paged_decode(
406
- query=query[:num_actual_tokens],
407
- key=key[:num_actual_tokens],
408
- value=value[:num_actual_tokens],
409
- output=output[:num_actual_tokens],
410
- kv_cache_dtype=self.kv_cache_dtype,
411
- key_cache=key_cache,
412
- value_cache=value_cache,
413
- block_table=block_table,
414
- query_start_loc=cu_seqlens_q,
415
- seq_lens=seqused_k,
416
- max_seq_len=max_seqlen_k,
417
- max_query_len=max_seqlen_q,
418
- k_scale=layer._k_scale,
419
- v_scale=layer._v_scale,
420
- alibi_slopes=self.alibi_slopes,
421
- sliding_window=self.sliding_window[0],
422
- sm_scale=self.scale,
423
- output_scale=output_scale,
424
- sinks=self.sinks,
425
- )
426
-
427
- else:
428
- descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
429
-
430
- self.unified_attention(
431
- q=query[:num_actual_tokens],
432
- k=key_cache,
433
- v=value_cache,
434
- out=output[:num_actual_tokens],
435
- cu_seqlens_q=cu_seqlens_q,
436
- max_seqlen_q=max_seqlen_q,
437
- seqused_k=seqused_k,
438
- max_seqlen_k=max_seqlen_k,
439
- softmax_scale=self.scale,
440
- causal=True,
441
- alibi_slopes=self.alibi_slopes,
442
- window_size=self.sliding_window,
443
- block_table=block_table,
444
- softcap=self.logits_soft_cap,
445
- q_descale=None, # Not supported
446
- k_descale=layer._k_scale.expand(descale_shape),
447
- v_descale=layer._v_scale.expand(descale_shape),
448
- sinks=self.sinks,
449
- output_scale=output_scale,
450
- )
348
+ # Compute attention and update output up to `num_actual_tokens`.
349
+ chunked_prefill_paged_decode(
350
+ query=query[:num_actual_tokens],
351
+ key=key[:num_actual_tokens],
352
+ value=value[:num_actual_tokens],
353
+ output=output[:num_actual_tokens],
354
+ kv_cache_dtype=self.kv_cache_dtype,
355
+ key_cache=key_cache,
356
+ value_cache=value_cache,
357
+ block_table=block_table,
358
+ query_start_loc=cu_seqlens_q,
359
+ seq_lens=seqused_k,
360
+ max_seq_len=max_seqlen_k,
361
+ max_query_len=max_seqlen_q,
362
+ k_scale=layer._k_scale,
363
+ v_scale=layer._v_scale,
364
+ alibi_slopes=self.alibi_slopes,
365
+ sliding_window=self.sliding_window[0],
366
+ sm_scale=self.scale,
367
+ output_scale=output_scale,
368
+ sinks=self.sinks,
369
+ )
370
371
return output