Improves attention bias numerical stability

Replaces $\\exp(A\\cdot\\mathrm{softplus}(\\Delta V))$ with $A\\cdot\\mathrm{softplus}(\\Delta V)$ to prevent overflow/NaNs in attention bias and stabilize training/inference.

Preserves tensor shape/dtype and adds a clarifying comment on the rationale.

Author: algo-home

examples/modeling/modeling_doge.py

@@ -217,7 +217,8 @@ def forward(
         dt_states = self.dt_proj(
             value_states.transpose(1, 2).reshape(value_states.shape[0], value_states.shape[-2], -1)
         )
-        attn_bias = torch.exp(self.A * F.softplus(dt_states)).transpose(-1, -2).unsqueeze(-2).to(hidden_states.dtype)
+        # original formula is exp(A * softplus(delta V)), but for numerical stability, it is changed to A * softplus(delta V)
+        attn_bias = self.A * F.softplus(dt_states).transpose(-1, -2).unsqueeze(-2).to(hidden_states.dtype)
 
         attention_interface: Callable = flash_dynamic_mask_attention_forward