执行摘要

• 一句话：修复 Mistral Small 4 MLA prefill 被排除的问题
• 推荐动作：低风险、易理解的 bugfix PR，值得快速合入。如果有其他 MLA 模型（如 DeepSeek 或其他变种）使用不同头部维度，可以借鉴此模式进行扩展。

功能与动机

Mistral Small 4 模型使用 MLA 头部维度 (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128)，但 FlashAttnPrefillBackend 此前未将其收录，导致该模型无法使用 FlashAttention MLA prefill backend，影响了推理性能。

实现拆解

1. 在 vllm/v1/attention/backends/mla/prefill/flash_attn.py 中，于 FlashAttnPrefillBackend.supports_mla_dimensions 方法内新增一个 MLADimensions 实例 dims_mistral_s4，包含 (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128)。
2. 修改返回值逻辑：FA4 版本从只接受 deepseek 改为接受 deepseek 和 mistral_s4；其他版本 (FA2/FA3) 从接受 deepseek 和 glm 改为接受 deepseek、glm 和 mistral_s4。
3. 在 tests/v1/attention/test_mla_prefill_selector.py 中，将 test_auto_selection_on_hopper 的参数化元组扩展，新增 (64, 128) 组合并赋予 id 'mistral_s4'。
4. 在 docs/design/attention_backends.md 中更新 FLASH_ATTN 后端说明，新增支持 (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) 维度组合。

关键文件：

• vllm/v1/attention/backends/mla/prefill/flash_attn.py（模块 注意力；类别 source；类型 core-logic；符号 supports_mla_dimensions）: 核心变更文件，新增了 Mistral Small 4 的 MLA 维度组合并更新了 supports_mla_dimensions 的返回值逻辑。
• tests/v1/attention/test_mla_prefill_selector.py（模块 测试；类别 test；类型 test-coverage；符号 test_auto_selection_on_hopper）: 测试文件，扩展了参数化测试用例以覆盖 Mistral Small 4 维度。
• docs/design/attention_backends.md（模块 文档；类别 docs；类型 documentation）: 文档文件，更新了 FLASH_ATTN 后端支持的 MLA 维度列表。

关键符号：supports_mla_dimensions, test_auto_selection_on_hopper

关键源码片段

vllm/v1/attention/backends/mla/prefill/flash_attn.py

核心变更文件，新增了 Mistral Small 4 的 MLA 维度组合并更新了 supports_mla_dimensions 的返回值逻辑。

# vllm/v1/attention/backends/mla/prefill/flash_attn.py
# 声名三个 MLA 维度配置：DeepSeek、GLM 和新增的 Mistral Small 4
@classmethod
def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:
    dims_deepseek = MLADimensions(
        qk_nope_head_dim=128,
        qk_rope_head_dim=64,
        v_head_dim=128,
    )
    dims_glm = MLADimensions(
        qk_nope_head_dim=192,
        qk_rope_head_dim=64,
        v_head_dim=256,
    )
    # 新增：Mistral Small 4 的 MLA 维度组合
    dims_mistral_s4 = MLADimensions(
        qk_nope_head_dim=64,
        qk_rope_head_dim=64,
        v_head_dim=128,
    )
    fa_version = get_flash_attn_version()
    if fa_version == 4:
        # FA4 目前仅支持 deepseek 和 mistral_s4
        return mla_dimensions in [dims_deepseek, dims_mistral_s4]
    else:
        # FA2/FA3 支持 deepseek、glm 和 mistral_s4
        return mla_dimensions in [dims_deepseek, dims_glm, dims_mistral_s4]

tests/v1/attention/test_mla_prefill_selector.py

测试文件，扩展了参数化测试用例以覆盖 Mistral Small 4 维度。

# tests/v1/attention/test_mla_prefill_selector.py
# 扩展参数化测试，覆盖 Mistral Small 4 的维度组合
@pytest.mark.parametrize(
    ("qk_nope_head_dim", "v_head_dim"),
    [
        (128, 128), # deepseek
        (192, 256), # glm
        (64, 128), # mistral_s4 — 新增
    ],
    ids=["deepseek", "glm", "mistral_s4"],
)
def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int):
    # 测试逻辑不变，用于验证自动选择时 Mistral Small 4 维度能被正确路由到 FLASH_ATTN 后端。
    # ...（后续代码保持不变）

评论区精华

仅有一条机器人评论：Claude 检测到 PR 来自 fork 自动 review 被禁用。后续由 NickLucche 审批通过。未发现设计争议。

• 暂无高价值评论线程

风险与影响

• 风险：变更局限在一个类方法和参数化测试中，回归风险低。FA2/FA3/FA4 下的条件判断逻辑保持一致，不会影响现有支持模型。文档仅作维度列举更新，无结构性风险。
• 影响：影响范围小：仅修正 Mistral Small 4 模型在 V1 引擎下使用 FlashAttention MLA prefill 的路径选择，不影响其他模型或后端。降低了对 Mistral Small 4 的推理性能开销。
• 风险标记：无显著风险

关联脉络

• PR #48642 [Bugfix] Sparse MLA: enable fp8_ds_mla dense prefill: 同为 MLA prefill 相关的 bugfix PR，涉及类似的后端选择逻辑。