diff --git a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py index 6276dc976..c018217bb 100644 --- a/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py +++ b/python/sglang/srt/layers/deep_gemm_wrapper/entrypoint.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 6cf1d4610..8969a3d8f 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 0fce6024e..90f29b841 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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": diff --git a/test/manual/dsv4/test_wo_a_fp8_sm90.py b/test/manual/dsv4/test_wo_a_fp8_sm90.py new file mode 100644 index 000000000..6258c666b --- /dev/null +++ b/test/manual/dsv4/test_wo_a_fp8_sm90.py @@ -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()