diff --git a/mindformers/pynative/transformers/nsa.py b/mindformers/pynative/transformers/nsa.py
index 851c1d858ab2f54669f3e26005a02b98acb540cc..af6c38a65e724116c6dc86d862472bc8499aaaa8 100644
--- a/mindformers/pynative/transformers/nsa.py
+++ b/mindformers/pynative/transformers/nsa.py
@@ -318,6 +318,22 @@ def _normalize_attention_mask(attention_mask: Tensor, dtype: ms.dtype) -> Tensor
     return attn
 
 
+def safe_softmax(x: Tensor, dim: int = -1) -> Tensor:
+    """
+    Numerically stable softmax that prevents overflow/underflow.
+    
+    Uses the identity: softmax(x) = softmax(x - max(x))
+    This prevents exp() overflow while maintaining mathematical equivalence.
+    """
+    x_max = mint.max(x, dim=dim, keepdim=True)[0]
+    x_shifted = x - x_max
+    exp_x = mint.exp(x_shifted)
+    sum_exp = mint.sum(exp_x, dim=dim, keepdim=True)
+    # Avoid division by zero
+    sum_exp = mint.where(sum_exp == 0.0, mint.ones_like(sum_exp), sum_exp)
+    return exp_x / sum_exp
+
+
 class NSAAttention(nn.Cell):
     """Native Sparse Attention core attention for MLA."""
 
@@ -407,7 +423,6 @@ class NSAAttention(nn.Cell):
             raise ValueError(f"Unsupported nsa_gate_mode: {self.gate_mode}")
 
         self.attn_dropout = Dropout(drop_prob=float(self.nsa_dropout))
-        self.softmax = mint.nn.functional.softmax
 
     def construct(
         self,
@@ -513,7 +528,7 @@ class NSAAttention(nn.Cell):
 
         # Gating
         if self.gate_mode == "static":
-            g = self.softmax(self.gate, dim=-1)  # (h, 3)
+            g = safe_softmax(self.gate, dim=-1)  # (h, 3)
             w = g.reshape((1, h, 1, 3))
         else:
             q_mean = q.mean(dim=-2)  # (b, h, d)
@@ -521,12 +536,12 @@ class NSAAttention(nn.Cell):
             gate_proj = self.gate_proj(q_mean_flat)[0]
             gate_proj = mint.reshape(gate_proj, (b, h, 3))
             g = gate_proj + self.gate.reshape((1, h, 3))
-            w = mint.unsqueeze(self.softmax(g, dim=-1), -2)  # (b, h, 1, 3)
+            w = mint.unsqueeze(safe_softmax(g, dim=-1), -2)  # (b, h, 1, 3)
 
         # Apply validity mask to selective branch gate weight instead of output
         # This prevents gradient flow issues when sel_valid is 0
-        # Redistribute the masked weight to other branches for stability
         w_sel = w[..., 2, None] * sel_valid_mask
+        # Redistribute the masked weight to other branches for stability
         w_redistrib = w[..., 2, None] * (1.0 - sel_valid_mask)
         w_local = w[..., 0, None] + w_redistrib * 0.5
         w_comp = w[..., 1, None] + w_redistrib * 0.5