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 15, 21:00 UTC
[core] Perf improvement for DSv3 on AMD GPUs (#13718)
Signed-off-by: qli88 <qiang.li2@amd.com>
qli88 committed
Feb 27, 2025 at 16:14 UTC
8294773e48a2d5cde4bb48b8607a10c14de6afbf
3 files changed
+210
-25
vllm/attention/backends/mla/common.py
+72
-20
@@ -237,14 +237,20 @@ from vllm.utils import async_tensor_h2d, cdiv, make_tensor_with_pad, round_down
237
238
try:
239
from vllm.vllm_flash_attn import flash_attn_varlen_func
240
+ is_vllm_fa = True
241
except ImportError:
242
# For rocm use upstream flash attention
243
from flash_attn import flash_attn_varlen_func
244
+ is_vllm_fa = False
245
+
246
+from vllm.attention.ops.triton_flash_attention import triton_attention
247
248
if TYPE_CHECKING:
249
from vllm.worker.model_runner import (ModelInputForGPUBuilder,
250
ModelInputForGPUWithSamplingMetadata)
251
252
+is_hip = current_platform.is_rocm()
253
+
254
255
class MLACommonBackend(AttentionBackend):
256
@@ -1046,12 +1052,13 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1052
self.q_proj = q_proj
1053
self.kv_b_proj = kv_b_proj
1054
self.o_proj = o_proj
1049
- self.vllm_flash_attn_version = get_flash_attn_version()
1055
+ self.triton_fa_func = triton_attention
1056
1057
# Handle the differences between the flash_attn_varlen from flash_attn
1058
# and the one from vllm_flash_attn. The former is used on RoCM and the
1059
# latter has an additional parameter to control FA2 vs FA3
1060
self.flash_attn_varlen_func = flash_attn_varlen_func
1061
+ self.vllm_flash_attn_version = get_flash_attn_version()
1062
if self.vllm_flash_attn_version is not None:
1063
self.flash_attn_varlen_func = \
1064
functools.partial(flash_attn_varlen_func,
@@ -1315,18 +1322,48 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1322
[0, q.shape[-1] - v.shape[-1]],
1323
value=0)
1324
1318
- attn_output, attn_softmax_lse = self.flash_attn_varlen_func(
1319
- q=q,
1320
- k=k,
1321
- v=v_padded,
1322
- cu_seqlens_q=prefill_metadata.query_start_loc,
1323
- cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],
1324
- max_seqlen_q=prefill_metadata.max_query_len,
1325
- max_seqlen_k=prefill_metadata.context_chunk_max_seq_lens[i],
1326
- softmax_scale=self.scale,
1327
- causal=False, # Context is unmasked
1328
- return_softmax_lse=True,
1329
- )
1325
+ if is_hip and envs.VLLM_USE_TRITON_FLASH_ATTN:
1326
+ attn_output, attn_softmax_lse = self.triton_fa_func(
1327
+ q,
1328
+ k,
1329
+ v_padded,
1330
+ None,
1331
+ prefill_metadata.query_start_loc,
1332
+ prefill_metadata.context_chunk_cu_seq_lens[i],
1333
+ prefill_metadata.max_query_len,
1334
+ prefill_metadata.context_chunk_max_seq_lens[i],
1335
+ False, # causal
1336
+ self.scale,
1337
+ None, # attn_mask is None unless applying ALiBi mask
1338
+ )
1339
+ elif is_vllm_fa:
1340
+ attn_output, attn_softmax_lse = self.flash_attn_varlen_func(
1341
+ q=q,
1342
+ k=k,
1343
+ v=v_padded,
1344
+ cu_seqlens_q=prefill_metadata.query_start_loc,
1345
+ cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],
1346
+ max_seqlen_q=prefill_metadata.max_query_len,
1347
+ max_seqlen_k=prefill_metadata.
1348
+ context_chunk_max_seq_lens[i],
1349
+ softmax_scale=self.scale,
1350
+ causal=False, # Context is unmasked
1351
+ return_softmax_lse=True,
1352
+ )
1353
+ else:
1354
+ attn_output, attn_softmax_lse, _ = self.flash_attn_varlen_func(
1355
+ q=q,
1356
+ k=k,
1357
+ v=v_padded,
1358
+ cu_seqlens_q=prefill_metadata.query_start_loc,
1359
+ cu_seqlens_k=prefill_metadata.context_chunk_cu_seq_lens[i],
1360
+ max_seqlen_q=prefill_metadata.max_query_len,
1361
+ max_seqlen_k=prefill_metadata.
1362
+ context_chunk_max_seq_lens[i],
1363
+ softmax_scale=self.scale,
1364
+ causal=False, # Context is unmasked
1365
+ return_attn_probs=True,
1366
+ )
1367
1368
if output is None:
1369
output = attn_output
@@ -1374,11 +1411,24 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1411
v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
1412
value=0)
1413
1377
- if has_context:
1378
- if not current_platform.is_cuda():
1379
- raise NotImplementedError(
1380
- "Chunked Prefill for MLA is not currently supported on"
1381
- "non-cuda platforms")
1414
+ if is_hip and envs.VLLM_USE_TRITON_FLASH_ATTN:
1415
+ output = self.triton_fa_func(
1416
+ q,
1417
+ k,
1418
+ v_padded,
1419
+ None,
1420
+ prefill_metadata.query_start_loc,
1421
+ prefill_metadata.query_start_loc,
1422
+ prefill_metadata.max_prefill_seq_len,
1423
+ prefill_metadata.max_prefill_seq_len,
1424
+ True, # causal
1425
+ self.scale,
1426
+ None, # attn_mask is None unless applying ALiBi mask
1427
+ )
1428
+ ## triton flash attention always return 2 objects
1429
+ if not has_context:
1430
+ output = output[0]
1431
+ elif is_vllm_fa:
1432
output = self.flash_attn_varlen_func(
1433
q=q,
1434
k=k,
@@ -1389,7 +1439,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1439
max_seqlen_k=prefill_metadata.max_prefill_seq_len,
1440
softmax_scale=self.scale,
1441
causal=True,
1392
- return_softmax_lse=True,
1442
+ return_softmax_lse=has_context,
1443
)
1444
else:
1445
output = self.flash_attn_varlen_func(
@@ -1402,10 +1452,12 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1452
max_seqlen_k=prefill_metadata.max_prefill_seq_len,
1453
softmax_scale=self.scale,
1454
causal=True,
1455
+ return_attn_probs=has_context,
1456
)
1457
1458
if has_context:
1408
- suffix_output, suffix_lse = output
1459
+ # ROCm flash_attn_varlen_func will return 3 objects instead of 2
1460
+ suffix_output, suffix_lse, *rest = output
1461
context_output, context_lse = self._compute_prefill_context( \
1462
q, kv_c_and_k_pe_cache, attn_metadata)
1463
vllm/attention/ops/triton_decode_attention.py
+10
-5
@@ -178,7 +178,8 @@ def _decode_att_m_fwd(
178
page_size,
179
logit_cap,
180
):
181
- BLOCK = 64
181
+ BLOCK = 64 if not is_hip_ else 8
182
+
183
NUM_KV_SPLITS = num_kv_splits
184
Lk = k_buffer.shape[-1]
185
Lv = v_buffer.shape[-1]
@@ -188,7 +189,9 @@ def _decode_att_m_fwd(
189
grid = (batch, head_num, NUM_KV_SPLITS)
190
kv_group_num = q.shape[1] // k_buffer.shape[-2]
191
191
- num_warps = 4 if kv_group_num == 1 else 2
192
+ num_warps = 4
193
+ if kv_group_num != 1:
194
+ num_warps = 1 if is_hip_ else 2
195
196
BLOCK_DMODEL = triton.next_power_of_2(Lk)
197
BLOCK_DV = triton.next_power_of_2(Lv)
@@ -418,14 +421,16 @@ def _decode_grouped_att_m_fwd(
421
)
422
423
extra_kargs = {}
424
+ num_stages = 2
425
if is_hip_:
422
- # https://rocm.docs.amd.com/en/docs-6.2.0/how-to/llm-fine-tuning-optimization/optimizing-triton-kernel.html
426
+ # https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization/workload.html#mi300x-triton-kernel-performance-optimization
427
# https://github.com/triton-lang/triton/blob/main/third_party/amd/backend/compiler.py
428
extra_kargs = {
425
- "waves_per_eu": 4,
429
+ "waves_per_eu": 1,
430
"matrix_instr_nonkdim": 16,
431
"kpack": 2
432
}
433
+ num_stages = 1
434
435
_fwd_grouped_kernel_stage1[grid](
436
q,
@@ -456,7 +461,7 @@ def _decode_grouped_att_m_fwd(
461
PAGE_SIZE=page_size,
462
logit_cap=logit_cap,
463
num_warps=4,
459
- num_stages=2,
464
+ num_stages=num_stages,
465
Lk=Lk,
466
Lv=Lv,
467
**extra_kargs,
vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128,128].json
new
+128
@@ -0,0 +1,128 @@
1
+{
2
+ "1": {
3
+ "BLOCK_SIZE_M": 16,
4
+ "BLOCK_SIZE_N": 32,
5
+ "BLOCK_SIZE_K": 256,
6
+ "GROUP_SIZE_M": 1,
7
+ "num_warps": 4,
8
+ "num_stages": 2,
9
+ "waves_per_eu": 0
10
+ },
11
+ "2": {
12
+ "BLOCK_SIZE_M": 32,
13
+ "BLOCK_SIZE_N": 16,
14
+ "BLOCK_SIZE_K": 256,
15
+ "GROUP_SIZE_M": 1,
16
+ "num_warps": 2,
17
+ "num_stages": 2,
18
+ "waves_per_eu": 0
19
+ },
20
+ "4": {
21
+ "BLOCK_SIZE_M": 16,
22
+ "BLOCK_SIZE_N": 64,
23
+ "BLOCK_SIZE_K": 256,
24
+ "GROUP_SIZE_M": 1,
25
+ "num_warps": 4,
26
+ "num_stages": 2,
27
+ "waves_per_eu": 0
28
+ },
29
+ "8": {
30
+ "BLOCK_SIZE_M": 16,
31
+ "BLOCK_SIZE_N": 128,
32
+ "BLOCK_SIZE_K": 128,
33
+ "GROUP_SIZE_M": 1,
34
+ "num_warps": 4,
35
+ "num_stages": 2,
36
+ "waves_per_eu": 0
37
+ },
38
+ "16": {
39
+ "BLOCK_SIZE_M": 16,
40
+ "BLOCK_SIZE_N": 64,
41
+ "BLOCK_SIZE_K": 128,
42
+ "GROUP_SIZE_M": 4,
43
+ "num_warps": 2,
44
+ "num_stages": 2,
45
+ "waves_per_eu": 0
46
+ },
47
+ "24": {
48
+ "BLOCK_SIZE_M": 16,
49
+ "BLOCK_SIZE_N": 128,
50
+ "BLOCK_SIZE_K": 128,
51
+ "GROUP_SIZE_M": 1,
52
+ "num_warps": 4,
53
+ "num_stages": 2,
54
+ "waves_per_eu": 0
55
+ },
56
+ "32": {
57
+ "BLOCK_SIZE_M": 16,
58
+ "BLOCK_SIZE_N": 64,
59
+ "BLOCK_SIZE_K": 128,
60
+ "GROUP_SIZE_M": 4,
61
+ "num_warps": 2,
62
+ "num_stages": 2,
63
+ "waves_per_eu": 0
64
+ },
65
+ "48": {
66
+ "BLOCK_SIZE_M": 16,
67
+ "BLOCK_SIZE_N": 64,
68
+ "BLOCK_SIZE_K": 128,
69
+ "GROUP_SIZE_M": 4,
70
+ "num_warps": 2,
71
+ "num_stages": 2,
72
+ "waves_per_eu": 0
73
+ },
74
+ "64": {
75
+ "BLOCK_SIZE_M": 16,
76
+ "BLOCK_SIZE_N": 64,
77
+ "BLOCK_SIZE_K": 128,
78
+ "GROUP_SIZE_M": 1,
79
+ "num_warps": 2,
80
+ "num_stages": 2,
81
+ "waves_per_eu": 0
82
+ },
83
+ "96": {
84
+ "BLOCK_SIZE_M": 16,
85
+ "BLOCK_SIZE_N": 64,
86
+ "BLOCK_SIZE_K": 128,
87
+ "GROUP_SIZE_M": 4,
88
+ "num_warps": 4,
89
+ "num_stages": 2,
90
+ "waves_per_eu": 0
91
+ },
92
+ "128": {
93
+ "BLOCK_SIZE_M": 16,
94
+ "BLOCK_SIZE_N": 64,
95
+ "BLOCK_SIZE_K": 256,
96
+ "GROUP_SIZE_M": 1,
97
+ "num_warps": 2,
98
+ "num_stages": 2,
99
+ "waves_per_eu": 0
100
+ },
101
+ "256": {
102
+ "BLOCK_SIZE_M": 16,
103
+ "BLOCK_SIZE_N": 64,
104
+ "BLOCK_SIZE_K": 128,
105
+ "GROUP_SIZE_M": 4,
106
+ "num_warps": 4,
107
+ "num_stages": 2,
108
+ "waves_per_eu": 0
109
+ },
110
+ "512": {
111
+ "BLOCK_SIZE_M": 32,
112
+ "BLOCK_SIZE_N": 256,
113
+ "BLOCK_SIZE_K": 128,
114
+ "GROUP_SIZE_M": 8,
115
+ "num_warps": 8,
116
+ "num_stages": 2,
117
+ "waves_per_eu": 0
118
+ },
119
+ "1024": {
120
+ "BLOCK_SIZE_M": 64,
121
+ "BLOCK_SIZE_N": 256,
122
+ "BLOCK_SIZE_K": 128,
123
+ "GROUP_SIZE_M": 8,
124
+ "num_warps": 8,
125
+ "num_stages": 2,
126
+ "waves_per_eu": 0
127
+ }
128
+}