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":