diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index eed096de1..999500280 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -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" diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index bff62c44e..ba474ac5f 100755 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -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, ...]: diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py index 4ab69ddfe..6e0b5581e 100644 --- a/python/sglang/srt/utils/weight_checker.py +++ b/python/sglang/srt/utils/weight_checker.py @@ -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) diff --git a/python/sglang/srt/utils/weight_checker_comparator.py b/python/sglang/srt/utils/weight_checker_comparator.py index 1f0471d91..3578e9c58 100644 --- a/python/sglang/srt/utils/weight_checker_comparator.py +++ b/python/sglang/srt/utils/weight_checker_comparator.py @@ -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): diff --git a/test/registered/kernel/quantization/test_fp8_utils_aiter.py b/test/registered/kernel/quantization/test_fp8_utils_aiter.py new file mode 100644 index 000000000..b35a46d66 --- /dev/null +++ b/test/registered/kernel/quantization/test_fp8_utils_aiter.py @@ -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() diff --git a/test/registered/unit/utils/test_weight_checker.py b/test/registered/unit/utils/test_weight_checker.py index 47c49eaa0..702923f51 100644 --- a/test/registered/unit/utils/test_weight_checker.py +++ b/test/registered/unit/utils/test_weight_checker.py @@ -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()