diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index d9c432aa7..fe3c14808 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -68,7 +68,9 @@ H200, and H100. Resident is the latency-oriented default; FSDP reduces DiT weight residency at the cost of per-block parameter collectives. On H200 it also selects the verified 2-node cross-node topology. **Online Quantization** appears only on B200 and B300. AMD keeps its resident AITER -recipe, while RTX 5090 uses its dedicated layerwise-offload profile. +recipe, while RTX 5090 uses its dedicated layerwise-offload profile. A +single 24 GB card (RTX 4090) uses the same offload knobs plus online +`kitchen_int8`; that recipe is documented below rather than in the picker. import { Deployment } from "/src/snippets/_deployment.jsx"; import { config } from "/src/snippets/configs/MiniMaxAI/minimax-h3.jsx"; @@ -168,6 +170,34 @@ transfer overhead. This exact recipe was validated on 2× RTX 5090 (32 GB each) and a 377 GiB host; use a 384 GiB-class machine. The latency and memory comparison is collected in the benchmark section below. +For a single 24 GB consumer card (RTX 4090), stream the DiT and text encoder +and quantize DiT linear layers online with `kitchen_int8`. Keep `vae` out of +`--layerwise-offload-components`: putting the VAE decoder in layerwise +offload re-streams about 9 GiB on each of 167 decode tiles. Default +attention stays `fa` (exact). Approximate backends are opt-in; see +[Attention Backends](/docs/sglang-diffusion/attention_backends#sage-then-sol-hybrid). +Install `comfy-kitchen` first (`pip install comfy-kitchen`). + +```bash 1×RTX 4090 24GB +sglang generate \ + --model-path MiniMaxAI/MiniMax-H3 \ + --model-variant fl2va \ + --quantization kitchen_int8 \ + --attention-backend fa \ + --performance-mode memory \ + --layerwise-offload-components dit,text_encoder \ + --dit-offload-prefetch-size 1 \ + --dit-layerwise-resident-layers 0 \ + --enable-torch-compile false \ + --prompt "A cat walking on a sunny beach, gentle waves." \ + --save-output +``` + +The same flags work on `sglang serve`. Drop `--quantization` for the BF16 +baseline; everything else stays identical. GPU peak stays about 18 GB +either way because streaming offload is set by the offload buffers and VAE +decode, not the weight dtype. + The first launch downloads the model through the selected Hub. If the Hugging Face repository requires authentication, export a Hugging Face token in the server environment. @@ -668,7 +698,7 @@ listed hardware and topology; it is not inherited by a similar GPU family. | Tensor parallelism | Verified: B200 TP2 + Ulysses4; H100 TP2 + Ulysses2 and TP4 + Ulysses1 | `--tp-size` may be combined with Ulysses when the TP-local head count remains divisible by the Ulysses degree. On 4×H100, TP2 + Ulysses2 is the measured speed default. | | FSDP inference | Verified: 4× B200 and 4× H100 + Ulysses4 | Preserves H3's mixed BF16/FP32 parameter policy. B200 completed the exact eager comparison; H100 completed consecutive real requests at about 57 GB peak memory per GPU. | | Resident components | Verified: B200, H200, 4×H100 with TP, and 1/2/4/8× MI300X and MI355X | This is the recommended single-request latency path when the complete workload fits. | -| CPU and layerwise offload | Verified: 2× RTX 5090 TP2 | The measured lossless recipe keeps 20 DiT blocks plus both VAE encoders resident, streams the remaining DiT blocks, text encoder, and video VAE decoder blocks, and leaves the small audio VAE resident. This status applies only to the listed topology. | +| CPU and layerwise offload | Verified: 2× RTX 5090 TP2; 1× RTX 4090 24 GB | The 5090 lossless recipe keeps 20 DiT blocks plus both VAE encoders resident, streams the remaining DiT blocks, text encoder, and video VAE decoder blocks, and leaves the small audio VAE resident. The 4090 recipe streams DiT and the text encoder with zero resident DiT layers and **omits `vae`** from `--layerwise-offload-components`. | | Breakable CUDA graph | Verified: B200 Ref2VA, opt-in | Matching eager output was observed for the captured signature, without a measured speedup. Re-capture for other shapes and reference sets. | | `torch.compile` | Measured: H200, opt-in | Steady-state benefit was below measurement noise, while startup increased and numerical output changed. Do not use it for consistency ground truth. | @@ -738,6 +768,40 @@ The picker exposes this option only on the B200 and B300 topologies used for real H3 validation runs. +On a single 24 GB card, use `kitchen_int8` instead of FP8. It quantizes the +four GEMMs per DiT block online from the Hub BF16 weights (data-free, no +calibration) and dispatches them through `comfy_kitchen.int8_linear`. +Quantization happens after H3's grouped `qkv` reorder, so do not load an +externally pre-quantized INT8 checkpoint here. + +```bash Command +sglang generate \ + --model-path MiniMaxAI/MiniMax-H3 \ + --model-variant fl2va \ + --quantization kitchen_int8 \ + --attention-backend fa \ + --performance-mode memory \ + --layerwise-offload-components dit,text_encoder \ + --dit-offload-prefetch-size 1 \ + --dit-layerwise-resident-layers 0 \ + --enable-torch-compile false \ + --prompt "A cat walking on a sunny beach, gentle waves." \ + --save-output +``` + +`fa` keeps exact attention. For a faster, approximate DiT path, use +`--attention-backend sol_attn` with +`--attention-backend-config dense_backend=sage_attn,dense_steps=10` and +`--component-attention-backends text_encoder=torch_sdpa,transformer=sol_attn`. +See [Quantization](/docs/sglang-diffusion/quantization#kitchen-int8-online-quantization) +and [Attention Backends](/docs/sglang-diffusion/attention_backends#sage-then-sol-hybrid). + + +`kitchen_int8` changes Linear numerics. `sol_attn` / `sage_attn` also change +the attention algorithm. Neither is a consistency ground-truth mode. The +BF16 path is unchanged when `comfy-kitchen` is not installed. + + The Qwen3-VL text encoder can be replaced independently of the DiT. To reduce its resident memory, point the text-encoder component at the serialized FP8 checkpoint used in validation: @@ -799,6 +863,7 @@ the configurations with collected measurements: | H100 | 4× TP2 + Ulysses2 resident | 4× TP4 + Ulysses1; 4× FSDP + Ulysses4 | | MI300X / MI355X | 8× Ulysses8 resident | 1×, 2×, and 4× scaling runs | | RTX 5090 | 2× TP2 + layerwise offload | — | +| RTX 4090 24 GB | 1× layerwise offload + `kitchen_int8` | Approximate attention backends are opt-in | ### B300 precision and encoder placement @@ -991,6 +1056,27 @@ peak per GPU. | prefetch 2, resident 20 | 43.37 s | 78.06 s | 27.5 GiB | No measurable gain | | Ulysses2, prefetch 2, resident 10 | Did not reach warmup | — | — | Rejected | +### RTX 4090 24 GB single-GPU run + +One RTX 4090 D 24 GB completed the 1344×768, 107-frame, 20-NFE T2VA +workload (euler, `torch.compile` and step caching disabled) with DiT and +text-encoder layerwise offload. Same process: load → warmup (seed 0) → +timed (seed 42); only the timed pass is reported. GPU peak stayed about +18 GB. + +| Config | Timed e2e | Denoise | vs BF16 | PSNR vs BF16 | +| --- | ---: | ---: | ---: | ---: | +| BF16 + FlashAttention | 405.6 s | 370.2 s | 1.00× | — | +| `kitchen_int8` + FA | 303.3 s | 273.7 s | 1.34× | 24.81 dB | +| `kitchen_int8` + `sol_attn` | 223.9 s | 203.9 s | 1.81× | 24.44 dB | +| `kitchen_int8` + `sage_attn` | 174.9 s | 154.2 s | 2.32× | 23.51 dB | +| `kitchen_int8` + Sage→Sol hybrid | 163.8 s | 143.1 s | 2.48× | 23.04 dB | + +`kitchen_int8` + FA changes Linear numerics only. The `sol_attn` / +`sage_attn` / hybrid rows also change the attention algorithm, so speed +and pixel fidelity rank in opposite orders there. Default remains +`kitchen_int8` + `fa`. + ### AMD Instinct task and scaling runs The AMD recipes keep the released BF16/FP32 precision policy and use AITER diff --git a/docs/docs/sglang-diffusion/attention_backends.mdx b/docs/docs/sglang-diffusion/attention_backends.mdx index 0e750f8d9..e0be8b6f5 100644 --- a/docs/docs/sglang-diffusion/attention_backends.mdx +++ b/docs/docs/sglang-diffusion/attention_backends.mdx @@ -398,6 +398,12 @@ Some backends require additional configuration. You can pass these parameters vi Layer indices kept dense, e.g. `0,1` or `0-2`. `0,1` + + `dense_backend` + `str` + Backend used for the dense prefix: `fa` (default) or `sage_attn`. `sage_attn` is approximate. + `fa` + `kv_splits` `int | str` @@ -626,6 +632,25 @@ sglang generate \ Component keys match pipeline module names from `model_index.json`, such as `text_encoder`, `text_encoder_2`, `transformer`, `transformer_2`, or `connectors`. +### Sage then Sol hybrid + +`sol_attn` keeps the first `dense_steps` steps dense. Set +`dense_backend=sage_attn` to run that prefix on SageAttention and the tail on +Sol sparse attention. Keep the text encoder on `torch_sdpa`: + +```bash +sglang generate \ + --model-path MiniMaxAI/MiniMax-H3 \ + --model-variant fl2va \ + --attention-backend sol_attn \ + --attention-backend-config dense_backend=sage_attn,dense_steps=10 \ + --component-attention-backends text_encoder=torch_sdpa,transformer=sol_attn \ + --prompt "A cat walking on a sunny beach, gentle waves." \ + --save-output +``` + +Both `sage_attn` and `sol_attn` are approximate. The default DiT backend remains `fa`. + ### Using Sliding Tile Attention (STA) ```bash diff --git a/docs/docs/sglang-diffusion/environment_variables.mdx b/docs/docs/sglang-diffusion/environment_variables.mdx index 6f74b92d2..44dde40f8 100644 --- a/docs/docs/sglang-diffusion/environment_variables.mdx +++ b/docs/docs/sglang-diffusion/environment_variables.mdx @@ -88,6 +88,16 @@ description: "Configure SGLang diffusion behavior with environment variables." true Use Run:AI model streamer for model loading + + SGLANG_KITCHEN_INT8_MAX_ROWS + 8192 + Max activation rows per kitchen_int8 fused GEMM. Set 0 to disable row splitting. + + + SGLANG_KITCHEN_INT8_MIN_SPLIT_N + 8192 + Minimum output features before kitchen_int8 splits large-M GEMMs. Narrow outputs stay as a single call. + diff --git a/docs/docs/sglang-diffusion/quantization.mdx b/docs/docs/sglang-diffusion/quantization.mdx index 94449742c..af4923f61 100644 --- a/docs/docs/sglang-diffusion/quantization.mdx +++ b/docs/docs/sglang-diffusion/quantization.mdx @@ -79,6 +79,14 @@ backend. MXFP4: aiter on ROCm MXFP4 requires ROCm and MI350+ (gfx95x). Weights quantized at load time, activations quantized to fp8 / mxfp4 dynamically. + + kitchen_int8 (online quantization) + Unquantized BF16/FP16 checkpoint + --quantization kitchen_int8 + MiniMax-H3 (validated on 1× RTX 4090 24 GB) + comfy-kitchen + Data-free INT8 ConvRot at load time via comfy_kitchen.int8_linear. Layers whose input dim is not divisible by the group size stay in BF16. Requires Turing+ (SM75). + fp8 (offline quantization) Quantized transformer component folder, or safetensors with quantization_config metadata @@ -260,6 +268,50 @@ sglang generate \ ``` **Note:** Requires `aiter` package with MXFP4 kernel support +### Kitchen INT8 Online Quantization + +`kitchen_int8` quantizes DiT linear weights online from the stock BF16 +checkpoint. Forward uses the fused `comfy_kitchen.int8_linear` op (rotation, +dynamic per-row activation quantization, INT8 GEMM, dequant, and bias). +Install the optional dependency first: + +```bash +pip install comfy-kitchen +``` + +```bash +sglang generate \ + --model-path MiniMaxAI/MiniMax-H3 \ + --model-variant fl2va \ + --quantization kitchen_int8 \ + --attention-backend fa \ + --performance-mode memory \ + --layerwise-offload-components dit,text_encoder \ + --dit-offload-prefetch-size 1 \ + --dit-layerwise-resident-layers 0 \ + --enable-torch-compile false \ + --prompt "A cat walking on a sunny beach, gentle waves." \ + --save-output +``` + +Quantization runs after the model's weight loaders, so MiniMax-H3's grouped +`qkv` reorder is already applied. Layers whose input dim is not divisible by +the group size (256) stay in BF16 instead of failing the load; H3's AdaLN +projections take that path. + + +`kitchen_int8` is approximate and is not a consistency ground-truth mode. +The BF16 path is unchanged when `comfy-kitchen` is not installed. See the +[MiniMax-H3 cookbook](/cookbook/diffusion/MiniMax/MiniMax-H3#7-runtime-feature-recipes) +for the 24 GB offload recipe, including why `vae` must stay out of +`--layerwise-offload-components`. + + +Large-M GEMMs (`rows > 8192` and `out_features >= 8192`) are split by rows so +the fused kernel stays on the data-parallel CUTLASS config. Override the +thresholds with `SGLANG_KITCHEN_INT8_MAX_ROWS` and +`SGLANG_KITCHEN_INT8_MIN_SPLIT_N`. + ### Skipping Layers By default, online quantization quantizes every linear layer in diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py index a63057a87..b7be98848 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/sol_attn.py @@ -2,6 +2,7 @@ from __future__ import annotations +import inspect import re import torch @@ -19,6 +20,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) _SOL_ATTN_HEAD_DIM = 128 +_DENSE_BACKENDS = {"fa", "sage_attn"} def _parse_layer_ranges(spec: str | int | None) -> frozenset[int]: @@ -57,7 +59,16 @@ def _resolve_kv_splits(q: torch.Tensor, kv_splits: int | str | None) -> int: def _get_sol_attn_runtime_config() -> dict: server_args = get_global_server_args() cfg = getattr(server_args, "attention_backend_config", None) or {} - dense_layers = cfg.get("dense_layers", "0,1") + dense_backend = ( + str(cfg.get("dense_backend", "fa")).strip().lower().replace("-", "_") + ) + if dense_backend in {"sage", "sageattention"}: + dense_backend = "sage_attn" + if dense_backend not in _DENSE_BACKENDS: + raise ValueError( + f"Unsupported sol_attn dense_backend={dense_backend!r}; " + f"expected one of {sorted(_DENSE_BACKENDS)}" + ) sink_start = cfg.get("sink_start", 0) return { "tau": float(cfg.get("tau", 1.0)), @@ -66,7 +77,8 @@ def _get_sol_attn_runtime_config() -> dict: "sink_tokens": int(cfg.get("sink_tokens", 0)), "sink_start": None if sink_start is None else int(sink_start), "dense_steps": int(cfg.get("dense_steps", 10)), - "dense_layers": _parse_layer_ranges(dense_layers), + "dense_layers": _parse_layer_ranges(cfg.get("dense_layers", "0,1")), + "dense_backend": dense_backend, } @@ -107,13 +119,12 @@ class SolAttnImpl(AttentionImpl): self.softmax_scale = softmax_scale self.prefix = prefix self.layer_idx = self._parse_layer_idx(prefix) + self._sol_params: frozenset[str] | None = None @staticmethod def _parse_layer_idx(prefix: str) -> int | None: match = re.search(r"blocks\.(\d+)", prefix) - if match is None: - return None - return int(match.group(1)) + return int(match.group(1)) if match else None def _should_use_dense(self) -> bool: cfg = _get_sol_attn_runtime_config() @@ -125,13 +136,11 @@ class SolAttnImpl(AttentionImpl): step = int(get_forward_context().current_timestep) except AssertionError: step = 0 - if step < cfg["dense_steps"]: - return True - if self.layer_idx is not None and self.layer_idx in cfg["dense_layers"]: - return True - return False + return step < cfg["dense_steps"] or ( + self.layer_idx is not None and self.layer_idx in cfg["dense_layers"] + ) - def _dense_varlen( + def _dense_fa( self, query: torch.Tensor, key: torch.Tensor, @@ -153,6 +162,46 @@ class SolAttnImpl(AttentionImpl): ) return output[0] if isinstance(output, tuple) else output + def _dense_sage( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + *, + cu_seqlens: torch.Tensor | None = None, + ) -> torch.Tensor: + """SageAttention dense path. + + - batched NHD ``[B, T, H, D]`` when ``cu_seqlens is None`` + - packed ``[total, H, D]`` when ``cu_seqlens`` is provided + """ + from sageattention import sageattn + + if cu_seqlens is None: + return sageattn( + query.contiguous(), + key.contiguous(), + value.contiguous(), + tensor_layout="NHD", + is_causal=self.causal, + sm_scale=self.softmax_scale, + ) + + bounds = [int(x) for x in cu_seqlens.tolist()] + output = torch.empty_like(query) + for start, stop in zip(bounds[:-1], bounds[1:]): + if start == stop: + continue + output[start:stop] = sageattn( + query[start:stop].unsqueeze(0).contiguous(), + key[start:stop].unsqueeze(0).contiguous(), + value[start:stop].unsqueeze(0).contiguous(), + tensor_layout="NHD", + is_causal=self.causal, + sm_scale=self.softmax_scale, + )[0] + return output + def _run_sol_attn_thd( self, query: torch.Tensor, @@ -167,17 +216,24 @@ class SolAttnImpl(AttentionImpl): v = value.unsqueeze(0).contiguous() if q.dtype != torch.bfloat16: raise TypeError(f"Sol-Attn requires bfloat16 activations, got {q.dtype}") - out = sol_attn( - q, - k, - v, - tau=cfg["tau"], - thresh_type=cfg["thresh_type"], - kv_splits=_resolve_kv_splits(q, cfg["kv_splits"]), - sink_start=cfg["sink_start"], - sink_tokens=cfg["sink_tokens"], - ) - return out.squeeze(0) + + if self._sol_params is None: + self._sol_params = frozenset(inspect.signature(sol_attn).parameters) + + kwargs = { + "tau": cfg["tau"], + "thresh_type": cfg["thresh_type"], + "kv_splits": _resolve_kv_splits(q, cfg["kv_splits"]), + "sink_start": cfg["sink_start"], + "sink_tokens": cfg["sink_tokens"], + } + # Wan2GP Ada port: INT8-QK Triton; official NVlabs API has no int8_qk. + if "int8_qk" in self._sol_params and tuple( + torch.cuda.get_device_capability(q.device) + ) >= (8, 9): + kwargs["int8_qk"] = True + kwargs = {k: v for k, v in kwargs.items() if k in self._sol_params} + return sol_attn(q, k, v, **kwargs).squeeze(0) def forward( self, @@ -187,14 +243,17 @@ class SolAttnImpl(AttentionImpl): attn_metadata: AttentionMetadata, ) -> torch.Tensor: del attn_metadata + # ``query`` is NHD: [B, T, H, D] if self._should_use_dense(): - q = query.transpose(1, 2).reshape( + if _get_sol_attn_runtime_config()["dense_backend"] == "sage_attn": + return self._dense_sage(query, key, value) + # NHD [B, T, H, D] → packed THD [B*T, H, D] (plain reshape; do not + # transpose — that would scramble token order for flash_attn_varlen). + q = query.reshape( query.shape[0] * query.shape[1], query.shape[2], query.shape[3] ) - k = key.transpose(1, 2).reshape( - key.shape[0] * key.shape[1], key.shape[2], key.shape[3] - ) - v = value.transpose(1, 2).reshape( + k = key.reshape(key.shape[0] * key.shape[1], key.shape[2], key.shape[3]) + v = value.reshape( value.shape[0] * value.shape[1], value.shape[2], value.shape[3] ) cu_seqlens = torch.arange( @@ -204,14 +263,11 @@ class SolAttnImpl(AttentionImpl): device=query.device, dtype=torch.int32, ) - out = self._dense_varlen( - q, - k, - v, - cu_seqlens=cu_seqlens, - max_seqlen=query.shape[1], + out = self._dense_fa( + q, k, v, cu_seqlens=cu_seqlens, max_seqlen=query.shape[1] ) return out.reshape(query.shape[0], query.shape[1], query.shape[2], -1) + q = query.reshape(query.shape[0] * query.shape[1], query.shape[2], -1) k = key.reshape(key.shape[0] * key.shape[1], key.shape[2], -1) v = value.reshape(value.shape[0] * value.shape[1], value.shape[2], -1) @@ -230,11 +286,9 @@ class SolAttnImpl(AttentionImpl): ) -> torch.Tensor: del cu_seqlens_host if self._should_use_dense(): - return self._dense_varlen( - query, - key, - value, - cu_seqlens=cu_seqlens, - max_seqlen=max_seqlen, + if _get_sol_attn_runtime_config()["dense_backend"] == "sage_attn": + return self._dense_sage(query, key, value, cu_seqlens=cu_seqlens) + return self._dense_fa( + query, key, value, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen ) return self._run_sol_attn_thd(query, key, value) diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py index b37e02a8a..25b1dc1da 100644 --- a/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/__init__.py @@ -8,6 +8,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.bitsandbytes import ( 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, +) from sglang.multimodal_gen.runtime.layers.quantization.fp8 import Fp8Config from sglang.multimodal_gen.runtime.layers.quantization.modelopt_fp8 import ( ModelOptFp8Config as ModelOptFp8DiffusionConfig, @@ -33,6 +36,7 @@ QuantizationMethods = Literal[ "mxfp8", "mxfp4", "mxfp4_npu", + "kitchen_int8", ] QUANTIZATION_METHODS: list[str] = list(get_args(QuantizationMethods)) @@ -48,6 +52,7 @@ _CUSTOMIZED_METHOD_TO_QUANT_CONFIG = { "mxfp4": Mxfp4Config, "mxfp8": MXFP8Config, "mxfp4_npu": NPUMXFP4Config, + "kitchen_int8": KitchenInt8Config, } diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_int8_config.py b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_int8_config.py new file mode 100644 index 000000000..ed1986739 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/configs/kitchen_int8_config.py @@ -0,0 +1,122 @@ +# 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``. +""" + +from __future__ import annotations + +from typing import Any + +import torch + +from sglang.multimodal_gen.runtime.layers.linear import UnquantizedLinearMethod +from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger +from sglang.srt.layers.quantization.utils import is_layer_skipped + +logger = init_logger(__name__) + +_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``. + """ + + def __init__( + self, + group_size: int = 256, + ignored_layers: list[str] | None = None, + packed_modules_mapping: dict[str, list[str]] | None = None, + ) -> None: + super().__init__() + if group_size not in _SUPPORTED_GROUP_SIZES: + raise ValueError( + f"kitchen_int8 group_size must be one of {_SUPPORTED_GROUP_SIZES}, " + f"got {group_size}" + ) + self.group_size = group_size + self.ignored_layers = ignored_layers or [] + self.packed_modules_mapping = packed_modules_mapping or {} + # 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] = [] + self.skipped: list[str] = [] + self._processed = 0 + self._quantized_bytes = 0 + + @classmethod + def get_name(cls) -> str: + return "kitchen_int8" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.float16] + + @classmethod + def get_min_capability(cls) -> int: + # INT8 tensor cores land on Turing. + return 75 + + @classmethod + def get_config_filenames(cls) -> list[str]: + return [] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> KitchenInt8Config: + return cls( + group_size=cls.get_from_keys_or(config, ["group_size"], 256), + ignored_layers=cls.get_from_keys_or(config, ["ignored_layers"], None), + ) + + def get_quant_method( + self, layer: torch.nn.Module, prefix: str + ) -> QuantizeMethodBase | None: + from sglang.multimodal_gen.runtime.layers.linear import LinearBase + from sglang.multimodal_gen.runtime.layers.quantization.kitchen_int8 import ( + KitchenInt8LinearMethod, + ) + + if not isinstance(layer, LinearBase): + return None + if is_layer_skipped( + prefix, self.ignored_layers, fused_mapping=self.packed_modules_mapping + ): + self.skipped.append(prefix) + return UnquantizedLinearMethod() + # The rotation partitions the input dim into fixed-size groups, so a + # layer whose input does not divide evenly simply stays in BF16 rather + # than failing the whole model. H3's adaln projections (in=2688) are + # the case this exists for, and they cost 0.2% of a step anyway. + if layer.input_size % self.group_size: + self.skipped.append(f"{prefix}(in={layer.input_size})") + return UnquantizedLinearMethod() + self.selected.append(prefix) + return KitchenInt8LinearMethod(self) + + def note_quantized(self, saved_bytes: int) -> None: + self._processed += 1 + self._quantized_bytes += saved_bytes + if self._processed == len(self.selected): + logger.info( + "kitchen_int8: quantized %d linear layers (%.2f GiB of BF16 weights " + "-> %.2f GiB INT8), left %d in BF16", + self._processed, + self._quantized_bytes / 1024**3, + self._quantized_bytes / 2 / 1024**3, + len(self.skipped), + ) + logger.debug("kitchen_int8: layers left in BF16: %s", self.skipped) + + def get_scaled_act_names(self) -> list[str]: + return [] diff --git a/python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_int8.py b/python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_int8.py new file mode 100644 index 000000000..53597d83b --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/layers/quantization/kitchen_int8.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +"""INT8 weight-only-storage linear backed by comfy_kitchen's fused ConvRot kernel. + +On Ada (RTX 4090) INT8 is only worth doing with the right kernel: on MiniMax H3 +shapes `torch._int_mm` measures 0.46-0.90x of BF16 (i.e. slower) and a Triton +INT8 GEMM roughly ties BF16, while `comfy_kitchen.int8_linear` reaches 2.49x. +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. +""" + +from __future__ import annotations + +import os + +import torch +from torch.nn.parameter import Parameter + +from sglang.multimodal_gen.runtime.layers.linear import LinearMethodBase +from sglang.multimodal_gen.runtime.layers.quantization.configs.kitchen_int8_config import ( + KitchenInt8Config, +) +from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs + +__all__ = ["KitchenInt8Config", "KitchenInt8LinearMethod"] + +# comfy_kitchen's dtype codes for the fused op's output. +_OUT_DTYPE_CODE = {torch.float32: 0, torch.float16: 1, torch.bfloat16: 2} + +# comfy_kitchen picks its CUTLASS tile configuration from a threshold tree +# (select_fused_int8_config in cutlass_gemm_int8.cu). Shapes whose N falls under +# its 24832 cutoff but whose M is large get a Stream-K schedule, which exists to +# balance load when there are too few tiles to fill the GPU. At H3's 32700 tokens +# qkv_proj already launches ~21k CTAs over the 4090's 128 SMs, so Stream-K's +# workspace and fixup reduction are pure overhead: 26.5 ms against 17.9 ms for +# the identical tile without it. Capping rows per call keeps the plain +# data-parallel config, and is bit-exact because splitting rows does not change +# any single row's arithmetic. +_MAX_ROWS_PER_CALL = int(os.environ.get("SGLANG_KITCHEN_INT8_MAX_ROWS", "8192")) +# Narrow outputs do not recover the cost of writing results back through a +# preallocated buffer; H3's out_proj and fc2 (N=5376) both measure slower split. +_MIN_SPLIT_OUTPUT = int(os.environ.get("SGLANG_KITCHEN_INT8_MIN_SPLIT_N", "8192")) + + +def _row_split(rows: int, out_features: int) -> int | None: + """Rows per `int8_linear` call, or None to issue one call for everything.""" + if _MAX_ROWS_PER_CALL <= 0 or rows <= _MAX_ROWS_PER_CALL: + return None + if out_features < _MIN_SPLIT_OUTPUT: + return None + return _MAX_ROWS_PER_CALL + + +def _load_comfy_kitchen(): + try: + import comfy_kitchen # noqa: F401 + except ImportError as exc: # pragma: no cover - depends on optional dep + raise ImportError( + "kitchen_int8 quantization requires the `comfy-kitchen` package " + "(pip install comfy-kitchen). It is a self-contained abi3 extension " + "and does not link against libtorch, so any torch version works." + ) from exc + if not hasattr(torch.ops.comfy_kitchen, "int8_linear"): + raise RuntimeError( + "comfy_kitchen is installed but did not register " + "torch.ops.comfy_kitchen.int8_linear" + ) + + +class KitchenInt8LinearMethod(LinearMethodBase): + """Quantizes BF16 weights to INT8 after load and runs the fused kernel.""" + + def __init__(self, quant_config: KitchenInt8Config) -> None: + self.quant_config = quant_config + _load_comfy_kitchen() + + def create_weights( + self, + layer: torch.nn.Module, + input_size_per_partition: int, + output_partition_sizes: list[int], + input_size: int, + output_size: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ) -> None: + # 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: + 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}" + ) + + # 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. + weight = Parameter( + torch.empty( + sum(output_partition_sizes), + input_size_per_partition, + dtype=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) + + 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 + return + + # 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. + home = weight.device + qdata, params = TensorWiseINT8Layout.quantize( + weight.to("cuda", non_blocking=True), + is_weight=True, + per_channel=True, + convrot=True, + convrot_groupsize=self.quant_config.group_size, + stochastic_rounding=0, + ) + layer.weight = Parameter(qdata.to(home), requires_grad=False) + layer.register_parameter( + "weight_scale", + Parameter( + params.scale.to(device=home, dtype=torch.float32), requires_grad=False + ), + ) + self.quant_config.note_quantized(weight.numel() * weight.element_size()) + del qdata, params + torch.cuda.empty_cache() + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + out_code = _OUT_DTYPE_CODE.get(x.dtype) + if out_code is None: + raise ValueError( + f"kitchen_int8 does not support activation dtype {x.dtype}" + ) + + # The kernel takes 2D activations; callers may pass [..., K]. + orig_shape = x.shape + if x.dim() != 2: + x = x.reshape(-1, orig_shape[-1]) + x = x.contiguous() + + def run(rows: torch.Tensor) -> torch.Tensor: + return torch.ops.comfy_kitchen.int8_linear( + rows, + layer.weight, + layer.weight_scale, + bias, + out_code, + True, # convrot + self.quant_config.group_size, + ) + + n_rows, n_out = x.shape[0], layer.weight.shape[0] + split = _row_split(n_rows, n_out) + if split is None: + out = run(x) + else: + # Row slices of a contiguous 2D tensor are themselves contiguous, so + # this splits without copying the activation. + out = torch.empty(n_rows, n_out, dtype=x.dtype, device=x.device) + for start in range(0, n_rows, split): + out[start : start + split] = run(x[start : start + split]) + + if len(orig_shape) != 2: + out = out.reshape(*orig_shape[:-1], out.shape[-1]) + return out diff --git a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py index 9967cd38e..052545dd1 100644 --- a/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/transformer_load_utils.py @@ -742,13 +742,13 @@ def _resolve_quant_config( if server_args.quantization == "modelslim": return get_quant_config(hf_config, component_model_path) - # Online-quant convention: for `fp8` and `mxfp4`, a no-arg - # QuantizationConfig() selects the post-load path -- weights load - # in source dtype and are quantized in + # Online-quant convention: for `fp8`, `mxfp4` and `kitchen_int8`, a + # no-arg QuantizationConfig() selects the post-load path -- weights + # load in source dtype and are quantized in # process_weights_after_loading. quant_cls = get_quantization_config(server_args.quantization) quant_kwargs = {} - if server_args.quantization in {"fp8", "mxfp4"}: + if server_args.quantization in {"fp8", "mxfp4", "kitchen_int8"}: quant_kwargs["ignored_layers"] = getattr( server_args, "quantization_ignored_layers", None ) diff --git a/python/sglang/multimodal_gen/test/unit/test_sol_attn_backend.py b/python/sglang/multimodal_gen/test/unit/test_sol_attn_backend.py index c9466ab33..b3910dfa4 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sol_attn_backend.py +++ b/python/sglang/multimodal_gen/test/unit/test_sol_attn_backend.py @@ -7,6 +7,7 @@ import torch from sglang.multimodal_gen.runtime.layers.attention.backends.sol_attn import ( SolAttnBackend, SolAttnImpl, + _get_sol_attn_runtime_config, _parse_layer_ranges, ) from sglang.multimodal_gen.runtime.platforms.cuda import CudaPlatformBase @@ -43,6 +44,32 @@ class TestSolAttnBackend(unittest.TestCase): def test_parse_layer_ranges(self): self.assertEqual(_parse_layer_ranges("0,1,3-5"), frozenset({0, 1, 3, 4, 5})) + def test_dense_backend_aliases(self): + for raw, expected in ( + ("fa", "fa"), + ("sage", "sage_attn"), + ("sage_attn", "sage_attn"), + ): + server_args = MagicMock() + server_args.attention_backend_config = {"dense_backend": raw} + with patch( + "sglang.multimodal_gen.runtime.layers.attention.backends.sol_attn.get_global_server_args", + return_value=server_args, + ): + self.assertEqual( + _get_sol_attn_runtime_config()["dense_backend"], expected + ) + server_args = MagicMock() + server_args.attention_backend_config = {"dense_backend": "torch_sdpa"} + with ( + patch( + "sglang.multimodal_gen.runtime.layers.attention.backends.sol_attn.get_global_server_args", + return_value=server_args, + ), + self.assertRaises(ValueError), + ): + _get_sol_attn_runtime_config() + def test_backend_head_size(self): self.assertEqual(SolAttnBackend.get_supported_head_sizes(), [128])