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, 08:30 UTC
[Bugfix][ROCm] fix the power of 2 exception from triton_unified_attention.py when running llama4 models and unit test fix (#18100)
Signed-off-by: Hongxia Yang <hongxia.yang@amd.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
Hongxia Yang committed
May 28, 2025 at 19:21 UTC
269d901734326432d5ef15deaca07364149f9b48
2 files changed
+54
-56
tests/kernels/attention/test_triton_unified_attention.py
+3
-1
@@ -13,7 +13,9 @@ HEAD_SIZES = [128, 256]
13
BLOCK_SIZES = [16, 32]
14
15
DTYPES = [torch.float16, torch.bfloat16]
16
-QDTYPES = [None, torch.float8_e4m3fn]
16
+QDTYPES = [None, torch.float8_e4m3fn] if not current_platform.is_rocm() else [
17
+ None, torch.float8_e4m3fnuz
18
+]
19
# one value large enough to test overflow in index calculation.
20
# one value small enough to test the schema op check
21
NUM_BLOCKS = [32768, 2048]
vllm/attention/ops/triton_unified_attention.py
+51
-55
@@ -29,41 +29,42 @@ def apply_softcap(S, x):
29
30
@triton.jit
31
def kernel_unified_attention_2d(
32
- output_ptr, # [num_tokens, num_query_heads, head_size]
33
- query_ptr, # [num_tokens, num_query_heads, head_size]
34
- key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
35
- value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
36
- block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
37
- seq_lens_ptr, # [num_seqs]
38
- alibi_slopes_ptr, # [num_query_heads]
39
- scale, # float32
40
- k_scale, # float32
41
- v_scale, # float32
42
- softcap, # float32
43
- num_query_heads: tl.constexpr, # int
44
- num_queries_per_kv: tl.constexpr, # int
45
- block_table_stride: tl.int64, # int
46
- query_stride_0: tl.int64, # int
47
- query_stride_1: tl.int64, # int, should be equal to head_size
48
- output_stride_0: tl.int64, # int
49
- output_stride_1: tl.int64, # int, should be equal to head_size
50
- BLOCK_SIZE: tl.constexpr, # int
51
- HEAD_SIZE: tl.constexpr, # int
52
- HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
53
- USE_ALIBI_SLOPES: tl.constexpr, # bool
54
- USE_SOFTCAP: tl.constexpr, # bool
55
- SLIDING_WINDOW: tl.constexpr, # int
56
- stride_k_cache_0: tl.int64, # int
57
- stride_k_cache_1: tl.int64, # int
58
- stride_k_cache_2: tl.int64, # int
59
- stride_k_cache_3: tl.constexpr, # int
60
- stride_v_cache_0: tl.int64, # int
61
- stride_v_cache_1: tl.int64, # int
62
- stride_v_cache_2: tl.int64, # int
63
- stride_v_cache_3: tl.constexpr, # int
64
- query_start_len_ptr, # [num_seqs+1]
65
- BLOCK_Q: tl.constexpr, # int
66
- num_seqs: tl.int32,
32
+ output_ptr, # [num_tokens, num_query_heads, head_size]
33
+ query_ptr, # [num_tokens, num_query_heads, head_size]
34
+ key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
35
+ value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]
36
+ block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]
37
+ seq_lens_ptr, # [num_seqs]
38
+ alibi_slopes_ptr, # [num_query_heads]
39
+ scale, # float32
40
+ k_scale, # float32
41
+ v_scale, # float32
42
+ softcap, # float32
43
+ num_query_heads: tl.constexpr, # int
44
+ num_queries_per_kv: tl.constexpr, # int
45
+ block_table_stride: tl.int64, # int
46
+ query_stride_0: tl.int64, # int
47
+ query_stride_1: tl.int64, # int, should be equal to head_size
48
+ output_stride_0: tl.int64, # int
49
+ output_stride_1: tl.int64, # int, should be equal to head_size
50
+ BLOCK_SIZE: tl.constexpr, # int
51
+ HEAD_SIZE: tl.constexpr, # int
52
+ HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2
53
+ USE_ALIBI_SLOPES: tl.constexpr, # bool
54
+ USE_SOFTCAP: tl.constexpr, # bool
55
+ SLIDING_WINDOW: tl.constexpr, # int
56
+ stride_k_cache_0: tl.int64, # int
57
+ stride_k_cache_1: tl.int64, # int
58
+ stride_k_cache_2: tl.int64, # int
59
+ stride_k_cache_3: tl.constexpr, # int
60
+ stride_v_cache_0: tl.int64, # int
61
+ stride_v_cache_1: tl.int64, # int
62
+ stride_v_cache_2: tl.int64, # int
63
+ stride_v_cache_3: tl.constexpr, # int
64
+ query_start_len_ptr, # [num_seqs+1]
65
+ BLOCK_Q: tl.constexpr, # int
66
+ num_seqs: tl.int32,
67
+ BLOCK_M: tl.constexpr, # int
68
):
69
70
q_block_global_idx = tl.program_id(0)
@@ -94,15 +95,13 @@ def kernel_unified_attention_2d(
95
if q_block_local_idx * BLOCK_Q >= cur_batch_query_len:
96
return
97
97
- offs_m = tl.arange(0, BLOCK_Q * num_queries_per_kv)
98
+ offs_m = tl.arange(0, BLOCK_M)
99
offs_d = tl.arange(0, HEAD_SIZE_PADDED)
99
-
100
query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv
101
102
query_offset_0 = cur_batch_in_all_start_index + query_pos
103
query_offset_1 = kv_head_idx * num_queries_per_kv + \
104
offs_m % num_queries_per_kv
105
-
105
query_offset = (query_offset_0[:, None] * query_stride_0 +
106
query_offset_1[:, None] * query_stride_1 + offs_d[None, :])
107
@@ -110,7 +109,7 @@ def kernel_unified_attention_2d(
109
query_mask_0 = tl.where(query_pos < cur_batch_query_len, 1, 0).to(tl.int1)
110
query_mask_1 = tl.where(query_offset_1 < num_query_heads, 1, 0).to(tl.int1)
111
113
- # Q : (BLOCK_Q * num_queries_per_kv, HEAD_SIZE,)
112
+ # Q : (BLOCK_M, HEAD_SIZE_PADDED)
113
Q = tl.load(
114
query_ptr + query_offset,
115
mask=dim_mask[None, :] & query_mask_0[:, None] & query_mask_1[:, None],
@@ -119,12 +118,9 @@ def kernel_unified_attention_2d(
118
119
block_table_offset = seq_idx * block_table_stride
120
122
- M = tl.full([BLOCK_Q * num_queries_per_kv],
123
- float("-inf"),
124
- dtype=tl.float32)
125
- L = tl.full([BLOCK_Q * num_queries_per_kv], 1.0, dtype=tl.float32)
126
- acc = tl.zeros([BLOCK_Q * num_queries_per_kv, HEAD_SIZE_PADDED],
127
- dtype=tl.float32)
121
+ M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
122
+ L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
123
+ acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)
124
125
# sequence len for this particular sequence
126
seq_len = tl.load(seq_lens_ptr + seq_idx)
@@ -183,13 +179,12 @@ def kernel_unified_attention_2d(
179
else:
180
V = V_load
181
186
- seq_offset = j * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
182
+ seq_offset = j * BLOCK_SIZE + offs_n
183
184
seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1
185
190
- # S : (BLOCK_Q * num_queries_per_kv, BLOCK_SIZE,)
191
- S = tl.zeros(shape=(BLOCK_Q * num_queries_per_kv, BLOCK_SIZE),
192
- dtype=tl.float32)
186
+ # S : (BLOCK_M, BLOCK_SIZE)
187
+ S = tl.zeros(shape=(BLOCK_M, BLOCK_SIZE), dtype=tl.float32)
188
189
S += scale * tl.dot(Q, K)
190
@@ -207,29 +202,29 @@ def kernel_unified_attention_2d(
202
S += alibi_slope[:, None] * (seq_offset - context_len)
203
204
# compute running maximum
210
- # m_j : (BLOCK_Q * num_queries_per_kv,)
205
+ # m_j : (BLOCK_M,)
206
m_j = tl.maximum(M, tl.max(S, axis=1))
207
# For sliding window there's a chance the max is -inf due to masking of
208
# the entire row. In this case we need to set m_j 0 to avoid NaN
209
m_j = tl.where(m_j > float("-inf"), m_j, 0.0)
210
216
- # P : (BLOCK_Q * num_queries_per_kv, BLOCK_SIZE,)
211
+ # P : (BLOCK_M, BLOCK_SIZE)
212
P = tl.exp(S - m_j[:, None])
213
219
- # l_j : (BLOCK_Q * num_queries_per_kv,)
214
+ # l_j : (BLOCK_M,)
215
l_j = tl.sum(P, axis=1)
216
222
- # alpha : (BLOCK_Q * num_queries_per_kv, )
217
+ # alpha : (BLOCK_M, )
218
alpha = tl.exp(M - m_j)
219
225
- # acc : (BLOCK_Q * num_queries_per_kv, BLOCK_SIZE,)
220
+ # acc : (BLOCK_M, HEAD_SIZE_PADDED)
221
acc = acc * alpha[:, None]
222
223
# update constants
224
L = L * alpha + l_j
225
M = m_j
226
232
- # acc : (BLOCK_Q * num_queries_per_kv, BLOCK_SIZE,)
227
+ # acc : (BLOCK_M, HEAD_SIZE_PADDED)
228
acc += tl.dot(P.to(V.dtype), V)
229
230
# epilogue
@@ -334,4 +329,5 @@ def unified_attention(
329
query_start_len_ptr=cu_seqlens_q,
330
BLOCK_Q=BLOCK_Q,
331
num_seqs=num_seqs,
332
+ BLOCK_M=BLOCK_M,
333
)