From dcf4072da9aeba343e4189b52a6e8d4228758d2d Mon Sep 17 00:00:00 2001 From: nemanjaudovic <152565955+nemanjaudovic@users.noreply.github.com> Date: Tue, 14 Jul 2026 05:48:54 +0200 Subject: [PATCH] [Perf][ROCm] Fix GDN KKT warmup regression on RDNA by avoiding fp32 tl.dot (#45000) Signed-off-by: Saeid Rostami Signed-off-by: nemanjaudovic --- .../layers/fla/ops/chunk_scaled_dot_kkt.py | 20 ++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py b/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py index a715504da6b..10ff970c809 100644 --- a/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py +++ b/vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py @@ -10,12 +10,23 @@ import torch +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from .index import prepare_chunk_indices from .op import exp from .utils import FLA_CHUNK_SIZE +# On RDNA (gfx11xx/gfx12xx) WMMA only +# accepts 16-bit/int inputs, so a widened (e.g. fp32) tl.dot is lowered to a +# software matmul (~190x amdgcn-stage blowup). There we cast both operands down +# to k's native storage dtype (bf16/fp16) so fast WMMA is used instead. +_CAST_DOT_TO_K_DTYPE = False +if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1x + + _CAST_DOT_TO_K_DTYPE = on_gfx1x() + @triton.heuristics( { @@ -48,6 +59,7 @@ def chunk_scaled_dot_kkt_fwd_kernel( BK: tl.constexpr, IS_VARLEN: tl.constexpr, USE_G: tl.constexpr, + CAST_DOT_TO_K_DTYPE: tl.constexpr, ): i_t, i_bh = tl.program_id(0), tl.program_id(1) i_b, i_h = i_bh // H, i_bh % H @@ -83,7 +95,12 @@ def chunk_scaled_dot_kkt_fwd_kernel( ) b_k = tl.load(p_k, boundary_check=(0, 1)) b_kb = b_k * b_beta[:, None] - b_A += tl.dot(b_kb, tl.trans(b_k).to(b_kb.dtype)) + if CAST_DOT_TO_K_DTYPE: + # RDNA: force operands to k's native dtype so WMMA is used. + b_A += tl.dot(b_kb.to(b_k.dtype), tl.trans(b_k)) + else: + # Keep the promoted precision of the beta-scaled operand (WGMMA/MFMA). + b_A += tl.dot(b_kb, tl.trans(b_k).to(b_kb.dtype)) if USE_G: 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( Hg=Hg, K=K, BT=BT, + CAST_DOT_TO_K_DTYPE=_CAST_DOT_TO_K_DTYPE, ) return A