diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md index 7add2f7d8..b20796726 100644 --- a/recipes/Robbyant/LingBot-World-2.0.md +++ b/recipes/Robbyant/LingBot-World-2.0.md @@ -213,6 +213,47 @@ accepts: `s`, `d`, `i`, `j`, `k`, and `l`; - `lingbot.camera_trajectory.v1` for explicit pose/intrinsics trajectories. +## Online FP8 linear layers (experimental) + +Pass `quantization_config={"method": "fp8"}` to the existing `AsyncOmni` +construction to use vLLM's online FP8 linear implementation. On the tested +Hopper-class GPU this dispatches to CUTLASS FP8 GEMM with online per-tensor +activation scaling; no LingBot-specific quantization kernel is introduced. + +Only the transformer's existing vLLM parallel linear layers are eligible: +self-attention, cross-attention, FFN, camera injectors, and C2WS projections. +The ordinary PyTorch linear layers (including the output head and +time/text embeddings), normalization, convolutions, VAE, text encoder, and +AR KV cache are not converted to FP8 by this option. + +Quality-sensitive projections can remain BF16 through fully qualified +`ignored_layers` names. For example, the following configuration tests +retaining the camera/C2WS path in BF16: + +```python +quantization_config = { + "method": "fp8", + "ignored_layers": [ + f"transformer.blocks.{i}.cam_injector_layer{j}" + for i in range(40) + for j in (1, 2) + ] + + [ + "transformer.c2ws_hidden_states_layer1", + "transformer.c2ws_hidden_states_layer2", + ], +} +``` + +Use `transformer.blocks..self_attn.qkv` to exclude the fused Q/K/V +projection as a unit. The example exclusion list is an ablation, not a +universal quality guarantee or a default. Validate the intended scene, seed, +camera trajectory, and session length against BF16 before selecting a policy. +An FP8 kernel speedup alone is not an end-to-end performance claim. +The long-session E8 quality gate in +[#7074](https://github.com/vllm-project/vllm-omni/issues/7074) remains required +for Tier-2 acceptance. + ## Validation Real-checkpoint validation uses 480x832 output, four DMD steps, and seed 42. @@ -246,7 +287,7 @@ tested commit. AR blocks in one request. `max_num_seqs` must be one in both cases. - Stateful streaming VAE decode is not implemented; the realtime example emits latent chunks. -- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization, +- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization methods other than online FP8, Cache-DiT, TeaCache, causal-pretrain, and the 1.3B checkpoint are not claimed. - No AMD GPU, Ascend NPU, or Intel GPU support is claimed. diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py index 7e12b3db3..d2d11538c 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py @@ -134,6 +134,8 @@ def _install_vllm_stubs() -> None: **kwargs, ) -> None: super().__init__() + self.quant_config = kwargs.pop("quant_config", None) + self.prefix = kwargs.pop("prefix", "") del kwargs self.return_bias = return_bias self.weight = nn.Parameter(torch.empty(output_size, input_size)) @@ -164,6 +166,8 @@ def _install_vllm_stubs() -> None: **kwargs, ) -> None: super().__init__() + self.quant_config = kwargs.pop("quant_config", None) + self.prefix = kwargs.pop("prefix", "") del kwargs self.num_heads = total_num_heads self.num_kv_heads = total_num_kv_heads or total_num_heads diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py index e11c34e35..7087e6171 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_transformer.py @@ -35,6 +35,8 @@ def _tiny_model( num_layers: int = 2, num_frames_per_block: int = 1, sliding_window_num_frames: int = 3, + quant_config=None, + prefix: str = "", attention_head_dim: int = 2, rope_max_seq_len: int = 16, ): @@ -55,6 +57,8 @@ def _tiny_model( num_frames_per_block=num_frames_per_block, sliding_window_num_frames=sliding_window_num_frames, local_attn_size=-1, + quant_config=quant_config, + prefix=prefix, ) @@ -485,11 +489,52 @@ def test_constructor_rejects_non_null_image_embedding_fields(field: str) -> None @pytest.mark.cpu -def test_constructor_rejects_unsupported_quantization_with_runtime_error() -> None: +@pytest.mark.parametrize("prefix", ["", "transformer"]) +@pytest.mark.parametrize("quantized", [False, True]) +def test_constructor_propagates_quantization_to_parallel_linears(prefix: str, quantized: bool) -> None: module = attention_tests._load_module() + quant_config = object() if quantized else None + model = _tiny_model(module, num_layers=1, quant_config=quant_config, prefix=prefix) + block = model.blocks[0] + quantized_linears = [ + model.c2ws_hidden_states_layer1, + model.c2ws_hidden_states_layer2, + block.self_attn.qkv, + block.self_attn.o, + block.cross_attn.q, + block.cross_attn.k, + block.cross_attn.v, + block.cross_attn.o, + block.ffn[0], + block.ffn[2], + block.cam_injector_layer1, + block.cam_injector_layer2, + ] - with pytest.raises(RuntimeError, match="quant_config.*not supported"): - module.CausalLingBotWorldTransformer3DModel(quant_config=object()) + assert len(quantized_linears) == 12 + assert all(layer.quant_config is quant_config for layer in quantized_linears) + expected_names = { + "c2ws_hidden_states_layer1", + "c2ws_hidden_states_layer2", + *( + f"blocks.0.{name}" + for name in ( + "self_attn.qkv", + "self_attn.o", + "cross_attn.q", + "cross_attn.k", + "cross_attn.v", + "cross_attn.o", + "ffn.0", + "ffn.2", + "cam_injector_layer1", + "cam_injector_layer2", + ) + ), + } + assert {layer.prefix for layer in quantized_linears} == { + f"{prefix}.{name}" if prefix else name for name in expected_names + } @pytest.mark.cpu diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index b5a57ce75..8c24b8562 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -889,13 +889,14 @@ def test_unsupported_sp_config_fails_before_component_loading(overrides): assert module._loader_state.prefetch_calls == [] -def test_unsupported_quantization_fails_before_component_loading() -> None: +def test_quantization_reaches_transformer_factory() -> None: module = _load_pipeline_module() + quant_config = object() - with pytest.raises(NotImplementedError, match="quantization"): - module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=object())) + pipeline = module.LingBotWorldCausalDMDPipeline(od_config=_od_config(quantization_config=quant_config)) - assert module._loader_state.prefetch_calls == [] + assert pipeline.transformer is not None + assert _FakeTransformerFactory.last_call[1:] == (quant_config, "transformer") def test_official_scheduler_config_matches_fixed_dmd_contract() -> None: diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index d40a9c226..3b0192a79 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -563,19 +563,19 @@ def test_from_cli_args_preserves_text_encoder_tp_size(): def test_text_encoder_tp_size_reaches_default_diffusion_parallel_config(): """The preserved CLI value must land in DiffusionParallelConfig. - Mirrors the library flow: from_cli_args -> kwargs -> the generic - diffusion fallback, which resolves the field through - ``DiffusionParallelConfig.from_stage_overrides``. + Forward the preserved explicit override to the generic diffusion + fallback, which resolves it through + ``DiffusionParallelConfig.from_stage_overrides``. Serializing all engine + defaults would also forward unrelated LLM-only fields to strict diffusion + ingress, unlike the explicit-kwargs library entrypoint. """ - from dataclasses import asdict - from vllm_omni.config.config_factory import StageConfigFactory engine_args = OmniEngineArgs.from_cli_args( SimpleNamespace(text_encoder_tp_size=2), ) stage_cfg = StageConfigFactory.create_default_diffusion( - {k: v for k, v in asdict(engine_args).items() if v is not None}, + {"text_encoder_tp_size": engine_args.text_encoder_tp_size}, )[0] parallel_config = stage_cfg["engine_args"]["parallel_config"] diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index 8bb83ef80..fec175e69 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -241,8 +241,6 @@ def _validate_scheduler_config(config: dict[str, Any]) -> None: def _validate_parallel_config(od_config: OmniDiffusionConfig) -> None: - if getattr(od_config, "quantization_config", None) is not None: - raise NotImplementedError("LingBot World v1 does not support quantization.") parallel_config = getattr(od_config, "parallel_config", None) if parallel_config is None: return diff --git a/vllm_omni/diffusion/models/lingbot_world/transformer.py b/vllm_omni/diffusion/models/lingbot_world/transformer.py index 9a270e611..5b36299d8 100644 --- a/vllm_omni/diffusion/models/lingbot_world/transformer.py +++ b/vllm_omni/diffusion/models/lingbot_world/transformer.py @@ -8,7 +8,7 @@ import math from collections.abc import Iterable from dataclasses import dataclass from types import SimpleNamespace -from typing import Any, Self, cast +from typing import TYPE_CHECKING, Any, Self, cast import torch import torch.nn as nn @@ -36,6 +36,9 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.paged_attention import ( paged_write_attn, ) +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + @dataclass class LingBotAttentionCache: @@ -170,6 +173,7 @@ class LingBotSelfAttention(nn.Module): num_heads: int, *, eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() @@ -188,6 +192,7 @@ class LingBotSelfAttention(nn.Module): head_size=self.head_dim, total_num_heads=num_heads, bias=True, + quant_config=quant_config, prefix=_projection_prefix(prefix, "qkv"), ) self.num_local_heads = self.qkv.num_heads @@ -205,6 +210,7 @@ class LingBotSelfAttention(nn.Module): bias=True, input_is_parallel=True, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "o"), ) self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps) @@ -416,6 +422,7 @@ class LingBotCrossAttention(nn.Module): num_heads: int, *, eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() @@ -443,6 +450,7 @@ class LingBotCrossAttention(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "q"), ) self.k = ColumnParallelLinear( @@ -451,6 +459,7 @@ class LingBotCrossAttention(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "k"), ) self.v = ColumnParallelLinear( @@ -459,6 +468,7 @@ class LingBotCrossAttention(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "v"), ) self.o = RowParallelLinear( @@ -467,6 +477,7 @@ class LingBotCrossAttention(nn.Module): bias=True, input_is_parallel=True, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "o"), ) self.norm_q = _LingBotRMSNorm(self.tp_inner_dim, eps) @@ -528,6 +539,7 @@ class LingBotAttentionBlock(nn.Module): ffn_dim: int | None = None, cross_attn_norm: bool = True, eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() @@ -538,12 +550,14 @@ class LingBotAttentionBlock(nn.Module): dim, num_heads, eps=eps, + quant_config=quant_config, prefix=_projection_prefix(prefix, "self_attn"), ) self.cross_attn = LingBotCrossAttention( dim, num_heads, eps=eps, + quant_config=quant_config, prefix=_projection_prefix(prefix, "cross_attn"), ) self.norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False) @@ -555,6 +569,7 @@ class LingBotAttentionBlock(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "ffn.0"), ), nn.GELU(approximate="tanh"), @@ -564,6 +579,7 @@ class LingBotAttentionBlock(nn.Module): bias=True, input_is_parallel=True, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "ffn.2"), ), ) @@ -574,6 +590,7 @@ class LingBotAttentionBlock(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "cam_injector_layer1"), ) self.cam_injector_layer2 = RowParallelLinear( @@ -582,6 +599,7 @@ class LingBotAttentionBlock(nn.Module): bias=True, input_is_parallel=True, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "cam_injector_layer2"), ) self.cam_scale_layer = nn.Linear(dim, dim) @@ -823,7 +841,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): num_frames_per_block: int = 3, sliding_window_num_frames: int = 18, local_attn_size: int = -1, - quant_config: object | None = None, + quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: super().__init__() @@ -853,11 +871,6 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): ): if field_value is not None: raise ValueError(f"{field_name} must be None because LingBot World v2 has no image embedding path.") - if quant_config is not None: - raise RuntimeError( - "quant_config is not supported by the LingBot World transformer; construct the unquantized model." - ) - dim = num_attention_heads * attention_head_dim self.dim = dim self.config = SimpleNamespace( @@ -898,6 +911,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): bias=True, gather_output=False, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "c2ws_hidden_states_layer1"), ) self.c2ws_hidden_states_layer2 = RowParallelLinear( @@ -906,6 +920,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): bias=True, input_is_parallel=True, return_bias=False, + quant_config=quant_config, prefix=_projection_prefix(prefix, "c2ws_hidden_states_layer2"), ) self.text_embedding = nn.Sequential( @@ -930,6 +945,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): ffn_dim=ffn_dim, cross_attn_norm=cross_attn_norm, eps=eps, + quant_config=quant_config, prefix=_projection_prefix(prefix, f"blocks.{index}"), ) for index in range(num_layers) @@ -976,7 +992,7 @@ class CausalLingBotWorldTransformer3DModel(nn.Module): cls, config: dict[str, Any], *, - quant_config: Any | None = None, + quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> Self: checkpoint_contract = {