执行摘要

• 一句话：为 QSA 索引器新增 FP8 e4m3 缓存支持，提升 Prefill 和 Decode 性能。
• 推荐动作：本 PR 是一个高质量的性能特性增强，具有详尽的基准测试和针对性的测试覆盖。建议技术专家重点审查 qsa_indexer.py 中 Triton 内核的动态配置逻辑以及 assert_fp8_within_one_ulp 的鲁棒性。对于追求极致吞吐的用户，这是一个非常值得关注的特性。

功能与动机

当前 QSA 索引器仅支持 BF16 缓存，这在长上下文推理场景下可能成为内存和计算瓶颈。PR body 中提供了详尽的微基准测试，显示在 TP4 GB300 上，FP8 路径相比 BF16 在 Decode 和 Prefill 阶段均有 1.1x 到 1.7x 的加速。目标是通过 FP8 量化减少显存占用并利用硬件 FP8 算力，提升 Qwen3.8-Flash-Next 的推理吞吐。

实现拆解

1. 数据契约扩展：修改 QSAStateBackend 类（vllm/models/qwen4_exp/common/qsa_cache.py），将 supported_kv_cache_dtypes 从仅 BF16 扩展为包含 fp8 和 fp8_e4m3。同时将 _QSAStateCache 中的类型校验逻辑从硬编码 torch.bfloat16 改为基于 self.dtype 的动态校验，以支持多类型缓存。
2. 配置与类型逻辑集成：在 QSAIndexer.__init__（vllm/models/qwen4_exp/nvidia/indexer_qsa.py）中引入 indexer_kv_dtype 配置读取。根据配置将 indexer_dtype 设置为 torch.float8_e4m3fn 或 torch.bfloat16，并将其传递给 QSACompressedKeyCache。在 Forward 路径中，确保 query 和 compressed key 在存储/计算前被转换为对应的 indexer_dtype。
3. Triton 内核优化：更新 vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py 中的 Decode 和 Prefill 内核。针对 FP8 数据调整 num_warps 和 TILE_R 等参数，并添加了基于数据类型的动态分支，以适配不同硬件（SM90/SM100）的 FP8 运算特性。
4. 测试配套：更新 test_qsa_pre_indexer.py 和 test_qsa_reference.py，新增 indexer_dtype 参数化。引入 assert_fp8_within_one_ulp 辅助函数，处理 FP8 舍入误差；在 Decode 选择测试中，针对 FP8 比较逻辑改为值集合比较（Multiset comparison），以应对 FP8 精度下的排序不稳定性。

关键文件：

• vllm/models/qwen4_exp/nvidia/indexer_qsa.py（模块 Qwen4Exp 索引器；类别 source；类型 core-logic；符号 QSAIndexer, indexer_dtype）: 核心逻辑入口，负责 FP8 类型的决策与 Forward 路径中类型的强制转换。
• vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py（模块 QSA Triton 内核；类别 infra；类型 infrastructure；符号 _prefill_logits, qsa_select_paged_decode）: 底层 Triton 内核实现，包含针对 FP8 的性能调优参数和设备适配逻辑。

关键符号：QSAIndexer.init, QSAIndexer.forward, _prefill_logits, qsa_select_paged_decode, assert_fp8_within_one_ulp

关键源码片段

vllm/models/qwen4_exp/nvidia/indexer_qsa.py

核心逻辑入口，负责 FP8 类型的决策与 Forward 路径中类型的强制转换。

# 在 QSAIndexer.__init__ 中根据配置选择索引器缓存的 dtype
# Q 和压缩 K 缓存在存储前会被 RMSNormed，然后直接进行 fp8 x fp8 的 dot 运算
self.indexer_kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype(
    "bf16"
)
if self.indexer_kv_dtype == "fp8":
    indexer_dtype = torch.float8_e4m3fn
elif self.indexer_kv_dtype == "bf16":
    indexer_dtype = torch.bfloat16
else:
    raise NotImplementedError(
        f"indexer_kv_dtype={self.indexer_kv_dtype!r} is not supported "
        "by the Qwen4Exp QSA indexer (only 'bf16' or 'fp8')."
    )
self.indexer_dtype = indexer_dtype

评论区精华

Reviewer: ZJY0516 提出当前的 FP8 精度校验不如 dsv4 测试严格。
Author: gau-nernst 回复称已尝试实现完全位相同步，但由于融合内核与非融合内核在中间累加路径上的微小差异（如 bf16 1ulp 偏差），无法保证 bit-identical 结果，当前方案已是可达到的最严标准。最终在 SM90 上放宽了 1e-3 的容差，以补偿 wgmma 指令的 reduced precision 累加特性。

• FP8 精度校验与位相同步的权衡 (correctness): 采用基于 uint8 code diff 和绝对差值的混合校验方案，并为 SM90 架构设置了 1e-3 的特定容差。
• Triton 内核跨架构适配性 (performance): 最终移除了平台门控，认为 fp8e4nv dots 在 SM120 上同样可用。

风险与影响

• 风险：
  1. 精度风险：FP8 e4m3 相比 BF16 信息损失较大。虽然测试显示 GSM8K 和 LongBench-v2 分数持平或略升，但在极长上下文或高精度要求场景下需持续监控。
  2. 硬件兼容性：虽然实现了 SM90/SM100 的优化，但早期提交中曾尝试针对 SM120 设置门控后又移除。依赖 Triton 对不同架构 FP8 指令的支持稳定性。
  3. 内核配置硬编码：Triton 内核中的 num_warps 和 TILE_R 是针对 GB300 调优的“硬编码”常量。如果在不同架构（如 H100 或 AMD）上运行，性能可能无法达到最优甚至退化。
     - 影响：该变更显著提升了 Qwen4Exp 模型在高性能硬件上的推理效率。通过支持 FP8，能够降低 KV 缓存带宽压力，为用户提供更高的并发支持（TTFT 和 TPOT 在高负载下有明显改善）。它进一步完善了 vLLM 对新硬件特性的适配能力。
     - 风险标记：硬件架构依赖, 内核参数硬编码, FP8 精度退化

关联脉络

• PR #55301 [Transformers] Generalize merged-column linear fusion: 属于 Qwen4Exp 模型架构演进系列的一部分，涉及类似的内核优化与模块重构。