minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
guzekai01
2026-07-16 11:16:32 +08:00
committed by GitHub
parent 7f9a902cd9
commit dee91c51cf
4 changed files with 217 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py
+2
View File
@@ -253,6 +253,8 @@ def _sanity_check_input(x_fp8: Tuple[torch.Tensor, torch.Tensor]):
if x_scale.dtype == torch.int:
return
if not DEEPGEMM_SCALE_UE8M0:
return
from sglang.srt.layers.quantization.fp8_utils import ceil_to_ue8m0
python/sglang/srt/models/deepseek_v4.py
+54 -12
View File
@@ -32,6 +32,9 @@ from sglang.jit_kernel.dsv4 import (
from sglang.kernels.ops.attention.deepseek_v4_rope import (
v4_rope_inplace_npu,
)
from sglang.kernels.ops.quantization.fp8_kernel import (
sglang_per_token_group_quant_fp8,
)
from sglang.srt.compilation.compilation_config import register_split_op
from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config
from sglang.srt.distributed import (
@@ -495,10 +498,14 @@ class MqaAttentionBase(nn.Module):
**({} if fp8 else {"params_dtype": torch.bfloat16}),
)
if fp8:
from sglang.srt.layers import deep_gemm_wrapper
assert hasattr(
self.wo_a, "weight_scale_inv"
), "FP8 quant_config must create weight_scale_inv"
self.wo_a.weight_scale_inv.format_ue8m0 = True
self.wo_a.weight_scale_inv.format_ue8m0 = (
deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
)
self.wo_b = RowParallelLinear(
self.n_groups * self.o_lora_rank,
self.hidden_size,
@@ -1225,16 +1232,31 @@ class MQALayer(MqaAttentionBase):
if _FP8_WO_A_GEMM:
import deep_gemm
from sglang.srt.layers import deep_gemm_wrapper
T, G, D = o.shape
R = self.o_lora_rank
o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o)
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
# sm100 (Blackwell): ue8m0 scales via the dedicated JIT kernel.
o_fp8, o_s = sglang_per_token_group_quant_fp8_dsv4_wo_a(o)
recipe = (1, 1, 128)
else:
# sm90 (Hopper): fp32 scales.
o_fp8, o_s = sglang_per_token_group_quant_fp8(
o.reshape(T * G, D).contiguous(),
group_size=128,
scale_ue8m0=False,
)
o_fp8 = o_fp8.view(T, G, D)
o_s = o_s.view(T, G, -1)
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, o_s),
(self.wo_a.weight.view(G, R, D), self.wo_a.weight_scale_inv.data),
output,
recipe=(1, 1, 128),
recipe=recipe,
)
o = output
else:
@@ -2512,7 +2534,10 @@ class DeepseekV4ForCausalLM(nn.Module):
)
def _setup_fp8_wo_a_scales(self, is_nextn: bool) -> None:
from deep_gemm import transform_sf_into_required_layout
from sglang.srt.layers import deep_gemm_wrapper
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
from deep_gemm import transform_sf_into_required_layout
if is_nextn:
layers = [self.model.decoder]
@@ -2528,14 +2553,19 @@ class DeepseekV4ForCausalLM(nn.Module):
D = attn.wo_a.weight.shape[1]
raw_scale = attn.wo_a.weight_scale_inv.data.view(G, R // 128, D // 128)
attn.wo_a.weight_scale_inv.data = transform_sf_into_required_layout(
raw_scale,
mn=R,
k=D,
recipe=(1, 128, 128),
num_groups=G,
is_sfa=False,
)
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0:
attn.wo_a.weight_scale_inv.data = transform_sf_into_required_layout(
raw_scale,
mn=R,
k=D,
recipe=(1, 128, 128),
num_groups=G,
is_sfa=False,
)
attn.wo_a.weight_scale_inv.format_ue8m0 = True
else:
attn.wo_a.weight_scale_inv.data = raw_scale.contiguous()
attn.wo_a.weight_scale_inv.format_ue8m0 = False
def post_load_weights(self, is_nextn=False, weight_names=None):
if _FP8_WO_A_GEMM:
@@ -2748,6 +2778,18 @@ class DeepseekV4ForCausalLM(nn.Module):
futures = []
weight_names = []
for name, loaded_weight in weights:
if (
_FP8_WO_A_GEMM
and name.endswith(".wo_a.weight")
and loaded_weight.dtype != torch.float8_e4m3fn
):
raise ValueError(
f"SGLANG_OPT_FP8_WO_A_GEMM is enabled but {name} has "
f"dtype {loaded_weight.dtype}, expected "
"torch.float8_e4m3fn. This checkpoint does not provide "
"a supported fp8-quantized wo_a; rerun with "
"SGLANG_OPT_FP8_WO_A_GEMM=0."
)
try:
use_async_loading = should_async_load(loaded_weight)
python/sglang/srt/server_args.py
+20 -6
View File
@@ -6231,15 +6231,29 @@ class ServerArgs:
"--enable-deepseek-v4-fp4-indexer requires SM100 GPUs with "
"DeepGEMM FP4 indexer support."
)
# FP8 W_o GEMM requires Blackwell (sm100+). Auto-disable on Hopper.
if is_cuda() and envs.SGLANG_OPT_FP8_WO_A_GEMM.get() and get_device_sm() < 100:
if envs.SGLANG_OPT_FP8_WO_A_GEMM.is_set():
# FP8 W_o GEMM needs DeepGEMM JIT. Enable exactly where the runtime can run
# it, mirroring the forward scale split: the ue8m0 path
# (DEEPGEMM_SCALE_UE8M0, true sm100, default on) or an sm90 opt-in
# fp32-scale path (use FP4 expert ckpt). Disable in every other case.
if is_cuda() and envs.SGLANG_OPT_FP8_WO_A_GEMM.get():
from sglang.srt.layers import deep_gemm_wrapper
sm = get_device_sm()
explicit = envs.SGLANG_OPT_FP8_WO_A_GEMM.is_set()
supported = deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 or (
deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and is_sm90_supported()
and explicit
)
if not supported and explicit:
logger.warning(
"Disabling SGLANG_OPT_FP8_WO_A_GEMM: requires sm100+ (Blackwell), "
"Disabling SGLANG_OPT_FP8_WO_A_GEMM: requires DeepGEMM JIT "
"and sm100+ (Blackwell), or explicit opt-in on sm90; "
"detected sm%d.",
get_device_sm(),
sm,
)
envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
if not supported:
envs.SGLANG_OPT_FP8_WO_A_GEMM.set(False)
def _handle_cache_compatibility(self):
if self.enable_session_radix_cache and self.radix_eviction_policy != "priority":
test/manual/dsv4/test_wo_a_fp8_sm90.py
+141
View File
@@ -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()