执行摘要

• 一句话：MiniMax-M3 支持 FP8 稀疏 GQA
• 推荐动作：建议精读此 PR，以了解如何在稀疏注意力模型中集成 FP8 KV cache 支持，以及如何将手动 fallback 路径融合到统一 kernel 中。select_main_impl_cls 的选择逻辑调整（从 is_quantized_kv_cache 到 dtype 精确判断）是一个值得关注的设计决策。

功能与动机

Reland of #45680 after M3 is merged. Add support for FP8 sparse GQA on NVIDIA: only KV is quantized, Q is not. 用户需要 FP8 KV cache 以减少显存占用，但此前 fused kernel 仅支持 BF16，需要手动 fallback 路径。本 PR 将 FP8 支持直接集成到 fused kernel 中。

实现拆解

1. 扩展 fused kernel 接口：在 csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu 中添加 kv_cache_dtype 字符串参数，根据是否为 "fp8" 决定是否调用量化存储函数（reshape_and_cache_flash），而非直接写 BF16 值。Python 绑定 vllm/_custom_ops.py 同步新增参数。

2. 移除手动 fallback 路径：同时修改 amd/model.py 和 nvidia/model.py 中的 MiniMaxM3SparseAttention 类，删除 _insert_kv 方法（原用于在 fused kernel 不支持 fp8 时通过 reshape_and_cache_flash 手动写入 cache）。forward 方法不再检查 _fp8_kv 标志，统一调用 fused kernel 并传入 kv_cache_dtype。

3. 调整后端选择逻辑：在 common/sparse_attention.py 的 select_main_impl_cls 中，原条件 not is_quantized_kv_cache(kv_cache_dtype) 排除了所有量化 KV cache 使用 MSA；改为 kv_cache_dtype != "fp8_e5m2"，即允许 FP8 E4M3 在 SM100 上继续使用 MSA 后端，仅阻止 E5M2。同时增加 logger.info_once 记录所选后端类型（MSA 或 Triton）。

4. 扩展测试覆盖：tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py 中 test_sparse_full 新增 kv_cache_dtype 参数化（"auto" 和 "fp8"），测试 FP8 路径下的 cache 插入正确性，包括张量 dtype 调整（torch.uint8）和通过 reshape_and_cache_flash 生成预期缓存对比。

5. 清理与编译调整：移除不再需要的导入（如 MiniMaxM3SparseMetadata），更新 cmake/external_projects/fmha_sm100.cmake 以适应 FP8 路径的编译。

关键文件：

• vllm/models/minimax_m3/nvidia/model.py（模块 NVIDIA 模型；类别 source；类型 core-logic；符号 _insert_kv）: 核心模型文件：删除 _insert_kv 方法，forward 中新增 kv_cache_dtype 参数后调用 fused kernel；调整导入以简化依赖。
• vllm/models/minimax_m3/amd/model.py（模块 AMD 模型；类别 source；类型 core-logic；符号 _insert_kv）: AMD 对应模型文件：与 nvidia 同步变更，删除 _fp8_kv 标志和 _insert_kv 方法，forward 传入 kv_cache_dtype。
• vllm/models/minimax_m3/common/sparse_attention.py（模块 稀疏注意力；类别 source；类型 core-logic）: 稀疏注意力后端选择核心模块：修改 select_main_impl_cls 以允许 FP8 E4M3 使用 MSA 后端，并添加日志。
• tests/kernels/test_fused_minimax_m3_qknorm_rope_kv_insert.py（模块 融合核测试；类别 test；类型 test-coverage；符号 test_sparse_full）: 测试覆盖：扩展 test_sparse_full 支持 kv_cache_dtype 参数化（auto 和 fp8），验证 fp8 下 cache 插入正确性。
• csrc/libtorch_stable/fused_minimax_m3_qknorm_rope_kv_insert_kernel.cu（模块 融合核；类别 other；类型 dependency-wiring）: C++ 融合核实现：添加 kv_cache_dtype 参数，实现 fp8 量化写入路径。

关键符号：_insert_kv, select_main_impl_cls, fused_minimax_m3_qknorm_rope_kv_insert, test_sparse_full

关键源码片段

vllm/models/minimax_m3/nvidia/model.py

核心模型文件：删除 _insert_kv 方法，forward 中新增 kv_cache_dtype 参数后调用 fused kernel；调整导入以简化依赖。

