Fix dt softplus in update function (h/t Junxiong Wang)

Author: tridao

mamba_ssm/ops/triton/selective_state_update.py

@@ -75,7 +75,7 @@ def _selective_scan_update_kernel(
     if HAS_DT_BIAS:
         dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
     if DT_SOFTPLUS:
-        dt = tl.log(1.0 + tl.exp(dt))
+        dt = tl.where(dt <= 20.0, tl.math.log1p(tl.exp(dt)), dt)
     A = tl.load(A_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0).to(tl.float32)
     dA = tl.exp(A * dt[:, None])
     B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)