# PR #35059 完整报告

- 仓库：`vllm-project/vllm`
- 标题：feat(cpu): add CPU support for Mamba ShortConv
- 合并时间：2026-07-07 10:47
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/35059

---

# 执行摘要

- 一句话：CPU 实现 Mamba ShortConv 前向，修复无输出的 bug
- 推荐动作：该 PR 值得精读，特别是 `forward_native` 的实现与调度模式。它展示了如何在 vLLM 的 V1 架构下为自定义算子添加 CPU 支持，同时保持与 GPU 路径的隔离。此外，review 过程中关于复用已有算子、平台判断方式等设计决策值得关注。

# 功能与动机

根据 Issue #25771，在 CPU 上运行 LiquidAI/LFM2 等 Mamba 模型时，ShortConv 的 forward_native 为空操作（仅 return），导致模型输出错误。PR body 明确说明需要实现该方法的 CPU 逻辑以修复此问题。

# 实现拆解

1. **核心逻辑实现**：在 `vllm/model_executor/layers/mamba/short_conv.py` 的 `ShortConv.forward_native` 中，导入 `vllm.model_executor.layers.mamba.ops.cpu.causal_conv1d` 提供的 `causal_conv1d_torch` 和 `causal_conv1d_update_torch` 作为参考实现。函数内部首先通过 `get_forward_context()` 获取 `ShortConvAttentionMetadata`，然后按照 prefill 和 decode 分段处理：prefill 阶段调用 `causal_conv1d_torch`（前向卷积，更新状态），decode 阶段调用 `causal_conv1d_update_torch`（单步卷积，直接从 conv_state 索引并更新）。最终输出经过 `out_proj` 线性投影。同时保留了 profile run 分支（attn_metadata 为 None）以保证内存预热的正确性。

2. **调度修改**：在同文件的 `short_conv` 自定义算子中，由原来的直接调用 `forward_cuda` 改为根据平台类型派发：若设备为 CPU 则调用 `forward_native`，否则调用 `forward_cuda`（使用 `current_platform.is_cpu()` 判断）。此修改确保了 CPU 路径使用原生实现，同时不干扰 GPU 路径。

3. **测试覆盖**：新增 `tests/kernels/mamba/test_cpu_short_conv.py`，包含三个测试用例：`test_short_conv_forward_native_prefill`（模拟 prefill 阶段，验证 KV cache 被更新）、`test_short_conv_forward_native_decode`（模拟 decode 阶段，验证多序列状态管理）和 `test_dispatch_cpu_unquantized_gemm_conv_layer`（验证卷积层不被错误地当作线性层处理）。测试文件使用 `pytest.fixture` 自动 mock 分布式环境，并通过 `current_platform.is_cpu()` 跳过非 CPU 后端。

4. **CI 集成**：将新增的测试文件注册到两个 CI 配置中：
 - `.buildkite/hardware_tests/cpu.yaml`：在 CPU-Kernel Tests 组中增加 `pytest -x -v -s tests/kernels/mamba/test_cpu_short_conv.py`。
 - `.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh`：在 ARM CI 执行列表中追加同一测试，确保 ARM 平台也能运行。

关键文件：
- `vllm/model_executor/layers/mamba/short_conv.py`（模块 短卷积层；类别 source；类型 core-logic；符号 forward_native, short_conv）: 核心修改文件：实现了 CPU 版的 forward_native，处理了 prefill/decode、KV cache 状态管理和 profile 分支，并修改了 short_conv 算子调度逻辑。
- `tests/kernels/mamba/test_cpu_short_conv.py`（模块 单元测试；类别 test；类型 test-coverage；符号 mock_dist, vllm_config, test_short_conv_forward_native_prefill, test_short_conv_forward_native_decode）: 新增测试文件，覆盖 prefill 和 decode 场景，验证 KV cache 更新和 dispatch 正确性，并自动跳过非 CPU 后端，确保回归安全。
- `.buildkite/hardware_tests/cpu.yaml`（模块 CI 配置；类别 test；类型 configuration）: 将新测试加入 CPU CI 的 kernel tests 列表中，确保在 Intel CPU 上持续运行。
- `.buildkite/scripts/hardware_ci/run-cpu-test-arm.sh`（模块 CI 脚本；类别 infra；类型 infrastructure）: 将新测试加入 ARM CI 执行列表，确保 ARM 平台也能正确运行。