# vllm/models/minimax_m3/nvidia/model.py (partial)
​
def forward(
    self,
    positions: torch.Tensor,
    hidden_states: torch.Tensor,
) -> torch.Tensor:
    qkv, _ = self.qkv_proj(hidden_states)
    # Fused per-head Gemma QK-norm + partial NeoX RoPE on q/k, in place.
    # kv_cache_dtype="auto" 表示不量化（BF16），但 fused kernel 也支持 "fp8"。
    ops.fused_minimax_m3_qknorm_rope_kv_insert(
        qkv,
        self.q_norm.weight,
        self.k_norm.weight,
        self.rotary_emb.cos_sin_cache,
        positions,
        self.num_heads,
        self.num_kv_heads,
        self.rotary_emb.rotary_dim,
        self.q_norm.variance_epsilon,
        kv_cache_dtype="auto", # 新增参数，控制 cache 写入格式
    )
    q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
    attn_output = self.attn(q, k, v)
    output, _ = self.o_proj(attn_output)
    return output

vllm/models/minimax_m3/common/sparse_attention.py

稀疏注意力后端选择核心模块：修改 select_main_impl_cls 以允许 FP8 E4M3 使用 MSA 后端，并添加日志。

# vllm/models/minimax_m3/common/sparse_attention.py (partial)
def select_main_impl_cls(
    *, topk_blocks: int, kv_cache_dtype: str,
) -> type[MiniMaxM3SparseImpl]:
    """Pick the main attend impl off the main KV-cache dtype.
​
    Blackwell (SM100) uses the MSA attend for supported top-k block counts
    when the KV cache is BF16 or FP8 E4M3; non-Blackwell and FP8 E5M2 fall
    back to Triton. The MSA module is imported lazily so AMD/non-SM100 never
    import fmha_sm100.
    """
    use_msa = (
        current_platform.is_cuda()
        and current_platform.is_device_capability_family(100)
        and topk_blocks in (4, 8, 16, 32)
        and kv_cache_dtype != "fp8_e5m2" # 允许 fp8_e4m3，仅排除 e5m2
    )
    selected = "MSA" if use_msa else "Triton"
    logger.info_once(
        "MiniMax M3 sparse attention selected %s (kv_cache_dtype=%s, topk_blocks=%s)",
        selected, kv_cache_dtype, topk_blocks,
    )
    if use_msa:
        from vllm.models.minimax_m3.nvidia.sparse_attention_msa import (
            MiniMaxM3SparseMSAImpl,
        )
        return MiniMaxM3SparseMSAImpl
    return MiniMaxM3SparseTritonImpl

评论区精华

本 PR 未产生实质 review 讨论；zyongye 给予 approve 并直接合并，评论为空。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. 后端依赖风险：FP8 KV cache 仅在 TRITON_ATTN 或 FLASHINFER (use_trtllm_attention=true) 后端下工作；FLASH_ATTN 不支持，用户需显式配置。
  2. FP8 精度风险：低精度存储可能影响生成质量；测试仅验证数值正确性，未进行端到端精度评估。
  3. 跨平台兼容性：AMD ROCm 路径 (amd/model.py) 统一传入 kv_cache_dtype="auto" (BF16)，FP8 路径尚未在 ROCm 启用，但 fused kernel 的 kv_cache_dtype 参数可能影响 ROCm 路径行为。
  4. C++ kernel 风险：新增字符串参数可能因符号错配导致运行时错误；依赖 reshape_and_cache_flash 的量化路径需保证 CUDA 版本兼容。
     - 影响：用户影响：MiniMax-M3 用户现可使用 --kv-cache-dtype fp8 降低显存占用，适用于长上下文场景；但必须配合正确的 attention backend。团队影响：移除独立 fallback 路径简化了代码维护，但仍需同步维护 AMD 与 NVIDIA 两个 model.py 文件。
     - 风险标记：FP8 精度风险, 依赖特定 attention 后端, 跨平台兼容性, C++ 核变更

关联脉络

• PR #45720 [Bugfix][ROCm] Fix MiniMax-M3 FP8 KV cache dtype: 之前修复了 M3 FP8 KV cache dtype 问题，为本 PR 在 ROCm 上的兼容性提供基础。
• PR #45743 [M3] Tune Triton indexer score decode for spec-decode: 同一作者对 M3 Triton kernel 的优化，与本 PR 的 FP8 路径可能存在交互（spec-decode 场景）。