# PR #48642 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[Bugfix] Sparse MLA: enable fp8_ds_mla dense prefill
- 合并时间：2026-07-17 06:44
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/48642

---

# 执行摘要

- 一句话：修复 fp8_ds_mla dense prefill 的越界和 gather 问题
- 推荐动作：建议精读 `flashmla_sparse.py` 中 metadata 重构和 `mla_attention.py` 中 prefill 路径切换的设计，理解如何将 decode-only metadata 与 dense prefill 解耦。对有 FP8 sparse attention 需求的团队，此 PR 提供了可参考的修复模式。

# 功能与动机

PR #47327 添加了 dense-MHA prefill 路由到 sparse MLA，暴露了两个 fp8_ds_mla 混合 batch bug：1）req_id_per_token 可能超过 top-k 张量长度导致越界写；2）dense MHA 无法收集 packed 656 字节 FP8 cache，且 FlashMLA 假定其 metadata 覆盖整个 batch。本 PR 解决这些问题以支持 fp8_ds_mla 下的 dense prefill。

# 实现拆解

1. **边界检查 **（`sparse_utils.py`）：在 `triton_convert_req_index_to_global_index` 中增加长度断言，拒绝 `req_id` 和 `token_indices` 形状不匹配的调用，防止 kernel grid 越界。
2. **FP8 gather kernel 扩展 **（`csrc/cache.h`、`csrc/libtorch_stable/ops.h`）：`cp_gather_and_upconvert_fp8_kv_cache` 新增 `seq_starts` 参数，支持从任意缓存位置开始 gather；同时移除 `seq_lens` 参数简化接口。所有调用点（`_custom_ops.py`、benchmark、测试）同步更新。
3. **Metadata 重构 **（`flashmla_sparse.py`）：在 `FlashMLASparseMetadataBuilder` 中引入 `require_uniform_decodes = True` 标志。`_build_fp8_separate_prefill_decode` 现在直接使用 metadata 中已解析的 decode/prefill 计数，而非再次调用 `split_decodes_and_prefills`，使得当 dense MHA 处理 prefill 时，decode-only MQA metadata 能正确构建。
4. **Prefill 路径适配 **（`mla_attention.py`）：当 `cache_dtype == "fp8_ds_mla"` 时，`chunked_prefill_workspace` dtype 固定为 `bfloat16`（而非 query dtype），因为 FP8 cache 上转换后为 BF16。在 `_compute_prefill_context` 和 `_context_parallel_compute_prefill_context` 中，对 `fp8_ds_mla` 调用 `cp_gather_and_upconvert_fp8_kv_cache` 进行 gather+ 上转换，替代原有 `gather_and_maybe_dequant_cache`。
5. **测试覆盖 **（`test_sparse_mla_backends.py`、`test_cp_gather_fp8.py`、`test_mla_backends.py`）：新增越界拒绝、decode slice 正确性、seq_starts gather、cache dtype 映射等测试。

关键文件：
- `vllm/v1/attention/backends/mla/flashmla_sparse.py`（模块 稀疏注意后端；类别 source；类型 core-logic；符号 FlashMLASparseBackend, FlashMLASparseMetadataBuilder, require_uniform_decodes, _build_fp8_separate_prefill_decode）: 核心逻辑变更：重构 metadata 构建，支持 decode-only MQA metadata 与 dense MHA prefill 分离，引入 require_uniform_decodes 标志，修改 _build_fp8_separate_prefill_decode 接口。
- `vllm/model_executor/layers/attention/mla_attention.py`（模块 MLA 预填充；类别 source；类型 data-contract；符号 _compute_prefill_context, MLAAttention.__init__）: prefill 路径调整：在 fp8_ds_mla 时使用 cp_gather_and_upconvert kernel 替换原有 gather，workspace dtype 固定为 bfloat16
- `tests/v1/attention/test_sparse_mla_backends.py`（模块 测试覆盖；类别 test；类型 test-coverage；符号 test_triton_convert_rejects_req_id_longer_than_token_indices, test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens）: 新增越界拒绝测试和 decode slice 测试，覆盖核心 bug 场景
- `tests/kernels/test_cp_gather_fp8.py`（模块 测试覆盖；类别 test；类型 test-coverage；符号 test_cp_gather_fp8_with_sequence_starts）: 新增 seq_starts gather 测试，验证任意起始 gather 正确性
- `vllm/v1/attention/backends/mla/sparse_utils.py`（模块 工具函数；类别 source；类型 core-logic；符号 triton_convert_req_index_to_global_index）: 添加 triton_convert_req_index_to_global_index 的长度断言，防止越界

关键符号：test_triton_convert_rejects_req_id_longer_than_token_indices, test_flashmla_forward_bf16_kv_slices_req_id_to_mqa_tokens, test_cp_gather_fp8_with_sequence_starts, test_mla_kv_cache_spec_uses_layer_cache_dtype, _compute_prefill_context, _build_fp8_separate_prefill_decode, triton_convert_req_index_to_global_index

