From 7c568f5154915b6325877b0b58ea1babe20a4234 Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Tue, 31 Mar 2026 11:50:36 +0530 Subject: [PATCH 01/43] feat: initial Granite4 enablement modifications Signed-off-by: Akash Kaothalkar --- .../layers/mamba/ops/causal_conv1d.py | 220 +++++++++++++++++- .../layers/mamba/ops/mamba_ssm.py | 73 ++++++ .../layers/mamba/ops/ssd_combined.py | 122 +++++++++- vllm/model_executor/layers/utils.py | 8 + 4 files changed, 420 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index a8efdc9f18b5..9de4c8d79864 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -8,7 +8,7 @@ import numpy as np import torch -from vllm.triton_utils import tl, triton +from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID @@ -466,6 +466,212 @@ def _causal_conv1d_fwd_kernel( # continuous batching tl.store(o_ptrs, acc, mask=mask_1d) +def _causal_conv1d_fn_cpu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, + activation: str | None = "silu", + pad_slot_id: int = PAD_SLOT_ID, +) -> torch.Tensor: + """Pure PyTorch CPU fallback for causal_conv1d_fn (prefill path). + + Implements depthwise causal 1D convolution over variable-length sequences + with conv state caching and optional SiLU activation. + """ + if isinstance(activation, bool) and activation: + activation = "silu" + + 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.empty_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 + if seq_len == 0: + continue + + # Determine cache index for this sequence + if cache_indices is not None: + cache_idx = cache_indices[b].item() + if cache_idx == pad_slot_id: + continue + else: + cache_idx = b + + # Get initial state if available + use_initial = (has_initial_state is not None + and has_initial_state[b].item()) + + # x_seq: (dim, seq_len) + x_seq = x[:, seq_start:seq_end] + + if use_initial: + # Prepend conv state to input for full convolution context + # conv_state: (dim, state_len) + init_state = conv_states[cache_idx, :, :state_len] + # Concatenate: (dim, state_len + seq_len) + x_padded = torch.cat([init_state, x_seq], dim=1) + else: + # Zero-pad on the left + x_padded = torch.nn.functional.pad(x_seq, (state_len, 0)) + + # Depthwise conv1d: for each channel independently + # weight: (dim, width) — apply as depthwise filter + # x_padded: (dim, state_len + seq_len) + # Output length = seq_len (causal: no padding on right) + out_seq = torch.zeros(dim, seq_len, dtype=x.dtype, device=x.device) + for i in range(seq_len): + # Extract window: (dim, width) + window = x_padded[:, i : i + width] + # Element-wise multiply and sum over width + val = (window * weight).sum(dim=1) # (dim,) + if bias is not None: + val = val + bias + out_seq[:, i] = val + + # Apply activation + if activation in ["silu", "swish"]: + out_seq = out_seq * torch.sigmoid(out_seq) + + out[:, seq_start:seq_end] = out_seq + + # Update conv state with last `state_len` tokens + if seq_len >= state_len: + conv_states[cache_idx, :, :state_len] = x_seq[:, -state_len:] + else: + # Shift and append + conv_states[cache_idx, :, : state_len - seq_len] = conv_states[ + cache_idx, :, seq_len:state_len + ] + conv_states[cache_idx, :, state_len - seq_len :] = x_seq + + return out.to(original_x_dtype) + + +def _causal_conv1d_update_cpu( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: bool | str | None = None, + conv_state_indices: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + pad_slot_id: int = PAD_SLOT_ID, +) -> torch.Tensor: + """Pure PyTorch CPU fallback for causal_conv1d_update (decode path). + + Vectorized: processes all batch elements simultaneously using + batched tensor operations instead of Python for-loops. + """ + if isinstance(activation, bool): + activation = "silu" if activation else None + + original_x_dtype = x.dtype + x = x.to(conv_state.dtype) + _, width = weight.shape + state_len = width - 1 + + if query_start_loc is None and x.dim() == 2: + x = x.unsqueeze(-1) + unsqueeze = True + else: + unsqueeze = False + + # Fast path: no varlen, handle all batches at once + if query_start_loc is None: + batch, dim, seqlen = x.shape + + if conv_state_indices is not None: + cache_idxs = conv_state_indices.flatten() + # Build valid mask (non-padded entries) + valid_mask = cache_idxs != pad_slot_id # (batch,) + else: + cache_idxs = torch.arange(batch, device=x.device) + valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) + + for t in range(seqlen): + x_t = x[:, :, t].clone() # (batch, dim) - clone to avoid aliasing + + # Gather states for all valid batch entries + states = conv_state[cache_idxs] # (batch, dim, state_len) + + # Build windows: [state, x_t] -> (batch, dim, width) + windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) + + # Compute: (batch, dim, width) * (dim, width) -> sum over width + # weight is (dim, width), broadcast over batch + val = (windows * weight.unsqueeze(0)).sum(dim=-1) # (batch, dim) + if bias is not None: + val = val + bias.unsqueeze(0) + + if activation in ["silu", "swish"]: + val = val * torch.sigmoid(val) + + # Mask out padded entries + val = val * valid_mask.unsqueeze(-1).to(val.dtype) + x[:, :, t] = val + + # Shift state left, push ORIGINAL x_t (not computed val) + new_state = torch.cat( + [states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1 + ) # (batch, dim, state_len) + conv_state[cache_idxs] = new_state + + out = x + if unsqueeze: + out = out.squeeze(-1) + return out.to(original_x_dtype) + + # Varlen path: must loop over sequences (different lengths) + batch = conv_state_indices.size(0) + dim = x.size(1) + out = x.clone() + + for b in range(batch): + cache_idx = conv_state_indices[b].item() + if cache_idx == pad_slot_id: + continue + + seq_start = query_start_loc[b].item() + seq_end = query_start_loc[b + 1].item() + seqlen_b = seq_end - seq_start + + if seqlen_b == 0: + continue + + local_state = conv_state[cache_idx].clone() # (dim, state_len) + + for t in range(seqlen_b): + x_t = x[seq_start + t, :] # (dim,) + + window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) + val = (window * weight).sum(dim=-1) + if bias is not None: + val = val + bias + if activation in ["silu", "swish"]: + val = val * torch.sigmoid(val) + + out[seq_start + t, :] = val + + if state_len > 1: + local_state[:, :-1] = local_state[:, 1:].clone() + local_state[:, -1] = x_t + + conv_state[cache_idx] = local_state + + return out.to(original_x_dtype) + + def causal_conv1d_fn( x: torch.Tensor, weight: torch.Tensor, @@ -536,6 +742,12 @@ def causal_conv1d_fn( The block size to align the cached states to out: same shape as `x` """ + if not HAS_TRITON: + return _causal_conv1d_fn_cpu( + x, weight, bias, conv_states, query_start_loc, + cache_indices, has_initial_state, activation, pad_slot_id, + ) + if isinstance(activation, bool) and activation: activation = "silu" @@ -1122,6 +1334,12 @@ def causal_conv1d_update( indices 0 and 3 out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x` """ + if not HAS_TRITON: + return _causal_conv1d_update_cpu( + x, conv_state, weight, bias, activation, + conv_state_indices, query_start_loc, null_block_id, + ) + if validate_data: assert null_block_id is not None assert x.stride(1) == 1 diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index c4a0ef385d1c..13557b8ec189 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -429,6 +429,79 @@ def selective_state_update( if num_accepted_tokens is not None: assert num_accepted_tokens.shape == (N,) + if not HAS_TRITON: + # CPU fallback: vectorized PyTorch implementation of SSM state update + nheads_ngroups_ratio = nheads // ngroups + for seq_idx in range(N): + if cu_seqlens is not None: + bos = cu_seqlens[seq_idx].item() + seq_len = cu_seqlens[seq_idx + 1].item() - bos + else: + bos = seq_idx + seq_len = 1 + + # Determine state indices + if state_batch_indices is not None: + state_idx = state_batch_indices[seq_idx, 0].item() + if state_idx == null_block_id: + continue + else: + state_idx = seq_idx + + if dst_state_batch_indices is not None: + dst_idx = dst_state_batch_indices[seq_idx, 0].item() + else: + dst_idx = state_idx + + # Load state for all heads at once: (nheads, dim, dstate) + s = state[state_idx].float() + + for t in range(seq_len): + token_idx = bos + t + + x_val = x[token_idx].float() # (nheads, dim) + dt_val = dt[token_idx].float() # (nheads, dim) + + if dt_bias is not None: + dt_val = dt_val + dt_bias.float() + if dt_softplus: + dt_val = torch.nn.functional.softplus(dt_val) + + A_val = A.float() # (nheads, dim, dstate) + + # Expand B and C from groups to heads + # B: (ngroups, dstate) -> (nheads, dstate) + B_val = B[token_idx].float() # (ngroups, dstate) + B_expanded = B_val.repeat_interleave( + nheads_ngroups_ratio, dim=0) # (nheads, dstate) + C_val = C[token_idx].float() # (ngroups, dstate) + C_expanded = C_val.repeat_interleave( + nheads_ngroups_ratio, dim=0) # (nheads, dstate) + + # SSM state update — all heads at once + # dA = exp(A * dt) -> (nheads, dim, dstate) + dA = torch.exp(A_val * dt_val.unsqueeze(-1)) + # dB*x: outer product -> (nheads, dim, dstate) + dBx = (B_expanded.unsqueeze(1) + * (x_val * dt_val).unsqueeze(-1)) + s = s * dA + dBx + + # Output: sum(state * C, dim=-1) -> (nheads, dim) + out_val = (s * C_expanded.unsqueeze(1)).sum(dim=-1) + + if D is not None: + out_val = out_val + x_val * D.float() + + if z is not None: + z_val = z[token_idx].float() # (nheads, dim) + out_val = out_val * z_val * torch.sigmoid(z_val) + + out[token_idx] = out_val.to(out.dtype) + + # Store final state for all heads + state[dst_idx] = s.to(state.dtype) + return + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), N, nheads) z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) state_batch_indices_strides = ( diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 4c93a768b629..4ce5eb67bc15 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -10,19 +10,122 @@ from einops import rearrange from packaging import version -from vllm.triton_utils import triton +from vllm.triton_utils import HAS_TRITON, triton from .ssd_bmm import _bmm_chunk_fwd from .ssd_chunk_scan import _chunk_scan_fwd from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd from .ssd_state_passing import _state_passing_fwd -TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0") +TRITON_22 = HAS_TRITON and version.parse(triton.__version__) >= version.parse("2.2.0") def is_int_pow_2(n): return isinstance(n, int) and n > 0 and (n & (n - 1)) == 0 +def _mamba_chunk_scan_combined_fwd_cpu( + x, # (seqlen, nheads, headdim) + dt, # (seqlen, nheads) + A, # (nheads,) + B, # (seqlen, ngroups, dstate) + C, # (seqlen, ngroups, dstate) + chunk_size, + out, # (seqlen, nheads, headdim) - preallocated + D=None, + z=None, + dt_bias=None, + initial_states=None, + return_intermediate_states=False, + seq_idx=None, + cu_seqlens=None, + cu_chunk_seqlens=None, + last_chunk_indices=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), + state_dtype=None, +): + """Pure PyTorch CPU fallback for the entire SSD scan pipeline. + + Replaces all 5 Triton sub-ops with a single sequential SSM scan. + For each token: state = state * exp(A*dt) + outer(x*dt, B), + y = (state * C).sum(-1) + D*x, optionally gated by z. + """ + seqlen, nheads, headdim = x.shape + _, ngroups, dstate = B.shape + nheads_per_group = nheads // ngroups + + assert cu_seqlens is not None + batch = cu_seqlens.size(0) - 1 + + # Prepare dt + dt_f = dt.float() + if dt_bias is not None: + dt_f = dt_f + dt_bias.float().unsqueeze(0) + if dt_softplus: + dt_f = torch.nn.functional.softplus(dt_f) + if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): + dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) + + # Preallocate final states: (batch, nheads, headdim, dstate) + all_states = torch.zeros( + batch, nheads, headdim, dstate, + dtype=torch.float32, device=x.device + ) + + for b_idx in range(batch): + seq_start = cu_seqlens[b_idx].item() + seq_end = cu_seqlens[b_idx + 1].item() + + # Initialize state: (nheads, headdim, dstate) + if initial_states is not None: + state = initial_states[b_idx].float() + else: + state = torch.zeros( + nheads, headdim, dstate, + dtype=torch.float32, device=x.device + ) + + for t in range(seq_start, seq_end): + x_t = x[t].float() # (nheads, headdim) + dt_t = dt_f[t] # (nheads,) + A_val = A.float() # (nheads,) + + # dA = exp(A * dt) -> (nheads, 1, 1) for broadcasting + dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) + + # Expand B,C from groups to heads: (ngroups, dstate) -> (nheads, dstate) + B_expanded = B[t].float().repeat_interleave( + nheads_per_group, dim=0) # (nheads, dstate) + C_expanded = C[t].float().repeat_interleave( + nheads_per_group, dim=0) # (nheads, dstate) + + # State update — all heads at once + # outer(x*dt, B) -> (nheads, headdim, dstate) + xdt = (x_t * dt_t.unsqueeze(-1)) # (nheads, headdim) + dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) # (nheads, headdim, dstate) + state = state * dA + dBx + + # Output: (state * C).sum(-1) -> (nheads, headdim) + y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) + + if D is not None: + y = y + x_t * D.float().unsqueeze(-1) if D.dim() == 1 else y + x_t * D.float() + + if z is not None: + z_t = z[t].float() # (nheads, headdim) + y = y * z_t * torch.sigmoid(z_t) + + out[t] = y.to(out.dtype) + + all_states[b_idx] = state.to(all_states.dtype) + + # Cast to expected output dtype (matches Triton path behavior) + out_dtype = state_dtype if state_dtype is not None else x.dtype + all_states = all_states.to(out_dtype) + + return all_states + + def _mamba_chunk_scan_combined_fwd( x, @@ -45,6 +148,21 @@ def _mamba_chunk_scan_combined_fwd( dt_limit=(0.0, float("inf")), state_dtype=None, ): + if not HAS_TRITON: + return _mamba_chunk_scan_combined_fwd_cpu( + x, dt, A, B, C, chunk_size, out, + D=D, z=z, dt_bias=dt_bias, + initial_states=initial_states, + return_intermediate_states=return_intermediate_states, + seq_idx=seq_idx, + cu_seqlens=cu_seqlens, + cu_chunk_seqlens=cu_chunk_seqlens, + last_chunk_indices=last_chunk_indices, + dt_softplus=dt_softplus, + dt_limit=dt_limit, + state_dtype=state_dtype, + ) + assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2" seqlen, nheads, headdim = x.shape _, ngroups, dstate = B.shape diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 4918c83bdc39..6947f14be590 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -228,6 +228,14 @@ def dispatch_cpu_unquantized_gemm( layer.cpu_linear = torch.nn.functional.linear return + # Skip CPU GEMM dispatch for non-2D weights (e.g. MoE 3D expert weights). + # These layers are handled by their own specialized methods. + if layer.weight.dim() != 2: + layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear( + x, weight, bias + ) + return + N, K = layer.weight.size() dtype = layer.weight.dtype From f81e3716ddeee65661e3ff6d2b5e5a63f04ac030 Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Tue, 31 Mar 2026 16:22:51 +0530 Subject: [PATCH 02/43] granite4 running on power Signed-off-by: Akash Kaothalkar Signed-off-by: Akash Kaothalkar --- .../layers/mamba/ops/causal_conv1d.py | 26 +++++++++---------- vllm/platforms/cpu.py | 2 +- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 9de4c8d79864..238e24d571e9 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -526,18 +526,18 @@ def _causal_conv1d_fn_cpu( x_padded = torch.nn.functional.pad(x_seq, (state_len, 0)) # Depthwise conv1d: for each channel independently - # weight: (dim, width) — apply as depthwise filter - # x_padded: (dim, state_len + seq_len) - # Output length = seq_len (causal: no padding on right) - out_seq = torch.zeros(dim, seq_len, dtype=x.dtype, device=x.device) - for i in range(seq_len): - # Extract window: (dim, width) - window = x_padded[:, i : i + width] - # Element-wise multiply and sum over width - val = (window * weight).sum(dim=1) # (dim,) - if bias is not None: - val = val + bias - out_seq[:, i] = val + # x_padded_batch: (1, dim, state_len + seq_len) + x_padded_batch = x_padded.unsqueeze(0) + # weight_conv: (dim, 1, width) + weight_conv = weight.unsqueeze(1) + + # F.conv1d performs the exact cross-correlation matching our previous for-loop + out_seq = torch.nn.functional.conv1d( + x_padded_batch, + weight=weight_conv, + bias=bias, + groups=dim + ).squeeze(0) # (dim, seq_len) # Apply activation if activation in ["silu", "swish"]: @@ -552,7 +552,7 @@ def _causal_conv1d_fn_cpu( # Shift and append conv_states[cache_idx, :, : state_len - seq_len] = conv_states[ cache_idx, :, seq_len:state_len - ] + ].clone() conv_states[cache_idx, :, state_len - seq_len :] = x_seq return out.to(original_x_dtype) diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 7fbad3e4c76e..848872c6a554 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -374,7 +374,7 @@ def check_and_update_config(cls, vllm_config: VllmConfig) -> None: def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: # TODO: CPU still sets block_size in check_and_update_config. # Move that logic here so block_size is chosen by the backend. - pass + super().update_block_size_for_backend(vllm_config) @classmethod def get_allowed_cpu_core_node_list(cls) -> tuple[list[int], list[LogicalCPUInfo]]: From fbfa31d856224036607fa5b11a462ff63fc3d473 Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Mon, 6 Apr 2026 19:24:29 +0530 Subject: [PATCH 03/43] granite4 cpu fallback support Signed-off-by: Akash Kaothalkar Signed-off-by: Akash Kaothalkar --- .../layers/mamba/ops/causal_conv1d.py | 259 ++---------- .../layers/mamba/ops/cpu_fallbacks.py | 382 ++++++++++++++++++ .../layers/mamba/ops/mamba_ssm.py | 83 +--- .../layers/mamba/ops/ssd_combined.py | 134 +----- 4 files changed, 447 insertions(+), 411 deletions(-) create mode 100644 vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 238e24d571e9..f339e1417efa 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -5,11 +5,10 @@ # Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py -import numpy as np -import torch - +from vllm.model_executor.custom_op import CustomOp from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID +from .cpu_fallbacks import _causal_conv1d_fn_cpu, _causal_conv1d_update_cpu @triton.jit() @@ -466,213 +465,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching tl.store(o_ptrs, acc, mask=mask_1d) -def _causal_conv1d_fn_cpu( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None, - conv_states: torch.Tensor, - query_start_loc: torch.Tensor, - cache_indices: torch.Tensor | None = None, - has_initial_state: torch.Tensor | None = None, - activation: str | None = "silu", - pad_slot_id: int = PAD_SLOT_ID, -) -> torch.Tensor: - """Pure PyTorch CPU fallback for causal_conv1d_fn (prefill path). - - Implements depthwise causal 1D convolution over variable-length sequences - with conv state caching and optional SiLU activation. - """ - if isinstance(activation, bool) and activation: - activation = "silu" - - 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.empty_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 - if seq_len == 0: - continue - - # Determine cache index for this sequence - if cache_indices is not None: - cache_idx = cache_indices[b].item() - if cache_idx == pad_slot_id: - continue - else: - cache_idx = b - - # Get initial state if available - use_initial = (has_initial_state is not None - and has_initial_state[b].item()) - - # x_seq: (dim, seq_len) - x_seq = x[:, seq_start:seq_end] - - if use_initial: - # Prepend conv state to input for full convolution context - # conv_state: (dim, state_len) - init_state = conv_states[cache_idx, :, :state_len] - # Concatenate: (dim, state_len + seq_len) - x_padded = torch.cat([init_state, x_seq], dim=1) - else: - # Zero-pad on the left - x_padded = torch.nn.functional.pad(x_seq, (state_len, 0)) - - # Depthwise conv1d: for each channel independently - # x_padded_batch: (1, dim, state_len + seq_len) - x_padded_batch = x_padded.unsqueeze(0) - # weight_conv: (dim, 1, width) - weight_conv = weight.unsqueeze(1) - - # F.conv1d performs the exact cross-correlation matching our previous for-loop - out_seq = torch.nn.functional.conv1d( - x_padded_batch, - weight=weight_conv, - bias=bias, - groups=dim - ).squeeze(0) # (dim, seq_len) - - # Apply activation - if activation in ["silu", "swish"]: - out_seq = out_seq * torch.sigmoid(out_seq) - - out[:, seq_start:seq_end] = out_seq - - # Update conv state with last `state_len` tokens - if seq_len >= state_len: - conv_states[cache_idx, :, :state_len] = x_seq[:, -state_len:] - else: - # Shift and append - conv_states[cache_idx, :, : state_len - seq_len] = conv_states[ - cache_idx, :, seq_len:state_len - ].clone() - conv_states[cache_idx, :, state_len - seq_len :] = x_seq - - return out.to(original_x_dtype) - - -def _causal_conv1d_update_cpu( - x: torch.Tensor, - conv_state: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None = None, - activation: bool | str | None = None, - conv_state_indices: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, - pad_slot_id: int = PAD_SLOT_ID, -) -> torch.Tensor: - """Pure PyTorch CPU fallback for causal_conv1d_update (decode path). - - Vectorized: processes all batch elements simultaneously using - batched tensor operations instead of Python for-loops. - """ - if isinstance(activation, bool): - activation = "silu" if activation else None - - original_x_dtype = x.dtype - x = x.to(conv_state.dtype) - _, width = weight.shape - state_len = width - 1 - - if query_start_loc is None and x.dim() == 2: - x = x.unsqueeze(-1) - unsqueeze = True - else: - unsqueeze = False - - # Fast path: no varlen, handle all batches at once - if query_start_loc is None: - batch, dim, seqlen = x.shape - - if conv_state_indices is not None: - cache_idxs = conv_state_indices.flatten() - # Build valid mask (non-padded entries) - valid_mask = cache_idxs != pad_slot_id # (batch,) - else: - cache_idxs = torch.arange(batch, device=x.device) - valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) - - for t in range(seqlen): - x_t = x[:, :, t].clone() # (batch, dim) - clone to avoid aliasing - - # Gather states for all valid batch entries - states = conv_state[cache_idxs] # (batch, dim, state_len) - - # Build windows: [state, x_t] -> (batch, dim, width) - windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) - - # Compute: (batch, dim, width) * (dim, width) -> sum over width - # weight is (dim, width), broadcast over batch - val = (windows * weight.unsqueeze(0)).sum(dim=-1) # (batch, dim) - if bias is not None: - val = val + bias.unsqueeze(0) - - if activation in ["silu", "swish"]: - val = val * torch.sigmoid(val) - - # Mask out padded entries - val = val * valid_mask.unsqueeze(-1).to(val.dtype) - x[:, :, t] = val - - # Shift state left, push ORIGINAL x_t (not computed val) - new_state = torch.cat( - [states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1 - ) # (batch, dim, state_len) - conv_state[cache_idxs] = new_state - - out = x - if unsqueeze: - out = out.squeeze(-1) - return out.to(original_x_dtype) - - # Varlen path: must loop over sequences (different lengths) - batch = conv_state_indices.size(0) - dim = x.size(1) - out = x.clone() - - for b in range(batch): - cache_idx = conv_state_indices[b].item() - if cache_idx == pad_slot_id: - continue - - seq_start = query_start_loc[b].item() - seq_end = query_start_loc[b + 1].item() - seqlen_b = seq_end - seq_start - - if seqlen_b == 0: - continue - - local_state = conv_state[cache_idx].clone() # (dim, state_len) - - for t in range(seqlen_b): - x_t = x[seq_start + t, :] # (dim,) - - window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) - val = (window * weight).sum(dim=-1) - if bias is not None: - val = val + bias - if activation in ["silu", "swish"]: - val = val * torch.sigmoid(val) - - out[seq_start + t, :] = val - - if state_len > 1: - local_state[:, :-1] = local_state[:, 1:].clone() - local_state[:, -1] = x_t - - conv_state[cache_idx] = local_state - - return out.to(original_x_dtype) - - -def causal_conv1d_fn( +def _causal_conv1d_fn_cuda( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, @@ -742,12 +535,6 @@ def causal_conv1d_fn( The block size to align the cached states to out: same shape as `x` """ - if not HAS_TRITON: - return _causal_conv1d_fn_cpu( - x, weight, bias, conv_states, query_start_loc, - cache_indices, has_initial_state, activation, pad_slot_id, - ) - if isinstance(activation, bool) and activation: activation = "silu" @@ -1281,7 +1068,7 @@ def _causal_conv1d_update_kernel( tl.store(o_ptrs, acc, mask=mask_1d) -def causal_conv1d_update( +def _causal_conv1d_update_cuda( x: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor, @@ -1334,12 +1121,6 @@ def causal_conv1d_update( indices 0 and 3 out: (batch, dim) or (batch, dim, seqlen) or (num_tokens, dim), same shape as `x` """ - if not HAS_TRITON: - return _causal_conv1d_update_cpu( - x, conv_state, weight, bias, activation, - conv_state_indices, query_start_loc, null_block_id, - ) - if validate_data: assert null_block_id is not None assert x.stride(1) == 1 @@ -1461,3 +1242,35 @@ def grid(META): if unsqueeze: out = out.squeeze(-1) return out.to(original_x_dtype) + + +class CausalConv1dFnOp(CustomOp): + def forward_native(self, *args, **kwargs): + return _causal_conv1d_fn_cpu(*args, **kwargs) + + def forward_cpu(self, *args, **kwargs): + return _causal_conv1d_fn_cpu(*args, **kwargs) + + def forward_cuda(self, *args, **kwargs): + return _causal_conv1d_fn_cuda(*args, **kwargs) + + +class CausalConv1dUpdateOp(CustomOp): + def forward_native(self, *args, **kwargs): + return _causal_conv1d_update_cpu(*args, **kwargs) + + def forward_cpu(self, *args, **kwargs): + return _causal_conv1d_update_cpu(*args, **kwargs) + + def forward_cuda(self, *args, **kwargs): + return _causal_conv1d_update_cuda(*args, **kwargs) + + +_causal_conv1d_fn_op = CausalConv1dFnOp() +_causal_conv1d_update_op = CausalConv1dUpdateOp() + +def causal_conv1d_fn(*args, **kwargs): + return _causal_conv1d_fn_op(*args, **kwargs) + +def causal_conv1d_update(*args, **kwargs): + return _causal_conv1d_update_op(*args, **kwargs) diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py new file mode 100644 index 000000000000..717eff1bec36 --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -0,0 +1,382 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from vllm.v1.attention.backends.utils import PAD_SLOT_ID, NULL_BLOCK_ID + +def _causal_conv1d_fn_cpu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, + activation: str | None = "silu", + pad_slot_id: int = PAD_SLOT_ID, +) -> torch.Tensor: + """Pure PyTorch CPU fallback for causal_conv1d_fwd.""" + if isinstance(activation, bool) and activation: + activation = "silu" + + 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 + + if seq_len == 0: + continue + + cache_idx = cache_indices[b].item() if cache_indices is not None else b + + if cache_idx == 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] # (dim, state_len) + 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 + + if state_len > 1: + state[:, :-1] = state[:, 1:].clone() + state[:, -1] = x_t + + if seq_len >= state_len: + conv_states[cache_idx, :, :state_len] = x_seq[:, -state_len:] + else: + conv_states[cache_idx, :, : state_len - seq_len] = conv_states[ + cache_idx, :, seq_len:state_len + ].clone() + conv_states[cache_idx, :, state_len - seq_len :] = x_seq + + return out.to(original_x_dtype) + + +def _causal_conv1d_update_cpu( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: bool | str | None = None, + conv_state_indices: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + pad_slot_id: int = PAD_SLOT_ID, +) -> torch.Tensor: + """Pure PyTorch CPU fallback for causal_conv1d_update (decode path).""" + if isinstance(activation, bool): + activation = "silu" if activation else None + + original_x_dtype = x.dtype + x = x.to(conv_state.dtype) + _, width = weight.shape + state_len = width - 1 + + if query_start_loc is None and x.dim() == 2: + x = x.unsqueeze(-1) + unsqueeze = True + else: + unsqueeze = False + + if query_start_loc is None: + batch, dim, seqlen = x.shape + + if conv_state_indices is not None: + cache_idxs = conv_state_indices.flatten() + valid_mask = cache_idxs != pad_slot_id + else: + cache_idxs = torch.arange(batch, device=x.device) + valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) + + for t in range(seqlen): + x_t = x[:, :, t].clone() + + states = conv_state[cache_idxs] + + windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) + + val = (windows * weight.unsqueeze(0)).sum(dim=-1) + if bias is not None: + val = val + bias.unsqueeze(0) + + if activation in ["silu", "swish"]: + val = val * torch.sigmoid(val) + + val = val * valid_mask.unsqueeze(-1).to(val.dtype) + x[:, :, t] = val + + new_state = torch.cat( + [states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1 + ) + conv_state[cache_idxs] = new_state + + out = x + if unsqueeze: + out = out.squeeze(-1) + return out.to(original_x_dtype) + + batch = conv_state_indices.size(0) + dim = x.size(1) + out = x.clone() + + for b in range(batch): + cache_idx = conv_state_indices[b].item() + if cache_idx == pad_slot_id: + continue + + seq_start = query_start_loc[b].item() + seq_end = query_start_loc[b + 1].item() + seqlen_b = seq_end - seq_start + + if seqlen_b == 0: + continue + + local_state = conv_state[cache_idx].clone() + + for t in range(seqlen_b): + x_t = x[seq_start + t, :] + + window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) + val = (window * weight).sum(dim=-1) + if bias is not None: + val = val + bias + if activation in ["silu", "swish"]: + val = val * torch.sigmoid(val) + + out[seq_start + t, :] = val + + if state_len > 1: + local_state[:, :-1] = local_state[:, 1:].clone() + local_state[:, -1] = x_t + + conv_state[cache_idx] = local_state + + return out.to(original_x_dtype) + + +def _selective_state_update_cpu( + state, + x, + dt, + A, + B, + C, + D=None, + z=None, + dt_bias=None, + dt_softplus=False, + state_batch_indices=None, + dst_state_batch_indices=None, + null_block_id=NULL_BLOCK_ID, + out=None, + num_accepted_tokens=None, + cu_seqlens=None, + is_blackwell=False, + enable_stochastic_rounding=False, + cache_philox_rounds=0, +): + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + if out.dim() == 2: + out = out.unsqueeze(1) + if state_batch_indices is not None and state_batch_indices.dim() == 1: + state_batch_indices = state_batch_indices.unsqueeze(1) + if dst_state_batch_indices is not None and dst_state_batch_indices.dim() == 1: + dst_state_batch_indices = dst_state_batch_indices.unsqueeze(1) + + _, nheads, dim, dstate = state.shape + batch = x.shape[0] + if cu_seqlens is not None: + N = len(cu_seqlens) - 1 + else: + N = batch + + ngroups = B.shape[1] + nheads_ngroups_ratio = nheads // ngroups + + for seq_idx in range(N): + if cu_seqlens is not None: + bos = cu_seqlens[seq_idx].item() + seq_len = cu_seqlens[seq_idx + 1].item() - bos + else: + bos = seq_idx + seq_len = 1 + + if state_batch_indices is not None: + state_idx = state_batch_indices[seq_idx, 0].item() + if state_idx == null_block_id: + continue + else: + state_idx = seq_idx + + if dst_state_batch_indices is not None: + dst_idx = dst_state_batch_indices[seq_idx, 0].item() + else: + dst_idx = state_idx + + s = state[state_idx].float() + + for t in range(seq_len): + token_idx = bos + t + + x_val = x[token_idx].float() + dt_val = dt[token_idx].float() + + if dt_bias is not None: + dt_val = dt_val + dt_bias.float() + if dt_softplus: + dt_val = torch.nn.functional.softplus(dt_val) + + A_val = A.float() + + B_val = B[token_idx].float() + B_expanded = B_val.repeat_interleave(nheads_ngroups_ratio, dim=0) + C_val = C[token_idx].float() + C_expanded = C_val.repeat_interleave(nheads_ngroups_ratio, dim=0) + + dA = torch.exp(A_val * dt_val.unsqueeze(-1)) + dBx = (B_expanded.unsqueeze(1) * (x_val * dt_val).unsqueeze(-1)) + s = s * dA + dBx + + out_val = (s * C_expanded.unsqueeze(1)).sum(dim=-1) + + if D is not None: + out_val = out_val + x_val * D.float() + + if z is not None: + z_val = z[token_idx].float() + out_val = out_val * z_val * torch.sigmoid(z_val) + + out[token_idx] = out_val.to(out.dtype) + + state[dst_idx] = s.to(state.dtype) + + +def _mamba_chunk_scan_combined_fwd_cpu( + x, + dt, + A, + B, + C, + chunk_size, + out, + D=None, + z=None, + dt_bias=None, + initial_states=None, + return_intermediate_states=False, + seq_idx=None, + cu_seqlens=None, + cu_chunk_seqlens=None, + last_chunk_indices=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), + state_dtype=None, +): + seqlen, nheads, headdim = x.shape + _, ngroups, dstate = B.shape + nheads_per_group = nheads // ngroups + + assert cu_seqlens is not None + batch = cu_seqlens.size(0) - 1 + + dt_f = dt.float() + if dt_bias is not None: + dt_f = dt_f + dt_bias.float().unsqueeze(0) + if dt_softplus: + dt_f = torch.nn.functional.softplus(dt_f) + if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): + dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) + + all_states = torch.zeros( + batch, nheads, headdim, dstate, + dtype=torch.float32, device=x.device + ) + + for b_idx in range(batch): + seq_start = cu_seqlens[b_idx].item() + seq_end = cu_seqlens[b_idx + 1].item() + + if initial_states is not None: + state = initial_states[b_idx].float() + else: + state = torch.zeros( + nheads, headdim, dstate, + dtype=torch.float32, device=x.device + ) + + for t in range(seq_start, seq_end): + x_t = x[t].float() + dt_t = dt_f[t] + A_val = A.float() + + dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) + + B_expanded = B[t].float().repeat_interleave( + nheads_per_group, dim=0) + C_expanded = C[t].float().repeat_interleave( + nheads_per_group, dim=0) + + xdt = (x_t * dt_t.unsqueeze(-1)) + dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) + state = state * dA + dBx + + y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) + + if D is not None: + y = y + x_t * D.float().unsqueeze(-1) if D.dim() == 1 else y + x_t * D.float() + + if z is not None: + z_t = z[t].float() + y = y * z_t * torch.sigmoid(z_t) + + out[t] = y.to(out.dtype) + + all_states[b_idx] = state.to(all_states.dtype) + + out_dtype = state_dtype if state_dtype is not None else x.dtype + all_states = all_states.to(out_dtype) + + return all_states diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index 13557b8ec189..bda563fc7a15 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -7,6 +7,7 @@ import torch from packaging import version +from vllm.model_executor.custom_op import CustomOp from vllm import _custom_ops as ops from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp from vllm.triton_utils import HAS_TRITON, tl, triton @@ -316,7 +317,10 @@ def _selective_scan_update_kernel( tl.store(dst_state_ptrs, state, mask=mask) -def selective_state_update( +from .cpu_fallbacks import _selective_state_update_cpu + + +def _selective_state_update_cuda( state, x, dt, @@ -430,77 +434,12 @@ def selective_state_update( assert num_accepted_tokens.shape == (N,) if not HAS_TRITON: - # CPU fallback: vectorized PyTorch implementation of SSM state update - nheads_ngroups_ratio = nheads // ngroups - for seq_idx in range(N): - if cu_seqlens is not None: - bos = cu_seqlens[seq_idx].item() - seq_len = cu_seqlens[seq_idx + 1].item() - bos - else: - bos = seq_idx - seq_len = 1 - - # Determine state indices - if state_batch_indices is not None: - state_idx = state_batch_indices[seq_idx, 0].item() - if state_idx == null_block_id: - continue - else: - state_idx = seq_idx - - if dst_state_batch_indices is not None: - dst_idx = dst_state_batch_indices[seq_idx, 0].item() - else: - dst_idx = state_idx - - # Load state for all heads at once: (nheads, dim, dstate) - s = state[state_idx].float() - - for t in range(seq_len): - token_idx = bos + t - - x_val = x[token_idx].float() # (nheads, dim) - dt_val = dt[token_idx].float() # (nheads, dim) - - if dt_bias is not None: - dt_val = dt_val + dt_bias.float() - if dt_softplus: - dt_val = torch.nn.functional.softplus(dt_val) - - A_val = A.float() # (nheads, dim, dstate) - - # Expand B and C from groups to heads - # B: (ngroups, dstate) -> (nheads, dstate) - B_val = B[token_idx].float() # (ngroups, dstate) - B_expanded = B_val.repeat_interleave( - nheads_ngroups_ratio, dim=0) # (nheads, dstate) - C_val = C[token_idx].float() # (ngroups, dstate) - C_expanded = C_val.repeat_interleave( - nheads_ngroups_ratio, dim=0) # (nheads, dstate) - - # SSM state update — all heads at once - # dA = exp(A * dt) -> (nheads, dim, dstate) - dA = torch.exp(A_val * dt_val.unsqueeze(-1)) - # dB*x: outer product -> (nheads, dim, dstate) - dBx = (B_expanded.unsqueeze(1) - * (x_val * dt_val).unsqueeze(-1)) - s = s * dA + dBx - - # Output: sum(state * C, dim=-1) -> (nheads, dim) - out_val = (s * C_expanded.unsqueeze(1)).sum(dim=-1) - - if D is not None: - out_val = out_val + x_val * D.float() - - if z is not None: - z_val = z[token_idx].float() # (nheads, dim) - out_val = out_val * z_val * torch.sigmoid(z_val) - - out[token_idx] = out_val.to(out.dtype) - - # Store final state for all heads - state[dst_idx] = s.to(state.dtype) - return + return _selective_state_update_cpu( + state, x, dt, A, B, C, D, z, dt_bias, dt_softplus, + state_batch_indices, dst_state_batch_indices, + null_block_id, out, num_accepted_tokens, cu_seqlens, + is_blackwell, enable_stochastic_rounding, cache_philox_rounds + ) grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), N, nheads) z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 4ce5eb67bc15..647cecb5977d 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -10,6 +10,7 @@ from einops import rearrange from packaging import version +from vllm.model_executor.custom_op import CustomOp from vllm.triton_utils import HAS_TRITON, triton from .ssd_bmm import _bmm_chunk_fwd @@ -23,111 +24,11 @@ def is_int_pow_2(n): return isinstance(n, int) and n > 0 and (n & (n - 1)) == 0 -def _mamba_chunk_scan_combined_fwd_cpu( - x, # (seqlen, nheads, headdim) - dt, # (seqlen, nheads) - A, # (nheads,) - B, # (seqlen, ngroups, dstate) - C, # (seqlen, ngroups, dstate) - chunk_size, - out, # (seqlen, nheads, headdim) - preallocated - D=None, - z=None, - dt_bias=None, - initial_states=None, - return_intermediate_states=False, - seq_idx=None, - cu_seqlens=None, - cu_chunk_seqlens=None, - last_chunk_indices=None, - dt_softplus=False, - dt_limit=(0.0, float("inf")), - state_dtype=None, -): - """Pure PyTorch CPU fallback for the entire SSD scan pipeline. - - Replaces all 5 Triton sub-ops with a single sequential SSM scan. - For each token: state = state * exp(A*dt) + outer(x*dt, B), - y = (state * C).sum(-1) + D*x, optionally gated by z. - """ - seqlen, nheads, headdim = x.shape - _, ngroups, dstate = B.shape - nheads_per_group = nheads // ngroups - - assert cu_seqlens is not None - batch = cu_seqlens.size(0) - 1 - - # Prepare dt - dt_f = dt.float() - if dt_bias is not None: - dt_f = dt_f + dt_bias.float().unsqueeze(0) - if dt_softplus: - dt_f = torch.nn.functional.softplus(dt_f) - if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): - dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) +from .cpu_fallbacks import _mamba_chunk_scan_combined_fwd_cpu - # Preallocate final states: (batch, nheads, headdim, dstate) - all_states = torch.zeros( - batch, nheads, headdim, dstate, - dtype=torch.float32, device=x.device - ) - - for b_idx in range(batch): - seq_start = cu_seqlens[b_idx].item() - seq_end = cu_seqlens[b_idx + 1].item() - - # Initialize state: (nheads, headdim, dstate) - if initial_states is not None: - state = initial_states[b_idx].float() - else: - state = torch.zeros( - nheads, headdim, dstate, - dtype=torch.float32, device=x.device - ) - - for t in range(seq_start, seq_end): - x_t = x[t].float() # (nheads, headdim) - dt_t = dt_f[t] # (nheads,) - A_val = A.float() # (nheads,) - - # dA = exp(A * dt) -> (nheads, 1, 1) for broadcasting - dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) - - # Expand B,C from groups to heads: (ngroups, dstate) -> (nheads, dstate) - B_expanded = B[t].float().repeat_interleave( - nheads_per_group, dim=0) # (nheads, dstate) - C_expanded = C[t].float().repeat_interleave( - nheads_per_group, dim=0) # (nheads, dstate) - - # State update — all heads at once - # outer(x*dt, B) -> (nheads, headdim, dstate) - xdt = (x_t * dt_t.unsqueeze(-1)) # (nheads, headdim) - dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) # (nheads, headdim, dstate) - state = state * dA + dBx - - # Output: (state * C).sum(-1) -> (nheads, headdim) - y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) - - if D is not None: - y = y + x_t * D.float().unsqueeze(-1) if D.dim() == 1 else y + x_t * D.float() - - if z is not None: - z_t = z[t].float() # (nheads, headdim) - y = y * z_t * torch.sigmoid(z_t) - out[t] = y.to(out.dtype) - all_states[b_idx] = state.to(all_states.dtype) - - # Cast to expected output dtype (matches Triton path behavior) - out_dtype = state_dtype if state_dtype is not None else x.dtype - all_states = all_states.to(out_dtype) - - return all_states - - - -def _mamba_chunk_scan_combined_fwd( +def _mamba_chunk_scan_combined_fwd_cuda( x, dt, A, @@ -148,20 +49,6 @@ def _mamba_chunk_scan_combined_fwd( dt_limit=(0.0, float("inf")), state_dtype=None, ): - if not HAS_TRITON: - return _mamba_chunk_scan_combined_fwd_cpu( - x, dt, A, B, C, chunk_size, out, - D=D, z=z, dt_bias=dt_bias, - initial_states=initial_states, - return_intermediate_states=return_intermediate_states, - seq_idx=seq_idx, - cu_seqlens=cu_seqlens, - cu_chunk_seqlens=cu_chunk_seqlens, - last_chunk_indices=last_chunk_indices, - dt_softplus=dt_softplus, - dt_limit=dt_limit, - state_dtype=state_dtype, - ) assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2" seqlen, nheads, headdim = x.shape @@ -271,6 +158,21 @@ def _mamba_chunk_scan_combined_fwd( else: return states[last_chunk_indices] +class MambaChunkScanCombinedFwdOp(CustomOp): + def forward_native(self, *args, **kwargs): + return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) + + def forward_cpu(self, *args, **kwargs): + return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) + + def forward_cuda(self, *args, **kwargs): + return _mamba_chunk_scan_combined_fwd_cuda(*args, **kwargs) + +_mamba_chunk_scan_combined_fwd_op = MambaChunkScanCombinedFwdOp() + +def _mamba_chunk_scan_combined_fwd(*args, **kwargs): + return _mamba_chunk_scan_combined_fwd_op(*args, **kwargs) + def mamba_chunk_scan_combined_varlen( x, From 9604a4e0fae17dd03b710dbfd136956bf5364ef7 Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Tue, 7 Apr 2026 10:06:04 +0530 Subject: [PATCH 04/43] fix torch import Signed-off-by: Akash Kaothalkar --- vllm/model_executor/layers/mamba/ops/causal_conv1d.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index f339e1417efa..bedd3e105a0c 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -5,6 +5,7 @@ # Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py +import torch from vllm.model_executor.custom_op import CustomOp from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID From 60e6d899816803c4eba147caba1e1db7f2d7469b Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Tue, 7 Apr 2026 10:13:00 +0530 Subject: [PATCH 05/43] fix CustomOp Register Signed-off-by: Akash Kaothalkar --- .../layers/mamba/ops/causal_conv1d.py | 2 ++ .../layers/mamba/ops/mamba_ssm.py | 17 +++++++++++++++++ .../layers/mamba/ops/ssd_combined.py | 1 + 3 files changed, 20 insertions(+) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index bedd3e105a0c..355cae6d5ba6 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -1245,6 +1245,7 @@ def grid(META): return out.to(original_x_dtype) +@CustomOp.register("causal_conv1d_fn") class CausalConv1dFnOp(CustomOp): def forward_native(self, *args, **kwargs): return _causal_conv1d_fn_cpu(*args, **kwargs) @@ -1256,6 +1257,7 @@ def forward_cuda(self, *args, **kwargs): return _causal_conv1d_fn_cuda(*args, **kwargs) +@CustomOp.register("causal_conv1d_update") class CausalConv1dUpdateOp(CustomOp): def forward_native(self, *args, **kwargs): return _causal_conv1d_update_cpu(*args, **kwargs) diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index bda563fc7a15..63e03813e909 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -668,3 +668,20 @@ def selective_scan_fn( return delta # output written inplace to delta else: return z # output written inplace to z + +@CustomOp.register("selective_state_update") +class SelectiveStateUpdateOp(CustomOp): + def forward_native(self, *args, **kwargs): + return _selective_state_update_cpu(*args, **kwargs) + + def forward_cpu(self, *args, **kwargs): + return _selective_state_update_cpu(*args, **kwargs) + + def forward_cuda(self, *args, **kwargs): + return _selective_state_update_cuda(*args, **kwargs) + +_selective_state_update_op = SelectiveStateUpdateOp() + +def selective_state_update(*args, **kwargs): + return _selective_state_update_op(*args, **kwargs) + diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 647cecb5977d..26691c4a663f 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -158,6 +158,7 @@ def _mamba_chunk_scan_combined_fwd_cuda( else: return states[last_chunk_indices] +@CustomOp.register("mamba_chunk_scan_combined_fwd") class MambaChunkScanCombinedFwdOp(CustomOp): def forward_native(self, *args, **kwargs): return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) From 4e9a3a25ad9fe739c3a4907904ea403357bbb9c3 Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Tue, 7 Apr 2026 10:17:36 +0530 Subject: [PATCH 06/43] fix kwargs Signed-off-by: Akash Kaothalkar --- vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py index 717eff1bec36..fd3eab6cec93 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -14,6 +14,7 @@ def _causal_conv1d_fn_cpu( has_initial_state: torch.Tensor | None = None, activation: str | None = "silu", pad_slot_id: int = PAD_SLOT_ID, + **kwargs, ) -> torch.Tensor: """Pure PyTorch CPU fallback for causal_conv1d_fwd.""" if isinstance(activation, bool) and activation: @@ -88,6 +89,7 @@ def _causal_conv1d_update_cpu( conv_state_indices: torch.Tensor | None = None, query_start_loc: torch.Tensor | None = None, pad_slot_id: int = PAD_SLOT_ID, + **kwargs, ) -> torch.Tensor: """Pure PyTorch CPU fallback for causal_conv1d_update (decode path).""" if isinstance(activation, bool): @@ -200,6 +202,7 @@ def _selective_state_update_cpu( is_blackwell=False, enable_stochastic_rounding=False, cache_philox_rounds=0, + **kwargs, ): if state.dim() == 3: state = state.unsqueeze(1) @@ -314,6 +317,7 @@ def _mamba_chunk_scan_combined_fwd_cpu( dt_softplus=False, dt_limit=(0.0, float("inf")), state_dtype=None, + **kwargs, ): seqlen, nheads, headdim = x.shape _, ngroups, dstate = B.shape From b5b047857b7f4cab325d0f7ce134dac25df3a1dc Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Tue, 7 Apr 2026 12:11:48 +0530 Subject: [PATCH 07/43] fix: address review comments Signed-off-by: Akash Kaothalkar --- .../layers/mamba/ops/cpu_fallbacks.py | 22 +++++++++++++------ 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py index fd3eab6cec93..b419e0f38033 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -47,7 +47,7 @@ def _causal_conv1d_fn_cpu( 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] # (dim, state_len) + state = conv_states[cache_idx].clone() # (dim, state_len) else: state = torch.zeros((dim, state_len), dtype=x.dtype, device=x.device) @@ -136,7 +136,7 @@ def _causal_conv1d_update_cpu( new_state = torch.cat( [states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1 ) - conv_state[cache_idxs] = new_state + conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] out = x if unsqueeze: @@ -254,10 +254,11 @@ def _selective_state_update_cpu( else: state_idx = seq_idx - if dst_state_batch_indices is not None: - dst_idx = dst_state_batch_indices[seq_idx, 0].item() - else: - dst_idx = state_idx + if num_accepted_tokens is None: + if dst_state_batch_indices is not None: + dst_idx = dst_state_batch_indices[seq_idx, 0].item() + else: + dst_idx = state_idx s = state[state_idx].float() @@ -283,6 +284,11 @@ def _selective_state_update_cpu( dBx = (B_expanded.unsqueeze(1) * (x_val * dt_val).unsqueeze(-1)) s = s * dA + dBx + if num_accepted_tokens is not None: + token_dst_idx = dst_state_batch_indices[seq_idx, t].item() + if token_dst_idx != null_block_id: + state[token_dst_idx] = s.to(state.dtype) + out_val = (s * C_expanded.unsqueeze(1)).sum(dim=-1) if D is not None: @@ -294,7 +300,9 @@ def _selective_state_update_cpu( out[token_idx] = out_val.to(out.dtype) - state[dst_idx] = s.to(state.dtype) + if num_accepted_tokens is None: + if dst_idx != null_block_id: + state[dst_idx] = s.to(state.dtype) def _mamba_chunk_scan_combined_fwd_cpu( From fcae6a751c8b296f9d3dc78530481dd8e61bbb97 Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Tue, 7 Apr 2026 12:57:58 +0530 Subject: [PATCH 08/43] fix precommit Signed-off-by: Akash Kaothalkar --- .../layers/mamba/ops/causal_conv1d.py | 8 ++- .../layers/mamba/ops/cpu_fallbacks.py | 68 +++++++++---------- .../layers/mamba/ops/mamba_ssm.py | 38 ++++++++--- .../layers/mamba/ops/ssd_combined.py | 7 +- 4 files changed, 70 insertions(+), 51 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 355cae6d5ba6..9c867e7f1e04 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -6,9 +6,13 @@ import torch +import torch.nn as nn +import numpy as np + from vllm.model_executor.custom_op import CustomOp -from vllm.triton_utils import HAS_TRITON, tl, triton +from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID + from .cpu_fallbacks import _causal_conv1d_fn_cpu, _causal_conv1d_update_cpu @@ -1272,8 +1276,10 @@ def forward_cuda(self, *args, **kwargs): _causal_conv1d_fn_op = CausalConv1dFnOp() _causal_conv1d_update_op = CausalConv1dUpdateOp() + def causal_conv1d_fn(*args, **kwargs): return _causal_conv1d_fn_op(*args, **kwargs) + def causal_conv1d_update(*args, **kwargs): return _causal_conv1d_update_op(*args, **kwargs) diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py index b419e0f38033..c33a574395ce 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -2,7 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import torch -from vllm.v1.attention.backends.utils import PAD_SLOT_ID, NULL_BLOCK_ID + +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID + def _causal_conv1d_fn_cpu( x: torch.Tensor, @@ -49,8 +51,7 @@ def _causal_conv1d_fn_cpu( if has_initial_state is not None and has_initial_state[b]: state = conv_states[cache_idx].clone() # (dim, state_len) else: - state = torch.zeros((dim, state_len), - dtype=x.dtype, device=x.device) + state = torch.zeros((dim, state_len), dtype=x.dtype, device=x.device) for t in range(seq_len): x_t = x_seq[:, t] # (dim,) @@ -133,9 +134,7 @@ def _causal_conv1d_update_cpu( val = val * valid_mask.unsqueeze(-1).to(val.dtype) x[:, :, t] = val - new_state = torch.cat( - [states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1 - ) + new_state = torch.cat([states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1) conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] out = x @@ -143,8 +142,9 @@ def _causal_conv1d_update_cpu( out = out.squeeze(-1) return out.to(original_x_dtype) + assert conv_state_indices is not None + assert query_start_loc is not None batch = conv_state_indices.size(0) - dim = x.size(1) out = x.clone() for b in range(batch): @@ -231,10 +231,7 @@ def _selective_state_update_cpu( _, nheads, dim, dstate = state.shape batch = x.shape[0] - if cu_seqlens is not None: - N = len(cu_seqlens) - 1 - else: - N = batch + N = len(cu_seqlens) - 1 if cu_seqlens is not None else batch ngroups = B.shape[1] nheads_ngroups_ratio = nheads // ngroups @@ -281,7 +278,7 @@ def _selective_state_update_cpu( C_expanded = C_val.repeat_interleave(nheads_ngroups_ratio, dim=0) dA = torch.exp(A_val * dt_val.unsqueeze(-1)) - dBx = (B_expanded.unsqueeze(1) * (x_val * dt_val).unsqueeze(-1)) + dBx = B_expanded.unsqueeze(1) * (x_val * dt_val).unsqueeze(-1) s = s * dA + dBx if num_accepted_tokens is not None: @@ -300,19 +297,18 @@ def _selective_state_update_cpu( out[token_idx] = out_val.to(out.dtype) - if num_accepted_tokens is None: - if dst_idx != null_block_id: - state[dst_idx] = s.to(state.dtype) + if num_accepted_tokens is None and dst_idx != null_block_id: + state[dst_idx] = s.to(state.dtype) def _mamba_chunk_scan_combined_fwd_cpu( - x, - dt, - A, - B, - C, + x, + dt, + A, + B, + C, chunk_size, - out, + out, D=None, z=None, dt_bias=None, @@ -343,8 +339,7 @@ def _mamba_chunk_scan_combined_fwd_cpu( dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) all_states = torch.zeros( - batch, nheads, headdim, dstate, - dtype=torch.float32, device=x.device + batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device ) for b_idx in range(batch): @@ -355,33 +350,34 @@ def _mamba_chunk_scan_combined_fwd_cpu( state = initial_states[b_idx].float() else: state = torch.zeros( - nheads, headdim, dstate, - dtype=torch.float32, device=x.device + nheads, headdim, dstate, dtype=torch.float32, device=x.device ) for t in range(seq_start, seq_end): - x_t = x[t].float() - dt_t = dt_f[t] - A_val = A.float() + x_t = x[t].float() + dt_t = dt_f[t] + A_val = A.float() dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) - B_expanded = B[t].float().repeat_interleave( - nheads_per_group, dim=0) - C_expanded = C[t].float().repeat_interleave( - nheads_per_group, dim=0) + B_expanded = B[t].float().repeat_interleave(nheads_per_group, dim=0) + C_expanded = C[t].float().repeat_interleave(nheads_per_group, dim=0) - xdt = (x_t * dt_t.unsqueeze(-1)) - dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) + xdt = x_t * dt_t.unsqueeze(-1) + dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) state = state * dA + dBx y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) if D is not None: - y = y + x_t * D.float().unsqueeze(-1) if D.dim() == 1 else y + x_t * D.float() + y = ( + y + x_t * D.float().unsqueeze(-1) + if D.dim() == 1 + else y + x_t * D.float() + ) if z is not None: - z_t = z[t].float() + z_t = z[t].float() y = y * z_t * torch.sigmoid(z_t) out[t] = y.to(out.dtype) diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index 63e03813e909..ec38e516a04d 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -7,11 +7,12 @@ import torch from packaging import version -from vllm.model_executor.custom_op import CustomOp from vllm import _custom_ops as ops +from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from .cpu_fallbacks import _selective_state_update_cpu TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0")) @@ -49,8 +50,9 @@ def convert_rs_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) @triton.heuristics( { - "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] - is not None + "HAS_STATE_BATCH_INDICES": lambda args: ( + args["state_batch_indices_ptr"] is not None + ) } ) @triton.heuristics( @@ -317,9 +319,6 @@ def _selective_scan_update_kernel( tl.store(dst_state_ptrs, state, mask=mask) -from .cpu_fallbacks import _selective_state_update_cpu - - def _selective_state_update_cuda( state, x, @@ -435,10 +434,25 @@ def _selective_state_update_cuda( if not HAS_TRITON: return _selective_state_update_cpu( - state, x, dt, A, B, C, D, z, dt_bias, dt_softplus, - state_batch_indices, dst_state_batch_indices, - null_block_id, out, num_accepted_tokens, cu_seqlens, - is_blackwell, enable_stochastic_rounding, cache_philox_rounds + state, + x, + dt, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + dst_state_batch_indices, + null_block_id, + out, + num_accepted_tokens, + cu_seqlens, + is_blackwell, + enable_stochastic_rounding, + cache_philox_rounds, ) grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), N, nheads) @@ -669,6 +683,7 @@ def selective_scan_fn( else: return z # output written inplace to z + @CustomOp.register("selective_state_update") class SelectiveStateUpdateOp(CustomOp): def forward_native(self, *args, **kwargs): @@ -680,8 +695,9 @@ def forward_cpu(self, *args, **kwargs): def forward_cuda(self, *args, **kwargs): return _selective_state_update_cuda(*args, **kwargs) + _selective_state_update_op = SelectiveStateUpdateOp() + def selective_state_update(*args, **kwargs): return _selective_state_update_op(*args, **kwargs) - diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 26691c4a663f..7128f93c6791 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -17,6 +17,7 @@ from .ssd_chunk_scan import _chunk_scan_fwd from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd from .ssd_state_passing import _state_passing_fwd +from .cpu_fallbacks import _mamba_chunk_scan_combined_fwd_cpu TRITON_22 = HAS_TRITON and version.parse(triton.__version__) >= version.parse("2.2.0") @@ -24,9 +25,6 @@ def is_int_pow_2(n): return isinstance(n, int) and n > 0 and (n & (n - 1)) == 0 -from .cpu_fallbacks import _mamba_chunk_scan_combined_fwd_cpu - - def _mamba_chunk_scan_combined_fwd_cuda( x, @@ -158,6 +156,7 @@ def _mamba_chunk_scan_combined_fwd_cuda( else: return states[last_chunk_indices] + @CustomOp.register("mamba_chunk_scan_combined_fwd") class MambaChunkScanCombinedFwdOp(CustomOp): def forward_native(self, *args, **kwargs): @@ -169,8 +168,10 @@ def forward_cpu(self, *args, **kwargs): def forward_cuda(self, *args, **kwargs): return _mamba_chunk_scan_combined_fwd_cuda(*args, **kwargs) + _mamba_chunk_scan_combined_fwd_op = MambaChunkScanCombinedFwdOp() + def _mamba_chunk_scan_combined_fwd(*args, **kwargs): return _mamba_chunk_scan_combined_fwd_op(*args, **kwargs) From a68aba48d9fa57c1f21a496b01b2c04cc941be3e Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 7 Apr 2026 03:25:57 -0500 Subject: [PATCH 09/43] Address pre-commit issues Signed-off-by: Akash kaothalkar --- vllm/model_executor/layers/mamba/ops/causal_conv1d.py | 3 +-- vllm/model_executor/layers/mamba/ops/mamba_ssm.py | 1 + vllm/model_executor/layers/mamba/ops/ssd_combined.py | 3 +-- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 9c867e7f1e04..31def2856e7b 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -5,9 +5,8 @@ # Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py -import torch -import torch.nn as nn import numpy as np +import torch from vllm.model_executor.custom_op import CustomOp from vllm.triton_utils import tl, triton diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index ec38e516a04d..62e626cd080e 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -12,6 +12,7 @@ from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + from .cpu_fallbacks import _selective_state_update_cpu TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0")) diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 7128f93c6791..4b2ad011f819 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -13,11 +13,11 @@ from vllm.model_executor.custom_op import CustomOp from vllm.triton_utils import HAS_TRITON, triton +from .cpu_fallbacks import _mamba_chunk_scan_combined_fwd_cpu from .ssd_bmm import _bmm_chunk_fwd from .ssd_chunk_scan import _chunk_scan_fwd from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd from .ssd_state_passing import _state_passing_fwd -from .cpu_fallbacks import _mamba_chunk_scan_combined_fwd_cpu TRITON_22 = HAS_TRITON and version.parse(triton.__version__) >= version.parse("2.2.0") @@ -47,7 +47,6 @@ def _mamba_chunk_scan_combined_fwd_cuda( dt_limit=(0.0, float("inf")), state_dtype=None, ): - assert is_int_pow_2(chunk_size), "chunk_size must be integer power of 2" seqlen, nheads, headdim = x.shape _, ngroups, dstate = B.shape From 5ecd3ba8eb9b4fab62e1ce5aae2c4d1191f47844 Mon Sep 17 00:00:00 2001 From: Akash Kaothalkar Date: Sat, 25 Apr 2026 18:01:30 +0530 Subject: [PATCH 10/43] Add mamba files Signed-off-by: Akash Kaothalkar --- cmake/cpu_extension.cmake | 3 + csrc/cpu/mamba_cpu.cpp | 226 +++++++++++++++ csrc/cpu/mamba_kernels.hpp | 261 ++++++++++++++++++ csrc/cpu/torch_bindings.cpp | 35 +++ .../layers/mamba/ops/causal_conv1d.py | 105 ++++--- .../layers/mamba/ops/mamba_ssm.py | 89 ++++-- 6 files changed, 667 insertions(+), 52 deletions(-) create mode 100644 csrc/cpu/mamba_cpu.cpp create mode 100644 csrc/cpu/mamba_kernels.hpp diff --git a/cmake/cpu_extension.cmake b/cmake/cpu_extension.cmake index 8d74d6d5d96c..8e5337a3b166 100644 --- a/cmake/cpu_extension.cmake +++ b/cmake/cpu_extension.cmake @@ -352,6 +352,7 @@ set(VLLM_EXT_SRC "csrc/cpu/layernorm.cpp" "csrc/cpu/mla_decode.cpp" "csrc/cpu/pos_encoding.cpp" + "csrc/cpu/mamba_cpu.cpp" "csrc/moe/dynamic_4bit_int_moe_cpu.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/torch_bindings.cpp") @@ -384,6 +385,7 @@ if (ENABLE_X86_ISA) "csrc/cpu/utils.cpp" "csrc/cpu/cpu_attn.cpp" "csrc/cpu/dnnl_kernels.cpp" + "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files "csrc/cpu/activation.cpp" @@ -395,6 +397,7 @@ if (ENABLE_X86_ISA) set(VLLM_EXT_SRC_AVX2 "csrc/cpu/utils.cpp" "csrc/cpu/cpu_attn.cpp" + "csrc/cpu/mamba_cpu.cpp" "csrc/cpu/torch_bindings.cpp" # TODO: Remove these files "csrc/cpu/activation.cpp" diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp new file mode 100644 index 000000000000..721d2eb07452 --- /dev/null +++ b/csrc/cpu/mamba_cpu.cpp @@ -0,0 +1,226 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// CPU at::Tensor wrappers for Mamba decode-step kernels defined in +// mamba_kernels.hpp. + +#include "cpu/mamba_kernels.hpp" + +#include +#include +#include + +// --------------------------------------------------------------------------- +// causal_conv1d_update +// +// x : (batch, dim) [decode] or (batch, dim, seqlen) or +// (total_tokens, dim) [varlen] conv_state : (num_cache, dim, state_len) – +// updated in-place weight : (dim, width) bias : (dim,) optional +// activation : "silu" / "swish" / None / True / False +// conv_state_indices : (batch,) int32 optional +// query_start_loc : (batch+1,) int32 optional [varlen mode] +// pad_slot_id : int +// +// Returns x (overwritten with output) cast back to original dtype. +// --------------------------------------------------------------------------- +at::Tensor causal_conv1d_update_cpu_impl( + at::Tensor& x, // modified in-place (re-typed to float32) + at::Tensor& conv_state, const at::Tensor& weight, + const c10::optional& bias, + const c10::optional& activation, + const c10::optional& conv_state_indices, + const c10::optional& query_start_loc, int64_t pad_slot_id) { + // ------------------------------------------------------------------ + // Resolve activation + // ------------------------------------------------------------------ + bool do_silu = false; + if (activation.has_value()) { + const std::string& act = activation.value(); + do_silu = (act == "silu" || act == "swish"); + } + + // ------------------------------------------------------------------ + // Keep original dtype; work in float32 + // ------------------------------------------------------------------ + at::ScalarType orig_dtype = x.scalar_type(); + at::Tensor x_f32 = x.to(at::kFloat); + at::Tensor state_f32 = conv_state.to(at::kFloat); + at::Tensor w_f32 = weight.to(at::kFloat); + at::Tensor bias_f32; + bool has_bias = bias.has_value() && bias.value().defined(); + if (has_bias) bias_f32 = bias.value().to(at::kFloat); + + // ------------------------------------------------------------------ + // Dimensions + // ------------------------------------------------------------------ + bool is_2d = (x_f32.dim() == 2); + bool is_varlen = query_start_loc.has_value(); + at::Tensor x_3d; + + if (is_varlen) { + // x: (total_tokens, dim) – treat as (dim, total_tokens) then wrap + // We reshape to (1, dim, total_tokens) and iterate via query_start_loc + // For simplicity, call the PyTorch fallback for varlen path. + // (Varlen is only used during prefill-style chunked decode, not the + // critical single-token decode path.) + TORCH_CHECK(false, + "causal_conv1d_update_cpu_impl: varlen mode not yet supported " + "in C++ kernel; use the PyTorch fallback."); + } + + if (is_2d) { + x_3d = x_f32.unsqueeze(-1); // (batch, dim, 1) + } else { + x_3d = x_f32; + } + + int64_t batch = x_3d.size(0); + int64_t dim = x_3d.size(1); + int64_t seqlen = x_3d.size(2); + int64_t width = w_f32.size(1); + int64_t state_len = state_f32.size(2); + + // Ensure contiguous + x_3d = x_3d.contiguous(); + state_f32 = state_f32.contiguous(); + at::Tensor w_cont = w_f32.contiguous(); + + // Output: same shape as x_3d (written in place) + at::Tensor out = x_3d.clone(); + + const int32_t* cache_idx_ptr = nullptr; + at::Tensor cache_idx_int; + if (conv_state_indices.has_value()) { + cache_idx_int = conv_state_indices.value().to(at::kInt).contiguous(); + cache_idx_ptr = cache_idx_int.data_ptr(); + } + + mamba_cpu::causal_conv1d_update_kernel( + x_3d.data_ptr(), state_f32.data_ptr(), + w_cont.data_ptr(), + has_bias ? bias_f32.contiguous().data_ptr() : nullptr, + out.data_ptr(), cache_idx_ptr, static_cast(pad_slot_id), + batch, dim, seqlen, width, state_len, do_silu); + + // Write back updated state + conv_state.copy_(state_f32.to(conv_state.scalar_type())); + + at::Tensor out_final = is_2d ? out.squeeze(-1) : out; + return out_final.to(orig_dtype); +} + +// --------------------------------------------------------------------------- +// selective_state_update +// +// Decode-step (single token or short varlen sequence) SSM recurrence. +// +// Tensor shapes follow _selective_state_update_cuda convention but all tensors +// are on CPU. +// --------------------------------------------------------------------------- +void selective_state_update_cpu_impl( + at::Tensor& state, // (batch_or_nstates, nheads, dim, dstate) – in-place + const at::Tensor& x, // (N, nheads, dim) where N == batch + // (cu_seqlens==None) or total tokens + const at::Tensor& dt, + const at::Tensor& A, // (nheads, dim, dstate) + const at::Tensor& B, // (N, ngroups, dstate) + const at::Tensor& C, // (N, ngroups, dstate) + const c10::optional& D, // (nheads, dim) + const c10::optional& z, // (N, nheads, dim) + const c10::optional& dt_bias, // (nheads, dim) + bool dt_softplus, + const c10::optional& state_batch_indices, // (N,) int32 + const c10::optional& dst_state_batch_indices, // (N,) int32 + int64_t null_block_id, + at::Tensor& out, // (N, nheads, dim) – written + const c10::optional& num_accepted_tokens, // (N,) int32 + const c10::optional& cu_seqlens // (N+1,) int32 +) { + // ------------------------------------------------------------------ + // Work in float32 + // ------------------------------------------------------------------ + at::Tensor state_f32 = state.to(at::kFloat); + at::Tensor x_f32 = x.to(at::kFloat).contiguous(); + at::Tensor dt_f32 = dt.to(at::kFloat).contiguous(); + at::Tensor A_f32 = A.to(at::kFloat).contiguous(); + at::Tensor B_f32 = B.to(at::kFloat).contiguous(); + at::Tensor C_f32 = C.to(at::kFloat).contiguous(); + at::Tensor out_f32 = at::zeros_like(x_f32); // (N, nheads, dim) + state_f32 = state_f32.contiguous(); + + at::Tensor D_f32, z_f32, dtbias_f32; + bool has_D = D.has_value() && D.value().defined(); + bool has_z = z.has_value() && z.value().defined(); + bool has_dt_bias = dt_bias.has_value() && dt_bias.value().defined(); + if (has_D) D_f32 = D.value().to(at::kFloat).contiguous(); + if (has_z) z_f32 = z.value().to(at::kFloat).contiguous(); + if (has_dt_bias) dtbias_f32 = dt_bias.value().to(at::kFloat).contiguous(); + + // ------------------------------------------------------------------ + // Dimensions + // ------------------------------------------------------------------ + // state_f32: (nstates, nheads, dim, dstate) + int64_t nstates = state_f32.size(0); + int64_t nheads = state_f32.size(1); + int64_t dim = state_f32.size(2); + int64_t dstate = state_f32.size(3); + int64_t ngroups = B_f32.size(1); + + int64_t N; + if (cu_seqlens.has_value()) { + N = cu_seqlens.value().size(0) - 1; + } else { + N = x_f32.size(0); + } + + // Strides (all contiguous after .contiguous()) + int64_t stride_state_n = state_f32.stride(0); + int64_t stride_state_h = state_f32.stride(1); + int64_t stride_state_d = state_f32.stride(2); + int64_t stride_xdt_n = x_f32.stride(0); + int64_t stride_xdt_h = x_f32.stride(1); + int64_t stride_BC_n = B_f32.stride(0); + int64_t stride_BC_g = B_f32.stride(1); + int64_t stride_out_n = out_f32.stride(0); + int64_t stride_out_h = out_f32.stride(1); + + // Optional pointer helpers + const int32_t* sbi_ptr = nullptr; + const int32_t* dsbi_ptr = nullptr; + const int32_t* nat_ptr = nullptr; + const int32_t* csl_ptr = nullptr; + + at::Tensor sbi_int, dsbi_int, nat_int, csl_int; + if (state_batch_indices.has_value()) { + sbi_int = state_batch_indices.value().to(at::kInt).contiguous(); + sbi_ptr = sbi_int.data_ptr(); + } + if (dst_state_batch_indices.has_value()) { + dsbi_int = dst_state_batch_indices.value().to(at::kInt).contiguous(); + dsbi_ptr = dsbi_int.data_ptr(); + } + if (num_accepted_tokens.has_value()) { + nat_int = num_accepted_tokens.value().to(at::kInt).contiguous(); + nat_ptr = nat_int.data_ptr(); + } + if (cu_seqlens.has_value()) { + csl_int = cu_seqlens.value().to(at::kInt).contiguous(); + csl_ptr = csl_int.data_ptr(); + } + + mamba_cpu::selective_state_update_kernel( + state_f32.data_ptr(), stride_state_n, stride_state_h, + stride_state_d, x_f32.data_ptr(), dt_f32.data_ptr(), + stride_xdt_n, stride_xdt_h, A_f32.data_ptr(), + B_f32.data_ptr(), C_f32.data_ptr(), stride_BC_n, + stride_BC_g, has_D ? D_f32.data_ptr() : nullptr, + has_z ? z_f32.data_ptr() : nullptr, + has_dt_bias ? dtbias_f32.data_ptr() : nullptr, + out_f32.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, dsbi_ptr, + static_cast(null_block_id), nat_ptr, csl_ptr, N, nheads, ngroups, + dim, dstate, dt_softplus); + + // Write back + state.copy_(state_f32.to(state.scalar_type())); + out.copy_(out_f32.to(out.scalar_type())); +} diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp new file mode 100644 index 000000000000..59d44a2b82a2 --- /dev/null +++ b/csrc/cpu/mamba_kernels.hpp @@ -0,0 +1,261 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Fused CPU vector kernels for Mamba decode-step hotspots: +// - causal_conv1d_update (depthwise 1-D conv state roll + compute) +// - selective_state_update (SSM recurrence, single-step) +// +// The kernels operate in float32 throughout (with BF16 input/output conversion +// at the boundary) so that the compiler can auto-vectorise the inner loops +// with -O3 on any ISA (AVX2, NEON, VSX, …) without ISA-specific intrinsics. + +#pragma once + +#include +#include +#include + +namespace mamba_cpu { + +// --------------------------------------------------------------------------- +// causal_conv1d_update +// +// For each batch sequence and each feature channel d: +// +// state[cache_idx, d, :] = [state[cache_idx, d, 1:], x[b, d]] (roll) +// out[b, d] = sum_k( state[cache_idx, d, k] * weight[d, k] ) + bias[d] +// if silu: out[b, d] = out[b, d] * sigmoid(out[b, d]) +// +// Tensors (host pointers, all contiguous): +// x_ptr : float*, shape (batch, dim) or (batch, dim, seqlen) +// state_ptr : float*, shape (num_cache, dim, state_len) +// weight_ptr : float*, shape (dim, width) +// bias_ptr : float* or nullptr, shape (dim,) +// out_ptr : float*, same shape as x +// cache_idxs : int32_t*, shape (batch,) – may be nullptr (use b) +// pad_slot_id : cache index value to skip +// batch, dim, seqlen, width, state_len – dimensions +// do_silu : apply SiLU activation +// --------------------------------------------------------------------------- +inline void causal_conv1d_update_kernel( + const float* __restrict__ x_ptr, float* __restrict__ state_ptr, + const float* __restrict__ weight_ptr, const float* __restrict__ bias_ptr, + float* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, + int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen, + int64_t width, int64_t state_len, bool do_silu) { + // state layout: [num_cache, dim, state_len] + // x / out layout: [batch, dim, seqlen] (seqlen == 1 for plain decode) + for (int64_t b = 0; b < batch; ++b) { + int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b; + if (cache_idx == pad_slot_id) continue; + + for (int64_t t = 0; t < seqlen; ++t) { + const float* x_b = x_ptr + (b * dim * seqlen + t); // stride seqlen + float* out_b = out_ptr + (b * dim * seqlen + t); // stride seqlen + float* s = state_ptr + cache_idx * dim * state_len; + + for (int64_t d = 0; d < dim; ++d) { + float x_val = x_b[d * seqlen]; + + // Roll state left, append x_val at the end + float* sd = s + d * state_len; // state for channel d + // shift left: sd[0..state_len-2] = sd[1..state_len-1] + // use memmove for contiguous float array + if (state_len > 1) { + std::memmove(sd, sd + 1, (state_len - 1) * sizeof(float)); + } + sd[state_len - 1] = x_val; + + // Dot product with weight[d, :] = weight + d*width + const float* w = weight_ptr + d * width; + float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f; + // state_len == width - 1, so the window is sd[0..state_len-1] + // and weight[d, 0..width-1] aligns as: w[0]*sd[0] + ... + + // w[width-1]*sd[state_len-1] because the new x_val already sits at + // sd[state_len-1] + for (int64_t k = 0; k < width; ++k) { + // sd holds [old_1, ..., old_{width-2}, x_new] after the memmove + acc += w[k] * sd[k]; + } + + if (do_silu) { + acc = acc / (1.0f + std::exp(-acc)); + } + out_b[d * seqlen] = acc; + } + } + } +} + +// --------------------------------------------------------------------------- +// selective_state_update (SSM recurrence, single decode step OR varlen decode) +// +// Algorithm (per sequence, per head, per token): +// +// if dt_bias: dt += dt_bias +// if dt_softplus: dt = log1p(exp(dt)) +// dA = exp(A * dt) shape (nheads, dim, dstate) +// B_exp = B.repeat_interleave(nheads_per_group) (nheads, dstate) +// C_exp = C.repeat_interleave(nheads_per_group) (nheads, dstate) +// dBx = B_exp * (x * dt) shape (nheads, dim, dstate) +// state = state * dA + dBx +// out = (state * C_exp).sum(-1) +// if D: out += x * D +// if z: out *= z * sigmoid(z) +// +// All tensors are in float32 (caller converts before passing). +// Strides: state (nheads, dim, dstate), rest as 1-D slices indexed by +// (seq_idx, head, dim, dstate) as needed. +// +// Parameters match those in the Python fallback _selective_state_update_cpu +// but receive raw pointers + sizes. +// --------------------------------------------------------------------------- +inline void selective_state_update_kernel( + // state: [nstates, nheads, dim, dstate] – modified in place + float* __restrict__ state_ptr, + int64_t stride_state_n, // stride along nstates dim + int64_t stride_state_h, // stride along nheads dim + int64_t stride_state_d, // stride along dim dim + // x, dt: [N, nheads, dim] + const float* __restrict__ x_ptr, const float* __restrict__ dt_ptr, + int64_t stride_xdt_n, // stride along N dim + int64_t stride_xdt_h, // stride along nheads dim + // A: [nheads, dim, dstate] + const float* __restrict__ A_ptr, + // B, C: [N, ngroups, dstate] + const float* __restrict__ B_ptr, const float* __restrict__ C_ptr, + int64_t stride_BC_n, // stride along N dim + int64_t stride_BC_g, // stride along ngroups dim + // D: [nheads, dim] or nullptr + const float* __restrict__ D_ptr, + // z: [N, nheads, dim] or nullptr + const float* __restrict__ z_ptr, + // dt_bias: [nheads, dim] or nullptr + const float* __restrict__ dt_bias_ptr, + // out: [N, nheads, dim] – written in place + float* __restrict__ out_ptr, int64_t stride_out_n, int64_t stride_out_h, + // state_batch_indices: [N] or nullptr (use seq_idx) + const int32_t* __restrict__ state_batch_indices, + // dst_state_batch_indices: [N] or nullptr + const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id, + // num_accepted_tokens: [N] or nullptr + const int32_t* __restrict__ num_accepted_tokens, + // cu_seqlens: [N+1] or nullptr + const int32_t* __restrict__ cu_seqlens, + // dimensions + int64_t N, // number of sequences (or batch) + int64_t nheads, int64_t ngroups, int64_t dim, int64_t dstate, + bool dt_softplus) { + int64_t nheads_per_group = nheads / ngroups; + + for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) { + int64_t bos, seq_len; + if (cu_seqlens != nullptr) { + bos = cu_seqlens[seq_idx]; + seq_len = cu_seqlens[seq_idx + 1] - bos; + } else { + bos = seq_idx; + seq_len = 1; + } + + // Determine state read index + int64_t state_read_idx; + if (state_batch_indices != nullptr) { + state_read_idx = state_batch_indices[seq_idx]; + if (state_read_idx == null_block_id) continue; + } else { + state_read_idx = seq_idx; + } + + // Determine state write index + int64_t state_write_idx; + if (num_accepted_tokens == nullptr) { + if (dst_state_batch_indices != nullptr) { + state_write_idx = dst_state_batch_indices[seq_idx]; + } else { + state_write_idx = state_read_idx; + } + } else { + state_write_idx = -1; // written per-token inside the loop + } + + // Per-sequence state buffer: copy to float32 scratch + // We work directly on state_ptr (it's already float32 at this layer) + float* s = state_ptr + state_read_idx * stride_state_n; + + for (int64_t t = 0; t < seq_len; ++t) { + int64_t token_idx = bos + t; + + const float* x_tok = x_ptr + token_idx * stride_xdt_n; + const float* dt_tok = dt_ptr + token_idx * stride_xdt_n; + const float* B_tok = B_ptr + token_idx * stride_BC_n; + const float* C_tok = C_ptr + token_idx * stride_BC_n; + float* out_tok = out_ptr + token_idx * stride_out_n; + + for (int64_t h = 0; h < nheads; ++h) { + int64_t g = h / nheads_per_group; + const float* x_h = x_tok + h * stride_xdt_h; + const float* dt_h = dt_tok + h * stride_xdt_h; + const float* B_g = B_tok + g * stride_BC_g; + const float* C_g = C_tok + g * stride_BC_g; + const float* A_h = A_ptr + h * dim * dstate; + const float* dt_bias_h = + (dt_bias_ptr != nullptr) ? dt_bias_ptr + h * dim : nullptr; + const float* D_h = (D_ptr != nullptr) ? D_ptr + h * dim : nullptr; + const float* z_h = + (z_ptr != nullptr) + ? z_ptr + token_idx * stride_xdt_n + h * stride_xdt_h + : nullptr; + float* out_h = out_tok + h * stride_out_h; + float* s_h = s + h * stride_state_h; + + for (int64_t d = 0; d < dim; ++d) { + float x_val = x_h[d]; + float dt_val = dt_h[d]; + if (dt_bias_h != nullptr) dt_val += dt_bias_h[d]; + if (dt_softplus) { + // log1p(exp(dt)) — numerically stable + dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; + } + + float out_val = 0.0f; + float* s_hd = s_h + d * dstate; + const float* A_hd = A_h + d * dstate; + for (int64_t n = 0; n < dstate; ++n) { + float dA = std::exp(A_hd[n] * dt_val); + float dBx = B_g[n] * x_val * dt_val; + float s_new = s_hd[n] * dA + dBx; + s_hd[n] = s_new; + out_val += s_new * C_g[n]; + } + + if (D_h != nullptr) out_val += x_val * D_h[d]; + if (z_h != nullptr) { + float z_val = z_h[d]; + out_val *= z_val / (1.0f + std::exp(-z_val)); + } + out_h[d] = out_val; + } + } + + // Handle spec-decoding token-level dst write + if (num_accepted_tokens != nullptr && + dst_state_batch_indices != nullptr) { + int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t]; + if (token_dst_idx != null_block_id) { + float* dst_s = state_ptr + token_dst_idx * stride_state_n; + std::memcpy(dst_s, s, nheads * stride_state_h * sizeof(float)); + } + } + } + + // Write final state + if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && + state_write_idx != state_read_idx) { + float* dst_s = state_ptr + state_write_idx * stride_state_n; + std::memcpy(dst_s, s, nheads * stride_state_h * sizeof(float)); + } + } +} + +} // namespace mamba_cpu diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 15b254662f0a..be089906c2c8 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -132,6 +132,24 @@ void compute_slot_mapping_kernel_impl(const torch::Tensor query_start_loc, torch::Tensor slot_mapping, const int64_t block_size); +at::Tensor causal_conv1d_update_cpu_impl( + at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight, + const c10::optional& bias, + const c10::optional& activation, + const c10::optional& conv_state_indices, + const c10::optional& query_start_loc, int64_t pad_slot_id); + +void selective_state_update_cpu_impl( + at::Tensor& state, const at::Tensor& x, const at::Tensor& dt, + const at::Tensor& A, const at::Tensor& B, const at::Tensor& C, + const c10::optional& D, const c10::optional& z, + const c10::optional& dt_bias, bool dt_softplus, + const c10::optional& state_batch_indices, + const c10::optional& dst_state_batch_indices, + int64_t null_block_id, at::Tensor& out, + const c10::optional& num_accepted_tokens, + const c10::optional& cu_seqlens); + TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // vLLM custom ops @@ -346,6 +364,23 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "positions, Tensor block_table, Tensor(a3!) slot_mapping, SymInt " "block_size) -> ()", &compute_slot_mapping_kernel_impl); + + // Mamba CPU kernels + ops.def( + "causal_conv1d_update_cpu(" + "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, " + "Tensor? bias, str? activation, Tensor? conv_state_indices, " + "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor", + &causal_conv1d_update_cpu_impl); + + ops.def( + "selective_state_update_cpu(" + "Tensor(a0!) state, Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, " + "Tensor? D, Tensor? z, Tensor? dt_bias, bool dt_softplus, " + "Tensor? state_batch_indices, Tensor? dst_state_batch_indices, " + "SymInt null_block_id, Tensor(a13!) out, " + "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()", + &selective_state_update_cpu_impl); } REGISTER_EXTENSION(TORCH_EXTENSION_NAME) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 31def2856e7b..5f9c317c9efa 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -8,13 +8,14 @@ import numpy as np import torch -from vllm.model_executor.custom_op import CustomOp +from vllm import _custom_ops as ops from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID from .cpu_fallbacks import _causal_conv1d_fn_cpu, _causal_conv1d_update_cpu + @triton.jit() def _causal_conv1d_fwd_kernel( # continuous batching # Pointers to matrices @@ -1248,37 +1249,75 @@ def grid(META): return out.to(original_x_dtype) -@CustomOp.register("causal_conv1d_fn") -class CausalConv1dFnOp(CustomOp): - def forward_native(self, *args, **kwargs): - return _causal_conv1d_fn_cpu(*args, **kwargs) - - def forward_cpu(self, *args, **kwargs): - return _causal_conv1d_fn_cpu(*args, **kwargs) - - def forward_cuda(self, *args, **kwargs): - return _causal_conv1d_fn_cuda(*args, **kwargs) - - -@CustomOp.register("causal_conv1d_update") -class CausalConv1dUpdateOp(CustomOp): - def forward_native(self, *args, **kwargs): - return _causal_conv1d_update_cpu(*args, **kwargs) - - def forward_cpu(self, *args, **kwargs): - return _causal_conv1d_update_cpu(*args, **kwargs) - - def forward_cuda(self, *args, **kwargs): - return _causal_conv1d_update_cuda(*args, **kwargs) - - -_causal_conv1d_fn_op = CausalConv1dFnOp() -_causal_conv1d_update_op = CausalConv1dUpdateOp() - - def causal_conv1d_fn(*args, **kwargs): - return _causal_conv1d_fn_op(*args, **kwargs) - + """Dispatch causal_conv1d_fn to CPU PyTorch fallback or CUDA Triton kernel.""" + x = args[0] if args else kwargs.get("x") + if x is not None and x.device.type == "cpu": + return _causal_conv1d_fn_cpu(*args, **kwargs) + return _causal_conv1d_fn_cuda(*args, **kwargs) + + +def causal_conv1d_update( + x, + conv_state, + weight, + bias=None, + activation=None, + conv_state_indices=None, + num_accepted_tokens=None, + query_start_loc=None, + max_query_len=-1, + null_block_id=NULL_BLOCK_ID, + block_idx_last_scheduled_token=None, + initial_state_idx=None, + validate_data=False, +): + """Dispatch causal_conv1d_update to CPU C++ kernel or CUDA Triton kernel.""" + if x.device.type == "cpu": + # The C++ kernel handles the standard (non-varlen, non-spec-decoding) + # decode path. Fall back to PyTorch for the complex varlen / + # spec-decoding paths. + if query_start_loc is not None or num_accepted_tokens is not None: + return _causal_conv1d_update_cpu( + x, + conv_state, + weight, + bias=bias, + activation=activation, + conv_state_indices=conv_state_indices, + query_start_loc=query_start_loc, + ) + # Determine activation string + act_str = None + if isinstance(activation, bool): + act_str = "silu" if activation else None + elif activation is not None: + act_str = activation + + pad_slot_id = int(NULL_BLOCK_ID) + return ops.causal_conv1d_update_cpu( + x, + conv_state, + weight, + bias, + act_str, + conv_state_indices, + None, # query_start_loc + pad_slot_id, + ) -def causal_conv1d_update(*args, **kwargs): - return _causal_conv1d_update_op(*args, **kwargs) + return _causal_conv1d_update_cuda( + x, + conv_state, + weight, + bias=bias, + activation=activation, + conv_state_indices=conv_state_indices, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=query_start_loc, + max_query_len=max_query_len, + null_block_id=null_block_id, + block_idx_last_scheduled_token=block_idx_last_scheduled_token, + initial_state_idx=initial_state_idx, + validate_data=validate_data, + ) diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index 62e626cd080e..9745bbe84d03 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -8,13 +8,10 @@ from packaging import version from vllm import _custom_ops as ops -from vllm.model_executor.custom_op import CustomOp from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp from vllm.triton_utils import HAS_TRITON, tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID -from .cpu_fallbacks import _selective_state_update_cpu - TRITON3 = HAS_TRITON and (version.parse(triton.__version__) >= version.parse("3.0.0")) if TRITON3: @@ -685,20 +682,74 @@ def selective_scan_fn( return z # output written inplace to z -@CustomOp.register("selective_state_update") -class SelectiveStateUpdateOp(CustomOp): - def forward_native(self, *args, **kwargs): - return _selective_state_update_cpu(*args, **kwargs) - - def forward_cpu(self, *args, **kwargs): - return _selective_state_update_cpu(*args, **kwargs) - - def forward_cuda(self, *args, **kwargs): - return _selective_state_update_cuda(*args, **kwargs) - -_selective_state_update_op = SelectiveStateUpdateOp() - - -def selective_state_update(*args, **kwargs): - return _selective_state_update_op(*args, **kwargs) +def selective_state_update( + state, + x, + dt, + A, + B, + C, + D=None, + z=None, + dt_bias=None, + dt_softplus=False, + state_batch_indices=None, + dst_state_batch_indices=None, + null_block_id=NULL_BLOCK_ID, + out=None, + num_accepted_tokens=None, + cu_seqlens=None, + is_blackwell=False, + enable_stochastic_rounding=False, + cache_philox_rounds=0, +): + """Dispatch selective_state_update to CPU C++ kernel or CUDA Triton kernel.""" + # Ensure out tensor exists + if out is None: + out = torch.empty_like(x if x.dim() == 2 else x) + + if x.device.type == "cpu": + # Reshape tensors from (batch, dim) -> (batch, 1, dim) if needed + # The C++ kernel expects (N, nheads, dim) layout + _state = state.unsqueeze(1) if state.dim() == 3 else state + _x = x.unsqueeze(1) if x.dim() == 2 else x + _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt + _A = A.unsqueeze(0) if A.dim() == 2 else A + _B = B.unsqueeze(1) if B.dim() == 2 else B + _C = C.unsqueeze(1) if C.dim() == 2 else C + _D = D.unsqueeze(0) if (D is not None and D.dim() == 1) else D + _z = z.unsqueeze(1) if (z is not None and z.dim() == 2) else z + _dt_bias = ( + dt_bias.unsqueeze(0) + if (dt_bias is not None and dt_bias.dim() == 1) + else dt_bias + ) + _out = out.unsqueeze(1) if out.dim() == 2 else out + # state_batch_indices and dst_state_batch_indices are 1D index arrays; + # do NOT reshape them. + _sbi = state_batch_indices + _dsbi = dst_state_batch_indices + ops.selective_state_update_cpu( + _state, _x, _dt, _A, _B, _C, + _D, _z, _dt_bias, dt_softplus, + _sbi, _dsbi, + null_block_id, _out, + num_accepted_tokens, cu_seqlens, + ) + return _out.squeeze(1) if out.dim() == 2 else _out + + return _selective_state_update_cuda( + state, x, dt, A, B, C, + D=D, z=z, dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + dst_state_batch_indices=dst_state_batch_indices, + null_block_id=null_block_id, + out=out, + num_accepted_tokens=num_accepted_tokens, + cu_seqlens=cu_seqlens, + is_blackwell=is_blackwell, + enable_stochastic_rounding=enable_stochastic_rounding, + cache_philox_rounds=cache_philox_rounds, + ) From 6b8010775c15475827fe67803615581ad3105e71 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 5 May 2026 12:21:59 +0530 Subject: [PATCH 11/43] fix utils.py Signed-off-by: Akash kaothalkar --- vllm/model_executor/layers/utils.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 47da445745ae..a8e7b47dc3c3 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -236,9 +236,7 @@ def dispatch_cpu_unquantized_gemm( # Skip CPU GEMM dispatch for non-2D weights (e.g. MoE 3D expert weights). # These layers are handled by their own specialized methods. if layer.weight.dim() != 2: - layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear( - x, weight, bias - ) + layer.cpu_linear = torch.nn.functional.linear return N, K = layer.weight.size() @@ -298,9 +296,7 @@ def dispatch_cpu_unquantized_gemm( ) # fallback case - layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear( - x, weight, bias - ) + layer.cpu_linear = torch.nn.functional.linear def cpu_unquantized_gemm( From 9c37ab53d1e8a0e8b7509ec23eab920dfc5317f4 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 5 May 2026 15:15:50 +0530 Subject: [PATCH 12/43] fix Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 10 +++- csrc/cpu/mamba_kernels.hpp | 39 ++++++------ vllm/_custom_ops.py | 60 +++++++++++++++++++ .../layers/fused_moe/cpu_fused_moe.py | 10 +++- .../layers/mamba/ops/causal_conv1d.py | 11 +++- 5 files changed, 104 insertions(+), 26 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 721d2eb07452..98e706881827 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -183,6 +183,10 @@ void selective_state_update_cpu_impl( int64_t stride_BC_g = B_f32.stride(1); int64_t stride_out_n = out_f32.stride(0); int64_t stride_out_h = out_f32.stride(1); + int64_t stride_A_h = (A_f32.size(0) == 1) ? 0 : A_f32.stride(0); + int64_t stride_D_h = has_D ? ((D_f32.size(0) == 1) ? 0 : D_f32.stride(0)) : 0; + int64_t stride_dtbias_h = + has_dt_bias ? ((dtbias_f32.size(0) == 1) ? 0 : dtbias_f32.stride(0)) : 0; // Optional pointer helpers const int32_t* sbi_ptr = nullptr; @@ -211,11 +215,11 @@ void selective_state_update_cpu_impl( mamba_cpu::selective_state_update_kernel( state_f32.data_ptr(), stride_state_n, stride_state_h, stride_state_d, x_f32.data_ptr(), dt_f32.data_ptr(), - stride_xdt_n, stride_xdt_h, A_f32.data_ptr(), + stride_xdt_n, stride_xdt_h, A_f32.data_ptr(), stride_A_h, B_f32.data_ptr(), C_f32.data_ptr(), stride_BC_n, - stride_BC_g, has_D ? D_f32.data_ptr() : nullptr, + stride_BC_g, has_D ? D_f32.data_ptr() : nullptr, stride_D_h, has_z ? z_f32.data_ptr() : nullptr, - has_dt_bias ? dtbias_f32.data_ptr() : nullptr, + has_dt_bias ? dtbias_f32.data_ptr() : nullptr, stride_dtbias_h, out_f32.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, dsbi_ptr, static_cast(null_block_id), nat_ptr, csl_ptr, N, nheads, ngroups, dim, dstate, dt_softplus); diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index 59d44a2b82a2..0508749a089d 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -57,26 +57,25 @@ inline void causal_conv1d_update_kernel( for (int64_t d = 0; d < dim; ++d) { float x_val = x_b[d * seqlen]; - // Roll state left, append x_val at the end float* sd = s + d * state_len; // state for channel d - // shift left: sd[0..state_len-2] = sd[1..state_len-1] - // use memmove for contiguous float array - if (state_len > 1) { - std::memmove(sd, sd + 1, (state_len - 1) * sizeof(float)); - } - sd[state_len - 1] = x_val; - // Dot product with weight[d, :] = weight + d*width + // Compute the dot product with weight[d, :] = weight + d*width const float* w = weight_ptr + d * width; float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f; - // state_len == width - 1, so the window is sd[0..state_len-1] - // and weight[d, 0..width-1] aligns as: w[0]*sd[0] + ... + - // w[width-1]*sd[state_len-1] because the new x_val already sits at - // sd[state_len-1] - for (int64_t k = 0; k < width; ++k) { - // sd holds [old_1, ..., old_{width-2}, x_new] after the memmove + + // The convolution window is [sd[0], ..., sd[state_len-1], x_val] + for (int64_t k = 0; k < state_len; ++k) { acc += w[k] * sd[k]; } + acc += w[state_len] * x_val; + + // Roll state left, append x_val at the end + if (state_len > 1) { + std::memmove(sd, sd + 1, (state_len - 1) * sizeof(float)); + } + if (state_len > 0) { + sd[state_len - 1] = x_val; + } if (do_silu) { acc = acc / (1.0f + std::exp(-acc)); @@ -121,17 +120,17 @@ inline void selective_state_update_kernel( int64_t stride_xdt_n, // stride along N dim int64_t stride_xdt_h, // stride along nheads dim // A: [nheads, dim, dstate] - const float* __restrict__ A_ptr, + const float* __restrict__ A_ptr, int64_t stride_A_h, // B, C: [N, ngroups, dstate] const float* __restrict__ B_ptr, const float* __restrict__ C_ptr, int64_t stride_BC_n, // stride along N dim int64_t stride_BC_g, // stride along ngroups dim // D: [nheads, dim] or nullptr - const float* __restrict__ D_ptr, + const float* __restrict__ D_ptr, int64_t stride_D_h, // z: [N, nheads, dim] or nullptr const float* __restrict__ z_ptr, // dt_bias: [nheads, dim] or nullptr - const float* __restrict__ dt_bias_ptr, + const float* __restrict__ dt_bias_ptr, int64_t stride_dtbias_h, // out: [N, nheads, dim] – written in place float* __restrict__ out_ptr, int64_t stride_out_n, int64_t stride_out_h, // state_batch_indices: [N] or nullptr (use seq_idx) @@ -198,10 +197,10 @@ inline void selective_state_update_kernel( const float* dt_h = dt_tok + h * stride_xdt_h; const float* B_g = B_tok + g * stride_BC_g; const float* C_g = C_tok + g * stride_BC_g; - const float* A_h = A_ptr + h * dim * dstate; + const float* A_h = A_ptr + h * stride_A_h; const float* dt_bias_h = - (dt_bias_ptr != nullptr) ? dt_bias_ptr + h * dim : nullptr; - const float* D_h = (D_ptr != nullptr) ? D_ptr + h * dim : nullptr; + (dt_bias_ptr != nullptr) ? dt_bias_ptr + h * stride_dtbias_h : nullptr; + const float* D_h = (D_ptr != nullptr) ? D_ptr + h * stride_D_h : nullptr; const float* z_h = (z_ptr != nullptr) ? z_ptr + token_idx * stride_xdt_n + h * stride_xdt_h diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index ffea35de8811..f23f69a5c34c 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2305,6 +2305,66 @@ def selective_scan_fwd( ) +def causal_conv1d_update_cpu( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: str | None = None, + conv_state_indices: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + pad_slot_id: int = 0, +) -> torch.Tensor: + return torch.ops._C.causal_conv1d_update_cpu( + x, + conv_state, + weight, + bias, + activation, + conv_state_indices, + query_start_loc, + pad_slot_id, + ) + + +def selective_state_update_cpu( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None, + z: torch.Tensor | None, + dt_bias: torch.Tensor | None, + dt_softplus: bool, + state_batch_indices: torch.Tensor | None, + dst_state_batch_indices: torch.Tensor | None, + null_block_id: int, + out: torch.Tensor, + num_accepted_tokens: torch.Tensor | None, + cu_seqlens: torch.Tensor | None, +): + torch.ops._C.selective_state_update_cpu( + state, + x, + dt, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + dst_state_batch_indices, + null_block_id, + out, + num_accepted_tokens, + cu_seqlens, + ) + + # ROCm skinny gemms def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor: return torch.ops._rocm_C.LLMM1(a, b, rows_per_block) diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 985f33e10098..3609d350da73 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -8,7 +8,6 @@ from vllm import _custom_ops as ops from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight -from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.quantization.utils.layer_utils import replace_parameter from vllm.utils.torch_utils import direct_register_custom_op @@ -41,11 +40,18 @@ def _gelu_and_mul( return F.gelu(x[..., :d], approximate="none") * x[..., d:] +def _silu_and_mul(x: torch.Tensor) -> torch.Tensor: + """Standalone SiluAndMul forward to avoid instantiating CustomOp + (which calls get_current_vllm_config()) at model-forward time.""" + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] + + # Map activation names to their native forward functions. # Uses static methods or standalone functions to avoid instantiating CustomOp # classes, which would call get_current_vllm_config() before config is set. _CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = { - MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x), + MoEActivation.SILU: _silu_and_mul, MoEActivation.SWIGLUOAI: _swigluoai_forward_native, MoEActivation.GELU: _gelu_and_mul, } diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 61193727ac3b..fad11e67384d 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -1290,6 +1290,15 @@ def causal_conv1d_update( elif activation is not None: act_str = activation + _conv_state_indices = conv_state_indices + if _conv_state_indices is not None and _conv_state_indices.dim() == 2: + if initial_state_idx is not None: + _conv_state_indices = _conv_state_indices.gather( + 1, initial_state_idx.unsqueeze(1) + ).squeeze(1) + else: + _conv_state_indices = _conv_state_indices[:, 0] + pad_slot_id = int(NULL_BLOCK_ID) return ops.causal_conv1d_update_cpu( x, @@ -1297,7 +1306,7 @@ def causal_conv1d_update( weight, bias, act_str, - conv_state_indices, + _conv_state_indices, None, # query_start_loc pad_slot_id, ) From 32f8ded6f6694933f420317a6307fb6c0bb618c3 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 5 May 2026 15:44:55 +0530 Subject: [PATCH 13/43] fix mamba kernels Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 53 +++++++++--------- csrc/cpu/mamba_kernels.hpp | 38 +++++++------ scratch/verify_cpu_ssu.py | 107 +++++++++++++++++++++++++++++++++++++ 3 files changed, 153 insertions(+), 45 deletions(-) create mode 100644 scratch/verify_cpu_ssu.py diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 98e706881827..4ce855e47495 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -137,16 +137,15 @@ void selective_state_update_cpu_impl( const c10::optional& cu_seqlens // (N+1,) int32 ) { // ------------------------------------------------------------------ - // Work in float32 + // Dimensions and Setup // ------------------------------------------------------------------ - at::Tensor state_f32 = state.to(at::kFloat); + // Work in float32 for parameters, but allow state to be in original dtype at::Tensor x_f32 = x.to(at::kFloat).contiguous(); at::Tensor dt_f32 = dt.to(at::kFloat).contiguous(); at::Tensor A_f32 = A.to(at::kFloat).contiguous(); at::Tensor B_f32 = B.to(at::kFloat).contiguous(); at::Tensor C_f32 = C.to(at::kFloat).contiguous(); at::Tensor out_f32 = at::zeros_like(x_f32); // (N, nheads, dim) - state_f32 = state_f32.contiguous(); at::Tensor D_f32, z_f32, dtbias_f32; bool has_D = D.has_value() && D.value().defined(); @@ -156,14 +155,9 @@ void selective_state_update_cpu_impl( if (has_z) z_f32 = z.value().to(at::kFloat).contiguous(); if (has_dt_bias) dtbias_f32 = dt_bias.value().to(at::kFloat).contiguous(); - // ------------------------------------------------------------------ - // Dimensions - // ------------------------------------------------------------------ - // state_f32: (nstates, nheads, dim, dstate) - int64_t nstates = state_f32.size(0); - int64_t nheads = state_f32.size(1); - int64_t dim = state_f32.size(2); - int64_t dstate = state_f32.size(3); + int64_t nheads = state.size(1); + int64_t dim = state.size(2); + int64_t dstate = state.size(3); int64_t ngroups = B_f32.size(1); int64_t N; @@ -173,10 +167,10 @@ void selective_state_update_cpu_impl( N = x_f32.size(0); } - // Strides (all contiguous after .contiguous()) - int64_t stride_state_n = state_f32.stride(0); - int64_t stride_state_h = state_f32.stride(1); - int64_t stride_state_d = state_f32.stride(2); + // Strides + int64_t stride_state_n = state.stride(0); + int64_t stride_state_h = state.stride(1); + int64_t stride_state_d = state.stride(2); int64_t stride_xdt_n = x_f32.stride(0); int64_t stride_xdt_h = x_f32.stride(1); int64_t stride_BC_n = B_f32.stride(0); @@ -212,19 +206,20 @@ void selective_state_update_cpu_impl( csl_ptr = csl_int.data_ptr(); } - mamba_cpu::selective_state_update_kernel( - state_f32.data_ptr(), stride_state_n, stride_state_h, - stride_state_d, x_f32.data_ptr(), dt_f32.data_ptr(), - stride_xdt_n, stride_xdt_h, A_f32.data_ptr(), stride_A_h, - B_f32.data_ptr(), C_f32.data_ptr(), stride_BC_n, - stride_BC_g, has_D ? D_f32.data_ptr() : nullptr, stride_D_h, - has_z ? z_f32.data_ptr() : nullptr, - has_dt_bias ? dtbias_f32.data_ptr() : nullptr, stride_dtbias_h, - out_f32.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, dsbi_ptr, - static_cast(null_block_id), nat_ptr, csl_ptr, N, nheads, ngroups, - dim, dstate, dt_softplus); - - // Write back - state.copy_(state_f32.to(state.scalar_type())); + VLLM_DISPATCH_FLOATING_TYPES(state.scalar_type(), "selective_state_update_kernel", [&] { + mamba_cpu::selective_state_update_kernel( + state.data_ptr(), stride_state_n, stride_state_h, + stride_state_d, x_f32.data_ptr(), dt_f32.data_ptr(), + stride_xdt_n, stride_xdt_h, A_f32.data_ptr(), stride_A_h, + B_f32.data_ptr(), C_f32.data_ptr(), stride_BC_n, + stride_BC_g, has_D ? D_f32.data_ptr() : nullptr, stride_D_h, + has_z ? z_f32.data_ptr() : nullptr, + has_dt_bias ? dtbias_f32.data_ptr() : nullptr, stride_dtbias_h, + out_f32.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, dsbi_ptr, + static_cast(null_block_id), nat_ptr, csl_ptr, N, nheads, ngroups, + dim, dstate, dt_softplus); + }); + out.copy_(out_f32.to(out.scalar_type())); } +} diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index 0508749a089d..c3bb2aef49f1 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -109,9 +109,10 @@ inline void causal_conv1d_update_kernel( // Parameters match those in the Python fallback _selective_state_update_cpu // but receive raw pointers + sizes. // --------------------------------------------------------------------------- +template inline void selective_state_update_kernel( // state: [nstates, nheads, dim, dstate] – modified in place - float* __restrict__ state_ptr, + scalar_t* __restrict__ state_ptr, int64_t stride_state_n, // stride along nstates dim int64_t stride_state_h, // stride along nheads dim int64_t stride_state_d, // stride along dim dim @@ -178,9 +179,8 @@ inline void selective_state_update_kernel( state_write_idx = -1; // written per-token inside the loop } - // Per-sequence state buffer: copy to float32 scratch - // We work directly on state_ptr (it's already float32 at this layer) - float* s = state_ptr + state_read_idx * stride_state_n; + // Per-sequence state buffer + scalar_t* s = state_ptr + state_read_idx * stride_state_n; for (int64_t t = 0; t < seq_len; ++t) { int64_t token_idx = bos + t; @@ -191,6 +191,7 @@ inline void selective_state_update_kernel( const float* C_tok = C_ptr + token_idx * stride_BC_n; float* out_tok = out_ptr + token_idx * stride_out_n; +#pragma omp parallel for for (int64_t h = 0; h < nheads; ++h) { int64_t g = h / nheads_per_group; const float* x_h = x_tok + h * stride_xdt_h; @@ -206,7 +207,7 @@ inline void selective_state_update_kernel( ? z_ptr + token_idx * stride_xdt_n + h * stride_xdt_h : nullptr; float* out_h = out_tok + h * stride_out_h; - float* s_h = s + h * stride_state_h; + scalar_t* s_h = s + h * stride_state_h; for (int64_t d = 0; d < dim; ++d) { float x_val = x_h[d]; @@ -214,24 +215,29 @@ inline void selective_state_update_kernel( if (dt_bias_h != nullptr) dt_val += dt_bias_h[d]; if (dt_softplus) { // log1p(exp(dt)) — numerically stable - dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; + dt_val = (dt_val <= 20.0f) ? std::log1pf(std::expf(dt_val)) : dt_val; } float out_val = 0.0f; - float* s_hd = s_h + d * dstate; + scalar_t* s_hd = s_h + d * dstate; const float* A_hd = A_h + d * dstate; for (int64_t n = 0; n < dstate; ++n) { - float dA = std::exp(A_hd[n] * dt_val); + float dA = std::expf(A_hd[n] * dt_val); float dBx = B_g[n] * x_val * dt_val; - float s_new = s_hd[n] * dA + dBx; - s_hd[n] = s_new; + float s_old = static_cast(s_hd[n]); + float s_new = s_old * dA + dBx; + s_hd[n] = static_cast(s_new); out_val += s_new * C_g[n]; } if (D_h != nullptr) out_val += x_val * D_h[d]; if (z_h != nullptr) { float z_val = z_h[d]; - out_val *= z_val / (1.0f + std::exp(-z_val)); + // Stable SiLU: z * sigmoid(z) + float sigmoid_z = (z_val >= 0) ? + 1.0f / (1.0f + std::expf(-z_val)) : + std::expf(z_val) / (1.0f + std::expf(z_val)); + out_val *= z_val * sigmoid_z; } out_h[d] = out_val; } @@ -241,9 +247,9 @@ inline void selective_state_update_kernel( if (num_accepted_tokens != nullptr && dst_state_batch_indices != nullptr) { int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t]; - if (token_dst_idx != null_block_id) { - float* dst_s = state_ptr + token_dst_idx * stride_state_n; - std::memcpy(dst_s, s, nheads * stride_state_h * sizeof(float)); + if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) { + scalar_t* dst_s = state_ptr + token_dst_idx * stride_state_n; + std::memmove(dst_s, s, nheads * stride_state_h * sizeof(scalar_t)); } } } @@ -251,8 +257,8 @@ inline void selective_state_update_kernel( // Write final state if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && state_write_idx != state_read_idx) { - float* dst_s = state_ptr + state_write_idx * stride_state_n; - std::memcpy(dst_s, s, nheads * stride_state_h * sizeof(float)); + scalar_t* dst_s = state_ptr + state_write_idx * stride_state_n; + std::memmove(dst_s, s, nheads * stride_state_h * sizeof(scalar_t)); } } } diff --git a/scratch/verify_cpu_ssu.py b/scratch/verify_cpu_ssu.py new file mode 100644 index 000000000000..c5ff0fe276bc --- /dev/null +++ b/scratch/verify_cpu_ssu.py @@ -0,0 +1,107 @@ +import torch +import torch.nn.functional as F +from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update +import time + +def selective_state_update_ref( + state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False +): + has_heads = state.dim() > 3 + if state.dim() == 3: + state = state.unsqueeze(1) + if x.dim() == 2: + x = x.unsqueeze(1) + if dt.dim() == 2: + dt = dt.unsqueeze(1) + if A.dim() == 2: + A = A.unsqueeze(0) + if B.dim() == 2: + B = B.unsqueeze(1) + if C.dim() == 2: + C = C.unsqueeze(1) + if D is not None and D.dim() == 1: + D = D.unsqueeze(0) + if z is not None and z.dim() == 2: + z = z.unsqueeze(1) + if dt_bias is not None and dt_bias.dim() == 1: + dt_bias = dt_bias.unsqueeze(0) + batch, nheads, dim, dstate = state.shape + + dt = dt.float() + if dt_bias is not None: + dt = dt + dt_bias.float() + dt = F.softplus(dt) if dt_softplus else dt + + dA = torch.exp(dt.unsqueeze(-1) * A.float()) # (batch, nheads, dim, dstate) + ngroups = B.shape[1] + B = B.repeat_interleave(nheads // ngroups, dim=1) + C = C.repeat_interleave(nheads // ngroups, dim=1) + + dBx = (dt.unsqueeze(-1) * B.float().unsqueeze(2)) * x.float().unsqueeze(-1) + + state.copy_(state.float() * dA + dBx) + + out = torch.einsum("bhdn,bhn->bhd", state.float(), C.float()) + if D is not None: + out += (x.float() * D.float()) + + if z is not None: + out = out * F.silu(z.float()) + + if not has_heads: + out = out.squeeze(1) + return out.to(x.dtype) + +def test_cpu_ssu(): + device = "cpu" + torch.manual_seed(42) + + # Mamba-2 style dimensions + batch_size = 2 + nheads = 128 + dim = 64 + dstate = 128 + itype = torch.bfloat16 + + print(f"Testing CPU selective_state_update with {itype}...") + + state = torch.randn(batch_size, nheads, dim, dstate, dtype=itype, device=device) + x = torch.randn(batch_size, nheads, dim, device=device, dtype=itype) + dt = torch.randn(batch_size, nheads, dim, device=device, dtype=itype) + dt_bias = torch.rand(nheads, dim, device=device) - 4.0 + A = -torch.rand(nheads, dim, dstate, device=device) - 1.0 + B = torch.randn(batch_size, 1, dstate, device=device, dtype=itype) + C = torch.randn(batch_size, 1, dstate, device=device, dtype=itype) + D = torch.randn(nheads, dim, device=device) + z = torch.randn_like(x) + + out = torch.empty_like(x) + state_ref = state.detach().clone() + + start = time.time() + selective_state_update( + state, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, dt_softplus=True, out=out + ) + end = time.time() + print(f"Kernel time: {end - start:.4f}s") + + start = time.time() + out_ref = selective_state_update_ref( + state_ref, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, dt_softplus=True + ) + end = time.time() + print(f"Ref time: {end - start:.4f}s") + + state_diff = (state.float() - state_ref.float()).abs().max().item() + out_diff = (out.float() - out_ref.float()).abs().max().item() + + print(f"State max diff: {state_diff}") + print(f"Out max diff: {out_diff}") + + if state_diff < 1e-2 and out_diff < 1e-2: + print("SUCCESS") + else: + print("FAILURE") + +if __name__ == "__main__": + test_cpu_ssu() From dde32c9dfff889b0dca37b9446ffdc12b32d8e00 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 5 May 2026 15:51:54 +0530 Subject: [PATCH 14/43] use cpu_types Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 3 +- csrc/cpu/mamba_kernels.hpp | 179 ++++++++++++------------------------- 2 files changed, 61 insertions(+), 121 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 4ce855e47495..a409fa81f500 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -10,6 +10,8 @@ #include #include +#include "cpu_types.hpp" + // --------------------------------------------------------------------------- // causal_conv1d_update // @@ -222,4 +224,3 @@ void selective_state_update_cpu_impl( out.copy_(out_f32.to(out.scalar_type())); } -} diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index c3bb2aef49f1..42c6d9d77378 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -4,38 +4,20 @@ // Fused CPU vector kernels for Mamba decode-step hotspots: // - causal_conv1d_update (depthwise 1-D conv state roll + compute) // - selective_state_update (SSM recurrence, single-step) -// -// The kernels operate in float32 throughout (with BF16 input/output conversion -// at the boundary) so that the compiler can auto-vectorise the inner loops -// with -O3 on any ISA (AVX2, NEON, VSX, …) without ISA-specific intrinsics. #pragma once +#include "cpu_types.hpp" #include #include #include +#include +#include namespace mamba_cpu { // --------------------------------------------------------------------------- // causal_conv1d_update -// -// For each batch sequence and each feature channel d: -// -// state[cache_idx, d, :] = [state[cache_idx, d, 1:], x[b, d]] (roll) -// out[b, d] = sum_k( state[cache_idx, d, k] * weight[d, k] ) + bias[d] -// if silu: out[b, d] = out[b, d] * sigmoid(out[b, d]) -// -// Tensors (host pointers, all contiguous): -// x_ptr : float*, shape (batch, dim) or (batch, dim, seqlen) -// state_ptr : float*, shape (num_cache, dim, state_len) -// weight_ptr : float*, shape (dim, width) -// bias_ptr : float* or nullptr, shape (dim,) -// out_ptr : float*, same shape as x -// cache_idxs : int32_t*, shape (batch,) – may be nullptr (use b) -// pad_slot_id : cache index value to skip -// batch, dim, seqlen, width, state_len – dimensions -// do_silu : apply SiLU activation // --------------------------------------------------------------------------- inline void causal_conv1d_update_kernel( const float* __restrict__ x_ptr, float* __restrict__ state_ptr, @@ -43,33 +25,29 @@ inline void causal_conv1d_update_kernel( float* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen, int64_t width, int64_t state_len, bool do_silu) { - // state layout: [num_cache, dim, state_len] - // x / out layout: [batch, dim, seqlen] (seqlen == 1 for plain decode) + +#pragma omp parallel for for (int64_t b = 0; b < batch; ++b) { int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b; if (cache_idx == pad_slot_id) continue; for (int64_t t = 0; t < seqlen; ++t) { - const float* x_b = x_ptr + (b * dim * seqlen + t); // stride seqlen - float* out_b = out_ptr + (b * dim * seqlen + t); // stride seqlen + const float* x_b = x_ptr + (b * dim * seqlen + t); + float* out_b = out_ptr + (b * dim * seqlen + t); float* s = state_ptr + cache_idx * dim * state_len; for (int64_t d = 0; d < dim; ++d) { float x_val = x_b[d * seqlen]; + float* sd = s + d * state_len; - float* sd = s + d * state_len; // state for channel d - - // Compute the dot product with weight[d, :] = weight + d*width const float* w = weight_ptr + d * width; float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f; - // The convolution window is [sd[0], ..., sd[state_len-1], x_val] for (int64_t k = 0; k < state_len; ++k) { acc += w[k] * sd[k]; } acc += w[state_len] * x_val; - // Roll state left, append x_val at the end if (state_len > 1) { std::memmove(sd, sd + 1, (state_len - 1) * sizeof(float)); } @@ -78,7 +56,10 @@ inline void causal_conv1d_update_kernel( } if (do_silu) { - acc = acc / (1.0f + std::exp(-acc)); + float sigmoid = (acc >= 0) ? + 1.0f / (1.0f + std::exp(-acc)) : + std::exp(acc) / (1.0f + std::exp(acc)); + acc *= sigmoid; } out_b[d * seqlen] = acc; } @@ -87,65 +68,31 @@ inline void causal_conv1d_update_kernel( } // --------------------------------------------------------------------------- -// selective_state_update (SSM recurrence, single decode step OR varlen decode) -// -// Algorithm (per sequence, per head, per token): -// -// if dt_bias: dt += dt_bias -// if dt_softplus: dt = log1p(exp(dt)) -// dA = exp(A * dt) shape (nheads, dim, dstate) -// B_exp = B.repeat_interleave(nheads_per_group) (nheads, dstate) -// C_exp = C.repeat_interleave(nheads_per_group) (nheads, dstate) -// dBx = B_exp * (x * dt) shape (nheads, dim, dstate) -// state = state * dA + dBx -// out = (state * C_exp).sum(-1) -// if D: out += x * D -// if z: out *= z * sigmoid(z) -// -// All tensors are in float32 (caller converts before passing). -// Strides: state (nheads, dim, dstate), rest as 1-D slices indexed by -// (seq_idx, head, dim, dstate) as needed. -// -// Parameters match those in the Python fallback _selective_state_update_cpu -// but receive raw pointers + sizes. +// selective_state_update // --------------------------------------------------------------------------- template inline void selective_state_update_kernel( - // state: [nstates, nheads, dim, dstate] – modified in place scalar_t* __restrict__ state_ptr, - int64_t stride_state_n, // stride along nstates dim - int64_t stride_state_h, // stride along nheads dim - int64_t stride_state_d, // stride along dim dim - // x, dt: [N, nheads, dim] + int64_t stride_state_n, int64_t stride_state_h, int64_t stride_state_d, const float* __restrict__ x_ptr, const float* __restrict__ dt_ptr, - int64_t stride_xdt_n, // stride along N dim - int64_t stride_xdt_h, // stride along nheads dim - // A: [nheads, dim, dstate] + int64_t stride_xdt_n, int64_t stride_xdt_h, const float* __restrict__ A_ptr, int64_t stride_A_h, - // B, C: [N, ngroups, dstate] const float* __restrict__ B_ptr, const float* __restrict__ C_ptr, - int64_t stride_BC_n, // stride along N dim - int64_t stride_BC_g, // stride along ngroups dim - // D: [nheads, dim] or nullptr + int64_t stride_BC_n, int64_t stride_BC_g, const float* __restrict__ D_ptr, int64_t stride_D_h, - // z: [N, nheads, dim] or nullptr const float* __restrict__ z_ptr, - // dt_bias: [nheads, dim] or nullptr const float* __restrict__ dt_bias_ptr, int64_t stride_dtbias_h, - // out: [N, nheads, dim] – written in place float* __restrict__ out_ptr, int64_t stride_out_n, int64_t stride_out_h, - // state_batch_indices: [N] or nullptr (use seq_idx) const int32_t* __restrict__ state_batch_indices, - // dst_state_batch_indices: [N] or nullptr const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id, - // num_accepted_tokens: [N] or nullptr const int32_t* __restrict__ num_accepted_tokens, - // cu_seqlens: [N+1] or nullptr const int32_t* __restrict__ cu_seqlens, - // dimensions - int64_t N, // number of sequences (or batch) - int64_t nheads, int64_t ngroups, int64_t dim, int64_t dstate, + int64_t N, int64_t nheads, int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) { + + using scalar_vec_t = vec_op::vec_t; + constexpr int VEC_ELEM_NUM = 8; + int64_t nheads_per_group = nheads / ngroups; for (int64_t seq_idx = 0; seq_idx < N; ++seq_idx) { @@ -158,33 +105,17 @@ inline void selective_state_update_kernel( seq_len = 1; } - // Determine state read index - int64_t state_read_idx; - if (state_batch_indices != nullptr) { - state_read_idx = state_batch_indices[seq_idx]; - if (state_read_idx == null_block_id) continue; - } else { - state_read_idx = seq_idx; - } + int64_t state_read_idx = (state_batch_indices != nullptr) ? + state_batch_indices[seq_idx] : seq_idx; + if (state_read_idx == null_block_id) continue; - // Determine state write index - int64_t state_write_idx; - if (num_accepted_tokens == nullptr) { - if (dst_state_batch_indices != nullptr) { - state_write_idx = dst_state_batch_indices[seq_idx]; - } else { - state_write_idx = state_read_idx; - } - } else { - state_write_idx = -1; // written per-token inside the loop - } + int64_t state_write_idx = (num_accepted_tokens == nullptr) ? + ((dst_state_batch_indices != nullptr) ? dst_state_batch_indices[seq_idx] : state_read_idx) : -1; - // Per-sequence state buffer scalar_t* s = state_ptr + state_read_idx * stride_state_n; for (int64_t t = 0; t < seq_len; ++t) { int64_t token_idx = bos + t; - const float* x_tok = x_ptr + token_idx * stride_xdt_n; const float* dt_tok = dt_ptr + token_idx * stride_xdt_n; const float* B_tok = B_ptr + token_idx * stride_BC_n; @@ -199,13 +130,9 @@ inline void selective_state_update_kernel( const float* B_g = B_tok + g * stride_BC_g; const float* C_g = C_tok + g * stride_BC_g; const float* A_h = A_ptr + h * stride_A_h; - const float* dt_bias_h = - (dt_bias_ptr != nullptr) ? dt_bias_ptr + h * stride_dtbias_h : nullptr; + const float* dt_bias_h = (dt_bias_ptr != nullptr) ? dt_bias_ptr + h * stride_dtbias_h : nullptr; const float* D_h = (D_ptr != nullptr) ? D_ptr + h * stride_D_h : nullptr; - const float* z_h = - (z_ptr != nullptr) - ? z_ptr + token_idx * stride_xdt_n + h * stride_xdt_h - : nullptr; + const float* z_h = (z_ptr != nullptr) ? z_ptr + token_idx * stride_xdt_n + h * stride_xdt_h : nullptr; float* out_h = out_tok + h * stride_out_h; scalar_t* s_h = s + h * stride_state_h; @@ -214,18 +141,37 @@ inline void selective_state_update_kernel( float dt_val = dt_h[d]; if (dt_bias_h != nullptr) dt_val += dt_bias_h[d]; if (dt_softplus) { - // log1p(exp(dt)) — numerically stable - dt_val = (dt_val <= 20.0f) ? std::log1pf(std::expf(dt_val)) : dt_val; + dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; } - float out_val = 0.0f; + vec_op::FP32Vec8 out_vec(0.0f); scalar_t* s_hd = s_h + d * dstate; const float* A_hd = A_h + d * dstate; - for (int64_t n = 0; n < dstate; ++n) { - float dA = std::expf(A_hd[n] * dt_val); + + vec_op::FP32Vec8 x_vec(x_val); + vec_op::FP32Vec8 dt_vec(dt_val); + + int64_t n = 0; + for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { + vec_op::FP32Vec8 A_v(A_hd + n); + vec_op::FP32Vec8 B_v(B_g + n); + vec_op::FP32Vec8 C_v(C_g + n); + + vec_op::FP32Vec8 s_v((scalar_vec_t(s_hd + n))); + + vec_op::FP32Vec8 dA = (A_v * dt_vec).exp(); + vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec; + vec_op::FP32Vec8 s_new = s_v * dA + dBx; + + scalar_vec_t(s_new).save(s_hd + n); + out_vec = out_vec + s_new * C_v; + } + + float out_val = out_vec.reduce_sum(); + for (; n < dstate; ++n) { + float dA = std::exp(A_hd[n] * dt_val); float dBx = B_g[n] * x_val * dt_val; - float s_old = static_cast(s_hd[n]); - float s_new = s_old * dA + dBx; + float s_new = static_cast(s_hd[n]) * dA + dBx; s_hd[n] = static_cast(s_new); out_val += s_new * C_g[n]; } @@ -233,19 +179,14 @@ inline void selective_state_update_kernel( if (D_h != nullptr) out_val += x_val * D_h[d]; if (z_h != nullptr) { float z_val = z_h[d]; - // Stable SiLU: z * sigmoid(z) - float sigmoid_z = (z_val >= 0) ? - 1.0f / (1.0f + std::expf(-z_val)) : - std::expf(z_val) / (1.0f + std::expf(z_val)); - out_val *= z_val * sigmoid_z; + float sigmoid = (z_val >= 0) ? 1.0f / (1.0f + std::exp(-z_val)) : std::exp(z_val) / (1.0f + std::exp(z_val)); + out_val *= z_val * sigmoid; } out_h[d] = out_val; } } - // Handle spec-decoding token-level dst write - if (num_accepted_tokens != nullptr && - dst_state_batch_indices != nullptr) { + if (num_accepted_tokens != nullptr && dst_state_batch_indices != nullptr) { int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t]; if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) { scalar_t* dst_s = state_ptr + token_dst_idx * stride_state_n; @@ -254,13 +195,11 @@ inline void selective_state_update_kernel( } } - // Write final state - if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && - state_write_idx != state_read_idx) { + if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && state_write_idx != state_read_idx) { scalar_t* dst_s = state_ptr + state_write_idx * stride_state_n; std::memmove(dst_s, s, nheads * stride_state_h * sizeof(scalar_t)); } } } -} // namespace mamba_cpu +} // namespace mamba_cpu From 41294d4de6316d378657c5b778a65bbfecfe391e Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 5 May 2026 16:02:56 +0530 Subject: [PATCH 15/43] optimize float loading Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 230 ++++++++++++------------------------- csrc/cpu/mamba_kernels.hpp | 87 +++++++------- 2 files changed, 115 insertions(+), 202 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index a409fa81f500..fb010f0916c3 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -14,16 +14,6 @@ // --------------------------------------------------------------------------- // causal_conv1d_update -// -// x : (batch, dim) [decode] or (batch, dim, seqlen) or -// (total_tokens, dim) [varlen] conv_state : (num_cache, dim, state_len) – -// updated in-place weight : (dim, width) bias : (dim,) optional -// activation : "silu" / "swish" / None / True / False -// conv_state_indices : (batch,) int32 optional -// query_start_loc : (batch+1,) int32 optional [varlen mode] -// pad_slot_id : int -// -// Returns x (overwritten with output) cast back to original dtype. // --------------------------------------------------------------------------- at::Tensor causal_conv1d_update_cpu_impl( at::Tensor& x, // modified in-place (re-typed to float32) @@ -32,63 +22,27 @@ at::Tensor causal_conv1d_update_cpu_impl( const c10::optional& activation, const c10::optional& conv_state_indices, const c10::optional& query_start_loc, int64_t pad_slot_id) { - // ------------------------------------------------------------------ - // Resolve activation - // ------------------------------------------------------------------ + bool do_silu = false; if (activation.has_value()) { const std::string& act = activation.value(); do_silu = (act == "silu" || act == "swish"); } - // ------------------------------------------------------------------ - // Keep original dtype; work in float32 - // ------------------------------------------------------------------ - at::ScalarType orig_dtype = x.scalar_type(); - at::Tensor x_f32 = x.to(at::kFloat); - at::Tensor state_f32 = conv_state.to(at::kFloat); - at::Tensor w_f32 = weight.to(at::kFloat); + // Causal conv still works in float32 for now (minimal overhead compared to SSM) + at::Tensor x_f32 = x.to(at::kFloat).contiguous(); + at::Tensor state_f32 = conv_state.to(at::kFloat).contiguous(); + at::Tensor w_f32 = weight.to(at::kFloat).contiguous(); at::Tensor bias_f32; - bool has_bias = bias.has_value() && bias.value().defined(); - if (has_bias) bias_f32 = bias.value().to(at::kFloat); - - // ------------------------------------------------------------------ - // Dimensions - // ------------------------------------------------------------------ - bool is_2d = (x_f32.dim() == 2); - bool is_varlen = query_start_loc.has_value(); - at::Tensor x_3d; - - if (is_varlen) { - // x: (total_tokens, dim) – treat as (dim, total_tokens) then wrap - // We reshape to (1, dim, total_tokens) and iterate via query_start_loc - // For simplicity, call the PyTorch fallback for varlen path. - // (Varlen is only used during prefill-style chunked decode, not the - // critical single-token decode path.) - TORCH_CHECK(false, - "causal_conv1d_update_cpu_impl: varlen mode not yet supported " - "in C++ kernel; use the PyTorch fallback."); - } - - if (is_2d) { - x_3d = x_f32.unsqueeze(-1); // (batch, dim, 1) - } else { - x_3d = x_f32; - } + if (bias.has_value() && bias.value().defined()) bias_f32 = bias.value().to(at::kFloat).contiguous(); - int64_t batch = x_3d.size(0); - int64_t dim = x_3d.size(1); - int64_t seqlen = x_3d.size(2); + int64_t batch = x_f32.size(0); + int64_t dim = x_f32.size(1); + int64_t seqlen = (x_f32.dim() == 3) ? x_f32.size(2) : 1; int64_t width = w_f32.size(1); int64_t state_len = state_f32.size(2); - // Ensure contiguous - x_3d = x_3d.contiguous(); - state_f32 = state_f32.contiguous(); - at::Tensor w_cont = w_f32.contiguous(); - - // Output: same shape as x_3d (written in place) - at::Tensor out = x_3d.clone(); + at::Tensor out_f32 = at::empty_like(x_f32); const int32_t* cache_idx_ptr = nullptr; at::Tensor cache_idx_int; @@ -98,129 +52,87 @@ at::Tensor causal_conv1d_update_cpu_impl( } mamba_cpu::causal_conv1d_update_kernel( - x_3d.data_ptr(), state_f32.data_ptr(), - w_cont.data_ptr(), - has_bias ? bias_f32.contiguous().data_ptr() : nullptr, - out.data_ptr(), cache_idx_ptr, static_cast(pad_slot_id), + x_f32.data_ptr(), state_f32.data_ptr(), + w_f32.data_ptr(), bias_f32.defined() ? bias_f32.data_ptr() : nullptr, + out_f32.data_ptr(), cache_idx_ptr, static_cast(pad_slot_id), batch, dim, seqlen, width, state_len, do_silu); - // Write back updated state conv_state.copy_(state_f32.to(conv_state.scalar_type())); - - at::Tensor out_final = is_2d ? out.squeeze(-1) : out; - return out_final.to(orig_dtype); + return out_f32.to(x.scalar_type()); } // --------------------------------------------------------------------------- // selective_state_update -// -// Decode-step (single token or short varlen sequence) SSM recurrence. -// -// Tensor shapes follow _selective_state_update_cuda convention but all tensors -// are on CPU. // --------------------------------------------------------------------------- void selective_state_update_cpu_impl( - at::Tensor& state, // (batch_or_nstates, nheads, dim, dstate) – in-place - const at::Tensor& x, // (N, nheads, dim) where N == batch - // (cu_seqlens==None) or total tokens + at::Tensor& state, // (nstates, nheads, dim, dstate) + const at::Tensor& x, // (N, nheads, dim) const at::Tensor& dt, - const at::Tensor& A, // (nheads, dim, dstate) - const at::Tensor& B, // (N, ngroups, dstate) - const at::Tensor& C, // (N, ngroups, dstate) - const c10::optional& D, // (nheads, dim) - const c10::optional& z, // (N, nheads, dim) - const c10::optional& dt_bias, // (nheads, dim) + const at::Tensor& A, + const at::Tensor& B, + const at::Tensor& C, + const c10::optional& D, + const c10::optional& z, + const c10::optional& dt_bias, bool dt_softplus, - const c10::optional& state_batch_indices, // (N,) int32 - const c10::optional& dst_state_batch_indices, // (N,) int32 + const c10::optional& state_batch_indices, + const c10::optional& dst_state_batch_indices, int64_t null_block_id, - at::Tensor& out, // (N, nheads, dim) – written - const c10::optional& num_accepted_tokens, // (N,) int32 - const c10::optional& cu_seqlens // (N+1,) int32 + at::Tensor& out, + const c10::optional& num_accepted_tokens, + const c10::optional& cu_seqlens ) { - // ------------------------------------------------------------------ - // Dimensions and Setup - // ------------------------------------------------------------------ - // Work in float32 for parameters, but allow state to be in original dtype - at::Tensor x_f32 = x.to(at::kFloat).contiguous(); - at::Tensor dt_f32 = dt.to(at::kFloat).contiguous(); - at::Tensor A_f32 = A.to(at::kFloat).contiguous(); - at::Tensor B_f32 = B.to(at::kFloat).contiguous(); - at::Tensor C_f32 = C.to(at::kFloat).contiguous(); - at::Tensor out_f32 = at::zeros_like(x_f32); // (N, nheads, dim) - - at::Tensor D_f32, z_f32, dtbias_f32; - bool has_D = D.has_value() && D.value().defined(); - bool has_z = z.has_value() && z.value().defined(); - bool has_dt_bias = dt_bias.has_value() && dt_bias.value().defined(); - if (has_D) D_f32 = D.value().to(at::kFloat).contiguous(); - if (has_z) z_f32 = z.value().to(at::kFloat).contiguous(); - if (has_dt_bias) dtbias_f32 = dt_bias.value().to(at::kFloat).contiguous(); - + // Optimization: No full-tensor conversions to float32 here. + // We pass original dtypes and let the kernel handle it via vec_op. + int64_t nheads = state.size(1); int64_t dim = state.size(2); int64_t dstate = state.size(3); - int64_t ngroups = B_f32.size(1); + int64_t N = x.size(0); + int64_t ngroups = B.size(1); - int64_t N; - if (cu_seqlens.has_value()) { - N = cu_seqlens.value().size(0) - 1; - } else { - N = x_f32.size(0); - } - - // Strides + // Strides (crucial for correctness) int64_t stride_state_n = state.stride(0); int64_t stride_state_h = state.stride(1); int64_t stride_state_d = state.stride(2); - int64_t stride_xdt_n = x_f32.stride(0); - int64_t stride_xdt_h = x_f32.stride(1); - int64_t stride_BC_n = B_f32.stride(0); - int64_t stride_BC_g = B_f32.stride(1); - int64_t stride_out_n = out_f32.stride(0); - int64_t stride_out_h = out_f32.stride(1); - int64_t stride_A_h = (A_f32.size(0) == 1) ? 0 : A_f32.stride(0); - int64_t stride_D_h = has_D ? ((D_f32.size(0) == 1) ? 0 : D_f32.stride(0)) : 0; - int64_t stride_dtbias_h = - has_dt_bias ? ((dtbias_f32.size(0) == 1) ? 0 : dtbias_f32.stride(0)) : 0; - - // Optional pointer helpers - const int32_t* sbi_ptr = nullptr; - const int32_t* dsbi_ptr = nullptr; - const int32_t* nat_ptr = nullptr; - const int32_t* csl_ptr = nullptr; - - at::Tensor sbi_int, dsbi_int, nat_int, csl_int; - if (state_batch_indices.has_value()) { - sbi_int = state_batch_indices.value().to(at::kInt).contiguous(); - sbi_ptr = sbi_int.data_ptr(); - } - if (dst_state_batch_indices.has_value()) { - dsbi_int = dst_state_batch_indices.value().to(at::kInt).contiguous(); - dsbi_ptr = dsbi_int.data_ptr(); - } - if (num_accepted_tokens.has_value()) { - nat_int = num_accepted_tokens.value().to(at::kInt).contiguous(); - nat_ptr = nat_int.data_ptr(); - } - if (cu_seqlens.has_value()) { - csl_int = cu_seqlens.value().to(at::kInt).contiguous(); - csl_ptr = csl_int.data_ptr(); - } - - VLLM_DISPATCH_FLOATING_TYPES(state.scalar_type(), "selective_state_update_kernel", [&] { - mamba_cpu::selective_state_update_kernel( - state.data_ptr(), stride_state_n, stride_state_h, - stride_state_d, x_f32.data_ptr(), dt_f32.data_ptr(), - stride_xdt_n, stride_xdt_h, A_f32.data_ptr(), stride_A_h, - B_f32.data_ptr(), C_f32.data_ptr(), stride_BC_n, - stride_BC_g, has_D ? D_f32.data_ptr() : nullptr, stride_D_h, - has_z ? z_f32.data_ptr() : nullptr, - has_dt_bias ? dtbias_f32.data_ptr() : nullptr, stride_dtbias_h, - out_f32.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, dsbi_ptr, - static_cast(null_block_id), nat_ptr, csl_ptr, N, nheads, ngroups, - dim, dstate, dt_softplus); + int64_t stride_xdt_n = x.stride(0); + int64_t stride_xdt_h = x.stride(1); + int64_t stride_A_h = A.stride(0); + int64_t stride_BC_n = B.stride(0); + int64_t stride_BC_g = B.stride(1); + int64_t stride_out_n = out.stride(0); + int64_t stride_out_h = out.stride(1); + int64_t stride_D_h = D.has_value() ? D.value().stride(0) : 0; + int64_t stride_dtbias_h = dt_bias.has_value() ? dt_bias.value().stride(0) : 0; + + // Optional pointers + const int32_t* sbi_ptr = state_batch_indices.has_value() ? state_batch_indices.value().data_ptr() : nullptr; + const int32_t* dsbi_ptr = dst_state_batch_indices.has_value() ? dst_state_batch_indices.value().data_ptr() : nullptr; + const int32_t* nat_ptr = num_accepted_tokens.has_value() ? num_accepted_tokens.value().data_ptr() : nullptr; + const int32_t* csl_ptr = cu_seqlens.has_value() ? cu_seqlens.value().data_ptr() : nullptr; + + // out is often bfloat16, but kernel math is float32. + // We use a temporary float32 buffer per token if needed, or just convert at the end. + // For simplicity and correctness, we use a float32 out buffer and copy back. + at::Tensor out_f32 = at::empty_like(out, at::kFloat); + + VLLM_DISPATCH_FLOATING_TYPES(state.scalar_type(), "ssu_state", [&] { + using state_t = scalar_t; + VLLM_DISPATCH_FLOATING_TYPES(x.scalar_type(), "ssu_input", [&] { + using input_t = scalar_t; + mamba_cpu::selective_state_update_kernel( + state.data_ptr(), stride_state_n, stride_state_h, stride_state_d, + x.data_ptr(), dt.data_ptr(), stride_xdt_n, stride_xdt_h, + A.data_ptr(), stride_A_h, + B.data_ptr(), C.data_ptr(), stride_BC_n, stride_BC_g, + D.has_value() ? D.value().data_ptr() : nullptr, stride_D_h, + z.has_value() ? z.value().data_ptr() : nullptr, + dt_bias.has_value() ? dt_bias.value().data_ptr() : nullptr, stride_dtbias_h, + out_f32.data_ptr(), stride_out_n, stride_out_h, + sbi_ptr, dsbi_ptr, static_cast(null_block_id), + nat_ptr, csl_ptr, N, nheads, ngroups, dim, dstate, dt_softplus); + }); }); - out.copy_(out_f32.to(out.scalar_type())); + out.copy_(out_f32); } diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index 42c6d9d77378..f351f9acade8 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -70,18 +70,18 @@ inline void causal_conv1d_update_kernel( // --------------------------------------------------------------------------- // selective_state_update // --------------------------------------------------------------------------- -template +template inline void selective_state_update_kernel( - scalar_t* __restrict__ state_ptr, + state_t* __restrict__ state_ptr, int64_t stride_state_n, int64_t stride_state_h, int64_t stride_state_d, - const float* __restrict__ x_ptr, const float* __restrict__ dt_ptr, + const input_t* __restrict__ x_ptr, const input_t* __restrict__ dt_ptr, int64_t stride_xdt_n, int64_t stride_xdt_h, - const float* __restrict__ A_ptr, int64_t stride_A_h, - const float* __restrict__ B_ptr, const float* __restrict__ C_ptr, + const input_t* __restrict__ A_ptr, int64_t stride_A_h, + const input_t* __restrict__ B_ptr, const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g, - const float* __restrict__ D_ptr, int64_t stride_D_h, - const float* __restrict__ z_ptr, - const float* __restrict__ dt_bias_ptr, int64_t stride_dtbias_h, + const input_t* __restrict__ D_ptr, int64_t stride_D_h, + const input_t* __restrict__ z_ptr, + const input_t* __restrict__ dt_bias_ptr, int64_t stride_dtbias_h, float* __restrict__ out_ptr, int64_t stride_out_n, int64_t stride_out_h, const int32_t* __restrict__ state_batch_indices, const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id, @@ -90,7 +90,8 @@ inline void selective_state_update_kernel( int64_t N, int64_t nheads, int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) { - using scalar_vec_t = vec_op::vec_t; + using state_vec_t = vec_op::vec_t; + using input_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = 8; int64_t nheads_per_group = nheads / ngroups; @@ -112,73 +113,73 @@ inline void selective_state_update_kernel( int64_t state_write_idx = (num_accepted_tokens == nullptr) ? ((dst_state_batch_indices != nullptr) ? dst_state_batch_indices[seq_idx] : state_read_idx) : -1; - scalar_t* s = state_ptr + state_read_idx * stride_state_n; + state_t* s = state_ptr + state_read_idx * stride_state_n; for (int64_t t = 0; t < seq_len; ++t) { int64_t token_idx = bos + t; - const float* x_tok = x_ptr + token_idx * stride_xdt_n; - const float* dt_tok = dt_ptr + token_idx * stride_xdt_n; - const float* B_tok = B_ptr + token_idx * stride_BC_n; - const float* C_tok = C_ptr + token_idx * stride_BC_n; + const input_t* x_tok = x_ptr + token_idx * stride_xdt_n; + const input_t* dt_tok = dt_ptr + token_idx * stride_xdt_n; + const input_t* B_tok = B_ptr + token_idx * stride_BC_n; + const input_t* C_tok = C_ptr + token_idx * stride_BC_n; float* out_tok = out_ptr + token_idx * stride_out_n; #pragma omp parallel for for (int64_t h = 0; h < nheads; ++h) { int64_t g = h / nheads_per_group; - const float* x_h = x_tok + h * stride_xdt_h; - const float* dt_h = dt_tok + h * stride_xdt_h; - const float* B_g = B_tok + g * stride_BC_g; - const float* C_g = C_tok + g * stride_BC_g; - const float* A_h = A_ptr + h * stride_A_h; - const float* dt_bias_h = (dt_bias_ptr != nullptr) ? dt_bias_ptr + h * stride_dtbias_h : nullptr; - const float* D_h = (D_ptr != nullptr) ? D_ptr + h * stride_D_h : nullptr; - const float* z_h = (z_ptr != nullptr) ? z_ptr + token_idx * stride_xdt_n + h * stride_xdt_h : nullptr; + const input_t* x_h = x_tok + h * stride_xdt_h; + const input_t* dt_h = dt_tok + h * stride_xdt_h; + const input_t* B_g = B_tok + g * stride_BC_g; + const input_t* C_g = C_tok + g * stride_BC_g; + const input_t* A_h = A_ptr + h * stride_A_h; + const input_t* dt_bias_h = (dt_bias_ptr != nullptr) ? dt_bias_ptr + h * stride_dtbias_h : nullptr; + const input_t* D_h = (D_ptr != nullptr) ? D_ptr + h * stride_D_h : nullptr; + const input_t* z_h = (z_ptr != nullptr) ? z_ptr + token_idx * stride_xdt_n + h * stride_xdt_h : nullptr; float* out_h = out_tok + h * stride_out_h; - scalar_t* s_h = s + h * stride_state_h; + state_t* s_h = s + h * stride_state_h; for (int64_t d = 0; d < dim; ++d) { - float x_val = x_h[d]; - float dt_val = dt_h[d]; - if (dt_bias_h != nullptr) dt_val += dt_bias_h[d]; + float x_val = static_cast(x_h[d]); + float dt_val = static_cast(dt_h[d]); + if (dt_bias_h != nullptr) dt_val += static_cast(dt_bias_h[d]); if (dt_softplus) { dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; } vec_op::FP32Vec8 out_vec(0.0f); - scalar_t* s_hd = s_h + d * dstate; - const float* A_hd = A_h + d * dstate; + state_t* s_hd = s_h + d * stride_state_d; + const input_t* A_hd = A_h + d * dstate; vec_op::FP32Vec8 x_vec(x_val); vec_op::FP32Vec8 dt_vec(dt_val); int64_t n = 0; for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { - vec_op::FP32Vec8 A_v(A_hd + n); - vec_op::FP32Vec8 B_v(B_g + n); - vec_op::FP32Vec8 C_v(C_g + n); + vec_op::FP32Vec8 A_v((input_vec_t(A_hd + n))); + vec_op::FP32Vec8 B_v((input_vec_t(B_g + n))); + vec_op::FP32Vec8 C_v((input_vec_t(C_g + n))); - vec_op::FP32Vec8 s_v((scalar_vec_t(s_hd + n))); + vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n))); vec_op::FP32Vec8 dA = (A_v * dt_vec).exp(); vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec; vec_op::FP32Vec8 s_new = s_v * dA + dBx; - scalar_vec_t(s_new).save(s_hd + n); + state_vec_t(s_new).save(s_hd + n); out_vec = out_vec + s_new * C_v; } float out_val = out_vec.reduce_sum(); for (; n < dstate; ++n) { - float dA = std::exp(A_hd[n] * dt_val); - float dBx = B_g[n] * x_val * dt_val; + float dA = std::exp(static_cast(A_hd[n]) * dt_val); + float dBx = static_cast(B_g[n]) * x_val * dt_val; float s_new = static_cast(s_hd[n]) * dA + dBx; - s_hd[n] = static_cast(s_new); - out_val += s_new * C_g[n]; + s_hd[n] = static_cast(s_new); + out_val += s_new * static_cast(C_g[n]); } - if (D_h != nullptr) out_val += x_val * D_h[d]; + if (D_h != nullptr) out_val += x_val * static_cast(D_h[d]); if (z_h != nullptr) { - float z_val = z_h[d]; + float z_val = static_cast(z_h[d]); float sigmoid = (z_val >= 0) ? 1.0f / (1.0f + std::exp(-z_val)) : std::exp(z_val) / (1.0f + std::exp(z_val)); out_val *= z_val * sigmoid; } @@ -189,15 +190,15 @@ inline void selective_state_update_kernel( if (num_accepted_tokens != nullptr && dst_state_batch_indices != nullptr) { int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t]; if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) { - scalar_t* dst_s = state_ptr + token_dst_idx * stride_state_n; - std::memmove(dst_s, s, nheads * stride_state_h * sizeof(scalar_t)); + state_t* dst_s = state_ptr + token_dst_idx * stride_state_n; + std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t)); } } } if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && state_write_idx != state_read_idx) { - scalar_t* dst_s = state_ptr + state_write_idx * stride_state_n; - std::memmove(dst_s, s, nheads * stride_state_h * sizeof(scalar_t)); + state_t* dst_s = state_ptr + state_write_idx * stride_state_n; + std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t)); } } } From 484dd4003cd1683a7b29c398032fdde028a7ef5d Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 6 May 2026 14:56:27 +0530 Subject: [PATCH 16/43] fix bf16 issues Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 75 +++++++++++++++++++++++++----------------- 1 file changed, 44 insertions(+), 31 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index fb010f0916c3..996d1b5a5189 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -82,28 +82,43 @@ void selective_state_update_cpu_impl( const c10::optional& num_accepted_tokens, const c10::optional& cu_seqlens ) { - // Optimization: No full-tensor conversions to float32 here. - // We pass original dtypes and let the kernel handle it via vec_op. + // To avoid "expected scalar type but found" errors, we must ensure all + // input/parameter tensors match the input_t we dispatch on. + at::ScalarType input_type = x.scalar_type(); - int64_t nheads = state.size(1); - int64_t dim = state.size(2); - int64_t dstate = state.size(3); + // Convert state to input_type to ensure type consistency + at::Tensor state_in = state.to(input_type); + + // No-op if already matching the correct type. + at::Tensor dt_in = dt.to(input_type); + at::Tensor A_in = A.to(input_type); + at::Tensor B_in = B.to(input_type); + at::Tensor C_in = C.to(input_type); + + at::Tensor D_in, z_in, dt_bias_in; + if (D.has_value() && D.value().defined()) D_in = D.value().to(input_type); + if (z.has_value() && z.value().defined()) z_in = z.value().to(input_type); + if (dt_bias.has_value() && dt_bias.value().defined()) dt_bias_in = dt_bias.value().to(input_type); + + int64_t nheads = state_in.size(1); + int64_t dim = state_in.size(2); + int64_t dstate = state_in.size(3); int64_t N = x.size(0); - int64_t ngroups = B.size(1); + int64_t ngroups = B_in.size(1); - // Strides (crucial for correctness) - int64_t stride_state_n = state.stride(0); - int64_t stride_state_h = state.stride(1); - int64_t stride_state_d = state.stride(2); + // Strides + int64_t stride_state_n = state_in.stride(0); + int64_t stride_state_h = state_in.stride(1); + int64_t stride_state_d = state_in.stride(2); int64_t stride_xdt_n = x.stride(0); int64_t stride_xdt_h = x.stride(1); - int64_t stride_A_h = A.stride(0); - int64_t stride_BC_n = B.stride(0); - int64_t stride_BC_g = B.stride(1); + int64_t stride_A_h = A_in.stride(0); + int64_t stride_BC_n = B_in.stride(0); + int64_t stride_BC_g = B_in.stride(1); int64_t stride_out_n = out.stride(0); int64_t stride_out_h = out.stride(1); - int64_t stride_D_h = D.has_value() ? D.value().stride(0) : 0; - int64_t stride_dtbias_h = dt_bias.has_value() ? dt_bias.value().stride(0) : 0; + int64_t stride_D_h = D_in.defined() ? D_in.stride(0) : 0; + int64_t stride_dtbias_h = dt_bias_in.defined() ? dt_bias_in.stride(0) : 0; // Optional pointers const int32_t* sbi_ptr = state_batch_indices.has_value() ? state_batch_indices.value().data_ptr() : nullptr; @@ -111,28 +126,26 @@ void selective_state_update_cpu_impl( const int32_t* nat_ptr = num_accepted_tokens.has_value() ? num_accepted_tokens.value().data_ptr() : nullptr; const int32_t* csl_ptr = cu_seqlens.has_value() ? cu_seqlens.value().data_ptr() : nullptr; - // out is often bfloat16, but kernel math is float32. - // We use a temporary float32 buffer per token if needed, or just convert at the end. - // For simplicity and correctness, we use a float32 out buffer and copy back. + // We write to a temporary float32 buffer to ensure high-precision accumulation + // before copying back to the original out tensor. at::Tensor out_f32 = at::empty_like(out, at::kFloat); - VLLM_DISPATCH_FLOATING_TYPES(state.scalar_type(), "ssu_state", [&] { - using state_t = scalar_t; - VLLM_DISPATCH_FLOATING_TYPES(x.scalar_type(), "ssu_input", [&] { - using input_t = scalar_t; - mamba_cpu::selective_state_update_kernel( - state.data_ptr(), stride_state_n, stride_state_h, stride_state_d, - x.data_ptr(), dt.data_ptr(), stride_xdt_n, stride_xdt_h, - A.data_ptr(), stride_A_h, - B.data_ptr(), C.data_ptr(), stride_BC_n, stride_BC_g, - D.has_value() ? D.value().data_ptr() : nullptr, stride_D_h, - z.has_value() ? z.value().data_ptr() : nullptr, - dt_bias.has_value() ? dt_bias.value().data_ptr() : nullptr, stride_dtbias_h, + VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] { + using input_t = scalar_t; + mamba_cpu::selective_state_update_kernel( + state_in.data_ptr(), stride_state_n, stride_state_h, stride_state_d, + x.data_ptr(), dt_in.data_ptr(), stride_xdt_n, stride_xdt_h, + A_in.data_ptr(), stride_A_h, + B_in.data_ptr(), C_in.data_ptr(), stride_BC_n, stride_BC_g, + D_in.defined() ? D_in.data_ptr() : nullptr, stride_D_h, + z_in.defined() ? z_in.data_ptr() : nullptr, + dt_bias_in.defined() ? dt_bias_in.data_ptr() : nullptr, stride_dtbias_h, out_f32.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, dsbi_ptr, static_cast(null_block_id), nat_ptr, csl_ptr, N, nheads, ngroups, dim, dstate, dt_softplus); - }); }); + // Copy the updated state back to the original state tensor + state.copy_(state_in); out.copy_(out_f32); } From 0a4dd5c9fbdc346efbfe42ed126c2c542138bb83 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 6 May 2026 15:53:15 +0530 Subject: [PATCH 17/43] fix bf16 to fp32 Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 90 ++++++++++++++++++++++++++---------------- 1 file changed, 55 insertions(+), 35 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 996d1b5a5189..716da80c315f 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -82,34 +82,45 @@ void selective_state_update_cpu_impl( const c10::optional& num_accepted_tokens, const c10::optional& cu_seqlens ) { - // To avoid "expected scalar type but found" errors, we must ensure all - // input/parameter tensors match the input_t we dispatch on. + // Use state's dtype as the primary type to avoid expensive conversions + // The kernel supports mixed types: state_t can be BFloat16 while input_t matches x + at::ScalarType state_type = state.scalar_type(); at::ScalarType input_type = x.scalar_type(); - // Convert state to input_type to ensure type consistency - at::Tensor state_in = state.to(input_type); + // Only convert/contiguous if needed to minimize overhead + auto ensure_type_and_contiguous = [input_type](const at::Tensor& t) -> at::Tensor { + if (t.scalar_type() != input_type) { + return t.to(input_type).contiguous(); + } + return t.is_contiguous() ? t : t.contiguous(); + }; - // No-op if already matching the correct type. - at::Tensor dt_in = dt.to(input_type); - at::Tensor A_in = A.to(input_type); - at::Tensor B_in = B.to(input_type); - at::Tensor C_in = C.to(input_type); + at::Tensor dt_in = ensure_type_and_contiguous(dt); + at::Tensor A_in = ensure_type_and_contiguous(A); + at::Tensor B_in = ensure_type_and_contiguous(B); + at::Tensor C_in = ensure_type_and_contiguous(C); at::Tensor D_in, z_in, dt_bias_in; - if (D.has_value() && D.value().defined()) D_in = D.value().to(input_type); - if (z.has_value() && z.value().defined()) z_in = z.value().to(input_type); - if (dt_bias.has_value() && dt_bias.value().defined()) dt_bias_in = dt_bias.value().to(input_type); + if (D.has_value() && D.value().defined()) { + D_in = ensure_type_and_contiguous(D.value()); + } + if (z.has_value() && z.value().defined()) { + z_in = ensure_type_and_contiguous(z.value()); + } + if (dt_bias.has_value() && dt_bias.value().defined()) { + dt_bias_in = ensure_type_and_contiguous(dt_bias.value()); + } - int64_t nheads = state_in.size(1); - int64_t dim = state_in.size(2); - int64_t dstate = state_in.size(3); + int64_t nheads = state.size(1); + int64_t dim = state.size(2); + int64_t dstate = state.size(3); int64_t N = x.size(0); int64_t ngroups = B_in.size(1); // Strides - int64_t stride_state_n = state_in.stride(0); - int64_t stride_state_h = state_in.stride(1); - int64_t stride_state_d = state_in.stride(2); + int64_t stride_state_n = state.stride(0); + int64_t stride_state_h = state.stride(1); + int64_t stride_state_d = state.stride(2); int64_t stride_xdt_n = x.stride(0); int64_t stride_xdt_h = x.stride(1); int64_t stride_A_h = A_in.stride(0); @@ -120,20 +131,27 @@ void selective_state_update_cpu_impl( int64_t stride_D_h = D_in.defined() ? D_in.stride(0) : 0; int64_t stride_dtbias_h = dt_bias_in.defined() ? dt_bias_in.stride(0) : 0; - // Optional pointers - const int32_t* sbi_ptr = state_batch_indices.has_value() ? state_batch_indices.value().data_ptr() : nullptr; - const int32_t* dsbi_ptr = dst_state_batch_indices.has_value() ? dst_state_batch_indices.value().data_ptr() : nullptr; - const int32_t* nat_ptr = num_accepted_tokens.has_value() ? num_accepted_tokens.value().data_ptr() : nullptr; - const int32_t* csl_ptr = cu_seqlens.has_value() ? cu_seqlens.value().data_ptr() : nullptr; - - // We write to a temporary float32 buffer to ensure high-precision accumulation - // before copying back to the original out tensor. - at::Tensor out_f32 = at::empty_like(out, at::kFloat); - - VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] { - using input_t = scalar_t; - mamba_cpu::selective_state_update_kernel( - state_in.data_ptr(), stride_state_n, stride_state_h, stride_state_d, + // Optional pointers - extract once + auto get_int32_ptr = [](const c10::optional& opt) -> const int32_t* { + return (opt.has_value() && opt.value().defined()) ? opt.value().data_ptr() : nullptr; + }; + + const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices); + const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices); + const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens); + const int32_t* csl_ptr = get_int32_ptr(cu_seqlens); + + // Optimize output buffer: only use float32 if output type is not already float32 + // This avoids an extra copy when out is already float32 + bool need_out_conversion = (out.scalar_type() != at::kFloat); + at::Tensor out_f32 = need_out_conversion ? at::empty_like(out, at::kFloat) : out; + + VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] { + using state_t = scalar_t; + VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] { + using input_t = scalar_t; + mamba_cpu::selective_state_update_kernel( + state.data_ptr(), stride_state_n, stride_state_h, stride_state_d, x.data_ptr(), dt_in.data_ptr(), stride_xdt_n, stride_xdt_h, A_in.data_ptr(), stride_A_h, B_in.data_ptr(), C_in.data_ptr(), stride_BC_n, stride_BC_g, @@ -143,9 +161,11 @@ void selective_state_update_cpu_impl( out_f32.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, dsbi_ptr, static_cast(null_block_id), nat_ptr, csl_ptr, N, nheads, ngroups, dim, dstate, dt_softplus); + }); }); - // Copy the updated state back to the original state tensor - state.copy_(state_in); - out.copy_(out_f32); + // Only copy back if we used a temporary buffer + if (need_out_conversion) { + out.copy_(out_f32); + } } From fde43b0d0e8591c82525b837ae0212caff6f4257 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 6 May 2026 16:01:04 +0530 Subject: [PATCH 18/43] fix bf16 to fp32 . Signed-off-by: Akash kaothalkar --- vllm/model_executor/layers/mamba/mamba_mixer.py | 3 ++- vllm/model_executor/layers/mamba/mamba_mixer2.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index 0e476755201e..7800a44eed0e 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -369,6 +369,7 @@ def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): time_proj_bias = self._time_proj_bias() # 4. Perform the recurrence y ← SSM(A, B, C, Δ)(x) + # Keep D in native dtype to avoid expensive conversions scan_out_p = selective_scan_fn( conv_out_p, ssm_state, @@ -376,7 +377,7 @@ def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): self.A, B_p.transpose(-2, -1), C_p.transpose(-2, -1), - self.D.float(), + self.D, gate_p, time_proj_bias, delta_softplus=True, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 2b4b1934f9b3..6ce8242dca25 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -876,10 +876,11 @@ def conv_ssm_forward( # 3. State Space Model sequence transformation n_groups = self.n_groups // self.tp_size + # Keep A in native dtype to avoid expensive conversions + # The SSU kernel handles mixed types efficiently A_d = ( self.A[:, None, ...][:, :, None] .expand(-1, self.head_dim, self.ssm_state_size) - .to(dtype=torch.float32) ) dt_d = dt_d[:, :, None].expand(-1, -1, self.head_dim) dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) From 0945608b0f4182839120c5d63cb4695d68150e72 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 13 May 2026 10:43:22 +0530 Subject: [PATCH 19/43] fix triton path Signed-off-by: Akash kaothalkar --- vllm/config/mamba.py | 23 ++++++- .../layers/mamba/ops/ssu_dispatch.py | 68 ++++++++++++++++++- vllm/utils/torch_utils.py | 7 ++ 3 files changed, 93 insertions(+), 5 deletions(-) diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py index 996478c36760..745779de4455 100644 --- a/vllm/config/mamba.py +++ b/vllm/config/mamba.py @@ -27,6 +27,10 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta): TRITON = "triton" FLASHINFER = "flashinfer" + # Pure-PyTorch fallback for CPU-only platforms (PowerPC, no CUDA, etc.). + # Avoids Triton JIT compilation which is unstable / unsupported on those + # architectures. + CPU = "cpu" @config @@ -34,7 +38,11 @@ class MambaConfig: """Configuration for Mamba SSM backends.""" backend: MambaBackendEnum = MambaBackendEnum.TRITON - """Mamba SSU backend to use.""" + """Mamba SSU backend to use. + + On CPU-only platforms (e.g. PowerPC, x86 without CUDA) the default is + automatically overridden to 'cpu' by ``__post_init__``. + """ enable_stochastic_rounding: bool = False """Enable stochastic rounding when writing SSM state to fp16 cache. @@ -54,9 +62,18 @@ def validate_backend_before(cls, value: Any) -> Any: return value def __post_init__(self): - if self.enable_stochastic_rounding: - from vllm.platforms import current_platform + from vllm.platforms import current_platform + # On CPU-only platforms, silently override the backend to 'cpu' unless + # the user explicitly chose a different backend. Triton JIT is unstable + # (and often unavailable) on platforms like PowerPC. + if ( + self.backend == MambaBackendEnum.TRITON + and current_platform.is_cpu() + ): + object.__setattr__(self, "backend", MambaBackendEnum.CPU) + + if self.enable_stochastic_rounding: if not current_platform.is_cuda(): raise ValueError( "Stochastic rounding for Mamba cache is only supported " diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 92258ef204bd..0f23a4ba2a0c 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -4,8 +4,10 @@ Dispatch module for Mamba selective state update (SSU) backends. Provides a unified `selective_state_update` function that dispatches to -either the Triton or FlashInfer backend based on the configured -`MambaBackendEnum`. Follows SGLang's dispatch pattern adapted for vLLM. +the Triton, FlashInfer, or CPU backend based on the configured +`MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA) +the backend defaults to 'cpu', which uses a pure-PyTorch fallback that +avoids Triton JIT compilation entirely. """ from abc import ABC, abstractmethod @@ -182,9 +184,71 @@ def __call__( ) +class CPUSSUBackend(MambaSSUBackend): + """Pure-PyTorch CPU SSU backend. + + Used on CPU-only platforms (PowerPC, x86 without CUDA, etc.) where + Triton JIT is unavailable or unstable. Delegates to the pure-PyTorch + reference implementation in cpu_fallbacks.py which is numerically + equivalent to the Triton path. + """ + + def __init__(self, mamba_config: MambaConfig): + super().__init__(mamba_config) + from vllm.model_executor.layers.mamba.ops.cpu_fallbacks import ( + _selective_state_update_cpu, + ) + + self._kernel = _selective_state_update_cpu + + @property + def name(self) -> str: + return "cpu" + + def __call__( + self, + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + dt_bias: torch.Tensor, + z: torch.Tensor | None = None, + dt_softplus: bool = False, + state_batch_indices: torch.Tensor | None = None, + dst_state_batch_indices: torch.Tensor | None = None, + null_block_id: int = NULL_BLOCK_ID, + out: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + cu_seqlens: torch.Tensor | None = None, + is_blackwell: bool = False, + ) -> None: + self._kernel( + state, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + dst_state_batch_indices=dst_state_batch_indices, + null_block_id=null_block_id, + out=out, + num_accepted_tokens=num_accepted_tokens, + cu_seqlens=cu_seqlens, + ) + + _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = { MambaBackendEnum.TRITON: TritonSSUBackend, MambaBackendEnum.FLASHINFER: FlashInferSSUBackend, + MambaBackendEnum.CPU: CPUSSUBackend, } _mamba_ssu_backend: MambaSSUBackend | None = None diff --git a/vllm/utils/torch_utils.py b/vllm/utils/torch_utils.py index 12ec5b0fcc66..a7d0ffb9a6ab 100644 --- a/vllm/utils/torch_utils.py +++ b/vllm/utils/torch_utils.py @@ -612,6 +612,13 @@ def async_tensor_h2d( pin_memory: bool = PIN_MEMORY, ) -> torch.Tensor: """Asynchronously create a tensor and copy it from host to device.""" + # On CPU-only platforms (e.g. PowerPC, RISC-V) there is no pinned-memory + # allocator, so torch.tensor(..., pin_memory=True) raises RuntimeError. + # PIN_MEMORY only guards against WSL; check the platform at call time too. + if pin_memory: + from vllm.platforms import current_platform + + pin_memory = current_platform.is_pin_memory_available() t = torch.tensor(data, dtype=dtype, pin_memory=pin_memory, device="cpu") return t.to(device=device, non_blocking=True) From d721f7d94fe05c9487d62c89cd0d6d00fd157d85 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 13 May 2026 00:46:34 -0500 Subject: [PATCH 20/43] Add cpu mamba ssu backend Signed-off-by: Akash kaothalkar --- vllm/config/mamba.py | 14 +++----------- .../layers/mamba/ops/ssu_dispatch.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 11 deletions(-) diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py index 745779de4455..421e5a942cb5 100644 --- a/vllm/config/mamba.py +++ b/vllm/config/mamba.py @@ -62,18 +62,9 @@ def validate_backend_before(cls, value: Any) -> Any: return value def __post_init__(self): - from vllm.platforms import current_platform - - # On CPU-only platforms, silently override the backend to 'cpu' unless - # the user explicitly chose a different backend. Triton JIT is unstable - # (and often unavailable) on platforms like PowerPC. - if ( - self.backend == MambaBackendEnum.TRITON - and current_platform.is_cpu() - ): - object.__setattr__(self, "backend", MambaBackendEnum.CPU) - if self.enable_stochastic_rounding: + from vllm.platforms import current_platform + if not current_platform.is_cuda(): raise ValueError( "Stochastic rounding for Mamba cache is only supported " @@ -91,3 +82,4 @@ def __post_init__(self): "specify `--enable-mamba-cache-stochastic-rounding`, " "or set `--mamba-backend flashinfer`." ) + diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 0f23a4ba2a0c..33a7b7189f4b 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -274,6 +274,20 @@ def initialize_mamba_ssu_backend( global _mamba_ssu_backend backend = mamba_config.backend + + # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is + # unstable or unavailable. Silently fall back to the pure-PyTorch CPU + # backend unless the user explicitly chose something other than "triton". + if backend == MambaBackendEnum.TRITON: + from vllm.platforms import current_platform + + if current_platform.is_cpu(): + logger.info( + "CPU platform detected: overriding Mamba SSU backend " + "from 'triton' to 'cpu' (pure-PyTorch fallback)." + ) + backend = MambaBackendEnum.CPU + if backend not in _BACKEND_REGISTRY: raise ValueError( f"Unknown Mamba SSU backend: {backend}. " @@ -340,3 +354,4 @@ def selective_state_update( cu_seqlens=cu_seqlens, is_blackwell=is_blackwell, ) + From 8f7b4c6106f873e843d414b0e02fcdaddeadad17 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 13 May 2026 02:30:13 -0500 Subject: [PATCH 21/43] fix Signed-off-by: Akash kaothalkar --- csrc/cpu/cpu_types_vsx.hpp | 16 ++- csrc/cpu/mamba_cpu.cpp | 136 ++++++++++-------- csrc/cpu/mamba_kernels.hpp | 130 ++++++++++------- .../layers/mamba/ops/ssu_dispatch.py | 94 ++++++++---- 4 files changed, 236 insertions(+), 140 deletions(-) diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index 87c7a9dd51f4..e1e4659da167 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -208,13 +208,14 @@ struct FP32Vec8 : public Vec { } float reduce_sum() const { - AliasReg ar; - ar.reg = reg; - float result = 0; - unroll_loop( - [&result, &ar](int i) { result += ar.values[i]; }); - - return result; + // VSX horizontal reduction: 3 vector ops instead of 8 scalar adds. + // Step 1: pairwise sum of the two 4-wide halves + __vector float s = vec_add(reg.val[0], reg.val[1]); + // Step 2: rotate by 8 bytes (2 floats) and add + s = vec_add(s, vec_sld(s, s, 8)); + // Step 3: rotate by 4 bytes (1 float) and add => all lanes hold total + s = vec_add(s, vec_sld(s, s, 4)); + return vec_extract(s, 0); } FP32Vec8 exp() const { @@ -797,3 +798,4 @@ inline void prefetch(const void* addr) { }; // namespace vec_op #endif + diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 716da80c315f..ee26c119e497 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -82,90 +82,112 @@ void selective_state_update_cpu_impl( const c10::optional& num_accepted_tokens, const c10::optional& cu_seqlens ) { - // Use state's dtype as the primary type to avoid expensive conversions - // The kernel supports mixed types: state_t can be BFloat16 while input_t matches x at::ScalarType state_type = state.scalar_type(); at::ScalarType input_type = x.scalar_type(); - - // Only convert/contiguous if needed to minimize overhead - auto ensure_type_and_contiguous = [input_type](const at::Tensor& t) -> at::Tensor { - if (t.scalar_type() != input_type) { - return t.to(input_type).contiguous(); - } - return t.is_contiguous() ? t : t.contiguous(); + + // x, B, C must be contiguous and match input_type + auto ensure_input = [input_type](const at::Tensor& t) -> at::Tensor { + at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t; + return r.is_contiguous() ? r : r.contiguous(); }; - - at::Tensor dt_in = ensure_type_and_contiguous(dt); - at::Tensor A_in = ensure_type_and_contiguous(A); - at::Tensor B_in = ensure_type_and_contiguous(B); - at::Tensor C_in = ensure_type_and_contiguous(C); - - at::Tensor D_in, z_in, dt_bias_in; - if (D.has_value() && D.value().defined()) { - D_in = ensure_type_and_contiguous(D.value()); - } - if (z.has_value() && z.value().defined()) { - z_in = ensure_type_and_contiguous(z.value()); - } - if (dt_bias.has_value() && dt_bias.value().defined()) { - dt_bias_in = ensure_type_and_contiguous(dt_bias.value()); + at::Tensor x_in = ensure_input(x); + at::Tensor B_in = ensure_input(B); + at::Tensor C_in = ensure_input(C); + at::Tensor z_in; + if (z.has_value() && z.value().defined()) z_in = ensure_input(z.value()); + + // A, D, dt_bias are float32 model parameters that arrive here as expanded + // tensors, e.g. A is (nheads, head_dim, dstate) with strides (1, 0, 0). + // We need just the scalar value per head as a (nheads,) 1-D array so that + // A_ptr[h] in the kernel correctly reads head h's value. + // + // Strategy: peel trailing expanded (stride=0) dims via .select(), which is + // a zero-copy view. For A: (nheads, head_dim, dstate) strides (1,0,0) + // → .select(2,0) → (nheads, head_dim) strides (1,0) + // → .select(1,0) → (nheads,) stride (1,) ← contiguous, free. + // No allocation, no type conversion (A is already float32). + auto to_per_head_1d_f32 = [](const at::Tensor& t) -> at::Tensor { + at::Tensor r = t; + // Peel trailing dimensions that are broadcast (stride=0 or size=1) + while (r.dim() > 1) r = r.select(r.dim() - 1, 0); + if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat); + return r.is_contiguous() ? r : r.contiguous(); + }; + + at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32 + at::Tensor D_f32, dt_bias_f32; + if (D.has_value() && D.value().defined()) + D_f32 = to_per_head_1d_f32(D.value()); + if (dt_bias.has_value() && dt_bias.value().defined()) + dt_bias_f32 = to_per_head_1d_f32(dt_bias.value()); + + // dt: reduce (N, nheads, head_dim) expanded tensor → (N, nheads) BEFORE + // the type conversion so we convert head_dim x fewer elements. + at::Tensor dt_f32; + { + // If dt was expanded to (N, nheads, head_dim) with stride-0 in dim 2, + // take a zero-copy view of index 0 along that dim first. + at::Tensor t2 = (dt.dim() == 3) ? dt.select(2, 0) : dt; // (N, nheads) + at::Tensor t3 = (t2.scalar_type() != at::kFloat) ? t2.to(at::kFloat) : t2; + dt_f32 = t3.is_contiguous() ? t3 : t3.contiguous(); } int64_t nheads = state.size(1); int64_t dim = state.size(2); int64_t dstate = state.size(3); - int64_t N = x.size(0); + int64_t N = x_in.size(0); int64_t ngroups = B_in.size(1); // Strides int64_t stride_state_n = state.stride(0); int64_t stride_state_h = state.stride(1); int64_t stride_state_d = state.stride(2); - int64_t stride_xdt_n = x.stride(0); - int64_t stride_xdt_h = x.stride(1); - int64_t stride_A_h = A_in.stride(0); + int64_t stride_x_n = x_in.stride(0); + int64_t stride_x_h = x_in.stride(1); + int64_t stride_dt_n = dt_f32.stride(0); // dt is (N, nheads) int64_t stride_BC_n = B_in.stride(0); int64_t stride_BC_g = B_in.stride(1); int64_t stride_out_n = out.stride(0); int64_t stride_out_h = out.stride(1); - int64_t stride_D_h = D_in.defined() ? D_in.stride(0) : 0; - int64_t stride_dtbias_h = dt_bias_in.defined() ? dt_bias_in.stride(0) : 0; - // Optional pointers - extract once + // Optional index pointers auto get_int32_ptr = [](const c10::optional& opt) -> const int32_t* { - return (opt.has_value() && opt.value().defined()) ? opt.value().data_ptr() : nullptr; + return (opt.has_value() && opt.value().defined()) ? + opt.value().data_ptr() : nullptr; }; - - const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices); + const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices); const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices); - const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens); - const int32_t* csl_ptr = get_int32_ptr(cu_seqlens); - - // Optimize output buffer: only use float32 if output type is not already float32 - // This avoids an extra copy when out is already float32 - bool need_out_conversion = (out.scalar_type() != at::kFloat); - at::Tensor out_f32 = need_out_conversion ? at::empty_like(out, at::kFloat) : out; + const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens); + const int32_t* csl_ptr = get_int32_ptr(cu_seqlens); + // Dispatch on (state_t, input_t, out_t): write directly into `out` + // without any intermediate float32 buffer. VLLM_DISPATCH_FLOATING_TYPES(state_type, "ssu_state", [&] { using state_t = scalar_t; VLLM_DISPATCH_FLOATING_TYPES(input_type, "ssu_input", [&] { using input_t = scalar_t; - mamba_cpu::selective_state_update_kernel( - state.data_ptr(), stride_state_n, stride_state_h, stride_state_d, - x.data_ptr(), dt_in.data_ptr(), stride_xdt_n, stride_xdt_h, - A_in.data_ptr(), stride_A_h, - B_in.data_ptr(), C_in.data_ptr(), stride_BC_n, stride_BC_g, - D_in.defined() ? D_in.data_ptr() : nullptr, stride_D_h, - z_in.defined() ? z_in.data_ptr() : nullptr, - dt_bias_in.defined() ? dt_bias_in.data_ptr() : nullptr, stride_dtbias_h, - out_f32.data_ptr(), stride_out_n, stride_out_h, - sbi_ptr, dsbi_ptr, static_cast(null_block_id), - nat_ptr, csl_ptr, N, nheads, ngroups, dim, dstate, dt_softplus); + VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] { + using out_t = scalar_t; + mamba_cpu::selective_state_update_kernel( + state.data_ptr(), + stride_state_n, stride_state_h, stride_state_d, + x_in.data_ptr(), + stride_x_n, stride_x_h, + dt_f32.data_ptr(), + stride_dt_n, + A_f32.data_ptr(), + B_in.data_ptr(), C_in.data_ptr(), + stride_BC_n, stride_BC_g, + D_f32.defined() ? D_f32.data_ptr() : nullptr, + z_in.defined() ? z_in.data_ptr() : nullptr, + dt_bias_f32.defined() ? dt_bias_f32.data_ptr() : nullptr, + out.data_ptr(), + stride_out_n, stride_out_h, + sbi_ptr, dsbi_ptr, static_cast(null_block_id), + nat_ptr, csl_ptr, N, nheads, ngroups, dim, dstate, dt_softplus); + }); }); }); - - // Only copy back if we used a temporary buffer - if (need_out_conversion) { - out.copy_(out_f32); - } } + + diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index f351f9acade8..9e62b4c8325d 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -11,7 +11,6 @@ #include #include #include -#include #include namespace mamba_cpu { @@ -25,7 +24,7 @@ inline void causal_conv1d_update_kernel( float* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen, int64_t width, int64_t state_len, bool do_silu) { - + #pragma omp parallel for for (int64_t b = 0; b < batch; ++b) { int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b; @@ -42,7 +41,7 @@ inline void causal_conv1d_update_kernel( const float* w = weight_ptr + d * width; float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f; - + for (int64_t k = 0; k < state_len; ++k) { acc += w[k] * sd[k]; } @@ -56,8 +55,8 @@ inline void causal_conv1d_update_kernel( } if (do_silu) { - float sigmoid = (acc >= 0) ? - 1.0f / (1.0f + std::exp(-acc)) : + float sigmoid = (acc >= 0) ? + 1.0f / (1.0f + std::exp(-acc)) : std::exp(acc) / (1.0f + std::exp(acc)); acc *= sigmoid; } @@ -69,27 +68,48 @@ inline void causal_conv1d_update_kernel( // --------------------------------------------------------------------------- // selective_state_update +// +// Template parameters: +// state_t - dtype of ssm_state cache (typically BFloat16) +// input_t - dtype of x, B, C (typically BFloat16) +// out_t - dtype of output tensor (typically BFloat16) +// Write directly — no float32 intermediate buffer needed. +// +// A, D, dt_bias are accepted as const float* (they are always float32 +// model parameters in Mamba2). This eliminates the per-call float32→BF16 +// conversion and the .contiguous() materialisation of the broadcast-expand. +// +// dt is accepted as a (N, nheads) scalar-per-head tensor, not as the +// (N, nheads, head_dim) expansion, so no .contiguous() copy is needed. // --------------------------------------------------------------------------- -template +template inline void selective_state_update_kernel( state_t* __restrict__ state_ptr, int64_t stride_state_n, int64_t stride_state_h, int64_t stride_state_d, - const input_t* __restrict__ x_ptr, const input_t* __restrict__ dt_ptr, - int64_t stride_xdt_n, int64_t stride_xdt_h, - const input_t* __restrict__ A_ptr, int64_t stride_A_h, + const input_t* __restrict__ x_ptr, + int64_t stride_x_n, int64_t stride_x_h, + // dt: (N, nheads) — scalar per head, NOT expanded to head_dim + const float* __restrict__ dt_ptr, + int64_t stride_dt_n, + // A: (nheads,) float32 — scalar per head + const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr, const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g, - const input_t* __restrict__ D_ptr, int64_t stride_D_h, + // D: (nheads,) float32 — scalar per head (nullptr if not used) + const float* __restrict__ D_ptr, + // z: same shape as x (optional) const input_t* __restrict__ z_ptr, - const input_t* __restrict__ dt_bias_ptr, int64_t stride_dtbias_h, - float* __restrict__ out_ptr, int64_t stride_out_n, int64_t stride_out_h, + // dt_bias: (nheads,) float32 — scalar per head (nullptr if not used) + const float* __restrict__ dt_bias_ptr, + out_t* __restrict__ out_ptr, + int64_t stride_out_n, int64_t stride_out_h, const int32_t* __restrict__ state_batch_indices, const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id, const int32_t* __restrict__ num_accepted_tokens, const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads, int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) { - + using state_vec_t = vec_op::vec_t; using input_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = 8; @@ -106,84 +126,96 @@ inline void selective_state_update_kernel( seq_len = 1; } - int64_t state_read_idx = (state_batch_indices != nullptr) ? + int64_t state_read_idx = (state_batch_indices != nullptr) ? state_batch_indices[seq_idx] : seq_idx; if (state_read_idx == null_block_id) continue; - int64_t state_write_idx = (num_accepted_tokens == nullptr) ? - ((dst_state_batch_indices != nullptr) ? dst_state_batch_indices[seq_idx] : state_read_idx) : -1; + int64_t state_write_idx = (num_accepted_tokens == nullptr) ? + ((dst_state_batch_indices != nullptr) ? + dst_state_batch_indices[seq_idx] : state_read_idx) : -1; state_t* s = state_ptr + state_read_idx * stride_state_n; for (int64_t t = 0; t < seq_len; ++t) { int64_t token_idx = bos + t; - const input_t* x_tok = x_ptr + token_idx * stride_xdt_n; - const input_t* dt_tok = dt_ptr + token_idx * stride_xdt_n; + const input_t* x_tok = x_ptr + token_idx * stride_x_n; + // dt: (N, nheads) — one float per head per token + const float* dt_tok = dt_ptr + token_idx * stride_dt_n; const input_t* B_tok = B_ptr + token_idx * stride_BC_n; const input_t* C_tok = C_ptr + token_idx * stride_BC_n; - float* out_tok = out_ptr + token_idx * stride_out_n; + out_t* out_tok = out_ptr + token_idx * stride_out_n; #pragma omp parallel for for (int64_t h = 0; h < nheads; ++h) { int64_t g = h / nheads_per_group; - const input_t* x_h = x_tok + h * stride_xdt_h; - const input_t* dt_h = dt_tok + h * stride_xdt_h; + const input_t* x_h = x_tok + h * stride_x_h; const input_t* B_g = B_tok + g * stride_BC_g; const input_t* C_g = C_tok + g * stride_BC_g; - const input_t* A_h = A_ptr + h * stride_A_h; - const input_t* dt_bias_h = (dt_bias_ptr != nullptr) ? dt_bias_ptr + h * stride_dtbias_h : nullptr; - const input_t* D_h = (D_ptr != nullptr) ? D_ptr + h * stride_D_h : nullptr; - const input_t* z_h = (z_ptr != nullptr) ? z_ptr + token_idx * stride_xdt_n + h * stride_xdt_h : nullptr; - float* out_h = out_tok + h * stride_out_h; + out_t* out_h = out_tok + h * stride_out_h; state_t* s_h = s + h * stride_state_h; + // Read scalars-per-head (A, dt, dt_bias, D) — no per-dim indexing + float dt_val = dt_tok[h]; + if (dt_bias_ptr != nullptr) dt_val += dt_bias_ptr[h]; + if (dt_softplus) { + dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; + } + const float A_val = A_ptr[h]; // scalar: same for all dim, dstate + const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f; + + const input_t* z_h = (z_ptr != nullptr) ? + z_ptr + token_idx * stride_x_n + h * stride_x_h : nullptr; + + vec_op::FP32Vec8 dt_vec(dt_val); + // dA = exp(A * dt): A and dt are SCALARS per head, so compute once + // and broadcast. This saves 7 redundant std::exp() calls that + // FP32Vec8::exp() would otherwise make on the broadcast vector. + const float dA_scalar = std::exp(A_val * dt_val); + vec_op::FP32Vec8 dA(dA_scalar); // broadcast + for (int64_t d = 0; d < dim; ++d) { float x_val = static_cast(x_h[d]); - float dt_val = static_cast(dt_h[d]); - if (dt_bias_h != nullptr) dt_val += static_cast(dt_bias_h[d]); - if (dt_softplus) { - dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; - } vec_op::FP32Vec8 out_vec(0.0f); state_t* s_hd = s_h + d * stride_state_d; - const input_t* A_hd = A_h + d * dstate; - + const input_t* B_g_base = B_g; + const input_t* C_g_base = C_g; + vec_op::FP32Vec8 x_vec(x_val); - vec_op::FP32Vec8 dt_vec(dt_val); + // dBx = B * x * dt — same dA for all dstate (A is scalar) + // s_new = s * dA + B * x * dt int64_t n = 0; for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { - vec_op::FP32Vec8 A_v((input_vec_t(A_hd + n))); - vec_op::FP32Vec8 B_v((input_vec_t(B_g + n))); - vec_op::FP32Vec8 C_v((input_vec_t(C_g + n))); - + vec_op::FP32Vec8 B_v((input_vec_t(B_g_base + n))); + vec_op::FP32Vec8 C_v((input_vec_t(C_g_base + n))); vec_op::FP32Vec8 s_v((state_vec_t(s_hd + n))); - - vec_op::FP32Vec8 dA = (A_v * dt_vec).exp(); + vec_op::FP32Vec8 dBx = B_v * x_vec * dt_vec; vec_op::FP32Vec8 s_new = s_v * dA + dBx; - + state_vec_t(s_new).save(s_hd + n); out_vec = out_vec + s_new * C_v; } - + float out_val = out_vec.reduce_sum(); for (; n < dstate; ++n) { - float dA = std::exp(static_cast(A_hd[n]) * dt_val); + // Reuse dA_scalar computed once per head — no exp() re-call float dBx = static_cast(B_g[n]) * x_val * dt_val; - float s_new = static_cast(s_hd[n]) * dA + dBx; + float s_new = static_cast(s_hd[n]) * dA_scalar + dBx; s_hd[n] = static_cast(s_new); out_val += s_new * static_cast(C_g[n]); } - if (D_h != nullptr) out_val += x_val * static_cast(D_h[d]); + if (D_ptr != nullptr) out_val += x_val * D_val; if (z_h != nullptr) { float z_val = static_cast(z_h[d]); - float sigmoid = (z_val >= 0) ? 1.0f / (1.0f + std::exp(-z_val)) : std::exp(z_val) / (1.0f + std::exp(z_val)); + float sigmoid = (z_val >= 0) ? + 1.0f / (1.0f + std::exp(-z_val)) : + std::exp(z_val) / (1.0f + std::exp(z_val)); out_val *= z_val * sigmoid; } - out_h[d] = out_val; + out_h[d] = static_cast(out_val); } } @@ -196,7 +228,8 @@ inline void selective_state_update_kernel( } } - if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && state_write_idx != state_read_idx) { + if (num_accepted_tokens == nullptr && state_write_idx != null_block_id && + state_write_idx != state_read_idx) { state_t* dst_s = state_ptr + state_write_idx * stride_state_n; std::memmove(dst_s, s, nheads * stride_state_h * sizeof(state_t)); } @@ -204,3 +237,4 @@ inline void selective_state_update_kernel( } } // namespace mamba_cpu + diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 33a7b7189f4b..9fcfbf16efd5 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -185,21 +185,37 @@ def __call__( class CPUSSUBackend(MambaSSUBackend): - """Pure-PyTorch CPU SSU backend. + """CPU SSU backend using the compiled C++ VSX/scalar kernel. - Used on CPU-only platforms (PowerPC, x86 without CUDA, etc.) where - Triton JIT is unavailable or unstable. Delegates to the pure-PyTorch - reference implementation in cpu_fallbacks.py which is numerically - equivalent to the Triton path. + On CPU-only platforms (PowerPC, x86 without CUDA) this dispatches to + the vectorized C++ kernel registered as ``torch.ops._C.selective_state_update_cpu``. + That kernel uses vec_op SIMD intrinsics (VSX on ppc64le, AVX2 on x86, + scalar fallback elsewhere) and is parallelised with OpenMP across heads. + + Falls back to the pure-PyTorch implementation only if the C++ op is + unavailable (e.g. a CPU-less build). """ def __init__(self, mamba_config: MambaConfig): super().__init__(mamba_config) - from vllm.model_executor.layers.mamba.ops.cpu_fallbacks import ( - _selective_state_update_cpu, - ) - - self._kernel = _selective_state_update_cpu + try: + from vllm import _custom_ops as ops + + # Verify the op is actually registered (CPU build required) + _ = ops.selective_state_update_cpu + self._use_cpp = True + self._cpp_kernel = ops.selective_state_update_cpu + logger.info("CPUSSUBackend: using compiled C++ selective_state_update kernel.") + except (ImportError, AttributeError): + from vllm.model_executor.layers.mamba.ops.cpu_fallbacks import ( + _selective_state_update_cpu, + ) + self._use_cpp = False + self._py_kernel = _selective_state_update_cpu + logger.warning( + "CPUSSUBackend: C++ selective_state_update op not available, " + "falling back to pure-PyTorch (slow). Rebuild with CPU extensions." + ) @property def name(self) -> str: @@ -225,24 +241,46 @@ def __call__( cu_seqlens: torch.Tensor | None = None, is_blackwell: bool = False, ) -> None: - self._kernel( - state, - x, - dt, - A, - B, - C, - D=D, - z=z, - dt_bias=dt_bias, - dt_softplus=dt_softplus, - state_batch_indices=state_batch_indices, - dst_state_batch_indices=dst_state_batch_indices, - null_block_id=null_block_id, - out=out, - num_accepted_tokens=num_accepted_tokens, - cu_seqlens=cu_seqlens, - ) + if self._use_cpp: + # C++ kernel: state shape expected as (nstates, nheads, dim, dstate) + # The kernel writes in-place into `out` and updates `state`. + self._cpp_kernel( + state, + x, + dt, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + dst_state_batch_indices, + null_block_id, + out, + num_accepted_tokens, + cu_seqlens, + ) + else: + self._py_kernel( + state, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=state_batch_indices, + dst_state_batch_indices=dst_state_batch_indices, + null_block_id=null_block_id, + out=out, + num_accepted_tokens=num_accepted_tokens, + cu_seqlens=cu_seqlens, + ) _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = { From 127ef5425eea79934d54c4de503669600f8e60e0 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 13 May 2026 02:45:06 -0500 Subject: [PATCH 22/43] fix bf16 to fp32 Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 61 +++++++++++++++++++++++++------------- csrc/cpu/mamba_kernels.hpp | 44 ++++++++++++++------------- 2 files changed, 65 insertions(+), 40 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index ee26c119e497..8c4eaa6a7524 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -16,33 +16,46 @@ // causal_conv1d_update // --------------------------------------------------------------------------- at::Tensor causal_conv1d_update_cpu_impl( - at::Tensor& x, // modified in-place (re-typed to float32) + at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight, const c10::optional& bias, const c10::optional& activation, const c10::optional& conv_state_indices, const c10::optional& query_start_loc, int64_t pad_slot_id) { - + bool do_silu = false; if (activation.has_value()) { const std::string& act = activation.value(); do_silu = (act == "silu" || act == "swish"); } - // Causal conv still works in float32 for now (minimal overhead compared to SSM) - at::Tensor x_f32 = x.to(at::kFloat).contiguous(); - at::Tensor state_f32 = conv_state.to(at::kFloat).contiguous(); - at::Tensor w_f32 = weight.to(at::kFloat).contiguous(); + at::ScalarType dtype = x.scalar_type(); + + // Ensure contiguous in native dtype — no float32 conversion. + at::Tensor x_c = x.is_contiguous() ? x : x.contiguous(); + // conv_state is modified in-place; make a contiguous working copy if needed. + bool state_copied = !conv_state.is_contiguous() || + conv_state.scalar_type() != dtype; + at::Tensor state_c = state_copied ? + conv_state.to(dtype).contiguous() : + conv_state; + // Weight: coerce to same dtype if needed (should match in practice) + at::Tensor w_c = (weight.scalar_type() != dtype) ? + weight.to(dtype).contiguous() : + (weight.is_contiguous() ? weight : weight.contiguous()); + + // Bias stays float32 (small tensor, used for accumulation only) at::Tensor bias_f32; - if (bias.has_value() && bias.value().defined()) bias_f32 = bias.value().to(at::kFloat).contiguous(); + if (bias.has_value() && bias.value().defined()) + bias_f32 = bias.value().to(at::kFloat).contiguous(); - int64_t batch = x_f32.size(0); - int64_t dim = x_f32.size(1); - int64_t seqlen = (x_f32.dim() == 3) ? x_f32.size(2) : 1; - int64_t width = w_f32.size(1); - int64_t state_len = state_f32.size(2); + int64_t batch = x_c.size(0); + int64_t dim = x_c.size(1); + int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1; + int64_t width = w_c.size(1); + int64_t state_len = state_c.size(2); - at::Tensor out_f32 = at::empty_like(x_f32); + at::Tensor out = at::empty_like(x_c); // native dtype, no float32 alloc const int32_t* cache_idx_ptr = nullptr; at::Tensor cache_idx_int; @@ -51,16 +64,24 @@ at::Tensor causal_conv1d_update_cpu_impl( cache_idx_ptr = cache_idx_int.data_ptr(); } - mamba_cpu::causal_conv1d_update_kernel( - x_f32.data_ptr(), state_f32.data_ptr(), - w_f32.data_ptr(), bias_f32.defined() ? bias_f32.data_ptr() : nullptr, - out_f32.data_ptr(), cache_idx_ptr, static_cast(pad_slot_id), - batch, dim, seqlen, width, state_len, do_silu); + VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] { + mamba_cpu::causal_conv1d_update_kernel( + x_c.data_ptr(), + state_c.data_ptr(), + w_c.data_ptr(), + bias_f32.defined() ? bias_f32.data_ptr() : nullptr, + out.data_ptr(), + cache_idx_ptr, static_cast(pad_slot_id), + batch, dim, seqlen, width, state_len, do_silu); + }); + + // Write state back only if we made a working copy + if (state_copied) conv_state.copy_(state_c); - conv_state.copy_(state_f32.to(conv_state.scalar_type())); - return out_f32.to(x.scalar_type()); + return out; } + // --------------------------------------------------------------------------- // selective_state_update // --------------------------------------------------------------------------- diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index 9e62b4c8325d..18fffc26e538 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -16,12 +16,15 @@ namespace mamba_cpu { // --------------------------------------------------------------------------- -// causal_conv1d_update +// causal_conv1d_update — templated for native BF16/FP32 +// x, state, weight, out are in scalar_t; accumulation in float32. +// Eliminates 5 dtype-conversion copies that the old float32-only path did. // --------------------------------------------------------------------------- +template inline void causal_conv1d_update_kernel( - const float* __restrict__ x_ptr, float* __restrict__ state_ptr, - const float* __restrict__ weight_ptr, const float* __restrict__ bias_ptr, - float* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, + const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr, + const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr, + scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen, int64_t width, int64_t state_len, bool do_silu) { @@ -31,36 +34,37 @@ inline void causal_conv1d_update_kernel( if (cache_idx == pad_slot_id) continue; for (int64_t t = 0; t < seqlen; ++t) { - const float* x_b = x_ptr + (b * dim * seqlen + t); - float* out_b = out_ptr + (b * dim * seqlen + t); - float* s = state_ptr + cache_idx * dim * state_len; + const scalar_t* x_b = x_ptr + (b * dim * seqlen + t); + scalar_t* out_b = out_ptr + (b * dim * seqlen + t); + scalar_t* s = state_ptr + cache_idx * dim * state_len; for (int64_t d = 0; d < dim; ++d) { - float x_val = x_b[d * seqlen]; - float* sd = s + d * state_len; + float x_val = static_cast(x_b[d * seqlen]); + scalar_t* sd = s + d * state_len; + const scalar_t* w = weight_ptr + d * width; - const float* w = weight_ptr + d * width; + // Accumulate in float32 for precision float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f; - for (int64_t k = 0; k < state_len; ++k) { - acc += w[k] * sd[k]; + acc += static_cast(w[k]) * static_cast(sd[k]); } - acc += w[state_len] * x_val; + acc += static_cast(w[state_len]) * x_val; + // Shift state left, append new input — native dtype, no copy if (state_len > 1) { - std::memmove(sd, sd + 1, (state_len - 1) * sizeof(float)); + std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t)); } if (state_len > 0) { - sd[state_len - 1] = x_val; + sd[state_len - 1] = static_cast(x_val); } if (do_silu) { - float sigmoid = (acc >= 0) ? - 1.0f / (1.0f + std::exp(-acc)) : - std::exp(acc) / (1.0f + std::exp(acc)); - acc *= sigmoid; + float sigmoid = (acc >= 0) ? + 1.0f / (1.0f + std::exp(-acc)) : + std::exp(acc) / (1.0f + std::exp(acc)); + acc *= sigmoid; } - out_b[d * seqlen] = acc; + out_b[d * seqlen] = static_cast(acc); } } } From 48da0cb684871d56a1bacca5f5f7c737d47d21da Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Thu, 14 May 2026 01:14:56 -0500 Subject: [PATCH 23/43] fix direct copy Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 33 +++++-- csrc/cpu/mamba_kernels.hpp | 43 ++++++--- .../layers/fused_moe/cpu_fused_moe.py | 95 +++++++++++++++++++ 3 files changed, 149 insertions(+), 22 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 8c4eaa6a7524..0d0a829c4eec 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -31,20 +31,23 @@ at::Tensor causal_conv1d_update_cpu_impl( at::ScalarType dtype = x.scalar_type(); - // Ensure contiguous in native dtype — no float32 conversion. + // Input x: contiguous in native dtype. at::Tensor x_c = x.is_contiguous() ? x : x.contiguous(); - // conv_state is modified in-place; make a contiguous working copy if needed. - bool state_copied = !conv_state.is_contiguous() || - conv_state.scalar_type() != dtype; - at::Tensor state_c = state_copied ? - conv_state.to(dtype).contiguous() : - conv_state; + + // conv_state: NEVER copy the full paged tensor just for layout reasons. + // If the dtype matches we work directly on conv_state (contiguous or not) + // by extracting strides and passing them to the kernel. + // Only a dtype-conversion copy is made when types differ (rare for BF16). + bool state_type_ok = (conv_state.scalar_type() == dtype); + at::Tensor state_c = state_type_ok ? conv_state : conv_state.to(dtype); + // state_c and conv_state may be non-contiguous — that is intentional. + // Weight: coerce to same dtype if needed (should match in practice) at::Tensor w_c = (weight.scalar_type() != dtype) ? weight.to(dtype).contiguous() : (weight.is_contiguous() ? weight : weight.contiguous()); - // Bias stays float32 (small tensor, used for accumulation only) + // Bias stays float32 (small scalar, used only for fp32 accumulation) at::Tensor bias_f32; if (bias.has_value() && bias.value().defined()) bias_f32 = bias.value().to(at::kFloat).contiguous(); @@ -55,6 +58,14 @@ at::Tensor causal_conv1d_update_cpu_impl( int64_t width = w_c.size(1); int64_t state_len = state_c.size(2); + // Extract strides — works for contiguous AND non-contiguous (transposed) state. + // stride(0): between cache slots (e.g. num_slots × dim × width-1 in contiguous) + // stride(1): between conv channels (dim stride) + // stride(2): between state elements (=1 when contiguous, =dim when transposed) + int64_t stride_s_slot = state_c.stride(0); + int64_t stride_s_dim = state_c.stride(1); + int64_t stride_s_state = state_c.stride(2); + at::Tensor out = at::empty_like(x_c); // native dtype, no float32 alloc const int32_t* cache_idx_ptr = nullptr; @@ -68,6 +79,7 @@ at::Tensor causal_conv1d_update_cpu_impl( mamba_cpu::causal_conv1d_update_kernel( x_c.data_ptr(), state_c.data_ptr(), + stride_s_slot, stride_s_dim, stride_s_state, w_c.data_ptr(), bias_f32.defined() ? bias_f32.data_ptr() : nullptr, out.data_ptr(), @@ -75,8 +87,9 @@ at::Tensor causal_conv1d_update_cpu_impl( batch, dim, seqlen, width, state_len, do_silu); }); - // Write state back only if we made a working copy - if (state_copied) conv_state.copy_(state_c); + // Write back only when a type-conversion copy was made. + // Layout-only non-contiguity is handled via strides above — no copy needed. + if (!state_type_ok) conv_state.copy_(state_c); return out; } diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index 18fffc26e538..b81313dfec2a 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -17,12 +17,23 @@ namespace mamba_cpu { // --------------------------------------------------------------------------- // causal_conv1d_update — templated for native BF16/FP32 -// x, state, weight, out are in scalar_t; accumulation in float32. -// Eliminates 5 dtype-conversion copies that the old float32-only path did. +// +// state_ptr may point to a NON-CONTIGUOUS paged KV cache tensor. +// Explicit strides are passed so the kernel writes directly into the +// correct memory locations without making a contiguous copy of the full +// paged tensor (which was the source of the 34-41% direct_copy_kernel). +// +// stride_s_slot = state.stride(0) — between cache slots +// stride_s_dim = state.stride(1) — between conv_dim channels +// stride_s_state = state.stride(2) — between state elements +// +// When stride_s_state == 1 (contiguous), the memmove fast path is used. // --------------------------------------------------------------------------- template inline void causal_conv1d_update_kernel( - const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr, + const scalar_t* __restrict__ x_ptr, + scalar_t* __restrict__ state_ptr, + int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state, const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr, scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen, @@ -36,26 +47,34 @@ inline void causal_conv1d_update_kernel( for (int64_t t = 0; t < seqlen; ++t) { const scalar_t* x_b = x_ptr + (b * dim * seqlen + t); scalar_t* out_b = out_ptr + (b * dim * seqlen + t); - scalar_t* s = state_ptr + cache_idx * dim * state_len; + // Base of this slot in the (possibly non-contiguous) paged state + scalar_t* s_base = state_ptr + cache_idx * stride_s_slot; for (int64_t d = 0; d < dim; ++d) { float x_val = static_cast(x_b[d * seqlen]); - scalar_t* sd = s + d * state_len; + scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state const scalar_t* w = weight_ptr + d * width; // Accumulate in float32 for precision float acc = (bias_ptr != nullptr) ? bias_ptr[d] : 0.0f; for (int64_t k = 0; k < state_len; ++k) { - acc += static_cast(w[k]) * static_cast(sd[k]); + acc += static_cast(w[k]) * + static_cast(sd[k * stride_s_state]); } acc += static_cast(w[state_len]) * x_val; - // Shift state left, append new input — native dtype, no copy - if (state_len > 1) { - std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t)); - } - if (state_len > 0) { - sd[state_len - 1] = static_cast(x_val); + // Shift state left and append new input. + // Use memmove when contiguous (stride==1); element loop otherwise. + if (stride_s_state == 1) { + if (state_len > 1) + std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t)); + if (state_len > 0) + sd[state_len - 1] = static_cast(x_val); + } else { + for (int64_t k = 0; k < state_len - 1; ++k) + sd[k * stride_s_state] = sd[(k + 1) * stride_s_state]; + if (state_len > 0) + sd[(state_len - 1) * stride_s_state] = static_cast(x_val); } if (do_silu) { diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 73b3745e29d6..de20ce2f2370 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -427,6 +427,100 @@ def forward_torch( return output + def forward_batched_gemm( + self, + layer: torch.nn.Module, + input: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + activation: MoEActivation, + global_num_experts: int = -1, + skip_weighted: bool = False, + ) -> torch.Tensor: + """Universal batched GEMM fallback for CPUs without AVX512 fused MoE. + + The per-expert F.linear loop in forward_torch issues N_experts separate + BLAS calls, each triggering a full weight-matrix packing pass + (sbgemm_incopy / similar). For decode (small T), packing dominates. + + This path instead: + 1. One torch.mm for gate+up — packs ALL expert weights in one pass. + 2. One torch.bmm for down — batched over experts. + → 2 BLAS weight-packing passes regardless of N_experts. + + Works on any CPU (x86 non-AVX512, ARM non-AMX, PowerPC, RISC-V, ...). + The trade-off (computing all experts densely) is beneficial whenever + decode batch size is small relative to expert weight size. + """ + if skip_weighted: + # Weight is pre-applied to input; must NOT re-apply in accumulation. + assert topk_ids.size(1) == 1, ( + "apply_router_weight_on_input is only supported for topk=1" + ) + input = input * topk_weights.to(input.dtype) + + num_tokens: int = input.shape[0] + hidden_dim: int = input.shape[1] + num_experts: int = layer.w13_weight.shape[0] + gate_up_dim: int = layer.w13_weight.shape[1] # gate + up concatenated + top_k: int = topk_ids.shape[1] + act_fn = _CPU_MOE_ACT_FN[activation] + + # ------------------------------------------------------------------ + # 1. Gate+Up — single large GEMM covering ALL experts at once. + # w13_weight: (E, gate+up, hidden) → (E*gate+up, hidden) + # mm: (T, hidden) @ (hidden, E*gate+up) → (T, E*gate+up) + # → view (T, E, gate+up) + # ------------------------------------------------------------------ + w13_2d = layer.w13_weight.view(num_experts * gate_up_dim, hidden_dim) + gate_up_all = torch.mm(input, w13_2d.T) # (T, E*gate+up) + if hasattr(layer, "w13_bias") and layer.w13_bias is not None: + gate_up_all = gate_up_all + layer.w13_bias.view(-1) + gate_up_all = gate_up_all.view(num_tokens, num_experts, gate_up_dim) + + # Apply gating activation: (T, E, gate+up) → (T, E, down_dim) + gate_up_act = act_fn(gate_up_all) + down_dim: int = gate_up_act.shape[-1] + + # ------------------------------------------------------------------ + # 2. Down — batched GEMM over experts. + # Permute to (E, T, down_dim) for bmm. + # w2_weight: (E, out, down_dim) → transpose → (E, down_dim, out) + # bmm: (E, T, down_dim) @ (E, down_dim, out) → (E, T, out) + # ------------------------------------------------------------------ + gate_up_t = gate_up_act.permute(1, 0, 2).contiguous() # (E, T, down) + down_all = torch.bmm( + gate_up_t, + layer.w2_weight.transpose(1, 2), + ) # (E, T, out_dim) + if hasattr(layer, "w2_bias") and layer.w2_bias is not None: + down_all = down_all + layer.w2_bias.unsqueeze(1) + down_all = down_all.permute(1, 0, 2) # (T, E, out_dim) + + # ------------------------------------------------------------------ + # 3. Weighted accumulation over top-k routed experts. + # output[t] += topk_weights[t,k] * down_all[t, topk_ids[t,k], :] + # ------------------------------------------------------------------ + out_dim: int = down_all.shape[-1] + output = torch.zeros( + num_tokens, out_dim, dtype=input.dtype, device=input.device + ) + t_idx = torch.arange(num_tokens, device=input.device) + if skip_weighted: + # Weight already baked into input — gather directly, no re-weighting. + expert_ids = topk_ids[:, 0].long() + output.add_(down_all[t_idx, expert_ids, :]) + else: + for k in range(top_k): + expert_ids = topk_ids[:, k].long() # (T,) + w_k = topk_weights[:, k].to(input.dtype) # (T,) + expert_out = down_all[t_idx, expert_ids, :] # (T, out) + output.addcmul_( + expert_out, w_k.unsqueeze(-1).expand_as(expert_out) + ) + + return output + def cpu_fused_moe_torch( layer_id: int, @@ -489,3 +583,4 @@ def cpu_fused_moe_torch( op_func=cpu_fused_moe_torch, mutates_args=["output"], ) + From 9b33af1a097f236fd79c757e58e3be5d09bc2c74 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Fri, 3 Jul 2026 06:09:30 -0500 Subject: [PATCH 24/43] decode improvements Signed-off-by: Akash kaothalkar --- vllm/_custom_ops.py | 22 +------------------ .../layers/mamba/ops/cpu/gdn_attention.py | 8 ++++--- .../layers/mamba/ops/ssd_combined.py | 21 +++++++++++++++++- 3 files changed, 26 insertions(+), 25 deletions(-) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 7b7c0e6a4b22..2740052801ac 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3395,27 +3395,6 @@ def causal_conv1d_fwd_cpu( ) -def causal_conv1d_update_cpu( - x: torch.Tensor, - conv_states: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None, - silu_activation: bool, - conv_state_indices: torch.Tensor | None, - is_vnni: bool, -) -> torch.Tensor: - return torch.ops._C.causal_conv1d_update_cpu( - x, - conv_states, - weight, - bias, - silu_activation, - None, - conv_state_indices, - -1, - is_vnni, - ) - class CPUDNNLGEMMHandler: def __init__(self) -> None: @@ -3870,3 +3849,4 @@ def _minimax_allreduce_rms_qk_fake( torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device), torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device), ) + diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index e5ef487ee9b5..bc1e0c120a40 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -100,12 +100,13 @@ def cpu_gdn_attention_core( if is_amx: decode_mixed_qkv = ops.causal_conv1d_update_cpu( x=decode_mixed_qkv, - conv_states=conv_state, + conv_state=conv_state, weight=layer.conv1d.weight, bias=layer.conv1d.bias, - silu_activation=layer.activation == "silu", + activation="silu" if layer.activation == "silu" else None, conv_state_indices=decode_state_indices, - is_vnni=True, + query_start_loc=None, + pad_slot_id=0, ) else: decode_conv_state = conv_state[decode_state_indices].contiguous() @@ -230,3 +231,4 @@ def register_cpu_gdn_attention_ops() -> None: fake_impl=cpu_gdn_attention_core_fake, ) _CPU_GDN_ATTENTION_OPS_REGISTERED = True + diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 4b2ad011f819..8a135e1f549f 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -158,6 +158,21 @@ def _mamba_chunk_scan_combined_fwd_cuda( @CustomOp.register("mamba_chunk_scan_combined_fwd") class MambaChunkScanCombinedFwdOp(CustomOp): + def __init__(self): + # Bypass dispatch_forward() / get_current_vllm_config() by directly + # selecting the forward method based on the platform. This op may be + # instantiated during Mamba2 kernel warm-up, which happens before the + # vLLM compilation config is available. + super(CustomOp, self).__init__() # nn.Module.__init__ only + from vllm.platforms import current_platform + + if current_platform.is_cpu(): + self._forward_method = self.forward_cpu + elif current_platform.is_rocm(): + self._forward_method = self.forward_hip + else: + self._forward_method = self.forward_cuda + def forward_native(self, *args, **kwargs): return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) @@ -168,10 +183,13 @@ def forward_cuda(self, *args, **kwargs): return _mamba_chunk_scan_combined_fwd_cuda(*args, **kwargs) -_mamba_chunk_scan_combined_fwd_op = MambaChunkScanCombinedFwdOp() +_mamba_chunk_scan_combined_fwd_op: "MambaChunkScanCombinedFwdOp | None" = None def _mamba_chunk_scan_combined_fwd(*args, **kwargs): + global _mamba_chunk_scan_combined_fwd_op + if _mamba_chunk_scan_combined_fwd_op is None: + _mamba_chunk_scan_combined_fwd_op = MambaChunkScanCombinedFwdOp() return _mamba_chunk_scan_combined_fwd_op(*args, **kwargs) @@ -246,3 +264,4 @@ def mamba_chunk_scan_combined_varlen( ) return varlen_states + From b39ec6f6f1b26585c5d269219ca0d7d9cbb4e596 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Fri, 3 Jul 2026 12:35:30 -0500 Subject: [PATCH 25/43] prefill method Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 70 ++++++++++ csrc/cpu/mamba_kernels.hpp | 124 ++++++++++++++++++ csrc/cpu/torch_bindings.cpp | 15 +++ vllm/_custom_ops.py | 27 ++++ .../layers/mamba/ops/cpu_fallbacks.py | 98 +++++++++----- 5 files changed, 301 insertions(+), 33 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 0d0a829c4eec..4824455a6040 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -225,3 +225,73 @@ void selective_state_update_cpu_impl( } +// --------------------------------------------------------------------------- +// mamba_chunk_scan_fwd_cpu +// --------------------------------------------------------------------------- +void mamba_chunk_scan_fwd_cpu_impl( + at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller + at::Tensor& final_states, // [batch, nheads, headdim, dstate] float32 contiguous + const at::Tensor& x, // [seqlen, nheads, headdim] + const at::Tensor& dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp) + const at::Tensor& A, // [nheads] float32 + const at::Tensor& B, // [seqlen, ngroups, dstate] + const at::Tensor& C, // [seqlen, ngroups, dstate] + const c10::optional& D, // [nheads] float32 (optional) + const c10::optional& z, // [seqlen, nheads, headdim] (optional) + const at::Tensor& cu_seqlens // [batch+1] int32 +) { + const at::ScalarType input_type = x.scalar_type(); + + auto ensure_contig = [input_type](const at::Tensor& t) -> at::Tensor { + at::Tensor r = (t.scalar_type() != input_type) ? t.to(input_type) : t; + return r.is_contiguous() ? r : r.contiguous(); + }; + at::Tensor x_in = ensure_contig(x); + at::Tensor B_in = ensure_contig(B); + at::Tensor C_in = ensure_contig(C); + at::Tensor z_in; + if (z.has_value() && z.value().defined()) z_in = ensure_contig(z.value()); + + // A and D are float32 model parameters, potentially broadcast-expanded. + // Strip trailing broadcast dims to get a contiguous (nheads,) array. + auto to_per_head_f32 = [](const at::Tensor& t) -> at::Tensor { + at::Tensor r = t; + while (r.dim() > 1) r = r.select(r.dim() - 1, 0); + if (r.scalar_type() != at::kFloat) r = r.to(at::kFloat); + return r.is_contiguous() ? r : r.contiguous(); + }; + at::Tensor A_f32 = to_per_head_f32(A); + at::Tensor D_f32; + if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value()); + + // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in Python. + at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous(); + if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat); + + at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous(); + + const int64_t batch = final_states.size(0); + const int64_t nheads = final_states.size(1); + const int64_t headdim = final_states.size(2); + const int64_t dstate = final_states.size(3); + const int64_t ngroups = B_in.size(1); + + TORCH_CHECK(final_states.is_contiguous(), + "mamba_chunk_scan_fwd_cpu: final_states must be contiguous"); + + VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] { + mamba_cpu::mamba_chunk_scan_fwd_kernel( + final_states.data_ptr(), + x_in.data_ptr(), + dt_c.data_ptr(), + A_f32.data_ptr(), + B_in.data_ptr(), + C_in.data_ptr(), + D_f32.defined() ? D_f32.data_ptr() : nullptr, + z_in.defined() ? z_in.data_ptr() : nullptr, + out.data_ptr(), + cu_int.data_ptr(), + batch, nheads, ngroups, headdim, dstate); + }); +} + diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index b81313dfec2a..8b632cee5638 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -259,5 +259,129 @@ inline void selective_state_update_kernel( } } +// --------------------------------------------------------------------------- +// mamba_chunk_scan_fwd +// +// Prefill SSM recurrence for Mamba2 / SSD models. +// +// Key difference from selective_state_update_kernel (decode path): +// - #pragma omp parallel for collapse(2) is OUTSIDE the time loop. +// Each thread owns a (batch, head) slice and runs the entire token +// sequence without any per-token OpenMP synchronisation overhead. +// For seqlen=256, this eliminates 256 thread-barrier launches per batch. +// +// `dt` arrives already processed (float32, after bias + softplus + clamp) +// to keep this kernel simple. Preprocessing is done in the Python wrapper. +// +// `states_ptr` points to the [batch, nheads, headdim, dstate] float32 output +// tensor, pre-initialised by the caller (zero or from initial_states). +// Each (b, h) slice is private to exactly one thread via collapse(2), so +// there are no write conflicts. +// +// D is treated as a scalar per head ([nheads] float32). +// --------------------------------------------------------------------------- +template +inline void mamba_chunk_scan_fwd_kernel( + float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32 + const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim] + const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed) + const float* __restrict__ A_ptr, // [nheads] f32 + const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate] + const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate] + const float* __restrict__ D_ptr, // [nheads] f32 (nullable) + const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable) + input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim] + const int32_t* __restrict__ cu_seqlens, // [batch+1] int32 + int64_t batch, int64_t nheads, int64_t ngroups, + int64_t headdim, int64_t dstate) { + + using input_vec_t = vec_op::vec_t; + constexpr int VEC_ELEM_NUM = 8; + + const int64_t nheads_per_group = nheads / ngroups; + // states layout: [batch, nheads, headdim, dstate] contiguous (caller guarantee) + const int64_t stride_s_b = nheads * headdim * dstate; + const int64_t stride_s_h = headdim * dstate; + // stride_s_d = dstate, stride_s_n = 1 + +#pragma omp parallel for collapse(2) schedule(dynamic) + for (int64_t b = 0; b < batch; ++b) { + for (int64_t h = 0; h < nheads; ++h) { + const int64_t seq_start = cu_seqlens[b]; + const int64_t seq_end = cu_seqlens[b + 1]; + const int64_t g = h / nheads_per_group; + + const float A_val = A_ptr[h]; + const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f; + + // Working state slice: states[b, h, :, :] — float32, headdim * dstate. + // Fits in L1/L2 for typical dims (e.g. 64*128*4 = 32 KB). + float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h; + + for (int64_t t = seq_start; t < seq_end; ++t) { + const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim; + const float* dt_h = dt_ptr + t * nheads + h; + const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate; + const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate; + const input_t* z_h = (z_ptr != nullptr) ? + z_ptr + t * nheads * headdim + h * headdim : nullptr; + input_t* out_h = out_ptr + t * nheads * headdim + h * headdim; + + const float dt_val = *dt_h; + const float dA_val = std::exp(A_val * dt_val); + const vec_op::FP32Vec8 dA_vec(dA_val); // broadcast scalar + const vec_op::FP32Vec8 dt_vec(dt_val); + + for (int64_t d = 0; d < headdim; ++d) { + const float x_val = static_cast(x_h[d]); + float* s_bhd = s_bh + d * dstate; // [dstate] contiguous float32 + + // Vectorised SSM update + readout over dstate: + // s_new = s * dA + x * dt * B + // y += s_new * C + int64_t n = 0; + vec_op::FP32Vec8 y_vec(0.0f); + const vec_op::FP32Vec8 x_vec(x_val); + + for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { + const vec_op::FP32Vec8 B_v((input_vec_t(B_g + n))); + const vec_op::FP32Vec8 C_v((input_vec_t(C_g + n))); + const vec_op::FP32Vec8 s_v(s_bhd + n); + + const vec_op::FP32Vec8 s_new = s_v * dA_vec + x_vec * dt_vec * B_v; + s_new.save(s_bhd + n); + y_vec = y_vec + s_new * C_v; + } + + float y_val = y_vec.reduce_sum(); + + // Scalar tail for remaining dstate elements + for (; n < dstate; ++n) { + const float B_n = static_cast(B_g[n]); + const float C_n = static_cast(C_g[n]); + const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n; + s_bhd[n] = s_new; + y_val += s_new * C_n; + } + + // D skip connection (scalar per head) + if (D_ptr != nullptr) y_val += x_val * D_val; + + // z gating: out = y * z * sigmoid(z) (SiLU) + if (z_h != nullptr) { + const float z_val = static_cast(z_h[d]); + const float sigmoid = (z_val >= 0.0f) ? + 1.0f / (1.0f + std::exp(-z_val)) : + std::exp(z_val) / (1.0f + std::exp(z_val)); + y_val *= z_val * sigmoid; + } + + out_h[d] = static_cast(y_val); + } + } + } + } +} + } // namespace mamba_cpu diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 3b62dd5b8baf..7a60ada32d7f 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -223,6 +223,13 @@ void selective_state_update_cpu_impl( const c10::optional& num_accepted_tokens, const c10::optional& cu_seqlens); +void mamba_chunk_scan_fwd_cpu_impl( + at::Tensor& out, at::Tensor& final_states, + const at::Tensor& x, const at::Tensor& dt, const at::Tensor& A, + const at::Tensor& B, const at::Tensor& C, + const c10::optional& D, const c10::optional& z, + const at::Tensor& cu_seqlens); + void init_cpu_memory_env(std::vector node_ids); namespace cpu_utils { @@ -611,6 +618,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor? num_accepted_tokens, Tensor? cu_seqlens) -> ()", &selective_state_update_cpu_impl); + ops.def( + "mamba_chunk_scan_fwd_cpu(" + "Tensor(a0!) out, Tensor(a1!) final_states, " + "Tensor x, Tensor dt, Tensor A, Tensor B, Tensor C, " + "Tensor? D, Tensor? z, Tensor cu_seqlens) -> ()", + &mamba_chunk_scan_fwd_cpu_impl); + ops.def("init_cpu_memory_env(SymInt[] node_ids) -> ()", &init_cpu_memory_env); // Speculative decoding kernels @@ -692,3 +706,4 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { } REGISTER_EXTENSION(TORCH_EXTENSION_NAME) + diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 2740052801ac..678591949f2d 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2189,6 +2189,33 @@ def selective_state_update_cpu( ) +def mamba_chunk_scan_fwd_cpu( + out: torch.Tensor, + final_states: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor | None, + z: torch.Tensor | None, + cu_seqlens: torch.Tensor, +) -> None: + """Prefill SSM scan kernel. out and final_states are written in-place.""" + torch.ops._C.mamba_chunk_scan_fwd_cpu( + out, + final_states, + x, + dt, + A, + B, + C, + D, + z, + cu_seqlens, + ) + + # ROCm skinny gemms def LLMM1(a: torch.Tensor, b: torch.Tensor, rows_per_block: int) -> torch.Tensor: return torch.ops._rocm_C.LLMM1(a, b, rows_per_block) diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py index c33a574395ce..6c886bc9e87e 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -3,6 +3,7 @@ import torch +import vllm._custom_ops as ops from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID @@ -325,11 +326,12 @@ def _mamba_chunk_scan_combined_fwd_cpu( ): seqlen, nheads, headdim = x.shape _, ngroups, dstate = B.shape - nheads_per_group = nheads // ngroups assert cu_seqlens is not None batch = cu_seqlens.size(0) - 1 + # Preprocess dt: apply bias, softplus, and clamp. + # Kept in Python to avoid duplicating this logic in C++. dt_f = dt.float() if dt_bias is not None: dt_f = dt_f + dt_bias.float().unsqueeze(0) @@ -338,53 +340,83 @@ def _mamba_chunk_scan_combined_fwd_cpu( if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) + # Allocate state output buffer (float32, contiguous — required by kernel). all_states = torch.zeros( batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device ) + if initial_states is not None: + all_states.copy_(initial_states.float()) + + # Use the C++ kernel when available (CPU build with mamba kernels compiled). + # Falls back to the pure-Python loop below if the op is not registered. + _use_cpp_kernel = hasattr(torch.ops._C, "mamba_chunk_scan_fwd_cpu") + + if _use_cpp_kernel: + # Ensure out is writable and contiguous. + if not out.is_contiguous(): + out = out.contiguous() + + # D: strip broadcast dims so the kernel sees (nheads,) float32. + D_1d = None + if D is not None: + d = D.float() + while d.dim() > 1: + d = d.select(d.dim() - 1, 0) + D_1d = d.contiguous() + + ops.mamba_chunk_scan_fwd_cpu( + out, + all_states, + x, + dt_f, + A, + B, + C, + D_1d, + z, + cu_seqlens.to(torch.int32), + ) + else: + # Pure-Python fallback (no compiled extension available). + for b_idx in range(batch): + seq_start = cu_seqlens[b_idx].item() + seq_end = cu_seqlens[b_idx + 1].item() - for b_idx in range(batch): - seq_start = cu_seqlens[b_idx].item() - seq_end = cu_seqlens[b_idx + 1].item() + state = all_states[b_idx] # shares storage — updated in-place - if initial_states is not None: - state = initial_states[b_idx].float() - else: - state = torch.zeros( - nheads, headdim, dstate, dtype=torch.float32, device=x.device - ) + for t in range(seq_start, seq_end): + x_t = x[t].float() + dt_t = dt_f[t] + A_val = A.float() - for t in range(seq_start, seq_end): - x_t = x[t].float() - dt_t = dt_f[t] - A_val = A.float() + dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) - dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) + B_expanded = B[t].float().repeat_interleave(nheads // ngroups, dim=0) + C_expanded = C[t].float().repeat_interleave(nheads // ngroups, dim=0) - B_expanded = B[t].float().repeat_interleave(nheads_per_group, dim=0) - C_expanded = C[t].float().repeat_interleave(nheads_per_group, dim=0) + xdt = x_t * dt_t.unsqueeze(-1) + dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) + state = state * dA + dBx - xdt = x_t * dt_t.unsqueeze(-1) - dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) - state = state * dA + dBx + y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) - y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) + if D is not None: + y = ( + y + x_t * D.float().unsqueeze(-1) + if D.dim() == 1 + else y + x_t * D.float() + ) - if D is not None: - y = ( - y + x_t * D.float().unsqueeze(-1) - if D.dim() == 1 - else y + x_t * D.float() - ) + if z is not None: + z_t = z[t].float() + y = y * z_t * torch.sigmoid(z_t) - if z is not None: - z_t = z[t].float() - y = y * z_t * torch.sigmoid(z_t) - - out[t] = y.to(out.dtype) + out[t] = y.to(out.dtype) - all_states[b_idx] = state.to(all_states.dtype) + all_states[b_idx] = state.to(all_states.dtype) out_dtype = state_dtype if state_dtype is not None else x.dtype all_states = all_states.to(out_dtype) return all_states + From fbb674a7f34f98724936e57bcf7f542d58ebbf7d Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Sat, 4 Jul 2026 01:40:00 -0500 Subject: [PATCH 26/43] error handling Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 2 ++ csrc/cpu/mamba_kernels.hpp | 2 +- .../layers/mamba/ops/cpu_fallbacks.py | 25 +++++++++++++------ 3 files changed, 21 insertions(+), 8 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 4824455a6040..3c240ae5d8b7 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -278,6 +278,8 @@ void mamba_chunk_scan_fwd_cpu_impl( TORCH_CHECK(final_states.is_contiguous(), "mamba_chunk_scan_fwd_cpu: final_states must be contiguous"); + TORCH_CHECK(out.is_contiguous(), + "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via raw data_ptr)"); VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] { mamba_cpu::mamba_chunk_scan_fwd_kernel( diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index 8b632cee5638..adcc23be1c43 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -304,7 +304,7 @@ inline void mamba_chunk_scan_fwd_kernel( const int64_t stride_s_h = headdim * dstate; // stride_s_d = dstate, stride_s_n = 1 -#pragma omp parallel for collapse(2) schedule(dynamic) +#pragma omp parallel for collapse(2) schedule(static) for (int64_t b = 0; b < batch; ++b) { for (int64_t h = 0; h < nheads; ++h) { const int64_t seq_start = cu_seqlens[b]; diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py index 6c886bc9e87e..2db5c7ad7286 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -352,16 +352,23 @@ def _mamba_chunk_scan_combined_fwd_cpu( _use_cpp_kernel = hasattr(torch.ops._C, "mamba_chunk_scan_fwd_cpu") if _use_cpp_kernel: - # Ensure out is writable and contiguous. - if not out.is_contiguous(): - out = out.contiguous() + # out must be contiguous — the C++ kernel writes via raw data_ptr(). + # Creating a contiguous copy here would discard the results, so we + # require the caller to pass a contiguous tensor. + assert out.is_contiguous(), ( + "_mamba_chunk_scan_combined_fwd_cpu: `out` must be " + "pre-allocated as a contiguous tensor" + ) - # D: strip broadcast dims so the kernel sees (nheads,) float32. + # D: strip trailing dims that are broadcast (stride == 0) so the + # kernel sees a (nheads,) float32 array. Only peel dims whose + # stride is 0 (broadcast-expanded); genuine multi-dim D is passed + # as-is and flattened by the C++ wrapper. D_1d = None if D is not None: d = D.float() - while d.dim() > 1: - d = d.select(d.dim() - 1, 0) + while d.dim() > 1 and d.stride(-1) == 0: + d = d.squeeze(-1) D_1d = d.contiguous() ops.mamba_chunk_scan_fwd_cpu( @@ -382,7 +389,10 @@ def _mamba_chunk_scan_combined_fwd_cpu( seq_start = cu_seqlens[b_idx].item() seq_end = cu_seqlens[b_idx + 1].item() - state = all_states[b_idx] # shares storage — updated in-place + # Note: basic indexing returns a view, but arithmetic ops + # below (state * dA + dBx) produce new tensors, so state is + # reassigned each token. all_states[b_idx] is committed at line 416. + state = all_states[b_idx].clone() # local working copy for t in range(seq_start, seq_end): x_t = x[t].float() @@ -420,3 +430,4 @@ def _mamba_chunk_scan_combined_fwd_cpu( return all_states + From ba5b43dbbf79545a77c695090897dc1788e1bfa8 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Sun, 5 Jul 2026 04:16:58 -0500 Subject: [PATCH 27/43] pre-commit fix Signed-off-by: Akash kaothalkar --- csrc/cpu/cpu_types_vsx.hpp | 1 - csrc/cpu/mamba_cpu.cpp | 150 ++++++++---------- csrc/cpu/mamba_kernels.hpp | 125 +++++++-------- csrc/cpu/torch_bindings.cpp | 14 +- scratch/verify_cpu_ssu.py | 43 ++--- vllm/_custom_ops.py | 2 - vllm/config/mamba.py | 1 - .../layers/fused_moe/cpu_fused_moe.py | 14 +- .../layers/mamba/mamba_mixer2.py | 5 +- .../layers/mamba/ops/causal_conv1d.py | 1 - .../layers/mamba/ops/cpu/gdn_attention.py | 1 - .../layers/mamba/ops/cpu_fallbacks.py | 2 - .../layers/mamba/ops/mamba_ssm.py | 36 ++++- .../layers/mamba/ops/ssd_combined.py | 1 - .../layers/mamba/ops/ssu_dispatch.py | 6 +- vllm/model_executor/layers/utils.py | 9 +- 16 files changed, 203 insertions(+), 208 deletions(-) diff --git a/csrc/cpu/cpu_types_vsx.hpp b/csrc/cpu/cpu_types_vsx.hpp index 323ed45ff6b4..57120679ea04 100644 --- a/csrc/cpu/cpu_types_vsx.hpp +++ b/csrc/cpu/cpu_types_vsx.hpp @@ -1015,4 +1015,3 @@ struct INT8Vec64 { }; } // namespace vec_op #endif - diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 3c240ae5d8b7..8c95947a9424 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -16,13 +16,11 @@ // causal_conv1d_update // --------------------------------------------------------------------------- at::Tensor causal_conv1d_update_cpu_impl( - at::Tensor& x, - at::Tensor& conv_state, const at::Tensor& weight, + at::Tensor& x, at::Tensor& conv_state, const at::Tensor& weight, const c10::optional& bias, const c10::optional& activation, const c10::optional& conv_state_indices, const c10::optional& query_start_loc, int64_t pad_slot_id) { - bool do_silu = false; if (activation.has_value()) { const std::string& act = activation.value(); @@ -43,27 +41,28 @@ at::Tensor causal_conv1d_update_cpu_impl( // state_c and conv_state may be non-contiguous — that is intentional. // Weight: coerce to same dtype if needed (should match in practice) - at::Tensor w_c = (weight.scalar_type() != dtype) ? - weight.to(dtype).contiguous() : - (weight.is_contiguous() ? weight : weight.contiguous()); + at::Tensor w_c = + (weight.scalar_type() != dtype) + ? weight.to(dtype).contiguous() + : (weight.is_contiguous() ? weight : weight.contiguous()); // Bias stays float32 (small scalar, used only for fp32 accumulation) at::Tensor bias_f32; if (bias.has_value() && bias.value().defined()) bias_f32 = bias.value().to(at::kFloat).contiguous(); - int64_t batch = x_c.size(0); - int64_t dim = x_c.size(1); - int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1; - int64_t width = w_c.size(1); + int64_t batch = x_c.size(0); + int64_t dim = x_c.size(1); + int64_t seqlen = (x_c.dim() == 3) ? x_c.size(2) : 1; + int64_t width = w_c.size(1); int64_t state_len = state_c.size(2); - // Extract strides — works for contiguous AND non-contiguous (transposed) state. - // stride(0): between cache slots (e.g. num_slots × dim × width-1 in contiguous) - // stride(1): between conv channels (dim stride) - // stride(2): between state elements (=1 when contiguous, =dim when transposed) - int64_t stride_s_slot = state_c.stride(0); - int64_t stride_s_dim = state_c.stride(1); + // Extract strides — works for contiguous AND non-contiguous (transposed) + // state. stride(0): between cache slots (e.g. num_slots × dim × width-1 in + // contiguous) stride(1): between conv channels (dim stride) stride(2): + // between state elements (=1 when contiguous, =dim when transposed) + int64_t stride_s_slot = state_c.stride(0); + int64_t stride_s_dim = state_c.stride(1); int64_t stride_s_state = state_c.stride(2); at::Tensor out = at::empty_like(x_c); // native dtype, no float32 alloc @@ -77,14 +76,12 @@ at::Tensor causal_conv1d_update_cpu_impl( VLLM_DISPATCH_FLOATING_TYPES(dtype, "causal_conv1d_update", [&] { mamba_cpu::causal_conv1d_update_kernel( - x_c.data_ptr(), - state_c.data_ptr(), - stride_s_slot, stride_s_dim, stride_s_state, - w_c.data_ptr(), + x_c.data_ptr(), state_c.data_ptr(), stride_s_slot, + stride_s_dim, stride_s_state, w_c.data_ptr(), bias_f32.defined() ? bias_f32.data_ptr() : nullptr, - out.data_ptr(), - cache_idx_ptr, static_cast(pad_slot_id), - batch, dim, seqlen, width, state_len, do_silu); + out.data_ptr(), cache_idx_ptr, + static_cast(pad_slot_id), batch, dim, seqlen, width, state_len, + do_silu); }); // Write back only when a type-conversion copy was made. @@ -94,28 +91,21 @@ at::Tensor causal_conv1d_update_cpu_impl( return out; } - // --------------------------------------------------------------------------- // selective_state_update // --------------------------------------------------------------------------- void selective_state_update_cpu_impl( at::Tensor& state, // (nstates, nheads, dim, dstate) const at::Tensor& x, // (N, nheads, dim) - const at::Tensor& dt, - const at::Tensor& A, - const at::Tensor& B, - const at::Tensor& C, - const c10::optional& D, + const at::Tensor& dt, const at::Tensor& A, const at::Tensor& B, + const at::Tensor& C, const c10::optional& D, const c10::optional& z, - const c10::optional& dt_bias, - bool dt_softplus, + const c10::optional& dt_bias, bool dt_softplus, const c10::optional& state_batch_indices, const c10::optional& dst_state_batch_indices, - int64_t null_block_id, - at::Tensor& out, + int64_t null_block_id, at::Tensor& out, const c10::optional& num_accepted_tokens, - const c10::optional& cu_seqlens -) { + const c10::optional& cu_seqlens) { at::ScalarType state_type = state.scalar_type(); at::ScalarType input_type = x.scalar_type(); @@ -148,7 +138,7 @@ void selective_state_update_cpu_impl( return r.is_contiguous() ? r : r.contiguous(); }; - at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32 + at::Tensor A_f32 = to_per_head_1d_f32(A); // (nheads,) float32 at::Tensor D_f32, dt_bias_f32; if (D.has_value() && D.value().defined()) D_f32 = to_per_head_1d_f32(D.value()); @@ -185,14 +175,16 @@ void selective_state_update_cpu_impl( int64_t stride_out_h = out.stride(1); // Optional index pointers - auto get_int32_ptr = [](const c10::optional& opt) -> const int32_t* { - return (opt.has_value() && opt.value().defined()) ? - opt.value().data_ptr() : nullptr; + auto get_int32_ptr = + [](const c10::optional& opt) -> const int32_t* { + return (opt.has_value() && opt.value().defined()) + ? opt.value().data_ptr() + : nullptr; }; - const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices); + const int32_t* sbi_ptr = get_int32_ptr(state_batch_indices); const int32_t* dsbi_ptr = get_int32_ptr(dst_state_batch_indices); - const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens); - const int32_t* csl_ptr = get_int32_ptr(cu_seqlens); + const int32_t* nat_ptr = get_int32_ptr(num_accepted_tokens); + const int32_t* csl_ptr = get_int32_ptr(cu_seqlens); // Dispatch on (state_t, input_t, out_t): write directly into `out` // without any intermediate float32 buffer. @@ -203,42 +195,37 @@ void selective_state_update_cpu_impl( VLLM_DISPATCH_FLOATING_TYPES(out.scalar_type(), "ssu_out", [&] { using out_t = scalar_t; mamba_cpu::selective_state_update_kernel( - state.data_ptr(), - stride_state_n, stride_state_h, stride_state_d, - x_in.data_ptr(), - stride_x_n, stride_x_h, - dt_f32.data_ptr(), - stride_dt_n, - A_f32.data_ptr(), - B_in.data_ptr(), C_in.data_ptr(), - stride_BC_n, stride_BC_g, - D_f32.defined() ? D_f32.data_ptr() : nullptr, + state.data_ptr(), stride_state_n, stride_state_h, + stride_state_d, x_in.data_ptr(), stride_x_n, stride_x_h, + dt_f32.data_ptr(), stride_dt_n, A_f32.data_ptr(), + B_in.data_ptr(), C_in.data_ptr(), stride_BC_n, + stride_BC_g, D_f32.defined() ? D_f32.data_ptr() : nullptr, z_in.defined() ? z_in.data_ptr() : nullptr, dt_bias_f32.defined() ? dt_bias_f32.data_ptr() : nullptr, - out.data_ptr(), - stride_out_n, stride_out_h, - sbi_ptr, dsbi_ptr, static_cast(null_block_id), - nat_ptr, csl_ptr, N, nheads, ngroups, dim, dstate, dt_softplus); + out.data_ptr(), stride_out_n, stride_out_h, sbi_ptr, + dsbi_ptr, static_cast(null_block_id), nat_ptr, csl_ptr, N, + nheads, ngroups, dim, dstate, dt_softplus); }); }); }); } - // --------------------------------------------------------------------------- // mamba_chunk_scan_fwd_cpu // --------------------------------------------------------------------------- void mamba_chunk_scan_fwd_cpu_impl( - at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller - at::Tensor& final_states, // [batch, nheads, headdim, dstate] float32 contiguous - const at::Tensor& x, // [seqlen, nheads, headdim] - const at::Tensor& dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp) - const at::Tensor& A, // [nheads] float32 - const at::Tensor& B, // [seqlen, ngroups, dstate] - const at::Tensor& C, // [seqlen, ngroups, dstate] - const c10::optional& D, // [nheads] float32 (optional) - const c10::optional& z, // [seqlen, nheads, headdim] (optional) - const at::Tensor& cu_seqlens // [batch+1] int32 + at::Tensor& out, // [seqlen, nheads, headdim] — pre-allocated by caller + at::Tensor& + final_states, // [batch, nheads, headdim, dstate] float32 contiguous + const at::Tensor& x, // [seqlen, nheads, headdim] + const at::Tensor& + dt, // [seqlen, nheads] float32 (preprocessed: bias+softplus+clamp) + const at::Tensor& A, // [nheads] float32 + const at::Tensor& B, // [seqlen, ngroups, dstate] + const at::Tensor& C, // [seqlen, ngroups, dstate] + const c10::optional& D, // [nheads] float32 (optional) + const c10::optional& z, // [seqlen, nheads, headdim] (optional) + const at::Tensor& cu_seqlens // [batch+1] int32 ) { const at::ScalarType input_type = x.scalar_type(); @@ -264,36 +251,33 @@ void mamba_chunk_scan_fwd_cpu_impl( at::Tensor D_f32; if (D.has_value() && D.value().defined()) D_f32 = to_per_head_f32(D.value()); - // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in Python. + // dt: [seqlen, nheads] float32 — caller has applied bias+softplus+clamp in + // Python. at::Tensor dt_c = dt.is_contiguous() ? dt : dt.contiguous(); if (dt_c.scalar_type() != at::kFloat) dt_c = dt_c.to(at::kFloat); at::Tensor cu_int = cu_seqlens.to(at::kInt).contiguous(); - const int64_t batch = final_states.size(0); - const int64_t nheads = final_states.size(1); + const int64_t batch = final_states.size(0); + const int64_t nheads = final_states.size(1); const int64_t headdim = final_states.size(2); - const int64_t dstate = final_states.size(3); + const int64_t dstate = final_states.size(3); const int64_t ngroups = B_in.size(1); TORCH_CHECK(final_states.is_contiguous(), - "mamba_chunk_scan_fwd_cpu: final_states must be contiguous"); + "mamba_chunk_scan_fwd_cpu: final_states must be contiguous"); TORCH_CHECK(out.is_contiguous(), - "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via raw data_ptr)"); + "mamba_chunk_scan_fwd_cpu: out must be contiguous (writes via " + "raw data_ptr)"); VLLM_DISPATCH_FLOATING_TYPES(input_type, "mamba_chunk_scan_fwd_cpu", [&] { mamba_cpu::mamba_chunk_scan_fwd_kernel( - final_states.data_ptr(), - x_in.data_ptr(), - dt_c.data_ptr(), - A_f32.data_ptr(), - B_in.data_ptr(), - C_in.data_ptr(), + final_states.data_ptr(), x_in.data_ptr(), + dt_c.data_ptr(), A_f32.data_ptr(), + B_in.data_ptr(), C_in.data_ptr(), D_f32.defined() ? D_f32.data_ptr() : nullptr, z_in.defined() ? z_in.data_ptr() : nullptr, - out.data_ptr(), - cu_int.data_ptr(), - batch, nheads, ngroups, headdim, dstate); + out.data_ptr(), cu_int.data_ptr(), batch, nheads, + ngroups, headdim, dstate); }); } - diff --git a/csrc/cpu/mamba_kernels.hpp b/csrc/cpu/mamba_kernels.hpp index adcc23be1c43..722dad97e3ca 100644 --- a/csrc/cpu/mamba_kernels.hpp +++ b/csrc/cpu/mamba_kernels.hpp @@ -29,16 +29,14 @@ namespace mamba_cpu { // // When stride_s_state == 1 (contiguous), the memmove fast path is used. // --------------------------------------------------------------------------- -template +template inline void causal_conv1d_update_kernel( - const scalar_t* __restrict__ x_ptr, - scalar_t* __restrict__ state_ptr, + const scalar_t* __restrict__ x_ptr, scalar_t* __restrict__ state_ptr, int64_t stride_s_slot, int64_t stride_s_dim, int64_t stride_s_state, const scalar_t* __restrict__ weight_ptr, const float* __restrict__ bias_ptr, scalar_t* __restrict__ out_ptr, const int32_t* __restrict__ cache_idxs, int32_t pad_slot_id, int64_t batch, int64_t dim, int64_t seqlen, int64_t width, int64_t state_len, bool do_silu) { - #pragma omp parallel for for (int64_t b = 0; b < batch; ++b) { int64_t cache_idx = (cache_idxs != nullptr) ? cache_idxs[b] : b; @@ -52,7 +50,7 @@ inline void causal_conv1d_update_kernel( for (int64_t d = 0; d < dim; ++d) { float x_val = static_cast(x_b[d * seqlen]); - scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state + scalar_t* sd = s_base + d * stride_s_dim; // start of this dim's state const scalar_t* w = weight_ptr + d * width; // Accumulate in float32 for precision @@ -68,8 +66,7 @@ inline void causal_conv1d_update_kernel( if (stride_s_state == 1) { if (state_len > 1) std::memmove(sd, sd + 1, (state_len - 1) * sizeof(scalar_t)); - if (state_len > 0) - sd[state_len - 1] = static_cast(x_val); + if (state_len > 0) sd[state_len - 1] = static_cast(x_val); } else { for (int64_t k = 0; k < state_len - 1; ++k) sd[k * stride_s_state] = sd[(k + 1) * stride_s_state]; @@ -78,9 +75,8 @@ inline void causal_conv1d_update_kernel( } if (do_silu) { - float sigmoid = (acc >= 0) ? - 1.0f / (1.0f + std::exp(-acc)) : - std::exp(acc) / (1.0f + std::exp(acc)); + float sigmoid = (acc >= 0) ? 1.0f / (1.0f + std::exp(-acc)) + : std::exp(acc) / (1.0f + std::exp(acc)); acc *= sigmoid; } out_b[d * seqlen] = static_cast(acc); @@ -107,32 +103,26 @@ inline void causal_conv1d_update_kernel( // --------------------------------------------------------------------------- template inline void selective_state_update_kernel( - state_t* __restrict__ state_ptr, - int64_t stride_state_n, int64_t stride_state_h, int64_t stride_state_d, - const input_t* __restrict__ x_ptr, - int64_t stride_x_n, int64_t stride_x_h, + state_t* __restrict__ state_ptr, int64_t stride_state_n, + int64_t stride_state_h, int64_t stride_state_d, + const input_t* __restrict__ x_ptr, int64_t stride_x_n, int64_t stride_x_h, // dt: (N, nheads) — scalar per head, NOT expanded to head_dim - const float* __restrict__ dt_ptr, - int64_t stride_dt_n, + const float* __restrict__ dt_ptr, int64_t stride_dt_n, // A: (nheads,) float32 — scalar per head - const float* __restrict__ A_ptr, - const input_t* __restrict__ B_ptr, const input_t* __restrict__ C_ptr, - int64_t stride_BC_n, int64_t stride_BC_g, + const float* __restrict__ A_ptr, const input_t* __restrict__ B_ptr, + const input_t* __restrict__ C_ptr, int64_t stride_BC_n, int64_t stride_BC_g, // D: (nheads,) float32 — scalar per head (nullptr if not used) const float* __restrict__ D_ptr, // z: same shape as x (optional) const input_t* __restrict__ z_ptr, // dt_bias: (nheads,) float32 — scalar per head (nullptr if not used) - const float* __restrict__ dt_bias_ptr, - out_t* __restrict__ out_ptr, + const float* __restrict__ dt_bias_ptr, out_t* __restrict__ out_ptr, int64_t stride_out_n, int64_t stride_out_h, const int32_t* __restrict__ state_batch_indices, const int32_t* __restrict__ dst_state_batch_indices, int32_t null_block_id, const int32_t* __restrict__ num_accepted_tokens, - const int32_t* __restrict__ cu_seqlens, - int64_t N, int64_t nheads, int64_t ngroups, int64_t dim, int64_t dstate, - bool dt_softplus) { - + const int32_t* __restrict__ cu_seqlens, int64_t N, int64_t nheads, + int64_t ngroups, int64_t dim, int64_t dstate, bool dt_softplus) { using state_vec_t = vec_op::vec_t; using input_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = 8; @@ -149,13 +139,16 @@ inline void selective_state_update_kernel( seq_len = 1; } - int64_t state_read_idx = (state_batch_indices != nullptr) ? - state_batch_indices[seq_idx] : seq_idx; + int64_t state_read_idx = (state_batch_indices != nullptr) + ? state_batch_indices[seq_idx] + : seq_idx; if (state_read_idx == null_block_id) continue; - int64_t state_write_idx = (num_accepted_tokens == nullptr) ? - ((dst_state_batch_indices != nullptr) ? - dst_state_batch_indices[seq_idx] : state_read_idx) : -1; + int64_t state_write_idx = (num_accepted_tokens == nullptr) + ? ((dst_state_batch_indices != nullptr) + ? dst_state_batch_indices[seq_idx] + : state_read_idx) + : -1; state_t* s = state_ptr + state_read_idx * stride_state_n; @@ -183,11 +176,12 @@ inline void selective_state_update_kernel( if (dt_softplus) { dt_val = (dt_val <= 20.0f) ? std::log1p(std::exp(dt_val)) : dt_val; } - const float A_val = A_ptr[h]; // scalar: same for all dim, dstate + const float A_val = A_ptr[h]; // scalar: same for all dim, dstate const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f; - const input_t* z_h = (z_ptr != nullptr) ? - z_ptr + token_idx * stride_x_n + h * stride_x_h : nullptr; + const input_t* z_h = + (z_ptr != nullptr) ? z_ptr + token_idx * stride_x_n + h * stride_x_h + : nullptr; vec_op::FP32Vec8 dt_vec(dt_val); // dA = exp(A * dt): A and dt are SCALARS per head, so compute once @@ -233,16 +227,17 @@ inline void selective_state_update_kernel( if (D_ptr != nullptr) out_val += x_val * D_val; if (z_h != nullptr) { float z_val = static_cast(z_h[d]); - float sigmoid = (z_val >= 0) ? - 1.0f / (1.0f + std::exp(-z_val)) : - std::exp(z_val) / (1.0f + std::exp(z_val)); + float sigmoid = (z_val >= 0) + ? 1.0f / (1.0f + std::exp(-z_val)) + : std::exp(z_val) / (1.0f + std::exp(z_val)); out_val *= z_val * sigmoid; } out_h[d] = static_cast(out_val); } } - if (num_accepted_tokens != nullptr && dst_state_batch_indices != nullptr) { + if (num_accepted_tokens != nullptr && + dst_state_batch_indices != nullptr) { int64_t token_dst_idx = dst_state_batch_indices[seq_idx * seq_len + t]; if (token_dst_idx != null_block_id && token_dst_idx != state_read_idx) { state_t* dst_s = state_ptr + token_dst_idx * stride_state_n; @@ -282,24 +277,24 @@ inline void selective_state_update_kernel( // --------------------------------------------------------------------------- template inline void mamba_chunk_scan_fwd_kernel( - float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32 - const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim] - const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed) - const float* __restrict__ A_ptr, // [nheads] f32 - const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate] - const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate] - const float* __restrict__ D_ptr, // [nheads] f32 (nullable) - const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable) - input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim] + float* __restrict__ states_ptr, // [batch, nheads, headdim, dstate] f32 + const input_t* __restrict__ x_ptr, // [seqlen, nheads, headdim] + const float* __restrict__ dt_ptr, // [seqlen, nheads] f32 (preprocessed) + const float* __restrict__ A_ptr, // [nheads] f32 + const input_t* __restrict__ B_ptr, // [seqlen, ngroups, dstate] + const input_t* __restrict__ C_ptr, // [seqlen, ngroups, dstate] + const float* __restrict__ D_ptr, // [nheads] f32 (nullable) + const input_t* __restrict__ z_ptr, // [seqlen, nheads, headdim] (nullable) + input_t* __restrict__ out_ptr, // [seqlen, nheads, headdim] const int32_t* __restrict__ cu_seqlens, // [batch+1] int32 - int64_t batch, int64_t nheads, int64_t ngroups, - int64_t headdim, int64_t dstate) { - + int64_t batch, int64_t nheads, int64_t ngroups, int64_t headdim, + int64_t dstate) { using input_vec_t = vec_op::vec_t; constexpr int VEC_ELEM_NUM = 8; const int64_t nheads_per_group = nheads / ngroups; - // states layout: [batch, nheads, headdim, dstate] contiguous (caller guarantee) + // states layout: [batch, nheads, headdim, dstate] contiguous (caller + // guarantee) const int64_t stride_s_b = nheads * headdim * dstate; const int64_t stride_s_h = headdim * dstate; // stride_s_d = dstate, stride_s_n = 1 @@ -308,8 +303,8 @@ inline void mamba_chunk_scan_fwd_kernel( for (int64_t b = 0; b < batch; ++b) { for (int64_t h = 0; h < nheads; ++h) { const int64_t seq_start = cu_seqlens[b]; - const int64_t seq_end = cu_seqlens[b + 1]; - const int64_t g = h / nheads_per_group; + const int64_t seq_end = cu_seqlens[b + 1]; + const int64_t g = h / nheads_per_group; const float A_val = A_ptr[h]; const float D_val = (D_ptr != nullptr) ? D_ptr[h] : 0.0f; @@ -319,12 +314,13 @@ inline void mamba_chunk_scan_fwd_kernel( float* s_bh = states_ptr + b * stride_s_b + h * stride_s_h; for (int64_t t = seq_start; t < seq_end; ++t) { - const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim; - const float* dt_h = dt_ptr + t * nheads + h; - const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate; - const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate; - const input_t* z_h = (z_ptr != nullptr) ? - z_ptr + t * nheads * headdim + h * headdim : nullptr; + const input_t* x_h = x_ptr + t * nheads * headdim + h * headdim; + const float* dt_h = dt_ptr + t * nheads + h; + const input_t* B_g = B_ptr + t * ngroups * dstate + g * dstate; + const input_t* C_g = C_ptr + t * ngroups * dstate + g * dstate; + const input_t* z_h = (z_ptr != nullptr) + ? z_ptr + t * nheads * headdim + h * headdim + : nullptr; input_t* out_h = out_ptr + t * nheads * headdim + h * headdim; const float dt_val = *dt_h; @@ -357,8 +353,8 @@ inline void mamba_chunk_scan_fwd_kernel( // Scalar tail for remaining dstate elements for (; n < dstate; ++n) { - const float B_n = static_cast(B_g[n]); - const float C_n = static_cast(C_g[n]); + const float B_n = static_cast(B_g[n]); + const float C_n = static_cast(C_g[n]); const float s_new = s_bhd[n] * dA_val + x_val * dt_val * B_n; s_bhd[n] = s_new; y_val += s_new * C_n; @@ -370,9 +366,9 @@ inline void mamba_chunk_scan_fwd_kernel( // z gating: out = y * z * sigmoid(z) (SiLU) if (z_h != nullptr) { const float z_val = static_cast(z_h[d]); - const float sigmoid = (z_val >= 0.0f) ? - 1.0f / (1.0f + std::exp(-z_val)) : - std::exp(z_val) / (1.0f + std::exp(z_val)); + const float sigmoid = + (z_val >= 0.0f) ? 1.0f / (1.0f + std::exp(-z_val)) + : std::exp(z_val) / (1.0f + std::exp(z_val)); y_val *= z_val * sigmoid; } @@ -383,5 +379,4 @@ inline void mamba_chunk_scan_fwd_kernel( } } -} // namespace mamba_cpu - +} // namespace mamba_cpu diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 7a60ada32d7f..1a0c8411cce1 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -223,12 +223,13 @@ void selective_state_update_cpu_impl( const c10::optional& num_accepted_tokens, const c10::optional& cu_seqlens); -void mamba_chunk_scan_fwd_cpu_impl( - at::Tensor& out, at::Tensor& final_states, - const at::Tensor& x, const at::Tensor& dt, const at::Tensor& A, - const at::Tensor& B, const at::Tensor& C, - const c10::optional& D, const c10::optional& z, - const at::Tensor& cu_seqlens); +void mamba_chunk_scan_fwd_cpu_impl(at::Tensor& out, at::Tensor& final_states, + const at::Tensor& x, const at::Tensor& dt, + const at::Tensor& A, const at::Tensor& B, + const at::Tensor& C, + const c10::optional& D, + const c10::optional& z, + const at::Tensor& cu_seqlens); void init_cpu_memory_env(std::vector node_ids); @@ -706,4 +707,3 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { } REGISTER_EXTENSION(TORCH_EXTENSION_NAME) - diff --git a/scratch/verify_cpu_ssu.py b/scratch/verify_cpu_ssu.py index c5ff0fe276bc..4b151ba83f56 100644 --- a/scratch/verify_cpu_ssu.py +++ b/scratch/verify_cpu_ssu.py @@ -1,7 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import time + import torch import torch.nn.functional as F + from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update -import time + def selective_state_update_ref( state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False @@ -26,45 +31,46 @@ def selective_state_update_ref( if dt_bias is not None and dt_bias.dim() == 1: dt_bias = dt_bias.unsqueeze(0) batch, nheads, dim, dstate = state.shape - + dt = dt.float() if dt_bias is not None: dt = dt + dt_bias.float() dt = F.softplus(dt) if dt_softplus else dt - + dA = torch.exp(dt.unsqueeze(-1) * A.float()) # (batch, nheads, dim, dstate) ngroups = B.shape[1] B = B.repeat_interleave(nheads // ngroups, dim=1) C = C.repeat_interleave(nheads // ngroups, dim=1) - + dBx = (dt.unsqueeze(-1) * B.float().unsqueeze(2)) * x.float().unsqueeze(-1) - + state.copy_(state.float() * dA + dBx) - + out = torch.einsum("bhdn,bhn->bhd", state.float(), C.float()) if D is not None: - out += (x.float() * D.float()) - + out += x.float() * D.float() + if z is not None: out = out * F.silu(z.float()) - + if not has_heads: out = out.squeeze(1) return out.to(x.dtype) + def test_cpu_ssu(): device = "cpu" torch.manual_seed(42) - + # Mamba-2 style dimensions batch_size = 2 nheads = 128 dim = 64 dstate = 128 itype = torch.bfloat16 - + print(f"Testing CPU selective_state_update with {itype}...") - + state = torch.randn(batch_size, nheads, dim, dstate, dtype=itype, device=device) x = torch.randn(batch_size, nheads, dim, device=device, dtype=itype) dt = torch.randn(batch_size, nheads, dim, device=device, dtype=itype) @@ -74,34 +80,35 @@ def test_cpu_ssu(): C = torch.randn(batch_size, 1, dstate, device=device, dtype=itype) D = torch.randn(nheads, dim, device=device) z = torch.randn_like(x) - + out = torch.empty_like(x) state_ref = state.detach().clone() - + start = time.time() selective_state_update( state, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, dt_softplus=True, out=out ) end = time.time() print(f"Kernel time: {end - start:.4f}s") - + start = time.time() out_ref = selective_state_update_ref( state_ref, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, dt_softplus=True ) end = time.time() print(f"Ref time: {end - start:.4f}s") - + state_diff = (state.float() - state_ref.float()).abs().max().item() out_diff = (out.float() - out_ref.float()).abs().max().item() - + print(f"State max diff: {state_diff}") print(f"Out max diff: {out_diff}") - + if state_diff < 1e-2 and out_diff < 1e-2: print("SUCCESS") else: print("FAILURE") + if __name__ == "__main__": test_cpu_ssu() diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 678591949f2d..54f977605e73 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3422,7 +3422,6 @@ def causal_conv1d_fwd_cpu( ) - class CPUDNNLGEMMHandler: def __init__(self) -> None: self.handler_tensor: torch.Tensor | None = None @@ -3876,4 +3875,3 @@ def _minimax_allreduce_rms_qk_fake( torch.empty([token_num, q_size], dtype=qkv.dtype, device=qkv.device), torch.empty([token_num, kv_size], dtype=qkv.dtype, device=qkv.device), ) - diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py index 421e5a942cb5..fa556723e846 100644 --- a/vllm/config/mamba.py +++ b/vllm/config/mamba.py @@ -82,4 +82,3 @@ def __post_init__(self): "specify `--enable-mamba-cache-stochastic-rounding`, " "or set `--mamba-backend flashinfer`." ) - diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index b8ded2fb14dc..27a9da5fbfe3 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -475,7 +475,7 @@ def forward_batched_gemm( num_tokens: int = input.shape[0] hidden_dim: int = input.shape[1] num_experts: int = layer.w13_weight.shape[0] - gate_up_dim: int = layer.w13_weight.shape[1] # gate + up concatenated + gate_up_dim: int = layer.w13_weight.shape[1] # gate + up concatenated top_k: int = topk_ids.shape[1] act_fn = _CPU_MOE_ACT_FN[activation] @@ -493,7 +493,6 @@ def forward_batched_gemm( # Apply gating activation: (T, E, gate+up) → (T, E, down_dim) gate_up_act = act_fn(gate_up_all) - down_dim: int = gate_up_act.shape[-1] # ------------------------------------------------------------------ # 2. Down — batched GEMM over experts. @@ -501,7 +500,7 @@ def forward_batched_gemm( # w2_weight: (E, out, down_dim) → transpose → (E, down_dim, out) # bmm: (E, T, down_dim) @ (E, down_dim, out) → (E, T, out) # ------------------------------------------------------------------ - gate_up_t = gate_up_act.permute(1, 0, 2).contiguous() # (E, T, down) + gate_up_t = gate_up_act.permute(1, 0, 2).contiguous() # (E, T, down) down_all = torch.bmm( gate_up_t, layer.w2_weight.transpose(1, 2), @@ -525,12 +524,10 @@ def forward_batched_gemm( output.add_(down_all[t_idx, expert_ids, :]) else: for k in range(top_k): - expert_ids = topk_ids[:, k].long() # (T,) + expert_ids = topk_ids[:, k].long() # (T,) w_k = topk_weights[:, k].to(input.dtype) # (T,) - expert_out = down_all[t_idx, expert_ids, :] # (T, out) - output.addcmul_( - expert_out, w_k.unsqueeze(-1).expand_as(expert_out) - ) + expert_out = down_all[t_idx, expert_ids, :] # (T, out) + output.addcmul_(expert_out, w_k.unsqueeze(-1).expand_as(expert_out)) return output @@ -596,4 +593,3 @@ def cpu_fused_moe_torch( op_func=cpu_fused_moe_torch, mutates_args=["output"], ) - diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index b193f703e2e1..fc690384a410 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -1010,9 +1010,8 @@ def conv_ssm_forward( n_groups = self.n_groups // self.tp_size # Keep A in native dtype to avoid expensive conversions # The SSU kernel handles mixed types efficiently - A_d = ( - self.A[:, None, ...][:, :, None] - .expand(-1, self.head_dim, self.ssm_state_size) + A_d = self.A[:, None, ...][:, :, None].expand( + -1, self.head_dim, self.ssm_state_size ) dt_d = dt_d[:, :, None].expand(-1, -1, self.head_dim) dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index f0e4361b1a0a..acbdb83b0c44 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -15,7 +15,6 @@ from .cpu_fallbacks import _causal_conv1d_fn_cpu, _causal_conv1d_update_cpu - @triton.jit() def _causal_conv1d_fwd_kernel( # continuous batching # Pointers to matrices diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index bc1e0c120a40..2e79e42e26f4 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -231,4 +231,3 @@ def register_cpu_gdn_attention_ops() -> None: fake_impl=cpu_gdn_attention_core_fake, ) _CPU_GDN_ATTENTION_OPS_REGISTERED = True - diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py index 2db5c7ad7286..be25b0946321 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -429,5 +429,3 @@ def _mamba_chunk_scan_combined_fwd_cpu( all_states = all_states.to(out_dtype) return all_states - - diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index 6607aef3e9d2..c9bd88f5dc34 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -16,6 +16,9 @@ import vllm.envs as envs from vllm import _custom_ops as ops from vllm.logger import init_logger +from vllm.model_executor.layers.mamba.ops.cpu_fallbacks import ( + _selective_state_update_cpu, +) from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton @@ -870,7 +873,6 @@ def selective_scan_fn( return z # output written inplace to z - def selective_state_update( state, x, @@ -919,17 +921,35 @@ def selective_state_update( _sbi = state_batch_indices _dsbi = dst_state_batch_indices ops.selective_state_update_cpu( - _state, _x, _dt, _A, _B, _C, - _D, _z, _dt_bias, dt_softplus, - _sbi, _dsbi, - null_block_id, _out, - num_accepted_tokens, cu_seqlens, + _state, + _x, + _dt, + _A, + _B, + _C, + _D, + _z, + _dt_bias, + dt_softplus, + _sbi, + _dsbi, + null_block_id, + _out, + num_accepted_tokens, + cu_seqlens, ) return _out.squeeze(1) if out.dim() == 2 else _out return _selective_state_update_cuda( - state, x, dt, A, B, C, - D=D, z=z, dt_bias=dt_bias, + state, + x, + dt, + A, + B, + C, + D=D, + z=z, + dt_bias=dt_bias, dt_softplus=dt_softplus, state_batch_indices=state_batch_indices, dst_state_batch_indices=dst_state_batch_indices, diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 8a135e1f549f..caa57211e3d9 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -264,4 +264,3 @@ def mamba_chunk_scan_combined_varlen( ) return varlen_states - diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 9fcfbf16efd5..eb46054d77f4 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -205,11 +205,14 @@ def __init__(self, mamba_config: MambaConfig): _ = ops.selective_state_update_cpu self._use_cpp = True self._cpp_kernel = ops.selective_state_update_cpu - logger.info("CPUSSUBackend: using compiled C++ selective_state_update kernel.") + logger.info( + "CPUSSUBackend: using compiled C++ selective_state_update kernel." + ) except (ImportError, AttributeError): from vllm.model_executor.layers.mamba.ops.cpu_fallbacks import ( _selective_state_update_cpu, ) + self._use_cpp = False self._py_kernel = _selective_state_update_cpu logger.warning( @@ -392,4 +395,3 @@ def selective_state_update( cu_seqlens=cu_seqlens, is_blackwell=is_blackwell, ) - diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index eb45627ba312..38ef3fa769af 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility methods for model layers.""" +import contextlib from collections.abc import Callable import torch @@ -239,8 +240,10 @@ def dispatch_cpu_unquantized_gemm( if layer.weight.ndim != 2: # this is not a linear layer # For now it should be a causal_conv1d op or MoE 3D expert weights - if torch.cpu._is_amx_tile_supported() and hasattr(ops, "causal_conv1d_weight_pack"): - try: + if torch.cpu._is_amx_tile_supported() and hasattr( + ops, "causal_conv1d_weight_pack" + ): + with contextlib.suppress(Exception): # prepack conv weight layer.weight.data = ops.causal_conv1d_weight_pack( layer.weight.view( @@ -248,8 +251,6 @@ def dispatch_cpu_unquantized_gemm( layer.weight.size(2), ) ) - except Exception: - pass layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear( x, weight, bias ) From f3529d794862edfabaeddd9efce268325953cace Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 8 Jul 2026 04:09:31 -0500 Subject: [PATCH 28/43] review fixes Signed-off-by: Akash kaothalkar --- csrc/cpu/torch_bindings.cpp | 5 +- scratch/verify_cpu_ssu.py | 114 ----------------- vllm/_custom_ops.py | 27 ++++ .../layers/mamba/mamba_mixer.py | 3 +- .../layers/mamba/mamba_mixer2.py | 8 +- .../layers/mamba/ops/cpu/gdn_attention.py | 9 +- .../layers/mamba/ops/cpu_fallbacks.py | 121 +----------------- .../layers/mamba/ops/mamba_ssm.py | 8 +- .../layers/mamba/ops/ssu_dispatch.py | 83 ++++-------- 9 files changed, 64 insertions(+), 314 deletions(-) delete mode 100644 scratch/verify_cpu_ssu.py diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 1a0c8411cce1..feaa9ae6c0bf 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -491,11 +491,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor"); ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu); ops.def( - "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor " + "causal_conv1d_update_cpu_amx(Tensor x, Tensor(a!) conv_states, Tensor " "weight, Tensor? bias, bool silu_activation," "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, " "bool is_vnni) -> Tensor"); - ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); + ops.impl("causal_conv1d_update_cpu_amx", torch::kCPU, + &causal_conv1d_update_cpu); #endif #if (defined(__AVX512BF16__) && defined(__AVX512F__) && \ diff --git a/scratch/verify_cpu_ssu.py b/scratch/verify_cpu_ssu.py deleted file mode 100644 index 4b151ba83f56..000000000000 --- a/scratch/verify_cpu_ssu.py +++ /dev/null @@ -1,114 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import time - -import torch -import torch.nn.functional as F - -from vllm.model_executor.layers.mamba.ops.mamba_ssm import selective_state_update - - -def selective_state_update_ref( - state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False -): - has_heads = state.dim() > 3 - if state.dim() == 3: - state = state.unsqueeze(1) - if x.dim() == 2: - x = x.unsqueeze(1) - if dt.dim() == 2: - dt = dt.unsqueeze(1) - if A.dim() == 2: - A = A.unsqueeze(0) - if B.dim() == 2: - B = B.unsqueeze(1) - if C.dim() == 2: - C = C.unsqueeze(1) - if D is not None and D.dim() == 1: - D = D.unsqueeze(0) - if z is not None and z.dim() == 2: - z = z.unsqueeze(1) - if dt_bias is not None and dt_bias.dim() == 1: - dt_bias = dt_bias.unsqueeze(0) - batch, nheads, dim, dstate = state.shape - - dt = dt.float() - if dt_bias is not None: - dt = dt + dt_bias.float() - dt = F.softplus(dt) if dt_softplus else dt - - dA = torch.exp(dt.unsqueeze(-1) * A.float()) # (batch, nheads, dim, dstate) - ngroups = B.shape[1] - B = B.repeat_interleave(nheads // ngroups, dim=1) - C = C.repeat_interleave(nheads // ngroups, dim=1) - - dBx = (dt.unsqueeze(-1) * B.float().unsqueeze(2)) * x.float().unsqueeze(-1) - - state.copy_(state.float() * dA + dBx) - - out = torch.einsum("bhdn,bhn->bhd", state.float(), C.float()) - if D is not None: - out += x.float() * D.float() - - if z is not None: - out = out * F.silu(z.float()) - - if not has_heads: - out = out.squeeze(1) - return out.to(x.dtype) - - -def test_cpu_ssu(): - device = "cpu" - torch.manual_seed(42) - - # Mamba-2 style dimensions - batch_size = 2 - nheads = 128 - dim = 64 - dstate = 128 - itype = torch.bfloat16 - - print(f"Testing CPU selective_state_update with {itype}...") - - state = torch.randn(batch_size, nheads, dim, dstate, dtype=itype, device=device) - x = torch.randn(batch_size, nheads, dim, device=device, dtype=itype) - dt = torch.randn(batch_size, nheads, dim, device=device, dtype=itype) - dt_bias = torch.rand(nheads, dim, device=device) - 4.0 - A = -torch.rand(nheads, dim, dstate, device=device) - 1.0 - B = torch.randn(batch_size, 1, dstate, device=device, dtype=itype) - C = torch.randn(batch_size, 1, dstate, device=device, dtype=itype) - D = torch.randn(nheads, dim, device=device) - z = torch.randn_like(x) - - out = torch.empty_like(x) - state_ref = state.detach().clone() - - start = time.time() - selective_state_update( - state, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, dt_softplus=True, out=out - ) - end = time.time() - print(f"Kernel time: {end - start:.4f}s") - - start = time.time() - out_ref = selective_state_update_ref( - state_ref, x, dt, A, B, C, D=D, z=z, dt_bias=dt_bias, dt_softplus=True - ) - end = time.time() - print(f"Ref time: {end - start:.4f}s") - - state_diff = (state.float() - state_ref.float()).abs().max().item() - out_diff = (out.float() - out_ref.float()).abs().max().item() - - print(f"State max diff: {state_diff}") - print(f"Out max diff: {out_diff}") - - if state_diff < 1e-2 and out_diff < 1e-2: - print("SUCCESS") - else: - print("FAILURE") - - -if __name__ == "__main__": - test_cpu_ssu() diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 54f977605e73..3f9bbdb0c6df 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3422,6 +3422,33 @@ def causal_conv1d_fwd_cpu( ) +def causal_conv1d_update_cpu_amx( + x: torch.Tensor, + conv_states: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + silu_activation: bool, + conv_state_indices: torch.Tensor | None, + is_vnni: bool, +) -> torch.Tensor: + from vllm.platforms import CpuArchEnum + + if current_platform.get_cpu_architecture() == CpuArchEnum.X86: + return torch.ops._C.causal_conv1d_update_cpu_amx( + x, + conv_states, + weight, + bias, + silu_activation, + None, + conv_state_indices, + -1, + is_vnni, + ) + else: + raise NotImplementedError("AMX conv1d is only for x86") + + class CPUDNNLGEMMHandler: def __init__(self) -> None: self.handler_tensor: torch.Tensor | None = None diff --git a/vllm/model_executor/layers/mamba/mamba_mixer.py b/vllm/model_executor/layers/mamba/mamba_mixer.py index 55437c0d0c64..1d3159d1e7b3 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer.py @@ -371,7 +371,6 @@ def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): time_proj_bias = self._time_proj_bias() # 4. Perform the recurrence y ← SSM(A, B, C, Δ)(x) - # Keep D in native dtype to avoid expensive conversions scan_out_p = selective_scan_fn( conv_out_p, ssm_state, @@ -379,7 +378,7 @@ def forward_impl(self, hidden_states: torch.Tensor, output: torch.Tensor): self.A, B_p.transpose(-2, -1), C_p.transpose(-2, -1), - self.D, + self.D.float(), gate_p, time_proj_bias, delta_softplus=True, diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index fc690384a410..a6524961ea92 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -1008,10 +1008,10 @@ def conv_ssm_forward( # 3. State Space Model sequence transformation n_groups = self.n_groups // self.tp_size - # Keep A in native dtype to avoid expensive conversions - # The SSU kernel handles mixed types efficiently - A_d = self.A[:, None, ...][:, :, None].expand( - -1, self.head_dim, self.ssm_state_size + A_d = ( + self.A[:, None, ...][:, :, None] + .expand(-1, self.head_dim, self.ssm_state_size) + .to(dtype=torch.float32) ) dt_d = dt_d[:, :, None].expand(-1, -1, self.head_dim) dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index 2e79e42e26f4..14f2a92861ff 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -98,15 +98,14 @@ def cpu_gdn_attention_core( decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] if is_amx: - decode_mixed_qkv = ops.causal_conv1d_update_cpu( + decode_mixed_qkv = ops.causal_conv1d_update_cpu_amx( x=decode_mixed_qkv, - conv_state=conv_state, + conv_states=conv_state, weight=layer.conv1d.weight, bias=layer.conv1d.bias, - activation="silu" if layer.activation == "silu" else None, + silu_activation=(layer.activation == "silu"), conv_state_indices=decode_state_indices, - query_start_loc=None, - pad_slot_id=0, + is_vnni=True, ) else: decode_conv_state = conv_state[decode_state_indices].contiguous() diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py index be25b0946321..54da43818c26 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -4,7 +4,7 @@ import torch import vllm._custom_ops as ops -from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID +from vllm.v1.attention.backends.utils import PAD_SLOT_ID def _causal_conv1d_fn_cpu( @@ -183,125 +183,6 @@ def _causal_conv1d_update_cpu( return out.to(original_x_dtype) -def _selective_state_update_cpu( - state, - x, - dt, - A, - B, - C, - D=None, - z=None, - dt_bias=None, - dt_softplus=False, - state_batch_indices=None, - dst_state_batch_indices=None, - null_block_id=NULL_BLOCK_ID, - out=None, - num_accepted_tokens=None, - cu_seqlens=None, - is_blackwell=False, - enable_stochastic_rounding=False, - cache_philox_rounds=0, - **kwargs, -): - if state.dim() == 3: - state = state.unsqueeze(1) - if x.dim() == 2: - x = x.unsqueeze(1) - if dt.dim() == 2: - dt = dt.unsqueeze(1) - if A.dim() == 2: - A = A.unsqueeze(0) - if B.dim() == 2: - B = B.unsqueeze(1) - if C.dim() == 2: - C = C.unsqueeze(1) - if D is not None and D.dim() == 1: - D = D.unsqueeze(0) - if z is not None and z.dim() == 2: - z = z.unsqueeze(1) - if dt_bias is not None and dt_bias.dim() == 1: - dt_bias = dt_bias.unsqueeze(0) - if out.dim() == 2: - out = out.unsqueeze(1) - if state_batch_indices is not None and state_batch_indices.dim() == 1: - state_batch_indices = state_batch_indices.unsqueeze(1) - if dst_state_batch_indices is not None and dst_state_batch_indices.dim() == 1: - dst_state_batch_indices = dst_state_batch_indices.unsqueeze(1) - - _, nheads, dim, dstate = state.shape - batch = x.shape[0] - N = len(cu_seqlens) - 1 if cu_seqlens is not None else batch - - ngroups = B.shape[1] - nheads_ngroups_ratio = nheads // ngroups - - for seq_idx in range(N): - if cu_seqlens is not None: - bos = cu_seqlens[seq_idx].item() - seq_len = cu_seqlens[seq_idx + 1].item() - bos - else: - bos = seq_idx - seq_len = 1 - - if state_batch_indices is not None: - state_idx = state_batch_indices[seq_idx, 0].item() - if state_idx == null_block_id: - continue - else: - state_idx = seq_idx - - if num_accepted_tokens is None: - if dst_state_batch_indices is not None: - dst_idx = dst_state_batch_indices[seq_idx, 0].item() - else: - dst_idx = state_idx - - s = state[state_idx].float() - - for t in range(seq_len): - token_idx = bos + t - - x_val = x[token_idx].float() - dt_val = dt[token_idx].float() - - if dt_bias is not None: - dt_val = dt_val + dt_bias.float() - if dt_softplus: - dt_val = torch.nn.functional.softplus(dt_val) - - A_val = A.float() - - B_val = B[token_idx].float() - B_expanded = B_val.repeat_interleave(nheads_ngroups_ratio, dim=0) - C_val = C[token_idx].float() - C_expanded = C_val.repeat_interleave(nheads_ngroups_ratio, dim=0) - - dA = torch.exp(A_val * dt_val.unsqueeze(-1)) - dBx = B_expanded.unsqueeze(1) * (x_val * dt_val).unsqueeze(-1) - s = s * dA + dBx - - if num_accepted_tokens is not None: - token_dst_idx = dst_state_batch_indices[seq_idx, t].item() - if token_dst_idx != null_block_id: - state[token_dst_idx] = s.to(state.dtype) - - out_val = (s * C_expanded.unsqueeze(1)).sum(dim=-1) - - if D is not None: - out_val = out_val + x_val * D.float() - - if z is not None: - z_val = z[token_idx].float() - out_val = out_val * z_val * torch.sigmoid(z_val) - - out[token_idx] = out_val.to(out.dtype) - - if num_accepted_tokens is None and dst_idx != null_block_id: - state[dst_idx] = s.to(state.dtype) - - def _mamba_chunk_scan_combined_fwd_cpu( x, dt, diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index c9bd88f5dc34..906575b00350 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -16,9 +16,6 @@ import vllm.envs as envs from vllm import _custom_ops as ops from vllm.logger import init_logger -from vllm.model_executor.layers.mamba.ops.cpu_fallbacks import ( - _selective_state_update_cpu, -) from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp from vllm.platforms import current_platform from vllm.triton_utils import HAS_TRITON, tl, triton @@ -609,7 +606,7 @@ def _selective_state_update_cuda( assert num_accepted_tokens.shape == (N,) if not HAS_TRITON: - return _selective_state_update_cpu( + return ops.selective_state_update_cpu( state, x, dt, @@ -626,9 +623,6 @@ def _selective_state_update_cuda( out, num_accepted_tokens, cu_seqlens, - is_blackwell, - enable_stochastic_rounding, - cache_philox_rounds, ) grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), N, nheads) diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index eb46054d77f4..1ceeac665f98 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -198,27 +198,10 @@ class CPUSSUBackend(MambaSSUBackend): def __init__(self, mamba_config: MambaConfig): super().__init__(mamba_config) - try: - from vllm import _custom_ops as ops - - # Verify the op is actually registered (CPU build required) - _ = ops.selective_state_update_cpu - self._use_cpp = True - self._cpp_kernel = ops.selective_state_update_cpu - logger.info( - "CPUSSUBackend: using compiled C++ selective_state_update kernel." - ) - except (ImportError, AttributeError): - from vllm.model_executor.layers.mamba.ops.cpu_fallbacks import ( - _selective_state_update_cpu, - ) + from vllm import _custom_ops as ops - self._use_cpp = False - self._py_kernel = _selective_state_update_cpu - logger.warning( - "CPUSSUBackend: C++ selective_state_update op not available, " - "falling back to pure-PyTorch (slow). Rebuild with CPU extensions." - ) + self._cpp_kernel = ops.selective_state_update_cpu + logger.info("CPUSSUBackend: using compiled C++ selective_state_update kernel.") @property def name(self) -> str: @@ -244,46 +227,26 @@ def __call__( cu_seqlens: torch.Tensor | None = None, is_blackwell: bool = False, ) -> None: - if self._use_cpp: - # C++ kernel: state shape expected as (nstates, nheads, dim, dstate) - # The kernel writes in-place into `out` and updates `state`. - self._cpp_kernel( - state, - x, - dt, - A, - B, - C, - D, - z, - dt_bias, - dt_softplus, - state_batch_indices, - dst_state_batch_indices, - null_block_id, - out, - num_accepted_tokens, - cu_seqlens, - ) - else: - self._py_kernel( - state, - x, - dt, - A, - B, - C, - D=D, - z=z, - dt_bias=dt_bias, - dt_softplus=dt_softplus, - state_batch_indices=state_batch_indices, - dst_state_batch_indices=dst_state_batch_indices, - null_block_id=null_block_id, - out=out, - num_accepted_tokens=num_accepted_tokens, - cu_seqlens=cu_seqlens, - ) + # C++ kernel: state shape expected as (nstates, nheads, dim, dstate) + # The kernel writes in-place into `out` and updates `state`. + self._cpp_kernel( + state, + x, + dt, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + dst_state_batch_indices, + null_block_id, + out, + num_accepted_tokens, + cu_seqlens, + ) _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = { From e00254ec13300c0ecd2422ee2896bf11d3745c0d Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Thu, 9 Jul 2026 05:59:23 -0500 Subject: [PATCH 29/43] review fixes and add tests Signed-off-by: Akash kaothalkar --- csrc/cpu/mamba_cpu.cpp | 4 +++- csrc/cpu/torch_bindings.cpp | 6 +++--- tests/kernels/mamba/test_causal_conv1d.py | 8 +++++--- tests/kernels/mamba/test_mamba_ssm.py | 11 +++++++++-- vllm/_custom_ops.py | 8 ++++---- vllm/model_executor/layers/mamba/ops/causal_conv1d.py | 2 +- .../layers/mamba/ops/cpu/gdn_attention.py | 2 +- vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py | 8 +------- 8 files changed, 27 insertions(+), 22 deletions(-) diff --git a/csrc/cpu/mamba_cpu.cpp b/csrc/cpu/mamba_cpu.cpp index 8c95947a9424..ebc3c3516286 100644 --- a/csrc/cpu/mamba_cpu.cpp +++ b/csrc/cpu/mamba_cpu.cpp @@ -159,7 +159,9 @@ void selective_state_update_cpu_impl( int64_t nheads = state.size(1); int64_t dim = state.size(2); int64_t dstate = state.size(3); - int64_t N = x_in.size(0); + int64_t N = (cu_seqlens.has_value() && cu_seqlens.value().defined()) + ? cu_seqlens.value().size(0) - 1 + : x_in.size(0); int64_t ngroups = B_in.size(1); // Strides diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index feaa9ae6c0bf..4b14f9d9cea8 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -491,11 +491,11 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "Tensor"); ops.impl("causal_conv1d_fwd_cpu", torch::kCPU, &causal_conv1d_fwd_cpu); ops.def( - "causal_conv1d_update_cpu_amx(Tensor x, Tensor(a!) conv_states, Tensor " + "causal_conv1d_update_cpu(Tensor x, Tensor(a!) conv_states, Tensor " "weight, Tensor? bias, bool silu_activation," "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, " "bool is_vnni) -> Tensor"); - ops.impl("causal_conv1d_update_cpu_amx", torch::kCPU, + ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); #endif @@ -605,7 +605,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { // Mamba CPU kernels ops.def( - "causal_conv1d_update_cpu(" + "causal_conv1d_update_cpu_vec(" "Tensor(a0!) x, Tensor(a1!) conv_state, Tensor weight, " "Tensor? bias, str? activation, Tensor? conv_state_indices, " "Tensor? query_start_loc, SymInt pad_slot_id) -> Tensor", diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py index c6554f131fed..9c3825487bcf 100644 --- a/tests/kernels/mamba/test_causal_conv1d.py +++ b/tests/kernels/mamba/test_causal_conv1d.py @@ -18,8 +18,9 @@ DEVICE = current_platform.device_type pytestmark = pytest.mark.skipif( - not (current_platform.is_cuda_alike() or current_platform.is_xpu()), - reason="causal_conv1d Triton kernels require CUDA-alike or XPU", + not (current_platform.is_cuda_alike() or current_platform.is_xpu() + or current_platform.is_cpu()), + reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU", ) @@ -284,7 +285,8 @@ def test_causal_conv1d_varlen( batch, with_padding, dim, seqlen, width, has_bias, silu_activation, itype ): device = DEVICE - torch.accelerator.empty_cache() + if not current_platform.is_cpu(): + torch.accelerator.empty_cache() rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3) if itype == torch.bfloat16: rtol, atol = 1e-2, 5e-2 diff --git a/tests/kernels/mamba/test_mamba_ssm.py b/tests/kernels/mamba/test_mamba_ssm.py index 7350b6465231..ddcde8bb8d3b 100644 --- a/tests/kernels/mamba/test_mamba_ssm.py +++ b/tests/kernels/mamba/test_mamba_ssm.py @@ -20,8 +20,9 @@ DEVICE = current_platform.device_type pytestmark = pytest.mark.skipif( - not (current_platform.is_cuda_alike() or current_platform.is_xpu()), - reason="mamba_ssm kernels require CUDA-alike or XPU", + not (current_platform.is_cuda_alike() or current_platform.is_xpu() + or current_platform.is_cpu()), + reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU", ) # selective_scan_fn is backed by the CUDA-only `ops.selective_scan_fwd` C++ op, @@ -342,6 +343,7 @@ def test_selective_scan( @pytest.mark.parametrize("has_z", [False, True]) @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) +@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") def test_selective_state_update(dim, dstate, has_z, itype): device = DEVICE rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2) @@ -432,6 +434,7 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) @pytest.mark.parametrize("max_seq_len", [1, 2, 4]) +@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len): device = DEVICE rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2) @@ -693,6 +696,7 @@ def test_selective_scan_varlen( @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) # tests correctness in case subset of the sequences are padded @pytest.mark.parametrize("with_padding", [True, False]) +@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") def test_selective_state_update_with_batch_indices( with_padding, dim, dstate, has_z, itype ): @@ -785,6 +789,7 @@ def test_selective_state_update_with_batch_indices( @pytest.mark.parametrize("ngroups", [1, 4]) @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) +@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") def test_selective_state_update_with_heads_with_batch_indices( dim, dstate, ngroups, has_z, tie_hdim, itype ): @@ -858,6 +863,7 @@ def test_selective_state_update_with_heads_with_batch_indices( @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) @pytest.mark.parametrize("max_seq_len", [2, 4]) +@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") def test_selective_state_update_with_num_accepted_tokens( dim, dstate, has_z, itype, max_seq_len ): @@ -984,6 +990,7 @@ def test_selective_state_update_with_num_accepted_tokens( @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) @pytest.mark.parametrize("max_seq_len", [2, 4]) +@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") def test_selective_state_update_varlen_with_num_accepted( dim, dstate, has_z, itype, max_seq_len ): diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 3f9bbdb0c6df..8f3c780b9212 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2129,7 +2129,7 @@ def selective_scan_fwd( ) -def causal_conv1d_update_cpu( +def causal_conv1d_update_cpu_vec( x: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor, @@ -2139,7 +2139,7 @@ def causal_conv1d_update_cpu( query_start_loc: torch.Tensor | None = None, pad_slot_id: int = 0, ) -> torch.Tensor: - return torch.ops._C.causal_conv1d_update_cpu( + return torch.ops._C.causal_conv1d_update_cpu_vec( x, conv_state, weight, @@ -3422,7 +3422,7 @@ def causal_conv1d_fwd_cpu( ) -def causal_conv1d_update_cpu_amx( +def causal_conv1d_update_cpu( x: torch.Tensor, conv_states: torch.Tensor, weight: torch.Tensor, @@ -3434,7 +3434,7 @@ def causal_conv1d_update_cpu_amx( from vllm.platforms import CpuArchEnum if current_platform.get_cpu_architecture() == CpuArchEnum.X86: - return torch.ops._C.causal_conv1d_update_cpu_amx( + return torch.ops._C.causal_conv1d_update_cpu( x, conv_states, weight, diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index acbdb83b0c44..022cdfc012c8 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -1297,7 +1297,7 @@ def causal_conv1d_update( _conv_state_indices = _conv_state_indices[:, 0] pad_slot_id = int(NULL_BLOCK_ID) - return ops.causal_conv1d_update_cpu( + return ops.causal_conv1d_update_cpu_vec( x, conv_state, weight, diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index 14f2a92861ff..538d2f5741c6 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -98,7 +98,7 @@ def cpu_gdn_attention_core( decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] if is_amx: - decode_mixed_qkv = ops.causal_conv1d_update_cpu_amx( + decode_mixed_qkv = ops.causal_conv1d_update_cpu( x=decode_mixed_qkv, conv_states=conv_state, weight=layer.conv1d.weight, diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py index 54da43818c26..50a3bb392f54 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py @@ -71,13 +71,7 @@ def _causal_conv1d_fn_cpu( state[:, :-1] = state[:, 1:].clone() state[:, -1] = x_t - if seq_len >= state_len: - conv_states[cache_idx, :, :state_len] = x_seq[:, -state_len:] - else: - conv_states[cache_idx, :, : state_len - seq_len] = conv_states[ - cache_idx, :, seq_len:state_len - ].clone() - conv_states[cache_idx, :, state_len - seq_len :] = x_seq + conv_states[cache_idx].copy_(state) return out.to(original_x_dtype) From 5ddbef6466385b4e37edf89b788d4bfb5a9ceba5 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Fri, 10 Jul 2026 00:48:27 -0500 Subject: [PATCH 30/43] pre-commit fix Signed-off-by: Akash kaothalkar --- csrc/cpu/torch_bindings.cpp | 3 +- tests/kernels/mamba/test_causal_conv1d.py | 7 +++-- tests/kernels/mamba/test_mamba_ssm.py | 37 ++++++++++++++++++----- 3 files changed, 35 insertions(+), 12 deletions(-) diff --git a/csrc/cpu/torch_bindings.cpp b/csrc/cpu/torch_bindings.cpp index 4b14f9d9cea8..8eedc2eaf234 100644 --- a/csrc/cpu/torch_bindings.cpp +++ b/csrc/cpu/torch_bindings.cpp @@ -495,8 +495,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { "weight, Tensor? bias, bool silu_activation," "Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, " "bool is_vnni) -> Tensor"); - ops.impl("causal_conv1d_update_cpu", torch::kCPU, - &causal_conv1d_update_cpu); + ops.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu); #endif #if (defined(__AVX512BF16__) && defined(__AVX512F__) && \ diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py index 9c3825487bcf..d2a2981edf2c 100644 --- a/tests/kernels/mamba/test_causal_conv1d.py +++ b/tests/kernels/mamba/test_causal_conv1d.py @@ -18,8 +18,11 @@ DEVICE = current_platform.device_type pytestmark = pytest.mark.skipif( - not (current_platform.is_cuda_alike() or current_platform.is_xpu() - or current_platform.is_cpu()), + not ( + current_platform.is_cuda_alike() + or current_platform.is_xpu() + or current_platform.is_cpu() + ), reason="causal_conv1d Triton kernels require CUDA-alike, XPU, or CPU", ) diff --git a/tests/kernels/mamba/test_mamba_ssm.py b/tests/kernels/mamba/test_mamba_ssm.py index ddcde8bb8d3b..fd9d412b0a04 100644 --- a/tests/kernels/mamba/test_mamba_ssm.py +++ b/tests/kernels/mamba/test_mamba_ssm.py @@ -20,8 +20,11 @@ DEVICE = current_platform.device_type pytestmark = pytest.mark.skipif( - not (current_platform.is_cuda_alike() or current_platform.is_xpu() - or current_platform.is_cpu()), + not ( + current_platform.is_cuda_alike() + or current_platform.is_xpu() + or current_platform.is_cpu() + ), reason="mamba_ssm kernels require CUDA-alike, XPU, or CPU", ) @@ -343,7 +346,10 @@ def test_selective_scan( @pytest.mark.parametrize("has_z", [False, True]) @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) -@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") +@pytest.mark.skipif( + current_platform.is_cpu(), + reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", +) def test_selective_state_update(dim, dstate, has_z, itype): device = DEVICE rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2) @@ -434,7 +440,10 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) @pytest.mark.parametrize("max_seq_len", [1, 2, 4]) -@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") +@pytest.mark.skipif( + current_platform.is_cpu(), + reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", +) def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len): device = DEVICE rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (5e-3, 1e-2) @@ -696,7 +705,10 @@ def test_selective_scan_varlen( @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) # tests correctness in case subset of the sequences are padded @pytest.mark.parametrize("with_padding", [True, False]) -@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") +@pytest.mark.skipif( + current_platform.is_cpu(), + reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", +) def test_selective_state_update_with_batch_indices( with_padding, dim, dstate, has_z, itype ): @@ -789,7 +801,10 @@ def test_selective_state_update_with_batch_indices( @pytest.mark.parametrize("ngroups", [1, 4]) @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) -@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") +@pytest.mark.skipif( + current_platform.is_cpu(), + reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", +) def test_selective_state_update_with_heads_with_batch_indices( dim, dstate, ngroups, has_z, tie_hdim, itype ): @@ -863,7 +878,10 @@ def test_selective_state_update_with_heads_with_batch_indices( @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) @pytest.mark.parametrize("max_seq_len", [2, 4]) -@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") +@pytest.mark.skipif( + current_platform.is_cpu(), + reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", +) def test_selective_state_update_with_num_accepted_tokens( dim, dstate, has_z, itype, max_seq_len ): @@ -990,7 +1008,10 @@ def test_selective_state_update_with_num_accepted_tokens( @pytest.mark.parametrize("dstate", [16, 64]) @pytest.mark.parametrize("dim", [2048, 4096]) @pytest.mark.parametrize("max_seq_len", [2, 4]) -@pytest.mark.skipif(current_platform.is_cpu(), reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.") +@pytest.mark.skipif( + current_platform.is_cpu(), + reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", +) def test_selective_state_update_varlen_with_num_accepted( dim, dstate, has_z, itype, max_seq_len ): From 25bdaea8cb95313ab38620ba8a4a7d5a80dca3b0 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Fri, 10 Jul 2026 01:22:50 -0500 Subject: [PATCH 31/43] pre-commit fix _ Signed-off-by: Akash kaothalkar --- tests/kernels/mamba/test_mamba_ssm.py | 30 +++++++++++++++++++++------ 1 file changed, 24 insertions(+), 6 deletions(-) diff --git a/tests/kernels/mamba/test_mamba_ssm.py b/tests/kernels/mamba/test_mamba_ssm.py index fd9d412b0a04..a30295cc3845 100644 --- a/tests/kernels/mamba/test_mamba_ssm.py +++ b/tests/kernels/mamba/test_mamba_ssm.py @@ -348,7 +348,10 @@ def test_selective_scan( @pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096]) @pytest.mark.skipif( current_platform.is_cpu(), - reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), ) def test_selective_state_update(dim, dstate, has_z, itype): device = DEVICE @@ -442,7 +445,10 @@ def test_selective_state_update_stochastic_rounding(dim, dstate, has_z, philox_r @pytest.mark.parametrize("max_seq_len", [1, 2, 4]) @pytest.mark.skipif( current_platform.is_cpu(), - reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), ) def test_selective_state_update_varlen(dim, dstate, has_z, itype, max_seq_len): device = DEVICE @@ -707,7 +713,10 @@ def test_selective_scan_varlen( @pytest.mark.parametrize("with_padding", [True, False]) @pytest.mark.skipif( current_platform.is_cpu(), - reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), ) def test_selective_state_update_with_batch_indices( with_padding, dim, dstate, has_z, itype @@ -803,7 +812,10 @@ def test_selective_state_update_with_batch_indices( @pytest.mark.parametrize("dim", [2048, 4096]) @pytest.mark.skipif( current_platform.is_cpu(), - reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), ) def test_selective_state_update_with_heads_with_batch_indices( dim, dstate, ngroups, has_z, tie_hdim, itype @@ -880,7 +892,10 @@ def test_selective_state_update_with_heads_with_batch_indices( @pytest.mark.parametrize("max_seq_len", [2, 4]) @pytest.mark.skipif( current_platform.is_cpu(), - reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), ) def test_selective_state_update_with_num_accepted_tokens( dim, dstate, has_z, itype, max_seq_len @@ -1010,7 +1025,10 @@ def test_selective_state_update_with_num_accepted_tokens( @pytest.mark.parametrize("max_seq_len", [2, 4]) @pytest.mark.skipif( current_platform.is_cpu(), - reason="CPU kernel for selective_state_update only supports Mamba 2 (scalar A/dt), not Mamba 1.", + reason=( + "CPU kernel for selective_state_update only supports " + "Mamba 2 (scalar A/dt), not Mamba 1." + ), ) def test_selective_state_update_varlen_with_num_accepted( dim, dstate, has_z, itype, max_seq_len From 5c33a7a4ade4a8c482e1a7e6a4f99f39842ed3e9 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 14 Jul 2026 13:56:23 +0530 Subject: [PATCH 32/43] mamba cpu updates and review comments Signed-off-by: Akash kaothalkar --- .buildkite/hardware_tests/cpu.yaml | 6 +- mamba_cpu_updates.patch | 968 ++++++++++++++++++ vllm/_custom_ops.py | 27 +- vllm/config/mamba.py | 9 +- .../layers/fused_moe/cpu_fused_moe.py | 90 -- .../layers/mamba/ops/causal_conv1d.py | 55 +- .../layers/mamba/ops/cpu/causal_conv1d.py | 177 +++- .../layers/mamba/ops/cpu/mamba_ssm.py | 149 +++ .../layers/mamba/ops/cpu_fallbacks.py | 306 ------ .../layers/mamba/ops/mamba_ssm.py | 47 +- .../layers/mamba/ops/ssd_combined.py | 22 +- .../layers/mamba/ops/ssu_dispatch.py | 7 +- vllm/model_executor/layers/utils.py | 28 +- 13 files changed, 1349 insertions(+), 542 deletions(-) create mode 100644 mamba_cpu_updates.patch create mode 100644 vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py delete mode 100644 vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml index ebfd1c7524ad..fd2d45be5a06 100644 --- a/.buildkite/hardware_tests/cpu.yaml +++ b/.buildkite/hardware_tests/cpu.yaml @@ -18,6 +18,8 @@ steps: - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py - tests/kernels/mamba/test_cpu_short_conv.py + - tests/kernels/mamba/test_causal_conv1d.py + - tests/kernels/mamba/test_mamba_ssm.py commands: - | bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " @@ -28,7 +30,9 @@ steps: pytest -x -v -s tests/kernels/test_onednn.py pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py - pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py" + pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py + pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py + pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py" # Note: SDE can't be downloaded from CI host because of AWS WAF # - label: CPU-Compatibility Tests diff --git a/mamba_cpu_updates.patch b/mamba_cpu_updates.patch new file mode 100644 index 000000000000..5f128bdf7768 --- /dev/null +++ b/mamba_cpu_updates.patch @@ -0,0 +1,968 @@ +diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml +index ebfd1c752..fd2d45be5 100644 +--- a/.buildkite/hardware_tests/cpu.yaml ++++ b/.buildkite/hardware_tests/cpu.yaml +@@ -18,6 +18,8 @@ steps: + - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py + - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py + - tests/kernels/mamba/test_cpu_short_conv.py ++ - tests/kernels/mamba/test_causal_conv1d.py ++ - tests/kernels/mamba/test_mamba_ssm.py + commands: + - | + bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " +@@ -28,7 +30,9 @@ steps: + pytest -x -v -s tests/kernels/test_onednn.py + pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py + pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py +- pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py" ++ pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py ++ pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py ++ pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py" + + # Note: SDE can't be downloaded from CI host because of AWS WAF + # - label: CPU-Compatibility Tests +diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py +index 4fb3a551b..82297fe84 100644 +--- a/vllm/_custom_ops.py ++++ b/vllm/_custom_ops.py +@@ -3429,22 +3429,17 @@ def causal_conv1d_update_cpu( + conv_state_indices: torch.Tensor | None, + is_vnni: bool, + ) -> torch.Tensor: +- from vllm.platforms import CpuArchEnum +- +- if current_platform.get_cpu_architecture() == CpuArchEnum.X86: +- return torch.ops._C.causal_conv1d_update_cpu( +- x, +- conv_states, +- weight, +- bias, +- silu_activation, +- None, +- conv_state_indices, +- -1, +- is_vnni, +- ) +- else: +- raise NotImplementedError("AMX conv1d is only for x86") ++ return torch.ops._C.causal_conv1d_update_cpu( ++ x, ++ conv_states, ++ weight, ++ bias, ++ silu_activation, ++ None, ++ conv_state_indices, ++ -1, ++ is_vnni, ++ ) + + + class CPUDNNLGEMMHandler: +diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py +index fa556723e..a68842f4b 100644 +--- a/vllm/config/mamba.py ++++ b/vllm/config/mamba.py +@@ -27,9 +27,6 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta): + + TRITON = "triton" + FLASHINFER = "flashinfer" +- # Pure-PyTorch fallback for CPU-only platforms (PowerPC, no CUDA, etc.). +- # Avoids Triton JIT compilation which is unstable / unsupported on those +- # architectures. + CPU = "cpu" + + +@@ -38,11 +35,7 @@ class MambaConfig: + """Configuration for Mamba SSM backends.""" + + backend: MambaBackendEnum = MambaBackendEnum.TRITON +- """Mamba SSU backend to use. +- +- On CPU-only platforms (e.g. PowerPC, x86 without CUDA) the default is +- automatically overridden to 'cpu' by ``__post_init__``. +- """ ++ """Mamba SSU backend to use.""" + + enable_stochastic_rounding: bool = False + """Enable stochastic rounding when writing SSM state to fp16 cache. +diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +index 27a9da5fb..352f25334 100644 +--- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py ++++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +@@ -440,96 +440,6 @@ class CPUFusedMOE: + + return output + +- def forward_batched_gemm( +- self, +- layer: torch.nn.Module, +- input: torch.Tensor, +- topk_weights: torch.Tensor, +- topk_ids: torch.Tensor, +- activation: MoEActivation, +- global_num_experts: int = -1, +- skip_weighted: bool = False, +- ) -> torch.Tensor: +- """Universal batched GEMM fallback for CPUs without AVX512 fused MoE. +- +- The per-expert F.linear loop in forward_torch issues N_experts separate +- BLAS calls, each triggering a full weight-matrix packing pass +- (sbgemm_incopy / similar). For decode (small T), packing dominates. +- +- This path instead: +- 1. One torch.mm for gate+up — packs ALL expert weights in one pass. +- 2. One torch.bmm for down — batched over experts. +- → 2 BLAS weight-packing passes regardless of N_experts. +- +- Works on any CPU (x86 non-AVX512, ARM non-AMX, PowerPC, RISC-V, ...). +- The trade-off (computing all experts densely) is beneficial whenever +- decode batch size is small relative to expert weight size. +- """ +- if skip_weighted: +- # Weight is pre-applied to input; must NOT re-apply in accumulation. +- assert topk_ids.size(1) == 1, ( +- "apply_router_weight_on_input is only supported for topk=1" +- ) +- input = input * topk_weights.to(input.dtype) +- +- num_tokens: int = input.shape[0] +- hidden_dim: int = input.shape[1] +- num_experts: int = layer.w13_weight.shape[0] +- gate_up_dim: int = layer.w13_weight.shape[1] # gate + up concatenated +- top_k: int = topk_ids.shape[1] +- act_fn = _CPU_MOE_ACT_FN[activation] +- +- # ------------------------------------------------------------------ +- # 1. Gate+Up — single large GEMM covering ALL experts at once. +- # w13_weight: (E, gate+up, hidden) → (E*gate+up, hidden) +- # mm: (T, hidden) @ (hidden, E*gate+up) → (T, E*gate+up) +- # → view (T, E, gate+up) +- # ------------------------------------------------------------------ +- w13_2d = layer.w13_weight.view(num_experts * gate_up_dim, hidden_dim) +- gate_up_all = torch.mm(input, w13_2d.T) # (T, E*gate+up) +- if hasattr(layer, "w13_bias") and layer.w13_bias is not None: +- gate_up_all = gate_up_all + layer.w13_bias.view(-1) +- gate_up_all = gate_up_all.view(num_tokens, num_experts, gate_up_dim) +- +- # Apply gating activation: (T, E, gate+up) → (T, E, down_dim) +- gate_up_act = act_fn(gate_up_all) +- +- # ------------------------------------------------------------------ +- # 2. Down — batched GEMM over experts. +- # Permute to (E, T, down_dim) for bmm. +- # w2_weight: (E, out, down_dim) → transpose → (E, down_dim, out) +- # bmm: (E, T, down_dim) @ (E, down_dim, out) → (E, T, out) +- # ------------------------------------------------------------------ +- gate_up_t = gate_up_act.permute(1, 0, 2).contiguous() # (E, T, down) +- down_all = torch.bmm( +- gate_up_t, +- layer.w2_weight.transpose(1, 2), +- ) # (E, T, out_dim) +- if hasattr(layer, "w2_bias") and layer.w2_bias is not None: +- down_all = down_all + layer.w2_bias.unsqueeze(1) +- down_all = down_all.permute(1, 0, 2) # (T, E, out_dim) +- +- # ------------------------------------------------------------------ +- # 3. Weighted accumulation over top-k routed experts. +- # output[t] += topk_weights[t,k] * down_all[t, topk_ids[t,k], :] +- # ------------------------------------------------------------------ +- out_dim: int = down_all.shape[-1] +- output = torch.zeros( +- num_tokens, out_dim, dtype=input.dtype, device=input.device +- ) +- t_idx = torch.arange(num_tokens, device=input.device) +- if skip_weighted: +- # Weight already baked into input — gather directly, no re-weighting. +- expert_ids = topk_ids[:, 0].long() +- output.add_(down_all[t_idx, expert_ids, :]) +- else: +- for k in range(top_k): +- expert_ids = topk_ids[:, k].long() # (T,) +- w_k = topk_weights[:, k].to(input.dtype) # (T,) +- expert_out = down_all[t_idx, expert_ids, :] # (T, out) +- output.addcmul_(expert_out, w_k.unsqueeze(-1).expand_as(expert_out)) +- +- return output + + + def cpu_fused_moe_torch( +diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +index 022cdfc01..bb1a3390e 100644 +--- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py ++++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +@@ -12,7 +12,6 @@ from vllm import _custom_ops as ops + from vllm.triton_utils import tl, triton + from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID + +-from .cpu_fallbacks import _causal_conv1d_fn_cpu, _causal_conv1d_update_cpu + + + @triton.jit() +@@ -1243,10 +1242,6 @@ def _causal_conv1d_update_cuda( + + + def causal_conv1d_fn(*args, **kwargs): +- """Dispatch causal_conv1d_fn to CPU PyTorch fallback or CUDA Triton kernel.""" +- x = args[0] if args else kwargs.get("x") +- if x is not None and x.device.type == "cpu": +- return _causal_conv1d_fn_cpu(*args, **kwargs) + return _causal_conv1d_fn_cuda(*args, **kwargs) + + +@@ -1266,47 +1261,6 @@ def causal_conv1d_update( + validate_data=False, + ): + """Dispatch causal_conv1d_update to CPU C++ kernel or CUDA Triton kernel.""" +- if x.device.type == "cpu": +- # The C++ kernel handles the standard (non-varlen, non-spec-decoding) +- # decode path. Fall back to PyTorch for the complex varlen / +- # spec-decoding paths. +- if query_start_loc is not None or num_accepted_tokens is not None: +- return _causal_conv1d_update_cpu( +- x, +- conv_state, +- weight, +- bias=bias, +- activation=activation, +- conv_state_indices=conv_state_indices, +- query_start_loc=query_start_loc, +- ) +- # Determine activation string +- act_str = None +- if isinstance(activation, bool): +- act_str = "silu" if activation else None +- elif activation is not None: +- act_str = activation +- +- _conv_state_indices = conv_state_indices +- if _conv_state_indices is not None and _conv_state_indices.dim() == 2: +- if initial_state_idx is not None: +- _conv_state_indices = _conv_state_indices.gather( +- 1, initial_state_idx.unsqueeze(1) +- ).squeeze(1) +- else: +- _conv_state_indices = _conv_state_indices[:, 0] +- +- pad_slot_id = int(NULL_BLOCK_ID) +- return ops.causal_conv1d_update_cpu_vec( +- x, +- conv_state, +- weight, +- bias, +- act_str, +- _conv_state_indices, +- None, # query_start_loc +- pad_slot_id, +- ) + + return _causal_conv1d_update_cuda( + x, +@@ -1323,3 +1277,12 @@ def causal_conv1d_update( + initial_state_idx=initial_state_idx, + validate_data=validate_data, + ) ++ ++ ++from vllm.platforms import current_platform ++ ++if current_platform.is_cpu(): ++ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( ++ causal_conv1d_fn_cpu as causal_conv1d_fn, ++ causal_conv1d_update_cpu as causal_conv1d_update, ++ ) +diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +index b047ca6d6..dba752624 100644 +--- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py ++++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +@@ -4,10 +4,182 @@ + from __future__ import annotations + + import torch +-import torch.nn.functional as F + ++from vllm.v1.attention.backends.utils import PAD_SLOT_ID ++ ++ ++def causal_conv1d_fn_cpu( ++ x: torch.Tensor, ++ weight: torch.Tensor, ++ bias: torch.Tensor | None, ++ conv_states: torch.Tensor, ++ query_start_loc: torch.Tensor, ++ cache_indices: torch.Tensor | None = None, ++ has_initial_state: torch.Tensor | None = None, ++ activation: str | None = "silu", ++ pad_slot_id: int = PAD_SLOT_ID, ++ **kwargs, ++) -> torch.Tensor: ++ """CPU implementation for causal_conv1d_fwd.""" ++ if isinstance(activation, bool) and activation: ++ activation = "silu" ++ ++ 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 ++ ++ if seq_len == 0: ++ continue ++ ++ cache_idx = cache_indices[b].item() if cache_indices is not None else b ++ ++ if cache_idx == 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) ++ 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 ++ ++ if state_len > 1: ++ state[:, :-1] = state[:, 1:].clone() ++ state[:, -1] = x_t ++ ++ conv_states[cache_idx].copy_(state) ++ ++ return out.to(original_x_dtype) ++ ++ ++def causal_conv1d_update_cpu( ++ x: torch.Tensor, ++ conv_state: torch.Tensor, ++ weight: torch.Tensor, ++ bias: torch.Tensor | None = None, ++ activation: bool | str | None = None, ++ conv_state_indices: torch.Tensor | None = None, ++ query_start_loc: torch.Tensor | None = None, ++ pad_slot_id: int = PAD_SLOT_ID, ++ **kwargs, ++) -> torch.Tensor: ++ """CPU implementation for causal_conv1d_update.""" ++ if isinstance(activation, bool): ++ activation = "silu" if activation else None ++ ++ original_x_dtype = x.dtype ++ x = x.to(conv_state.dtype) ++ _, width = weight.shape ++ state_len = width - 1 ++ ++ if query_start_loc is None and x.dim() == 2: ++ x = x.unsqueeze(-1) ++ unsqueeze = True ++ else: ++ unsqueeze = False ++ ++ if query_start_loc is None: ++ batch, dim, seqlen = x.shape ++ ++ if conv_state_indices is not None: ++ cache_idxs = conv_state_indices.flatten() ++ valid_mask = cache_idxs != pad_slot_id ++ else: ++ cache_idxs = torch.arange(batch, device=x.device) ++ valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) ++ ++ for t in range(seqlen): ++ x_t = x[:, :, t].clone() ++ ++ states = conv_state[cache_idxs] ++ ++ windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) ++ ++ val = (windows * weight.unsqueeze(0)).sum(dim=-1) ++ if bias is not None: ++ val = val + bias.unsqueeze(0) ++ ++ if activation in ["silu", "swish"]: ++ val = val * torch.sigmoid(val) ++ ++ val = val * valid_mask.unsqueeze(-1).to(val.dtype) ++ x[:, :, t] = val ++ ++ new_state = torch.cat([states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1) ++ conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] ++ ++ out = x ++ if unsqueeze: ++ out = out.squeeze(-1) ++ return out.to(original_x_dtype) ++ ++ assert conv_state_indices is not None ++ assert query_start_loc is not None ++ batch = conv_state_indices.size(0) ++ out = x.clone() ++ ++ for b in range(batch): ++ cache_idx = conv_state_indices[b].item() ++ if cache_idx == pad_slot_id: ++ continue ++ ++ seq_start = query_start_loc[b].item() ++ seq_end = query_start_loc[b + 1].item() ++ seqlen_b = seq_end - seq_start ++ ++ if seqlen_b == 0: ++ continue ++ ++ local_state = conv_state[cache_idx].clone() ++ ++ for t in range(seqlen_b): ++ x_t = x[seq_start + t, :] ++ ++ window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) ++ val = (window * weight).sum(dim=-1) ++ if bias is not None: ++ val = val + bias ++ if activation in ["silu", "swish"]: ++ val = val * torch.sigmoid(val) ++ ++ out[seq_start + t, :] = val ++ ++ if state_len > 1: ++ local_state[:, :-1] = local_state[:, 1:].clone() ++ local_state[:, -1] = x_t ++ ++ conv_state[cache_idx] = local_state ++ ++ return out.to(original_x_dtype) ++ ++ ++import torch.nn.functional as F + +-# for prefill + def causal_conv1d_torch( + x: torch.Tensor, + weight: torch.Tensor, +@@ -60,7 +232,6 @@ def causal_conv1d_torch( + return out + + +-# for decode + def causal_conv1d_update_torch( + x: torch.Tensor, + conv_state: torch.Tensor, +diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +deleted file mode 100644 +index 50a3bb392..000000000 +--- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py ++++ /dev/null +@@ -1,306 +0,0 @@ +-# SPDX-License-Identifier: Apache-2.0 +-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +- +-import torch +- +-import vllm._custom_ops as ops +-from vllm.v1.attention.backends.utils import PAD_SLOT_ID +- +- +-def _causal_conv1d_fn_cpu( +- x: torch.Tensor, +- weight: torch.Tensor, +- bias: torch.Tensor | None, +- conv_states: torch.Tensor, +- query_start_loc: torch.Tensor, +- cache_indices: torch.Tensor | None = None, +- has_initial_state: torch.Tensor | None = None, +- activation: str | None = "silu", +- pad_slot_id: int = PAD_SLOT_ID, +- **kwargs, +-) -> torch.Tensor: +- """Pure PyTorch CPU fallback for causal_conv1d_fwd.""" +- if isinstance(activation, bool) and activation: +- activation = "silu" +- +- 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 +- +- if seq_len == 0: +- continue +- +- cache_idx = cache_indices[b].item() if cache_indices is not None else b +- +- if cache_idx == 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) +- 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 +- +- if state_len > 1: +- state[:, :-1] = state[:, 1:].clone() +- state[:, -1] = x_t +- +- conv_states[cache_idx].copy_(state) +- +- return out.to(original_x_dtype) +- +- +-def _causal_conv1d_update_cpu( +- x: torch.Tensor, +- conv_state: torch.Tensor, +- weight: torch.Tensor, +- bias: torch.Tensor | None = None, +- activation: bool | str | None = None, +- conv_state_indices: torch.Tensor | None = None, +- query_start_loc: torch.Tensor | None = None, +- pad_slot_id: int = PAD_SLOT_ID, +- **kwargs, +-) -> torch.Tensor: +- """Pure PyTorch CPU fallback for causal_conv1d_update (decode path).""" +- if isinstance(activation, bool): +- activation = "silu" if activation else None +- +- original_x_dtype = x.dtype +- x = x.to(conv_state.dtype) +- _, width = weight.shape +- state_len = width - 1 +- +- if query_start_loc is None and x.dim() == 2: +- x = x.unsqueeze(-1) +- unsqueeze = True +- else: +- unsqueeze = False +- +- if query_start_loc is None: +- batch, dim, seqlen = x.shape +- +- if conv_state_indices is not None: +- cache_idxs = conv_state_indices.flatten() +- valid_mask = cache_idxs != pad_slot_id +- else: +- cache_idxs = torch.arange(batch, device=x.device) +- valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) +- +- for t in range(seqlen): +- x_t = x[:, :, t].clone() +- +- states = conv_state[cache_idxs] +- +- windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) +- +- val = (windows * weight.unsqueeze(0)).sum(dim=-1) +- if bias is not None: +- val = val + bias.unsqueeze(0) +- +- if activation in ["silu", "swish"]: +- val = val * torch.sigmoid(val) +- +- val = val * valid_mask.unsqueeze(-1).to(val.dtype) +- x[:, :, t] = val +- +- new_state = torch.cat([states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1) +- conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] +- +- out = x +- if unsqueeze: +- out = out.squeeze(-1) +- return out.to(original_x_dtype) +- +- assert conv_state_indices is not None +- assert query_start_loc is not None +- batch = conv_state_indices.size(0) +- out = x.clone() +- +- for b in range(batch): +- cache_idx = conv_state_indices[b].item() +- if cache_idx == pad_slot_id: +- continue +- +- seq_start = query_start_loc[b].item() +- seq_end = query_start_loc[b + 1].item() +- seqlen_b = seq_end - seq_start +- +- if seqlen_b == 0: +- continue +- +- local_state = conv_state[cache_idx].clone() +- +- for t in range(seqlen_b): +- x_t = x[seq_start + t, :] +- +- window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) +- val = (window * weight).sum(dim=-1) +- if bias is not None: +- val = val + bias +- if activation in ["silu", "swish"]: +- val = val * torch.sigmoid(val) +- +- out[seq_start + t, :] = val +- +- if state_len > 1: +- local_state[:, :-1] = local_state[:, 1:].clone() +- local_state[:, -1] = x_t +- +- conv_state[cache_idx] = local_state +- +- return out.to(original_x_dtype) +- +- +-def _mamba_chunk_scan_combined_fwd_cpu( +- x, +- dt, +- A, +- B, +- C, +- chunk_size, +- out, +- D=None, +- z=None, +- dt_bias=None, +- initial_states=None, +- return_intermediate_states=False, +- seq_idx=None, +- cu_seqlens=None, +- cu_chunk_seqlens=None, +- last_chunk_indices=None, +- dt_softplus=False, +- dt_limit=(0.0, float("inf")), +- state_dtype=None, +- **kwargs, +-): +- seqlen, nheads, headdim = x.shape +- _, ngroups, dstate = B.shape +- +- assert cu_seqlens is not None +- batch = cu_seqlens.size(0) - 1 +- +- # Preprocess dt: apply bias, softplus, and clamp. +- # Kept in Python to avoid duplicating this logic in C++. +- dt_f = dt.float() +- if dt_bias is not None: +- dt_f = dt_f + dt_bias.float().unsqueeze(0) +- if dt_softplus: +- dt_f = torch.nn.functional.softplus(dt_f) +- if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): +- dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) +- +- # Allocate state output buffer (float32, contiguous — required by kernel). +- all_states = torch.zeros( +- batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device +- ) +- if initial_states is not None: +- all_states.copy_(initial_states.float()) +- +- # Use the C++ kernel when available (CPU build with mamba kernels compiled). +- # Falls back to the pure-Python loop below if the op is not registered. +- _use_cpp_kernel = hasattr(torch.ops._C, "mamba_chunk_scan_fwd_cpu") +- +- if _use_cpp_kernel: +- # out must be contiguous — the C++ kernel writes via raw data_ptr(). +- # Creating a contiguous copy here would discard the results, so we +- # require the caller to pass a contiguous tensor. +- assert out.is_contiguous(), ( +- "_mamba_chunk_scan_combined_fwd_cpu: `out` must be " +- "pre-allocated as a contiguous tensor" +- ) +- +- # D: strip trailing dims that are broadcast (stride == 0) so the +- # kernel sees a (nheads,) float32 array. Only peel dims whose +- # stride is 0 (broadcast-expanded); genuine multi-dim D is passed +- # as-is and flattened by the C++ wrapper. +- D_1d = None +- if D is not None: +- d = D.float() +- while d.dim() > 1 and d.stride(-1) == 0: +- d = d.squeeze(-1) +- D_1d = d.contiguous() +- +- ops.mamba_chunk_scan_fwd_cpu( +- out, +- all_states, +- x, +- dt_f, +- A, +- B, +- C, +- D_1d, +- z, +- cu_seqlens.to(torch.int32), +- ) +- else: +- # Pure-Python fallback (no compiled extension available). +- for b_idx in range(batch): +- seq_start = cu_seqlens[b_idx].item() +- seq_end = cu_seqlens[b_idx + 1].item() +- +- # Note: basic indexing returns a view, but arithmetic ops +- # below (state * dA + dBx) produce new tensors, so state is +- # reassigned each token. all_states[b_idx] is committed at line 416. +- state = all_states[b_idx].clone() # local working copy +- +- for t in range(seq_start, seq_end): +- x_t = x[t].float() +- dt_t = dt_f[t] +- A_val = A.float() +- +- dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) +- +- B_expanded = B[t].float().repeat_interleave(nheads // ngroups, dim=0) +- C_expanded = C[t].float().repeat_interleave(nheads // ngroups, dim=0) +- +- xdt = x_t * dt_t.unsqueeze(-1) +- dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) +- state = state * dA + dBx +- +- y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) +- +- if D is not None: +- y = ( +- y + x_t * D.float().unsqueeze(-1) +- if D.dim() == 1 +- else y + x_t * D.float() +- ) +- +- if z is not None: +- z_t = z[t].float() +- y = y * z_t * torch.sigmoid(z_t) +- +- out[t] = y.to(out.dtype) +- +- all_states[b_idx] = state.to(all_states.dtype) +- +- out_dtype = state_dtype if state_dtype is not None else x.dtype +- all_states = all_states.to(out_dtype) +- +- return all_states +diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +index d7055fbd7..9d463e570 100644 +--- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py ++++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +@@ -894,47 +894,6 @@ def selective_state_update( + if out is None: + out = torch.empty_like(x if x.dim() == 2 else x) + +- if x.device.type == "cpu": +- # Reshape tensors from (batch, dim) -> (batch, 1, dim) if needed +- # The C++ kernel expects (N, nheads, dim) layout +- _state = state.unsqueeze(1) if state.dim() == 3 else state +- _x = x.unsqueeze(1) if x.dim() == 2 else x +- _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt +- _A = A.unsqueeze(0) if A.dim() == 2 else A +- _B = B.unsqueeze(1) if B.dim() == 2 else B +- _C = C.unsqueeze(1) if C.dim() == 2 else C +- _D = D.unsqueeze(0) if (D is not None and D.dim() == 1) else D +- _z = z.unsqueeze(1) if (z is not None and z.dim() == 2) else z +- _dt_bias = ( +- dt_bias.unsqueeze(0) +- if (dt_bias is not None and dt_bias.dim() == 1) +- else dt_bias +- ) +- _out = out.unsqueeze(1) if out.dim() == 2 else out +- # state_batch_indices and dst_state_batch_indices are 1D index arrays; +- # do NOT reshape them. +- _sbi = state_batch_indices +- _dsbi = dst_state_batch_indices +- ops.selective_state_update_cpu( +- _state, +- _x, +- _dt, +- _A, +- _B, +- _C, +- _D, +- _z, +- _dt_bias, +- dt_softplus, +- _sbi, +- _dsbi, +- null_block_id, +- _out, +- num_accepted_tokens, +- cu_seqlens, +- ) +- return _out.squeeze(1) if out.dim() == 2 else _out +- + return _selective_state_update_cuda( + state, + x, +@@ -956,3 +915,9 @@ def selective_state_update( + enable_stochastic_rounding=enable_stochastic_rounding, + cache_philox_rounds=cache_philox_rounds, + ) ++ ++ ++if current_platform.is_cpu(): ++ from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( ++ selective_state_update, ++ ) +diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py +index caa57211e..89c8bd814 100644 +--- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py ++++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py +@@ -13,7 +13,7 @@ from packaging import version + from vllm.model_executor.custom_op import CustomOp + from vllm.triton_utils import HAS_TRITON, triton + +-from .cpu_fallbacks import _mamba_chunk_scan_combined_fwd_cpu ++ + from .ssd_bmm import _bmm_chunk_fwd + from .ssd_chunk_scan import _chunk_scan_fwd + from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd +@@ -166,19 +166,11 @@ class MambaChunkScanCombinedFwdOp(CustomOp): + super(CustomOp, self).__init__() # nn.Module.__init__ only + from vllm.platforms import current_platform + +- if current_platform.is_cpu(): +- self._forward_method = self.forward_cpu +- elif current_platform.is_rocm(): ++ if current_platform.is_rocm(): + self._forward_method = self.forward_hip + else: + self._forward_method = self.forward_cuda + +- def forward_native(self, *args, **kwargs): +- return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) +- +- def forward_cpu(self, *args, **kwargs): +- return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) +- + def forward_cuda(self, *args, **kwargs): + return _mamba_chunk_scan_combined_fwd_cuda(*args, **kwargs) + +@@ -260,7 +252,15 @@ def mamba_chunk_scan_combined_varlen( + last_chunk_indices=last_chunk_indices, + dt_softplus=dt_softplus, + dt_limit=dt_limit, ++ return_intermediate_states=return_intermediate_states, + state_dtype=state_dtype, + ) +- + return varlen_states ++ ++ ++from vllm.platforms import current_platform ++ ++if current_platform.is_cpu(): ++ from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( ++ _mamba_chunk_scan_combined_fwd_cpu as _mamba_chunk_scan_combined_fwd, ++ ) +diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +index 1ceeac665..1795a036b 100644 +--- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py ++++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +@@ -6,8 +6,7 @@ Dispatch module for Mamba selective state update (SSU) backends. + Provides a unified `selective_state_update` function that dispatches to + the Triton, FlashInfer, or CPU backend based on the configured + `MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA) +-the backend defaults to 'cpu', which uses a pure-PyTorch fallback that +-avoids Triton JIT compilation entirely. ++the backend defaults to 'cpu'. + """ + + from abc import ABC, abstractmethod +@@ -280,7 +279,7 @@ def initialize_mamba_ssu_backend( + backend = mamba_config.backend + + # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is +- # unstable or unavailable. Silently fall back to the pure-PyTorch CPU ++ # unstable or unavailable. Silently fall back to the CPU + # backend unless the user explicitly chose something other than "triton". + if backend == MambaBackendEnum.TRITON: + from vllm.platforms import current_platform +@@ -288,7 +287,7 @@ def initialize_mamba_ssu_backend( + if current_platform.is_cpu(): + logger.info( + "CPU platform detected: overriding Mamba SSU backend " +- "from 'triton' to 'cpu' (pure-PyTorch fallback)." ++ "from 'triton' to 'cpu'." + ) + backend = MambaBackendEnum.CPU + +diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py +index 450c74a60..f0d3e6ead 100644 +--- a/vllm/model_executor/layers/utils.py ++++ b/vllm/model_executor/layers/utils.py +@@ -243,23 +243,19 @@ def dispatch_cpu_unquantized_gemm( + if torch.cpu._is_amx_tile_supported() and hasattr( + ops, "causal_conv1d_weight_pack" + ): +- with contextlib.suppress(Exception): +- # prepack conv weight +- unpacked = ( +- layer.weight.view( +- layer.weight.size(0), +- layer.weight.size(2), +- ) +- .contiguous() +- .clone() ++ # prepack conv weight ++ unpacked = ( ++ layer.weight.view( ++ layer.weight.size(0), ++ layer.weight.size(2), + ) +- # Stash the un-packed (dim, width) weight so the speculative-decode +- # GDN path (which uses torch conv, not the AMX kernel) can use it. +- layer._cpu_unpacked_conv_weight = unpacked +- layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) +- layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear( +- x, weight, bias +- ) ++ .contiguous() ++ .clone() ++ ) ++ # Stash the un-packed (dim, width) weight so the speculative-decode ++ # GDN path (which uses torch conv, not the AMX kernel) can use it. ++ layer._cpu_unpacked_conv_weight = unpacked ++ layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) + return + + N, K = layer.weight.size() diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 4fb3a551b1e6..82297fe84296 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3429,22 +3429,17 @@ def causal_conv1d_update_cpu( conv_state_indices: torch.Tensor | None, is_vnni: bool, ) -> torch.Tensor: - from vllm.platforms import CpuArchEnum - - if current_platform.get_cpu_architecture() == CpuArchEnum.X86: - return torch.ops._C.causal_conv1d_update_cpu( - x, - conv_states, - weight, - bias, - silu_activation, - None, - conv_state_indices, - -1, - is_vnni, - ) - else: - raise NotImplementedError("AMX conv1d is only for x86") + return torch.ops._C.causal_conv1d_update_cpu( + x, + conv_states, + weight, + bias, + silu_activation, + None, + conv_state_indices, + -1, + is_vnni, + ) class CPUDNNLGEMMHandler: diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py index fa556723e846..a68842f4bfee 100644 --- a/vllm/config/mamba.py +++ b/vllm/config/mamba.py @@ -27,9 +27,6 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta): TRITON = "triton" FLASHINFER = "flashinfer" - # Pure-PyTorch fallback for CPU-only platforms (PowerPC, no CUDA, etc.). - # Avoids Triton JIT compilation which is unstable / unsupported on those - # architectures. CPU = "cpu" @@ -38,11 +35,7 @@ class MambaConfig: """Configuration for Mamba SSM backends.""" backend: MambaBackendEnum = MambaBackendEnum.TRITON - """Mamba SSU backend to use. - - On CPU-only platforms (e.g. PowerPC, x86 without CUDA) the default is - automatically overridden to 'cpu' by ``__post_init__``. - """ + """Mamba SSU backend to use.""" enable_stochastic_rounding: bool = False """Enable stochastic rounding when writing SSM state to fp16 cache. diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 27a9da5fbfe3..352f25334069 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -440,96 +440,6 @@ def forward_torch( return output - def forward_batched_gemm( - self, - layer: torch.nn.Module, - input: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: MoEActivation, - global_num_experts: int = -1, - skip_weighted: bool = False, - ) -> torch.Tensor: - """Universal batched GEMM fallback for CPUs without AVX512 fused MoE. - - The per-expert F.linear loop in forward_torch issues N_experts separate - BLAS calls, each triggering a full weight-matrix packing pass - (sbgemm_incopy / similar). For decode (small T), packing dominates. - - This path instead: - 1. One torch.mm for gate+up — packs ALL expert weights in one pass. - 2. One torch.bmm for down — batched over experts. - → 2 BLAS weight-packing passes regardless of N_experts. - - Works on any CPU (x86 non-AVX512, ARM non-AMX, PowerPC, RISC-V, ...). - The trade-off (computing all experts densely) is beneficial whenever - decode batch size is small relative to expert weight size. - """ - if skip_weighted: - # Weight is pre-applied to input; must NOT re-apply in accumulation. - assert topk_ids.size(1) == 1, ( - "apply_router_weight_on_input is only supported for topk=1" - ) - input = input * topk_weights.to(input.dtype) - - num_tokens: int = input.shape[0] - hidden_dim: int = input.shape[1] - num_experts: int = layer.w13_weight.shape[0] - gate_up_dim: int = layer.w13_weight.shape[1] # gate + up concatenated - top_k: int = topk_ids.shape[1] - act_fn = _CPU_MOE_ACT_FN[activation] - - # ------------------------------------------------------------------ - # 1. Gate+Up — single large GEMM covering ALL experts at once. - # w13_weight: (E, gate+up, hidden) → (E*gate+up, hidden) - # mm: (T, hidden) @ (hidden, E*gate+up) → (T, E*gate+up) - # → view (T, E, gate+up) - # ------------------------------------------------------------------ - w13_2d = layer.w13_weight.view(num_experts * gate_up_dim, hidden_dim) - gate_up_all = torch.mm(input, w13_2d.T) # (T, E*gate+up) - if hasattr(layer, "w13_bias") and layer.w13_bias is not None: - gate_up_all = gate_up_all + layer.w13_bias.view(-1) - gate_up_all = gate_up_all.view(num_tokens, num_experts, gate_up_dim) - - # Apply gating activation: (T, E, gate+up) → (T, E, down_dim) - gate_up_act = act_fn(gate_up_all) - - # ------------------------------------------------------------------ - # 2. Down — batched GEMM over experts. - # Permute to (E, T, down_dim) for bmm. - # w2_weight: (E, out, down_dim) → transpose → (E, down_dim, out) - # bmm: (E, T, down_dim) @ (E, down_dim, out) → (E, T, out) - # ------------------------------------------------------------------ - gate_up_t = gate_up_act.permute(1, 0, 2).contiguous() # (E, T, down) - down_all = torch.bmm( - gate_up_t, - layer.w2_weight.transpose(1, 2), - ) # (E, T, out_dim) - if hasattr(layer, "w2_bias") and layer.w2_bias is not None: - down_all = down_all + layer.w2_bias.unsqueeze(1) - down_all = down_all.permute(1, 0, 2) # (T, E, out_dim) - - # ------------------------------------------------------------------ - # 3. Weighted accumulation over top-k routed experts. - # output[t] += topk_weights[t,k] * down_all[t, topk_ids[t,k], :] - # ------------------------------------------------------------------ - out_dim: int = down_all.shape[-1] - output = torch.zeros( - num_tokens, out_dim, dtype=input.dtype, device=input.device - ) - t_idx = torch.arange(num_tokens, device=input.device) - if skip_weighted: - # Weight already baked into input — gather directly, no re-weighting. - expert_ids = topk_ids[:, 0].long() - output.add_(down_all[t_idx, expert_ids, :]) - else: - for k in range(top_k): - expert_ids = topk_ids[:, k].long() # (T,) - w_k = topk_weights[:, k].to(input.dtype) # (T,) - expert_out = down_all[t_idx, expert_ids, :] # (T, out) - output.addcmul_(expert_out, w_k.unsqueeze(-1).expand_as(expert_out)) - - return output def cpu_fused_moe_torch( diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 022cdfc012c8..bb1a3390e98f 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -12,7 +12,6 @@ from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID -from .cpu_fallbacks import _causal_conv1d_fn_cpu, _causal_conv1d_update_cpu @triton.jit() @@ -1243,10 +1242,6 @@ def grid(META): def causal_conv1d_fn(*args, **kwargs): - """Dispatch causal_conv1d_fn to CPU PyTorch fallback or CUDA Triton kernel.""" - x = args[0] if args else kwargs.get("x") - if x is not None and x.device.type == "cpu": - return _causal_conv1d_fn_cpu(*args, **kwargs) return _causal_conv1d_fn_cuda(*args, **kwargs) @@ -1266,47 +1261,6 @@ def causal_conv1d_update( validate_data=False, ): """Dispatch causal_conv1d_update to CPU C++ kernel or CUDA Triton kernel.""" - if x.device.type == "cpu": - # The C++ kernel handles the standard (non-varlen, non-spec-decoding) - # decode path. Fall back to PyTorch for the complex varlen / - # spec-decoding paths. - if query_start_loc is not None or num_accepted_tokens is not None: - return _causal_conv1d_update_cpu( - x, - conv_state, - weight, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - query_start_loc=query_start_loc, - ) - # Determine activation string - act_str = None - if isinstance(activation, bool): - act_str = "silu" if activation else None - elif activation is not None: - act_str = activation - - _conv_state_indices = conv_state_indices - if _conv_state_indices is not None and _conv_state_indices.dim() == 2: - if initial_state_idx is not None: - _conv_state_indices = _conv_state_indices.gather( - 1, initial_state_idx.unsqueeze(1) - ).squeeze(1) - else: - _conv_state_indices = _conv_state_indices[:, 0] - - pad_slot_id = int(NULL_BLOCK_ID) - return ops.causal_conv1d_update_cpu_vec( - x, - conv_state, - weight, - bias, - act_str, - _conv_state_indices, - None, # query_start_loc - pad_slot_id, - ) return _causal_conv1d_update_cuda( x, @@ -1323,3 +1277,12 @@ def causal_conv1d_update( initial_state_idx=initial_state_idx, validate_data=validate_data, ) + + +from vllm.platforms import current_platform + +if current_platform.is_cpu(): + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_fn_cpu as causal_conv1d_fn, + causal_conv1d_update_cpu as causal_conv1d_update, + ) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py index b047ca6d6169..dba752624bba 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py @@ -4,10 +4,182 @@ from __future__ import annotations import torch -import torch.nn.functional as F +from vllm.v1.attention.backends.utils import PAD_SLOT_ID + + +def causal_conv1d_fn_cpu( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, + activation: str | None = "silu", + pad_slot_id: int = PAD_SLOT_ID, + **kwargs, +) -> torch.Tensor: + """CPU implementation for causal_conv1d_fwd.""" + if isinstance(activation, bool) and activation: + activation = "silu" + + 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 + + if seq_len == 0: + continue + + cache_idx = cache_indices[b].item() if cache_indices is not None else b + + if cache_idx == 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) + 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 + + if state_len > 1: + state[:, :-1] = state[:, 1:].clone() + state[:, -1] = x_t + + conv_states[cache_idx].copy_(state) + + return out.to(original_x_dtype) + + +def causal_conv1d_update_cpu( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: bool | str | None = None, + conv_state_indices: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + pad_slot_id: int = PAD_SLOT_ID, + **kwargs, +) -> torch.Tensor: + """CPU implementation for causal_conv1d_update.""" + if isinstance(activation, bool): + activation = "silu" if activation else None + + original_x_dtype = x.dtype + x = x.to(conv_state.dtype) + _, width = weight.shape + state_len = width - 1 + + if query_start_loc is None and x.dim() == 2: + x = x.unsqueeze(-1) + unsqueeze = True + else: + unsqueeze = False + + if query_start_loc is None: + batch, dim, seqlen = x.shape + + if conv_state_indices is not None: + cache_idxs = conv_state_indices.flatten() + valid_mask = cache_idxs != pad_slot_id + else: + cache_idxs = torch.arange(batch, device=x.device) + valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) + + for t in range(seqlen): + x_t = x[:, :, t].clone() + + states = conv_state[cache_idxs] + + windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) + + val = (windows * weight.unsqueeze(0)).sum(dim=-1) + if bias is not None: + val = val + bias.unsqueeze(0) + + if activation in ["silu", "swish"]: + val = val * torch.sigmoid(val) + + val = val * valid_mask.unsqueeze(-1).to(val.dtype) + x[:, :, t] = val + + new_state = torch.cat([states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1) + conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] + + out = x + if unsqueeze: + out = out.squeeze(-1) + return out.to(original_x_dtype) + + assert conv_state_indices is not None + assert query_start_loc is not None + batch = conv_state_indices.size(0) + out = x.clone() + + for b in range(batch): + cache_idx = conv_state_indices[b].item() + if cache_idx == pad_slot_id: + continue + + seq_start = query_start_loc[b].item() + seq_end = query_start_loc[b + 1].item() + seqlen_b = seq_end - seq_start + + if seqlen_b == 0: + continue + + local_state = conv_state[cache_idx].clone() + + for t in range(seqlen_b): + x_t = x[seq_start + t, :] + + window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) + val = (window * weight).sum(dim=-1) + if bias is not None: + val = val + bias + if activation in ["silu", "swish"]: + val = val * torch.sigmoid(val) + + out[seq_start + t, :] = val + + if state_len > 1: + local_state[:, :-1] = local_state[:, 1:].clone() + local_state[:, -1] = x_t + + conv_state[cache_idx] = local_state + + return out.to(original_x_dtype) + + +import torch.nn.functional as F -# for prefill def causal_conv1d_torch( x: torch.Tensor, weight: torch.Tensor, @@ -60,7 +232,6 @@ def causal_conv1d_torch( return out -# for decode def causal_conv1d_update_torch( x: torch.Tensor, conv_state: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py new file mode 100644 index 000000000000..aa1c122731ea --- /dev/null +++ b/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch + +import vllm._custom_ops as ops +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + + +def _mamba_chunk_scan_combined_fwd_cpu( + x, + dt, + A, + B, + C, + chunk_size, + out, + D=None, + z=None, + dt_bias=None, + initial_states=None, + return_intermediate_states=False, + seq_idx=None, + cu_seqlens=None, + cu_chunk_seqlens=None, + last_chunk_indices=None, + dt_softplus=False, + dt_limit=(0.0, float("inf")), + state_dtype=None, + **kwargs, +): + seqlen, nheads, headdim = x.shape + _, ngroups, dstate = B.shape + + assert cu_seqlens is not None + batch = cu_seqlens.size(0) - 1 + + + dt_f = dt.float() + if dt_bias is not None: + dt_f = dt_f + dt_bias.float().unsqueeze(0) + if dt_softplus: + dt_f = torch.nn.functional.softplus(dt_f) + if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): + dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) + + + all_states = torch.zeros( + batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device + ) + if initial_states is not None: + all_states.copy_(initial_states.float()) + + + assert out.is_contiguous(), ( + "_mamba_chunk_scan_combined_fwd_cpu: `out` must be " + "pre-allocated as a contiguous tensor" + ) + + + D_1d = None + if D is not None: + d = D.float() + while d.dim() > 1 and d.stride(-1) == 0: + d = d.squeeze(-1) + D_1d = d.contiguous() + + ops.mamba_chunk_scan_fwd_cpu( + out, + all_states, + x, + dt_f, + A, + B, + C, + D_1d, + z, + cu_seqlens.to(torch.int32), + ) + + out_dtype = state_dtype if state_dtype is not None else x.dtype + all_states = all_states.to(out_dtype) + + return all_states + + +def selective_state_update( + state, + x, + dt, + A, + B, + C, + D=None, + z=None, + dt_bias=None, + dt_softplus=False, + state_batch_indices=None, + dst_state_batch_indices=None, + null_block_id=NULL_BLOCK_ID, + out=None, + num_accepted_tokens=None, + cu_seqlens=None, + is_blackwell=False, + enable_stochastic_rounding=False, + cache_philox_rounds=0, +): + """CPU implementation for selective_state_update.""" + # Ensure out tensor exists + if out is None: + out = torch.empty_like(x if x.dim() == 2 else x) + + + _state = state.unsqueeze(1) if state.dim() == 3 else state + _x = x.unsqueeze(1) if x.dim() == 2 else x + _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt + _A = A.unsqueeze(0) if A.dim() == 2 else A + _B = B.unsqueeze(1) if B.dim() == 2 else B + _C = C.unsqueeze(1) if C.dim() == 2 else C + _D = D.unsqueeze(0) if (D is not None and D.dim() == 1) else D + _z = z.unsqueeze(1) if (z is not None and z.dim() == 2) else z + _dt_bias = ( + dt_bias.unsqueeze(0) + if (dt_bias is not None and dt_bias.dim() == 1) + else dt_bias + ) + _out = out.unsqueeze(1) if out.dim() == 2 else out + + _sbi = state_batch_indices + _dsbi = dst_state_batch_indices + ops.selective_state_update_cpu( + _state, + _x, + _dt, + _A, + _B, + _C, + _D, + _z, + _dt_bias, + dt_softplus, + _sbi, + _dsbi, + null_block_id, + _out, + num_accepted_tokens, + cu_seqlens, + ) + return _out.squeeze(1) if out.dim() == 2 else _out diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py deleted file mode 100644 index 50a3bb392f54..000000000000 --- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py +++ /dev/null @@ -1,306 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import torch - -import vllm._custom_ops as ops -from vllm.v1.attention.backends.utils import PAD_SLOT_ID - - -def _causal_conv1d_fn_cpu( - x: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None, - conv_states: torch.Tensor, - query_start_loc: torch.Tensor, - cache_indices: torch.Tensor | None = None, - has_initial_state: torch.Tensor | None = None, - activation: str | None = "silu", - pad_slot_id: int = PAD_SLOT_ID, - **kwargs, -) -> torch.Tensor: - """Pure PyTorch CPU fallback for causal_conv1d_fwd.""" - if isinstance(activation, bool) and activation: - activation = "silu" - - 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 - - if seq_len == 0: - continue - - cache_idx = cache_indices[b].item() if cache_indices is not None else b - - if cache_idx == 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) - 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 - - if state_len > 1: - state[:, :-1] = state[:, 1:].clone() - state[:, -1] = x_t - - conv_states[cache_idx].copy_(state) - - return out.to(original_x_dtype) - - -def _causal_conv1d_update_cpu( - x: torch.Tensor, - conv_state: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None = None, - activation: bool | str | None = None, - conv_state_indices: torch.Tensor | None = None, - query_start_loc: torch.Tensor | None = None, - pad_slot_id: int = PAD_SLOT_ID, - **kwargs, -) -> torch.Tensor: - """Pure PyTorch CPU fallback for causal_conv1d_update (decode path).""" - if isinstance(activation, bool): - activation = "silu" if activation else None - - original_x_dtype = x.dtype - x = x.to(conv_state.dtype) - _, width = weight.shape - state_len = width - 1 - - if query_start_loc is None and x.dim() == 2: - x = x.unsqueeze(-1) - unsqueeze = True - else: - unsqueeze = False - - if query_start_loc is None: - batch, dim, seqlen = x.shape - - if conv_state_indices is not None: - cache_idxs = conv_state_indices.flatten() - valid_mask = cache_idxs != pad_slot_id - else: - cache_idxs = torch.arange(batch, device=x.device) - valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) - - for t in range(seqlen): - x_t = x[:, :, t].clone() - - states = conv_state[cache_idxs] - - windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) - - val = (windows * weight.unsqueeze(0)).sum(dim=-1) - if bias is not None: - val = val + bias.unsqueeze(0) - - if activation in ["silu", "swish"]: - val = val * torch.sigmoid(val) - - val = val * valid_mask.unsqueeze(-1).to(val.dtype) - x[:, :, t] = val - - new_state = torch.cat([states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1) - conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] - - out = x - if unsqueeze: - out = out.squeeze(-1) - return out.to(original_x_dtype) - - assert conv_state_indices is not None - assert query_start_loc is not None - batch = conv_state_indices.size(0) - out = x.clone() - - for b in range(batch): - cache_idx = conv_state_indices[b].item() - if cache_idx == pad_slot_id: - continue - - seq_start = query_start_loc[b].item() - seq_end = query_start_loc[b + 1].item() - seqlen_b = seq_end - seq_start - - if seqlen_b == 0: - continue - - local_state = conv_state[cache_idx].clone() - - for t in range(seqlen_b): - x_t = x[seq_start + t, :] - - window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) - val = (window * weight).sum(dim=-1) - if bias is not None: - val = val + bias - if activation in ["silu", "swish"]: - val = val * torch.sigmoid(val) - - out[seq_start + t, :] = val - - if state_len > 1: - local_state[:, :-1] = local_state[:, 1:].clone() - local_state[:, -1] = x_t - - conv_state[cache_idx] = local_state - - return out.to(original_x_dtype) - - -def _mamba_chunk_scan_combined_fwd_cpu( - x, - dt, - A, - B, - C, - chunk_size, - out, - D=None, - z=None, - dt_bias=None, - initial_states=None, - return_intermediate_states=False, - seq_idx=None, - cu_seqlens=None, - cu_chunk_seqlens=None, - last_chunk_indices=None, - dt_softplus=False, - dt_limit=(0.0, float("inf")), - state_dtype=None, - **kwargs, -): - seqlen, nheads, headdim = x.shape - _, ngroups, dstate = B.shape - - assert cu_seqlens is not None - batch = cu_seqlens.size(0) - 1 - - # Preprocess dt: apply bias, softplus, and clamp. - # Kept in Python to avoid duplicating this logic in C++. - dt_f = dt.float() - if dt_bias is not None: - dt_f = dt_f + dt_bias.float().unsqueeze(0) - if dt_softplus: - dt_f = torch.nn.functional.softplus(dt_f) - if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): - dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) - - # Allocate state output buffer (float32, contiguous — required by kernel). - all_states = torch.zeros( - batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device - ) - if initial_states is not None: - all_states.copy_(initial_states.float()) - - # Use the C++ kernel when available (CPU build with mamba kernels compiled). - # Falls back to the pure-Python loop below if the op is not registered. - _use_cpp_kernel = hasattr(torch.ops._C, "mamba_chunk_scan_fwd_cpu") - - if _use_cpp_kernel: - # out must be contiguous — the C++ kernel writes via raw data_ptr(). - # Creating a contiguous copy here would discard the results, so we - # require the caller to pass a contiguous tensor. - assert out.is_contiguous(), ( - "_mamba_chunk_scan_combined_fwd_cpu: `out` must be " - "pre-allocated as a contiguous tensor" - ) - - # D: strip trailing dims that are broadcast (stride == 0) so the - # kernel sees a (nheads,) float32 array. Only peel dims whose - # stride is 0 (broadcast-expanded); genuine multi-dim D is passed - # as-is and flattened by the C++ wrapper. - D_1d = None - if D is not None: - d = D.float() - while d.dim() > 1 and d.stride(-1) == 0: - d = d.squeeze(-1) - D_1d = d.contiguous() - - ops.mamba_chunk_scan_fwd_cpu( - out, - all_states, - x, - dt_f, - A, - B, - C, - D_1d, - z, - cu_seqlens.to(torch.int32), - ) - else: - # Pure-Python fallback (no compiled extension available). - for b_idx in range(batch): - seq_start = cu_seqlens[b_idx].item() - seq_end = cu_seqlens[b_idx + 1].item() - - # Note: basic indexing returns a view, but arithmetic ops - # below (state * dA + dBx) produce new tensors, so state is - # reassigned each token. all_states[b_idx] is committed at line 416. - state = all_states[b_idx].clone() # local working copy - - for t in range(seq_start, seq_end): - x_t = x[t].float() - dt_t = dt_f[t] - A_val = A.float() - - dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) - - B_expanded = B[t].float().repeat_interleave(nheads // ngroups, dim=0) - C_expanded = C[t].float().repeat_interleave(nheads // ngroups, dim=0) - - xdt = x_t * dt_t.unsqueeze(-1) - dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) - state = state * dA + dBx - - y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) - - if D is not None: - y = ( - y + x_t * D.float().unsqueeze(-1) - if D.dim() == 1 - else y + x_t * D.float() - ) - - if z is not None: - z_t = z[t].float() - y = y * z_t * torch.sigmoid(z_t) - - out[t] = y.to(out.dtype) - - all_states[b_idx] = state.to(all_states.dtype) - - out_dtype = state_dtype if state_dtype is not None else x.dtype - all_states = all_states.to(out_dtype) - - return all_states diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index d7055fbd741a..9d463e570147 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -894,47 +894,6 @@ def selective_state_update( if out is None: out = torch.empty_like(x if x.dim() == 2 else x) - if x.device.type == "cpu": - # Reshape tensors from (batch, dim) -> (batch, 1, dim) if needed - # The C++ kernel expects (N, nheads, dim) layout - _state = state.unsqueeze(1) if state.dim() == 3 else state - _x = x.unsqueeze(1) if x.dim() == 2 else x - _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt - _A = A.unsqueeze(0) if A.dim() == 2 else A - _B = B.unsqueeze(1) if B.dim() == 2 else B - _C = C.unsqueeze(1) if C.dim() == 2 else C - _D = D.unsqueeze(0) if (D is not None and D.dim() == 1) else D - _z = z.unsqueeze(1) if (z is not None and z.dim() == 2) else z - _dt_bias = ( - dt_bias.unsqueeze(0) - if (dt_bias is not None and dt_bias.dim() == 1) - else dt_bias - ) - _out = out.unsqueeze(1) if out.dim() == 2 else out - # state_batch_indices and dst_state_batch_indices are 1D index arrays; - # do NOT reshape them. - _sbi = state_batch_indices - _dsbi = dst_state_batch_indices - ops.selective_state_update_cpu( - _state, - _x, - _dt, - _A, - _B, - _C, - _D, - _z, - _dt_bias, - dt_softplus, - _sbi, - _dsbi, - null_block_id, - _out, - num_accepted_tokens, - cu_seqlens, - ) - return _out.squeeze(1) if out.dim() == 2 else _out - return _selective_state_update_cuda( state, x, @@ -956,3 +915,9 @@ def selective_state_update( enable_stochastic_rounding=enable_stochastic_rounding, cache_philox_rounds=cache_philox_rounds, ) + + +if current_platform.is_cpu(): + from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( + selective_state_update, + ) diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index caa57211e3d9..89c8bd814d27 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -13,7 +13,7 @@ from vllm.model_executor.custom_op import CustomOp from vllm.triton_utils import HAS_TRITON, triton -from .cpu_fallbacks import _mamba_chunk_scan_combined_fwd_cpu + from .ssd_bmm import _bmm_chunk_fwd from .ssd_chunk_scan import _chunk_scan_fwd from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd @@ -166,19 +166,11 @@ def __init__(self): super(CustomOp, self).__init__() # nn.Module.__init__ only from vllm.platforms import current_platform - if current_platform.is_cpu(): - self._forward_method = self.forward_cpu - elif current_platform.is_rocm(): + if current_platform.is_rocm(): self._forward_method = self.forward_hip else: self._forward_method = self.forward_cuda - def forward_native(self, *args, **kwargs): - return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) - - def forward_cpu(self, *args, **kwargs): - return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) - def forward_cuda(self, *args, **kwargs): return _mamba_chunk_scan_combined_fwd_cuda(*args, **kwargs) @@ -260,7 +252,15 @@ def mamba_chunk_scan_combined_varlen( last_chunk_indices=last_chunk_indices, dt_softplus=dt_softplus, dt_limit=dt_limit, + return_intermediate_states=return_intermediate_states, state_dtype=state_dtype, ) - return varlen_states + + +from vllm.platforms import current_platform + +if current_platform.is_cpu(): + from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( + _mamba_chunk_scan_combined_fwd_cpu as _mamba_chunk_scan_combined_fwd, + ) diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 1ceeac665f98..1795a036b124 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -6,8 +6,7 @@ Provides a unified `selective_state_update` function that dispatches to the Triton, FlashInfer, or CPU backend based on the configured `MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA) -the backend defaults to 'cpu', which uses a pure-PyTorch fallback that -avoids Triton JIT compilation entirely. +the backend defaults to 'cpu'. """ from abc import ABC, abstractmethod @@ -280,7 +279,7 @@ def initialize_mamba_ssu_backend( backend = mamba_config.backend # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is - # unstable or unavailable. Silently fall back to the pure-PyTorch CPU + # unstable or unavailable. Silently fall back to the CPU # backend unless the user explicitly chose something other than "triton". if backend == MambaBackendEnum.TRITON: from vllm.platforms import current_platform @@ -288,7 +287,7 @@ def initialize_mamba_ssu_backend( if current_platform.is_cpu(): logger.info( "CPU platform detected: overriding Mamba SSU backend " - "from 'triton' to 'cpu' (pure-PyTorch fallback)." + "from 'triton' to 'cpu'." ) backend = MambaBackendEnum.CPU diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index 450c74a6095b..f0d3e6eaddfe 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -243,23 +243,19 @@ def dispatch_cpu_unquantized_gemm( if torch.cpu._is_amx_tile_supported() and hasattr( ops, "causal_conv1d_weight_pack" ): - with contextlib.suppress(Exception): - # prepack conv weight - unpacked = ( - layer.weight.view( - layer.weight.size(0), - layer.weight.size(2), - ) - .contiguous() - .clone() + # prepack conv weight + unpacked = ( + layer.weight.view( + layer.weight.size(0), + layer.weight.size(2), ) - # Stash the un-packed (dim, width) weight so the speculative-decode - # GDN path (which uses torch conv, not the AMX kernel) can use it. - layer._cpu_unpacked_conv_weight = unpacked - layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) - layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear( - x, weight, bias - ) + .contiguous() + .clone() + ) + # Stash the un-packed (dim, width) weight so the speculative-decode + # GDN path (which uses torch conv, not the AMX kernel) can use it. + layer._cpu_unpacked_conv_weight = unpacked + layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) return N, K = layer.weight.size() From 9dc951f6b2e4db27d8db9ea89f8db5b343913472 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 14 Jul 2026 14:49:37 +0530 Subject: [PATCH 33/43] pre-commit fixes Signed-off-by: Akash kaothalkar --- vllm/model_executor/layers/mamba/ops/causal_conv1d.py | 4 ++-- vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py | 2 +- vllm/model_executor/layers/mamba/ops/mamba_ssm.py | 2 +- vllm/model_executor/layers/mamba/ops/ssd_combined.py | 5 ++--- 4 files changed, 6 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index bb1a3390e98f..ab4b224b2ab7 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -1279,10 +1279,10 @@ def causal_conv1d_update( ) -from vllm.platforms import current_platform +from vllm.platforms import current_platform # noqa: F401, E402 if current_platform.is_cpu(): - from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( # noqa: F401, E402 causal_conv1d_fn_cpu as causal_conv1d_fn, causal_conv1d_update_cpu as causal_conv1d_update, ) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py index dba752624bba..ffe7f675c520 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py @@ -4,6 +4,7 @@ from __future__ import annotations import torch +import torch.nn.functional as F from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -178,7 +179,6 @@ def causal_conv1d_update_cpu( return out.to(original_x_dtype) -import torch.nn.functional as F def causal_conv1d_torch( x: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index 9d463e570147..bfe883596d2d 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -918,6 +918,6 @@ def selective_state_update( if current_platform.is_cpu(): - from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( + from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( # noqa: F401, E402 selective_state_update, ) diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 89c8bd814d27..bae8ea32d73c 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -252,15 +252,14 @@ def mamba_chunk_scan_combined_varlen( last_chunk_indices=last_chunk_indices, dt_softplus=dt_softplus, dt_limit=dt_limit, - return_intermediate_states=return_intermediate_states, state_dtype=state_dtype, ) return varlen_states -from vllm.platforms import current_platform +from vllm.platforms import current_platform # noqa: F401, E402 if current_platform.is_cpu(): - from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( + from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( # noqa: F401, E402 _mamba_chunk_scan_combined_fwd_cpu as _mamba_chunk_scan_combined_fwd, ) From b285ec61d4aa113fc5dc555c11d645c3f45af443 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 14 Jul 2026 07:28:01 -0500 Subject: [PATCH 34/43] mypy fix Signed-off-by: Akash kaothalkar --- .../model_executor/layers/fused_moe/cpu_fused_moe.py | 1 - .../model_executor/layers/mamba/ops/causal_conv1d.py | 12 +++++------- .../layers/mamba/ops/cpu/causal_conv1d.py | 1 - .../model_executor/layers/mamba/ops/cpu/mamba_ssm.py | 5 ----- vllm/model_executor/layers/mamba/ops/mamba_ssm.py | 6 +++--- vllm/model_executor/layers/mamba/ops/ssd_combined.py | 9 ++++----- vllm/model_executor/layers/utils.py | 1 - 7 files changed, 12 insertions(+), 23 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py index 352f25334069..1a0acff058cc 100644 --- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py @@ -441,7 +441,6 @@ def forward_torch( return output - def cpu_fused_moe_torch( layer_id: int, output: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index ab4b224b2ab7..7ba730e92400 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -8,12 +8,10 @@ import numpy as np import torch -from vllm import _custom_ops as ops from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID - @triton.jit() def _causal_conv1d_fwd_kernel( # continuous batching # Pointers to matrices @@ -1279,10 +1277,10 @@ def causal_conv1d_update( ) -from vllm.platforms import current_platform # noqa: F401, E402 +from vllm.platforms import current_platform # noqa: E402 if current_platform.is_cpu(): - from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( # noqa: F401, E402 - causal_conv1d_fn_cpu as causal_conv1d_fn, - causal_conv1d_update_cpu as causal_conv1d_update, - ) + import vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d as cpu_conv1d # noqa: E402 + + causal_conv1d_fn = cpu_conv1d.causal_conv1d_fn_cpu # type: ignore + causal_conv1d_update = cpu_conv1d.causal_conv1d_update_cpu # type: ignore diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py index ffe7f675c520..3715cfd6a2d9 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py @@ -179,7 +179,6 @@ def causal_conv1d_update_cpu( return out.to(original_x_dtype) - def causal_conv1d_torch( x: torch.Tensor, weight: torch.Tensor, diff --git a/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py index aa1c122731ea..a65793d79245 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/mamba_ssm.py @@ -35,7 +35,6 @@ def _mamba_chunk_scan_combined_fwd_cpu( assert cu_seqlens is not None batch = cu_seqlens.size(0) - 1 - dt_f = dt.float() if dt_bias is not None: dt_f = dt_f + dt_bias.float().unsqueeze(0) @@ -44,20 +43,17 @@ def _mamba_chunk_scan_combined_fwd_cpu( if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) - all_states = torch.zeros( batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device ) if initial_states is not None: all_states.copy_(initial_states.float()) - assert out.is_contiguous(), ( "_mamba_chunk_scan_combined_fwd_cpu: `out` must be " "pre-allocated as a contiguous tensor" ) - D_1d = None if D is not None: d = D.float() @@ -110,7 +106,6 @@ def selective_state_update( if out is None: out = torch.empty_like(x if x.dim() == 2 else x) - _state = state.unsqueeze(1) if state.dim() == 3 else state _x = x.unsqueeze(1) if x.dim() == 2 else x _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index bfe883596d2d..a2b279d23831 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -918,6 +918,6 @@ def selective_state_update( if current_platform.is_cpu(): - from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( # noqa: F401, E402 - selective_state_update, - ) + import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_ssm # noqa: E402 + + selective_state_update = cpu_ssm.selective_state_update diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index bae8ea32d73c..626792afd0f3 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -13,7 +13,6 @@ from vllm.model_executor.custom_op import CustomOp from vllm.triton_utils import HAS_TRITON, triton - from .ssd_bmm import _bmm_chunk_fwd from .ssd_chunk_scan import _chunk_scan_fwd from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd @@ -257,9 +256,9 @@ def mamba_chunk_scan_combined_varlen( return varlen_states -from vllm.platforms import current_platform # noqa: F401, E402 +from vllm.platforms import current_platform # noqa: E402 if current_platform.is_cpu(): - from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( # noqa: F401, E402 - _mamba_chunk_scan_combined_fwd_cpu as _mamba_chunk_scan_combined_fwd, - ) + import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_ssm # noqa: E402 + + _mamba_chunk_scan_combined_fwd = cpu_ssm._mamba_chunk_scan_combined_fwd_cpu diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py index f0d3e6eaddfe..2b19b3d3c4e3 100644 --- a/vllm/model_executor/layers/utils.py +++ b/vllm/model_executor/layers/utils.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Utility methods for model layers.""" -import contextlib from collections.abc import Callable import torch From e9b6c3c01226046decc18289f56bd940dd1e237e Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Tue, 14 Jul 2026 22:44:22 -0500 Subject: [PATCH 35/43] remove redundant file Signed-off-by: Akash kaothalkar --- mamba_cpu_updates.patch | 968 ---------------------------------------- 1 file changed, 968 deletions(-) delete mode 100644 mamba_cpu_updates.patch diff --git a/mamba_cpu_updates.patch b/mamba_cpu_updates.patch deleted file mode 100644 index 5f128bdf7768..000000000000 --- a/mamba_cpu_updates.patch +++ /dev/null @@ -1,968 +0,0 @@ -diff --git a/.buildkite/hardware_tests/cpu.yaml b/.buildkite/hardware_tests/cpu.yaml -index ebfd1c752..fd2d45be5 100644 ---- a/.buildkite/hardware_tests/cpu.yaml -+++ b/.buildkite/hardware_tests/cpu.yaml -@@ -18,6 +18,8 @@ steps: - - tests/kernels/quantization/test_cpu_fp8_scaled_mm.py - - tests/kernels/mamba/cpu/test_cpu_gdn_ops.py - - tests/kernels/mamba/test_cpu_short_conv.py -+ - tests/kernels/mamba/test_causal_conv1d.py -+ - tests/kernels/mamba/test_mamba_ssm.py - commands: - - | - bash .buildkite/scripts/hardware_ci/run-cpu-test.sh 30m " -@@ -28,7 +30,9 @@ steps: - pytest -x -v -s tests/kernels/test_onednn.py - pytest -x -v -s tests/kernels/test_awq_int4_to_int8.py - pytest -x -v -s tests/kernels/quantization/test_cpu_fp8_scaled_mm.py -- pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py" -+ pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py -+ pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py -+ pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py" - - # Note: SDE can't be downloaded from CI host because of AWS WAF - # - label: CPU-Compatibility Tests -diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py -index 4fb3a551b..82297fe84 100644 ---- a/vllm/_custom_ops.py -+++ b/vllm/_custom_ops.py -@@ -3429,22 +3429,17 @@ def causal_conv1d_update_cpu( - conv_state_indices: torch.Tensor | None, - is_vnni: bool, - ) -> torch.Tensor: -- from vllm.platforms import CpuArchEnum -- -- if current_platform.get_cpu_architecture() == CpuArchEnum.X86: -- return torch.ops._C.causal_conv1d_update_cpu( -- x, -- conv_states, -- weight, -- bias, -- silu_activation, -- None, -- conv_state_indices, -- -1, -- is_vnni, -- ) -- else: -- raise NotImplementedError("AMX conv1d is only for x86") -+ return torch.ops._C.causal_conv1d_update_cpu( -+ x, -+ conv_states, -+ weight, -+ bias, -+ silu_activation, -+ None, -+ conv_state_indices, -+ -1, -+ is_vnni, -+ ) - - - class CPUDNNLGEMMHandler: -diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py -index fa556723e..a68842f4b 100644 ---- a/vllm/config/mamba.py -+++ b/vllm/config/mamba.py -@@ -27,9 +27,6 @@ class MambaBackendEnum(Enum, metaclass=_MambaBackendEnumMeta): - - TRITON = "triton" - FLASHINFER = "flashinfer" -- # Pure-PyTorch fallback for CPU-only platforms (PowerPC, no CUDA, etc.). -- # Avoids Triton JIT compilation which is unstable / unsupported on those -- # architectures. - CPU = "cpu" - - -@@ -38,11 +35,7 @@ class MambaConfig: - """Configuration for Mamba SSM backends.""" - - backend: MambaBackendEnum = MambaBackendEnum.TRITON -- """Mamba SSU backend to use. -- -- On CPU-only platforms (e.g. PowerPC, x86 without CUDA) the default is -- automatically overridden to 'cpu' by ``__post_init__``. -- """ -+ """Mamba SSU backend to use.""" - - enable_stochastic_rounding: bool = False - """Enable stochastic rounding when writing SSM state to fp16 cache. -diff --git a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py -index 27a9da5fb..352f25334 100644 ---- a/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py -+++ b/vllm/model_executor/layers/fused_moe/cpu_fused_moe.py -@@ -440,96 +440,6 @@ class CPUFusedMOE: - - return output - -- def forward_batched_gemm( -- self, -- layer: torch.nn.Module, -- input: torch.Tensor, -- topk_weights: torch.Tensor, -- topk_ids: torch.Tensor, -- activation: MoEActivation, -- global_num_experts: int = -1, -- skip_weighted: bool = False, -- ) -> torch.Tensor: -- """Universal batched GEMM fallback for CPUs without AVX512 fused MoE. -- -- The per-expert F.linear loop in forward_torch issues N_experts separate -- BLAS calls, each triggering a full weight-matrix packing pass -- (sbgemm_incopy / similar). For decode (small T), packing dominates. -- -- This path instead: -- 1. One torch.mm for gate+up — packs ALL expert weights in one pass. -- 2. One torch.bmm for down — batched over experts. -- → 2 BLAS weight-packing passes regardless of N_experts. -- -- Works on any CPU (x86 non-AVX512, ARM non-AMX, PowerPC, RISC-V, ...). -- The trade-off (computing all experts densely) is beneficial whenever -- decode batch size is small relative to expert weight size. -- """ -- if skip_weighted: -- # Weight is pre-applied to input; must NOT re-apply in accumulation. -- assert topk_ids.size(1) == 1, ( -- "apply_router_weight_on_input is only supported for topk=1" -- ) -- input = input * topk_weights.to(input.dtype) -- -- num_tokens: int = input.shape[0] -- hidden_dim: int = input.shape[1] -- num_experts: int = layer.w13_weight.shape[0] -- gate_up_dim: int = layer.w13_weight.shape[1] # gate + up concatenated -- top_k: int = topk_ids.shape[1] -- act_fn = _CPU_MOE_ACT_FN[activation] -- -- # ------------------------------------------------------------------ -- # 1. Gate+Up — single large GEMM covering ALL experts at once. -- # w13_weight: (E, gate+up, hidden) → (E*gate+up, hidden) -- # mm: (T, hidden) @ (hidden, E*gate+up) → (T, E*gate+up) -- # → view (T, E, gate+up) -- # ------------------------------------------------------------------ -- w13_2d = layer.w13_weight.view(num_experts * gate_up_dim, hidden_dim) -- gate_up_all = torch.mm(input, w13_2d.T) # (T, E*gate+up) -- if hasattr(layer, "w13_bias") and layer.w13_bias is not None: -- gate_up_all = gate_up_all + layer.w13_bias.view(-1) -- gate_up_all = gate_up_all.view(num_tokens, num_experts, gate_up_dim) -- -- # Apply gating activation: (T, E, gate+up) → (T, E, down_dim) -- gate_up_act = act_fn(gate_up_all) -- -- # ------------------------------------------------------------------ -- # 2. Down — batched GEMM over experts. -- # Permute to (E, T, down_dim) for bmm. -- # w2_weight: (E, out, down_dim) → transpose → (E, down_dim, out) -- # bmm: (E, T, down_dim) @ (E, down_dim, out) → (E, T, out) -- # ------------------------------------------------------------------ -- gate_up_t = gate_up_act.permute(1, 0, 2).contiguous() # (E, T, down) -- down_all = torch.bmm( -- gate_up_t, -- layer.w2_weight.transpose(1, 2), -- ) # (E, T, out_dim) -- if hasattr(layer, "w2_bias") and layer.w2_bias is not None: -- down_all = down_all + layer.w2_bias.unsqueeze(1) -- down_all = down_all.permute(1, 0, 2) # (T, E, out_dim) -- -- # ------------------------------------------------------------------ -- # 3. Weighted accumulation over top-k routed experts. -- # output[t] += topk_weights[t,k] * down_all[t, topk_ids[t,k], :] -- # ------------------------------------------------------------------ -- out_dim: int = down_all.shape[-1] -- output = torch.zeros( -- num_tokens, out_dim, dtype=input.dtype, device=input.device -- ) -- t_idx = torch.arange(num_tokens, device=input.device) -- if skip_weighted: -- # Weight already baked into input — gather directly, no re-weighting. -- expert_ids = topk_ids[:, 0].long() -- output.add_(down_all[t_idx, expert_ids, :]) -- else: -- for k in range(top_k): -- expert_ids = topk_ids[:, k].long() # (T,) -- w_k = topk_weights[:, k].to(input.dtype) # (T,) -- expert_out = down_all[t_idx, expert_ids, :] # (T, out) -- output.addcmul_(expert_out, w_k.unsqueeze(-1).expand_as(expert_out)) -- -- return output - - - def cpu_fused_moe_torch( -diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py -index 022cdfc01..bb1a3390e 100644 ---- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py -+++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py -@@ -12,7 +12,6 @@ from vllm import _custom_ops as ops - from vllm.triton_utils import tl, triton - from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID - --from .cpu_fallbacks import _causal_conv1d_fn_cpu, _causal_conv1d_update_cpu - - - @triton.jit() -@@ -1243,10 +1242,6 @@ def _causal_conv1d_update_cuda( - - - def causal_conv1d_fn(*args, **kwargs): -- """Dispatch causal_conv1d_fn to CPU PyTorch fallback or CUDA Triton kernel.""" -- x = args[0] if args else kwargs.get("x") -- if x is not None and x.device.type == "cpu": -- return _causal_conv1d_fn_cpu(*args, **kwargs) - return _causal_conv1d_fn_cuda(*args, **kwargs) - - -@@ -1266,47 +1261,6 @@ def causal_conv1d_update( - validate_data=False, - ): - """Dispatch causal_conv1d_update to CPU C++ kernel or CUDA Triton kernel.""" -- if x.device.type == "cpu": -- # The C++ kernel handles the standard (non-varlen, non-spec-decoding) -- # decode path. Fall back to PyTorch for the complex varlen / -- # spec-decoding paths. -- if query_start_loc is not None or num_accepted_tokens is not None: -- return _causal_conv1d_update_cpu( -- x, -- conv_state, -- weight, -- bias=bias, -- activation=activation, -- conv_state_indices=conv_state_indices, -- query_start_loc=query_start_loc, -- ) -- # Determine activation string -- act_str = None -- if isinstance(activation, bool): -- act_str = "silu" if activation else None -- elif activation is not None: -- act_str = activation -- -- _conv_state_indices = conv_state_indices -- if _conv_state_indices is not None and _conv_state_indices.dim() == 2: -- if initial_state_idx is not None: -- _conv_state_indices = _conv_state_indices.gather( -- 1, initial_state_idx.unsqueeze(1) -- ).squeeze(1) -- else: -- _conv_state_indices = _conv_state_indices[:, 0] -- -- pad_slot_id = int(NULL_BLOCK_ID) -- return ops.causal_conv1d_update_cpu_vec( -- x, -- conv_state, -- weight, -- bias, -- act_str, -- _conv_state_indices, -- None, # query_start_loc -- pad_slot_id, -- ) - - return _causal_conv1d_update_cuda( - x, -@@ -1323,3 +1277,12 @@ def causal_conv1d_update( - initial_state_idx=initial_state_idx, - validate_data=validate_data, - ) -+ -+ -+from vllm.platforms import current_platform -+ -+if current_platform.is_cpu(): -+ from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( -+ causal_conv1d_fn_cpu as causal_conv1d_fn, -+ causal_conv1d_update_cpu as causal_conv1d_update, -+ ) -diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py -index b047ca6d6..dba752624 100644 ---- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py -+++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py -@@ -4,10 +4,182 @@ - from __future__ import annotations - - import torch --import torch.nn.functional as F - -+from vllm.v1.attention.backends.utils import PAD_SLOT_ID -+ -+ -+def causal_conv1d_fn_cpu( -+ x: torch.Tensor, -+ weight: torch.Tensor, -+ bias: torch.Tensor | None, -+ conv_states: torch.Tensor, -+ query_start_loc: torch.Tensor, -+ cache_indices: torch.Tensor | None = None, -+ has_initial_state: torch.Tensor | None = None, -+ activation: str | None = "silu", -+ pad_slot_id: int = PAD_SLOT_ID, -+ **kwargs, -+) -> torch.Tensor: -+ """CPU implementation for causal_conv1d_fwd.""" -+ if isinstance(activation, bool) and activation: -+ activation = "silu" -+ -+ 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 -+ -+ if seq_len == 0: -+ continue -+ -+ cache_idx = cache_indices[b].item() if cache_indices is not None else b -+ -+ if cache_idx == 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) -+ 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 -+ -+ if state_len > 1: -+ state[:, :-1] = state[:, 1:].clone() -+ state[:, -1] = x_t -+ -+ conv_states[cache_idx].copy_(state) -+ -+ return out.to(original_x_dtype) -+ -+ -+def causal_conv1d_update_cpu( -+ x: torch.Tensor, -+ conv_state: torch.Tensor, -+ weight: torch.Tensor, -+ bias: torch.Tensor | None = None, -+ activation: bool | str | None = None, -+ conv_state_indices: torch.Tensor | None = None, -+ query_start_loc: torch.Tensor | None = None, -+ pad_slot_id: int = PAD_SLOT_ID, -+ **kwargs, -+) -> torch.Tensor: -+ """CPU implementation for causal_conv1d_update.""" -+ if isinstance(activation, bool): -+ activation = "silu" if activation else None -+ -+ original_x_dtype = x.dtype -+ x = x.to(conv_state.dtype) -+ _, width = weight.shape -+ state_len = width - 1 -+ -+ if query_start_loc is None and x.dim() == 2: -+ x = x.unsqueeze(-1) -+ unsqueeze = True -+ else: -+ unsqueeze = False -+ -+ if query_start_loc is None: -+ batch, dim, seqlen = x.shape -+ -+ if conv_state_indices is not None: -+ cache_idxs = conv_state_indices.flatten() -+ valid_mask = cache_idxs != pad_slot_id -+ else: -+ cache_idxs = torch.arange(batch, device=x.device) -+ valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) -+ -+ for t in range(seqlen): -+ x_t = x[:, :, t].clone() -+ -+ states = conv_state[cache_idxs] -+ -+ windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) -+ -+ val = (windows * weight.unsqueeze(0)).sum(dim=-1) -+ if bias is not None: -+ val = val + bias.unsqueeze(0) -+ -+ if activation in ["silu", "swish"]: -+ val = val * torch.sigmoid(val) -+ -+ val = val * valid_mask.unsqueeze(-1).to(val.dtype) -+ x[:, :, t] = val -+ -+ new_state = torch.cat([states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1) -+ conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] -+ -+ out = x -+ if unsqueeze: -+ out = out.squeeze(-1) -+ return out.to(original_x_dtype) -+ -+ assert conv_state_indices is not None -+ assert query_start_loc is not None -+ batch = conv_state_indices.size(0) -+ out = x.clone() -+ -+ for b in range(batch): -+ cache_idx = conv_state_indices[b].item() -+ if cache_idx == pad_slot_id: -+ continue -+ -+ seq_start = query_start_loc[b].item() -+ seq_end = query_start_loc[b + 1].item() -+ seqlen_b = seq_end - seq_start -+ -+ if seqlen_b == 0: -+ continue -+ -+ local_state = conv_state[cache_idx].clone() -+ -+ for t in range(seqlen_b): -+ x_t = x[seq_start + t, :] -+ -+ window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) -+ val = (window * weight).sum(dim=-1) -+ if bias is not None: -+ val = val + bias -+ if activation in ["silu", "swish"]: -+ val = val * torch.sigmoid(val) -+ -+ out[seq_start + t, :] = val -+ -+ if state_len > 1: -+ local_state[:, :-1] = local_state[:, 1:].clone() -+ local_state[:, -1] = x_t -+ -+ conv_state[cache_idx] = local_state -+ -+ return out.to(original_x_dtype) -+ -+ -+import torch.nn.functional as F - --# for prefill - def causal_conv1d_torch( - x: torch.Tensor, - weight: torch.Tensor, -@@ -60,7 +232,6 @@ def causal_conv1d_torch( - return out - - --# for decode - def causal_conv1d_update_torch( - x: torch.Tensor, - conv_state: torch.Tensor, -diff --git a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py b/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py -deleted file mode 100644 -index 50a3bb392..000000000 ---- a/vllm/model_executor/layers/mamba/ops/cpu_fallbacks.py -+++ /dev/null -@@ -1,306 +0,0 @@ --# SPDX-License-Identifier: Apache-2.0 --# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -- --import torch -- --import vllm._custom_ops as ops --from vllm.v1.attention.backends.utils import PAD_SLOT_ID -- -- --def _causal_conv1d_fn_cpu( -- x: torch.Tensor, -- weight: torch.Tensor, -- bias: torch.Tensor | None, -- conv_states: torch.Tensor, -- query_start_loc: torch.Tensor, -- cache_indices: torch.Tensor | None = None, -- has_initial_state: torch.Tensor | None = None, -- activation: str | None = "silu", -- pad_slot_id: int = PAD_SLOT_ID, -- **kwargs, --) -> torch.Tensor: -- """Pure PyTorch CPU fallback for causal_conv1d_fwd.""" -- if isinstance(activation, bool) and activation: -- activation = "silu" -- -- 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 -- -- if seq_len == 0: -- continue -- -- cache_idx = cache_indices[b].item() if cache_indices is not None else b -- -- if cache_idx == 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) -- 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 -- -- if state_len > 1: -- state[:, :-1] = state[:, 1:].clone() -- state[:, -1] = x_t -- -- conv_states[cache_idx].copy_(state) -- -- return out.to(original_x_dtype) -- -- --def _causal_conv1d_update_cpu( -- x: torch.Tensor, -- conv_state: torch.Tensor, -- weight: torch.Tensor, -- bias: torch.Tensor | None = None, -- activation: bool | str | None = None, -- conv_state_indices: torch.Tensor | None = None, -- query_start_loc: torch.Tensor | None = None, -- pad_slot_id: int = PAD_SLOT_ID, -- **kwargs, --) -> torch.Tensor: -- """Pure PyTorch CPU fallback for causal_conv1d_update (decode path).""" -- if isinstance(activation, bool): -- activation = "silu" if activation else None -- -- original_x_dtype = x.dtype -- x = x.to(conv_state.dtype) -- _, width = weight.shape -- state_len = width - 1 -- -- if query_start_loc is None and x.dim() == 2: -- x = x.unsqueeze(-1) -- unsqueeze = True -- else: -- unsqueeze = False -- -- if query_start_loc is None: -- batch, dim, seqlen = x.shape -- -- if conv_state_indices is not None: -- cache_idxs = conv_state_indices.flatten() -- valid_mask = cache_idxs != pad_slot_id -- else: -- cache_idxs = torch.arange(batch, device=x.device) -- valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) -- -- for t in range(seqlen): -- x_t = x[:, :, t].clone() -- -- states = conv_state[cache_idxs] -- -- windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) -- -- val = (windows * weight.unsqueeze(0)).sum(dim=-1) -- if bias is not None: -- val = val + bias.unsqueeze(0) -- -- if activation in ["silu", "swish"]: -- val = val * torch.sigmoid(val) -- -- val = val * valid_mask.unsqueeze(-1).to(val.dtype) -- x[:, :, t] = val -- -- new_state = torch.cat([states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1) -- conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] -- -- out = x -- if unsqueeze: -- out = out.squeeze(-1) -- return out.to(original_x_dtype) -- -- assert conv_state_indices is not None -- assert query_start_loc is not None -- batch = conv_state_indices.size(0) -- out = x.clone() -- -- for b in range(batch): -- cache_idx = conv_state_indices[b].item() -- if cache_idx == pad_slot_id: -- continue -- -- seq_start = query_start_loc[b].item() -- seq_end = query_start_loc[b + 1].item() -- seqlen_b = seq_end - seq_start -- -- if seqlen_b == 0: -- continue -- -- local_state = conv_state[cache_idx].clone() -- -- for t in range(seqlen_b): -- x_t = x[seq_start + t, :] -- -- window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) -- val = (window * weight).sum(dim=-1) -- if bias is not None: -- val = val + bias -- if activation in ["silu", "swish"]: -- val = val * torch.sigmoid(val) -- -- out[seq_start + t, :] = val -- -- if state_len > 1: -- local_state[:, :-1] = local_state[:, 1:].clone() -- local_state[:, -1] = x_t -- -- conv_state[cache_idx] = local_state -- -- return out.to(original_x_dtype) -- -- --def _mamba_chunk_scan_combined_fwd_cpu( -- x, -- dt, -- A, -- B, -- C, -- chunk_size, -- out, -- D=None, -- z=None, -- dt_bias=None, -- initial_states=None, -- return_intermediate_states=False, -- seq_idx=None, -- cu_seqlens=None, -- cu_chunk_seqlens=None, -- last_chunk_indices=None, -- dt_softplus=False, -- dt_limit=(0.0, float("inf")), -- state_dtype=None, -- **kwargs, --): -- seqlen, nheads, headdim = x.shape -- _, ngroups, dstate = B.shape -- -- assert cu_seqlens is not None -- batch = cu_seqlens.size(0) - 1 -- -- # Preprocess dt: apply bias, softplus, and clamp. -- # Kept in Python to avoid duplicating this logic in C++. -- dt_f = dt.float() -- if dt_bias is not None: -- dt_f = dt_f + dt_bias.float().unsqueeze(0) -- if dt_softplus: -- dt_f = torch.nn.functional.softplus(dt_f) -- if dt_limit[0] > 0.0 or dt_limit[1] < float("inf"): -- dt_f = dt_f.clamp(min=dt_limit[0], max=dt_limit[1]) -- -- # Allocate state output buffer (float32, contiguous — required by kernel). -- all_states = torch.zeros( -- batch, nheads, headdim, dstate, dtype=torch.float32, device=x.device -- ) -- if initial_states is not None: -- all_states.copy_(initial_states.float()) -- -- # Use the C++ kernel when available (CPU build with mamba kernels compiled). -- # Falls back to the pure-Python loop below if the op is not registered. -- _use_cpp_kernel = hasattr(torch.ops._C, "mamba_chunk_scan_fwd_cpu") -- -- if _use_cpp_kernel: -- # out must be contiguous — the C++ kernel writes via raw data_ptr(). -- # Creating a contiguous copy here would discard the results, so we -- # require the caller to pass a contiguous tensor. -- assert out.is_contiguous(), ( -- "_mamba_chunk_scan_combined_fwd_cpu: `out` must be " -- "pre-allocated as a contiguous tensor" -- ) -- -- # D: strip trailing dims that are broadcast (stride == 0) so the -- # kernel sees a (nheads,) float32 array. Only peel dims whose -- # stride is 0 (broadcast-expanded); genuine multi-dim D is passed -- # as-is and flattened by the C++ wrapper. -- D_1d = None -- if D is not None: -- d = D.float() -- while d.dim() > 1 and d.stride(-1) == 0: -- d = d.squeeze(-1) -- D_1d = d.contiguous() -- -- ops.mamba_chunk_scan_fwd_cpu( -- out, -- all_states, -- x, -- dt_f, -- A, -- B, -- C, -- D_1d, -- z, -- cu_seqlens.to(torch.int32), -- ) -- else: -- # Pure-Python fallback (no compiled extension available). -- for b_idx in range(batch): -- seq_start = cu_seqlens[b_idx].item() -- seq_end = cu_seqlens[b_idx + 1].item() -- -- # Note: basic indexing returns a view, but arithmetic ops -- # below (state * dA + dBx) produce new tensors, so state is -- # reassigned each token. all_states[b_idx] is committed at line 416. -- state = all_states[b_idx].clone() # local working copy -- -- for t in range(seq_start, seq_end): -- x_t = x[t].float() -- dt_t = dt_f[t] -- A_val = A.float() -- -- dA = torch.exp(A_val * dt_t).unsqueeze(-1).unsqueeze(-1) -- -- B_expanded = B[t].float().repeat_interleave(nheads // ngroups, dim=0) -- C_expanded = C[t].float().repeat_interleave(nheads // ngroups, dim=0) -- -- xdt = x_t * dt_t.unsqueeze(-1) -- dBx = xdt.unsqueeze(-1) * B_expanded.unsqueeze(1) -- state = state * dA + dBx -- -- y = (state * C_expanded.unsqueeze(1)).sum(dim=-1) -- -- if D is not None: -- y = ( -- y + x_t * D.float().unsqueeze(-1) -- if D.dim() == 1 -- else y + x_t * D.float() -- ) -- -- if z is not None: -- z_t = z[t].float() -- y = y * z_t * torch.sigmoid(z_t) -- -- out[t] = y.to(out.dtype) -- -- all_states[b_idx] = state.to(all_states.dtype) -- -- out_dtype = state_dtype if state_dtype is not None else x.dtype -- all_states = all_states.to(out_dtype) -- -- return all_states -diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py -index d7055fbd7..9d463e570 100644 ---- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py -+++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py -@@ -894,47 +894,6 @@ def selective_state_update( - if out is None: - out = torch.empty_like(x if x.dim() == 2 else x) - -- if x.device.type == "cpu": -- # Reshape tensors from (batch, dim) -> (batch, 1, dim) if needed -- # The C++ kernel expects (N, nheads, dim) layout -- _state = state.unsqueeze(1) if state.dim() == 3 else state -- _x = x.unsqueeze(1) if x.dim() == 2 else x -- _dt = dt.unsqueeze(1) if dt.dim() == 2 else dt -- _A = A.unsqueeze(0) if A.dim() == 2 else A -- _B = B.unsqueeze(1) if B.dim() == 2 else B -- _C = C.unsqueeze(1) if C.dim() == 2 else C -- _D = D.unsqueeze(0) if (D is not None and D.dim() == 1) else D -- _z = z.unsqueeze(1) if (z is not None and z.dim() == 2) else z -- _dt_bias = ( -- dt_bias.unsqueeze(0) -- if (dt_bias is not None and dt_bias.dim() == 1) -- else dt_bias -- ) -- _out = out.unsqueeze(1) if out.dim() == 2 else out -- # state_batch_indices and dst_state_batch_indices are 1D index arrays; -- # do NOT reshape them. -- _sbi = state_batch_indices -- _dsbi = dst_state_batch_indices -- ops.selective_state_update_cpu( -- _state, -- _x, -- _dt, -- _A, -- _B, -- _C, -- _D, -- _z, -- _dt_bias, -- dt_softplus, -- _sbi, -- _dsbi, -- null_block_id, -- _out, -- num_accepted_tokens, -- cu_seqlens, -- ) -- return _out.squeeze(1) if out.dim() == 2 else _out -- - return _selective_state_update_cuda( - state, - x, -@@ -956,3 +915,9 @@ def selective_state_update( - enable_stochastic_rounding=enable_stochastic_rounding, - cache_philox_rounds=cache_philox_rounds, - ) -+ -+ -+if current_platform.is_cpu(): -+ from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( -+ selective_state_update, -+ ) -diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py -index caa57211e..89c8bd814 100644 ---- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py -+++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py -@@ -13,7 +13,7 @@ from packaging import version - from vllm.model_executor.custom_op import CustomOp - from vllm.triton_utils import HAS_TRITON, triton - --from .cpu_fallbacks import _mamba_chunk_scan_combined_fwd_cpu -+ - from .ssd_bmm import _bmm_chunk_fwd - from .ssd_chunk_scan import _chunk_scan_fwd - from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd -@@ -166,19 +166,11 @@ class MambaChunkScanCombinedFwdOp(CustomOp): - super(CustomOp, self).__init__() # nn.Module.__init__ only - from vllm.platforms import current_platform - -- if current_platform.is_cpu(): -- self._forward_method = self.forward_cpu -- elif current_platform.is_rocm(): -+ if current_platform.is_rocm(): - self._forward_method = self.forward_hip - else: - self._forward_method = self.forward_cuda - -- def forward_native(self, *args, **kwargs): -- return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) -- -- def forward_cpu(self, *args, **kwargs): -- return _mamba_chunk_scan_combined_fwd_cpu(*args, **kwargs) -- - def forward_cuda(self, *args, **kwargs): - return _mamba_chunk_scan_combined_fwd_cuda(*args, **kwargs) - -@@ -260,7 +252,15 @@ def mamba_chunk_scan_combined_varlen( - last_chunk_indices=last_chunk_indices, - dt_softplus=dt_softplus, - dt_limit=dt_limit, -+ return_intermediate_states=return_intermediate_states, - state_dtype=state_dtype, - ) -- - return varlen_states -+ -+ -+from vllm.platforms import current_platform -+ -+if current_platform.is_cpu(): -+ from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( -+ _mamba_chunk_scan_combined_fwd_cpu as _mamba_chunk_scan_combined_fwd, -+ ) -diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py -index 1ceeac665..1795a036b 100644 ---- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py -+++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py -@@ -6,8 +6,7 @@ Dispatch module for Mamba selective state update (SSU) backends. - Provides a unified `selective_state_update` function that dispatches to - the Triton, FlashInfer, or CPU backend based on the configured - `MambaBackendEnum`. On CPU-only platforms (PowerPC, x86 without CUDA) --the backend defaults to 'cpu', which uses a pure-PyTorch fallback that --avoids Triton JIT compilation entirely. -+the backend defaults to 'cpu'. - """ - - from abc import ABC, abstractmethod -@@ -280,7 +279,7 @@ def initialize_mamba_ssu_backend( - backend = mamba_config.backend - - # On CPU-only platforms (PowerPC, x86 without CUDA) Triton JIT is -- # unstable or unavailable. Silently fall back to the pure-PyTorch CPU -+ # unstable or unavailable. Silently fall back to the CPU - # backend unless the user explicitly chose something other than "triton". - if backend == MambaBackendEnum.TRITON: - from vllm.platforms import current_platform -@@ -288,7 +287,7 @@ def initialize_mamba_ssu_backend( - if current_platform.is_cpu(): - logger.info( - "CPU platform detected: overriding Mamba SSU backend " -- "from 'triton' to 'cpu' (pure-PyTorch fallback)." -+ "from 'triton' to 'cpu'." - ) - backend = MambaBackendEnum.CPU - -diff --git a/vllm/model_executor/layers/utils.py b/vllm/model_executor/layers/utils.py -index 450c74a60..f0d3e6ead 100644 ---- a/vllm/model_executor/layers/utils.py -+++ b/vllm/model_executor/layers/utils.py -@@ -243,23 +243,19 @@ def dispatch_cpu_unquantized_gemm( - if torch.cpu._is_amx_tile_supported() and hasattr( - ops, "causal_conv1d_weight_pack" - ): -- with contextlib.suppress(Exception): -- # prepack conv weight -- unpacked = ( -- layer.weight.view( -- layer.weight.size(0), -- layer.weight.size(2), -- ) -- .contiguous() -- .clone() -+ # prepack conv weight -+ unpacked = ( -+ layer.weight.view( -+ layer.weight.size(0), -+ layer.weight.size(2), - ) -- # Stash the un-packed (dim, width) weight so the speculative-decode -- # GDN path (which uses torch conv, not the AMX kernel) can use it. -- layer._cpu_unpacked_conv_weight = unpacked -- layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) -- layer.cpu_linear = lambda x, weight, bias: torch.nn.functional.linear( -- x, weight, bias -- ) -+ .contiguous() -+ .clone() -+ ) -+ # Stash the un-packed (dim, width) weight so the speculative-decode -+ # GDN path (which uses torch conv, not the AMX kernel) can use it. -+ layer._cpu_unpacked_conv_weight = unpacked -+ layer.weight.data = ops.causal_conv1d_weight_pack(unpacked) - return - - N, K = layer.weight.size() From 1928a66b23837b23bd9d792e9af1b669a64fd3b5 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 15 Jul 2026 00:33:44 -0500 Subject: [PATCH 36/43] combine causal_conv1d_fn_cpu Signed-off-by: Akash kaothalkar --- .../layers/mamba/ops/cpu/causal_conv1d.py | 130 ++++++------------ 1 file changed, 42 insertions(+), 88 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py index 3715cfd6a2d9..924a5a8e661a 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +++ b/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( From 8e1bab35ded70c544b8891ed7f1ac164269ecd2f Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 15 Jul 2026 00:51:58 -0500 Subject: [PATCH 37/43] update causal_conv1d to use vec function Signed-off-by: Akash kaothalkar --- .../layers/mamba/ops/cpu/causal_conv1d.py | 123 ++---------------- 1 file changed, 10 insertions(+), 113 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py index 924a5a8e661a..45b5aae9f038 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py @@ -6,6 +6,7 @@ import torch import torch.nn.functional as F +from vllm._custom_ops import causal_conv1d_update_cpu_vec from vllm.v1.attention.backends.utils import PAD_SLOT_ID @@ -96,117 +97,13 @@ def causal_conv1d_update_cpu( if isinstance(activation, bool): activation = "silu" if activation else None - original_x_dtype = x.dtype - x = x.to(conv_state.dtype) - _, width = weight.shape - state_len = width - 1 - - if query_start_loc is None and x.dim() == 2: - x = x.unsqueeze(-1) - unsqueeze = True - else: - unsqueeze = False - - if query_start_loc is None: - batch, dim, seqlen = x.shape - - if conv_state_indices is not None: - cache_idxs = conv_state_indices.flatten() - valid_mask = cache_idxs != pad_slot_id - else: - cache_idxs = torch.arange(batch, device=x.device) - valid_mask = torch.ones(batch, dtype=torch.bool, device=x.device) - - for t in range(seqlen): - x_t = x[:, :, t].clone() - - states = conv_state[cache_idxs] - - windows = torch.cat([states, x_t.unsqueeze(-1)], dim=-1) - - val = (windows * weight.unsqueeze(0)).sum(dim=-1) - if bias is not None: - val = val + bias.unsqueeze(0) - - if activation in ["silu", "swish"]: - val = val * torch.sigmoid(val) - - val = val * valid_mask.unsqueeze(-1).to(val.dtype) - x[:, :, t] = val - - new_state = torch.cat([states[:, :, 1:], x_t.unsqueeze(-1)], dim=-1) - conv_state[cache_idxs[valid_mask]] = new_state[valid_mask] - - out = x - if unsqueeze: - out = out.squeeze(-1) - return out.to(original_x_dtype) - - assert conv_state_indices is not None - assert query_start_loc is not None - batch = conv_state_indices.size(0) - out = x.clone() - - for b in range(batch): - cache_idx = conv_state_indices[b].item() - if cache_idx == pad_slot_id: - continue - - seq_start = query_start_loc[b].item() - seq_end = query_start_loc[b + 1].item() - seqlen_b = seq_end - seq_start - - if seqlen_b == 0: - continue - - local_state = conv_state[cache_idx].clone() - - for t in range(seqlen_b): - x_t = x[seq_start + t, :] - - window = torch.cat([local_state, x_t.unsqueeze(-1)], dim=-1) - val = (window * weight).sum(dim=-1) - if bias is not None: - val = val + bias - if activation in ["silu", "swish"]: - val = val * torch.sigmoid(val) - - out[seq_start + t, :] = val - - if state_len > 1: - local_state[:, :-1] = local_state[:, 1:].clone() - local_state[:, -1] = x_t - - conv_state[cache_idx] = local_state - - return out.to(original_x_dtype) - - - - - -def causal_conv1d_update_torch( - x: torch.Tensor, - conv_state: torch.Tensor, - weight: torch.Tensor, - bias: torch.Tensor | None = None, - activation: str | None = None, -) -> torch.Tensor: - assert activation in {None, "silu", "swish"} - - _, dim, seq_len = x.shape - state_len = conv_state.shape[-1] - - x_new = torch.cat([conv_state, x], dim=-1).to(weight.dtype) - conv_state.copy_(x_new[:, :, -state_len:]) - - out = F.conv1d( - x_new, - weight.unsqueeze(1), + return causal_conv1d_update_cpu_vec( + x, + conv_state, + weight, bias, - padding=0, - groups=dim, - )[:, :, -seq_len:] - if activation in ("silu", "swish"): - out = F.silu(out) - return out + activation, + conv_state_indices, + query_start_loc, + pad_slot_id, + ) From 2380372c5169e945e649880a1867bba6539780c0 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 15 Jul 2026 01:02:05 -0500 Subject: [PATCH 38/43] pre-commit fix Signed-off-by: Akash kaothalkar --- .../layers/mamba/ops/cpu/causal_conv1d.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py index 45b5aae9f038..4c138585ec41 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py @@ -40,18 +40,20 @@ def causal_conv1d_fn_cpu( 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 - - slot = int(cache_indices[seq_idx].item()) if cache_indices is not None else seq_idx - + + slot = ( + int(cache_indices[seq_idx].item()) if cache_indices is not None else seq_idx + ) + if slot == pad_slot_id: continue 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: @@ -72,7 +74,7 @@ def causal_conv1d_fn_cpu( 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) From 1c1f4b49749714f3733b4d42bada904fe6431e75 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 15 Jul 2026 01:07:53 -0500 Subject: [PATCH 39/43] remove torch backend from short_conv Signed-off-by: Akash kaothalkar --- .../layers/mamba/ops/cpu/gdn_attention.py | 30 ++++++++----------- .../model_executor/layers/mamba/short_conv.py | 17 +++++------ 2 files changed, 21 insertions(+), 26 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index d99c2490ee2e..4f3d96643d8c 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -10,8 +10,8 @@ from vllm.forward_context import ForwardContext, get_forward_context from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( - causal_conv1d_torch, - causal_conv1d_update_torch, + causal_conv1d_fn_cpu as causal_conv1d_torch, + causal_conv1d_update_cpu, ) from vllm.utils.torch_utils import ( LayerNameType, @@ -145,16 +145,14 @@ def _cpu_gdn_attention_nonspec( is_vnni=True, ) else: - decode_conv_state = conv_state[decode_state_indices].contiguous() - decode_mixed_qkv = causal_conv1d_update_torch( - # [B, dim] -> [B, dim, 1] - x=decode_mixed_qkv.unsqueeze(-1), - conv_state=decode_conv_state, + decode_mixed_qkv = causal_conv1d_update_cpu( + x=decode_mixed_qkv, + conv_state=conv_state, weight=conv_weights, bias=layer.conv1d.bias, activation=layer.activation, - ).squeeze(-1) - conv_state[decode_state_indices] = decode_conv_state + conv_state_indices=decode_state_indices, + ) query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) @@ -495,17 +493,14 @@ def _spec_aware_nonspec( decode_a = a[:num_decode_tokens] decode_state_indices = state_indices_tensor[:num_decodes] # Only the first ``width-1`` columns hold the real conv state. - decode_conv_state = conv_buf[decode_state_indices][ - :, :, : width - 1 - ].contiguous() - decode_mixed_qkv = causal_conv1d_update_torch( - x=decode_mixed_qkv.unsqueeze(-1), - conv_state=decode_conv_state, + decode_mixed_qkv = causal_conv1d_update_cpu( + x=decode_mixed_qkv, + conv_state=conv_buf[:, :, : width - 1], weight=conv_weights, bias=layer.conv1d.bias, activation=layer.activation, - ).squeeze(-1) - conv_buf[decode_state_indices, :, : width - 1] = decode_conv_state + conv_state_indices=decode_state_indices, + ) query, key, value = layer.rearrange_mixed_qkv(decode_mixed_qkv) # rearrange_mixed_qkv can return views whose last dim is not @@ -663,3 +658,4 @@ def register_cpu_gdn_attention_ops() -> None: fake_impl=cpu_gdn_attention_core_fake, ) _CPU_GDN_ATTENTION_OPS_REGISTERED = True + diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index e7e36f2fc538..66830c9586c2 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -94,8 +94,8 @@ def forward_native( # Reference torch causal conv1d; runs on all CPU platforms. AMX kernels # for causal conv can be plugged in here later. from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( - causal_conv1d_torch, - causal_conv1d_update_torch, + causal_conv1d_fn_cpu as causal_conv1d_torch, + causal_conv1d_update_cpu, ) forward_context = get_forward_context() @@ -164,17 +164,15 @@ def forward_native( if has_decode: assert attn_metadata.state_indices_tensor_d is not None state_indices_d = attn_metadata.state_indices_tensor_d.flatten() - Bx_d = (B_d * x_d).unsqueeze(-1) # (num_decodes, dim, 1) - # Advanced indexing returns a copy; update in-place then scatter back - gathered = conv_state[state_indices_d] # (num_decodes, dim, state_len) - out_d = causal_conv1d_update_torch( + Bx_d = (B_d * x_d) # (num_decodes, dim) + out_d = causal_conv1d_update_cpu( Bx_d, - gathered, + conv_state, conv_weights, self.conv.bias, activation=None, - ).squeeze(-1) # (num_decodes, dim) - conv_state[state_indices_d] = gathered + conv_state_indices=state_indices_d, + ) conv_output_list.insert(0, C_d * out_d) hidden_states_out = torch.vstack(conv_output_list) @@ -343,3 +341,4 @@ def short_conv_fake( mutates_args=["output"], fake_impl=short_conv_fake, ) + From c8b763131998ec4a0b56ffe8e893546946db05ab Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 15 Jul 2026 01:19:56 -0500 Subject: [PATCH 40/43] pre-commit fix Signed-off-by: Akash kaothalkar --- vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py | 3 ++- vllm/model_executor/layers/mamba/short_conv.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py index 4f3d96643d8c..2749544b4957 100644 --- a/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py +++ b/vllm/model_executor/layers/mamba/ops/cpu/gdn_attention.py @@ -11,6 +11,8 @@ from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_fn_cpu as causal_conv1d_torch, +) +from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_update_cpu, ) from vllm.utils.torch_utils import ( @@ -658,4 +660,3 @@ def register_cpu_gdn_attention_ops() -> None: fake_impl=cpu_gdn_attention_core_fake, ) _CPU_GDN_ATTENTION_OPS_REGISTERED = True - diff --git a/vllm/model_executor/layers/mamba/short_conv.py b/vllm/model_executor/layers/mamba/short_conv.py index 66830c9586c2..9be9ada36378 100644 --- a/vllm/model_executor/layers/mamba/short_conv.py +++ b/vllm/model_executor/layers/mamba/short_conv.py @@ -95,6 +95,8 @@ def forward_native( # for causal conv can be plugged in here later. from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_fn_cpu as causal_conv1d_torch, + ) + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( causal_conv1d_update_cpu, ) @@ -164,7 +166,7 @@ def forward_native( if has_decode: assert attn_metadata.state_indices_tensor_d is not None state_indices_d = attn_metadata.state_indices_tensor_d.flatten() - Bx_d = (B_d * x_d) # (num_decodes, dim) + Bx_d = B_d * x_d # (num_decodes, dim) out_d = causal_conv1d_update_cpu( Bx_d, conv_state, @@ -341,4 +343,3 @@ def short_conv_fake( mutates_args=["output"], fake_impl=short_conv_fake, ) - From db72a7cfce1ef79bcffc6a90f586f7036922d99a Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 15 Jul 2026 02:23:44 -0500 Subject: [PATCH 41/43] revert gpu wrappers Signed-off-by: Akash kaothalkar --- .../layers/mamba/ops/causal_conv1d.py | 54 ++---------- .../layers/mamba/ops/mamba_ssm.py | 83 ++----------------- .../layers/mamba/ops/ssd_combined.py | 45 ++-------- 3 files changed, 24 insertions(+), 158 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 7ba730e92400..84a3653b68c0 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -465,7 +465,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching tl.store(o_ptrs, acc, mask=mask_1d) -def _causal_conv1d_fn_cuda( +def causal_conv1d_fn( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, @@ -1066,7 +1066,7 @@ def _causal_conv1d_update_kernel( tl.store(o_ptrs, acc, mask=mask_1d) -def _causal_conv1d_update_cuda( +def causal_conv1d_update( x: torch.Tensor, conv_state: torch.Tensor, weight: torch.Tensor, @@ -1238,49 +1238,13 @@ def grid(META): out = out.squeeze(-1) return out.to(original_x_dtype) - -def causal_conv1d_fn(*args, **kwargs): - return _causal_conv1d_fn_cuda(*args, **kwargs) - - -def causal_conv1d_update( - x, - conv_state, - weight, - bias=None, - activation=None, - conv_state_indices=None, - num_accepted_tokens=None, - query_start_loc=None, - max_query_len=-1, - null_block_id=NULL_BLOCK_ID, - block_idx_last_scheduled_token=None, - initial_state_idx=None, - validate_data=False, -): - """Dispatch causal_conv1d_update to CPU C++ kernel or CUDA Triton kernel.""" - - return _causal_conv1d_update_cuda( - x, - conv_state, - weight, - bias=bias, - activation=activation, - conv_state_indices=conv_state_indices, - num_accepted_tokens=num_accepted_tokens, - query_start_loc=query_start_loc, - max_query_len=max_query_len, - null_block_id=null_block_id, - block_idx_last_scheduled_token=block_idx_last_scheduled_token, - initial_state_idx=initial_state_idx, - validate_data=validate_data, - ) - - -from vllm.platforms import current_platform # noqa: E402 +from vllm.platforms import current_platform if current_platform.is_cpu(): - import vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d as cpu_conv1d # noqa: E402 + from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( + causal_conv1d_fn_cpu, + causal_conv1d_update_cpu, + ) - causal_conv1d_fn = cpu_conv1d.causal_conv1d_fn_cpu # type: ignore - causal_conv1d_update = cpu_conv1d.causal_conv1d_update_cpu # type: ignore + causal_conv1d_fn = causal_conv1d_fn_cpu # type: ignore + causal_conv1d_update = causal_conv1d_update_cpu # type: ignore diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index a2b279d23831..9bea5b0e21d7 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -226,9 +226,8 @@ def convert_rs_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: @triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) @triton.heuristics( { - "HAS_STATE_BATCH_INDICES": lambda args: ( - args["state_batch_indices_ptr"] is not None - ) + "HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] + is not None } ) @triton.heuristics( @@ -495,7 +494,7 @@ def _selective_scan_update_kernel( tl.store(dst_state_ptrs, state, mask=mask) -def _selective_state_update_cuda( +def selective_state_update( state, x, dt, @@ -606,26 +605,6 @@ def _selective_state_update_cuda( if num_accepted_tokens is not None: assert num_accepted_tokens.shape == (N,) - if not HAS_TRITON: - return ops.selective_state_update_cpu( - state, - x, - dt, - A, - B, - C, - D, - z, - dt_bias, - dt_softplus, - state_batch_indices, - dst_state_batch_indices, - null_block_id, - out, - num_accepted_tokens, - cu_seqlens, - ) - grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), N, nheads) z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) state_batch_indices_strides = ( @@ -867,57 +846,11 @@ def selective_scan_fn( else: return z # output written inplace to z - -def selective_state_update( - state, - x, - dt, - A, - B, - C, - D=None, - z=None, - dt_bias=None, - dt_softplus=False, - state_batch_indices=None, - dst_state_batch_indices=None, - null_block_id=NULL_BLOCK_ID, - out=None, - num_accepted_tokens=None, - cu_seqlens=None, - is_blackwell=False, - enable_stochastic_rounding=False, - cache_philox_rounds=0, -): - """Dispatch selective_state_update to CPU C++ kernel or CUDA Triton kernel.""" - # Ensure out tensor exists - if out is None: - out = torch.empty_like(x if x.dim() == 2 else x) - - return _selective_state_update_cuda( - state, - x, - dt, - A, - B, - C, - D=D, - z=z, - dt_bias=dt_bias, - dt_softplus=dt_softplus, - state_batch_indices=state_batch_indices, - dst_state_batch_indices=dst_state_batch_indices, - null_block_id=null_block_id, - out=out, - num_accepted_tokens=num_accepted_tokens, - cu_seqlens=cu_seqlens, - is_blackwell=is_blackwell, - enable_stochastic_rounding=enable_stochastic_rounding, - cache_philox_rounds=cache_philox_rounds, - ) - +from vllm.platforms import current_platform if current_platform.is_cpu(): - import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_ssm # noqa: E402 + from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( + selective_state_update as selective_state_update_cpu, + ) - selective_state_update = cpu_ssm.selective_state_update + selective_state_update = selective_state_update_cpu # type: ignore diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 626792afd0f3..3d8e354cd80f 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -10,22 +10,21 @@ from einops import rearrange from packaging import version -from vllm.model_executor.custom_op import CustomOp -from vllm.triton_utils import HAS_TRITON, triton +from vllm.triton_utils import triton from .ssd_bmm import _bmm_chunk_fwd from .ssd_chunk_scan import _chunk_scan_fwd from .ssd_chunk_state import _chunk_cumsum_fwd, _chunk_state_fwd from .ssd_state_passing import _state_passing_fwd -TRITON_22 = HAS_TRITON and version.parse(triton.__version__) >= version.parse("2.2.0") +TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0") def is_int_pow_2(n): return isinstance(n, int) and n > 0 and (n & (n - 1)) == 0 -def _mamba_chunk_scan_combined_fwd_cuda( +def _mamba_chunk_scan_combined_fwd( x, dt, A, @@ -155,35 +154,6 @@ def _mamba_chunk_scan_combined_fwd_cuda( return states[last_chunk_indices] -@CustomOp.register("mamba_chunk_scan_combined_fwd") -class MambaChunkScanCombinedFwdOp(CustomOp): - def __init__(self): - # Bypass dispatch_forward() / get_current_vllm_config() by directly - # selecting the forward method based on the platform. This op may be - # instantiated during Mamba2 kernel warm-up, which happens before the - # vLLM compilation config is available. - super(CustomOp, self).__init__() # nn.Module.__init__ only - from vllm.platforms import current_platform - - if current_platform.is_rocm(): - self._forward_method = self.forward_hip - else: - self._forward_method = self.forward_cuda - - def forward_cuda(self, *args, **kwargs): - return _mamba_chunk_scan_combined_fwd_cuda(*args, **kwargs) - - -_mamba_chunk_scan_combined_fwd_op: "MambaChunkScanCombinedFwdOp | None" = None - - -def _mamba_chunk_scan_combined_fwd(*args, **kwargs): - global _mamba_chunk_scan_combined_fwd_op - if _mamba_chunk_scan_combined_fwd_op is None: - _mamba_chunk_scan_combined_fwd_op = MambaChunkScanCombinedFwdOp() - return _mamba_chunk_scan_combined_fwd_op(*args, **kwargs) - - def mamba_chunk_scan_combined_varlen( x, dt, @@ -253,12 +223,11 @@ def mamba_chunk_scan_combined_varlen( dt_limit=dt_limit, state_dtype=state_dtype, ) - return varlen_states + return varlen_states -from vllm.platforms import current_platform # noqa: E402 +from vllm.platforms import current_platform if current_platform.is_cpu(): - import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_ssm # noqa: E402 - - _mamba_chunk_scan_combined_fwd = cpu_ssm._mamba_chunk_scan_combined_fwd_cpu + import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_mamba_ssm + _mamba_chunk_scan_combined_fwd = cpu_mamba_ssm._mamba_chunk_scan_combined_fwd_cpu # type: ignore From e211f6ec2970f973136755327e75322154a9fa81 Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 15 Jul 2026 04:05:28 -0500 Subject: [PATCH 42/43] precommit fix Signed-off-by: Akash kaothalkar --- vllm/model_executor/layers/mamba/ops/causal_conv1d.py | 3 ++- vllm/model_executor/layers/mamba/ops/mamba_ssm.py | 3 ++- vllm/model_executor/layers/mamba/ops/ssd_combined.py | 4 +++- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 84a3653b68c0..15f08f265504 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -1238,7 +1238,8 @@ def grid(META): out = out.squeeze(-1) return out.to(original_x_dtype) -from vllm.platforms import current_platform + +from vllm.platforms import current_platform # noqa: E402 if current_platform.is_cpu(): from vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d import ( diff --git a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py index 9bea5b0e21d7..af45467886ea 100644 --- a/vllm/model_executor/layers/mamba/ops/mamba_ssm.py +++ b/vllm/model_executor/layers/mamba/ops/mamba_ssm.py @@ -846,7 +846,8 @@ def selective_scan_fn( else: return z # output written inplace to z -from vllm.platforms import current_platform + +from vllm.platforms import current_platform # noqa: E402 if current_platform.is_cpu(): from vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm import ( diff --git a/vllm/model_executor/layers/mamba/ops/ssd_combined.py b/vllm/model_executor/layers/mamba/ops/ssd_combined.py index 3d8e354cd80f..8c645574b9e7 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_combined.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_combined.py @@ -226,8 +226,10 @@ def mamba_chunk_scan_combined_varlen( return varlen_states -from vllm.platforms import current_platform + +from vllm.platforms import current_platform # noqa: E402 if current_platform.is_cpu(): import vllm.model_executor.layers.mamba.ops.cpu.mamba_ssm as cpu_mamba_ssm + _mamba_chunk_scan_combined_fwd = cpu_mamba_ssm._mamba_chunk_scan_combined_fwd_cpu # type: ignore From df90c14ba19b30d24adadd95133f10340c28f3cf Mon Sep 17 00:00:00 2001 From: Akash kaothalkar Date: Wed, 15 Jul 2026 05:48:31 -0500 Subject: [PATCH 43/43] add cpu arm test Signed-off-by: Akash kaothalkar --- .buildkite/scripts/hardware_ci/run-cpu-test-arm.sh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh index 2d11dd477eac..09396a9697a2 100755 --- a/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh +++ b/.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh @@ -40,7 +40,9 @@ function cpu_tests() { pytest -x -v -s tests/kernels/moe/test_cpu_fused_moe.py pytest -x -v -s tests/kernels/mamba/cpu/test_cpu_gdn_ops.py pytest -x -v -s tests/kernels/moe/test_cpu_int4_moe.py - pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py" + pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py + pytest -x -v -s tests/kernels/mamba/test_causal_conv1d.py + pytest -x -v -s tests/kernels/mamba/test_mamba_ssm.py" # skip tests requiring model downloads if HF_TOKEN is not set # due to rate-limits @@ -97,3 +99,4 @@ function cpu_tests() { # All of CPU tests are expected to be finished less than 40 mins. export -f cpu_tests timeout 2h bash -c cpu_tests +