Skip to content

Commit 75d0dcb

Browse files
huanghua1994huanghua1994
committed
Bugfix: support qk_head_dim != v_head_dim in FMHA
The original implementation does not support qk_head_dim != v_head_dim, which is needed in Multi-head Latent Attention. Also fix some test code logic. Signed-off-by: Hua Huang <huah@nvidia.com>
1 parent 3912dda commit 75d0dcb

5 files changed

Lines changed: 58 additions & 45 deletions

File tree

‎samples/AttentionFMHA.py‎

Lines changed: 31 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,8 @@
3131
def fmha_kernel(Q, K, V, Out,
3232
qk_scale: float,
3333
input_pos: int,
34-
TILE_D: ConstInt, # TILE_D = hidden_size
34+
Dqk: ConstInt, # Head dimension of Q and K
35+
Dv: ConstInt, # Head dimension of V
3536
H: ConstInt,
3637
TILE_M: ConstInt,
3738
TILE_N: ConstInt,
@@ -64,12 +65,12 @@ def fmha_kernel(Q, K, V, Out,
6465
# Initialize online softmax accumulators in float32 for stability
6566
m_i = ct.full((TILE_M, 1), -np.inf, dtype=np.float32)
6667
l_i = ct.full((TILE_M, 1), 0.0, dtype=np.float32)
67-
acc = ct.full((TILE_M, TILE_D), 0.0, dtype=np.float32)
68+
acc = ct.full((TILE_M, Dv), 0.0, dtype=np.float32)
6869

6970
# Load query tile for this batch, head, and M-chunk
7071
q = ct.load(
71-
Q, index=(batch_idx, head_idx, bid_x, 0), shape=(1, 1, TILE_M, TILE_D)
72-
).reshape((TILE_M, TILE_D)) # [TILE_M, TILE_D]
72+
Q, index=(batch_idx, head_idx, bid_x, 0), shape=(1, 1, TILE_M, Dqk)
73+
).reshape((TILE_M, Dqk)) # [TILE_M, Dqk]
7374

7475
# loop over k, v and update accumulator
7576
m_end = input_pos + (bid_x + 1) * TILE_M
@@ -88,11 +89,11 @@ def fmha_kernel(Q, K, V, Out,
8889
for j in range(0, Tc):
8990
# --- Compute QK product ---
9091
k = ct.load(
91-
K, index=(batch_idx, off_kv_h, 0, j), shape=(1, 1, TILE_D, TILE_N),
92+
K, index=(batch_idx, off_kv_h, 0, j), shape=(1, 1, Dqk, TILE_N),
9293
order=(0, 1, 3, 2),
9394
latency=2,
9495
)
95-
k = k.reshape((TILE_D, TILE_N)) # [TILE_D, TILE_N]
96+
k = k.reshape((Dqk, TILE_N)) # [Dqk, TILE_N]
9697
qk = ct.full((TILE_M, TILE_N), 0., dtype=np.float32)
9798
qk = ct.mma(q, k, qk) # [TILE_M, TILE_N]
9899

@@ -125,16 +126,16 @@ def fmha_kernel(Q, K, V, Out,
125126

126127
# --- Compute PV product ---
127128
v = ct.load(
128-
V, index=(batch_idx, off_kv_h, j, 0), shape=(1, 1, TILE_N, TILE_D),
129+
V, index=(batch_idx, off_kv_h, j, 0), shape=(1, 1, TILE_N, Dv),
129130
latency=4,
130-
).reshape((TILE_N, TILE_D)) # [TILE_N, TILE_D]
131+
).reshape((TILE_N, Dv)) # [TILE_N, Dv]
131132
p = p.astype(Q.dtype)
132133
acc = ct.mma(p, v, acc) # [TILE_M, TILE_N]
133134
m_i = m_ij # [TILE_M, 1]
134135

135136
# --- Final Normalization and Store ---
136137
acc = ct.truediv(acc, l_i, flush_to_zero=True, rounding_mode=RMd.APPROX)
137-
acc = acc.reshape((1, 1, TILE_M, TILE_D)).astype(Out.dtype)
138+
acc = acc.reshape((1, 1, TILE_M, Dv)).astype(Out.dtype)
138139
ct.store(Out, index=(batch_idx, head_idx, bid_x, 0), tile=acc)
139140

140141

@@ -202,6 +203,7 @@ def cutile_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
202203
qk_scale,
203204
input_pos,
204205
D_k,
206+
D_v,
205207
Heads,
206208
tile_m,
207209
tile_n,
@@ -273,12 +275,18 @@ def cutile_autotune_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
273275

274276

275277
def torch_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
276-
is_causal: bool, enable_gqa: bool) -> torch.Tensor:
277-
backend = SDPBackend.CUDNN_ATTENTION \
278-
if (Q.shape[2] == K.shape[2]) \
279-
else SDPBackend.FLASH_ATTENTION
280-
with sdpa_kernel(backend):
281-
ret = scaled_dot_product_attention(Q, K, V,
278+
is_causal: bool, enable_gqa: bool,
279+
use_backend_selection_rule: bool = False) -> torch.Tensor:
280+
if use_backend_selection_rule:
281+
backend = SDPBackend.CUDNN_ATTENTION \
282+
if (Q.shape[2] == K.shape[2]) \
283+
else SDPBackend.FLASH_ATTENTION
284+
with sdpa_kernel(backend):
285+
ret = scaled_dot_product_attention(Q, K, V,
286+
is_causal=is_causal,
287+
enable_gqa=enable_gqa)
288+
else:
289+
ret = scaled_dot_product_attention(Q, K, V,
282290
is_causal=is_causal,
283291
enable_gqa=enable_gqa)
284292
return ret
@@ -296,13 +304,14 @@ def torch_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
296304

297305
# --- User Configuration ---
298306
BATCH_SIZE = 2
299-
NUM_HEADS = 8
307+
NUM_HEADS = 32
300308
SEQ_LEN_Q = 128
301-
SEQ_LEN_KV = 128
302-
D_K = 64
309+
SEQ_LEN_KV = 256
310+
D_K = 128
303311
D_V = 64
304312

305-
QUERY_GROUP_SIZE = 1
313+
QUERY_GROUP_SIZE = 8
314+
enable_gqa = QUERY_GROUP_SIZE > 1
306315

307316
DTYPE = torch.float16
308317

@@ -336,7 +345,7 @@ def torch_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
336345
dtype:{output_fmha_cutile_non_causal.dtype}""")
337346
if args.correctness_check:
338347
ref_fmha = torch_fmha(Q_input, K_input, V_input,
339-
is_causal=False, enable_gqa=False)
348+
is_causal=False, enable_gqa=enable_gqa)
340349
torch.testing.assert_close(output_fmha_cutile_non_causal, ref_fmha, atol=1e-3, rtol=1e-3)
341350
print("Correctness check passed")
342351
else:
@@ -354,7 +363,7 @@ def torch_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
354363
dtype: {output_fmha_cutile_causal.dtype}""")
355364
if args.correctness_check:
356365
ref_fmha = torch_fmha(Q_input, K_input, V_input,
357-
is_causal=True, enable_gqa=False)
366+
is_causal=True, enable_gqa=enable_gqa)
358367
torch.testing.assert_close(output_fmha_cutile_causal, ref_fmha, atol=1e-3, rtol=1e-3)
359368
print("Correctness check passed")
360369
else:
@@ -394,7 +403,7 @@ def torch_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
394403
dtype: {output_fmha_cutile_autotune_causal.dtype}""")
395404
print(f"Tuned config: {tuned_config}")
396405
if args.correctness_check:
397-
ref_fmha = torch_fmha(Q_input, K_input, V_input, is_causal=True, enable_gqa=False)
406+
ref_fmha = torch_fmha(Q_input, K_input, V_input, is_causal=True, enable_gqa=enable_gqa)
398407
torch.testing.assert_close(
399408
output_fmha_cutile_autotune_causal, ref_fmha, atol=1e-2, rtol=5e-2
400409
)

‎samples/templates/AttentionFMHA.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ def cutile_fmha(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor,
8383
qk_scale,
8484
input_pos,
8585
D_k,
86+
D_v,
8687
Heads,
8788
tile_m,
8889
tile_n,

‎test/bench_attention.py‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -74,10 +74,11 @@ def bench_fmha(qkv_shape, dtype, backend, benchmark):
7474
rounds=rounds, warmup_rounds=warmup_rounds, iterations=iterations,
7575
)
7676

77-
B, H, L, D = q.shape
78-
# first gemm mma(q, k): 2 * B * H * L * L * D
79-
# second gemm mma(p, v): 2 * B * H * L * L * D
80-
flop_count = 4 * B * H * L * L * D
77+
B, H, L, Dqk = q.shape
78+
_, _, _, Dv = v.shape
79+
# first gemm mma(q, k): 2 * B * H * L * L * Dqk
80+
# second gemm mma(p, v): 2 * B * H * L * L * Dv
81+
flop_count = 2 * B * H * L * L * (Dqk + Dv)
8182

8283
if is_causal:
8384
flop_count /= 2
@@ -88,9 +89,10 @@ def bench_fmha(qkv_shape, dtype, backend, benchmark):
8889

8990

9091
def cutile_fmha(q, k, v, o, is_causal, enable_gqa):
91-
b, qh, q_len, d = q.shape
92+
b, qh, q_len, dqk = q.shape
9293
_, kh, k_len, _ = k.shape
93-
qk_scale = 1 / sqrt(d)
94+
_, _, _, dv = v.shape
95+
qk_scale = 1 / sqrt(dqk)
9496
TILE_M, TILE_N = (256, 128) if is_causal else (64, 128)
9597
query_group_size = qh // kh
9698
grid = (ceil(q_len / TILE_M), b * qh, 1)
@@ -100,7 +102,7 @@ def cutile_fmha(q, k, v, o, is_causal, enable_gqa):
100102
(q, k, v, o,
101103
qk_scale,
102104
input_pos,
103-
d, qh,
105+
dqk, dv, qh,
104106
TILE_M, TILE_N,
105107
query_group_size, is_causal, EVEN_K))
106108

‎test/kernels/attention.py‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,8 @@
2121
def fmha_kernel(Q, K, V, Out,
2222
qk_scale: float,
2323
input_pos: int,
24-
TILE_D: ConstInt, # TILE_D = hidden_size
24+
Dqk: ConstInt, # Head dimension of Q and K
25+
Dv: ConstInt, # Head dimension of V
2526
H: ConstInt,
2627
TILE_M: ConstInt,
2728
TILE_N: ConstInt,
@@ -54,12 +55,12 @@ def fmha_kernel(Q, K, V, Out,
5455
# Initialize online softmax accumulators in float32 for stability
5556
m_i = ct.full((TILE_M, 1), -np.inf, dtype=np.float32)
5657
l_i = ct.full((TILE_M, 1), 0.0, dtype=np.float32)
57-
acc = ct.full((TILE_M, TILE_D), 0.0, dtype=np.float32)
58+
acc = ct.full((TILE_M, Dv), 0.0, dtype=np.float32)
5859

5960
# Load query tile for this batch, head, and M-chunk
6061
q = ct.load(
61-
Q, index=(batch_idx, head_idx, bid_x, 0), shape=(1, 1, TILE_M, TILE_D)
62-
).reshape((TILE_M, TILE_D)) # [TILE_M, TILE_D]
62+
Q, index=(batch_idx, head_idx, bid_x, 0), shape=(1, 1, TILE_M, Dqk)
63+
).reshape((TILE_M, Dqk)) # [TILE_M, Dqk]
6364

6465
# loop over k, v and update accumulator
6566
m_end = input_pos + (bid_x + 1) * TILE_M
@@ -78,11 +79,11 @@ def fmha_kernel(Q, K, V, Out,
7879
for j in range(0, Tc):
7980
# --- Compute QK product ---
8081
k = ct.load(
81-
K, index=(batch_idx, off_kv_h, 0, j), shape=(1, 1, TILE_D, TILE_N),
82+
K, index=(batch_idx, off_kv_h, 0, j), shape=(1, 1, Dqk, TILE_N),
8283
order=(0, 1, 3, 2),
8384
latency=2,
8485
)
85-
k = k.reshape((TILE_D, TILE_N)) # [TILE_D, TILE_N]
86+
k = k.reshape((Dqk, TILE_N)) # [Dqk, TILE_N]
8687
qk = ct.full((TILE_M, TILE_N), 0., dtype=np.float32)
8788
qk = ct.mma(q, k, qk) # [TILE_M, TILE_N]
8889

@@ -115,14 +116,14 @@ def fmha_kernel(Q, K, V, Out,
115116

116117
# --- Compute PV product ---
117118
v = ct.load(
118-
V, index=(batch_idx, off_kv_h, j, 0), shape=(1, 1, TILE_N, TILE_D),
119+
V, index=(batch_idx, off_kv_h, j, 0), shape=(1, 1, TILE_N, Dv),
119120
latency=4,
120-
).reshape((TILE_N, TILE_D)) # [TILE_N, TILE_D]
121+
).reshape((TILE_N, Dv)) # [TILE_N, Dv]
121122
p = p.astype(Q.dtype)
122123
acc = ct.mma(p, v, acc) # [TILE_M, TILE_N]
123124
m_i = m_ij # [TILE_M, 1]
124125

125126
# --- Final Normalization and Store ---
126127
acc = ct.truediv(acc, l_i, flush_to_zero=True, rounding_mode=RMd.APPROX)
127-
acc = acc.reshape((1, 1, TILE_M, TILE_D)).astype(Out.dtype)
128+
acc = acc.reshape((1, 1, TILE_M, Dv)).astype(Out.dtype)
128129
ct.store(Out, index=(batch_idx, head_idx, bid_x, 0), tile=acc)

‎test/test_attention.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -15,21 +15,21 @@
1515
@pytest.mark.parametrize("k_heads", [8])
1616
@pytest.mark.parametrize("q_len", [1, 15, 32])
1717
@pytest.mark.parametrize("k_len", [32, 63])
18-
@pytest.mark.parametrize("hidden_size", [32])
18+
@pytest.mark.parametrize("head_dim", [32])
1919
@pytest.mark.parametrize("tile_size", [(8, 16)])
2020
@pytest.mark.parametrize("is_causal", [True, False])
2121
@pytest.mark.parametrize("use_input_pos", [True, False])
2222
def test_flash_attention(batch_size, q_heads, k_heads,
2323
q_len, k_len,
24-
hidden_size, tile_size, is_causal,
24+
head_dim, tile_size, is_causal,
2525
use_input_pos,
2626
float_dtype):
2727
query_group_size = q_heads // k_heads
2828
TILE_M, TILE_N = tile_size
29-
qk_scale = 1 / math.sqrt(hidden_size)
30-
q = torch.randn((batch_size, q_heads, q_len, hidden_size), dtype=float_dtype, device='cuda')
31-
k = torch.randn((batch_size, k_heads, k_len, hidden_size), dtype=float_dtype, device='cuda')
32-
v = torch.randn((batch_size, k_heads, k_len, hidden_size), dtype=float_dtype, device='cuda')
29+
qk_scale = 1 / math.sqrt(head_dim)
30+
q = torch.randn((batch_size, q_heads, q_len, head_dim), dtype=float_dtype, device='cuda')
31+
k = torch.randn((batch_size, k_heads, k_len, head_dim), dtype=float_dtype, device='cuda')
32+
v = torch.randn((batch_size, k_heads, k_len, head_dim), dtype=float_dtype, device='cuda')
3333
o = torch.zeros_like(q)
3434
grid = (math.ceil(q_len / TILE_M), batch_size * q_heads, 1)
3535
if use_input_pos:
@@ -43,7 +43,7 @@ def test_flash_attention(batch_size, q_heads, k_heads,
4343
(q, k, v, o,
4444
qk_scale,
4545
input_pos,
46-
hidden_size, q_heads,
46+
head_dim, head_dim, q_heads,
4747
TILE_M, TILE_N,
4848
query_group_size, is_causal, EVEN_K))
4949
if is_causal:

0 commit comments

Comments
 (0)