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

[AMD] Fix weight checking for AITER-shuffled block FP8 weights (#34330)

Browse Source
This commit is contained in:
Xinyu Kang
2026-09-11 15:49:56 -07:00
committed by GitHub
parent f963d7a27c
commit 6671cfc775
6 changed files with 211 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/fp8.py
+15
View File
@@ -66,6 +66,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
normalize_e4m3fn_to_e4m3fnuz,
requant_block_scale_ue8m0_for_deepgemm,
resolve_mxfp8_dense_gemm_backend,
unshuffle_aiter_fp8_weight,
)
from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod
from sglang.srt.layers.quantization.marlin_utils_fp8 import prepare_fp8_layer_for_marlin
@@ -146,6 +147,13 @@ def _require_fp4_dtype():
return fp4_dtype
def unshuffle_fp8_weight(weight: torch.Tensor) -> torch.Tensor:
"""Restore the logical layout of a backend-shuffled FP8 weight."""
if not _use_aiter:
raise RuntimeError("FP8 weight unshuffle requires AITER")
return unshuffle_aiter_fp8_weight(weight)
if _use_aiter or _use_hip_int4:
from aiter.ops.shuffle import (
moe_shuffle_scale,
@@ -771,6 +779,7 @@ class Fp8LinearMethod(LinearMethodBase):
# it so a consumer that needs the row-major layout can assert
# instead of silently reading a permuted weight.
layer.aiter_bpreshuffled = True
layer.weight.is_shuffled = True
def _process_mxfp8_linear_weight_scale(self, layer: Module) -> None:
if not self.use_mxfp8:
@@ -1666,6 +1675,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
layer.w2_weight.data = shuffle_weight(
layer.w2_weight.contiguous(), (16, 16)
)
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
return
elif self.use_mxfp8 and get_moe_a2a_backend().is_flashinfer_megamoe():
from sglang.srt.layers.moe.flashinfer_megamoe import (
@@ -1711,6 +1722,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
layer.w2_weight.data = shuffle_weight(
layer.w2_weight.contiguous(), (16, 16)
)
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
elif _use_aiter:
# Pre-shuffle weights
t = shuffle_weight(layer.w13_weight, (16, 16))
@@ -1719,6 +1732,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
t = shuffle_weight(layer.w2_weight, (16, 16))
layer.w2_weight.copy_(t)
del t
layer.w13_weight.is_shuffled = True
layer.w2_weight.is_shuffled = True
elif _is_cpu:
assert _is_cpu_amx_available, (
"Fp8MoEMethod on CPU requires that CPU has AMX support"
python/sglang/srt/layers/quantization/fp8_utils.py
+21
View File
@@ -137,6 +137,27 @@ def view_aiter_fused_rms_transposed_fp8_scale(scale: torch.Tensor) -> torch.Tens
return torch.as_strided(scale, scale.shape, (1, scale.shape[0]))
def unshuffle_aiter_fp8_weight(weight: torch.Tensor) -> torch.Tensor:
"""Undo AITER ``shuffle_weight(..., layout=(16, 16))`` for FP8 weights."""
if weight.element_size() != 1:
raise ValueError("AITER FP8 unshuffle requires a one-byte element type")
shape = weight.shape
n, k = shape[-2:]
if n % 16 != 0 or k % 32 != 0:
raise ValueError(
"AITER (16, 16) FP8 layout requires N % 16 == 0 and K % 32 == 0, "
f"got shape {tuple(shape)}"
)
return (
weight.reshape(-1, n // 16, k // 32, 2, 16, 16)
.permute(0, 1, 4, 2, 3, 5)
.contiguous()
.reshape(shape)
)
def materialize_bpreshuffle_fp8_scale_tuple(
value: Tuple[torch.Tensor, ...],
) -> Tuple[torch.Tensor, ...]:
python/sglang/srt/utils/weight_checker.py
+15 -4
View File
@@ -49,10 +49,13 @@ class CheckEntry(NamedTuple):
class QuantizedWeight(NamedTuple):
comparable_cls: type[ComparableWeight]
scale_name: str
is_shuffled: bool = False
_NON_PERSISTENT_BUFFER_PATTERNS = (
"cos_sin_cache",
"cos_cache",
"sin_cache",
"inv_freq",
"freqs_cis",
"expert_mask_gpu",
@@ -267,12 +270,14 @@ def _build_quantized_set(model) -> Dict[str, QuantizedWeight]:
if comparable_cls is None:
continue
prefix = f"{module_name}." if module_name else ""
own = {name for name, _ in module.named_parameters(recurse=False)}
for name in own:
own = dict(module.named_parameters(recurse=False))
for name, parameter in own.items():
scale = name.replace("weight", "weight_scale_inv")
if name.endswith("weight") and scale in own:
quantized_set[prefix + name] = QuantizedWeight(
comparable_cls, prefix + scale
comparable_cls,
prefix + scale,
getattr(parameter, "is_shuffled", False),
)
return quantized_set
@@ -293,7 +298,13 @@ def _build_check_entries(
continue # compared via its weight's comparable
if name in quantized_set:
qw = quantized_set[name]
yield CheckEntry(name, True, qw.comparable_cls(tensor, raw[qw.scale_name]))
yield CheckEntry(
name,
True,
qw.comparable_cls(
tensor, raw[qw.scale_name], is_shuffled=qw.is_shuffled
),
)
else:
should_compare = name not in skip_compare_names and (
not _is_non_persistent_buffer_name(name)
python/sglang/srt/utils/weight_checker_comparator.py
+22 -4
View File
@@ -2,7 +2,11 @@ from typing import Iterable, NamedTuple, Optional, Tuple
import torch
from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod, Fp8MoEMethod
from sglang.srt.layers.quantization.fp8 import (
Fp8LinearMethod,
Fp8MoEMethod,
unshuffle_fp8_weight,
)
from sglang.srt.layers.quantization.fp8_utils import (
block_quant_dequant,
inverse_transform_scale_ue8m0,
@@ -48,12 +52,21 @@ class ComparableWeight:
class Fp8BlockComparable(ComparableWeight):
"""Deepseek-style FP8 quantization."""
def __init__(self, w_q: torch.Tensor, w_s: torch.Tensor):
def __init__(
self,
w_q: torch.Tensor,
w_s: torch.Tensor,
is_shuffled: bool = False,
):
self.w_q = w_q
self.w_s = w_s
self.is_shuffled = is_shuffled
def __repr__(self) -> str:
return f"fp8_block(shape={tuple(self.w_q.shape)} dtype={self.w_q.dtype})"
return (
f"fp8_block(shape={tuple(self.w_q.shape)} dtype={self.w_q.dtype} "
f"is_shuffled={self.is_shuffled})"
)
@staticmethod
def _normalize_scale(w_q: torch.Tensor, w_s: torch.Tensor) -> torch.Tensor:
@@ -90,6 +103,8 @@ class Fp8BlockComparable(ComparableWeight):
s, block_size = self._scale_and_block_size()
for q, s_chunk in self._iter_quant_chunks(self.w_q, s, block_size[0]):
q, s_chunk = q.cuda(), s_chunk.cuda()
if self.is_shuffled:
q = unshuffle_fp8_weight(q)
yield (
block_quant_dequant(q, s_chunk, block_size, dtype=torch.bfloat16),
block_quant_dequant(
@@ -99,7 +114,10 @@ class Fp8BlockComparable(ComparableWeight):
def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
s, block_size = self._scale_and_block_size()
return block_quant_dequant(self.w_q, s, block_size, dtype=dtype)
w_q = self.w_q
if self.is_shuffled:
w_q = unshuffle_fp8_weight(w_q)
return block_quant_dequant(w_q, s, block_size, dtype=dtype)
class RawComparable(ComparableWeight):
test/registered/kernel/quantization/test_fp8_utils_aiter.py
+50
View File
@@ -0,0 +1,50 @@
# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import unittest
import torch
from sglang.srt.layers.quantization.fp8_utils import unshuffle_aiter_fp8_weight
from sglang.srt.utils import is_hip
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=5, stage="jit-kernel-unit", runner_config="amd")
@unittest.skipUnless(is_hip(), "requires ROCm AITER")
class TestAiterFp8Utils(CustomTestCase):
def test_unshuffle_weight_round_trip(self):
from aiter.ops.shuffle import shuffle_weight
for shape in ((32, 64), (2, 32, 64)):
with self.subTest(shape=shape):
logical = (
torch.arange(
torch.Size(shape).numel(), device="cuda", dtype=torch.float32
)
.remainder(7)
.to(torch.float8_e4m3fn)
.reshape(shape)
)
shuffled = shuffle_weight(logical, layout=(16, 16))
torch.testing.assert_close(
unshuffle_aiter_fp8_weight(shuffled), logical
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/utils/test_weight_checker.py
+88 -9
View File
@@ -25,6 +25,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
quant_weight_ue8m0,
transform_scale_ue8m0,
)
from sglang.srt.utils import is_hip
from sglang.srt.utils.weight_checker import (
CheckEntry,
ChecksumInfo,
@@ -42,9 +43,10 @@ from sglang.srt.utils.weight_checker_comparator import (
Fp8BlockComparable,
RawComparable,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd")
register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
@@ -75,6 +77,7 @@ def _assert_entries_close(
torch.testing.assert_close(
a_ref.w_s, e_ref.w_s, msg=f"[{i}] w_s {a_name!r}"
)
assert a_ref.is_shuffled == e_ref.is_shuffled
else:
torch.testing.assert_close(
a_ref.tensor, e_ref.tensor, msg=f"[{i}] tensor {a_name!r}"
@@ -82,18 +85,61 @@ def _assert_entries_close(
def _build_fp8_quant_pair(device: str = "cuda"):
"""Construct a real fp8-quantized weight + matching fp32 + ue8m0-packed scales.
"""Construct a real fp8-quantized weight and matching fp32 scales.
Returns (qweight, sf_fp32, sf_packed_int32) so callers can pick which scale dtype
drives the _build_check_entries branch under test.
Returns (qweight, sf_fp32).
"""
weight_bf16 = torch.randn((256, 128), dtype=torch.bfloat16, device=device)
block_size = [128, 128]
qweight, sf_fp32 = quant_weight_ue8m0(
weight_dequant=weight_bf16, weight_block_size=block_size
)
sf_packed_int32 = transform_scale_ue8m0(sf_fp32, mn=qweight.shape[-2])
return qweight, sf_fp32, sf_packed_int32
return qweight, sf_fp32
# ---------------------------------------------------------------------------
# Shuffled FP8 integration
# ---------------------------------------------------------------------------
class TestShuffledFp8Comparable(CustomTestCase):
def test_iter_chunks_unshuffles_before_dequantization(self):
shuffled = torch.zeros((32, 64), dtype=torch.float8_e4m3fn)
scale = torch.ones((2, 4), dtype=torch.float32)
comparable = Fp8BlockComparable(shuffled, scale, is_shuffled=True)
with (
patch(
"sglang.srt.utils.weight_checker_comparator.unshuffle_fp8_weight",
side_effect=lambda weight: weight,
) as unshuffle,
patch(
"sglang.srt.utils.weight_checker_comparator.block_quant_dequant",
side_effect=lambda weight, *_args, **_kwargs: weight,
),
):
next(iter(comparable.iter_chunks()))
unshuffle.assert_called_once()
def test_dequantize_unshuffles_before_checksum(self):
shuffled = torch.zeros((32, 64), dtype=torch.float8_e4m3fn)
scale = torch.ones((2, 4), dtype=torch.float32)
comparable = Fp8BlockComparable(shuffled, scale, is_shuffled=True)
with (
patch(
"sglang.srt.utils.weight_checker_comparator.unshuffle_fp8_weight",
side_effect=lambda weight: weight,
) as unshuffle,
patch(
"sglang.srt.utils.weight_checker_comparator.block_quant_dequant",
side_effect=lambda weight, *_args, **_kwargs: weight,
),
):
comparable.dequantize()
unshuffle.assert_called_once_with(shuffled)
# ---------------------------------------------------------------------------
@@ -113,6 +159,8 @@ class _TinyModel(nn.Module):
self.register_buffer("running_mean", torch.zeros(4))
# Buffer names used to exercise weight checker's hard-coded filters.
self.register_buffer("rotary_emb_cos_sin_cache", torch.full((8,), 3.14))
self.register_buffer("rotary_emb_cos_cache", torch.full((8,), 1.62))
self.register_buffer("rotary_emb_sin_cache", torch.full((8,), 0.58))
self.register_buffer("rotary_emb_freqs_cis", torch.full((8,), 2.71))
self.register_buffer("gate_proj_weight_fp32_cache", torch.full((8,), 1.41))
@@ -260,8 +308,10 @@ class TestPostprocessTensors(CustomTestCase):
# --- fp8 quant pair (real dequant on real fp8 tensors) ---
@unittest.skipIf(is_hip(), "DeepGEMM is not supported on ROCm")
def test_fp8_quant_pair_yields_lazy_pair(self):
qweight, sf_fp32, sf_packed_int32 = _build_fp8_quant_pair()
qweight, sf_fp32 = _build_fp8_quant_pair()
sf_packed_int32 = transform_scale_ue8m0(sf_fp32, mn=qweight.shape[-2])
raw = {"x.weight": qweight, "x.weight_scale_inv": sf_packed_int32}
ref = Fp8BlockComparable(qweight, sf_packed_int32)
@@ -273,8 +323,24 @@ class TestPostprocessTensors(CustomTestCase):
[("x.weight", True, ref)],
)
def test_fp8_quant_pair_preserves_shuffled_flag(self):
qweight = torch.zeros((128, 128), dtype=torch.float8_e4m3fn)
scale = torch.ones((1, 1), dtype=torch.float32)
raw = {"x.weight": qweight, "x.weight_scale_inv": scale}
quantized_set = {
"x.weight": QuantizedWeight(
Fp8BlockComparable,
"x.weight_scale_inv",
is_shuffled=True,
)
}
_assert_entries_close(
_build_check_entries(raw, set(), quantized_set),
[("x.weight", True, Fp8BlockComparable(qweight, scale, True))],
)
def test_fp8_quant_pair_yield_order_alongside_other_entries(self):
qweight, sf_fp32, _ = _build_fp8_quant_pair()
qweight, sf_fp32 = _build_fp8_quant_pair()
bias = torch.ones(4, device="cuda")
raw = {
"x.weight": qweight,
@@ -453,11 +519,12 @@ class TestBuildQuantizedSet(CustomTestCase):
model.proj.register_parameter(
"weight_scale_inv", nn.Parameter(torch.zeros(1, 1), requires_grad=False)
)
model.proj.weight.is_shuffled = True
self.assertEqual(
_build_quantized_set(model),
{
"proj.weight": QuantizedWeight(
Fp8BlockComparable, "proj.weight_scale_inv"
Fp8BlockComparable, "proj.weight_scale_inv", is_shuffled=True
)
},
)
@@ -496,6 +563,8 @@ class TestSnapshot(_WeightCheckerTestBase):
"b",
"running_mean",
"rotary_emb_cos_sin_cache",
"rotary_emb_cos_cache",
"rotary_emb_sin_cache",
"rotary_emb_freqs_cis",
"gate_proj_weight_fp32_cache",
}
@@ -526,6 +595,16 @@ class TestResetTensors(_WeightCheckerTestBase):
self.checker._reset_tensors()
torch.testing.assert_close(self.model.rotary_emb_cos_sin_cache, before)
def test_skips_cos_cache(self):
before = self.model.rotary_emb_cos_cache.clone()
self.checker._reset_tensors()
torch.testing.assert_close(self.model.rotary_emb_cos_cache, before)
def test_skips_sin_cache(self):
before = self.model.rotary_emb_sin_cache.clone()
self.checker._reset_tensors()
torch.testing.assert_close(self.model.rotary_emb_sin_cache, before)
def test_skips_freqs_cis(self):
before = self.model.rotary_emb_freqs_cis.clone()
self.checker._reset_tensors()