From ea4b6b10bc38954f6624359922a215f5b397a52d Mon Sep 17 00:00:00 2001
From: biondizzle <biondizzle@gmail.com>
Date: Sun, 24 May 2026 03:20:24 +0000
Subject: [PATCH] =?UTF-8?q?fix:=20LSE=20type=20mismatch=20Float32=E2=86=92?=
 =?UTF-8?q?BFloat16?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 dsv4/kernels/attention/fmha.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/dsv4/kernels/attention/fmha.py b/dsv4/kernels/attention/fmha.py
index 39405779..154ccd1f 100644
--- a/dsv4/kernels/attention/fmha.py
+++ b/dsv4/kernels/attention/fmha.py
@@ -443,7 +443,7 @@ class FmhaKernel:
             if sfw_idx == 0:
                 _ln2 = Float32(0.6931471805599453)  # ln(2)
                 lse_val = cute.math.log(row_sum, fastmath=True) + _row_max_safe * _ln2
-                mLSE[0] = lse_val
+                mLSE[0] = lse_val.to(self.q_dtype)
 
             tmem.relinquish_alloc_permit()
             tmem.free(tmem_ptr)