关键符号：ShortConv.forward_native, short_conv (custom op wrapper), test_short_conv_forward_native_prefill, test_short_conv_forward_native_decode, test_dispatch_cpu_unquantized_gemm_conv_layer, mock_dist, vllm_config

## 关键源码片段

### `tests/kernels/mamba/test_cpu_short_conv.py`

新增测试文件，覆盖 prefill 和 decode 场景，验证 KV cache 更新和 dispatch 正确性，并自动跳过非 CPU 后端，确保回归安全。

```python
# tests/kernels/mamba/test_cpu_short_conv.py

# 模块级跳过非 CPU 平台
if not current_platform.is_cpu():
    pytest.skip("skipping CPU-only tests", allow_module_level=True)


@pytest.fixture(autouse=True)
def mock_dist():
    """Mock 分布式环境（TP rank=0, world_size=1），避免 ShortConv 初始化时需要真实分布式上下文。"""
    with (
        patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0),
        patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1),
        patch("vllm.distributed.parallel_state.model_parallel_is_initialized", return_value=True),
        patch("vllm.distributed.parallel_state.get_tp_group", return_value=MagicMock(rank_in_group=0)),
    ):
        yield


@pytest.fixture
def vllm_config():
    # 使用最小化 VllmConfig（仅 compilation_config），避免 ModelConfig 的 mock 开销
    return VllmConfig(compilation_config=CompilationConfig())


def test_short_conv_forward_native_prefill(vllm_config):
    """验证 prefill 阶段 forward_native 能正确更新 KV cache（conv_state）。"""
    prefix = "test_layer"
    config = SimpleNamespace(conv_L_cache=4, conv_bias=True)
    dim = 16

    with set_current_vllm_config(vllm_config):
        layer = ShortConv(config=config, dim=dim, layer_idx=0, prefix=prefix)
    layer.to("cpu")
    # 将线性层替换为 CPU 原生实现（移除 AMX 预处理）
    dispatch_cpu_unquantized_gemm(layer.in_proj, remove_weight=False)
    dispatch_cpu_unquantized_gemm(layer.out_proj, remove_weight=False)

    # 构造 prefill 元数据：1 个请求，5 个 token
    attn_metadata = ShortConvAttentionMetadata(
        num_prefills=1, num_prefill_tokens=5, num_decodes=0, num_decode_tokens=0,
        num_reqs=1, query_start_loc_p=torch.tensor([0, 5], dtype=torch.int32),
        has_initial_states_p=torch.tensor([False]),
        state_indices_tensor_p=torch.tensor([0], dtype=torch.int32),
        state_indices_tensor_d=torch.empty((0, 1), dtype=torch.int32),
        seq_lens=torch.tensor([5]),
        # 其余字段设为 None 或空
        ...
    )
    # KV cache 初始化为零
    conv_state = torch.zeros((1, config.conv_L_cache - 1, dim))
    layer.kv_cache = (conv_state,)

    hidden_states = torch.randn((5, dim))
    output = torch.zeros_like(hidden_states)
    attn_metadata_dict = {prefix: attn_metadata}
    with set_forward_context(attn_metadata=attn_metadata_dict, vllm_config=vllm_config):
        layer.forward_native(hidden_states, output)

    # 验证 KV cache 不为零（被写入）
    assert not torch.allclose(conv_state, torch.zeros_like(conv_state))

```

# 评论区精华

