[Hardware][CPU] Enable granite-4 model on cpu#47641
Conversation
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@dhcp-9-123-5-76.bl1-in.ibm.com> Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan> Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@akashs-mbp.bl1-in.ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash Kaothalkar <akashkaothalkar@Akashs-MBP.lan>
# Conflicts: # csrc/cpu/torch_bindings.cpp
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
|
Hi @bigPYJ1151, I have addressed the review comments, requesting you to review once. |
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
![depthfirst-app[bot]](https://avatars.githubusercontent.com/in/1021002?s=60&v=4){kind=small}
| acc += static_cast<float>(w[k]) * | ||
| static_cast<float>(sd[k * stride_s_state]); | ||
| } | ||
| acc += static_cast<float>(w[state_len]) * x_val; |
There was a problem hiding this comment.
⚪ Severity: LOW
The weight array w has width elements (indices 0 to width-1), but state_len is read independently from state_c.size(2). Unlike the Triton kernel which safely derives state_len = width - 1, this C++ wrapper passes both values from separate tensors without any TORCH_CHECK(state_len < width). If state_len >= width, w[state_len] performs a heap out-of-bounds read.
Helpful? Add 👍 / 👎
💡 Fix Suggestion
Suggestion: Add a TORCH_CHECK in csrc/cpu/mamba_cpu.cpp after line 58 (where both width and state_len are computed) to validate that state_len < width before calling the kernel. This ensures w[state_len] at line 62 of mamba_kernels.hpp cannot perform an out-of-bounds read. Add the following line after int64_t state_len = state_c.size(2); (line 58 of mamba_cpu.cpp):
TORCH_CHECK(state_len < width,
"causal_conv1d_update: conv_state state_len (", state_len,
") must be < weight width (", width, ")");This mirrors the safety guarantee that the Triton kernel gets for free by deriving state_len = width - 1.
|
Hi @bigPYJ1151 , did you get a chance to look at the changes ? |
|
This pull request has merge conflicts that must be resolved before it can be |
| from vllm.v1.attention.backends.utils import PAD_SLOT_ID | ||
|
|
||
|
|
||
| def _causal_conv1d_fn_cpu( |
There was a problem hiding this comment.
Kernels in this file should be moved to mamba/ops/cpu
And there is a torch fallback conv_1d kernel https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/mamba/ops/cpu/causal_conv1d.py , is it possible to reuse and extend?
| # 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". |
| from vllm.platforms import CpuArchEnum | ||
|
|
||
| if current_platform.get_cpu_architecture() == CpuArchEnum.X86: |
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
| auto get_int32_ptr = | ||
| [](const c10::optional<at::Tensor>& opt) -> const int32_t* { | ||
| return (opt.has_value() && opt.value().defined()) | ||
| ? opt.value().data_ptr<int32_t>() |
There was a problem hiding this comment.
🟡 Severity: MEDIUM
Unlike causal_conv1d_update_cpu_impl (line 73) and mamba_chunk_scan_fwd_cpu_impl (line 265) which safely convert with .to(at::kInt).contiguous(), this get_int32_ptr lambda calls data_ptr<int32_t>() without dtype validation. In release builds, PyTorch skips type-checking in data_ptr<T>(), so if any caller passes int64 tensors (e.g., cu_seqlens, state_batch_indices), memory is silently reinterpreted as int32, producing corrupted index values used for state pointer arithmetic—causing out-of-bounds writes.
Helpful? Add 👍 / 👎
💡 Fix Suggestion
Suggestion: Refactor the get_int32_ptr lambda to first convert the optional tensor to at::kInt (int32) and make it contiguous before calling data_ptr<int32_t>(), consistent with the safe pattern used in causal_conv1d_update_cpu_impl (line 73) and mamba_chunk_scan_fwd_cpu_impl (line 261). Because the converted tensor must outlive the pointer, change the lambda to return an at::Tensor instead, keep the converted tensors alive in local variables, and then extract raw pointers. Replace lines 180–189 with:
auto safe_to_int32 =
[](const c10::optional<at::Tensor>& opt) -> at::Tensor {
if (opt.has_value() && opt.value().defined())
return opt.value().to(at::kInt).contiguous();
return {};
};
at::Tensor sbi_t = safe_to_int32(state_batch_indices);
at::Tensor dsbi_t = safe_to_int32(dst_state_batch_indices);
at::Tensor nat_t = safe_to_int32(num_accepted_tokens);
at::Tensor csl_t = safe_to_int32(cu_seqlens);
const int32_t* sbi_ptr = sbi_t.defined() ? sbi_t.data_ptr<int32_t>() : nullptr;
const int32_t* dsbi_ptr = dsbi_t.defined() ? dsbi_t.data_ptr<int32_t>() : nullptr;
const int32_t* nat_ptr = nat_t.defined() ? nat_t.data_ptr<int32_t>() : nullptr;
const int32_t* csl_ptr = csl_t.defined() ? csl_t.data_ptr<int32_t>() : nullptr;This ensures that even if callers pass int64 tensors (PyTorch's default integer dtype), they are safely converted to int32 before reinterpretation, preventing silent memory corruption and potential out-of-bounds access.
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
| for (; n <= dstate - VEC_ELEM_NUM; n += VEC_ELEM_NUM) { | ||
| 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))); |
There was a problem hiding this comment.
🟡 Severity: MEDIUM
The selective_state_update_kernel accesses state elements as s_hd[n] (pointer+offset), implicitly assuming state.stride(3)==1 (dstate dimension is contiguous). However, selective_state_update_cpu_impl only captures strides for dimensions 0–2 and never validates dimension-3 contiguity. Unlike mamba_chunk_scan_fwd_cpu_impl (which has a TORCH_CHECK(final_states.is_contiguous())) and causal_conv1d_update_kernel (which explicitly passes and uses stride_s_state), this kernel has no such protection, leading to heap OOB reads/writes if state.stride(3) != 1.
Helpful? Add 👍 / 👎
💡 Fix Suggestion
Suggestion: Add a TORCH_CHECK in the selective_state_update_cpu_impl function in csrc/cpu/mamba_cpu.cpp, after capturing the state strides (around line 170), to validate that the dstate dimension is contiguous. For example, add:
TORCH_CHECK(state.stride(3) == 1,
"selective_state_update_cpu: state must be contiguous in the "
"dstate (last) dimension, got stride ", state.stride(3));This mirrors the pattern used in mamba_chunk_scan_fwd_cpu_impl (line 269) which has TORCH_CHECK(final_states.is_contiguous(), ...). Alternatively, you could check full contiguity with state.is_contiguous() if the memmove on line 244 (which also assumes contiguity across heads) should be protected too.
|
Hi @bigPYJ1151 , I have addressed review comments, please take a look. |
| from vllm.v1.attention.backends.utils import PAD_SLOT_ID | ||
|
|
||
|
|
||
| def causal_conv1d_fn_cpu( |
There was a problem hiding this comment.
Is it possible to combine this with causal_conv1d_torch?
There was a problem hiding this comment.
Done! We merged the fast F.conv1d logic into causal_conv1d_fn_cpu and removed the redundant function.
| """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) |
There was a problem hiding this comment.
Why not use causal_conv1d_update_cpu_vec?
| @@ -1066,7 +1066,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, | |||
| @@ -1237,3 +1237,50 @@ def grid(META): | |||
| if unsqueeze: | |||
| 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, | |||
There was a problem hiding this comment.
These changes can be reverted.
| @@ -494,7 +495,7 @@ def _selective_scan_update_kernel( | |||
| tl.store(dst_state_ptrs, state, mask=mask) | |||
|
|
|||
|
|
|||
| def selective_state_update( | |||
| def _selective_state_update_cuda( | |||
| state, | |||
| x, | |||
| dt, | |||
| @@ -605,6 +606,26 @@ def selective_state_update( | |||
| 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 = ( | |||
| @@ -845,3 +866,58 @@ def selective_scan_fn( | |||
| return delta # output written inplace to delta | |||
| 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, | |||
| ) | |||
There was a problem hiding this comment.
These changes can be reverted.
| @@ -154,6 +155,35 @@ def _mamba_chunk_scan_combined_fwd( | |||
| 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) | |||
|
|
|||
|
|
|||
There was a problem hiding this comment.
These changes can be reverted.
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
Signed-off-by: Akash kaothalkar <akash.kaothalkar@ibm.com>
2 participants
Purpose
Enables the IBM Granite 4.0 model family (8 variants: h-micro, h-micro-base, h-small, h-small-base, h-tiny, h-tiny-base, tiny-preview, tiny-base-preview) to run correctly on CPU (ppc64le/Power10) with vLLM. This addresses [. #27971 — CPU inference failure for granite-4.0-h-tiny].
Test Plan
Test Result
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.