[Model] Enable online FP8 linears for LingBot World (#7549)

Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>

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.<i>.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.
 

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

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

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:

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"]

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

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 = {