From 716a6bf10ca35cab33309d5177961c954546f2c4 Mon Sep 17 00:00:00 2001 From: guzekai01 Date: Mon, 24 Aug 2026 18:59:27 +0800 Subject: [PATCH] feat(humming): support native W4AFP8 checkpoint schemas (#32033) --- .../sglang/srt/layers/quantization/humming.py | 352 +++++++++++++++++- .../test_humming_w4afp8_schemas.py | 176 +++++++++ 2 files changed, 516 insertions(+), 12 deletions(-) create mode 100644 test/registered/unit/layers/quantization/test_humming_w4afp8_schemas.py diff --git a/python/sglang/srt/layers/quantization/humming.py b/python/sglang/srt/layers/quantization/humming.py index 3f976b1eb..a758f801c 100644 --- a/python/sglang/srt/layers/quantization/humming.py +++ b/python/sglang/srt/layers/quantization/humming.py @@ -194,12 +194,286 @@ def compressed_tensors_get_config(config: dict[str, Any], key: str): return target_group_config +class _CheckpointWeightSchema: + def process_loaded_weight(self, tensor: torch.Tensor, name: str) -> torch.Tensor: + return tensor + + def get_padded_tensors_attrs( + self, + shape_n: int, + shape_k: int, + param_dtype: torch.dtype, + num_experts: int | None = None, + has_bias: bool = False, + pad_n_to_multiple: int = 1, + pad_k_to_multiple: int = 1, + stack_size: int = 1, + ) -> dict[str, dict[str, Any]]: + _lazy_import_humming() + return BaseWeightSchema.get_padded_tensors_attrs( + self, + shape_n=shape_n, + shape_k=shape_k, + param_dtype=param_dtype, + num_experts=num_experts, + has_bias=has_bias, + pad_n_to_multiple=pad_n_to_multiple, + pad_k_to_multiple=pad_k_to_multiple, + stack_size=stack_size, + ) + + +class _W4AFp8CheckpointWeightSchema(_CheckpointWeightSchema): + quant_method = "w4afp8" + + def __init__(self, group_size: int = 128): + if ( + not isinstance(group_size, int) + or isinstance(group_size, bool) + or group_size <= 0 + ): + raise ValueError( + f"W4AFP8 group_size must be a positive integer, got {group_size!r}." + ) + self.group_size = group_size + + def get_tensors_attrs( + self, + shape_n: int, + shape_k: int, + param_dtype: torch.dtype, + num_experts: int | None = None, + has_bias: bool = False, + stack_size: int = 1, + ) -> dict[str, dict[str, Any]]: + if shape_k % self.group_size != 0: + raise ValueError( + f"W4AFP8 shape_k = {shape_k} must be divisible by group_size = " + f"{self.group_size}. Choose a tensor-parallel configuration whose " + "local K dimension preserves quantization groups." + ) + if shape_k % 8 != 0: + raise ValueError( + f"W4AFP8 shape_k = {shape_k} must be divisible by 8 for int32 " + "packed-weight storage." + ) + + tensors_attrs = { + "weight": { + "shape": (shape_n, shape_k // 2), + "dtype": torch.int8, + "extra_attrs": {"output_dim": 0, "input_dim": 1}, + }, + "weight_scale_inv": { + "shape": (shape_n, shape_k // self.group_size), + "dtype": param_dtype, + "extra_attrs": { + "output_dim": 0, + "input_dim": 1, + "scale_type": "group", + }, + }, + } + if has_bias: + tensors_attrs["bias"] = { + "shape": (shape_n,), + "dtype": param_dtype, + "extra_attrs": {"output_dim": 0}, + } + _lazy_import_humming() + return BaseWeightSchema.may_add_expert_dim(tensors_attrs, num_experts) + + def convert_humming( + self, + tensors: dict[str, torch.Tensor], + shape_n_stacks: list[int], + shape_k_stacks: list[int], + param_dtype: torch.dtype, + num_experts: int | None = None, + ): + _lazy_import_humming() + schema = HummingWeightSchema( + b_dtype=DataType.from_str("uint4"), + bs_dtype=DataType.from_torch_dtype(param_dtype), + weight_scale_group_size=self.group_size, + ) + weight = tensors["weight"].view(torch.uint8).bitwise_xor(0x88) + output_tensors = { + "weight": weight.contiguous().view(torch.int32), + "weight_scale": tensors["weight_scale_inv"].to(param_dtype), + } + if "bias" in tensors: + output_tensors["bias"] = tensors["bias"] + return schema, output_tensors + + +class _StackedBlockFp8CheckpointWeightSchema(_CheckpointWeightSchema): + def __init__(self, schema): + self.schema = schema + self.quant_method = schema.quant_method + weight_block_size = tuple(schema.weight_block_size) + if len(weight_block_size) != 2 or any( + not isinstance(size, int) or isinstance(size, bool) or size <= 0 + for size in weight_block_size + ): + raise ValueError( + "FP8 weight_block_size must contain two positive integers, " + f"got {schema.weight_block_size!r}." + ) + self.weight_block_size = weight_block_size + self.weight_scale_key = schema.weight_scale_key + + def get_tensors_attrs( + self, + shape_n: int, + shape_k: int, + param_dtype: torch.dtype, + num_experts: int | None = None, + has_bias: bool = False, + stack_size: int = 1, + ) -> dict[str, dict[str, Any]]: + tensors_attrs = self.schema.get_tensors_attrs( + shape_n=shape_n, + shape_k=shape_k, + param_dtype=param_dtype, + num_experts=num_experts, + has_bias=has_bias, + stack_size=stack_size, + ) + block_n, block_k = self.weight_block_size + scale_shape = (math.ceil(shape_n / block_n), math.ceil(shape_k / block_k)) + if num_experts: + scale_shape = (num_experts,) + scale_shape + tensors_attrs[self.weight_scale_key]["shape"] = scale_shape + return tensors_attrs + + def get_stacked_tensors_attrs( + self, + shape_n_stacks: list[int], + shape_k: int, + param_dtype: torch.dtype, + has_bias: bool = False, + ) -> dict[str, dict[str, Any]]: + tensors_attrs = self.get_tensors_attrs( + shape_n=sum(shape_n_stacks), + shape_k=shape_k, + param_dtype=param_dtype, + has_bias=has_bias, + stack_size=len(shape_n_stacks), + ) + block_n, block_k = self.weight_block_size + tensors_attrs[self.weight_scale_key]["shape"] = ( + sum(math.ceil(shape_n / block_n) for shape_n in shape_n_stacks), + math.ceil(shape_k / block_k), + ) + return tensors_attrs + + def convert_humming( + self, + tensors: dict[str, torch.Tensor], + shape_n_stacks: list[int], + shape_k_stacks: list[int], + param_dtype: torch.dtype, + num_experts: int | None = None, + ): + schema, output_tensors = self.schema.convert_humming( + tensors=tensors, + shape_n_stacks=shape_n_stacks, + shape_k_stacks=shape_k_stacks, + param_dtype=param_dtype, + num_experts=num_experts, + ) + if len(shape_n_stacks) == 1: + return schema, output_tensors + + block_n, _ = self.weight_block_size + scale_rows = [math.ceil(shape_n / block_n) for shape_n in shape_n_stacks] + scale_stacks = tensors[self.weight_scale_key].split(scale_rows, dim=-2) + output_tensors["weight_scale"] = ( + torch.cat( + [ + scale.repeat_interleave(block_n, -2)[..., :shape_n, :] + for scale, shape_n in zip(scale_stacks, shape_n_stacks, strict=True) + ], + dim=-2, + ) + .to(param_dtype) + .contiguous() + ) + return schema, output_tensors + + +def _validate_block_fp8_partition_shape( + layer: torch.nn.Module, + weight_schema, + input_size: int, + output_size: int, + input_size_per_partition: int, + output_partition_sizes: list[int], + skip_block_quant_check: bool = False, +) -> None: + if skip_block_quant_check or not isinstance( + weight_schema, _StackedBlockFp8CheckpointWeightSchema + ): + return + + from sglang.srt.layers.quantization.fp8_utils import validate_fp8_block_shape + + validate_fp8_block_shape( + layer=layer, + input_size=input_size, + output_size=output_size, + input_size_per_partition=input_size_per_partition, + output_partition_sizes=output_partition_sizes, + block_size=list(weight_schema.weight_block_size), + ) + + +def _build_checkpoint_weight_schema(layer_config: dict[str, Any]): + quant_method = layer_config.get("quant_method") + if quant_method is None: + return None + if quant_method == "w4afp8": + return _W4AFp8CheckpointWeightSchema( + group_size=layer_config.get("group_size", 128) + ) + + schema = BaseWeightSchema.from_config(layer_config) + if quant_method == "fp8" and getattr(schema, "weight_block_size", None): + return _StackedBlockFp8CheckpointWeightSchema(schema) + return schema + + class HummingConfig(QuantizationConfig): packed_modules_mapping = {} def __init__(self, full_config: dict[str, Any] | None = None): _lazy_import_humming() self.full_config: dict[str, Any] = full_config or {} + self._w4afp8_config = None + if self.full_config.get("quant_method") == "w4afp8": + from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config + + self._w4afp8_config = W4AFp8Config.from_config(self.full_config) + # W4AFp8Config.from_config() hardcodes group_size and + # weight_block_size; carry the declared values (including an + # explicit null) instead of silently quantizing with the defaults. + # Unsupported values are rejected by checkpoint schema validation. + if "group_size" in self.full_config: + self._w4afp8_config.group_size = self.full_config["group_size"] + if "weight_block_size" in self.full_config: + self._w4afp8_config.weight_block_size = self.full_config[ + "weight_block_size" + ] + # DeepSeek's MLA weight post-processing reads the dense FP8 block size + # from the model-level quantization config before per-layer Humming + # post-processing runs. Keep that checkpoint metadata available here + # as well as on HummingLayerQuantizationConfig. + self.weight_block_size = ( + self._w4afp8_config.weight_block_size + if self._w4afp8_config is not None + else self.full_config.get("weight_block_size") + ) self.is_fp4_experts: bool = False @classmethod @@ -282,9 +556,7 @@ class HummingConfig(QuantizationConfig): layer_config.update(override_config) break - if "quant_method" in layer_config: - return BaseWeightSchema.from_config(layer_config) - return None + return _build_checkpoint_weight_schema(layer_config) def get_layer_input_schema(self, config: dict[str, Any], prefix: str): if self.is_layer_skipped(config, prefix): @@ -299,14 +571,46 @@ class HummingConfig(QuantizationConfig): return BaseInputSchema.from_config(config) return None + def get_checkpoint_configs_for_layer( + self, layer_type: str + ) -> tuple[dict[str, Any], dict[str, Any]]: + if self._w4afp8_config is None: + return self.full_config, self.full_config + + activation_scheme = ( + self._w4afp8_config.moe_activation_scheme + if layer_type == "moe" + else self._w4afp8_config.linear_activation_scheme + ) + fp8_config = { + "quant_method": "fp8", + "activation_scheme": activation_scheme, + "weight_block_size": self._w4afp8_config.weight_block_size, + "ignored_layers": self._w4afp8_config.ignored_layers, + } + weight_config = ( + { + "quant_method": "w4afp8", + "group_size": self._w4afp8_config.group_size, + } + if layer_type == "moe" + else fp8_config + ) + return weight_config, fp8_config + def get_quant_config_for_layer( self, prefix: str, layer_type: str ) -> "HummingLayerQuantizationConfig | None": weight_schema: BaseWeightSchema | None = None force_weight_schema: HummingWeightSchema | None = None - if self.full_config: - weight_schema = self.get_layer_weight_schema(self.full_config, prefix) + checkpoint_weight_config, checkpoint_input_config = ( + self.get_checkpoint_configs_for_layer(layer_type) + ) + if checkpoint_weight_config: + weight_schema = self.get_layer_weight_schema( + checkpoint_weight_config, prefix + ) is_online_quant = False online_quant_config = envs.SGLANG_HUMMING_ONLINE_QUANT_CONFIG.get() or {} @@ -325,8 +629,10 @@ class HummingConfig(QuantizationConfig): input_schema = None force_input_schema = None - if self.full_config: - input_schema = self.get_layer_input_schema(self.full_config, prefix) + if checkpoint_input_config: + input_schema = self.get_layer_input_schema( + checkpoint_input_config, prefix + ) if envs.SGLANG_HUMMING_INPUT_QUANT_CONFIG.get(): quant_config = envs.SGLANG_HUMMING_INPUT_QUANT_CONFIG.get().copy() @@ -510,6 +816,7 @@ class HummingLinearMethod(LinearMethodBase): input_size: int, output_size: int, params_dtype: torch.dtype, + skip_block_quant_check: bool = False, **extra_weight_attrs, ): from sglang.srt.model_loader.weight_utils import default_weight_loader @@ -523,6 +830,16 @@ class HummingLinearMethod(LinearMethodBase): layer.output_partition_sizes = output_partition_sizes layer.extra_weight_attrs = extra_weight_attrs.copy() + _validate_block_fp8_partition_shape( + layer=layer, + weight_schema=self.weight_schema, + input_size=input_size, + output_size=output_size, + input_size_per_partition=input_size_per_partition, + output_partition_sizes=output_partition_sizes, + skip_block_quant_check=skip_block_quant_check, + ) + weight_loader = extra_weight_attrs.get("weight_loader", default_weight_loader) new_weight_loader = self.prepare_weight_loader(layer, weight_loader) extra_weight_attrs["weight_loader"] = new_weight_loader @@ -532,12 +849,23 @@ class HummingLinearMethod(LinearMethodBase): if block_size is not None: layer.weight_block_size = block_size - weight_tensor_attrs = self.weight_schema.get_tensors_attrs( - shape_n=layer.output_partition_sizes_sum, - shape_k=layer.input_size_per_partition, - param_dtype=params_dtype, - stack_size=len(layer.output_partition_sizes), + get_stacked_tensors_attrs = getattr( + self.weight_schema, "get_stacked_tensors_attrs", None ) + if get_stacked_tensors_attrs is not None: + weight_tensor_attrs = get_stacked_tensors_attrs( + shape_n_stacks=layer.output_partition_sizes, + shape_k=layer.input_size_per_partition, + param_dtype=params_dtype, + has_bias=False, + ) + else: + weight_tensor_attrs = self.weight_schema.get_tensors_attrs( + shape_n=layer.output_partition_sizes_sum, + shape_k=layer.input_size_per_partition, + param_dtype=params_dtype, + stack_size=len(layer.output_partition_sizes), + ) input_tensor_attrs = self.input_schema.get_tensors_attrs( shape_k=layer.input_size_per_partition, diff --git a/test/registered/unit/layers/quantization/test_humming_w4afp8_schemas.py b/test/registered/unit/layers/quantization/test_humming_w4afp8_schemas.py new file mode 100644 index 000000000..ceeb07629 --- /dev/null +++ b/test/registered/unit/layers/quantization/test_humming_w4afp8_schemas.py @@ -0,0 +1,176 @@ +"""Unit coverage for Humming W4AFP8 packing and stacked block-FP8 scale shapes.""" + +from __future__ import annotations + +import math +import unittest +from types import SimpleNamespace + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel + +maybe_stub_sgl_kernel() + +from sglang.srt.layers.quantization.humming import ( # noqa: E402 + HummingConfig, + _StackedBlockFp8CheckpointWeightSchema, + _W4AFp8CheckpointWeightSchema, +) + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestW4AFp8CheckpointSchema(CustomTestCase): + def test_packed_tensor_shapes(self): + schema = _W4AFp8CheckpointWeightSchema(group_size=128) + attrs = schema.get_tensors_attrs( + shape_n=4096, + shape_k=6144, + param_dtype=torch.bfloat16, + num_experts=8, + ) + # int8 storage packs two int4 values per byte along K. + self.assertEqual(attrs["weight"]["shape"], (8, 4096, 6144 // 2)) + self.assertEqual(attrs["weight"]["dtype"], torch.int8) + self.assertEqual(attrs["weight_scale_inv"]["shape"], (8, 4096, 6144 // 128)) + self.assertEqual(attrs["weight_scale_inv"]["dtype"], torch.bfloat16) + + def test_group_size_validation(self): + for bad in (None, 0, -128, True, "128", 12.8): + with self.subTest(bad=bad): + with self.assertRaises(ValueError): + _W4AFp8CheckpointWeightSchema(group_size=bad) + + def test_shape_k_must_preserve_groups(self): + schema = _W4AFp8CheckpointWeightSchema(group_size=128) + with self.assertRaises(ValueError): + schema.get_tensors_attrs( + shape_n=64, shape_k=192, param_dtype=torch.bfloat16 + ) + + def test_config_carries_declared_group_size(self): + """HummingConfig must not silently fall back to the default group size. + + W4AFp8Config.from_config() ignores the checkpoint's group_size; the + Humming wrapper is responsible for carrying the declared value into + the per-layer weight config. + """ + config = HummingConfig( + { + "quant_method": "w4afp8", + "group_size": 64, + "weight_block_size": [128, 128], + } + ) + weight_config, _ = config.get_checkpoint_configs_for_layer("moe") + self.assertEqual(weight_config["group_size"], 64) + + default_config = HummingConfig( + {"quant_method": "w4afp8", "weight_block_size": [128, 128]} + ) + weight_config, _ = default_config.get_checkpoint_configs_for_layer("moe") + self.assertEqual(weight_config["group_size"], 128) + + # An explicit null must be carried (and later rejected by schema + # validation), not silently replaced with the default. + null_config = HummingConfig( + { + "quant_method": "w4afp8", + "group_size": None, + "weight_block_size": [128, 128], + } + ) + weight_config, _ = null_config.get_checkpoint_configs_for_layer("moe") + self.assertIsNone(weight_config["group_size"]) + with self.assertRaises(ValueError): + _W4AFp8CheckpointWeightSchema(group_size=weight_config["group_size"]) + + def test_config_carries_declared_weight_block_size(self): + """Checkpoint-declared block geometry must survive config translation. + + Scale shapes and the MLA post-processing both read this value; losing a + non-default declaration quantizes with the wrong block layout. + """ + config = HummingConfig( + { + "quant_method": "w4afp8", + "weight_block_size": [64, 64], + } + ) + self.assertEqual(config.weight_block_size, [64, 64]) + _, fp8_config = config.get_checkpoint_configs_for_layer("moe") + self.assertEqual(fp8_config["weight_block_size"], [64, 64]) + + default_config = HummingConfig({"quant_method": "w4afp8"}) + self.assertEqual(default_config.weight_block_size, [128, 128]) + + +class TestStackedBlockFp8Schema(CustomTestCase): + @staticmethod + def _make_schema(weight_block_size=(128, 128)): + base = SimpleNamespace( + quant_method="fp8", + weight_block_size=list(weight_block_size), + weight_scale_key="weight_scale_inv", + get_tensors_attrs=lambda **kwargs: { + "weight": { + "shape": (kwargs["shape_n"], kwargs["shape_k"]), + "dtype": torch.float8_e4m3fn, + "extra_attrs": {}, + }, + "weight_scale_inv": { + "shape": (), + "dtype": torch.float32, + "extra_attrs": {}, + }, + }, + ) + return _StackedBlockFp8CheckpointWeightSchema(base) + + def test_block_size_validation(self): + for bad in ([128], [128, 0], [128, -1], [128, True], [128.0, 128]): + with self.subTest(bad=bad): + with self.assertRaises(ValueError): + self._make_schema(bad) + + def test_stacked_scale_rows_use_per_stack_ceil(self): + """Unequal output partitions each round up to their own block count. + + A gate of 96 rows and an up of 32 rows both occupy one 128-row scale + block; the stacked checkpoint therefore stores 2 scale rows, not + ceil((96+32)/128) == 1. Reading the naive shape would misalign every + scale after the first stack. + """ + schema = self._make_schema() + attrs = schema.get_stacked_tensors_attrs( + shape_n_stacks=[96, 32], + shape_k=256, + param_dtype=torch.bfloat16, + ) + self.assertEqual(attrs["weight_scale_inv"]["shape"], (2, 2)) + + # Equal partitions that align with the block size collapse to the + # naive shape, so the distinction only shows up on unequal stacks. + aligned = schema.get_stacked_tensors_attrs( + shape_n_stacks=[128, 128], + shape_k=256, + param_dtype=torch.bfloat16, + ) + self.assertEqual(aligned["weight_scale_inv"]["shape"], (2, 2)) + self.assertEqual( + aligned["weight_scale_inv"]["shape"][0], + sum(math.ceil(n / 128) for n in [128, 128]), + ) + + def test_single_tensor_scale_shape_uses_ceil(self): + schema = self._make_schema() + attrs = schema.get_tensors_attrs( + shape_n=96, shape_k=384, param_dtype=torch.bfloat16, num_experts=4 + ) + self.assertEqual(attrs["weight_scale_inv"]["shape"], (4, 1, 3)) + + +if __name__ == "__main__": + unittest.main()