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

[diffusion] feat: support serialized comfy convrot int8 dits (#35994)

Browse Source
This commit is contained in:
Mick
2026-08-23 09:29:59 +08:00
committed by GitHub
parent 7f30d66045
commit a36c0746b9
11 changed files with 253 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/layers/quantization/comfy_fp8.py
+2
View File
@@ -88,6 +88,8 @@ class ComfyFullPrecisionFp8LinearMethod(LinearMethodBase):
class ComfyFp8Config(QuantizationConfig):
"""Dispatch each Linear according to its serialized ``comfy_quant`` marker."""
checkpoint_uses_native_qkv_layout = True
def __init__(self, layer_markers: dict[str, dict[str, Any]]) -> None:
super().__init__()
self.layer_markers = layer_markers
python/sglang/multimodal_gen/runtime/layers/quantization/configs/base_config.py
+1
View File
@@ -67,6 +67,7 @@ class QuantizationConfig(ABC):
# for quantization frameworks with a separate quantized model provided, e.g. Nunchaku
quantized_model_path: str | None = None
checkpoint_uses_native_qkv_layout: bool = False
def __init__(self):
super().__init__()
python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_int8_config.py
+48 -13
View File
@@ -1,11 +1,5 @@
# SPDX-License-Identifier: Apache-2.0
"""Config for online INT8 ConvRot quantization via comfy_kitchen.
A no-arg ``KitchenInt8Config()`` is the only supported form: weights load in
their source dtype and are quantized in ``process_weights_after_loading``.
Registered CLI name: ``kitchen_int8``.
"""
"""Config for online or serialized INT8 ConvRot via comfy_kitchen."""
from __future__ import annotations
@@ -27,17 +21,14 @@ _SUPPORTED_GROUP_SIZES = (16, 64, 256)
class KitchenInt8Config(QuantizationConfig):
"""Config for online INT8 ConvRot quantization via comfy_kitchen.
A no-arg ``KitchenInt8Config()`` is the only supported form: weights load in
their source dtype and are quantized in ``process_weights_after_loading``.
"""
"""Dispatch online quantization or serialized Comfy ConvRot layers."""
def __init__(
self,
group_size: int = 256,
ignored_layers: list[str] | None = None,
packed_modules_mapping: dict[str, list[str]] | None = None,
layer_markers: dict[str, dict[str, Any]] | None = None,
) -> None:
super().__init__()
if group_size not in _SUPPORTED_GROUP_SIZES:
@@ -48,6 +39,30 @@ class KitchenInt8Config(QuantizationConfig):
self.group_size = group_size
self.ignored_layers = ignored_layers or []
self.packed_modules_mapping = packed_modules_mapping or {}
self.layer_markers = layer_markers
self.is_checkpoint_int8_serialized = layer_markers is not None
self.checkpoint_uses_native_qkv_layout = self.is_checkpoint_int8_serialized
self._serialized_group_sizes: dict[str, int] = {}
if layer_markers is not None:
for prefix, marker in layer_markers.items():
if marker.get("format") != "int8_tensorwise":
raise ValueError(
f"Unsupported Comfy INT8 format for {prefix!r}: "
f"{marker.get('format')!r}"
)
if marker.get("convrot") is not True:
raise ValueError(
f"Serialized kitchen_int8 layer {prefix!r} must set "
"convrot=true"
)
marker_group_size = marker.get("convrot_groupsize")
if marker_group_size not in _SUPPORTED_GROUP_SIZES:
raise ValueError(
f"Serialized kitchen_int8 layer {prefix!r} must declare "
f"convrot_groupsize in {_SUPPORTED_GROUP_SIZES}, got "
f"{marker_group_size!r}"
)
self._serialized_group_sizes[prefix] = marker_group_size
# Which layers actually got quantized is worth stating plainly in the
# log: a silent fallback to BF16 looks exactly like a slow kernel.
self.selected: list[str] = []
@@ -89,6 +104,22 @@ class KitchenInt8Config(QuantizationConfig):
if not isinstance(layer, LinearBase):
return None
if self.layer_markers is not None:
marker_group_size = self._serialized_group_sizes.get(prefix)
if marker_group_size is None:
return UnquantizedLinearMethod()
if layer.input_size % marker_group_size:
raise ValueError(
f"Serialized kitchen_int8 layer {prefix!r} has input size "
f"{layer.input_size}, which is not divisible by its "
f"ConvRot group size {marker_group_size}"
)
self.selected.append(prefix)
return KitchenInt8LinearMethod(
self,
group_size=marker_group_size,
is_checkpoint_serialized=True,
)
if is_layer_skipped(
prefix, self.ignored_layers, fused_mapping=self.packed_modules_mapping
):
@@ -102,7 +133,11 @@ class KitchenInt8Config(QuantizationConfig):
self.skipped.append(f"{prefix}(in={layer.input_size})")
return UnquantizedLinearMethod()
self.selected.append(prefix)
return KitchenInt8LinearMethod(self)
return KitchenInt8LinearMethod(
self,
group_size=self.group_size,
is_checkpoint_serialized=False,
)
def note_quantized(self, saved_bytes: int) -> None:
self._processed += 1
python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_int8.py
+36 -18
View File
@@ -8,10 +8,9 @@ The difference is that it is a single fused op -- it takes a BF16 activation and
does the Hadamard rotation, dynamic per-row activation quantization, IMMA GEMM,
dequantization and bias add without ever materializing the intermediates.
Quantization is data-free (group-wise Hadamard rotation + per-output-channel
absmax), so weights are quantized here after loading rather than read from a
pre-quantized checkpoint. That keeps this usable with the stock BF16 checkpoint
and avoids depending on any external file layout.
The online path applies data-free group-wise Hadamard rotation and per-output
channel scaling after loading a stock BF16 checkpoint. Compatible serialized
Comfy checkpoints instead load their INT8 weights and row scales directly.
"""
from __future__ import annotations
@@ -73,10 +72,18 @@ def _load_comfy_kitchen():
class KitchenInt8LinearMethod(LinearMethodBase):
"""Quantizes BF16 weights to INT8 after load and runs the fused kernel."""
"""Loads or creates ConvRot INT8 weights and runs the fused kernel."""
def __init__(self, quant_config: KitchenInt8Config) -> None:
def __init__(
self,
quant_config: KitchenInt8Config,
*,
group_size: int,
is_checkpoint_serialized: bool,
) -> None:
self.quant_config = quant_config
self.group_size = group_size
self.is_checkpoint_serialized = is_checkpoint_serialized
_load_comfy_kitchen()
def create_weights(
@@ -92,36 +99,47 @@ class KitchenInt8LinearMethod(LinearMethodBase):
# get_quant_method already screened the unsharded input size, so this
# only fires under TP > 1, where a row-parallel layer splits the very
# dimension the rotation groups over.
if input_size_per_partition % self.quant_config.group_size:
if input_size_per_partition % self.group_size:
raise ValueError(
f"kitchen_int8 needs input_size_per_partition "
f"({input_size_per_partition}) divisible by group_size "
f"{self.quant_config.group_size}"
f"{self.group_size}"
)
# Deliberately identical to UnquantizedLinearMethod: weights load as
# BF16 through the model's existing loaders (H3 for instance installs a
# custom qkv loader that reorders the grouped checkpoint layout), and
# only then get replaced by their quantized form.
# The online path initially matches UnquantizedLinearMethod so the
# source weights load in BF16 before quantization. Serialized weights
# allocate their final INT8 storage immediately.
weight = Parameter(
torch.empty(
sum(output_partition_sizes),
input_size_per_partition,
dtype=params_dtype,
dtype=(torch.int8 if self.is_checkpoint_serialized else params_dtype),
),
requires_grad=False,
)
set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
layer.register_parameter("weight", weight)
set_weight_attrs(weight, extra_weight_attrs)
if self.is_checkpoint_serialized:
weight_scale = Parameter(
torch.empty(
sum(output_partition_sizes),
1,
dtype=torch.float32,
),
requires_grad=False,
)
set_weight_attrs(weight_scale, {"output_dim": 0})
set_weight_attrs(weight_scale, extra_weight_attrs)
layer.register_parameter("weight_scale", weight_scale)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
from comfy_kitchen.tensor.int8 import TensorWiseINT8Layout
weight = layer.weight.data
if weight.dtype == torch.int8: # already processed
if self.is_checkpoint_serialized or weight.dtype == torch.int8:
return
from comfy_kitchen.tensor.int8 import TensorWiseINT8Layout
# Quantization runs on CUDA, but the model may still be staged on CPU
# for offload. Round-trip one layer at a time rather than relying on
# the loader's whole-model device move, which would not fit in VRAM.
@@ -131,7 +149,7 @@ class KitchenInt8LinearMethod(LinearMethodBase):
is_weight=True,
per_channel=True,
convrot=True,
convrot_groupsize=self.quant_config.group_size,
convrot_groupsize=self.group_size,
stochastic_rounding=0,
)
layer.weight = Parameter(qdata.to(home), requires_grad=False)
@@ -171,7 +189,7 @@ class KitchenInt8LinearMethod(LinearMethodBase):
bias,
out_code,
True, # convrot
self.quant_config.group_size,
self.group_size,
)
n_rows, n_out = x.shape[0], layer.weight.shape[0]
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+8 -5
View File
@@ -280,10 +280,10 @@ class TransformerLoader(ComponentLoader):
or cpu_offload_flag
)
use_fsdp = server_args.should_use_fsdp_for_component(component_name)
if quant_spec.is_comfy_fp8 and use_fsdp:
if quant_spec.uses_comfy_layer_markers and use_fsdp:
raise ValueError(
"MiniMax-H3 Comfy FP8 does not support FSDP inference; use TP "
"and/or sequence parallelism instead"
"Comfy quantized checkpoints do not support FSDP "
"inference; use TP and/or sequence parallelism instead"
)
if quant_spec.gguf_file is not None:
@@ -308,7 +308,7 @@ class TransformerLoader(ComponentLoader):
"quant_config": quant_spec.runtime_quant_config,
}
checkpoint_key_filter: Callable[[str], bool] | None = (
comfy_quant_key_filter if quant_spec.is_comfy_fp8 else None
comfy_quant_key_filter if quant_spec.uses_comfy_layer_markers else None
)
adaln_cache_path = component_server_args.minimax_h3_adaln_cache_path
if adaln_cache_path is not None:
@@ -362,7 +362,10 @@ class TransformerLoader(ComponentLoader):
local_torch_device,
component_starts_on_cpu=component_starts_on_cpu,
runtime_quant_config=quant_spec.runtime_quant_config,
quantized_cpu_load_supported=quant_spec.gguf_file is not None,
quantized_cpu_load_supported=(
quant_spec.gguf_file is not None
or quant_spec.is_serialized_kitchen_int8
),
)
)
direct_gpu_weight_loading = bool(
python/sglang/multimodal_gen/runtime/loader/minimax_h3_weights.py
+41 -16
View File
@@ -10,6 +10,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import ComfyFp8
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
QuantizationConfig,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import (
KitchenInt8Config,
)
def comfy_quant_key_filter(name: str) -> bool:
@@ -23,6 +26,7 @@ def inspect_minimax_h3_safetensors(
adaln_curve_shape = None
layer_markers: dict[str, dict[str, Any]] = {}
checkpoint_keys: set[str] = set()
checkpoint_meta: dict[str, tuple[str, tuple[int, ...]]] = {}
fp8_weight_prefixes: set[str] = set()
for path in safetensors_list:
@@ -44,11 +48,15 @@ def inspect_minimax_h3_safetensors(
adaln_curve_shape = shape
for key in keys:
if (
key.endswith(".weight")
and checkpoint.get_slice(key).get_dtype() == "F8_E4M3"
):
fp8_weight_prefixes.add(key.removesuffix(".weight"))
if key.endswith((".weight", ".weight_scale")):
tensor_slice = checkpoint.get_slice(key)
dtype = tensor_slice.get_dtype()
checkpoint_meta[key] = (
dtype,
tuple(tensor_slice.get_shape()),
)
if key.endswith(".weight") and dtype == "F8_E4M3":
fp8_weight_prefixes.add(key.removesuffix(".weight"))
if not key.endswith(".comfy_quant"):
continue
try:
@@ -78,17 +86,39 @@ def inspect_minimax_h3_safetensors(
)
for prefix, marker in layer_markers.items():
if marker.get("format") != "float8_e4m3fn":
continue
marker_format = marker.get("format")
required = {f"{prefix}.weight", f"{prefix}.weight_scale"}
if not marker.get("full_precision_matrix_mult", False):
if marker_format == "float8_e4m3fn" and not marker.get(
"full_precision_matrix_mult", False
):
required.add(f"{prefix}.input_scale")
if marker_format not in ("float8_e4m3fn", "int8_tensorwise"):
continue
missing = required - checkpoint_keys
if missing:
raise ValueError(
f"MiniMax-H3 Comfy FP8 layer {prefix!r} is missing checkpoint "
f"MiniMax-H3 Comfy layer {prefix!r} is missing checkpoint "
f"tensors: {sorted(missing)}"
)
if marker_format == "int8_tensorwise":
weight_dtype, weight_shape = checkpoint_meta[f"{prefix}.weight"]
scale_dtype, scale_shape = checkpoint_meta[f"{prefix}.weight_scale"]
if weight_dtype != "I8" or scale_dtype != "F32":
raise ValueError(
f"MiniMax-H3 Comfy INT8 layer {prefix!r} needs I8 weights "
f"and F32 scales, got {weight_dtype} and {scale_dtype}"
)
if len(weight_shape) != 2:
raise ValueError(
f"MiniMax-H3 Comfy INT8 layer {prefix!r} needs a 2D weight, "
f"got {weight_shape}"
)
expected_scale_shape = (weight_shape[0], 1)
if scale_shape != expected_scale_shape:
raise ValueError(
f"MiniMax-H3 Comfy INT8 layer {prefix!r} needs scale shape "
f"{expected_scale_shape}, got {scale_shape}"
)
return adaln_curve_shape, layer_markers
@@ -100,13 +130,8 @@ def resolve_minimax_h3_checkpoint_quantization(
return None
formats = sorted({str(marker.get("format")) for marker in layer_markers.values()})
if "int8_tensorwise" in formats:
raise NotImplementedError(
"MiniMax-H3 pruned_int8_convrot is not supported yet. Its "
"int8_tensorwise weights require an online regular-Hadamard ConvRot "
"and dynamic INT8 activation quantization kernel; loading them as "
"ordinary INT8/BF16 weights would produce incorrect output."
)
if formats == ["int8_tensorwise"]:
return KitchenInt8Config(layer_markers=layer_markers)
if formats == ["float8_e4m3fn"]:
return ComfyFp8Config(layer_markers)
raise NotImplementedError(
python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py
+17
View File
@@ -18,6 +18,9 @@ from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME
from torch import nn
from sglang.multimodal_gen.runtime.layers.quantization import QuantizationConfig
from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import (
KitchenInt8Config,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
NunchakuConfig,
_patch_nunchaku_scales,
@@ -164,6 +167,17 @@ class TransformerQuantLoadSpec:
def is_comfy_fp8(self) -> bool:
return _get_quant_config_name(self.quant_config) == "comfy_fp8"
@property
def is_serialized_kitchen_int8(self) -> bool:
return (
isinstance(self.quant_config, KitchenInt8Config)
and self.quant_config.is_checkpoint_int8_serialized
)
@property
def uses_comfy_layer_markers(self) -> bool:
return self.is_comfy_fp8 or self.is_serialized_kitchen_int8
class _TransformerQuantAdapter:
def prepare(self) -> None:
@@ -844,6 +858,9 @@ def _needs_device_weight_postprocess(
quant_name = _get_quant_config_name(quant_config)
if quant_name in ("modelopt_fp8", "comfy_fp8"):
return True
if quant_name == "kitchen_int8":
assert isinstance(quant_config, KitchenInt8Config)
return not quant_config.is_checkpoint_int8_serialized
serialized_flag_by_quant_name = {
"fp8": "is_checkpoint_fp8_serialized",
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+7 -4
View File
@@ -612,10 +612,13 @@ class MiniMaxH3Attention(nn.Module):
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
# The reorder below translates the *safetensors* checkpoint layout. A
# GGUF checkpoint already stores qkv as [q_all, k_all, v_all], and its
# packed parameter is `qweight`, so there is nothing to reorder.
if quant_config is None or quant_config.get_name() != "gguf":
# Official safetensors interleave Q/K/V by head. Comfy and GGUF
# checkpoints already store [q_all, k_all, v_all].
checkpoint_qkv_is_native = quant_config is not None and (
quant_config.get_name() == "gguf"
or quant_config.checkpoint_uses_native_qkv_layout
)
if not checkpoint_qkv_is_native:
self._install_qkv_weight_loader(arch)
self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
python/sglang/multimodal_gen/test/unit/test_transformer_quant.py
+64 -4
View File
@@ -47,12 +47,16 @@ sys.modules.setdefault("partial_json_parser.core.options", partial_json_parser_o
from sglang.multimodal_gen.runtime.layers.linear import (
LinearBase,
ReplicatedLinear,
UnquantizedLinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.comfy_fp8 import (
ComfyFp8Config,
ComfyFullPrecisionFp8LinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import (
KitchenInt8Config,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config import (
NunchakuConfig,
)
@@ -246,11 +250,19 @@ class TestTransformerQuantHelpers(unittest.TestCase):
mock_download.reset_mock()
def test_inspect_minimax_h3_safetensors_detects_curve_and_comfy_format(self):
marker = json.dumps({"format": "int8_tensorwise", "convrot": True}).encode()
marker = json.dumps(
{
"format": "int8_tensorwise",
"convrot": True,
"convrot_groupsize": 256,
}
).encode()
with tempfile.NamedTemporaryFile(suffix=".safetensors") as f:
save_file(
{
"adaln_t_table": torch.zeros((1025, 8)),
"blocks.0.mlp.fc1.weight": torch.ones((2, 256), dtype=torch.int8),
"blocks.0.mlp.fc1.weight_scale": torch.ones((2, 1)),
"blocks.0.mlp.fc1.comfy_quant": torch.tensor(
list(marker), dtype=torch.uint8
),
@@ -282,13 +294,60 @@ class TestTransformerQuantHelpers(unittest.TestCase):
self.assertEqual(layer_markers["blocks.0.mlp.fc1"], {"format": "float8_e4m3fn"})
def test_minimax_h3_comfy_int8_fails_before_weight_loading(self):
with self.assertRaisesRegex(NotImplementedError, "regular-Hadamard"):
def test_minimax_h3_comfy_int8_resolves_serialized_kitchen(self):
config = resolve_minimax_h3_checkpoint_quantization(
{
"blocks.0.mlp.fc1": {
"format": "int8_tensorwise",
"convrot": True,
"convrot_groupsize": 256,
}
}
)
self.assertIsInstance(config, KitchenInt8Config)
self.assertTrue(config.is_checkpoint_int8_serialized)
self.assertTrue(config.checkpoint_uses_native_qkv_layout)
self.assertFalse(KitchenInt8Config().checkpoint_uses_native_qkv_layout)
self.assertFalse(_needs_device_weight_postprocess(config))
@patch(
"sglang.multimodal_gen.runtime.layers.quantization.kitchen_int8."
"_load_comfy_kitchen"
)
def test_serialized_kitchen_constructs_int8_weight_and_row_scale(self, _load):
config = KitchenInt8Config(
layer_markers={
"proj": {
"format": "int8_tensorwise",
"convrot": True,
"convrot_groupsize": 256,
}
}
)
layer = ReplicatedLinear(
256,
3,
bias=False,
params_dtype=torch.bfloat16,
quant_config=config,
prefix="proj",
)
self.assertEqual(layer.weight.dtype, torch.int8)
self.assertEqual(layer.weight.shape, (3, 256))
self.assertEqual(layer.weight_scale.dtype, torch.float32)
self.assertEqual(layer.weight_scale.shape, (3, 1))
def test_serialized_kitchen_rejects_non_convrot_marker(self):
with self.assertRaisesRegex(ValueError, "convrot=true"):
resolve_minimax_h3_checkpoint_quantization(
{
"blocks.0.mlp.fc1": {
"format": "int8_tensorwise",
"convrot": True,
"convrot": False,
"convrot_groupsize": 256,
}
}
)
@@ -305,6 +364,7 @@ class TestTransformerQuantHelpers(unittest.TestCase):
)
self.assertIsInstance(config, ComfyFp8Config)
self.assertTrue(config.checkpoint_uses_native_qkv_layout)
layer = LinearBase(input_size=1, output_size=1)
self.assertIsInstance(
config.get_quant_method(layer, "blocks.0.mlp.fc2"),