chore: clean up masking

chunkable_ssd_combined.py

@@ -189,10 +189,13 @@ class ChunkableMambaSplitConv1dScanCombinedFn(torch.autograd.Function):
         )
         assert dt_bias.shape == (nheads,)
         assert A.shape == (nheads,)
+        seq_idx = seq_idx.contiguous() if seq_idx is not None else None
+        mask = (seq_idx.unsqueeze(-1) >= 0).to(zxbcdt.dtype) if seq_idx is not None else None
+        if mask is not None:
+            zxbcdt = zxbcdt * mask
         zx0, z, xBC, dt = torch.split(
             zxbcdt, [2 * d_nonssm, dim, dim + ngroups * dstate * 2, nheads], dim=-1
         )
-        seq_idx = seq_idx.contiguous() if seq_idx is not None else None
         final_conv_states = (
             torch.empty(
                 (batch, conv1d_weight.shape[1] - 1, dim + ngroups * dstate * 2),
@@ -206,8 +209,6 @@ class ChunkableMambaSplitConv1dScanCombinedFn(torch.autograd.Function):
         # Additionally, there is a bug in causal_conv1d_fwd_function when seq_idx is used causing illegal memory access:
         # - Issue: https://github.com/Dao-AILab/causal-conv1d/issues/67
         # - PR: https://github.com/Dao-AILab/causal-conv1d/pull/101
-        if seq_idx is not None and initial_conv_states is not None:
-            xBC = xBC * (seq_idx.unsqueeze(-1) >= 0).to(xBC.dtype)
         xBC_conv = rearrange(
             causal_conv1d_fwd_function(
                 rearrange_and_update_stride(xBC, "b s d -> b d s"),
@@ -220,8 +221,8 @@ class ChunkableMambaSplitConv1dScanCombinedFn(torch.autograd.Function):
             ),
             "b d s -> b s d",
         )
-        if seq_idx is not None and initial_conv_states is not None:
-            xBC_conv = xBC_conv * (seq_idx.unsqueeze(-1) >= 0).to(xBC_conv.dtype)
+        if mask is not None:
+            xBC_conv = xBC_conv * mask
         x, B, C = torch.split(
             xBC_conv, [dim, ngroups * dstate, ngroups * dstate], dim=-1
         )
@@ -229,12 +230,6 @@ class ChunkableMambaSplitConv1dScanCombinedFn(torch.autograd.Function):
         B = rearrange(B, "b l (g n) -> b l g n", g=ngroups)
         C = rearrange(C, "b l (g n) -> b l g n", g=ngroups)
         z = rearrange(z, "b l (h p) -> b l h p", h=nheads) if z is not None else None
-        # Moved bias addition and softplus calculation out of kernel to mask padding dt values.
-        if dt_bias is not None:
-            dt = dt + dt_bias[None, None, :]
-        dt = torch.where(dt <= 20.0, torch.log(dt.exp() + 1), dt)
-        if seq_idx is not None:
-            dt = dt * (seq_idx.unsqueeze(-1) >= 0).to(dt.dtype)
         if rmsnorm_weight is None:
             out, out_x, dt_out, dA_cumsum, states, final_ssm_states = (
                 _mamba_chunk_scan_combined_fwd(
@@ -246,10 +241,10 @@ class ChunkableMambaSplitConv1dScanCombinedFn(torch.autograd.Function):
                     chunk_size=chunk_size,
                     D=D,
                     z=z,
-                    dt_bias=None,
+                    dt_bias=dt_bias,
                     initial_states=initial_ssm_states,
                     seq_idx=seq_idx,
-                    dt_softplus=False,
+                    dt_softplus=True,
                     dt_limit=dt_limit,
                 )
             )
@@ -268,10 +263,10 @@ class ChunkableMambaSplitConv1dScanCombinedFn(torch.autograd.Function):
                     chunk_size=chunk_size,
                     D=D,
                     z=None,
-                    dt_bias=None,
+                    dt_bias=dt_bias,
                     initial_states=initial_ssm_states,
                     seq_idx=seq_idx,
-                    dt_softplus=False,
+                    dt_softplus=True,
                     dt_limit=dt_limit,
                 )
             )
@@ -317,8 +312,8 @@ class ChunkableMambaSplitConv1dScanCombinedFn(torch.autograd.Function):
             out = F.linear(out, outproj_weight, outproj_bias)
         else:
             assert outproj_bias is None
-        if out is not None and seq_idx is not None:
-            out = out * (seq_idx.unsqueeze(-1) >= 0).to(out.dtype)
+        if out is not None and mask is not None:
+            out = out * mask
         ctx.save_for_backward(
             zxbcdt,
             conv1d_weight,

modeling_chunkable_mamba2.py

@@ -41,7 +41,6 @@ class ChunkableMamba2Mixer(Mamba2Mixer):
             )
 
         # 1. Gated MLP's linear projection
-        hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask[:, -hidden_states.size(1):])
         projected_states = self.in_proj(hidden_states)
 
         A = -torch.exp(self.A_log.float())  # (num_heads) or (intermediate_size, state_size)

safe_layernorm_gated.py

@@ -1,16 +0,0 @@
-# Copyright (c) 2024, Tri Dao.
-# Based on the Triton LayerNorm tutorial: https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html
-# For the backward pass, we keep weight_grad and bias_grad in registers and accumulate.
-# This backward pass is faster for dimensions up to 8k, but after that it's much slower due to register spilling.
-# The models we train have hidden dim up to 8k anyway (e.g. Llama 70B), so this is fine.
-
-import math
-
-import torch
-import torch.nn.functional as F
-
-import triton
-import triton.language as tl
-
-from einops import rearrange
-