[Perf][ROCm] Fix GDN KKT warmup regression on RDNA by avoiding fp32 tl.dot (#45000)
Signed-off-by: Saeid Rostami <srostami@amd.com> Signed-off-by: nemanjaudovic <nudovic@amd.com>
nemanjaudovic committed
Jul 14, 2026 at 05:48 UTC
dcf4072da9aeba343e4189b52a6e8d4228758d2d
1 file changed
+19
-1
vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py
+19
-1
@@ -10,12 +10,23 @@
10
11
import torch
12
13
+from vllm.platforms import current_platform
14
from vllm.triton_utils import tl, triton
15
16
from .index import prepare_chunk_indices
17
from .op import exp
18
from .utils import FLA_CHUNK_SIZE
19
20
+# On RDNA (gfx11xx/gfx12xx) WMMA only
21
+# accepts 16-bit/int inputs, so a widened (e.g. fp32) tl.dot is lowered to a
22
+# software matmul (~190x amdgcn-stage blowup). There we cast both operands down
23
+# to k's native storage dtype (bf16/fp16) so fast WMMA is used instead.
24
+_CAST_DOT_TO_K_DTYPE = False
25
+if current_platform.is_rocm():
26
+ from vllm.platforms.rocm import on_gfx1x
27
+
28
+ _CAST_DOT_TO_K_DTYPE = on_gfx1x()
29
+
30
31
@triton.heuristics(
32
{
@@ -48,6 +59,7 @@ def chunk_scaled_dot_kkt_fwd_kernel(
59
BK: tl.constexpr,
60
IS_VARLEN: tl.constexpr,
61
USE_G: tl.constexpr,
62
+ CAST_DOT_TO_K_DTYPE: tl.constexpr,
63
):
64
i_t, i_bh = tl.program_id(0), tl.program_id(1)
65
i_b, i_h = i_bh // H, i_bh % H
@@ -83,7 +95,12 @@ def chunk_scaled_dot_kkt_fwd_kernel(
95
)
96
b_k = tl.load(p_k, boundary_check=(0, 1))
97
b_kb = b_k * b_beta[:, None]
86
- b_A += tl.dot(b_kb, tl.trans(b_k).to(b_kb.dtype))
98
+ if CAST_DOT_TO_K_DTYPE:
99
+ # RDNA: force operands to k's native dtype so WMMA is used.
100
+ b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k))
101
+ else:
102
+ # Keep the promoted precision of the beta-scaled operand (WGMMA/MFMA).
103
+ b_A += tl.dot(b_kb, tl.trans(b_k).to(b_kb.dtype))
104
105
if USE_G:
106
p_g = tl.make_block_ptr(g + bos * H + i_h, (T,), (H,), (i_t * BT,), (BT,), (0,))
@@ -154,5 +171,6 @@ def chunk_scaled_dot_kkt_fwd(
171
Hg=Hg,
172
K=K,
173
BT=BT,
174
+ CAST_DOT_TO_K_DTYPE=_CAST_DOT_TO_K_DTYPE,
175
)
176
return A