From 516560be29a80f91209f50d496a82d1510ba917e Mon Sep 17 00:00:00 2001
From: Ason93 <18817617225@163.com>
Date: Tue, 2 Dec 2025 18:06:20 +0800
Subject: [PATCH] [KUNLUNXIN] Fix Softmax accuracy

---
 src/flag_gems/runtime/backend/_kunlunxin/ops/softmax.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/flag_gems/runtime/backend/_kunlunxin/ops/softmax.py b/src/flag_gems/runtime/backend/_kunlunxin/ops/softmax.py
index 2faedf55d..76026b9c5 100755
--- a/src/flag_gems/runtime/backend/_kunlunxin/ops/softmax.py
+++ b/src/flag_gems/runtime/backend/_kunlunxin/ops/softmax.py
@@ -228,6 +228,7 @@ def softmax(self, dim, half_to_float=False):
                 M * K,
                 N,
                 buffer_size_limit=2048,
+                is_use_mask_zero=True,
             )
 
             # 将输出恢复到原始布局
@@ -247,6 +248,7 @@ def softmax(self, dim, half_to_float=False):
                 N,
                 buffer_size_limit=2048,
                 isCloseVectorization=True,
+                is_use_mask_zero=True,
             )
     return out