anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: tsinghua-code <zr-liu25@mails.tsinghua.edu.cn>
This commit is contained in:
tsinghua-code authored and GitHub committed 2026-09-21 12:21:46 +00:00
1 parent 5e28163ba5
commit e5c0afc994
7 files changed
+129 -24
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/models/lingbot_world/pipeline.py
-2
View File
@@ -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
+24 -8
View File
@@ -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 = {