diff --git a/mindformers/pynative/transformers/nsa.py b/mindformers/pynative/transformers/nsa.py
index af6c38a65e724116c6dc86d862472bc8499aaaa8..f302d7db105acedc2aa0d0c32c12364545b71b2f 100644
--- a/mindformers/pynative/transformers/nsa.py
+++ b/mindformers/pynative/transformers/nsa.py
@@ -569,7 +569,7 @@ class NSAAttention(nn.Cell):
         # Clip scores to prevent overflow/underflow in softmax
         # This is critical for training stability
         scores = mint.clamp(scores, min=-50.0, max=50.0)
-        attn = self.softmax(scores, dim=-1)
+        attn = safe_softmax(scores, dim=-1)
         # Replace NaN values with uniform distribution if they occur
         # This prevents cascading NaN issues during training
         attn = mint.where(mint.isnan(attn), mint.full_like(attn, 1.0 / attn.shape[-1]), attn)
@@ -673,7 +673,7 @@ class NSAAttention(nn.Cell):
             logits = mint.einsum("bhid,bhjd->bhij", q_det, kc_mean) * self.softmax_scale
             logits = logits + _bool_to_score_mask(valid, logits.dtype)
 
-            attn = self.softmax(logits, dim=-1)
+            attn = safe_softmax(logits, dim=-1)
             topk_k = min(self.topk_blocks, blk_total)
             topk_values, topk_indices = mint.topk(attn, topk_k, dim=-1)
             updates = (topk_values > 1e-3).astype(ms.float32)
@@ -745,14 +745,14 @@ class NSAAttention(nn.Cell):
         true_logits = true_logits + _bool_to_score_mask(valid, true_logits.dtype)
 
         # Target distribution: softmax, sum across heads, L1-normalise
-        target = self.softmax(true_logits, dim=-1)          # (b, h, n, blk)
+        target = safe_softmax(true_logits, dim=-1)          # (b, h, n, blk)
         target = target.sum(axis=1)                          # (b, n, blk)
         target = target / (target.sum(axis=-1, keepdims=True) + 1e-10)
 
         # Indexer distribution: softmax of block_scores (with same mask)
         valid_2d = valid.reshape((1, n, blk_total))
         idx_logits = block_scores + _bool_to_score_mask(valid_2d, block_scores.dtype)
-        idx_dist = self.softmax(ops.cast(idx_logits, ms.float32), dim=-1)
+        idx_dist = safe_softmax(ops.cast(idx_logits, ms.float32), dim=-1)
 
         # KL(target || idx_dist)
         kl = target * (mint.log(target + 1e-10) - mint.log(idx_dist + 1e-10))