## 关键源码片段

### `vllm/v1/attention/backends/mla/flashmla_sparse.py`

核心逻辑变更：重构 metadata 构建，支持 decode-only MQA metadata 与 dense MHA prefill 分离，引入 require_uniform_decodes 标志，修改 _build_fp8_separate_prefill_decode 接口。

```python
# FlashMLASparseMetadataBuilder 中 require_uniform_decodes 标志的引入

class FlashMLASparseMetadataBuilder(SparseMLACommonMetadataBuilder[FlashMLASparseMetadata]):
    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
    # 当 require_uniform_decodes=True 时，metadata 构建器会跳过对 decode 的 split_decodes_and_prefills
    # 调用，直接使用已解析的 metadata 字段，从而支持 decode-only MQA kernel metadata
    require_uniform_decodes: ClassVar[bool] = True

    def _build_fp8_separate_prefill_decode(
        self,
        common_attn_metadata: CommonAttentionMetadata,
        metadata: FlashMLASparseMetadata,  # 传入已构建的 metadata，避免重复 split
    ) -> "FlashMLASparseMetadata.FP8SeparatePrefillDecode":
        num_tokens = common_attn_metadata.num_actual_tokens
        # 直接从 metadata 获取，而非再次调用 split_decodes_and_prefills
        num_decodes = metadata.num_decodes
        num_prefills = metadata.num_prefills
        num_decode_tokens = metadata.num_decode_tokens
        num_prefill_tokens = num_tokens - num_decode_tokens
        # ... 构建 decode 和 prefill 子结构 ...

```

### `vllm/model_executor/layers/attention/mla_attention.py`

prefill 路径调整：在 fp8_ds_mla 时使用 cp_gather_and_upconvert kernel 替换原有 gather，workspace dtype 固定为 bfloat16

```python
# _compute_prefill_context 中 fp8_ds_mla 分支

def _compute_prefill_context(self, q, kv_c_and_k_pe_cache, attn_metadata, k_scale):
    # ... 前置代码 ...
    for i in range(iters):
        toks = prefill_metadata.chunked_context.seq_tot[i]
        if self.kv_cache_dtype == 'fp8_ds_mla':
            # 使用支持任意 seq_starts 的 gather+upconvert kernel
            ops.cp_gather_and_upconvert_fp8_kv_cache(
                src_cache=kv_c_and_k_pe_cache,
                dst=workspace[:toks],
                block_table=prefill_metadata.block_table,
                workspace_starts=prefill_metadata.chunked_context.cu_seq_lens[i],
                batch_size=attn_metadata.num_prefills,
                seq_starts=prefill_metadata.chunked_context.starts[i],  # 新增：每行的起始偏移
            )
        elif not use_fp8_prefill:
            ops.gather_and_maybe_dequant_cache(...)  # 原有 BF16 路径
        else:
            # fp8 query prefill 路径不变
            ...

```

# 评论区精华

审核人 mgoin 认为设计合理，依赖 CI 验证。drakosha 在相关 Issue 的评论中确认了该分支在 4×H200 上通过 DCP 前缀缓存测试，decode 吞吐与封闭构建持平，且 gather 正确性无问题。

- DCP 前缀缓存验证 (testing): DCP 路径兼容性得到确认，无回归。

# 风险与影响

- 风险：主要风险包括：
 1) `cp_gather_and_upconvert_fp8_kv_cache` 移除了 `seq_lens` 参数，改为 `seq_starts`，若其他未发现的调用点未更新将引发编译或运行时错误；
 2) metadata 重构改变了 `_build_fp8_separate_prefill_decode` 的输入依赖，可能在某些未测试的 batch 组合下产生不一致；
 3) workspace dtype 从 `q_data_type` 改为固定 `bfloat16`，在非 fp8_ds_mla 模式下可能引入额外精度损失或性能变化；
 4) 所有变更仅在 CUDA SM90+（Blackwell）下测试，其他平台可能无法正常工作。但测试覆盖较全面，基准测试也同步更新，风险可控。
 - 影响：直接影响：使用 `fp8_ds_mla` 和 sparse MLA 的用户（如 DeepSeek 模型）的 dense prefill 路由不再崩溃，混合 batch 准确率与 sparse decode 一致。间接影响：kernel 接口变更可能影响内部其他分支或未来集成（如 DCP #46514）。团队需关注后续 merge 的 DCP 变更是否与此处 `seq_starts` 逻辑兼容。
 - 风险标记：越界风险修复 , kernel 接口变更 , metadata 重构 , 仅 CUDA SM90+ 测试

# 关联脉络

- PR #47327 dense-MHA prefill routing for sparse MLA: 引入本 PR 修复的两个 bug
- PR #48612 superseded combined fix: 本 PR supersedes #48612
- PR #46514 DCP change for KV offload: drakosha 验证本 PR 与 DCP 分支兼容，可能需要集成