执行摘要

• 一句话：修复 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 编译逻辑，是修复的核心。

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。

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

模型前向入口，替换了原有的量化调用，是精确保复的关键一环。

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 的方法。