Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2a47d901a0 |
No files matched your search
@@ -12,20 +12,20 @@ For the internal architecture and backend extension points, see the
|
||||
|
||||
| Mode | Guide | Description | Methods |
|
||||
|------|-------|-------------|---------|
|
||||
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, MXFP8 W8A8, MXFP4 W4A4 |
|
||||
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8, BitsAndBytes W4, MXFP8 W8A8, MXFP4 W4A4 |
|
||||
| Runtime attention quantization | [Quantized KV Cache](quantized_kvcache.md) | vLLM-Omni dynamically quantizes eligible diffusion Flash Attention tensors during inference. | FP8 FA |
|
||||
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, GGUF, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
|
||||
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | ModelOpt, AutoRound, msModelSlim, serialized Int8, offline MXFP8, offline MXFP4 DualScale |
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | GGUF | AutoRound | msModelSlim |
|
||||
|--------|----------|-----------|----------|------------|------------|------|-----------|-------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ✅ | ❌ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ⭕ | ⭕ | ❌ |
|
||||
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| Ascend NPU | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ |
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | BitsAndBytes W4 | ModelOpt | MXFP8 W8A8 | MXFP4 W4A4 | AutoRound | msModelSlim |
|
||||
|--------|----------|-----------|------------------|----------|------------|------------|-----------|-------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ✅ | ⭕ | ❌ |
|
||||
| Intel XPU | ⭕ | ⭕ | ❌ | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
@@ -43,10 +43,10 @@ otherwise.
|
||||
|--------|-------|------|----------------|--------|
|
||||
| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
|
||||
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
|
||||
| BitsAndBytes W4 | [BitsAndBytes](bitsandbytes.md) | Online W4 weight-only | Z-Image-Turbo; Qwen-Image/Wan2.2 not validated | Validated for Z-Image-Turbo |
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Pre-quantized FP8 checkpoints | Qwen-Image, Z-Image, FLUX.2, HunyuanImage-3.0 | Validated for ModelOpt FP8 diffusion checkpoints |
|
||||
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Online W8A8 or offline pre-quantized | Wan2.2-T2V-A14B, I2V-A14B, TI2V-5B | Ascend NPU only; validated for Wan2.2 |
|
||||
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | `mxfp4`: online single-scale only; `mxfp4_dualscale`: online or offline dual-scale (offline recommended) | Wan2.2-T2V-A14B, I2V-A14B | Ascend NPU only; validated for Wan2.2 A14B cascade models; TI2V-5B not supported; offline `mxfp4_dualscale` uses calibrated `mul_scale` for best accuracy |
|
||||
| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
|
||||
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
|
||||
|
||||
@@ -61,9 +61,9 @@ in BF16 unless the model guide explicitly adds support.
|
||||
|--------|-------|-------|----------------|--------|
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
|
||||
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| MXFP8 | [MXFP8](mxfp8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| MXFP4 | [MXFP4](mxfp4.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
|
||||
@@ -76,10 +76,10 @@ attached to the intended stage rather than applied globally.
|
||||
|--------|-------|-------|----------------|--------|
|
||||
| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| BitsAndBytes | [BitsAndBytes](bitsandbytes.md) | Stage-specific transformer or DiT module | BAGEL, GLM-Image | Not validated |
|
||||
| ModelOpt | [ModelOpt](modelopt.md) | Checkpoint-defined diffusion stage | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| MXFP8 | [MXFP8](mxfp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
|
||||
| MXFP4 | [MXFP4](mxfp4.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Not validated |
|
||||
| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
|
||||
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
|
||||
|
||||
@@ -106,7 +106,7 @@ config = build_quant_config({
|
||||
|
||||
| Component | Default quantized? | Notes |
|
||||
|-----------|--------------------|-------|
|
||||
| Diffusion transformer | Yes | Primary target for FP8, Int8, ModelOpt, MXFP8, MXFP4, GGUF, AutoRound, and msModelSlim |
|
||||
| Diffusion transformer | Yes | Primary target for FP8, Int8, BitsAndBytes, ModelOpt, MXFP8, MXFP4, AutoRound, and msModelSlim |
|
||||
| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
|
||||
| VAE | No | Keep BF16; storage-only paths are method-specific |
|
||||
| Scheduler/tokenizer | No | Loaded from the base model repository |
|
||||
@@ -139,8 +139,8 @@ from vllm_omni.quantization import build_quant_config
|
||||
|
||||
build_quant_config("fp8")
|
||||
build_quant_config({"method": "fp8", "activation_scheme": "static"})
|
||||
build_quant_config("bitsandbytes")
|
||||
build_quant_config("auto-round", bits=4, group_size=128)
|
||||
build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
|
||||
build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
|
||||
build_quant_config(None)
|
||||
```
|
||||
|
||||
@@ -6,7 +6,6 @@ import torch.nn as nn
|
||||
|
||||
import vllm_omni.diffusion.compile as compile_module
|
||||
from vllm_omni.diffusion.compile import regionally_compile
|
||||
from vllm_omni.diffusion.hooks import HookRegistry, ModelHook
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
@@ -77,37 +76,6 @@ def test_regionally_compile_does_not_partially_mutate_on_setup_failure(monkeypat
|
||||
assert [block.forward.__func__ for block in model.transformer_blocks] == original_forwards
|
||||
|
||||
|
||||
def test_regionally_compile_keeps_hook_dispatch_outside_compiled_graph(monkeypatch):
|
||||
model = _ModelWithWrappedRepeatedBlocks()
|
||||
block = model.transformer_blocks[0]
|
||||
registry = HookRegistry.get_or_create(block)
|
||||
hook = ModelHook()
|
||||
registry.register_hook("test", hook)
|
||||
|
||||
wrapped_forward = block.forward
|
||||
original_forward = block._omni_original_forward
|
||||
compile_calls = []
|
||||
|
||||
def _compile(fn, *args, **kwargs):
|
||||
compile_calls.append(fn)
|
||||
|
||||
def _compiled(*fn_args, **fn_kwargs):
|
||||
return f"compiled:{fn(*fn_args, **fn_kwargs)}"
|
||||
|
||||
return _compiled
|
||||
|
||||
monkeypatch.setattr(compile_module.torch, "compile", _compile)
|
||||
|
||||
regionally_compile(model)
|
||||
|
||||
assert compile_calls[0] is original_forward
|
||||
assert compile_calls[0] is not wrapped_forward
|
||||
assert block.forward is wrapped_forward
|
||||
assert block._omni_original_forward is not original_forward
|
||||
assert hook.fn_ref.original_forward is block._omni_original_forward
|
||||
assert block("ok") == "compiled:ok"
|
||||
|
||||
|
||||
def test_compiled_block_preserves_forward_signature_for_inspection(monkeypatch):
|
||||
"""cache-dit matches blocks via inspect.signature(block.forward).
|
||||
|
||||
|
||||
@@ -56,50 +56,26 @@ def regionally_compile(
|
||||
# Build all compiled callables before mutating the model. This keeps setup
|
||||
# failures atomic: callers can safely continue with the uncompiled model if
|
||||
# torch.compile raises synchronously for any repeated block.
|
||||
# Keep the pre-hook callable separately. DLO and the ordinary layerwise
|
||||
# offloader install a HookRegistry wrapper in ``module.forward``. That
|
||||
# wrapper performs stream/event synchronization and storage rebinding, so
|
||||
# compiling it pulls offload control flow into the graph. Compile the
|
||||
# original block compute instead and leave the wrapper outside the graph.
|
||||
compiled_forwards: list[tuple[nn.Module, Any | None, Any]] = []
|
||||
compiled_forwards: list[tuple[nn.Module, Any]] = []
|
||||
for name, submod in model.named_modules():
|
||||
if _matches_repeated_block(name, submod, repeated_blocks, repeated_block_attrs):
|
||||
# Compile the block compute while keeping nn.Module.__call__ hooks
|
||||
# outside the compiled graph. If a HookRegistry is already
|
||||
# installed, ``submod.forward`` is the hook dispatcher and the
|
||||
# original callable is stored in ``_omni_original_forward``.
|
||||
# NOTE: anything that wraps this callable must stay
|
||||
# signature-transparent — cache-dit's BlockAdapter matches blocks
|
||||
# by inspecting forward's parameter names and return annotation,
|
||||
# which torch.compile preserves.
|
||||
original_forward = getattr(submod, "_omni_original_forward", None)
|
||||
forward_to_compile = original_forward if original_forward is not None else submod.forward
|
||||
# outside the compiled graph. NOTE: anything that wraps this
|
||||
# callable must stay signature-transparent — cache-dit's
|
||||
# BlockAdapter matches blocks by inspecting forward's parameter
|
||||
# names and return annotation, which torch.compile preserves.
|
||||
compiled_forwards.append(
|
||||
(
|
||||
submod,
|
||||
original_forward,
|
||||
torch.compile(forward_to_compile, *compile_args, **compile_kwargs),
|
||||
torch.compile(submod.forward, *compile_args, **compile_kwargs),
|
||||
)
|
||||
)
|
||||
|
||||
if not compiled_forwards:
|
||||
logger.warning(f"Regional compilation skipped because {repeated_blocks} classes are not found in the model.")
|
||||
else:
|
||||
for submod, original_forward, compiled_forward in compiled_forwards:
|
||||
if original_forward is None:
|
||||
submod.forward = compiled_forward
|
||||
continue
|
||||
|
||||
# Keep HookRegistry's dispatcher as ``forward`` and replace only
|
||||
# the callable it dispatches to. Hooks such as MagCache retain a
|
||||
# reference to the original callable as well, so update those
|
||||
# references when they still point at the pre-compile function.
|
||||
submod._omni_original_forward = compiled_forward
|
||||
registry = getattr(submod, "_hook_registry", None)
|
||||
for hook in getattr(registry, "_hooks", {}).values():
|
||||
fn_ref = getattr(hook, "fn_ref", None)
|
||||
if fn_ref is not None and fn_ref.original_forward is original_forward:
|
||||
fn_ref.original_forward = compiled_forward
|
||||
for submod, compiled_forward in compiled_forwards:
|
||||
submod.forward = compiled_forward
|
||||
logger.info(
|
||||
"Regional compilation applied to %d module(s) for repeated blocks %s.",
|
||||
len(compiled_forwards),
|
||||
|
||||
Reference in new issue
Block a user