combine causal_conv1d_fn_cpu

Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>

Author: Akash kaothalkar

vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py

@@ -24,56 +24,59 @@ def causal_conv1d_fn_cpu(
     """CPU implementation for causal_conv1d_fwd."""
     if isinstance(activation, bool) and activation:
         activation = "silu"
+    elif isinstance(activation, bool):
+        activation = None
 
     original_x_dtype = x.dtype
     x = x.to(conv_states.dtype)
 
-    dim, cu_seqlen = x.shape
-    _, width = weight.shape
-    state_len = width - 1
-
-    out = torch.zeros_like(x)
-
-    batch = query_start_loc.size(0) - 1
-
-    for b in range(batch):
-        seq_start = query_start_loc[b].item()
-        seq_end = query_start_loc[b + 1].item()
-        seq_len = seq_end - seq_start
+    out = torch.empty_like(x)
+    state_len = weight.shape[1] - 1
+    assert activation in {None, "silu", "swish"}
 
-        if seq_len == 0:
+    seq_begin_end_idx = [
+        (int(query_start_loc[idx].item()), int(query_start_loc[idx + 1].item()))
+        for idx in range(query_start_loc.shape[0] - 1)
+    ]
+    weight = weight.unsqueeze(1)
+    
+    for seq_idx, (bos, eos) in enumerate(seq_begin_end_idx):
+        if bos == eos:
             continue
-
-        cache_idx = cache_indices[b].item() if cache_indices is not None else b
-
-        if cache_idx == pad_slot_id:
+            
+        slot = int(cache_indices[seq_idx].item()) if cache_indices is not None else seq_idx
+        
+        if slot == pad_slot_id:
             continue
 
-        x_seq = x[:, seq_start:seq_end]  # (dim, seq_len)
-
-        if has_initial_state is not None and has_initial_state[b]:
-            state = conv_states[cache_idx].clone()  # (dim, state_len)
+        seq_x = x[:, bos:eos].unsqueeze(0)
+        
+        if has_initial_state is not None and bool(has_initial_state[seq_idx].item()):
+            initial_state = conv_states[slot, :, :state_len].unsqueeze(0)
         else:
-            state = torch.zeros((dim, state_len), dtype=x.dtype, device=x.device)
-
-        for t in range(seq_len):
-            x_t = x_seq[:, t]  # (dim,)
-
-            window = torch.cat([state, x_t.unsqueeze(1)], dim=1)  # (dim, width)
-            val = (window * weight).sum(dim=1)  # (dim,)
-
-            if bias is not None:
-                val = val + bias
-            if activation in ["silu", "swish"]:
-                val = val * torch.sigmoid(val)
-
-            out[:, seq_start + t] = val
+            initial_state = torch.zeros(
+                1,
+                weight.shape[0],
+                state_len,
+                device=seq_x.device,
+                dtype=seq_x.dtype,
+            )
 
-            if state_len > 1:
-                state[:, :-1] = state[:, 1:].clone()
-            state[:, -1] = x_t
+        conv_input = torch.cat([initial_state, seq_x], dim=-1).to(weight.dtype)
+        seq_out = F.conv1d(
+            conv_input,
+            weight,
+            bias,
+            padding=0,
+            groups=weight.shape[0],
+        )
+        seq_out = seq_out[..., -seq_x.shape[-1] :].to(dtype=x.dtype)
+        
+        if activation in ("silu", "swish"):
+            seq_out = F.silu(seq_out)
 
-        conv_states[cache_idx].copy_(state)
+        out[:, bos:eos] = seq_out.squeeze(0)
+        conv_states[slot, :, :state_len].copy_(conv_input[..., -state_len:].squeeze(0))
 
     return out.to(original_x_dtype)
 
@@ -179,56 +182,7 @@ def causal_conv1d_update_cpu(
     return out.to(original_x_dtype)
 
 
-def causal_conv1d_torch(
-    x: torch.Tensor,
-    weight: torch.Tensor,
-    bias: torch.Tensor | None,
-    conv_states: torch.Tensor,
-    query_start_loc: torch.Tensor,
-    cache_indices: torch.Tensor,
-    has_initial_state: torch.Tensor,
-    activation: str | None = "silu",
-) -> torch.Tensor:
-    out = torch.empty_like(x)
-    state_len = weight.shape[1] - 1
-    assert activation in {None, "silu", "swish"}
-
-    seq_begin_end_idx = [
-        (int(query_start_loc[idx].item()), int(query_start_loc[idx + 1].item()))
-        for idx in range(query_start_loc.shape[0] - 1)
-    ]
-    weight = weight.unsqueeze(1)
-    for seq_idx, (bos, eos) in enumerate(seq_begin_end_idx):
-        slot = int(cache_indices[seq_idx].item())
-
-        seq_x = x[:, bos:eos].unsqueeze(0)
-        if bool(has_initial_state[seq_idx].item()):
-            initial_state = conv_states[slot, :, :state_len].unsqueeze(0)
-        else:
-            initial_state = torch.zeros(
-                1,
-                weight.shape[0],
-                state_len,
-                device=seq_x.device,
-                dtype=seq_x.dtype,
-            )
-
-        conv_input = torch.cat([initial_state, seq_x], dim=-1).to(weight.dtype)
-        seq_out = F.conv1d(
-            conv_input,
-            weight,
-            bias,
-            padding=0,
-            groups=weight.shape[0],
-        )
-        seq_out = seq_out[..., -seq_x.shape[-1] :].to(dtype=x.dtype)
-        if activation in ("silu", "swish"):
-            seq_out = F.silu(seq_out)
-
-        out[:, bos:eos] = seq_out.squeeze(0)
-        conv_states[slot, :, :state_len].copy_(conv_input[..., -state_len:].squeeze(0))
 
-    return out
 
 
 def causal_conv1d_update_torch(