From adb6d96516cb0cd6d1dec985cb5e5cdc20ed418b Mon Sep 17 00:00:00 2001 From: Kermit <497950719@qq.com> Date: Sat, 9 May 2026 21:08:46 +0800 Subject: [PATCH] [Bugfix] Fix GDN KKT precision loss on Hopper GPUs by aligning tl.dot operand layout with WGMMA (#42076) Signed-off-by: kermit --- vllm/model_executor/layers/fla/ops/chunk_scaled_dot_kkt.py | 2 +- 1 file changed, 1 insertion(+), 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 3f7628487d6..a715504da6b 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 @@ -83,7 +83,7 @@ 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.to(b_k.dtype), tl.trans(b_k)) + 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,))