Prhub

#56893 [Model][DSv4.1] Store the whole KV in MXFP8 (FlashMLA V4.1 record)

原始 PR 作者 zyongye 合并时间 2026-09-15 11:12 文件变更 14 提交数 4 评论 9 代码增减 +709 / -110

执行摘要

DSv4.1 在 SM100 改用 FlashMLA 528 B MXFP8 KV 记录,KV 缩小 9.6%

PR body 指出:DeepSeek-V4.1 一直在复用 DeepSeek-V4 的分页 fp8_ds_mla 记录,而 FlashMLA 有一个 V4.1 专用记录,会把 RoPE 维一起量化。切换后 data row 从「448 B fp8 NoPE + 128 B bf16 RoPE」变成「512 B fp8 e4m3 全量化」,scale row 从「7 个 64 维 UE8M0 + 1 pad」变成「16 个 32 维 UE8M0(MXFP8)」,每 token 从 584 B 降到 528 B(-9.6%),页 TMA stride 从 576 B 降到 512 B;由于 block_size 恒为 32 的倍数,V4.1 每页正好是 512 的整数倍,V4 需要的 padding 也一并消失。作者强调「No new cache dtype and no new flag」,--kv-cache-dtype fp8 仍解析为 fp8_ds_mla,记录随架构自动选择。

值得精读,尤其是维护 KV cache 布局、attention 后端与跨内核一致性的人。

  1. 重点看 vllm/models/deepseek_v41/attention.py 里 _use_v41_mxfp8_kv_record() 与 kv_bytes_per_token / kv_page_alignment 的绑定关系:把「记录宽度」和「页对齐」变成一组由单一 gate 派生的字段,是这次能在多处复用而不漂移的关键设计。
  2. 看 k_cache.shape[-1] 作为记录标识的做法:用已经发布的 KV cache spec 宽度做隐式契约,好处是 V4 路径零改动地跑在同一条代码路径上,代价是魔法数字散落在分派点与测试里——这正是 reviewer 要重构的点。
  3. 看索引器对齐这个「每个 spec 单独看都正确、合起来 block stride 不整除」的排障故事:它解释了为什么 alignment=576 if ... and not _use_v41_mxfp8_kv_record() else 512 必须存在,也提醒后续给同一 block 组加页类型的改动要先算总和。
  4. 若要跟进:优先按 reviewer 建议把 576/512/528/584 收敛成记录描述对象;在合入 #56344 时确认索引器对齐修复被带走;并为 SM100 上非 FlashMLA sparse 后端的记录选择补一条显式断言或测试。
讨论亮点
  1. 对齐魔数是否该收敛(ivanium,vllm/v1/attention/backends/mla/sparse_swa.py:127):行内评论指出「given that we now have self.packed_page_alignment, putting 512 here looks a bit hacky」——既然新增了 packed_page_alignment 参数,else 512 的兜底仍写死数字显得突兀。在整体 review 中进一步表态:「KV cache part LGTM. I'd prefer we can later refactor the alignment part to avoid too much hardcoded numbers like 576 and 512」。结论:本 PR 内未重构,reviewer 明确把常量收敛挂到后续工作,并以 APPROVED 放行。
  2. 自动 review 被禁用:claude[bot] 说明这是 fork 提交的 PR,自动 review 关闭,需要 maintainer 手动 @claude review——本 PR 实际只有一名人类 reviewer 过目,而改动横跨 CUDA/Triton/CuteDSL 三套内核。
  3. CI 往返与分支新鲜度:作者连续四次 /ci run,其中一次被拒绝:「This PR is 9 commits behind upstream main. Your branch must contain every commit currently on upstream main.」需要先 rebase 才能重跑,最终 38465d553f40 触发 Buildkite #88975。

