fixing softplus bug with _chunk_cumsum_bwd_kernel() triton kernel (#574)

Author: stephen-youn

mamba_ssm/ops/triton/selective_state_update.py

@@ -85,15 +85,15 @@ 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 = softplus(dt)
+            dt = tl.where(dt <= 20.0, softplus(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])
     else:
         dt = tl.load(dt_ptr).to(tl.float32)
         if HAS_DT_BIAS:
             dt += tl.load(dt_bias_ptr).to(tl.float32)
         if DT_SOFTPLUS:
-            dt = softplus(dt)
+            dt = tl.where(dt <= 20.0, softplus(dt), dt)
         A = tl.load(A_ptr).to(tl.float32)
         dA = tl.exp(A * dt)  # scalar, not a matrix
 

mamba_ssm/ops/triton/softplus.py

@@ -8,10 +8,8 @@
 if TRITON3:
     @triton.jit
     def softplus(dt):
-        dt = tl.where(dt <= 20.0, tl.math.log(tl.math.exp(dt) + 1), dt)
-        return dt
+        return tl.math.log(tl.math.exp(dt) + 1)
 else:
     @triton.jit
     def softplus(dt):
-        dt = tl.where(dt <= 20.0, tl.math.log1p(tl.exp(dt)), dt)
-        return dt
+        return tl.math.log1p(tl.exp(dt))

mamba_ssm/ops/triton/ssd_chunk_state.py

@@ -68,7 +68,7 @@ def _chunk_cumsum_fwd_kernel(
         dt_bias = tl.load(dt_bias_ptr + offs_h * stride_dt_bias_head, mask=offs_h < nheads, other=0.0).to(tl.float32)
         dt += dt_bias[:, None]
     if DT_SOFTPLUS:
-        dt = softplus(dt)
+        dt = tl.where(dt <= 20.0, softplus(dt), dt)
     # As of Triton 2.2.0, tl.clamp is not available yet
     # dt = tl.clamp(dt, dt_min, dt_max)
     dt = tl.minimum(tl.maximum(dt, dt_min), dt_max)
@@ -141,7 +141,7 @@ def _chunk_cumsum_bwd_kernel(
         dt += dt_bias[:, None]
     if DT_SOFTPLUS:
         dt_presoftplus = dt
-        dt = softplus(dt)
+        dt = tl.where(dt <= 20.0, softplus(dt), ddt)
     clamp_mask = (dt < dt_min) | (dt > dt_max)
     # As of Triton 2.2.0, tl.clamp is not available yet
     # dt = tl.clamp(dt, dt_min, dt_max)