perf(deepseek_v4): enable SGLANG_OPT_FP8_WO_A_GEMM on sm90 (Hopper) (#28983)

2026-07-16 11:16:32 +08:00
parent 7f9a902cd9
commit dee91c51cf
4 changed files with 217 additions and 18 deletions
@@ -0,0 +1,141 @@
+"""Manual correctness check for DeepSeek-V4 fp8 wo_a (deep_gemm.fp8_einsum path).
+
+Mirrors models/deepseek_v4.py MQALayer wo_a: quantize the token-major attention
+output [T, G, D] per-token-group(128) to fp8, then run the grouped matmul over the
+group/head dim via deep_gemm.fp8_einsum("bhr,hdr->bhd") -> [T, G, R], and compare
+against a bf16 einsum reference. sm100 (Blackwell) uses ue8m0 scales + recipe
+(1,1,128); sm90 (Hopper) uses fp32 scales + recipe (1,128,128). Covers the Flash
+(G=8) and Pro (G=16) shapes for both prefill (T=1024) and decode (small T).
+
+    CUDA_VISIBLE_DEVICES=0 python3 test/manual/dsv4/test_wo_a_fp8_sm90.py
+"""
+
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.layers import deep_gemm_wrapper
+
+
+@dataclass(frozen=True)
+class WoACase:
+    name: str
+    groups: int
+    tokens: int
+    k: int = 4096
+    n: int = 1024
+
+
+def cosine(a: torch.Tensor, b: torch.Tensor) -> float:
+    return F.cosine_similarity(a.float().flatten(), b.float().flatten(), dim=0).item()
+
+
+def quantize_weight_by_group(
+    weight: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-block (128x128) fp8 cast of the per-group wo_a weight [G, N, K]."""
+    from deep_gemm.utils import per_block_cast_to_fp8
+
+    groups, n, k = weight.shape
+    weight_fp8 = torch.empty_like(weight, dtype=torch.float8_e4m3fn)
+    weight_scale = torch.empty(
+        (groups, n // 128, k // 128), device=weight.device, dtype=torch.float32
+    )
+    for group in range(groups):
+        weight_fp8[group], weight_scale[group] = per_block_cast_to_fp8(
+            weight[group], use_ue8m0=False, gran_k=128
+        )
+    return weight_fp8, weight_scale
+
+
+def run_wo_a_einsum(
+    o: torch.Tensor,  # [T, G, D] bf16
+    weight_fp8: torch.Tensor,  # [G, N, K] fp8
+    weight_scale: torch.Tensor,  # [G, N/128, K/128] fp32
+) -> torch.Tensor:
+    """Mirror models/deepseek_v4.py MQALayer wo_a fp8 einsum path."""
+    import deep_gemm
+
+    from sglang.kernels.ops.quantization.fp8_kernel import (
+        sglang_per_token_group_quant_fp8,
+    )
+
+    T, G, D = o.shape
+    _, R, _ = weight_fp8.shape
+    o_fp8, o_s = sglang_per_token_group_quant_fp8(
+        o.reshape(T * G, D).contiguous(),
+        group_size=128,
+    )
+
+    recipe = (1, 128, 128)
+    output = torch.empty(T, G, R, device=o.device, dtype=torch.bfloat16)
+    deep_gemm.fp8_einsum(
+        "bhr,hdr->bhd",
+        (o_fp8.view(T, G, D), o_s.view(T, G, -1)),
+        (weight_fp8, weight_scale),
+        output,
+        recipe=recipe,
+    )
+    return output
+
+
+def check(case: WoACase, args: argparse.Namespace) -> None:
+    device = torch.device(args.device)
+    torch.manual_seed(args.seed)
+    o = (
+        torch.randn(
+            case.tokens, case.groups, case.k, device=device, dtype=torch.bfloat16
+        )
+        * 0.1
+    )
+    weight = (
+        torch.randn(case.groups, case.n, case.k, device=device, dtype=torch.bfloat16)
+        * 0.05
+    )
+    weight_fp8, weight_scale = quantize_weight_by_group(weight)
+
+    out = run_wo_a_einsum(o, weight_fp8, weight_scale)
+    bf16_ref = torch.einsum("tgd,grd->tgr", o.float(), weight.float()).to(
+        torch.bfloat16
+    )
+    torch.cuda.synchronize()
+
+    cb = cosine(out, bf16_ref)
+    print(f"{case.name}: G={case.groups} T={case.tokens} cos_bf16={cb:.6f}")
+    if cb <= args.cos_gate:
+        raise AssertionError(f"{case.name} cos_bf16 {cb} <= {args.cos_gate}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--cos-gate", type=float, default=0.999)
+    parser.add_argument("--decode-tokens", type=int, default=16)
+    args = parser.parse_args()
+
+    if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
+        print(
+            "SKIP: this manual test validates the sm90 fp32-scale wo_a einsum path; "
+            "the sm100/Blackwell production path uses ue8m0 scales + a weight-scale "
+            "transform not reproduced here."
+        )
+        return
+
+    cases = [
+        WoACase("flash prefill", groups=8, tokens=1024),
+        WoACase("flash decode", groups=8, tokens=args.decode_tokens),
+        WoACase("pro prefill", groups=16, tokens=1024),
+        WoACase("pro decode", groups=16, tokens=args.decode_tokens),
+    ]
+    for case in cases:
+        check(case, args)
+    print("ALL OK")
+
+
+if __name__ == "__main__":
+    main()