- **初始实现缺失状态处理**：gemini-code-assist[bot] 指出初始版本未使用 KV cache，导致 decode 阶段错误。开发者后续补充了完整的状态管理与连续批处理逻辑。
- **复用已有 CPU 算子**：fadara01 建议使用 `ops/cpu/causal_conv1d.py` 中的参考实现，避免双份代码。开发者采纳后重写了 `forward_native`，仅依赖 `causal_conv1d_torch` / `causal_conv1d_update_torch`。
- **AMX 兼容性**：fadara01 质疑是否需要 `is_amx()` 检查，认为 PyTorch 实现已可跨平台运行。最终版本删除了 AMX 分支，仅保留纯 torch 实现，并在注释中说明 AMX 加速可后续添加。
- **dispatch 方式**：bigPYJ1151 建议使用 `current_platform.is_cpu()` 代替设备类型判断，避免影响其他后端，最终照此修改。
- **测试文件重复与位置**：bigPYJ1151 指出存在重复测试文件，最终合并到 `tests/kernels/mamba/test_cpu_short_conv.py` 并调整了 CI 分组。
- **端到端精度验证**：fadara01 要求提供 `lm_eval` + `gsm8k` 结果以证明正确性，开发者提交了 vllm 与 hf baseline 的对齐数据（acc match），获得批准。

 - 初始实现缺少状态处理 (correctness): 开发者补充了完整的 conv_state 读取、更新逻辑，并正确实现 prefill/decode 分支。
 - 复用已有 CPU causal_conv1d 算子 (design): 开发者将 forward_native 实现改为调用 `causal_conv1d_torch` / `causal_conv1d_update_torch`。
 - dispatch 方式应使用 platform 判断 (design): 修改为 `current_platform.is_cpu()`，确保 GPU 路径不变。
 - 新增测试文件应当注册到 CI (testing): 测试被添加到两个 CI 配置中，并在 x86 和 ARM 上运行。
- AMX 兼容性分支取舍 (design): 移除了 AMX 分支，保留纯 torch 实现，后续可优化。

# 风险与影响

- 风险：
 - **数值精度风险**：当前 CPU 实现使用 PyTorch 参考算子，与 GPU 端的 Triton 或 cuDNN 优化可能存在数值偏差（例如浮点累积顺序不同），但 `lm_eval` 结果对齐表明在典型任务上可接受。
 - **性能风险**：纯 PyTorch 实现未利用 AMX/AVX 指令，在长序列或大批量情况下可能成为瓶颈，需要后续定制 kernel 优化。
 - **ARM 兼容性**：虽然已加入 ARM CI，但 PyTorch 对 ARM 的算子支持可能有细微差异，需观察 CI 实际运行结果。
 - **回归风险**：`short_conv` 调度函数的分支逻辑改动影响了所有后端，但已有 GPU CI 覆盖，且修改仅区分 `is_cpu()`，风险低。

- 影响：
 - **用户影响**：修复了 CPU 上 Mamba 模型（如 LiquidAI/LFM2）推理崩溃或输出错误的问题，使这些模型能在 CPU 上正常运行。用户无需额外配置，升级后自动生效。
 - **系统影响**：增加了约 180 行测试代码和少量 CI 配置，对构建时间影响极小。引入的 `causal_conv1d` 算子依赖已存在于代码库中，无新增依赖。
 - **团队影响**：明确了在 vLLM 中添加 CPU 算子的模式（CustomOp + forward_native + platform dispatch），为后续 Mamba 其他组件（如 SSD）的 CPU 支持提供了参考。

- 风险标记：CPU 特定路径 , 缺少端到端测试 , 纯参考实现性能慢 , ARM 兼容待观察

# 关联脉络

- PR #41025 [Mamba] Gated DeltaNet support for CPU and GPU: reviewer fadara01 要求参考该 PR 的测试方法（lm_eval+gsm8k）来验证端到端精度，且该 PR 也涉及 Mamba 和相关算子的修改，可能与本 PR 有交互。
- PR #47848 [Bugfix][CPU] Fix ARM test failure related to ShortConv: 本 PR 的测试在 ARM 上失败后，作者提交了此 hotfix 解决 ARM 兼容问题，属于后续修复。