# PR #27926 完整报告 - 仓库：`sgl-project/sglang` - 标题：[DSV4] perf: Make FP8 quant output tensor contiguous - 合并时间：2026-07-08 08:41 - 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/27926 --- # 执行摘要 - 一句话：修复 DSV4 FP8 wo_a 量化 scale 布局错误并优化内存连续性 - 推荐动作：该 PR 修复了影响 DeepSeek-V4 FP8 推理正确性的关键 bug，设计上采用了更干净的专用 kernel 方案，测试也较为充分。建议合并，并考虑后续将 `scale_tma_aligned` 等遗留参数清理到独立 PR。 # 功能与动机 DeepGEMM 的 FP8 einsum 在输入张量非连续时会静默回退到未优化路径，且原有的 flat 量化方式（将 [T, G, D] 展平为 [T*G, D]）会导致 scale 与 group 的错误关联，在 Blackwell 上引入精度损失（详见 Issue #29038）。本 PR 通过新的专用量化 kernel，直接输出 group-major 布局的 scale，既避免 fallback 又修复了 scale 对应关系。 # 实现拆解 1. **新增专用 JIT 量化 kernel**(`python/sglang/jit_kernel/dsv4/fp8_wo_a.py` 与 `csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh`)：实现 `fp8_wo_a_group_major_quant_ue8m0` CUDA kernel，输入 [T, G, D] 张量，输出连续 fp8 codes 和 group-major 布局的 scale（逻辑 [T, G, D/128] 但底层存储为 [G, T, D/128]）。 2. **封装 Python 入口**：`sglang_per_token_group_quant_fp8_dsv4_wo_a` 函数创建连续输出张量，调用 JIT kernel，最后 transpose scale 为 [T, G, D/128] 满足 DeepGEMM 的消费需求。 3. **集成到模型前向 **(`python/sglang/srt/models/deepseek_v4.py`)：替换原有的 `sglang_per_token_group_quant_fp8` 调用为新函数，移除不再需要的 reshape/view 操作，简化 `deep_gemm.fp8_einsum` 的参数传递。 4. **导出模块符号 **(`python/sglang/jit_kernel/dsv4/__init__.py`)：添加 `sglang_per_token_group_quant_fp8_dsv4_wo_a` 到导出列表。 5. **新增单元测试 **(`test/registered/jit/deepseek_v4/test_fp8_wo_a.py`)：使用 flat 量化参考实现验证新量化的 bit-exact 等价性，覆盖连续 / 非连续输入、空 token 维度以及大批量场景；并注册到 B200 CI。 6. **配套的 CUDA kernel 文件 **(`python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh`)：实现 group-major scale 输出的 CUDA 核逻辑，利用 warp 规约、PDL 等待等优化。 关键文件： - `python/sglang/jit_kernel/dsv4/fp8_wo_a.py`（模块 JIT 内核；类别 source；类型 core-logic；符号 _jit_module, _fp8_wo_a_group_major_quant_ue8m0_custom_op, fp8_wo_a_group_major_quant_ue8m0, sglang_per_token_group_quant_fp8_dsv4_wo_a）: 新增的 DSV4 专用 wo_a 量化 JIT kernel 核心实现，定义了 Python 入口和 JIT 编译逻辑，是修复的核心。 - `test/registered/jit/deepseek_v4/test_fp8_wo_a.py`（模块 测试套件；类别 test；类型 test-coverage；符号 TestDeepSeekV4FP8WoA, setUpClass, _flat_reference, _strided_tgd）: 完备的单元测试，验证新量化结果与 flat 参考实现的 bit-exact 等价性，覆盖多种边界条件，并注册到 CI。 - `python/sglang/srt/models/deepseek_v4.py`（模块 模型层；类别 source；类型 data-contract）: 模型前向入口，替换了原有的量化调用，是精确保复的关键一环。 关键符号：sglang_per_token_group_quant_fp8_dsv4_wo_a, _fp8_wo_a_group_major_quant_ue8m0_custom_op, fp8_wo_a_group_major_quant_ue8m0, _flat_reference, _assert_matches_flat_reference, test_dsv4_wo_a_quant_matches_flat_reference ## 关键源码片段 ### `python/sglang/jit_kernel/dsv4/fp8_wo_a.py` 新增的 DSV4 专用 wo_a 量化 JIT kernel 核心实现，定义了 Python 入口和 JIT 编译逻辑，是修复的核心。 ```python from __future__ import annotations from typing import TYPE_CHECKING, Tuple import torch from sglang.jit_kernel.utils import ( cache_once, is_arch_support_pdl, load_jit, make_cpp_args, ) from sglang.kernel_api_logging import debug_kernel_api from sglang.srt.utils.custom_op import register_custom_op from .utils import make_name if TYPE_CHECKING: from tvm_ffi.module import Module _GROUP_SIZE = 128 @cache_once def _jit_module(in_dtype: torch.dtype, use_pdl: bool) -> Module: # 构建并缓存 JIT 编译的 CUDA 模块，使用 fast-math 确保 FP8 四舍五入与 AOT 路径一致 args = make_cpp_args(in_dtype, use_pdl) return load_jit( make_name("fp8_wo_a_group_major_quant_ue8m0"), *args, cuda_files=["deepseek_v4/fp8_wo_a_group_major_quant.cuh"], cuda_wrappers=[ ( "fp8_wo_a_group_major_quant_ue8m0", f"FP8WoAGroupMajorQuantUE8M0Kernel<{args}>::run", ) ], extra_cuda_cflags=["--use_fast_math"], ) @register_custom_op( op_name="fp8_wo_a_group_major_quant_ue8m0", mutates_args=["output_q", "output_s"], ) def _fp8_wo_a_group_major_quant_ue8m0_custom_op( input: torch.Tensor, output_q: torch.Tensor, output_s: torch.Tensor, ) -> None: """Opaque custom-op 边界，直接调用 JIT kernel 完成量化。""" assert input.dtype in (torch.bfloat16, torch.float16) module = _jit_module(input.dtype, is_arch_support_pdl()) module.fp8_wo_a_group_major_quant_ue8m0(input, output_q, output_s) @debug_kernel_api def fp8_wo_a_group_major_quant_ue8m0( input: torch.Tensor, output_q: torch.Tensor, output_s: torch.Tensor, ) -> None: _fp8_wo_a_group_major_quant_ue8m0_custom_op(input, output_q, output_s) def sglang_per_token_group_quant_fp8_dsv4_wo_a( x: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """为 DeepGEMM fp8_einsum 量化 DSV4 wo_a 激活值。 输入是 [T, G, D] bf16/fp16 张量，hidden 维度连续。 输出 fp8 codes 连续 [T, G, D]。 Scale 张量逻辑形状 [T, G, D/128]，但底层存储为 [G, T, D/128]， 使得每个 group/head 的 [T, S] 面板对 DeepGEMM recipe=(1,1,128) 消费者连续。 """ num_tokens, num_groups, hidden = x.shape hidden_groups = hidden // _GROUP_SIZE x_q = torch.empty(x.shape, device=x.device, dtype=torch.float8_e4m3fn) x_s_storage = torch.empty( (num_groups, num_tokens, hidden_groups), device=x.device, dtype=torch.float32, ) if x.numel() > 0: fp8_wo_a_group_major_quant_ue8m0(x, x_q, x_s_storage) # 转置为 DeepGEMM 期望的布局 : [T, G, D/128] return x_q, x_s_storage.transpose(0, 1) ``` ### `test/registered/jit/deepseek_v4/test_fp8_wo_a.py` 完备的单元测试，验证新量化结果与 flat 参考实现的 bit-exact 等价性，覆盖多种边界条件，并注册到 CI。 ```python class TestDeepSeekV4FP8WoA(CustomTestCase): @classmethod def setUpClass(cls): # 跳过不支持 deep_gemm 或 SM<100 的环境 if not torch.cuda.is_available(): raise unittest.SkipTest("CUDA is not available") if get_device_sm() < 100: raise unittest.SkipTest("Test requires CUDA SM 100 or higher") try: import deep_gemm except ImportError as exc: raise unittest.SkipTest("deep_gemm is required") from exc cls.deep_gemm = deep_gemm def _flat_reference(self, o): # 使用通用的 flat 量化作为参考，将所有 group 展平后量化再 reshape T, G, D = o.shape q_ref, s_ref = sglang_per_token_group_quant_fp8( o.contiguous().view(T * G, D), _GROUP_SIZE, scale_ue8m0=True, ) return q_ref.view(T, G, D), s_ref.view(T, G, D // _GROUP_SIZE) def _assert_matches_flat_reference(self, o, o_fp8, o_s): T, G, D = o.shape q_ref, s_ref = self._flat_reference(o) torch.cuda.synchronize() # 验证形状、数据类型、步长 self.assertEqual(o_fp8.shape, (T, G, D)) self.assertEqual(o_fp8.dtype, fp8_dtype) self.assertEqual(o_s.shape, (T, G, D // _GROUP_SIZE)) self.assertEqual(o_s.dtype, torch.float32) self.assertEqual(o_s.stride(), (D // _GROUP_SIZE, T * (D // _GROUP_SIZE), 1)) self.assertTrue(o_s[:, 0, :].is_contiguous()) # 验证 fp8 codes 和 scales 逐元素相等 self.assertTrue( torch.equal(o_fp8.view(torch.int8), q_ref.view(torch.int8)), "fp8 codes differ", ) self.assertTrue(torch.equal(o_s, s_ref), "scales differ") def test_dsv4_wo_a_quant_matches_flat_reference(self): # 测试连续和非连续输入下量化结果与 flat 参考一致 device = torch.device("cuda") for dtype, T, G, D in [(torch.bfloat16, 9, 5, 384), (torch.float16, 7, 3, 512)]: with self.subTest(dtype=dtype, T=T, G=G, D=D): o = (torch.randn(T, G, D, device=device, dtype=torch.float32) * 0.25).to(dtype) o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) self._assert_matches_flat_reference(o, o_fp8, o_s) # 非连续输入：通过切片创建 o = self._strided_tgd(T, G, D, dtype, device) o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) self._assert_matches_flat_reference(o, o_fp8, o_s) ``` ### `python/sglang/srt/models/deepseek_v4.py` 模型前向入口，替换了原有的量化调用，是精确保复的关键一环。 ```python from sglang.jit_kernel.dsv4 import ( fused_norm_rope_inplace, fused_q_norm_rope, fused_rope_inplace, sglang_per_token_group_quant_fp8_dsv4_wo_a, # 新增导入 ) # ... 在 forward 方法内部，FP8 wo_a 分支： if _FP8_WO_A_GEMM: import deep_gemm T, G, D = o.shape R = self.o_lora_rank # 使用专用量化函数，无需手动 reshape/view o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o) output = torch.empty(T, G, R, device=o.device, dtype=torch.bfloat16) deep_gemm.fp8_einsum( "bhr,hdr->bhd", (o_fp8, o_s), # 直接传入，无需额外 view (self.wo_a.weight.view(G, R, D), self.wo_a.weight_scale_inv.data), output, recipe=(1, 1, 128), ) o = output ``` # 评论区精华 核心讨论集中于设计选择： - Fridge003 建议创建专用 kernel 而非在通用 kernel 中添加 `scale_outer_major` 参数，以避免通用 kernel 中的 if-else 逻辑膨胀。Mattteochen 采纳该建议，将逻辑移至新的 `fp8_wo_a.py`。 - 对于 `scale_tma_aligned` 参数，Mattteochen 解释其为禁用 DeepGEMM TMA scaling kernel 所需，承诺后续 PR 中单独清理。 - 自动化 bot 提醒测试文件中无条件导入 `deep_gemm` 可能在其他平台导致 ImportError，Mattteochen 添加了 try-except 和 skip 守卫。 - 其他讨论包括清理遗留注释、移除不必要的参数等，均已解决。 - 建议使用专用 kernel 替代通用 kernel 参数扩展 (design): 采用专用 kernel 方案，移除了通用 kernel 中新增的 scale_outer_major 参数。 - 导入 deep_gemm 的保护 (testing): 测试中添加了 try-except 和 skip 逻辑，确保没有 deep_gemm 时跳过测试。 - scale_tma_aligned 参数的必要性 (question): 暂时保留，后续通过独立 PR 清理。 # 风险与影响 - 风险： 1. **平台依赖性**：新 kernel 使用了 SM100+ 的 PDL 指令且依赖 `deep_gemm`，非 Blackwell GPU 上无法运行，但通过测试中的 skip 守卫已做保护。 2. **兼容性风险**：修改了 `deepseek_v4.py` 模型前向，若其他分支或配置未使用 `_FP8_WO_A_GEMM` 则无影响；但 New kernel 的 scale 布局与 DeepGEMM 版本强相关，升级 DeepGEMM 可能需要同步调整。 3. **测试覆盖**：虽已包含多种形状的单元测试，但缺少端到端精度对比测试（如与原始 flat 路径的 logit 一致性），仅依赖 GPQA 等外部评测。 4. **性能回退**：新增的 transpose 操作增加了微小开销，但整体吞吐提升 2-4%，权衡有利。 - 影响：**影响范围**：仅影响 DeepSeek-V4 模型在开启 FP8 wo_a 量化路径时的推理行为，主要受益硬件为 Blackwell (SM100+) GPU。 **影响程度**：恢复了下游基准（GPQA、AIME25、SWE-Bench）的精度，消除了 v0.5.12 版本以来的精度退化；同时带来 2-4% 的吞吐提升，TTFT 略有改善。 **团队影响**：为后续维护提供了一个清晰的专用 kernel 模式，可推广到其他类似精度敏感的量化场景。 - 风险标记：平台依赖 (SM100+), DeepGEMM 版本兼容性 , 未覆盖端到端 logit 对比 , 遗留参数待清理 # 关联脉络 - PR #27680 [DSV4] Alternative implementation for FP8 quant layout: 本 PR 的替代实现方案，目标相同但方法不同（修改通用 kernel vs 新增专用 kernel）。 - PR #29036 [BUG] Fix DeepSeek-V4 FP8 wo_a accuracy on Blackwell: 同一 issue #29038 的修复 PR，与本 PR 解决相同问题，采用了类似专用 kernel 的方法。