[diffusion] quant: support gguf (#35370)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
zijiexia
2026-08-20 15:46:34 +08:00
committed by GitHub
co-authored by Claude Fable 5 Mick
parent ae23423b46
commit 21c88f8625
23 changed files with 1788 additions and 45 deletions
@@ -198,6 +198,50 @@ baseline; everything else stays identical. GPU peak stays about 18 GB
either way because streaming offload is set by the offload buffers and VAE either way because streaming offload is set by the offload buffers and VAE
decode, not the weight dtype. decode, not the weight dtype.
### Pre-quantized GGUF transformer
Use `--transformer-weights-path` to replace only the DiT with a GGUF file; the
base repository continues to provide the text encoder, VAEs, scheduler, and
tokenizers. Do not also pass `--quantization gguf`.
```bash 1×RTX 5090 Q4_K_M
sglang serve \
--model-path MiniMaxAI/MiniMax-H3 \
--model-variant fl2va \
--transformer-weights-path \
leejet/MiniMax-H3-GGUF/minimax_h3_fl2va-Q4_K_M.gguf \
--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 \
--port 30010
```
The loader also recognizes pruned checkpoints that replace the timestep MLP
with `adaln_t_table`. Repositories containing both FL2VA and Ref2VA variants
need a full file reference:
```bash Pruned FL2VA Q4_K
sglang serve \
--model-path MiniMaxAI/MiniMax-H3 \
--model-variant fl2va \
--transformer-weights-path \
unsloth/MiniMax-H3-GGUF/minimax_h3_fl2va_pruned-Q4_K.gguf \
--performance-mode memory \
--layerwise-offload-components dit,text_encoder \
--enable-torch-compile false \
--port 30010
```
The linear adapter reuses SRT's GGUF type definitions and CUDA dequantization,
then runs the native GEMM. SRT's fused MMVQ/MMQ kernels target the low-token LLM
regime and are slower at diffusion sequence lengths. TP is supported when each
row-parallel input shard remains GGML-block aligned; incompatible degrees fail
during model construction. FSDP, LoRA merging, and the separate MiniMax-H3
AdaLN cache flags are not compatible with packed GGUF weights.
The first launch downloads the model through the selected Hub. If the Hugging The first launch downloads the model through the selected Hub. If the Hugging
Face repository requires authentication, export a Hugging Face token in the Face repository requires authentication, export a Hugging Face token in the
server environment. server environment.
+2 -1
View File
@@ -125,7 +125,8 @@ For quantized transformer checkpoints, prefer:
- `--model-path` for the base pipeline - `--model-path` for the base pipeline
- `--transformer-path` for a quantized `transformers` transformer component folder - `--transformer-path` for a quantized `transformers` transformer component folder
- `--transformer-weights-path` for a quantized safetensors file, directory, or repo - `--transformer-weights-path` for a quantized safetensors file, directory,
repo, or a supported GGUF transformer file
- `--quantization` for online quantization (apply quantization to unquantized models at load time, activations are quantized dynamically) - `--quantization` for online quantization (apply quantization to unquantized models at load time, activations are quantized dynamically)
- `--quantization-ignored-layers` layer name patterns to keep unquantized (e.g. `attention.to_`) - `--quantization-ignored-layers` layer name patterns to keep unquantized (e.g. `attention.to_`)
@@ -135,7 +135,7 @@ Rows are grouped when a family shares the same runtime path or optimization supp
<td>MiniMax-H3</td> <td>MiniMax-H3</td>
<td><div className="sgd-id-list"><code>MiniMaxAI/MiniMax-H3</code></div></td> <td><div className="sgd-id-list"><code>MiniMaxAI/MiniMax-H3</code></div></td>
<td>T2VA / FL2VA / Ref2VA, 768p at 24 fps with synchronized audio</td> <td>T2VA / FL2VA / Ref2VA, 768p at 24 fps with synchronized audio</td>
<td><span className="sgd-chip">Cache-DiT</span><span className="sgd-chip">Sage</span><span className="sgd-chip">Online FP8</span></td> <td><span className="sgd-chip">Cache-DiT</span><span className="sgd-chip">Sage</span><span className="sgd-chip">Online FP8</span><span className="sgd-chip">GGUF</span></td>
</tr> </tr>
<tr> <tr>
<td>Wan2.1 Fun</td> <td>Wan2.1 Fun</td>
+133
View File
@@ -135,6 +135,14 @@ backend.
<td>None</td> <td>None</td>
<td>Currently only compatible with the Ascend NPU family and supports <code>mxfp8</code>, <code>mxfp4</code>, <code>w8a8</code>, and <code>w4a4</code></td> <td>Currently only compatible with the Ascend NPU family and supports <code>mxfp8</code>, <code>mxfp4</code>, <code>w8a8</code>, and <code>w4a4</code></td>
</tr> </tr>
<tr>
<td><code>gguf</code></td>
<td>A single community <code>.gguf</code> holding the transformer</td>
<td><code>--transformer-weights-path</code></td>
<td>MiniMax-H3 <code>fl2va</code> (original and pruned AdaLN curve)</td>
<td>None</td>
<td>CUDA only, no FSDP. TP requires GGML-aligned shard boundaries. Shrinks the download and the host memory offload pins (17.5 vs 61.7 GiB for H3) rather than peak VRAM, which offload already bounds. Dequantized per use, so it is not faster. See <a href="#gguf">GGUF</a>.</td>
</tr>
</tbody> </tbody>
</table> </table>
@@ -672,6 +680,131 @@ sglang generate \
`quant_algo=NVFP4`; the `modelopt-nvfp4` label here is again a documentation `quant_algo=NVFP4`; the `modelopt-nvfp4` label here is again a documentation
family name rather than a serialized config key. family name rather than a serialized config key.
## GGUF
GGUF loads a community-quantized transformer from a single `.gguf` file while
the rest of the pipeline — VAE, text encoder, scheduler, tokenizer — keeps
loading from the base model.
GGUF primarily reduces checkpoint, host-memory, and resident-weight size. For
example, MiniMax-H3's transformer is 17.5 GiB as Q4_K_M versus 61.7 GiB as
BF16. With full layerwise offload, VAE decode and offload buffers can still
dominate peak GPU memory, but each streamed DiT layer also transfers fewer
bytes.
Packed linears reuse SRT's GGUF type definitions and CUDA dequantization, then
run the native GEMM. SRT's fused MMVQ/MMQ kernels target the low-token LLM
regime and are slower at diffusion sequence lengths. GGUF remains a
capacity-oriented option; latency depends on the quantization type, activation
shape, and placement policy.
Layers the checkpoint stores unquantized (F32/F16/BF16) take the ordinary linear
path rather than the packed one. Their precision is then whatever the model
declares for that layer, exactly as on the safetensors path — a checkpoint
cannot raise a layer above the model's own dtype by storing it wider. For
MiniMax-H3 the two agree: the layers it pins to FP32 are the ones the validated
checkpoint leaves unquantized, and `post_load_weights` fails the load if any of
them ends up narrower.
### Usage
No extra install: `gguf` is already a core SGLang dependency.
`--model-path` stays the base model; `--transformer-weights-path` takes the
GGUF. A local path, `owner/repo/file.gguf`, or `owner/repo:QUANT_TYPE` all work.
The quant-type shorthand is accepted only when exactly one repository file
matches it; otherwise SGLang lists the candidates and asks for a full path.
```bash
sglang serve \
--model-path MiniMaxAI/MiniMax-H3 \
--model-variant fl2va \
--transformer-weights-path \
leejet/MiniMax-H3-GGUF/minimax_h3_fl2va-Q4_K_M.gguf \
--num-gpus 1 \
--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 \
--port 30010
```
Note that `--quantization gguf` is not the selector — the quantization is read
from the file itself, so passing the file is what enables the path.
MiniMax-H3 GGUF checkpoints may use either the original timestep MLP or the
pruned AdaLN curve architecture. Repositories often contain both FL2VA and
Ref2VA files, so use the full Hub file reference instead of an ambiguous
`owner/repo:QUANT_TYPE` selector:
```bash
sglang serve \
--model-path MiniMaxAI/MiniMax-H3 \
--model-variant fl2va \
--transformer-weights-path \
unsloth/MiniMax-H3-GGUF/minimax_h3_fl2va_pruned-Q4_K.gguf \
--performance-mode memory \
--layerwise-offload-components dit,text_encoder \
--enable-torch-compile false \
--port 30010
```
The pruned architecture keeps its sampled curve and reduced AdaLN projections
in FP32, matching the published checkpoint implementation. This precision
island deliberately bypasses the BF16-only fused modulation kernels.
#### Whether to offload the VAE
Adding `vae` to `--layerwise-offload-components` trades a lot of latency for
some peak VRAM, because the video VAE decoder is re-streamed per decode tile.
Measured on 1× RTX 5090 with this checkpoint, 1344×768 × 107 frames:
| `--layerwise-offload-components` | Peak VRAM | Denoise | VAE decode |
| --- | ---: | ---: | ---: |
| `dit,text_encoder` | 26.3 GiB | 39.0 s | **9.5 s** |
| `dit,text_encoder,vae` | **19.6 GiB** | 39.0 s | 57.2 s |
Denoise is unaffected, and the output is bit-identical either way. Leave the VAE
resident unless the 6.7 GiB matters — on a 24 GB card by this measurement it
does, and the 6× slower decode is the price of fitting.
### Constraints
| Constraint | Reason |
| --- | --- |
| TP shard boundaries must align to GGML blocks | Column-parallel rows shard directly; row-parallel packed columns require each local input partition to contain whole quantization blocks |
| No `--use-fsdp-inference` | FSDP does not preserve the GGUF packed-block layout |
| CUDA only | The reused SRT GGML dequantization kernel currently ships for CUDA |
| Native byte order only | A quantized block embeds its scales, so a non-native file cannot be byte-swapped as a whole |
| No LoRA, and none of the H3 AdaLN cache flags | An adapter cannot be merged into packed blocks, and the AdaLN paths read the transformer's safetensors |
| No `--quantization` | The checkpoint fixes the quantization; the flag would be a second, conflicting selector |
Every constraint above fails at startup with an explanatory error rather than
silently producing wrong output.
Sequence parallelism (`--ulysses-degree` / `--ring-degree`) remains available
because it shards activations rather than packed weights. TP is also available;
startup rejects a degree that cuts a row-parallel matrix inside a GGML block.
### Validated scope
| Model | Checkpoint | Hardware | Result |
| --- | --- | --- | --- |
| MiniMax-H3 `fl2va` | [`leejet/MiniMax-H3-GGUF`](https://huggingface.co/leejet/MiniMax-H3-GGUF) `minimax_h3_fl2va-Q4_K_M.gguf` (17.5 GiB) | 1x RTX 5090 (32 GiB) | t2va 1344x768, 107 frames, video + audio; 19.6-26.3 GiB peak depending on VAE offload |
| MiniMax-H3 `fl2va` | `minimax_h3_fl2va-Q4_K_M.gguf` (17.5 GiB) | 1x H200 (141 GiB) | 2-step t2va 1344x768, 107 frames, H.264 + AAC; 10.13 s and 17.17 GiB peak |
| MiniMax-H3 `fl2va`, pruned AdaLN curve | [`unsloth/MiniMax-H3-GGUF`](https://huggingface.co/unsloth/MiniMax-H3-GGUF) `minimax_h3_fl2va_pruned-Q4_K.gguf` (10.7 GiB loaded DiT) | 1x H200 (141 GiB) | 2-step t2va 1344x768, 107 frames, H.264 + AAC |
| MiniMax-H3 `fl2va`, pruned AdaLN curve | `minimax_h3_fl2va_pruned-Q4_K.gguf` | 1x GB300 (CUDA 13, PyTorch 2.13) | 50-step t2va 1344x768, 107 frames, H.264 + AAC; 105.38 s and 80.88 GB peak |
| MiniMax-H3 `fl2va`, pruned AdaLN curve | `minimax_h3_fl2va_pruned-Q4_K.gguf` | 2x GB300, TP2 (CUDA 13, PyTorch 2.13) | 2-step t2va 1344x768, 107 frames, H.264 + AAC; 7.55 s and 51.90 GB peak per rank |
The DiT loads at 17.5 GiB against 61.7 GiB for the BF16 checkpoint. Weight
fidelity was checked tensor-by-tensor against the BF16 reference: cosine
1.00000 for the F32/BF16 tensors and 0.9973 for Q4_K/Q4_0.
Not validated in the measurements above: any other quantization type, the
`ref2va` partition, and a BF16-vs-GGUF output comparison.
## Nunchaku (SVDQuant) ## Nunchaku (SVDQuant)
### Install ### Install
@@ -65,6 +65,7 @@ These options **trade output quality** for speed or VRAM savings. Results will d
| **Progressive Resolution** | `--progressive-mode dct_rewind --progressive-levels N --progressive-delta D` | Runs early denoising at lower latent resolution, then spectrally upsamples and switches to the target resolution. | Model- and schedule-dependent | Approximate and pipeline-specific. Keep the switch schedule fixed and compare detail, composition, and temporal stability. | | **Progressive Resolution** | `--progressive-mode dct_rewind --progressive-levels N --progressive-delta D` | Runs early denoising at lower latent resolution, then spectrally upsamples and switches to the target resolution. | Model- and schedule-dependent | Approximate and pipeline-specific. Keep the switch schedule fixed and compare detail, composition, and temporal stability. |
| **Causal KV-Cache Quantization** | `--kv-cache-quant int4\|int2` plus optional `--kv-cache-quant-*` controls | Compresses completed causal KV-cache chunks with Quant-VideoGen PRQ while keeping the mutable/current chunk and recent chunks in BF16. | Primarily a long-session memory saving | Currently limited to LingBot World realtime causal serving; requires `quant-videogen`. INT4 is the starting point; INT2 saves more memory with more error. It quantizes cache state, not checkpoint weights. | | **Causal KV-Cache Quantization** | `--kv-cache-quant int4\|int2` plus optional `--kv-cache-quant-*` controls | Compresses completed causal KV-cache chunks with Quant-VideoGen PRQ while keeping the mutable/current chunk and recent chunks in BF16. | Primarily a long-session memory saving | Currently limited to LingBot World realtime causal serving; requires `quant-videogen`. INT4 is the starting point; INT2 saves more memory with more error. It quantizes cache state, not checkpoint weights. |
| **Quantized Models (Nunchaku / SVDQuant)** | `--enable-svdquant --transformer-weights-path <path>` + optional `--quantization-precision int4\|nvfp4`, `--quantization-rank 32` | W4A4-style quantization via [Nunchaku](https://nunchaku.tech). Reduces DiT weight memory by ~4x. Precision/rank can be auto-inferred from weight filename or set explicitly. | ~1.5–2x compute speedup | Lossy quantization; quality depends on rank and precision. Requires pre-quantized weights. Ampere (SM8x) or SM12x only (no Hopper SM90). Higher rank = better quality but more memory. | | **Quantized Models (Nunchaku / SVDQuant)** | `--enable-svdquant --transformer-weights-path <path>` + optional `--quantization-precision int4\|nvfp4`, `--quantization-rank 32` | W4A4-style quantization via [Nunchaku](https://nunchaku.tech). Reduces DiT weight memory by ~4x. Precision/rank can be auto-inferred from weight filename or set explicitly. | ~1.5–2x compute speedup | Lossy quantization; quality depends on rank and precision. Requires pre-quantized weights. Ampere (SM8x) or SM12x only (no Hopper SM90). Higher rank = better quality but more memory. |
| **GGUF Transformer** | `--transformer-weights-path <file.gguf\|owner/repo:QUANT>` | Loads a community-quantized DiT from one `.gguf`; other components stay on the base model. **Shrinks the checkpoint, not the peak VRAM** — offload already bounds peak, so reach for this when the *download* or the host RAM offload pins is the problem (MiniMax-H3 17.5 vs 61.7 GiB), not when VRAM is. For a 24 GB card `kitchen_int8` is the faster option if you can afford the full BF16 checkpoint on disk. | None; expect a small slowdown from per-step dequantization | Lossy (4-bit families ~0.997 cosine vs BF16). CUDA only, `--tp-size 1`, no FSDP, no LoRA, no `--quantization`, no `--enable-svdquant`, and mutually exclusive with the H3 AdaLN cache/online flags — each rejected at startup. Validated on MiniMax-H3 `fl2va` Q4_K_M, 1 GPU. |
| **Pre-quantized Transformer Override** | `--transformer-path <dir-or-repo>` / `--transformer-weights-path <path>` | Load a quantized transformer component or raw transformer weights. For converted ModelOpt FP8/NVFP4 directories, prefer `--transformer-path`; use `--transformer-weights-path` for weight-only artifacts the model loader expects. | ~1.3–1.5x compute (dtype dependent) | Requires a validated quantized transformer override, such as one produced by the ModelOpt helper tools. Quality is usually slightly worse than BF16 and depends on the format, fallback layers, and calibration scope. | | **Pre-quantized Transformer Override** | `--transformer-path <dir-or-repo>` / `--transformer-weights-path <path>` | Load a quantized transformer component or raw transformer weights. For converted ModelOpt FP8/NVFP4 directories, prefer `--transformer-path`; use `--transformer-weights-path` for weight-only artifacts the model loader expects. | ~1.3–1.5x compute (dtype dependent) | Requires a validated quantized transformer override, such as one produced by the ModelOpt helper tools. Quality is usually slightly worse than BF16 and depends on the format, fallback layers, and calibration scope. |
| **Component Precision Override** | `--dit-precision fp16`, `--vae-precision fp16\|bf16` | On-the-fly dtype conversion for individual components. E.g. convert a BF16 model to FP16 at load time, or run VAE in BF16 instead of FP32. | Reduces memory; FP16 can be faster on some GPUs | May affect numerical stability. VAE is FP32 by default for accuracy; lowering it is lossy. DiT defaults to BF16. | | **Component Precision Override** | `--dit-precision fp16`, `--vae-precision fp16\|bf16` | On-the-fly dtype conversion for individual components. E.g. convert a BF16 model to FP16 at load time, or run VAE in BF16 instead of FP32. | Reduces memory; FP16 can be faster on some GPUs | May affect numerical stability. VAE is FP32 by default for accuracy; lowering it is lossy. DiT defaults to BF16. |
| **Fewer Inference Steps** | `--num-inference-steps N` (sampling param) | Reduces the number of denoising steps. Fewer steps = faster. | Linear speedup | Quality degrades with too few steps. Model-dependent optimal range. | | **Fewer Inference Steps** | `--num-inference-steps N` (sampling param) | Reduces the number of denoising steps. Fewer steps = faster. | Linear speedup | Quality degrades with too few steps. Model-dependent optimal range. |
@@ -77,6 +77,8 @@ class MiniMaxH3DiTArchConfig(DiTArchConfig):
timestep_input_dim: int = 256 timestep_input_dim: int = 256
time_embed_hidden_size: int = 5376 time_embed_hidden_size: int = 5376
time_embed_dim: int = 2688 time_embed_dim: int = 2688
# Pruned checkpoints replace the timestep MLP with a sampled AdaLN curve.
adaln_curve_grid: int | None = None
adaln_out_features: int = 18 * 5376 adaln_out_features: int = 18 * 5376
final_adaln_out_features: int = 2 * 5376 final_adaln_out_features: int = 2 * 5376
rope_inv_freq_len: int = 16 rope_inv_freq_len: int = 16
@@ -132,6 +132,7 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
"num_gpus": server_args.num_gpus, "num_gpus": server_args.num_gpus,
"performance_mode": server_args.performance_mode, "performance_mode": server_args.performance_mode,
"quantization": server_args.quantization, "quantization": server_args.quantization,
"transformer_weights_path": server_args.transformer_weights_path,
"text_encoder_quantization": text_encoder_quantization, "text_encoder_quantization": text_encoder_quantization,
"regional_compile": server_args.regional_compile, "regional_compile": server_args.regional_compile,
"ring_degree": server_args.ring_degree, "ring_degree": server_args.ring_degree,
@@ -154,6 +155,7 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
"num_gpus": 4, "num_gpus": 4,
"performance_mode": "speed", "performance_mode": "speed",
"quantization": None, "quantization": None,
"transformer_weights_path": None,
"text_encoder_quantization": None, "text_encoder_quantization": None,
"regional_compile": False, "regional_compile": False,
"ring_degree": 1, "ring_degree": 1,
@@ -0,0 +1,133 @@
# SPDX-License-Identifier: Apache-2.0
"""Diffusion Linear adapter for SRT's GGUF kernels."""
from __future__ import annotations
from typing import Any
import gguf
import torch
from torch import nn
from sglang.multimodal_gen.runtime.layers.linear import (
LinearBase,
LinearMethodBase,
UnquantizedLinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.configs.base_config import (
QuantizationConfig,
QuantizeMethodBase,
)
from sglang.multimodal_gen.runtime.loader.gguf_weights import GGUFTensorMeta
from sglang.multimodal_gen.runtime.utils.weight_attrs import set_weight_attrs
from sglang.srt.layers.quantization.gguf import (
DEQUANT_TYPES,
UNQUANTIZED_TYPES,
dequantize_gguf_weight,
)
class GGUFConfig(QuantizationConfig):
"""Select a GGUF method from each checkpoint tensor's metadata."""
def __init__(self, gguf_file: str, tensor_meta: dict[str, GGUFTensorMeta]):
super().__init__()
self.gguf_file = gguf_file
self.tensor_meta = tensor_meta
@classmethod
def get_name(cls) -> str:
return "gguf"
@classmethod
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
return [torch.float32, torch.float16, torch.bfloat16]
@classmethod
def get_min_capability(cls) -> int:
return 60
@staticmethod
def get_config_filenames() -> list[str]:
return []
@classmethod
def from_config(cls, config: dict[str, Any]) -> GGUFConfig:
raise ValueError("GGUFConfig must be constructed from a GGUF checkpoint")
def get_quant_method(
self, layer: nn.Module, prefix: str
) -> QuantizeMethodBase | None:
if not isinstance(layer, LinearBase):
return None
metadata = self.tensor_meta.get(f"{prefix}.weight")
if metadata is None:
raise ValueError(
f"Linear layer {prefix!r} has no weight in the GGUF checkpoint "
f"{self.gguf_file!r}"
)
weight_type = metadata.weight_type
if weight_type in UNQUANTIZED_TYPES:
return UnquantizedLinearMethod()
if weight_type not in DEQUANT_TYPES:
raise ValueError(
f"GGUF tensor {prefix}.weight uses unsupported type {weight_type}"
)
return GGUFLinearMethod(metadata, prefix)
class GGUFLinearMethod(LinearMethodBase):
"""Register TP-local packed weights and reuse SRT dequantization."""
def __init__(self, metadata: GGUFTensorMeta, prefix: str) -> None:
self.metadata = metadata
self.prefix = prefix
self.weight_type = metadata.weight_type
def create_weights(
self,
layer: 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: Any,
) -> None:
if self.metadata.logical_shape != (output_size, input_size):
raise ValueError(
f"GGUF tensor {self.prefix}.weight has logical shape "
f"{self.metadata.logical_shape}, expected {(output_size, input_size)}"
)
block_size, type_size = gguf.GGML_QUANT_SIZES[self.weight_type]
if input_size_per_partition % block_size:
raise ValueError(
f"GGUF tensor {self.prefix}.weight cannot be TP-sharded: input "
f"partition {input_size_per_partition} is not aligned to "
f"quantization block size {block_size}"
)
qweight = nn.Parameter(
torch.empty(
sum(output_partition_sizes),
input_size_per_partition // block_size * type_size,
dtype=torch.uint8,
),
requires_grad=False,
)
set_weight_attrs(qweight, {"input_dim": 1, "output_dim": 0})
set_weight_attrs(qweight, extra_weight_attrs)
layer.register_parameter("qweight", qweight)
def apply(
self,
layer: nn.Module,
x: torch.Tensor,
bias: torch.Tensor | None = None,
) -> torch.Tensor:
weight = dequantize_gguf_weight(layer.qweight, self.weight_type, x.dtype)
return nn.functional.linear(x, weight, bias)
__all__ = ["GGUFConfig", "GGUFLinearMethod"]
@@ -16,8 +16,10 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader imp
ComponentLoader, ComponentLoader,
) )
from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model from sglang.multimodal_gen.runtime.loader.fsdp_load import maybe_load_fsdp_model
from sglang.multimodal_gen.runtime.loader.gguf_weights import gguf_weights_iterator
from sglang.multimodal_gen.runtime.loader.transformer_load_utils import ( from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
TransformerQuantLoadSpec, TransformerQuantLoadSpec,
resolve_transformer_gguf_to_load,
resolve_transformer_quant_load_spec, resolve_transformer_quant_load_spec,
resolve_transformer_safetensors_to_load, resolve_transformer_safetensors_to_load,
) )
@@ -45,8 +47,11 @@ def _resolve_checkpoint_load_device(
*, *,
component_starts_on_cpu: bool, component_starts_on_cpu: bool,
runtime_quant_config: object | None, runtime_quant_config: object | None,
quantized_cpu_load_supported: bool = False,
) -> torch.device: ) -> torch.device:
if component_starts_on_cpu and runtime_quant_config is None: if component_starts_on_cpu and (
runtime_quant_config is None or quantized_cpu_load_supported
):
return torch.device("cpu") return torch.device("cpu")
return runtime_device return runtime_device
@@ -177,9 +182,17 @@ class TransformerLoader(ComponentLoader):
# 1. hf config # 1. hf config
config = get_diffusers_component_config(component_path=component_model_path) config = get_diffusers_component_config(component_path=component_model_path)
safetensors_list = resolve_transformer_safetensors_to_load( gguf_file = resolve_transformer_gguf_to_load(
component_server_args, component_model_path component_server_args, component_name
) )
if gguf_file is not None:
# A GGUF file holds the whole transformer; the remaining components
# still load from the base model path.
safetensors_list = []
else:
safetensors_list = resolve_transformer_safetensors_to_load(
component_server_args, component_model_path
)
# 2. dit config # 2. dit config
# Config from Diffusers supersedes sgl_diffusion's model config # Config from Diffusers supersedes sgl_diffusion's model config
@@ -209,7 +222,21 @@ class TransformerLoader(ComponentLoader):
model_cls=model_cls, model_cls=model_cls,
cls_name=cls_name, cls_name=cls_name,
component_name=component_name, component_name=component_name,
gguf_file=gguf_file,
) )
if quant_spec.gguf_file is not None and cls_name == "MiniMaxH3DiTModel":
assert quant_spec.quant_config is not None
curve = quant_spec.quant_config.tensor_meta.get("adaln_t_table")
if curve is not None:
if curve.is_quantized or len(curve.logical_shape) != 2:
raise ValueError(
"MiniMax-H3 adaln_t_table must be an unquantized 2D tensor"
)
curve_grid, time_embed_dim = curve.logical_shape
if curve_grid < 2:
raise ValueError("MiniMax-H3 adaln_t_table needs at least two rows")
dit_config.arch_config.adaln_curve_grid = curve_grid
dit_config.arch_config.time_embed_dim = time_embed_dim
# Quantization adapters may require resident weights, so placement must # Quantization adapters may require resident weights, so placement must
# be resolved after they have validated the component configuration. # be resolved after they have validated the component configuration.
component_starts_on_cpu = ( component_starts_on_cpu = (
@@ -218,13 +245,21 @@ class TransformerLoader(ComponentLoader):
) )
use_fsdp = server_args.should_use_fsdp_for_component(component_name) use_fsdp = server_args.should_use_fsdp_for_component(component_name)
logger.info( if quant_spec.gguf_file is not None:
"Loading %s from %s safetensors file(s) %s, param_dtype: %s", logger.info(
cls_name, "Loading %s from GGUF file %s, param_dtype: %s",
len(safetensors_list), cls_name,
f": {safetensors_list}" if get_log_level() == logging.DEBUG else "", quant_spec.gguf_file,
quant_spec.param_dtype, quant_spec.param_dtype,
) )
else:
logger.info(
"Loading %s from %s safetensors file(s) %s, param_dtype: %s",
cls_name,
len(safetensors_list),
f": {safetensors_list}" if get_log_level() == logging.DEBUG else "",
quant_spec.param_dtype,
)
# prepare init_param # prepare init_param
init_params: dict[str, Any] = { init_params: dict[str, Any] = {
"config": dit_config, "config": dit_config,
@@ -283,6 +318,7 @@ class TransformerLoader(ComponentLoader):
local_torch_device, local_torch_device,
component_starts_on_cpu=component_starts_on_cpu, component_starts_on_cpu=component_starts_on_cpu,
runtime_quant_config=quant_spec.runtime_quant_config, runtime_quant_config=quant_spec.runtime_quant_config,
quantized_cpu_load_supported=quant_spec.gguf_file is not None,
) )
) )
direct_gpu_weight_loading = bool( direct_gpu_weight_loading = bool(
@@ -343,6 +379,15 @@ class TransformerLoader(ComponentLoader):
strict=False, strict=False,
weight_load_plan=weight_load_plan, weight_load_plan=weight_load_plan,
checkpoint_key_filter=checkpoint_key_filter, checkpoint_key_filter=checkpoint_key_filter,
weights_iterator=(
gguf_weights_iterator(
quant_spec.gguf_file,
quant_spec.quant_config.tensor_meta,
key_filter=checkpoint_key_filter,
)
if quant_spec.gguf_file is not None
else None
),
) )
# post-hooks (e.g., patch scales (nunchaku)) # post-hooks (e.g., patch scales (nunchaku))
@@ -238,6 +238,7 @@ def maybe_load_fsdp_model(
strict: bool = True, strict: bool = True,
weight_load_plan: WeightLoadPlan | None = None, weight_load_plan: WeightLoadPlan | None = None,
checkpoint_key_filter: Callable[[str], bool] | None = None, checkpoint_key_filter: Callable[[str], bool] | None = None,
weights_iterator: Generator[tuple[str, torch.Tensor], None, None] | None = None,
) -> torch.nn.Module: ) -> torch.nn.Module:
"""Load a model with optional FSDP (Fully Sharded Data Parallel) support. """Load a model with optional FSDP (Fully Sharded Data Parallel) support.
@@ -255,6 +256,9 @@ def maybe_load_fsdp_model(
Runtime residency strategies move it to the compute device before use. Runtime residency strategies move it to the compute device before use.
strict: If True, enforce strict state dict loading (all keys must match). strict: If True, enforce strict state dict loading (all keys must match).
weight_load_plan: Optional checkpoint/postprocess device plan for this load. weight_load_plan: Optional checkpoint/postprocess device plan for this load.
weights_iterator: Optional pre-built ``(name, tensor)`` source, used
instead of reading ``weight_dir_list`` as safetensors. Set by callers
whose checkpoint is not safetensors at all, such as GGUF.
""" """
# NOTE(will): cast_forward_inputs=True shouldn't be needed as we are # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
# manually casting the inputs to the model # manually casting the inputs to the model
@@ -354,6 +358,7 @@ def maybe_load_fsdp_model(
not weight_load_plan.load_full_state_dict_on_device not weight_load_plan.load_full_state_dict_on_device
and use_fsdp and use_fsdp
and weight_dir_list and weight_dir_list
and weights_iterator is None
and preprocess_loaded_state_dict is None and preprocess_loaded_state_dict is None
and checkpoint_key_filter is None and checkpoint_key_filter is None
and not is_bnb_quantized and not is_bnb_quantized
@@ -369,6 +374,7 @@ def maybe_load_fsdp_model(
not weight_load_plan.load_full_state_dict_on_device not weight_load_plan.load_full_state_dict_on_device
and not use_fsdp and not use_fsdp
and weight_dir_list and weight_dir_list
and weights_iterator is None
and preprocess_loaded_state_dict is None and preprocess_loaded_state_dict is None
and checkpoint_key_filter is None and checkpoint_key_filter is None
and not is_bnb_quantized and not is_bnb_quantized
@@ -382,7 +388,9 @@ def maybe_load_fsdp_model(
) )
if preconverted_state_dict is None: if preconverted_state_dict is None:
if weight_load_plan.load_full_state_dict_on_device: if weights_iterator is not None:
weight_iterator = weights_iterator
elif weight_load_plan.load_full_state_dict_on_device:
weight_iterator = safetensors_weights_iterator( weight_iterator = safetensors_weights_iterator(
weight_dir_list, weight_dir_list,
key_filter=checkpoint_key_filter, key_filter=checkpoint_key_filter,
@@ -0,0 +1,185 @@
# SPDX-License-Identifier: Apache-2.0
"""Diffusion-specific GGUF tensor layout and iteration."""
from __future__ import annotations
import math
import os
import warnings
from collections.abc import Callable, Generator
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import torch
from sglang.srt.utils.hf_transformers import check_gguf_file
if TYPE_CHECKING:
import gguf
from gguf import GGMLQuantizationType as WeightType
_GGML_F32, _GGML_F16, _GGML_BF16 = 0, 1, 30
_UNQUANTIZED_TYPES = {_GGML_F32, _GGML_F16, _GGML_BF16}
# SRT has no batched MMQ kernel for I-matrix types and may dequantize them.
_SUPER_BLOCK_DEQUANT_TYPES = {16, 17, 18, 19, 20, 21, 22, 23, 29}
_GGML_SUPER_BLOCK = 256
@dataclass(frozen=True)
class GGUFTensorMeta:
"""Logical and packed layouts read before constructing a DiT."""
ggml_type: int
logical_shape: tuple[int, ...]
stored_shape: tuple[int, ...]
stored_dtype: torch.dtype
param_name: str
@property
def weight_type(self) -> WeightType:
from gguf import GGMLQuantizationType as WeightType
return WeightType(self.ggml_type)
@property
def is_quantized(self) -> bool:
return self.ggml_type not in _UNQUANTIZED_TYPES
def _gguf_module() -> Any:
try:
import gguf
except ImportError as exc:
raise ImportError(
"Reading a GGUF checkpoint requires the `gguf` package"
) from exc
return gguf
def _open_reader(gguf_file: str) -> gguf.GGUFReader:
gguf = _gguf_module()
try:
reader = gguf.GGUFReader(gguf_file)
except Exception as exc:
size = os.path.getsize(gguf_file) if os.path.isfile(gguf_file) else 0
raise ValueError(
f"Failed to read GGUF {gguf_file} ({size} bytes). An incomplete or "
f"corrupt download is the usual cause. Underlying error: {exc}"
) from exc
if reader.byte_order == "S":
raise ValueError(
f"GGUF file {gguf_file} uses the opposite byte order from this host"
)
return reader
def read_gguf_tensor_meta(gguf_file: str) -> dict[str, GGUFTensorMeta]:
"""Read the exact packed shape required by diffusion parameters."""
gguf = _gguf_module()
WeightType = gguf.GGMLQuantizationType
reader = _open_reader(gguf_file)
metadata: dict[str, GGUFTensorMeta] = {}
for tensor in reader.tensors:
weight_type = WeightType(tensor.tensor_type)
logical_shape = tuple(int(dim) for dim in reversed(tensor.shape))
is_quantized = int(weight_type) not in _UNQUANTIZED_TYPES
if is_quantized:
if len(logical_shape) != 2 or not tensor.name.endswith(".weight"):
raise ValueError(
f"GGUF tensor {tensor.name} is quantized, but diffusion GGUF "
"currently supports packed data only for 2D linear .weight "
"tensors"
)
block_size, type_size = gguf.GGML_QUANT_SIZES[weight_type]
inner_dim = logical_shape[-1]
if inner_dim % block_size:
raise ValueError(
f"GGUF tensor {tensor.name} has inner dimension {inner_dim}, "
f"which is not a multiple of block size {block_size}"
)
stored_shape = (
*logical_shape[:-1],
inner_dim // block_size * type_size,
)
if (
int(weight_type) in _SUPER_BLOCK_DEQUANT_TYPES
and math.prod(logical_shape) % _GGML_SUPER_BLOCK
):
raise ValueError(
f"GGUF tensor {tensor.name} is not aligned to "
f"{_GGML_SUPER_BLOCK}-element super blocks"
)
stored_dtype = torch.uint8
else:
stored_shape = logical_shape
stored_dtype = {
_GGML_F32: torch.float32,
_GGML_F16: torch.float16,
_GGML_BF16: torch.bfloat16,
}[int(weight_type)]
param_name = (
f"{tensor.name.removesuffix('.weight')}.qweight"
if is_quantized
else tensor.name
)
metadata[tensor.name] = GGUFTensorMeta(
ggml_type=int(weight_type),
logical_shape=logical_shape,
stored_shape=stored_shape,
stored_dtype=stored_dtype,
param_name=param_name,
)
return metadata
def _tensor_to_torch(tensor, metadata: GGUFTensorMeta) -> torch.Tensor:
with warnings.catch_warnings():
warnings.filterwarnings(
"ignore",
message="The given NumPy array is not writable",
category=UserWarning,
)
value = torch.from_numpy(tensor.data)
if metadata.ggml_type == _GGML_BF16:
return value.view(torch.bfloat16).reshape(metadata.stored_shape).clone()
value = value.reshape(metadata.stored_shape)
return value.clone() if not metadata.is_quantized else value
def gguf_weights_iterator(
gguf_file: str,
tensor_meta: dict[str, GGUFTensorMeta],
key_filter: Callable[[str], bool] | None = None,
) -> Generator[tuple[str, torch.Tensor], None, None]:
"""Yield checkpoint tensors under their diffusion parameter names."""
reader = _open_reader(gguf_file)
for tensor in reader.tensors:
if key_filter is not None and not key_filter(tensor.name):
continue
metadata = tensor_meta[tensor.name]
yield metadata.param_name, _tensor_to_torch(tensor, metadata)
def names_gguf_checkpoint(reference: str) -> bool:
"""Recognize an explicit local or Hub GGUF reference without downloading."""
if not reference:
return False
if check_gguf_file(reference):
return True
if os.path.exists(reference):
return False
if os.path.isabs(reference) or reference.startswith((".", "~")):
return reference.endswith(".gguf")
if ":" in reference:
repo_id, _, quant_type = reference.rpartition(":")
return repo_id.count("/") == 1 and bool(quant_type)
return reference.endswith(".gguf") and len(reference.strip("/").split("/")) >= 3
__all__ = [
"GGUFTensorMeta",
"gguf_weights_iterator",
"names_gguf_checkpoint",
"read_gguf_tensor_meta",
]
@@ -22,6 +22,10 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
NunchakuConfig, NunchakuConfig,
_patch_nunchaku_scales, _patch_nunchaku_scales,
) )
from sglang.multimodal_gen.runtime.loader.gguf_weights import (
names_gguf_checkpoint,
read_gguf_tensor_meta,
)
from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files from sglang.multimodal_gen.runtime.loader.utils import _list_safetensors_files
from sglang.multimodal_gen.runtime.loader.weight_utils import ( from sglang.multimodal_gen.runtime.loader.weight_utils import (
filter_duplicate_safetensors_files, filter_duplicate_safetensors_files,
@@ -30,6 +34,7 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency
COMPONENT_OFFLOAD, COMPONENT_OFFLOAD,
ComponentResidencyError, ComponentResidencyError,
) )
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import ( from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
maybe_download_model, maybe_download_model,
@@ -43,6 +48,10 @@ from sglang.multimodal_gen.runtime.utils.quantization_utils import (
get_quant_config, get_quant_config,
get_quant_config_from_safetensors_metadata, get_quant_config_from_safetensors_metadata,
) )
from sglang.srt.utils.hf_transformers import (
check_gguf_file,
resolve_hf_gguf_reference,
)
logger = init_logger(__name__) logger = init_logger(__name__)
@@ -132,6 +141,8 @@ class TransformerQuantLoadSpec:
param_dtype: Optional[torch.dtype] param_dtype: Optional[torch.dtype]
needs_device_weight_postprocess: bool = False needs_device_weight_postprocess: bool = False
post_load_hooks: list[PostLoadHook] = field(default_factory=list) post_load_hooks: list[PostLoadHook] = field(default_factory=list)
# Set instead of ``safetensors_list`` when the transformer comes from GGUF.
gguf_file: Optional[str] = None
@property @property
def runtime_quant_config(self) -> Optional[object]: def runtime_quant_config(self) -> Optional[object]:
@@ -440,6 +451,114 @@ class _BitsAndBytes4BitAdapter(_TransformerQuantAdapter):
) )
def _validate_gguf_runtime_support(
server_args: ServerArgs, component_name: str | None = None
) -> None:
"""Reject configurations a GGUF transformer cannot serve.
Called before the checkpoint is downloaded or read, so an unsupported
combination costs a second rather than a multi-gigabyte fetch.
``component_name`` selects the FSDP decision to check. FSDP is resolved per
component, so a globally enabled ``--use-fsdp-inference`` does not shard a
transformer that is offloaded; only the component actually holding the
packed weights matters.
"""
# The quantization comes from the file, so an explicit --quantization is
# either redundant (gguf) or a conflicting request that would otherwise be
# dropped without a word.
if server_args.quantization == "gguf":
raise ValueError(
"GGUF is selected by passing the checkpoint itself, not "
"`--quantization gguf`. Drop the flag; "
"`--transformer-weights-path <file.gguf>` is what enables it."
)
if server_args.quantization is not None:
raise ValueError(
f"--quantization {server_args.quantization} cannot be combined with "
"a GGUF transformer, whose quantization is fixed by the checkpoint. "
"Drop the flag, or use an unquantized checkpoint to quantize online."
)
# Nunchaku shares --transformer-weights-path with GGUF, and the GGUF plan is
# resolved first, so without this the SVDQuant request would be dropped in
# silence rather than refused.
if server_args.nunchaku_config is not None:
raise ValueError(
"--enable-svdquant cannot be combined with a GGUF transformer: both "
"supply the transformer weights. Point "
"--transformer-weights-path at either an SVDQuant checkpoint or a "
".gguf, not one while requesting the other."
)
if not current_platform.is_cuda():
raise ValueError(
"GGUF diffusion checkpoints require CUDA; the GGML kernels have no "
f"{current_platform.device_type} implementation."
)
uses_fsdp = (
server_args.should_use_fsdp_for_component(component_name)
if component_name is not None
else server_args.use_fsdp_inference
)
if uses_fsdp:
raise ValueError(
"GGUF diffusion checkpoints are incompatible with FSDP inference. "
"Run without --use-fsdp-inference, or keep this component offloaded "
"so FSDP does not manage it."
)
if server_args.lora_path is not None:
raise ValueError(
"LoRA is not supported on a GGUF transformer: an adapter cannot be "
"merged into packed GGML blocks. Use the unquantized checkpoint to "
"serve LoRA."
)
# H3's AdaLN paths read the transformer's safetensors directly -- the cache
# builder needs unquantized weights, and the online rebuild is handed the
# safetensors file list, which is empty for a GGUF load.
if server_args.minimax_h3_adaln_online:
raise ValueError(
"--minimax-h3-adaln-online rebuilds AdaLN outputs from the "
"safetensors checkpoint and cannot read a GGUF transformer."
)
if server_args.minimax_h3_adaln_cache_path is not None:
raise ValueError(
"--minimax-h3-adaln-cache-path requires the unquantized "
"transformer and cannot be combined with a GGUF checkpoint."
)
def resolve_transformer_gguf_to_load(
server_args: ServerArgs, component_name: str | None = None
) -> Optional[str]:
"""Resolve ``--transformer-weights-path`` to a local ``.gguf``, if it is one.
Returns ``None`` when the override is absent or is not GGUF, so the caller
falls through to the safetensors path.
"""
override = server_args.transformer_weights_path
if not override:
return None
# A `~` can reach us unexpanded from a config file or a quoted argument.
override = os.path.expanduser(override)
if not names_gguf_checkpoint(override):
return None
# Before any download: a Hub reference would otherwise fetch gigabytes and
# only then hit an unsupported-configuration error.
_validate_gguf_runtime_support(server_args, component_name)
is_local_reference = os.path.isabs(override) or override.startswith(".")
resolved = (
override
if is_local_reference
else resolve_hf_gguf_reference(override, revision=server_args.revision)
or override
)
if not check_gguf_file(resolved):
raise ValueError(f"Resolved GGUF path is not a GGUF file: {resolved}")
logger.info("using GGUF transformer weights from: %s", resolved)
return resolved
def resolve_transformer_safetensors_to_load( def resolve_transformer_safetensors_to_load(
server_args: ServerArgs, component_model_path: str server_args: ServerArgs, component_model_path: str
) -> list[str]: ) -> list[str]:
@@ -574,7 +693,16 @@ def resolve_transformer_quant_load_spec(
model_cls: type[nn.Module], model_cls: type[nn.Module],
cls_name: str, cls_name: str,
component_name: str | None = None, component_name: str | None = None,
gguf_file: str | None = None,
) -> TransformerQuantLoadSpec: ) -> TransformerQuantLoadSpec:
if gguf_file is not None:
return _resolve_gguf_quant_load_spec(
gguf_file=gguf_file,
server_args=server_args,
model_cls=model_cls,
component_name=component_name,
)
if getattr(model_cls, "handles_checkpoint_quantization", False): if getattr(model_cls, "handles_checkpoint_quantization", False):
quant_config = None quant_config = None
else: else:
@@ -626,6 +754,40 @@ def resolve_transformer_quant_load_spec(
) )
def _resolve_gguf_quant_load_spec(
*,
gguf_file: str,
server_args: ServerArgs,
model_cls: type[nn.Module],
component_name: str | None = None,
) -> TransformerQuantLoadSpec:
"""Build the load plan for a GGUF transformer checkpoint."""
from sglang.multimodal_gen.runtime.layers.quantization.gguf import GGUFConfig
_validate_gguf_runtime_support(server_args, component_name)
quant_config = GGUFConfig(
gguf_file=gguf_file,
tensor_meta=read_gguf_tensor_meta(gguf_file),
)
packed = getattr(model_cls, "packed_modules_mapping", None)
if packed:
quant_config.packed_modules_mapping = packed
return TransformerQuantLoadSpec(
safetensors_list=[],
quant_config=quant_config,
nunchaku_config=None,
# No single dtype for the load: each parameter keeps the dtype the model
# declared for it, which the generic loader casts to. Packed weights are
# registered uint8, so that cast is a no-op for them. Note this matches
# every other quant path -- _resolve_target_param_dtype returns None
# whenever a quant_config is present.
param_dtype=None,
gguf_file=gguf_file,
)
def _needs_device_weight_postprocess( def _needs_device_weight_postprocess(
quant_config: Optional[QuantizationConfig], quant_config: Optional[QuantizationConfig],
) -> bool: ) -> bool:
@@ -742,6 +904,16 @@ def _resolve_quant_config(
if server_args.quantization == "modelslim": if server_args.quantization == "modelslim":
return get_quant_config(hf_config, component_model_path) return get_quant_config(hf_config, component_model_path)
# GGUF is selected by pointing at the file, not by this flag: the config
# has to be built from that file's header.
if server_args.quantization == "gguf":
raise ValueError(
"GGUF is selected by passing the checkpoint itself, not "
"`--quantization gguf`. Use "
"`--transformer-weights-path <file.gguf>` (or a Hub reference "
"such as owner/repo:Q4_K_M)."
)
# Online-quant convention: for `fp8`, `mxfp4` and `kitchen_int8`, a # Online-quant convention: for `fp8`, `mxfp4` and `kitchen_int8`, a
# no-arg QuantizationConfig() selects the post-load path -- weights # no-arg QuantizationConfig() selects the post-load path -- weights
# load in source dtype and are quantized in # load in source dtype and are quantized in
@@ -579,7 +579,11 @@ class MiniMaxH3Attention(nn.Module):
quant_config=quant_config, quant_config=quant_config,
prefix=f"{prefix}.qkv_proj", prefix=f"{prefix}.qkv_proj",
) )
self._install_qkv_weight_loader(arch) # 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":
self._install_qkv_weight_loader(arch)
self.q_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps) 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) self.k_norm = _norm(arch.attention_head_dim, eps=arch.qk_norm_eps)
# cache width covers cos/sin for temporal, height, and width frequencies # cache width covers cos/sin for temporal, height, and width frequencies
@@ -882,7 +886,9 @@ class MiniMaxH3MLP(nn.Module):
quant_config=quant_config, quant_config=quant_config,
prefix=f"{prefix}.fc2", prefix=f"{prefix}.fc2",
) )
self.reuse_fc1_activation = quant_config is None self.reuse_fc1_activation = quant_config is None or (
quant_config.get_name() == "gguf"
)
def forward(self, x: torch.Tensor) -> torch.Tensor: def forward(self, x: torch.Tensor) -> torch.Tensor:
if x.device.type == "mps": if x.device.type == "mps":
@@ -931,12 +937,17 @@ class MiniMaxH3AdalnProj(nn.Module):
self.expand_ratio = expand_ratio self.expand_ratio = expand_ratio
self.modality_num = modality_num self.modality_num = modality_num
self.hidden_size = arch.hidden_size self.hidden_size = arch.hidden_size
# Curve checkpoints store both the sampled curve and their reduced
# AdaLN projections in FP32. Preserve that precision island to match
# the published pruned implementation; these outputs intentionally do
# not enter the BF16-only fused modulation kernels.
params_dtype = _FP32_DTYPE if arch.adaln_curve_grid is not None else _BF16_DTYPE
self.linear = ColumnParallelLinear( self.linear = ColumnParallelLinear(
arch.time_embed_dim, arch.time_embed_dim,
out_features, out_features,
bias=True, bias=True,
gather_output=False, gather_output=False,
params_dtype=_BF16_DTYPE, params_dtype=params_dtype,
quant_config=quant_config, quant_config=quant_config,
prefix=f"{prefix}.linear", prefix=f"{prefix}.linear",
) )
@@ -951,7 +962,7 @@ class MiniMaxH3AdalnProj(nn.Module):
return tuple(x.chunk(self.expand_ratio, dim=-1)) return tuple(x.chunk(self.expand_ratio, dim=-1))
def forward(self, adaln_input: torch.Tensor) -> tuple[torch.Tensor, ...]: def forward(self, adaln_input: torch.Tensor) -> tuple[torch.Tensor, ...]:
"""adaln_input: SiLU(t_emb) BF16 -> expand_ratio tensors of [M*modality_num, H].""" """Project the post-SiLU embedding in its checkpoint-defined dtype."""
x = self.project_local(adaln_input) x = self.project_local(adaln_input)
if get_tp_world_size() > 1: if get_tp_world_size() > 1:
x = tensor_model_parallel_all_gather(x) x = tensor_model_parallel_all_gather(x)
@@ -1659,16 +1670,23 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
adaln_plan_width: int = MINIMAX_H3_ADALN_MAX_PLAN_WIDTH, adaln_plan_width: int = MINIMAX_H3_ADALN_MAX_PLAN_WIDTH,
) -> None: ) -> None:
super().__init__(config=config, hf_config=hf_config) super().__init__(config=config, hf_config=hf_config)
arch = self.config
if ( if (
adaln_cache_path is not None or adaln_weight_files is not None adaln_cache_path is not None or adaln_weight_files is not None
) and quant_config is not None: ) and quant_config is not None:
raise ValueError( raise ValueError(
"MiniMax H3 AdaLN cache is only compatible with unquantized weights" "MiniMax H3 AdaLN cache is only compatible with unquantized weights"
) )
if arch.adaln_curve_grid is not None and (
adaln_cache_path is not None or adaln_weight_files is not None
):
raise ValueError(
"MiniMax H3 pruned curve checkpoints cannot use a separate "
"AdaLN cache"
)
self._adaln_precomputed = ( self._adaln_precomputed = (
adaln_cache_path is not None or adaln_weight_files is not None adaln_cache_path is not None or adaln_weight_files is not None
) )
arch = self.config
self.arch = arch self.arch = arch
self.hidden_size = arch.hidden_size self.hidden_size = arch.hidden_size
self.num_attention_heads = arch.num_attention_heads self.num_attention_heads = arch.num_attention_heads
@@ -1713,10 +1731,22 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
quant_config=quant_config, quant_config=quant_config,
prefix="condition_proj", prefix="condition_proj",
) )
self.time_embedder = MiniMaxH3TimeEmbedder( if arch.adaln_curve_grid is None:
arch, self.time_embedder = MiniMaxH3TimeEmbedder(
prefix="time_embedder", arch,
) prefix="time_embedder",
)
self.register_parameter("adaln_t_table", None)
else:
self.time_embedder = None
self.adaln_t_table = nn.Parameter(
torch.empty(
arch.adaln_curve_grid,
arch.time_embed_dim,
dtype=_FP32_DTYPE,
),
requires_grad=False,
)
self.rope = MiniMaxH3Rope(arch.rope_inv_freq_len) self.rope = MiniMaxH3Rope(arch.rope_inv_freq_len)
self.token_refiner = MiniMaxH3TokenRefiner( self.token_refiner = MiniMaxH3TokenRefiner(
arch, arch,
@@ -1789,12 +1819,26 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
param.missing_param_init = "error" param.missing_param_init = "error"
def post_load_weights(self) -> None: def post_load_weights(self) -> None:
for name in _MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER: fp32_param_names = list(_MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER)
if self.adaln_t_table is not None:
fp32_param_names = [
name
for name in fp32_param_names
if not name.startswith("time_embedder.")
]
fp32_param_names.append("adaln_t_table")
for name in fp32_param_names:
param = self.get_parameter(name) param = self.get_parameter(name)
if param.dtype != _FP32_DTYPE: if param.dtype != _FP32_DTYPE:
raise ValueError( raise ValueError(
f"{name} must stay fp32 after load, got {param.dtype}." f"{name} must stay fp32 after load, got {param.dtype}."
) )
if self.adaln_t_table is not None:
for name, param in self.named_parameters():
if ".adaln_proj.linear." in name and param.dtype != _FP32_DTYPE:
raise ValueError(
f"{name} must stay fp32 with curve AdaLN, got {param.dtype}."
)
# assign=True loading may re-register this persistent buffer as a parameter # assign=True loading may re-register this persistent buffer as a parameter
rope_inv_freq = self.rope.inv_freq rope_inv_freq = self.rope.inv_freq
if rope_inv_freq.dtype != _FP32_DTYPE: if rope_inv_freq.dtype != _FP32_DTYPE:
@@ -1804,6 +1848,19 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
if self.adaln_cache is not None: if self.adaln_cache is not None:
self.adaln_cache.load(self.video_patch_proj.weight.device) self.adaln_cache.load(self.video_patch_proj.weight.device)
def _time_embedding(self, timesteps: torch.Tensor) -> torch.Tensor:
if self.adaln_t_table is None:
assert self.time_embedder is not None
return self.time_embedder(timesteps)
grid = self.adaln_t_table.shape[0]
position = timesteps.to(_FP32_DTYPE).clamp(0, 1) * (grid - 1)
lower = position.floor().clamp(max=grid - 2).to(torch.long)
fraction = (position - lower).unsqueeze(-1)
lower_value = self.adaln_t_table.index_select(0, lower)
upper_value = self.adaln_t_table.index_select(0, lower + 1)
return torch.lerp(lower_value, upper_value, fraction)
@staticmethod @staticmethod
def _pos_ids(pos_info: Any, key: str) -> torch.Tensor: def _pos_ids(pos_info: Any, key: str) -> torch.Tensor:
if isinstance(pos_info, dict): if isinstance(pos_info, dict):
@@ -2059,7 +2116,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
audio_embed.to(_BF16_DTYPE), audio_embed.to(_BF16_DTYPE),
) )
t_emb = self.time_embedder(unique_timesteps) t_emb = self._time_embedding(unique_timesteps)
return embeddings, t_emb return embeddings, t_emb
def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]: def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]:
@@ -2219,7 +2276,11 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
) )
self.release_mps_non_layer_weights(*_MPS_EMBED_WEIGHT_PREFIXES) self.release_mps_non_layer_weights(*_MPS_EMBED_WEIGHT_PREFIXES)
# request-step AdaLN input shared by all blocks # request-step AdaLN input shared by all blocks
adaln_input = nn.functional.silu(t_emb).to(_BF16_DTYPE) adaln_input = (
t_emb
if self.adaln_t_table is not None
else nn.functional.silu(t_emb).to(_BF16_DTYPE)
)
inverse_indices = inverse_indices.to(device) inverse_indices = inverse_indices.to(device)
block_inverse = inverse_indices[row_start:row_stop] block_inverse = inverse_indices[row_start:row_stop]
if block_token_tags is None: if block_token_tags is None:
@@ -184,7 +184,7 @@ class LoRAPipeline(ComposedPipelineBase):
) # type: ignore ) # type: ignore
def is_target_layer(self, module_name: str) -> bool: def is_target_layer(self, module_name: str) -> bool:
if self.lora_target_modules is None: if getattr(self, "lora_target_modules", None) is None:
return True return True
return any( return any(
target_name in module_name for target_name in self.lora_target_modules target_name in module_name for target_name in self.lora_target_modules
@@ -348,12 +348,38 @@ class LoRAPipeline(ComposedPipelineBase):
return converted_count return converted_count
def _reject_lora_on_packed_weights(self) -> None:
"""Fail before any layer is replaced if a target has no plain weight.
``BaseLayerWithLoRA`` reads ``base_layer.weight``, which a
weight-packing quantization (GGUF) does not expose -- it registers
``qweight``. Checking up front keeps a rejected request from leaving the
model half converted, and covers the dynamic ``set_lora`` API as well as
a startup ``--lora-path``.
"""
for module_name in ("transformer", "transformer_2"):
module = self.modules.get(module_name)
if module is None:
continue
for name, layer in module.named_modules():
if not self.is_target_layer(name):
continue
params = dict(layer.named_parameters(recurse=False))
if "weight" not in params and "qweight" in params:
raise ValueError(
f"LoRA is not supported on {module_name}.{name}: its "
"weights are stored packed (GGUF), which an adapter "
"cannot be merged into or applied alongside. Serve the "
"unquantized checkpoint to use LoRA."
)
def convert_to_lora_layers(self) -> None: def convert_to_lora_layers(self) -> None:
""" """
Unified method to convert the transformer to a LoRA transformer. Unified method to convert the transformer to a LoRA transformer.
""" """
if self.lora_initialized: if self.lora_initialized:
return return
self._reject_lora_on_packed_weights()
self.lora_initialized = True self.lora_initialized = True
# Convert transformer # Convert transformer
@@ -896,6 +922,11 @@ class LoRAPipeline(ComposedPipelineBase):
f"Invalid target(s): {invalid_targets}. Valid targets: {self.VALID_TARGETS}" f"Invalid target(s): {invalid_targets}. Valid targets: {self.VALID_TARGETS}"
) )
# Checked before disabling offload, which materializes every layer: on a
# memory-constrained deployment that would OOM instead of returning the
# unsupported-LoRA error. Offloaded placeholders still carry the name.
self._reject_lora_on_packed_weights()
# Disable layerwise offload before convert_to_lora_layers to ensure weights are accessible # Disable layerwise offload before convert_to_lora_layers to ensure weights are accessible
# This is critical because convert_to_lora_layers needs to save cpu_weight from actual weights, # This is critical because convert_to_lora_layers needs to save cpu_weight from actual weights,
# not from offloaded placeholder tensors # not from offloaded placeholder tensors
@@ -299,7 +299,7 @@ class ServerArgs(DisaggServerArgsMixin):
# Optional LTX-2.5 decoder is large enough to load only when requested. # Optional LTX-2.5 decoder is large enough to load only when requested.
load_diffusion_decoder: bool = False load_diffusion_decoder: bool = False
# path to pre-quantized transformer weights (single .safetensors or directory). # Pre-quantized transformer weights: safetensors file/directory or GGUF file.
transformer_weights_path: str | None = None transformer_weights_path: str | None = None
# path to precomputed MiniMax H3 AdaLN outputs for inference-only serving. # path to precomputed MiniMax H3 AdaLN outputs for inference-only serving.
minimax_h3_adaln_cache_path: str | None = None minimax_h3_adaln_cache_path: str | None = None
@@ -53,6 +53,7 @@ from sglang.multimodal_gen.runtime.utils.quantization_utils import (
normalize_flat_modelopt_quant_config, normalize_flat_modelopt_quant_config,
) )
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.utils.hf_transformers import check_gguf_file
from sglang.utils import is_in_ci from sglang.utils import is_in_ci
logger = init_logger(__name__) logger = init_logger(__name__)
@@ -597,19 +598,6 @@ def attach_additional_stop_token_ids(tokenizer):
tokenizer.additional_stop_token_ids = None tokenizer.additional_stop_token_ids = None
def check_gguf_file(model: str | os.PathLike) -> bool:
"""Check if the file is a GGUF model."""
model = Path(model)
if not model.is_file():
return False
elif model.suffix == ".gguf":
return True
with open(model, "rb") as f:
header = f.read(4)
return header == b"GGUF"
def maybe_download_lora( def maybe_download_lora(
model_name_or_path: str, model_name_or_path: str,
local_dir: str | None = None, local_dir: str | None = None,
@@ -0,0 +1,793 @@
"""CPU unit tests for the diffusion GGUF load path.
These cover the two things that must hold before any GPU run is meaningful:
the header-derived tensor layout (which is what lets the generic weight loader
work unchanged) and the quant-method selection per layer.
"""
import struct
import sys
import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch
import numpy as np
import torch
from gguf import GGMLQuantizationType as WeightType
from sglang.multimodal_gen.runtime.layers.linear import (
ColumnParallelLinear,
MergedColumnParallelLinear,
ReplicatedLinear,
RowParallelLinear,
UnquantizedLinearMethod,
)
from sglang.multimodal_gen.runtime.layers.quantization.gguf import (
GGUFConfig,
GGUFLinearMethod,
)
from sglang.multimodal_gen.runtime.loader.gguf_weights import (
GGUFTensorMeta,
gguf_weights_iterator,
names_gguf_checkpoint,
read_gguf_tensor_meta,
)
from sglang.srt.layers.quantization.gguf import UNQUANTIZED_TYPES
from sglang.srt.utils.hf_transformers import check_gguf_file
_F32 = WeightType.F32
_BF16 = WeightType.BF16
_Q4_K = WeightType.Q4_K
_Q4_K_BLOCK, _Q4_K_TYPE_SIZE = 256, 144
def _kv_string(key: str, value: str, bo: str = "<") -> bytes:
out = struct.pack(f"{bo}Q", len(key)) + key.encode()
out += struct.pack(f"{bo}I", 8) # value type: string
out += struct.pack(f"{bo}Q", len(value)) + value.encode()
return out
def _write_gguf(
path: Path,
tensors: list[tuple[str, list[int], int, bytes]],
byte_order: str = "<",
) -> None:
"""Write a minimal GGUF v3 file containing ``tensors``.
Each entry is ``(name, ne_dims, ggml_type, payload)`` where ``ne_dims`` is in
GGUF order (fastest-varying first). ``byte_order`` is a struct prefix; pass
``">"`` on a little-endian host to produce a file gguf-py reports as
swapped.
"""
bo = byte_order
header = b"GGUF" + struct.pack(f"{bo}I", 3)
header += struct.pack(f"{bo}QQ", len(tensors), 1)
header += _kv_string("general.architecture", "test", bo)
# Tensor info blocks, then padded data.
infos = b""
offset = 0
alignment = 32
payloads = []
for name, dims, ggml_type, payload in tensors:
infos += struct.pack(f"{bo}Q", len(name)) + name.encode()
infos += struct.pack(f"{bo}I", len(dims))
infos += b"".join(struct.pack(f"{bo}Q", d) for d in dims)
infos += struct.pack(f"{bo}I", ggml_type)
infos += struct.pack(f"{bo}Q", offset)
payloads.append(payload)
padded = (len(payload) + alignment - 1) // alignment * alignment
offset += padded
body = header + infos
pad = (alignment - len(body) % alignment) % alignment
body += b"\0" * pad
for payload in payloads:
padded = (len(payload) + alignment - 1) // alignment * alignment
body += payload + b"\0" * (padded - len(payload))
path.write_bytes(body)
class TestGGUFTensorMeta(unittest.TestCase):
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.tmp = Path(self._tmp.name)
self.addCleanup(self._tmp.cleanup)
def test_quantized_layout_is_packed_rows(self):
"""A quantized weight is registered as (out_features, row_bytes) uint8."""
out_features, in_features = 4, 512
row_bytes = in_features // _Q4_K_BLOCK * _Q4_K_TYPE_SIZE
payload = bytes(out_features * row_bytes)
path = self.tmp / "q.gguf"
# GGUF stores ne fastest-varying first: [in, out].
_write_gguf(path, [("w.weight", [in_features, out_features], _Q4_K, payload)])
meta = read_gguf_tensor_meta(str(path))["w.weight"]
self.assertEqual(meta.ggml_type, _Q4_K)
self.assertTrue(meta.is_quantized)
# logical_shape is torch order, i.e. reversed from the file.
self.assertEqual(meta.logical_shape, (out_features, in_features))
self.assertEqual(meta.stored_shape, (out_features, row_bytes))
self.assertEqual(meta.stored_dtype, torch.uint8)
# The layer registers `qweight`, so that is what the iterator must yield.
self.assertEqual(meta.param_name, "w.qweight")
def test_unquantized_layout_matches_logical_shape(self):
out_features, in_features = 3, 8
payload = np.zeros((out_features, in_features), dtype=np.float32).tobytes()
path = self.tmp / "f32.gguf"
_write_gguf(path, [("w.weight", [in_features, out_features], _F32, payload)])
meta = read_gguf_tensor_meta(str(path))["w.weight"]
self.assertFalse(meta.is_quantized)
self.assertEqual(meta.logical_shape, (out_features, in_features))
self.assertEqual(meta.stored_shape, (out_features, in_features))
self.assertEqual(meta.stored_dtype, torch.float32)
self.assertEqual(meta.param_name, "w.weight")
def test_pruned_adaln_curve_shape_is_available_before_model_init(self):
grid, width = 1025, 8
path = self.tmp / "pruned.gguf"
_write_gguf(
path,
[
(
"adaln_t_table",
[width, grid],
_F32,
bytes(grid * width * 4),
)
],
)
metadata = read_gguf_tensor_meta(str(path))["adaln_t_table"]
self.assertEqual(metadata.logical_shape, (grid, width))
self.assertFalse(metadata.is_quantized)
def test_non_block_aligned_inner_dim_is_rejected(self):
"""A row that is not a whole number of blocks must not load.
gguf-py validates this first, so the message comes from there rather
than from read_gguf_tensor_meta's own guard; either way the invariant
row_bytes depends on is enforced.
"""
# 100 is not a multiple of the Q4_K block size (256).
path = self.tmp / "bad.gguf"
_write_gguf(path, [("w.weight", [100, 2], _Q4_K, bytes(64))])
with self.assertRaisesRegex(ValueError, "not a multiple of.*block size"):
read_gguf_tensor_meta(str(path))
def test_super_block_element_count_is_enforced(self):
"""IQ4_NL has 32-element blocks but a 256-element super-block kernel with
no bounds check, so a tensor that is not a whole number of super blocks
would read or write past the buffer."""
_IQ4_NL, block, type_size = 20, 32, 18
# in=1440 is a multiple of 32 (so the row check passes) and out=4 makes
# numel=5760, which is not a multiple of 256.
out_features, in_features = 4, 1440
row_bytes = in_features // block * type_size
path = self.tmp / "iq4nl.gguf"
_write_gguf(
path,
[
(
"w.weight",
[in_features, out_features],
_IQ4_NL,
bytes(row_bytes * out_features),
)
],
)
with self.assertRaisesRegex(ValueError, "super\\s+blocks"):
read_gguf_tensor_meta(str(path))
def test_standard_quant_type_does_not_require_super_block_alignment(self):
"""Q4_0 has native MMVQ/MMQ kernels, so 256 must not be required."""
_Q4_0, block, type_size = 2, 32, 18
out_features, in_features = 1, 32 # numel 32: not a super block
row_bytes = in_features // block * type_size
path = self.tmp / "q40.gguf"
_write_gguf(
path,
[
(
"w.weight",
[in_features, out_features],
_Q4_0,
bytes(row_bytes * out_features),
)
],
)
meta = read_gguf_tensor_meta(str(path))["w.weight"]
self.assertEqual(meta.stored_shape, (out_features, row_bytes))
def test_bf16_is_reinterpreted_not_cast(self):
"""gguf-py returns BF16 as raw bytes; a cast would corrupt the values."""
values = torch.tensor([1.0, -2.5, 3.75], dtype=torch.bfloat16)
payload = values.view(torch.uint8).numpy().tobytes()
path = self.tmp / "bf16.gguf"
_write_gguf(path, [("norm.weight", [3], _BF16, payload)])
meta = read_gguf_tensor_meta(str(path))
loaded = dict(gguf_weights_iterator(str(path), meta))["norm.weight"]
self.assertEqual(loaded.dtype, torch.bfloat16)
torch.testing.assert_close(loaded, values)
def test_iterator_yields_stored_shapes(self):
out_features, in_features = 2, 256
row_bytes = in_features // _Q4_K_BLOCK * _Q4_K_TYPE_SIZE
payload = bytes(range(row_bytes)) * out_features
path = self.tmp / "iter.gguf"
_write_gguf(path, [("w.weight", [in_features, out_features], _Q4_K, payload)])
meta = read_gguf_tensor_meta(str(path))
loaded = dict(gguf_weights_iterator(str(path), meta))["w.qweight"]
self.assertEqual(loaded.dtype, torch.uint8)
self.assertEqual(tuple(loaded.shape), (out_features, row_bytes))
def test_swapped_endian_file_is_rejected(self):
"""A quantized block embeds its scales, so a swapped file cannot just be
byte-swapped whole; refuse it instead of dequantizing garbage."""
if sys.byteorder != "little":
self.skipTest("test builds a big-endian file to be the swapped one")
path = self.tmp / "be.gguf"
_write_gguf(path, [("w.weight", [4], _F32, bytes(16))], byte_order=">")
with self.assertRaisesRegex(ValueError, "opposite byte order"):
read_gguf_tensor_meta(str(path))
def test_truncated_file_names_the_likely_cause(self):
"""An interrupted download is common for a multi-GiB file; the bare
reshape error gguf-py raises reads like a layout bug instead."""
row_bytes = 256 // _Q4_K_BLOCK * _Q4_K_TYPE_SIZE
full = self.tmp / "full.gguf"
_write_gguf(full, [("w.weight", [256, 4], _Q4_K, bytes(row_bytes * 4))])
# Keep the header, drop most of the tensor data.
cut = self.tmp / "cut.gguf"
cut.write_bytes(full.read_bytes()[: -row_bytes * 3])
with self.assertRaisesRegex(ValueError, "incomplete or\\s+corrupt download"):
read_gguf_tensor_meta(str(cut))
def test_is_gguf_file_detects_by_magic(self):
path = self.tmp / "no-suffix.bin"
_write_gguf(path, [("w.weight", [4], _F32, bytes(16))])
self.assertTrue(check_gguf_file(str(path)))
other = self.tmp / "other.bin"
other.write_bytes(b"NOTGGUF")
self.assertFalse(check_gguf_file(str(other)))
self.assertFalse(check_gguf_file(str(self.tmp / "missing.gguf")))
def test_quantized_non_linear_tensor_is_rejected(self):
path = self.tmp / "bad-norm.gguf"
_write_gguf(path, [("norm.weight", [256], _Q4_K, bytes(144))])
with self.assertRaisesRegex(ValueError, "only for 2D linear"):
read_gguf_tensor_meta(str(path))
class TestGGUFQuantMethodSelection(unittest.TestCase):
def _config(self, **metas):
return GGUFConfig(gguf_file="/dev/null", tensor_meta=dict(metas))
def _meta(self, ggml_type, out_features, in_features, stored_shape=None):
return GGUFTensorMeta(
ggml_type=ggml_type,
logical_shape=(out_features, in_features),
stored_shape=stored_shape or (out_features, in_features),
stored_dtype=(
torch.float32 if ggml_type in UNQUANTIZED_TYPES else torch.uint8
),
param_name=("w.weight" if ggml_type in UNQUANTIZED_TYPES else "w.qweight"),
)
def test_quantized_layer_gets_gguf_method(self):
config = self._config(**{"w.weight": self._meta(_Q4_K, 4, 512, (4, 288))})
layer = ReplicatedLinear(512, 4, bias=False, quant_config=config, prefix="w")
self.assertIsInstance(layer.quant_method, GGUFLinearMethod)
self.assertEqual(layer.qweight.dtype, torch.uint8)
self.assertEqual(tuple(layer.qweight.shape), (4, 288))
self.assertEqual(layer.quant_method.weight_type, _Q4_K)
def test_unquantized_layer_falls_back(self):
"""H3 keeps its FP32 projections unquantized inside the same file."""
config = self._config(**{"w.weight": self._meta(_F32, 4, 8)})
layer = ReplicatedLinear(8, 4, bias=False, quant_config=config, prefix="w")
self.assertIsInstance(layer.quant_method, UnquantizedLinearMethod)
def test_missing_tensor_fails_fast(self):
config = self._config()
with self.assertRaisesRegex(ValueError, "no weight in the GGUF checkpoint"):
ReplicatedLinear(8, 4, bias=False, quant_config=config, prefix="absent")
def test_shape_mismatch_fails_fast(self):
config = self._config(**{"w.weight": self._meta(_Q4_K, 8, 512, (8, 288))})
with self.assertRaisesRegex(ValueError, "logical shape"):
ReplicatedLinear(512, 4, bias=False, quant_config=config, prefix="w")
@patch(
"sglang.multimodal_gen.runtime.layers.quantization.gguf.dequantize_gguf_weight"
)
def test_apply_reuses_srt_dequantization(self, dequantize):
config = self._config(**{"w.weight": self._meta(_Q4_K, 4, 512, (4, 288))})
layer = ReplicatedLinear(512, 4, bias=False, quant_config=config, prefix="w")
dequantize.return_value = torch.ones(4, 512)
output, _ = layer(torch.ones(2, 4, 512))
dequantize.assert_called_once_with(layer.qweight, _Q4_K, torch.float32)
self.assertEqual(tuple(output.shape), (2, 4, 4))
torch.testing.assert_close(output, torch.full_like(output, 512.0))
class TestGGUFTensorParallelLoading(unittest.TestCase):
def setUp(self):
self.group = SimpleNamespace(world_size=2, rank_in_group=1)
self.meta = GGUFTensorMeta(
ggml_type=int(_Q4_K),
logical_shape=(8, 512),
stored_shape=(8, 288),
stored_dtype=torch.uint8,
param_name="w.qweight",
)
self.config = GGUFConfig("/dev/null", {"w.weight": self.meta})
values = torch.arange(8 * 288, dtype=torch.int64).remainder(251)
self.loaded = values.to(torch.uint8).reshape(8, 288)
def test_column_parallel_slices_output_rows(self):
layer = ColumnParallelLinear(
512,
8,
bias=False,
quant_config=self.config,
prefix="w",
tp_group=self.group,
)
layer.weight_loader(layer.qweight, self.loaded)
torch.testing.assert_close(layer.qweight, self.loaded[4:])
def test_row_parallel_slices_packed_input_blocks(self):
layer = RowParallelLinear(
512,
8,
bias=False,
quant_config=self.config,
prefix="w",
tp_group=self.group,
)
layer.weight_loader(layer.qweight, self.loaded)
torch.testing.assert_close(layer.qweight, self.loaded[:, 144:])
def test_merged_column_parallel_slices_each_output_group(self):
layer = MergedColumnParallelLinear(
512,
[4, 4],
bias=False,
quant_config=self.config,
prefix="w",
tp_group=self.group,
)
layer.weight_loader(layer.qweight, self.loaded)
expected = torch.cat((self.loaded[2:4], self.loaded[6:8]))
torch.testing.assert_close(layer.qweight, expected)
def test_row_parallel_rejects_unaligned_partition(self):
metadata = GGUFTensorMeta(
ggml_type=int(_Q4_K),
logical_shape=(8, 256),
stored_shape=(8, 144),
stored_dtype=torch.uint8,
param_name="w.qweight",
)
config = GGUFConfig("/dev/null", {"w.weight": metadata})
with self.assertRaisesRegex(ValueError, "not aligned"):
RowParallelLinear(
256,
8,
bias=False,
quant_config=config,
prefix="w",
tp_group=self.group,
)
class TestGGUFIncompatibleOptions(unittest.TestCase):
"""Combinations that cannot work must fail at startup, not mid-run."""
def setUp(self):
self._tmp = tempfile.TemporaryDirectory()
self.tmp = Path(self._tmp.name)
self.addCleanup(self._tmp.cleanup)
self.gguf = self.tmp / "t.gguf"
_write_gguf(self.gguf, [("w.weight", [8, 2], _F32, bytes(64))])
def _resolve(self, **overrides):
"""Resolve with CUDA mocked in, so each test exercises its own guard.
These are CPU-only tests; without the mock the CUDA guard fires first
and every case would report the wrong reason.
"""
from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
_resolve_gguf_quant_load_spec,
)
server_args = Mock()
server_args.tp_size = 1
server_args.use_fsdp_inference = False
server_args.lora_path = None
server_args.minimax_h3_adaln_online = False
server_args.minimax_h3_adaln_cache_path = None
server_args.quantization = None
server_args.nunchaku_config = None
for key, value in overrides.items():
setattr(server_args, key, value)
with patch(
"sglang.multimodal_gen.runtime.loader.transformer_load_utils.current_platform"
) as platform:
platform.is_cuda.return_value = True
return _resolve_gguf_quant_load_spec(
gguf_file=str(self.gguf),
server_args=server_args,
model_cls=Mock(packed_modules_mapping=None),
)
def test_accepts_tp_without_fsdp(self):
spec = self._resolve(tp_size=2)
self.assertEqual(spec.gguf_file, str(self.gguf))
self.assertEqual(spec.safetensors_list, [])
# Each tensor keeps its checkpoint dtype rather than a single cast.
self.assertIsNone(spec.param_dtype)
def test_rejects_non_cuda_platform(self):
"""Fail before reading a multi-GiB checkpoint, not at the first linear."""
from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
_resolve_gguf_quant_load_spec,
)
server_args = Mock()
server_args.tp_size = 1
server_args.use_fsdp_inference = False
server_args.lora_path = None
server_args.minimax_h3_adaln_online = False
server_args.minimax_h3_adaln_cache_path = None
server_args.quantization = None
server_args.nunchaku_config = None
with patch(
"sglang.multimodal_gen.runtime.loader.transformer_load_utils.current_platform"
) as platform:
platform.is_cuda.return_value = False
platform.device_type = "rocm"
with self.assertRaisesRegex(ValueError, "require CUDA"):
_resolve_gguf_quant_load_spec(
gguf_file=str(self.gguf),
server_args=server_args,
model_cls=Mock(packed_modules_mapping=None),
)
def test_rejects_fsdp_when_this_component_is_fsdp_managed(self):
server_args_kwargs = {"use_fsdp_inference": True}
with self.assertRaisesRegex(ValueError, "FSDP"):
self._resolve(**server_args_kwargs)
def test_allows_global_fsdp_when_this_component_is_offloaded(self):
"""FSDP is per component: an offloaded transformer is never sharded.
Rejecting on the global flag would block FSDP on other resident
components for no reason.
"""
from sglang.multimodal_gen.runtime.loader.transformer_load_utils import (
_resolve_gguf_quant_load_spec,
)
server_args = Mock()
server_args.tp_size = 1
server_args.use_fsdp_inference = True
# ...but not for this component.
server_args.should_use_fsdp_for_component.return_value = False
server_args.lora_path = None
server_args.minimax_h3_adaln_online = False
server_args.minimax_h3_adaln_cache_path = None
server_args.quantization = None
server_args.nunchaku_config = None
with patch(
"sglang.multimodal_gen.runtime.loader.transformer_load_utils.current_platform"
) as platform:
platform.is_cuda.return_value = True
spec = _resolve_gguf_quant_load_spec(
gguf_file=str(self.gguf),
server_args=server_args,
model_cls=Mock(packed_modules_mapping=None),
component_name="transformer",
)
self.assertEqual(spec.gguf_file, str(self.gguf))
server_args.should_use_fsdp_for_component.assert_called_once_with("transformer")
def test_rejects_quantization_gguf_flag(self):
"""The file selects GGUF; the flag would be a second, silent selector."""
with self.assertRaisesRegex(ValueError, "not\\s+`--quantization gguf`"):
self._resolve(quantization="gguf")
def test_rejects_conflicting_quantization_flag(self):
"""--quantization fp8 with a GGUF file must not be silently dropped."""
with self.assertRaisesRegex(ValueError, "cannot be combined"):
self._resolve(quantization="fp8")
def test_rejects_svdquant(self):
"""Nunchaku shares --transformer-weights-path; a silent drop is worse."""
with self.assertRaisesRegex(ValueError, "svdquant"):
self._resolve(nunchaku_config=object())
def test_rejects_lora(self):
with self.assertRaisesRegex(ValueError, "LoRA"):
self._resolve(lora_path="some/adapter")
def test_rejects_h3_adaln_online(self):
with self.assertRaisesRegex(ValueError, "adaln-online"):
self._resolve(minimax_h3_adaln_online=True)
def test_rejects_h3_adaln_cache(self):
with self.assertRaisesRegex(ValueError, "adaln-cache-path"):
self._resolve(minimax_h3_adaln_cache_path="/tmp/cache.safetensors")
class TestGGUFKeyFilter(unittest.TestCase):
def test_key_filter_matches_checkpoint_names(self):
"""The filter sees checkpoint names, like the safetensors iterator."""
tmp = tempfile.TemporaryDirectory()
self.addCleanup(tmp.cleanup)
path = Path(tmp.name) / "f.gguf"
row_bytes = 256 // _Q4_K_BLOCK * _Q4_K_TYPE_SIZE
_write_gguf(
path,
[
("keep.weight", [256, 1], _Q4_K, bytes(row_bytes)),
("drop.weight", [256, 1], _Q4_K, bytes(row_bytes)),
],
)
meta = read_gguf_tensor_meta(str(path))
loaded = dict(
gguf_weights_iterator(
str(path), meta, key_filter=lambda name: name.startswith("keep")
)
)
# Filtered on the checkpoint name, yielded under the param name.
self.assertEqual(sorted(loaded), ["keep.qweight"])
class TestGGUFPreDownloadValidation(unittest.TestCase):
"""A Hub reference must be rejected before it is fetched."""
def test_hub_reference_is_recognized_without_io(self):
self.assertTrue(names_gguf_checkpoint("owner/repo:Q4_K_M"))
self.assertTrue(names_gguf_checkpoint("owner/repo/sub/model.gguf"))
self.assertFalse(names_gguf_checkpoint("owner/repo"))
self.assertFalse(names_gguf_checkpoint("owner/repo/model.safetensors"))
self.assertFalse(names_gguf_checkpoint(""))
# A local safetensors override must keep taking the safetensors path.
self.assertFalse(names_gguf_checkpoint("/models/transformer.safetensors"))
def test_missing_local_path_is_not_sent_to_the_hub(self):
"""A typo'd local path must report itself, not become a repo lookup.
Recognition must not depend on directory depth: /a/x.gguf and
/a/b/c/x.gguf are both local paths.
"""
from sglang.multimodal_gen.runtime.loader import transformer_load_utils
server_args = Mock(
transformer_weights_path="/models/missing.gguf",
revision=None,
tp_size=1,
use_fsdp_inference=False,
lora_path=None,
minimax_h3_adaln_online=False,
minimax_h3_adaln_cache_path=None,
quantization=None,
nunchaku_config=None,
)
with (
patch.object(
transformer_load_utils, "resolve_hf_gguf_reference"
) as resolve,
patch.object(transformer_load_utils, "current_platform") as platform,
self.assertRaisesRegex(ValueError, "/models/missing.gguf"),
):
platform.is_cuda.return_value = True
transformer_load_utils.resolve_transformer_gguf_to_load(server_args)
resolve.assert_not_called()
def test_home_relative_path_is_expanded(self):
"""A `~` can arrive unexpanded from a config file or quoted argument."""
import os
from sglang.multimodal_gen.runtime.loader import transformer_load_utils
# Put a real GGUF where an unexpanded "~" would miss it.
tmp = tempfile.TemporaryDirectory()
self.addCleanup(tmp.cleanup)
home = Path(tmp.name) / "home"
(home / "models").mkdir(parents=True)
real = home / "models" / "h3.gguf"
_write_gguf(real, [("w.weight", [4], _F32, bytes(16))])
server_args = Mock()
server_args.transformer_weights_path = "~/models/h3.gguf"
server_args.revision = None
server_args.tp_size = 1
server_args.use_fsdp_inference = False
server_args.lora_path = None
server_args.minimax_h3_adaln_online = False
server_args.minimax_h3_adaln_cache_path = None
server_args.quantization = None
server_args.nunchaku_config = None
with (
patch.dict(os.environ, {"HOME": str(home)}),
patch.object(transformer_load_utils, "current_platform") as platform,
):
platform.is_cuda.return_value = True
resolved = transformer_load_utils.resolve_transformer_gguf_to_load(
server_args
)
self.assertEqual(resolved, str(real))
def test_revision_is_forwarded_to_the_hub_resolver(self):
"""--revision must pin the GGUF download, not be silently dropped."""
from sglang.multimodal_gen.runtime.loader import transformer_load_utils
server_args = Mock()
server_args.transformer_weights_path = "owner/repo:Q4_K_M"
server_args.revision = "abc123"
server_args.tp_size = 1
server_args.use_fsdp_inference = False
server_args.lora_path = None
server_args.minimax_h3_adaln_online = False
server_args.minimax_h3_adaln_cache_path = None
server_args.quantization = None
server_args.nunchaku_config = None
with (
patch.object(
transformer_load_utils, "resolve_hf_gguf_reference", return_value=None
) as resolve,
patch.object(transformer_load_utils, "current_platform") as platform,
):
platform.is_cuda.return_value = True
# Resolution returns None, so the override itself is checked and
# rejected as a non-file; the call arguments are what matters here.
with self.assertRaises(ValueError):
transformer_load_utils.resolve_transformer_gguf_to_load(server_args)
resolve.assert_called_once_with("owner/repo:Q4_K_M", revision="abc123")
def test_unsupported_config_rejected_before_download(self):
from sglang.multimodal_gen.runtime.loader import transformer_load_utils
server_args = Mock()
# A Hub reference: resolving it would download the whole checkpoint.
server_args.transformer_weights_path = "owner/repo:Q4_K_M"
server_args.tp_size = 1
server_args.use_fsdp_inference = True
server_args.lora_path = None
server_args.minimax_h3_adaln_online = False
server_args.minimax_h3_adaln_cache_path = None
server_args.quantization = None
server_args.nunchaku_config = None
with (
patch.object(
transformer_load_utils, "resolve_hf_gguf_reference"
) as resolve,
patch.object(transformer_load_utils, "current_platform") as platform,
):
platform.is_cuda.return_value = True
with self.assertRaisesRegex(ValueError, "FSDP"):
transformer_load_utils.resolve_transformer_gguf_to_load(server_args)
resolve.assert_not_called()
class TestGGUFRejectsLoraConversion(unittest.TestCase):
"""The dynamic set_lora path must refuse before replacing any layer."""
def _pipeline_with(self, layer):
from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import (
LoRAPipeline,
)
transformer = torch.nn.Module()
transformer.add_module("blocks", torch.nn.Module())
transformer.blocks.add_module("attn", torch.nn.Module())
transformer.blocks.attn.add_module("qkv_proj", layer)
# LoRAPipeline is abstract; the guard only needs `modules` and
# `is_target_layer`, so a minimal concrete subclass keeps the test on
# the real method rather than a copy of it.
class _Pipeline(LoRAPipeline):
def create_pipeline_stages(self, *args, **kwargs):
raise NotImplementedError
pipeline = _Pipeline.__new__(_Pipeline)
pipeline.modules = {"transformer": transformer}
pipeline.lora_initialized = False
pipeline.is_target_layer = lambda name: name.endswith("qkv_proj")
return pipeline
def _gguf_like_layer(self):
layer = torch.nn.Module()
# uint8 cannot require grad, which is why GGUFLinearMethod passes
# requires_grad=False when it registers the real parameter.
layer.register_parameter(
"qweight",
torch.nn.Parameter(
torch.zeros(4, 8, dtype=torch.uint8), requires_grad=False
),
)
return layer
def _plain_layer(self):
layer = torch.nn.Module()
layer.register_parameter("weight", torch.nn.Parameter(torch.zeros(4, 8)))
return layer
def test_packed_weights_are_rejected(self):
pipeline = self._pipeline_with(self._gguf_like_layer())
with self.assertRaisesRegex(ValueError, "LoRA is not supported"):
pipeline._reject_lora_on_packed_weights()
# Nothing was converted, so a later retry on an unquantized model works.
self.assertFalse(pipeline.lora_initialized)
def test_plain_weights_are_accepted(self):
pipeline = self._pipeline_with(self._plain_layer())
pipeline._reject_lora_on_packed_weights()
def test_rejects_while_still_offloaded(self):
"""The check must not need the weights materialized.
Layerwise offload swaps `.data` for a 1-element placeholder but keeps the
Parameter and its name, so the guard can run before offload is disabled --
which is the point: disabling it would materialize the whole DiT first.
"""
layer = self._gguf_like_layer()
layer.qweight.data = torch.empty((1,), dtype=torch.uint8)
pipeline = self._pipeline_with(layer)
with self.assertRaisesRegex(ValueError, "LoRA is not supported"):
pipeline._reject_lora_on_packed_weights()
def test_set_lora_rejects_before_disabling_offload(self):
"""Ordering matters: disabling offload first would OOM a memory-limited
deployment instead of returning the unsupported-LoRA error."""
from unittest.mock import MagicMock
pipeline = self._pipeline_with(self._gguf_like_layer())
pipeline._resolve_lora_merge_mode = lambda *a, **k: "auto"
pipeline._normalize_lora_params = lambda *a, **k: (
["n"],
["p"],
[1.0],
["all"],
[None],
)
entered = MagicMock(
side_effect=AssertionError("offload was disabled before the check")
)
pipeline._temporarily_disable_offload = entered
with self.assertRaisesRegex(ValueError, "LoRA is not supported"):
pipeline.set_lora("n", lora_path="p")
entered.assert_not_called()
if __name__ == "__main__":
unittest.main()
@@ -254,6 +254,7 @@ def _quality_server_args():
is_dit_layerwise_offload_selected=False, is_dit_layerwise_offload_selected=False,
performance_mode="speed", performance_mode="speed",
quantization=None, quantization=None,
transformer_weights_path=None,
regional_compile=False, regional_compile=False,
ring_degree=1, ring_degree=1,
sp_degree=4, sp_degree=4,
@@ -263,6 +264,24 @@ def _quality_server_args():
) )
def test_high_quality_deployment_rejects_transformer_weight_override():
config = MiniMaxH3PipelineConfig()
server_args = _quality_server_args()
server_args.transformer_weights_path = "model.gguf"
with (
patch.object(current_platform, "is_cuda", return_value=True),
patch.object(current_platform, "get_device_name", return_value="NVIDIA H200"),
patch.object(
current_platform,
"get_device_capability",
return_value=_HopperCapability(),
),
pytest.raises(ValueError, match="transformer_weights_path"),
):
config.validate_quality_deployment(server_args)
def test_high_quality_request_warns_when_bcg_suppresses_cache_dit(): def test_high_quality_request_warns_when_bcg_suppresses_cache_dit():
stage = MiniMaxH3DenoisingStage.__new__(MiniMaxH3DenoisingStage) stage = MiniMaxH3DenoisingStage.__new__(MiniMaxH3DenoisingStage)
stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=True) stage.server_args = SimpleNamespace(enable_breakable_cuda_graph=True)
@@ -129,6 +129,25 @@ def test_native_weight_names_and_grouped_qkv_reorder():
) )
def test_pruned_adaln_curve_interpolates_without_timestep_mlp():
model = MiniMaxH3DiTModel.__new__(MiniMaxH3DiTModel)
torch.nn.Module.__init__(model)
model.time_embedder = None
model.adaln_t_table = torch.nn.Parameter(
torch.tensor([[0.0, 2.0], [2.0, 4.0], [4.0, 6.0]]),
requires_grad=False,
)
result = model._time_embedding(torch.tensor([0.0, 0.25, 1.0]))
torch.testing.assert_close(
result,
torch.tensor([[0.0, 2.0], [1.0, 3.0], [4.0, 6.0]]),
rtol=0,
atol=0,
)
class _KwargIdentity(torch.nn.Module): class _KwargIdentity(torch.nn.Module):
def forward(self, x, **_kwargs): def forward(self, x, **_kwargs):
return x return x
@@ -294,6 +313,27 @@ def test_meta_model_enforces_mixed_precision_weight_contract():
assert tensor.dtype == torch.bfloat16, name assert tensor.dtype == torch.bfloat16, name
def test_pruned_meta_model_preserves_curve_adaln_fp32_island():
_ensure_single_process_parallel_runtime()
config = MiniMaxH3DiTConfig(
arch_config=MiniMaxH3DiTArchConfig(
adaln_curve_grid=1025,
time_embed_dim=8,
)
)
with torch.device("meta"):
model = MiniMaxH3DiTModel(
config=config,
hf_config={},
quant_config=None,
)
assert model.time_embedder is None
assert model.adaln_t_table.dtype == torch.float32
assert model.blocks[0].adaln_proj.linear.weight.dtype == torch.float32
assert model.final_layer.adaln_proj.linear.weight.dtype == torch.float32
def test_online_fp8_keeps_fp32_boundaries_and_ignored_layers_unquantized(): def test_online_fp8_keeps_fp32_boundaries_and_ignored_layers_unquantized():
_ensure_single_process_parallel_runtime() _ensure_single_process_parallel_runtime()
with torch.device("meta"): with torch.device("meta"):
@@ -288,6 +288,16 @@ class TestTransformerQuantHelpers(unittest.TestCase):
self.assertEqual(device, runtime_device) self.assertEqual(device, runtime_device)
def test_gguf_cpu_offload_loads_packed_checkpoint_on_cpu(self):
device = _resolve_checkpoint_load_device(
torch.device("cuda:0"),
component_starts_on_cpu=True,
runtime_quant_config=object(),
quantized_cpu_load_supported=True,
)
self.assertEqual(device, torch.device("cpu"))
def test_resident_transformer_loads_checkpoint_on_runtime_device(self): def test_resident_transformer_loads_checkpoint_on_runtime_device(self):
runtime_device = torch.device("cuda:0") runtime_device = torch.device("cuda:0")
device = _resolve_checkpoint_load_device( device = _resolve_checkpoint_load_device(
+10 -3
View File
@@ -169,6 +169,15 @@ MMVQ_QUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES
MMQ_QUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES MMQ_QUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES
def dequantize_gguf_weight(
qweight: torch.Tensor, qweight_type: int, dtype: torch.dtype
) -> torch.Tensor:
"""Dequantize a packed GGUF matrix using its inferred logical shape."""
block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type]
shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size)
return ggml_dequantize(qweight, qweight_type, *shape, dtype)
def fused_mul_mat_gguf( def fused_mul_mat_gguf(
x: torch.Tensor, qweight: torch.Tensor, qweight_type: int x: torch.Tensor, qweight: torch.Tensor, qweight_type: int
) -> torch.Tensor: ) -> torch.Tensor:
@@ -191,9 +200,7 @@ def fused_mul_mat_gguf(
y = ggml_mul_mat_a8(qweight, x, qweight_type, qweight.shape[0]) y = ggml_mul_mat_a8(qweight, x, qweight_type, qweight.shape[0])
# If there is no available MMQ kernel, fallback to dequantize # If there is no available MMQ kernel, fallback to dequantize
elif qweight_type in DEQUANT_TYPES: elif qweight_type in DEQUANT_TYPES:
block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type] weight = dequantize_gguf_weight(qweight, qweight_type, x.dtype)
shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size)
weight = ggml_dequantize(qweight, qweight_type, *shape, x.dtype)
y = x @ weight.T y = x @ weight.T
else: else:
# Raise an error if the quantization type is not supported. # Raise an error if the quantization type is not supported.
@@ -294,6 +294,7 @@ def resolve_hf_gguf_reference(
"""Download a .gguf named by Hub reference and return its local path. """Download a .gguf named by Hub reference and return its local path.
owner/repo/path/inside/repo.gguf -> exactly that file owner/repo/path/inside/repo.gguf -> exactly that file
owner/repo:QUANT_TYPE -> the only matching quantization
owner/repo -> the only .gguf in the repo owner/repo -> the only .gguf in the repo
""" """
from sglang.srt.utils import is_remote_url from sglang.srt.utils import is_remote_url
@@ -301,12 +302,41 @@ def resolve_hf_gguf_reference(
if not model or os.path.exists(model) or is_remote_url(model): if not model or os.path.exists(model) or is_remote_url(model):
return None return None
from huggingface_hub import hf_hub_download
if ":" in model:
repo_id, _, quant_type = model.rpartition(":")
if repo_id.count("/") != 1 or not quant_type:
return None
from huggingface_hub import HfApi
files = [
sibling.rfilename
for sibling in HfApi().repo_info(repo_id, revision=revision).siblings
]
suffix = f"-{quant_type}.gguf"
candidates = [filename for filename in files if filename.endswith(suffix)]
if not candidates:
available = sorted(
filename for filename in files if filename.endswith(".gguf")
)
raise ValueError(
f"No file matching quant type {quant_type!r} in {repo_id}. "
f"Available GGUF files: {available}"
)
if len(candidates) > 1:
raise ValueError(
f"Quant type {quant_type!r} is ambiguous in {repo_id}: "
f"{sorted(candidates)}. Pass the full owner/repo/path/file.gguf "
"reference instead."
)
return hf_hub_download(repo_id, candidates[0], revision=revision)
parts = model.strip("/").split("/") parts = model.strip("/").split("/")
if len(parts) < 2: if len(parts) < 2:
return None return None
from huggingface_hub import hf_hub_download
if len(parts) > 2 and model.endswith(".gguf"): if len(parts) > 2 and model.endswith(".gguf"):
repo_id = "/".join(parts[:2]) repo_id = "/".join(parts[:2])
filename = "/".join(parts[2:]) filename = "/".join(parts[2:])
@@ -25,6 +25,7 @@ from sglang.srt.utils.hf_transformers.common import (
get_context_length, get_context_length,
get_hf_text_config, get_hf_text_config,
get_rope_config, get_rope_config,
resolve_hf_gguf_reference,
) )
from sglang.srt.utils.hf_transformers.tokenizer import _fix_special_tokens_pattern from sglang.srt.utils.hf_transformers.tokenizer import _fix_special_tokens_pattern
from sglang.srt.utils.hf_transformers_patches import normalize_rope_scaling_compat from sglang.srt.utils.hf_transformers_patches import normalize_rope_scaling_compat
@@ -362,6 +363,43 @@ class TestCheckGgufFile(unittest.TestCase):
self.assertFalse(check_gguf_file(d)) self.assertFalse(check_gguf_file(d))
class TestResolveHfGgufReference(unittest.TestCase):
@patch("huggingface_hub.hf_hub_download", return_value="/cache/model-Q4_K.gguf")
@patch("huggingface_hub.HfApi")
def test_resolves_quant_type(self, api_cls, download):
api_cls.return_value.repo_info.return_value.siblings = [
SimpleNamespace(rfilename="model-Q4_K.gguf"),
SimpleNamespace(rfilename="model-Q8_0.gguf"),
]
resolved = resolve_hf_gguf_reference("owner/repo:Q4_K", revision="revision")
self.assertEqual(resolved, "/cache/model-Q4_K.gguf")
download.assert_called_once_with(
"owner/repo", "model-Q4_K.gguf", revision="revision"
)
@patch("huggingface_hub.HfApi")
def test_rejects_ambiguous_quant_type(self, api_cls):
api_cls.return_value.repo_info.return_value.siblings = [
SimpleNamespace(rfilename="fl2va-Q4_K.gguf"),
SimpleNamespace(rfilename="ref2va-Q4_K.gguf"),
]
with self.assertRaisesRegex(ValueError, "ambiguous"):
resolve_hf_gguf_reference("owner/repo:Q4_K")
@patch("huggingface_hub.HfApi")
def test_reports_available_files_when_quant_type_is_missing(self, api_cls):
api_cls.return_value.repo_info.return_value.siblings = [
SimpleNamespace(rfilename="model-Q4_K.gguf"),
SimpleNamespace(rfilename="README.md"),
]
with self.assertRaisesRegex(ValueError, "model-Q4_K.gguf"):
resolve_hf_gguf_reference("owner/repo:Q8_0")
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# _is_deepseek_ocr_model / _is_deepseek_ocr2_model # _is_deepseek_ocr_model / _is_deepseek_ocr2_model
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------