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

[Fix] compressed-tensors block FP8: requantize weight scales to UE8M0 for DeepGEMM on Blackwell (#28662)

Browse Source
This commit is contained in:
Jimmy Shong
2026-06-26 21:41:18 +00:00
committed by GitHub
parent 7b02eab7a6
commit e745b3af22
5 changed files with 364 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/quant/test_compressed_tensors_fp8_block_ue8m0.py
+108
View File
@@ -0,0 +1,108 @@
"""Regression test for the Blackwell DeepGEMM UE8M0 weight-scale bug in the
compressed-tensors block-FP8 scheme (sgl-project/sglang#28662).
On Blackwell (SM100) block-wise FP8 dispatches to DeepGEMM, which quantizes
activations to UE8M0 scales and expects the weight scales to be UE8M0-packed
as well. If ``CompressedTensorsW8A8Fp8`` leaves the raw float32 block scales,
DeepGEMM combines UE8M0 activation scales with float32 weight scales and emits
NaN logits. ``process_weights_after_loading`` must requantize the weight scales
to UE8M0 when DeepGEMM is the active runner.
Manual / Blackwell-only: requires SM100 + DeepGEMM (DEEPGEMM_SCALE_UE8M0).
"""
import unittest
import torch
from sglang.test.test_utils import CustomTestCase
class TestCompressedTensorsFp8BlockUE8M0(CustomTestCase):
@unittest.skipUnless(torch.cuda.is_available(), "needs CUDA")
def test_block_fp8_weight_scales_requantized_to_ue8m0(self):
from sglang.srt.layers import deep_gemm_wrapper
if not (
deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
):
self.skipTest("requires Blackwell DeepGEMM (DEEPGEMM_SCALE_UE8M0)")
from compressed_tensors.quantization import (
QuantizationArgs,
QuantizationStrategy,
)
from sglang.srt.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 import (
CompressedTensorsW8A8Fp8,
)
torch.manual_seed(0)
device = "cuda"
block = [128, 128]
N, K = 512, 1024 # N % 64 == 0 and K % 128 == 0 -> DeepGEMM-supported shape
# A bf16 reference weight, block-quantized to fp8 + float32 block scales,
# exactly as a compressed-tensors `float-quantized` checkpoint stores it.
w_ref = torch.randn(N, K, device=device, dtype=torch.bfloat16) * 0.1
w_view = w_ref.float().view(N // 128, 128, K // 128, 128)
scale = w_view.abs().amax(dim=(1, 3)).clamp(min=1e-4) / 448.0 # (N/128, K/128)
w_fp8 = (w_view / scale[:, None, :, None]).to(torch.float8_e4m3fn).view(N, K)
w_deq = (
(w_fp8.float().view(N // 128, 128, K // 128, 128) * scale[:, None, :, None])
.view(N, K)
.to(torch.bfloat16)
)
weight_quant = QuantizationArgs(
num_bits=8,
type="float",
strategy=QuantizationStrategy.BLOCK,
symmetric=True,
dynamic=False,
block_structure=block,
)
scheme = CompressedTensorsW8A8Fp8(
weight_quant=weight_quant, is_static_input_scheme=False
)
from sglang.srt.layers.quantization.fp8_utils import (
deepgemm_w8a8_block_fp8_linear_with_fallback,
)
if (
scheme.w8a8_block_fp8_linear
is not deepgemm_w8a8_block_fp8_linear_with_fallback
):
self.skipTest("DeepGEMM is not the active block-FP8 runner")
# Build the layer directly (create_weights would require an initialized
# tensor-parallel group); the fix lives in process_weights_after_loading.
layer = torch.nn.Module()
layer.weight = torch.nn.Parameter(w_fp8, requires_grad=False)
layer.weight_scale = torch.nn.Parameter(
scale.to(torch.float32), requires_grad=False
)
layer.orig_dtype = torch.bfloat16
scheme.process_weights_after_loading(layer)
# The fix: weight scales must be requantized to UE8M0 for the DeepGEMM runner.
self.assertTrue(
getattr(layer.weight_scale, "format_ue8m0", False),
"block-FP8 weight scales were not requantized to UE8M0",
)
x = torch.randn(8, K, device=device, dtype=torch.bfloat16) * 0.1
y = scheme.apply_weights(layer, x)
# Pre-fix this path produced NaN; post-fix it matches the bf16 dequant ref.
self.assertTrue(torch.isfinite(y).all(), "output contains NaN/Inf")
y_ref = x.float() @ w_deq.float().t()
rel_err = ((y.float() - y_ref).norm() / y_ref.norm()).item()
self.assertLess(rel_err, 0.05, f"relative error too high: {rel_err}")
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/quantization/test_deepgemm_ue8m0_requant.py
+167
View File
@@ -0,0 +1,167 @@
"""CPU guards for DeepGEMM UE8M0 weight-scale requantization decisions."""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
import unittest
from unittest.mock import patch
import torch
from compressed_tensors.quantization import QuantizationStrategy
import sglang.srt.layers.quantization.fp8_utils as fp8_utils
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 import (
CompressedTensorsW8A8Fp8,
)
from sglang.test.test_utils import CustomTestCase
BLOCK_SIZE = [128, 128]
def _make_params(n: int = 64, k: int = 128):
weight = torch.nn.Parameter(torch.zeros((n, k)), requires_grad=False)
weight_scale = torch.nn.Parameter(torch.ones((1, 1)), requires_grad=False)
weight_scale.format_ue8m0 = False
return weight, weight_scale
class TestDeepGemmUE8M0Requant(CustomTestCase):
def _enabled_deepgemm_ue8m0(self):
return patch.multiple(
deep_gemm_wrapper,
ENABLE_JIT_DEEPGEMM=True,
DEEPGEMM_SCALE_UE8M0=True,
)
def test_helper_requants_supported_deepgemm_bf16_once(self):
weight, weight_scale = _make_params()
with self._enabled_deepgemm_ue8m0(), patch.object(
fp8_utils, "requant_weight_ue8m0_inplace"
) as requant:
fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
weight,
weight_scale,
BLOCK_SIZE,
use_deepgemm_runner=True,
output_dtype=torch.bfloat16,
weight_shape=weight.shape,
)
fired_again = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
weight,
weight_scale,
BLOCK_SIZE,
use_deepgemm_runner=True,
output_dtype=torch.bfloat16,
weight_shape=weight.shape,
)
self.assertTrue(fired)
self.assertFalse(fired_again)
self.assertTrue(weight_scale.format_ue8m0)
requant.assert_called_once_with(weight, weight_scale, BLOCK_SIZE)
def test_helper_skips_non_bf16_output(self):
weight, weight_scale = _make_params()
with self._enabled_deepgemm_ue8m0(), patch.object(
fp8_utils, "requant_weight_ue8m0_inplace"
) as requant:
fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
weight,
weight_scale,
BLOCK_SIZE,
use_deepgemm_runner=True,
output_dtype=torch.float16,
weight_shape=weight.shape,
)
self.assertFalse(fired)
self.assertFalse(weight_scale.format_ue8m0)
requant.assert_not_called()
def test_helper_skips_shape_deepgemm_will_not_run(self):
weight, weight_scale = _make_params(n=96, k=128)
with self._enabled_deepgemm_ue8m0(), patch.object(
fp8_utils, "requant_weight_ue8m0_inplace"
) as requant:
fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
weight,
weight_scale,
BLOCK_SIZE,
use_deepgemm_runner=True,
output_dtype=torch.bfloat16,
weight_shape=weight.shape,
)
self.assertFalse(fired)
self.assertFalse(weight_scale.format_ue8m0)
requant.assert_not_called()
def test_helper_skips_non_deepgemm_runner(self):
weight, weight_scale = _make_params()
with self._enabled_deepgemm_ue8m0(), patch.object(
fp8_utils, "requant_weight_ue8m0_inplace"
) as requant:
fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
weight,
weight_scale,
BLOCK_SIZE,
use_deepgemm_runner=False,
output_dtype=torch.bfloat16,
weight_shape=weight.shape,
)
self.assertFalse(fired)
self.assertFalse(weight_scale.format_ue8m0)
requant.assert_not_called()
def test_helper_skips_unsupported_block_size(self):
weight, weight_scale = _make_params()
unsupported_block_size = [128, 256]
with self._enabled_deepgemm_ue8m0(), patch.object(
fp8_utils, "requant_weight_ue8m0_inplace"
) as requant:
fired = fp8_utils.requant_block_scale_ue8m0_for_deepgemm(
weight,
weight_scale,
unsupported_block_size,
use_deepgemm_runner=True,
output_dtype=torch.bfloat16,
weight_shape=weight.shape,
)
self.assertFalse(fired)
self.assertFalse(weight_scale.format_ue8m0)
requant.assert_not_called()
def test_compressed_tensors_block_processing_preserves_ue8m0_marker(self):
scheme = CompressedTensorsW8A8Fp8.__new__(CompressedTensorsW8A8Fp8)
scheme.strategy = QuantizationStrategy.BLOCK
scheme.is_static_input_scheme = False
scheme.weight_block_size = BLOCK_SIZE
scheme.w8a8_block_fp8_linear = (
fp8_utils.deepgemm_w8a8_block_fp8_linear_with_fallback
)
layer = torch.nn.Module()
layer.weight, layer.weight_scale = _make_params()
layer.orig_dtype = torch.bfloat16
with self._enabled_deepgemm_ue8m0(), patch.object(
fp8_utils, "requant_weight_ue8m0_inplace"
) as requant:
scheme.process_weights_after_loading(layer)
scheme.process_weights_after_loading(layer)
self.assertTrue(layer.weight_scale.format_ue8m0)
requant.assert_called_once()
if __name__ == "__main__":
unittest.main(verbosity=3)