[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>

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