added cli mamba option and clean up

Signed-off-by: Rishi Astra <[email protected]>

Author: RishiAstra

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -466,6 +466,9 @@ def __init__(
         # The tuple is (conv_state, ssm_state)
         self.kv_cache = (torch.tensor([]), torch.tensor([]))
 
+        overrides = getattr(model_config, "hf_overrides", {}) or {}
+        self.use_fast_kernel = bool(overrides.get("mamba2_fast_kernel", False))
+
         self.model_config = model_config
         self.cache_config = cache_config
         self.prefix = prefix
@@ -712,6 +715,7 @@ def forward_cuda(
                 dt_limit=(0.0, float("inf")),
                 out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim),
                 state_dtype=ssm_state.dtype,
+                use_fused_kernel=self.use_fast_kernel,
             )
 
             if prefix_caching_enabled:

vllm/model_executor/layers/mamba/ops/ssd_combined.py

@@ -45,7 +45,7 @@ def _mamba_chunk_scan_combined_fwd(
     dt_softplus=False,
     dt_limit=(0.0, float("inf")),
     state_dtype=None,
-    fused=False,
+    use_fused_kernel=False,
 ):
     assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2"
     seqlen, nheads, headdim = x.shape
@@ -80,8 +80,8 @@ def _mamba_chunk_scan_combined_fwd(
     if initial_states is not None:
         assert initial_states.shape == (len(cu_seqlens) - 1, nheads, headdim, dstate)
 
-    if fused:  # all 5 kernels fused
-        _, states, _, dA_cumsum, dt = _fused5_ssd(
+    if use_fused_kernel:  # all 5 kernels fused
+        _, states, dA_cumsum, dt = _fused5_ssd(
             x,
             dt,
             A,
@@ -201,6 +201,7 @@ def mamba_chunk_scan_combined_varlen(
     dt_limit=(0.0, float("inf")),
     return_intermediate_states=False,
     state_dtype=None,
+    use_fused_kernel=False,
 ):
     """
     Argument:
@@ -249,6 +250,7 @@ def mamba_chunk_scan_combined_varlen(
         dt_softplus=dt_softplus,
         dt_limit=dt_limit,
         state_dtype=state_dtype,
+        use_fused_kernel=use_fused_kernel,
     )
 
     return varlen_states

vllm/model_executor/layers/mamba/ops/ssd_fused5.py

@@ -1051,9 +1051,4 @@ def _fused5_ssd(
         CB_COMP_FP32=cb_comp_fp32,
     )
 
-    # states_G holds both states and final states
-    # TODO: can skip this copy if copied outside of function
-    final_states = states_G[nchunks].to(
-        states_G.dtype, copy=True
-    )  # copy and convert to expected dtype
-    return out_x, states_G[1:], final_states, dA_cumsum, dt_out
+    return out_x, states_G[1:], dA_cumsum, dt_out