Fix mamba cumsum padded calculations (#1022)

Cherry-pick #1009

---------

Signed-off-by: Jan Kaniecki <jkaniecki@habana.ai>
Co-authored-by: Artur Fierka <artur.fierka@intel.com>

Author: jkaniecki

vllm_gaudi/ops/hpu_mamba_mixer2.py

@@ -437,6 +437,7 @@ def conv_ssm_forward(
                 dt_limit=(0.0, float("inf")),
                 out=output.view(output.shape[0], -1, self.head_dim),
                 state_dtype=ssm_state.dtype,
+                padding_mask=padding_mask_flat,
             )[last_chunk_indices_p]
             output = output * padding_mask_flat.view(output.shape[0], 1)
 

vllm_gaudi/ops/pytorch_implementation.py

@@ -5,7 +5,13 @@
 from einops import rearrange, repeat
 
 
-def new_chunk_cumsum(dt, A, chunk_size, dt_bias=None, dt_softplus=False, dt_limit=(0.0, float("inf"))):
+def new_chunk_cumsum(dt,
+                     A,
+                     chunk_size,
+                     dt_bias=None,
+                     dt_softplus=False,
+                     dt_limit=(0.0, float("inf")),
+                     padding_mask=None):
     """
     Arguments:
         dt: Tensor - (seqlen, nheads)
@@ -14,6 +20,7 @@ def new_chunk_cumsum(dt, A, chunk_size, dt_bias=None, dt_softplus=False, dt_limi
         dt_bias: Optional Tensor - (nheads)
         dt_softplus: bool
         dt_limit: tuple - (min: float, max: float)
+        padding_mask: Optional Tensor - (seqlen, 1) or (seqlen,)
 
     Return:
         dA_cumsum: Tensor - (nheads, nchunks, chunk_size)
@@ -32,6 +39,10 @@ def new_chunk_cumsum(dt, A, chunk_size, dt_bias=None, dt_softplus=False, dt_limi
         dt = torch.where(dt <= 20.0, F.softplus(dt), dt)
 
     dt = torch.clamp(dt, dt_min, dt_max)
+
+    if padding_mask is not None:
+        dt = dt * padding_mask.view(seqlen, 1).float()
+
     dA = dt * A.view(1, nheads)
     dA = dA.transpose(0, 1).reshape(nheads, nchunks, chunk_size)
     dt = dt.transpose(0, 1).reshape(nheads, nchunks, chunk_size)

vllm_gaudi/ops/ssd_combined.py

@@ -36,6 +36,7 @@ def _mamba_chunk_scan_combined_fwd(
         dt_softplus=False,
         dt_limit=(0.0, float("inf")),
         state_dtype=None,
+        padding_mask=None,
 ):
     assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2"
     seqlen, nheads, headdim = x.shape
@@ -82,6 +83,7 @@ def _mamba_chunk_scan_combined_fwd(
         dt_bias=dt_bias,
         dt_softplus=dt_softplus,
         dt_limit=dt_limit,
+        padding_mask=padding_mask,
     )
 
     # 2. Compute the state for each intra-chunk
@@ -143,6 +145,7 @@ def hpu_mamba_chunk_scan_combined_varlen(
         dt_softplus=False,
         dt_limit=(0.0, float("inf")),
         state_dtype=None,
+        padding_mask=None,
 ):
     """
     Argument:
@@ -185,6 +188,7 @@ def hpu_mamba_chunk_scan_combined_varlen(
         dt_softplus=dt_softplus,
         dt_limit=dt_limit,
         state_dtype=state_dtype,
+        padding_mask=padding_mask,
     )
 
     return varlen_states