实现拆解

  1. 记录门控与页几何(vllm/models/deepseek_v41/attention.py):新增 _use_v41_mxfp8_kv_record(),用 current_platform.is_device_capability_family(100) 判断是否启用 FlashMLA 的 V4.1 记录;DeepseekV4Attention.__init__ 由此派生 self.kv_mxfp8、self.kv_bytes_per_token(528/584)、self.kv_page_alignment(512/576),并把后两者注入 DeepseekV4SWACache。同一个 spec 里 state_content_bytes 和 alignment 改为从这两个字段取值,_fused_qnorm_rope_kv_insert 调 op 时多传一个 self.kv_mxfp8。效果是「记录宽度」和「页对齐」被绑定成一组,不再在 spec 与内核里各写一份常量。

  2. 索引器页对齐同步(同文件 DeepseekV4IndexerCache.get_kv_cache_spec):alignment 从固定 576 if uses_fp8_ds_mla_layout else 512 改为 576 if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record() else 512。原因是 _get_bytes_per_block 把组内每个页的 stride 相加得到块 stride,只有所有页 padding 一致时整除才成立;PR body 记录了保持 576 会产生 212160 的块 stride,在干净加载权重后的启动 profile 阶段触发 FlashMLA 的 block_stride_bytes % TMA_K_STRIDE == 0 断言,而「每个 spec 单独看都正确」。

  3. 写入端两套内核:SWA 路径给 torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert 增加 kv_mxfp8 参数,csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu 与 ops.h / torch_bindings.cpp 同步签名,MXFP8 分支改成每 lane pair 一个 UE8M0 scale、RoPE lane 参与量化;compressed-KV 路径在 vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py 新增 _rope_quant_insert_mxfp8_kernel,rope_quant_insert 按 kv_cache.shape[-1] in (584, 528) 选内核,MXFP8 版本先做 RoPE 再选 scale(RoPE 维这次也要量化)、只在压缩组边界写入。vllm/models/deepseek_v41/{amd,nvidia}/dspark.py 的 _insert_context_kv 把 attn.kv_mxfp8 透传给同一个 op。

  4. 读取端参数化记录:vllm/models/deepseek_v41/common/ops/cache_utils.py 新增 _quantize_and_insert_k_mxfp8_kernel 与 _dequantize_and_gather_k_mxfp8_kernel,入口用 bytes_per_token 参数与 k_cache.shape[-1] 分派,常量集中在模块头部(V4_BYTES_PER_TOKEN=584、V41_BYTES_PER_TOKEN=528、V41_QUANT_BLOCK=32、V41_NUM_SCALES=16);vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py 把 fp8_dim / group_size 从类常量改成 CompileKey 字段,scale_bytes = head_dim // group_size、bf16 tail 用 bf16_dim > 0 守卫,使同一份 CuteDSL 代码同时支持 448/64 与 512/32;vllm/v1/attention/backends/mla/flashmla_sparse.py 补写 V4.1 记录格式与「页内先 data 区、后 scale 区」的文档。

  5. 构建与测试配套:cmake/external_projects/flashmla.cmake 把 pin 换到 vllm-project/FlashMLA@b6ef01c,跟随 fork 目录重组(csrc/sm* → csrc/kernels/sm*)并补齐 api.cpp 注册的 fused_norm_rope_attn_rope_cast_* 与两个 permute kernel 实例化(镜像 fork 的 setup.py,去掉训练用 dense backward);commit f5f74909 再 pin 一次以修 nvcc 13.0.88 在 SM100 pass 上实例化被丢弃 else 分支里 static_assert 的构建失败(FlashMLA#24 改为 else if constexpr,SASS 不变)。测试侧在 tests/kernels/test_compressor_kv_cache.py 新增 _mxfp8_record_reference 参考量化器、test_v41_rope_insert_mxfp8_record(byte-exact,覆盖负 slot 与页 padding)、test_v41_mxfp8_cache_round_trip(insert→gather 往返并按 tile 尺度给误差上界),在 tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py 新增 test_kv_path_mxfp8_record_matches_reference 拿 CUDA 融合 op 对比 Triton 编码器并断言 NoPE 维 byte-identical。

文件 模块 状态 重要度
vllm/models/deepseek_v41/attention.py 注意力层 modified 7.67
vllm/models/deepseek_v41/common/ops/cache_utils.py 缓存算子 modified 7.18
vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py 融合算子 modified 6.2
csrc/libtorch_stable/fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu 算子内核 modified 5.26
vllm/v1/attention/backends/mla/sparse_swa.py 滑窗缓存 modified 6.25
tests/kernels/test_compressor_kv_cache.py 缓存测试 modified 6.79
tests/kernels/test_fused_deepseek_v4_qnorm_rope_kv_insert.py 算子测试 modified 6.26
vllm/models/deepseek_v4/nvidia/ops/dequant_gather_k_cutedsl.py 解量化内核 modified 5.39
vllm/v1/attention/backends/mla/flashmla_sparse.py 解码后端 modified 5.2
cmake/external_projects/flashmla.cmake 构建配置 modified 4.51
vllm/models/deepseek_v41/nvidia/dspark.py 上下文插入 modified 4.73
csrc/libtorch_stable/ops.h 算子注册 modified 4.4

关键符号

_use_v41_mxfp8_kv_record DeepseekV4Attention.__init__ DeepseekV4Attention._fused_qnorm_rope_kv_insert DeepseekV4Attention.get_kv_cache_spec DeepseekV4IndexerCache.get_kv_cache_spec DeepseekV4SWACache.__init__ DeepseekV4SWACache.get_kv_cache_spec _quantize_and_insert_k_mxfp8_kernel _dequantize_and_gather_k_mxfp8_kernel quantize_and_insert_k_cache dequantize_and_gather_k_cache_triton _rope_quant_insert_mxfp8_kernel rope_quant_insert DequantGatherKCacheKernel.dispatch DequantGatherKCacheKernel.host_entrypoint _mxfp8_record_reference test_v41_rope_insert_mxfp8_record test_v41_mxfp8_cache_round_trip test_kv_path_mxfp8_record_matches_reference

关键源码片段

vllm/models/deepseek_v41/attention.py data-contract

本次记录切换的唯一门控与页几何来源:新增 `_use_v41_mxfp8_kv_record()` 与 `kv_mxfp8` / `kv_bytes_per_token` / `kv_page_alignment`,注入 SWA 缓存、透传给融合 insert、改写压缩 KV spec,并携带索引器页对齐修复。

# vllm/models/deepseek_v41/attention.py
# V4.1 的 MXFP8 记录只在 SM100 的 FlashMLA 稀疏解码内核里可用:
# 记录宽度、页对齐、RoPE 是否量化,全部由这一个 gate 决定。
def _use_v41_mxfp8_kv_record() -> bool:
    return current_platform.is_device_capability_family(100)
​
​
class DeepseekV4Attention(...):
    def __init__(self, ...):
        ...
        # 三个量必须成组变化:528 B/512 B 页(V4.1)或 584 B/576 B 页(V4)。
        # 单独改任何一个都会让页 stride 与解码内核的 TMA stride 不再整除。
        self.kv_mxfp8 = _use_v41_mxfp8_kv_record()
        self.kv_bytes_per_token = 528 if self.kv_mxfp8 else 584
        self.kv_page_alignment = 512 if self.kv_mxfp8 else 576
​
        self.swa_cache_layer = DeepseekV4SWACache(
            head_dim=self.head_dim,
            window_size=self.window_size,
            dtype=self.kv_cache_torch_dtype,
            prefix=f"{prefix}.swa_cache",
            cache_config=cache_config,
            backend_cls=self.swa_backend_cls,
            block_size=32,
            packed_bytes_per_token=self.kv_bytes_per_token,
            packed_page_alignment=self.kv_page_alignment,
        )
​
    def _fused_qnorm_rope_kv_insert(self, q, kv, positions):
        # uint8 是 fp8_ds_mla 分页路径;最后两个参数告诉 CUDA 内核写哪条记录。
        if swa_kv_cache.dtype == torch.uint8:
            swa_kv_cache_2d = swa_kv_cache.view(swa_kv_cache.shape[0], -1)
            return torch.ops._C.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert(
                q,
                kv,
                swa_kv_cache_2d,
                swa_metadata.slot_mapping,
                positions,
                cos_sin_cache,
                self.padded_heads,
                self.eps,
                swa_metadata.block_size,
                False,
                self.kv_mxfp8, # True 时 RoPE lane 也量化,scale 密度翻倍
            )
​
    def get_kv_cache_spec(self, vllm_config):
        if not self.is_kv_source:
            return None
        uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
        return MLAAttentionSpec(
            block_size=vllm_config.cache_config.block_size,
            num_kv_heads=1,
            head_size=self.head_dim,
            dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
            tokens_per_state=self.compress_ratio,
            cache_dtype_str=self.kv_cache_dtype,
            # 页对齐随记录切换,索引器页必须与之同步(见下方 IndexerCache)。
            alignment=self.kv_page_alignment if uses_fp8_ds_mla_layout else 512,
            model_version="deepseek_v4",
            kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
            # 打包后每 token 的字节宽度;head_size 仍是语义上的 512。
            state_content_bytes=(
                self.kv_bytes_per_token if uses_fp8_ds_mla_layout else None
            ),
        )
​
​
class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
    def get_kv_cache_spec(self, vllm_config):
        # 索引器页与主 KV 页共用同一个物理 block,block stride 是组内各页之和,
        # 所以索引器必须跟着主缓存一起从 576 B 迁到 512 B,否则启动 profile
        # 阶段会触发 FlashMLA 的 block_stride_bytes % TMA_K_STRIDE 断言。
        return MLAAttentionSpec(
            ...,
            alignment=(
                576
                if uses_fp8_ds_mla_layout and not _use_v41_mxfp8_kv_record()
                else 512
            ),
        )
vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py core-logic

压缩 KV 缓存的写入端:新增 `_rope_quant_insert_mxfp8_kernel`,`rope_quant_insert` 按 cache 宽度 584/528 选内核;MXFP8 内核必须「先旋转再定 scale」,是本次量化语义变化最直接的体现。

# vllm/models/deepseek_v41/common/ops/fused_compress_quant_cache.py
# uint8 分页缓存有两种记录宽度,宽度即选择器:528 -> V4.1 MXFP8 内核。
if kv_cache.dtype == torch.uint8:
    assert kv_cache.shape[-1] in (584, 528), (
        f"unsupported paged KV record width {kv_cache.shape[-1]}"
    )
    kernel = (
        _rope_quant_insert_mxfp8_kernel
        if kv_cache.shape[-1] == 528
        else _rope_quant_insert_kernel
    )
    kernel[(num_tokens,)](
        latent, positions, cos_sin_cache, kv_cache, cache_slots, ...
    )
​
​
@triton.jit
def _rope_quant_insert_mxfp8_kernel(
    latent, positions, cos_sin, cache, cache_slots,
    COS_STRIDE: tl.constexpr,
    CACHE_STRIDE: tl.constexpr,
    CACHE_BLOCK: tl.constexpr,
    COMPRESS_RATIO: tl.constexpr,
    SANITIZE_CACHE_NANS: tl.constexpr,
):
    """V4.1 记录:先旋转,再对全部 512 维做 MXFP8 量化。
​
    RoPE 维在这里也被量化,所以必须「先旋转、后定 scale」;
    V4 内核则是先量化 NoPE,再把 RoPE 以 bf16 存下来。
    """
    t = tl.program_id(0)
    slot = tl.load(cache_slots + t)
    if slot < 0:
        return
    position = tl.load(positions + t)
    if (position + 1) % COMPRESS_RATIO != 0: # 只在压缩组边界落盘
        return
​
    d = tl.arange(0, 512)
    normed = tl.load(latent + t.to(tl.int64) * 512 + d).to(tl.float32)
​
    # NoPE 的 pair 取 (cos, sin) = (1, 0),旋转即恒等;
    # 只有 pair >= 0 的尾部 64 维才真正参与旋转。
    even, odd = tl.split(tl.reshape(normed, (256, 2)))
    pair = tl.arange(0, 256) - 224
    cs = cos_sin + (position // COMPRESS_RATIO * COMPRESS_RATIO) * COS_STRIDE
    c = tl.load(cs + tl.maximum(pair, 0), pair >= 0, other=1.0).to(tl.float32)
    s = tl.load(cs + 32 + tl.maximum(pair, 0), pair >= 0, other=0.0).to(tl.float32)
    rotated = tl.interleave(even * c - odd * s, odd * c + even * s)
    if SANITIZE_CACHE_NANS:
        rotated = tl.where(rotated == rotated, rotated, 0.0)
​
    page = cache + (slot // CACHE_BLOCK).to(tl.int64) * CACHE_STRIDE
    values = page + (slot % CACHE_BLOCK) * 512
    scales = page + CACHE_BLOCK * 512 + (slot % CACHE_BLOCK) * 16
​
    quant = tl.reshape(rotated, (16, 32))
    amax = tl.maximum(tl.max(tl.abs(quant), 1), 1e-4)
    exponent = tl.ceil(tl.log2(amax * (1.0 / 448.0)))
    scaled = quant * tl.reshape(tl.exp2(-exponent), (16, 1))
    fp8 = tl.clamp(scaled, -448.0, 448.0).to(tl.float8e4nv)
    tl.store(values + d, tl.reshape(fp8.to(tl.uint8, bitcast=True), (512,)))
​
    # NaN 清理模式下把指数码点压到 254,避免 scale 值被解释成 NaN。
    max_encoded: tl.constexpr = 254.0 if SANITIZE_CACHE_NANS else 255.0
    encoded = tl.minimum(tl.maximum(exponent + 127.0, 0.0), max_encoded)
    tl.store(scales + tl.arange(0, 16), encoded.to(tl.uint8))

评论区精华

对齐常量该不该收敛(packed_page_alignment 旁边仍写死 512) 设计

ivanium 在 `vllm/v1/attention/backends/mla/sparse_swa.py:127` 行内评论:「given that we now have self.packed_page_alignment, putting 512 here looks a bit hacky」;整体 review 里进一步说「I'd prefer we can later refactor the alignment part to avoid too much hardcoded numbers like 576 and 512」。

结论:作者未在本 PR 内重构,reviewer 以「later」措辞接受现状并给出 APPROVED,把常量收敛记为后续工作。 · approved_with_followup

fork PR 的自动 review 被禁用 other

claude[bot] 留言:「This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.」

结论:没有人手动触发,PR 实际只有 ivanium 一名人类 reviewer 覆盖了横跨 CUDA/Triton/CuteDSL 的改动。 · informational

CI 分支新鲜度校验拦下一次 /ci run other

作者连续四次 `/ci run`,其中一次被拒绝:「This PR is 9 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started.」

结论:需要先 rebase 到最新 main 再重跑,最终 commit `38465d553f40` 成功触发 Buildkite #88975。 · 已解决

FlashMLA pin 在 nvcc 13.0 上的构建失败与二次 pin other

commit `f5f74909` 说明:`b6ef01c` 在 vLLM CI 的 CUDA 13.0.3 工具链上构建失败——本 PR 新增的 16 个 `core_attn` 实例化会让 nvcc 13.0.88 实例化被丢弃 `else` 分支里的两条 `static_assert`;vllm-project/FlashMLA#24 改为 `else if constexpr`,SASS 逐字节不变。

结论:通过再 pin 解决,记录与数值不变;顺带纳入 FlashMLA#23 的可选 FP8 输出缓冲(vLLM 暂未调用)。 · 已解决

风险与影响

  1. RoPE 量化精度:64 维 RoPE 从 bf16 改为 fp8 e4m3 + UE8M0 scale,是本次唯一涉及数值行为的改动。作者 gsm8k(温度 0,n=1319)从 0.9272/0.9310 提升到 0.9333/0.9340,但 GPQA 基线用的是 deepseekv41-flash-0909 镜像且 epoch_num: 1,本 PR 是本地构建且 epoch_num: 4,作者自己也标注「the arms differ in image and epoch count as well as in the record」,因此长上下文、低温度下的 RoPE 量化影响缺少干净的专项评测。
  2. 门控粒度只到架构:_use_v41_mxfp8_kv_record() 只看 is_device_capability_family(100),与 backend_cls 无关。PR body 只声明 SM90、ROCm 与 FlashInfer SM120 保留 V4;若在 SM100 上落到非 FlashMLA 的 sparse 后端,528 B 记录能否被对端正确解码没有在本 PR 中验证。
  3. 记录宽度成为隐式契约:分派依赖 kv_cache.shape[-1] in (584, 528) 与 k_cache.shape[-1] == 528,测试里同样硬编码 584/528。一旦出现第三种记录宽度,或 block_size 不再是 32 的倍数(math.ceil(block_size * 528 / 512) * 512 的页几何前提被破坏),会走错分支或直接触发断言。
  4. 页对齐耦合且报错距离根因远:索引器与主 KV 共享同一物理 block,alignment 分别在 DeepseekV4Attention.get_kv_cache_spec 和 DeepseekV4IndexerCache.get_kv_cache_spec 两处计算(后者还重新调用 gate),任何一页 padding 不一致只会在启动 profile 阶段以 FlashMLA 的 TMA 整除断言暴露,属于难定位的启动期故障。
  5. 外部依赖与构建:pin 到 fork 的 b6ef01c 已在 CI 的 CUDA 13.0.3 上暴露 nvcc 13.0.88 实例化被丢弃 else 分支 static_assert 的构建失败,只能靠再 pin(f5f74909)绕过;flashmla.cmake 的源码清单是手工镜像 fork 的 setup.py,fork 再调整目录结构时容易漏编内核导致链接错误。
  6. 在飞 PR 冲突:#56344 改同一份 flashmla.cmake 与 V4.1 writer/reader(且其 pin c112cc1 与本 PR 只差一个提交),#56441 改同样三个 csrc/libtorch_stable/ 文件并会在 op 签名与绑定上文本冲突;PR body 明确指出 #56344 没带索引器对齐修复,需要人工带过去。
  1. 用户与推理效果:SM100(GB300/B200)上的 DeepSeek-V4.1 每 token KV 从 584 B 降到 528 B,且页不再需要 576 B padding,等效可缓存 token 数增加;作者实测启动路径分配 161.73 GiB、11,589,848 token,gsm8k 无退化。使用方式完全不变:--kv-cache-dtype fp8 仍解析成 fp8_ds_mla,不需要新 flag 或新 dtype。非 SM100 平台(SM90、ROCm、FlashInfer SM120)行为与之前逐字节一致。
  2. 系统与契约:KV cache 布局是跨 CUDA 融合 op、Triton 编码/解码、CuteDSL gather、KV cache spec 与索引器对齐的共同契约,本 PR 让同一份代码同时理解 V4(584 B / 576 B 页)与 V4.1(528 B / 512 B 页)两种记录。KV cache 不持久化,因此不存在 checkpoint 迁移问题;但任何假设「V4.1 行宽 584 B」的外部工具或下游脚本会失效。
  3. 团队维护成本:三套内核都要为 MXFP8 记录单独维护一条路径,576 / 512 / 584 / 528 分散在 attention.py、sparse_swa.py、flashmla_sparse.py、cache_utils.py 与两个测试文件中,reviewer 已明确要求后续收敛;同时与 #56344、#56441 的 rebase 协调需要作者持续跟进。
KV 布局契约变更 RoPE 量化精度风险 页对齐魔数分散 三套内核需同步 外部依赖 pin 到 fork 门控只看架构不看后端

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论