Merge main into dev/vllm-align for vLLM 0.31 release

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

vllm_omni/diffusion/attention/backends/abstract.py

@@ -7,6 +7,10 @@ from typing import Any, Generic, Literal, TypeVar
 
 import torch
 
+from vllm_omni.diffusion.attention.capabilities import (
+    ExecutionContext,
+    ExecutionPathResult,
+)
 from vllm_omni.platforms import current_omni_platform
 
 
@@ -132,6 +136,15 @@ class AttentionBackend(ABC):
 
         return False
 
+    @classmethod
+    def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult:
+        """Return conservative pre-construction capabilities.
+
+        Unmigrated backends remain advisory so introducing the contract cannot
+        reject an execution path that works today.
+        """
+        return ExecutionPathResult.unmigrated(cls.get_name(), context)
+
 
 @dataclass(frozen=True, slots=True)
 class QueryRange:
@@ -218,6 +231,9 @@ class AttentionMetadata:
     #   "valid_kv_length": int — contiguous valid K/V prefix length. Local
     #     backends may slice it directly; ring backends trim each circulated
     #     K/V block from the same global prefix.
+    #   "attention_mask_mode": "none" | "padding" | "arbitrary" — producer
+    #     assertion used by capability planning. "none" means a present mask
+    #     is semantically a no-op; absent means its contents are runtime-dependent.
     #   "npu_attn_varlen": bool — model opt-in for the NPU packed varlen path
     #     (TND npu_fusion_attention driven by cu_seqlens, mask never read).
     #     Requires the [real, pad] two-document packing contract; see
@@ -277,6 +293,21 @@ class AttentionImpl(ABC, Generic[T]):
             return True
         return kv_cache_dtype in cls._supported_kv_cache_dtypes.get(platform_key, set())
 
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: T | None,
+    ) -> ExecutionPathResult:
+        """Resolve the complete path after backend initialization.
+
+        Subclasses report verified results; this inspection API does not enforce them.
+        """
+        del query, key, value, attn_metadata
+        return ExecutionPathResult.unmigrated(type(self).__name__, context)
+
     def forward(
         self,
         query: torch.Tensor,

vllm_omni/diffusion/attention/backends/flash_attn.py

@@ -3,7 +3,9 @@
 
 import os
 from collections.abc import Callable
+from dataclasses import replace
 from functools import cache, partial
+from typing import NamedTuple
 
 import torch
 from vllm.logger import init_logger
@@ -14,6 +16,16 @@ from vllm_omni.diffusion.attention.backends.utils.piecewise_attn import (
     piecewise_attn,
     run_paged_piecewise_plan,
 )
+from vllm_omni.diffusion.attention.capabilities import (
+    CapabilityResult,
+    CompilationMode,
+    ExecutionContext,
+    ExecutionPathResult,
+    MaskMode,
+    PackingMode,
+    ParallelStrategy,
+    SupportStatus,
+)
 from vllm_omni.diffusion.config import get_current_diffusion_config_or_none
 from vllm_omni.platforms import current_omni_platform
 
@@ -29,6 +41,133 @@ def _get_npu_compressed_causal_mask(device: torch.device) -> torch.Tensor:
     ).contiguous()
 
 
+if not hasattr(torch.ops.vllm_omni, "fa4_dense_attention"):
+
+    @torch.library.custom_op("vllm_omni::fa4_dense_attention", mutates_args=())
+    def _fa4_dense_attention_op(
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        softmax_scale: float,
+        causal: bool,
+        deterministic: bool,
+    ) -> torch.Tensor:
+        from vllm_omni.diffusion.attention.backends.utils.fa import flash_attn_func
+
+        if flash_attn_func is None:
+            raise RuntimeError("CuTe FlashAttention-4 is unavailable")
+        kwargs = {
+            "causal": causal,
+            "softmax_scale": softmax_scale,
+        }
+        if deterministic:
+            kwargs["deterministic"] = True
+        out = flash_attn_func(query, key, value, **kwargs)
+        return out[0] if isinstance(out, tuple) else out
+
+    @_fa4_dense_attention_op.register_fake
+    def _fa4_dense_attention_fake(
+        query,
+        key,
+        value,
+        softmax_scale,
+        causal,
+        deterministic,
+    ):
+        # FA4 returns contiguous output even when Q is noncontiguous; V sets its head dimension.
+        return query.new_empty((*query.shape[:-1], value.shape[-1]))
+
+
+_fa4_dense_attention_op = torch.ops.vllm_omni.fa4_dense_attention
+
+
+_PACKED_KEYS = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
+
+
+class _FlashAttentionMetadataPlan(NamedTuple):
+    attention_mask: torch.Tensor | None
+    mask_mode: MaskMode
+    full_attn_spans: list[list[tuple[int, int]]] | None
+    extra: dict
+    packing_mode: PackingMode
+
+
+def _normalize_flash_attention_metadata(
+    attn_metadata: AttentionMetadata | None,
+) -> _FlashAttentionMetadataPlan:
+    attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
+    full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
+    extra = attn_metadata.extra if attn_metadata is not None else {}
+    published_mask_mode = extra.get("attention_mask_mode")
+    if attention_mask is None:
+        mask_mode = MaskMode.NONE
+    elif published_mask_mode is None:
+        mask_mode = MaskMode.UNKNOWN
+    else:
+        try:
+            mask_mode = MaskMode(published_mask_mode)
+        except ValueError as error:
+            raise ValueError(f"Unknown attention_mask_mode {published_mask_mode!r}") from error
+        if mask_mode is MaskMode.UNKNOWN:
+            raise ValueError("attention_mask_mode='unknown' is reserved for unpublished semantics")
+    # Piecewise dispatch takes precedence and does not consume packed metadata.
+    present_packed_keys = [key for key in _PACKED_KEYS if key in extra] if full_attn_spans is None else []
+    if present_packed_keys and len(present_packed_keys) != len(_PACKED_KEYS):
+        missing = sorted(set(_PACKED_KEYS) - set(present_packed_keys))
+        raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
+
+    packing_mode = PackingMode.NONE
+    if present_packed_keys:
+        cu_seqlens_q = extra["cu_seqlens_q"]
+        packing_mode = PackingMode.MULTI_DOCUMENT if cu_seqlens_q.shape[0] > 3 else PackingMode.PACKED_PADDING
+    return _FlashAttentionMetadataPlan(
+        attention_mask=attention_mask,
+        mask_mode=mask_mode,
+        full_attn_spans=full_attn_spans,
+        extra=extra,
+        packing_mode=packing_mode,
+    )
+
+
+def _flash_attention_execution_path(
+    context: ExecutionContext,
+) -> ExecutionPathResult:
+    if context.mask_mode is MaskMode.UNKNOWN:
+        return ExecutionPathResult.unmigrated(
+            "FLASH_ATTN",
+            context,
+            path="runtime_mask_dependent",
+        )
+    if not (
+        context.platform == "cuda"
+        and context.kernel_variant == "fa4"
+        and context.dtype in {"bfloat16", "torch.bfloat16"}
+        and context.causal is False
+        and context.mask_mode is MaskMode.NONE
+        and context.packing_mode is PackingMode.NONE
+        and not context.piecewise
+        and not context.paged_kv
+        and context.kv_cache_dtype is None
+        and context.parallel_strategy is ParallelStrategy.NONE
+        and not context.outer_boundaries
+    ):
+        return ExecutionPathResult.unmigrated(
+            "FLASH_ATTN",
+            context,
+            path="unverified",
+        )
+
+    return ExecutionPathResult(
+        backend="FLASH_ATTN",
+        path="fa4_dense",
+        support=CapabilityResult.supported(),
+        compilation_mode=CompilationMode.CUSTOM_OP,
+        platform=context.platform,
+        kernel_variant=context.kernel_variant,
+        parallel_strategy=context.parallel_strategy,
+    )
+
+
 class FlashAttentionBackend(AttentionBackend):
     accept_output_buffer: bool = True
     supports_piecewise_spans: bool = True
@@ -74,6 +213,10 @@ class FlashAttentionBackend(AttentionBackend):
     def get_impl_cls() -> type["FlashAttentionImpl"]:
         return FlashAttentionImpl
 
+    @classmethod
+    def resolve_capabilities(cls, context: ExecutionContext) -> ExecutionPathResult:
+        return _flash_attention_execution_path(replace(context, kernel_variant=None))
+
 
 class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
     # Per-platform attention quantization support.
@@ -108,11 +251,135 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         self.softmax_scale = softmax_scale
         self.qkv_layout = qkv_layout
         self.is_cross_attn = role == "cross"
+        from vllm_omni.diffusion.attention.backends.utils.fa import IS_AITER, IS_FLASH_ATTN_4
+
+        self._kernel_variant = "fa4" if IS_FLASH_ATTN_4 else "aiter" if IS_AITER else None
         cfg = get_current_diffusion_config_or_none()
         self.fa_deterministic = bool(getattr(cfg, "fa_deterministic", False)) if cfg is not None else False
         if backend_kwargs:
             logger.warning("FlashAttentionImpl ignoring backend_kwargs: %s", list(backend_kwargs.keys()))
 
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        if context.platform == "npu":
+            return self._resolve_npu_execution_path(context, query, key, attn_metadata)
+        metadata = _normalize_flash_attention_metadata(attn_metadata)
+        context = replace(
+            context,
+            kernel_variant=self._kernel_variant,
+            dtype=str(query.dtype).removeprefix("torch."),
+            causal=self.causal,
+            mask_mode=metadata.mask_mode,
+            packing_mode=metadata.packing_mode,
+            piecewise=metadata.full_attn_spans is not None,
+            kv_cache_dtype=metadata.extra.get("kv_cache_dtype"),
+        )
+        if context.platform == "rocm":
+            return self._resolve_rocm_execution_path(context)
+        result = _flash_attention_execution_path(context)
+        if result.support.status is not SupportStatus.SUPPORTED:
+            return result
+
+        from vllm_omni.diffusion.attention.backends.utils.fa import validate_fa4_head_dims
+
+        try:
+            if query.dtype != key.dtype or query.dtype != value.dtype:
+                raise ValueError("Q, K, and V dtypes must match")
+            if query.device != key.device or query.device != value.device:
+                raise ValueError("Q, K, and V devices must match")
+            if query.shape[-1] != key.shape[-1]:
+                raise ValueError("Q and K head dimensions must match")
+            verified = validate_fa4_head_dims(query.shape[-1], value.shape[-1], 16 // value.element_size())
+        except (AssertionError, ValueError) as error:
+            return replace(
+                result,
+                support=CapabilityResult.unsupported(f"FA4: {error} Select compatible inputs or another backend."),
+            )
+        if not verified:
+            return replace(
+                result,
+                support=CapabilityResult.unmigrated("Selected FA4 kernel has no available head-dimension validator"),
+                compilation_mode=CompilationMode.EAGER_ONLY,
+            )
+        return result
+
+    @staticmethod
+    def _resolve_rocm_execution_path(context: ExecutionContext) -> ExecutionPathResult:
+        """Describe AITER dispatch; device and compiler validation is pending."""
+        result = ExecutionPathResult.unmigrated("FLASH_ATTN", context, path="rocm_unverified")
+        if (
+            context.kernel_variant != "aiter"
+            or context.parallel_strategy is not ParallelStrategy.NONE
+            or context.outer_boundaries
+            or context.paged_kv
+            or context.piecewise
+            or context.kv_cache_dtype is not None
+        ):
+            return result
+        if context.packing_mode is not PackingMode.NONE:
+            path = "rocm_packed_varlen"
+        elif context.mask_mode is MaskMode.UNKNOWN:
+            path = "rocm_runtime_mask_dependent"
+        elif context.mask_mode is not MaskMode.NONE:
+            path = "rocm_masked_varlen"
+        else:
+            path = "rocm_dense"
+        return replace(
+            result,
+            path=path,
+            support=CapabilityResult.unmigrated("AITER routing only; kernel and compilation validation require ROCm"),
+        )
+
+    def _resolve_npu_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        """Describe existing MindIE routing without claiming hardware validation.
+
+        NPU packing has its own opt-in and [real, pad] contract. In particular,
+        incomplete packed metadata can fall back to a mask, so CUDA metadata
+        normalization must not run here. No tensor values are inspected.
+        """
+        context = replace(context, kernel_variant="mindiesd")
+        result = ExecutionPathResult.unmigrated("FLASH_ATTN", context, path="npu_unverified")
+        extra = attn_metadata.extra if attn_metadata else {}
+        if (
+            context.parallel_strategy is not ParallelStrategy.NONE
+            or context.outer_boundaries
+            or context.paged_kv
+            or self.causal
+            or extra.get("kv_cache_dtype") not in (None, "float", "auto")
+            or (attn_metadata is not None and attn_metadata.full_attn_spans is not None)
+        ):
+            return result
+
+        path = "npu_masked" if attn_metadata is not None and attn_metadata.attn_mask is not None else "npu_dense"
+        if extra.get("npu_attn_varlen", False):
+            if self._resolve_packed_seq_npu(query, key, extra) is not None:
+                path = (
+                    "npu_prefix_kv_slice"
+                    if os.environ.get("MINDIE_SD_FA_TYPE") == "ascend_laser_attention"
+                    else "npu_packed_varlen"
+                )
+            else:
+                # This route may reject the request if neither an explicit mask
+                # nor usable valid_kv_length is supplied; it is not support.
+                path = "npu_masked_fallback"
+        return replace(
+            result,
+            path=path,
+            support=CapabilityResult.unmigrated("NPU routing only; kernel and compilation validation require Ascend"),
+        )
+
     def _warn_fa_deterministic_non_dense(self, path: str) -> None:
         if not self.fa_deterministic:
             return
@@ -401,6 +668,7 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
         """CUDA/ROCm/MUSA flash attention implementation."""
         from vllm_omni.diffusion.attention.backends.utils.fa import (
             HAS_FLASH_ATTN,
+            IS_FLASH_ATTN_4,
             flash_attn_func,
             flash_attn_varlen_func,
         )
@@ -412,9 +680,10 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 "Otherwise, use SDPA backend by setting DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA"
             )
 
-        attention_mask = attn_metadata.attn_mask if attn_metadata is not None else None
-        full_attn_spans = attn_metadata.full_attn_spans if attn_metadata is not None else None
-        extra = attn_metadata.extra if attn_metadata is not None else {}
+        metadata_plan = _normalize_flash_attention_metadata(attn_metadata)
+        attention_mask = metadata_plan.attention_mask
+        full_attn_spans = metadata_plan.full_attn_spans
+        extra = metadata_plan.extra
 
         # Try piecewise attention
         if full_attn_spans is not None:
@@ -443,12 +712,7 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 query_ranges=None if attn_metadata is None else attn_metadata.query_ranges,
             )
 
-        packed_keys = ("cu_seqlens_q", "cu_seqlens_k", "max_seqlen_q", "max_seqlen_k")
-        present_packed_keys = [key for key in packed_keys if key in extra]
-        if present_packed_keys:
-            if len(present_packed_keys) != len(packed_keys):
-                missing = sorted(set(packed_keys) - set(present_packed_keys))
-                raise ValueError(f"Incomplete packed FlashAttention metadata; missing {missing}")
+        if metadata_plan.packing_mode is not PackingMode.NONE:
             self._warn_fa_deterministic_non_dense("packed-varlen")
             return self._forward_varlen_packed(
                 query,
@@ -460,7 +724,14 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
                 max_seqlen_k=extra["max_seqlen_k"],
             )
 
-        if attention_mask is not None and torch.any(~attention_mask):
+        use_masked_path = False
+        if attention_mask is not None:
+            use_masked_path = (
+                torch.any(~attention_mask)
+                if metadata_plan.mask_mode is MaskMode.UNKNOWN
+                else metadata_plan.mask_mode is not MaskMode.NONE
+            )
+        if use_masked_path:
             self._warn_fa_deterministic_non_dense("masked-varlen")
             return self._forward_varlen_masked(
                 query,
@@ -470,6 +741,15 @@ class FlashAttentionImpl(AttentionImpl[AttentionMetadata]):
             )
 
         if flash_attn_func is not None:
+            if IS_FLASH_ATTN_4:
+                return _fa4_dense_attention_op(
+                    query,
+                    key,
+                    value,
+                    self.softmax_scale,
+                    self.causal,
+                    self.fa_deterministic,
+                )
             fa_kwargs = {
                 "causal": self.causal,
                 "softmax_scale": self.softmax_scale,

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -59,6 +59,8 @@ def is_flash_attn_4_available() -> bool:
 # Bind via aliases so mypy does not treat each candidate import as a redefinition.
 flash_attn_func: FlashAttnFn | None = None
 flash_attn_varlen_func: FlashAttnFn | None = None
+IS_FLASH_ATTN_4 = False
+IS_AITER = False
 
 if current_omni_platform.is_rocm():
     # ROCm: try Aiter first
@@ -71,6 +73,7 @@ if current_omni_platform.is_rocm():
 
             flash_attn_func = _fa_func
             flash_attn_varlen_func = _fa_varlen
+            IS_AITER = True
     except (ImportError, ModuleNotFoundError):
         pass
 elif current_omni_platform.is_xpu():
@@ -102,6 +105,7 @@ else:
 
                 flash_attn_func = _fa_func
                 flash_attn_varlen_func = _fa_varlen
+                IS_FLASH_ATTN_4 = True
                 logger.info("Using CuTe FlashAttention-4 on Blackwell")
             except Exception as exc:
                 # Optional FA4 dependencies may be present but ABI-incompatible
@@ -166,6 +170,27 @@ else:
 HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not None
 
 
+def validate_fa4_head_dims(head_dim: int, head_dim_v: int, alignment: int) -> bool:
+    """Use the selected FA4 implementation's constraints during path resolution.
+
+    Return False when its validator is unavailable; propagate kernel validation
+    errors for the backend to report. Keep the private FA4 API dependency here,
+    rather than duplicating architecture/version-specific dimension rules.
+    This is a planning-time call, never part of compiled tensor execution.
+    """
+    try:
+        from flash_attn.cute.interface import _get_device_arch, _validate_head_dims
+    except ImportError:
+        return False
+
+    arch = _get_device_arch() // 10
+    # FA4 routes SM80/SM120 through separate kernels and bypasses this validator.
+    if arch not in (9, 10, 11):
+        return False
+    _validate_head_dims(head_dim, head_dim_v, arch, alignment)
+    return True
+
+
 def _choose_vllm_flash_attn_version(
     device_major: int,
     requested: str | int | None,

vllm_omni/diffusion/attention/capabilities.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+
+class SupportStatus(str, Enum):
+    SUPPORTED = "supported"
+    UNSUPPORTED = "unsupported"
+    UNMIGRATED = "unmigrated"
+
+
+class CompilationMode(str, Enum):
+    TRACEABLE = "traceable"
+    CUSTOM_OP = "custom_op"
+    EAGER_ONLY = "eager_only"
+
+
+class PackingMode(str, Enum):
+    NONE = "none"
+    PACKED_PADDING = "packed_padding"
+    MULTI_DOCUMENT = "multi_document"
+
+
+class MaskMode(str, Enum):
+    NONE = "none"
+    PADDING = "padding"
+    ARBITRARY = "arbitrary"
+    UNKNOWN = "unknown"
+
+
+class ParallelStrategy(str, Enum):
+    NONE = "none"
+    ULYSSES = "ulysses"
+    RING = "ring"
+    HYBRID_ULYSSES_RING = "hybrid_ulysses_ring"
+    ALLGATHER_KV = "allgather_kv"
+
+
+class OuterBoundary(str, Enum):
+    HSDP = "hsdp"
+
+
+@dataclass(frozen=True, slots=True)
+class CapabilityResult:
+    status: SupportStatus
+    reason: str | None = None
+
+    def __post_init__(self) -> None:
+        if self.status is SupportStatus.UNSUPPORTED and not self.reason:
+            raise ValueError("UNSUPPORTED capability results require an actionable reason")
+
+    @classmethod
+    def supported(cls) -> CapabilityResult:
+        return cls(SupportStatus.SUPPORTED)
+
+    @classmethod
+    def unsupported(cls, reason: str) -> CapabilityResult:
+        return cls(SupportStatus.UNSUPPORTED, reason)
+
+    @classmethod
+    def unmigrated(cls, reason: str | None = None) -> CapabilityResult:
+        return cls(SupportStatus.UNMIGRATED, reason)
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionContext:
+    platform: str
+    kernel_variant: str | None = None
+    dtype: str | None = None
+    causal: bool | None = None
+    mask_mode: MaskMode = MaskMode.NONE
+    packing_mode: PackingMode = PackingMode.NONE
+    piecewise: bool = False
+    paged_kv: bool = False
+    kv_cache_dtype: str | None = None
+    parallel_strategy: ParallelStrategy = ParallelStrategy.NONE
+    require_fullgraph: bool = False
+    outer_boundaries: frozenset[OuterBoundary] = frozenset()
+
+
+@dataclass(frozen=True, slots=True)
+class ExecutionPathResult:
+    backend: str
+    path: str
+    support: CapabilityResult
+    compilation_mode: CompilationMode
+    platform: str
+    kernel_variant: str | None
+    parallel_strategy: ParallelStrategy
+
+    @classmethod
+    def unmigrated(
+        cls,
+        backend: str,
+        context: ExecutionContext,
+        *,
+        path: str = "unmigrated",
+    ) -> ExecutionPathResult:
+        return cls(
+            backend=backend,
+            path=path,
+            support=CapabilityResult.unmigrated(),
+            compilation_mode=CompilationMode.EAGER_ONLY,
+            platform=context.platform,
+            kernel_variant=context.kernel_variant,
+            parallel_strategy=context.parallel_strategy,
+        )
+
+    def requested_support(self, context: ExecutionContext) -> CapabilityResult:
+        """Validate request-level guarantees without changing path migration state."""
+        if context.require_fullgraph and self.support.status is SupportStatus.UNMIGRATED:
+            return CapabilityResult.unsupported(
+                f"{self.backend} path {self.path!r} has no verified fullgraph declaration; "
+                "use a verified path or disable the guaranteed-fullgraph request"
+            )
+        if self.support.status is not SupportStatus.SUPPORTED:
+            return self.support
+        if context.require_fullgraph and self.compilation_mode is CompilationMode.EAGER_ONLY:
+            return CapabilityResult.unsupported(
+                f"{self.backend} path {self.path!r} does not support guaranteed fullgraph compilation"
+            )
+        return CapabilityResult.supported()

vllm_omni/diffusion/attention/layer.py

@@ -19,7 +19,13 @@ from vllm.model_executor.models.utils import extract_layer_index
 from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheSpec
 
 from vllm_omni.diffusion.attention.backends.abstract import AttentionBackend, AttentionImpl, AttentionMetadata
-from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend, SDPAImpl
+from vllm_omni.diffusion.attention.capabilities import (
+    ExecutionContext,
+    ExecutionPathResult,
+    OuterBoundary,
+    ParallelStrategy,
+)
 from vllm_omni.diffusion.attention.parallel import build_parallel_attention_strategy
 from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
 from vllm_omni.diffusion.attention.parallel.ring import RingParallelAttention
@@ -119,6 +125,8 @@ class Attention(nn.Module):
 
         config = get_current_diffusion_config_or_none()
         attention_config = config.diffusion_attention_config if config is not None else None
+        parallel_config = getattr(config, "parallel_config", None)
+        self._hsdp_compile_boundary_enabled = bool(getattr(parallel_config, "use_hsdp", False))
 
         from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
 
@@ -131,6 +139,9 @@ class Attention(nn.Module):
             is DiffusionKVCacheMode.PAGED_SCHEDULER
         )
         self._scheduler_paged_kv = scheduler_paged_kv
+        self.attn_backend: type[AttentionBackend] | None
+        self.attention: AttentionImpl | nn.Module
+        self.sdpa_fallback: SDPAImpl | None
         if custom_attention is None:
             attn_backend_cls, spec = get_attn_backend_for_role(
                 role=role,
@@ -166,7 +177,6 @@ class Attention(nn.Module):
                     attn_backend_cls.get_name(),
                     dense_backend_name,
                 )
-            parallel_config = getattr(config, "parallel_config", None)
             allgather_degree = getattr(parallel_config, "allgather_degree", 1)
             # TODO: Move AllGather-KV compatibility into an AttentionBackend capability
             # so validation does not depend on backend names.
@@ -279,6 +289,7 @@ class Attention(nn.Module):
 
         if self.paged_kv_cache_role is None:
             return None
+        assert self.attn_backend is not None  # Custom attention cannot opt into paged KV.
         dtype = self.paged_kv_cache_dtype or vllm_config.model_config.dtype
         # Keep backend layout discovery under the same config context used by
         # upstream vLLM's attention-spec collector.  vLLM 0.29 moved the
@@ -378,17 +389,68 @@ class Attention(nn.Module):
         value: torch.Tensor,
         attn_metadata: AttentionMetadata | None = None,
     ) -> torch.Tensor:
-        if torch.compiler.is_compiling() and is_forward_context_available():
-            od_config = get_forward_context().omni_diffusion_config
-            parallel_config = getattr(od_config, "parallel_config", None)
-            if getattr(parallel_config, "use_hsdp", False):
-                # Keep HSDP/FSDP2 parameter all-gather outside Inductor's
-                # attention graph; otherwise scheduler dependency analysis can
-                # fail on the fused attention region.
-                return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata)
+        if torch.compiler.is_compiling() and self._uses_hsdp_compile_boundary():
+            # Keep HSDP/FSDP2 parameter all-gather outside Inductor's
+            # attention graph; otherwise scheduler dependency analysis can
+            # fail on the fused attention region.
+            return self._forward_hsdp_compile_boundary(query, key, value, attn_metadata)
 
         return self._forward_impl(query, key, value, attn_metadata)
 
+    def _uses_hsdp_compile_boundary(self) -> bool:
+        if self._hsdp_compile_boundary_enabled:
+            return True
+        if not is_forward_context_available():
+            return False
+        od_config = get_forward_context().omni_diffusion_config
+        parallel_config = getattr(od_config, "parallel_config", None)
+        return bool(getattr(parallel_config, "use_hsdp", False))
+
+    def resolve_execution_path(
+        self,
+        context: ExecutionContext,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AttentionMetadata | None,
+    ) -> ExecutionPathResult:
+        """Compose backend capabilities with outer Attention boundaries."""
+        boundaries = set(context.outer_boundaries)
+        if self._uses_hsdp_compile_boundary():
+            boundaries.add(OuterBoundary.HSDP)
+        active_strategy = self._get_active_parallel_strategy()
+        strategy_name = active_strategy.name
+        if self.use_ring and active_strategy.enabled and strategy_name == "ulysses":
+            parallel_strategy = ParallelStrategy.HYBRID_ULYSSES_RING
+        elif self.use_ring and active_strategy.enabled:
+            parallel_strategy = ParallelStrategy.RING
+        else:
+            parallel_strategy = {
+                "allgather_kv": ParallelStrategy.ALLGATHER_KV,
+                "ulysses": ParallelStrategy.ULYSSES,
+            }.get(strategy_name, ParallelStrategy.NONE)
+        resolved_context = replace(
+            context,
+            outer_boundaries=frozenset(boundaries),
+            paged_kv=self.is_paged_kv_active(),
+            parallel_strategy=parallel_strategy,
+        )
+        resolver = getattr(self.attention, "resolve_execution_path", None)
+        if not callable(resolver):
+            return ExecutionPathResult.unmigrated(
+                type(self.attention).__name__,
+                resolved_context,
+            )
+        if not resolved_context.paged_kv:
+            attn_metadata = self._with_kv_cache_dtype(attn_metadata)
+        return resolver(
+            resolved_context,
+            query,
+            key,
+            value,
+            attn_metadata,
+        )
+
     @torch.compiler.disable
     def _forward_hsdp_compile_boundary(
         self,
@@ -422,10 +484,11 @@ class Attention(nn.Module):
             )
         use_paged_attention = paged_adapter is not None and self.paged_kv_cache_role is not None
         if use_paged_attention and not getattr(self.attn_backend, "supports_paged_kv", False):
-            assert self.attn_backend is not None
+            backend_name = (
+                self.attn_backend.get_name() if self.attn_backend is not None else type(self.attention).__name__
+            )
             raise NotImplementedError(
-                f"Diffusion paged KV requires an Omni backend with paged support; "
-                f"selected {self.attn_backend.get_name()}"
+                f"Diffusion paged KV requires an Omni backend with paged support; selected {backend_name}"
             )
         if use_paged_attention and strategy is not self._no_parallel_strategy:
             strategy_name = strategy.name
@@ -554,9 +617,11 @@ class Attention(nn.Module):
             self._scheduler_paged_kv
             and self.paged_kv_cache_role is not None
             and in_kv_memory_profile
-            and cast(type[AttentionBackend], self.attn_backend).get_name() == "FLASH_ATTN"
+            and self.attn_backend is not None
+            and self.attn_backend.get_name() == "FLASH_ATTN"
             and not current_omni_platform.supports_diffusion_dense_flash_attention()
         ):
+            assert self.sdpa_fallback is not None
             logger.warning_once(
                 "The startup KV memory profile is using SDPA because dense FLASH_ATTN is unavailable. "
                 "Formal paged requests still use the platform-native paged attention backend."

vllm_omni/diffusion/data.py

@@ -983,6 +983,9 @@ class OmniDiffusionConfig:
     enable_broadcast_weight_load: bool = False
     num_weight_load_threads: int = 4
 
+    # Shard meta parameters before reading rank-local HF safetensors slices.
+    hsdp_weight_load_strategy: str = "full"
+
     # Enable sleep mode
     enable_sleep_mode: bool = False
 
@@ -1196,6 +1199,10 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if self.hsdp_weight_load_strategy not in {"full", "pre_sharded"}:
+            raise ValueError(
+                f"hsdp_weight_load_strategy must be 'full' or 'pre_sharded', got {self.hsdp_weight_load_strategy!r}"
+            )
         from vllm_omni.diffusion.offloader.config import (
             OffloadStrategy,
             materialize_legacy_offload_flags,

vllm_omni/diffusion/distributed/hsdp.py

@@ -46,6 +46,14 @@ class HSDPInferenceConfig:
     reshard_after_forward: bool = True
 
 
+@dataclass
+class HSDPShardContext:
+    """Reusable state for incrementally sharding one HSDP model."""
+
+    hsdp_kwargs: dict[str, Any]
+    ignored_params: set[nn.Parameter] | None = None
+
+
 def _create_hsdp_mesh(
     device_type: str,
     replicate_size: int,
@@ -111,6 +119,38 @@ def apply_hsdp_to_model(
     Returns:
         HSDP-wrapped model ready for inference
     """
+    context = prepare_hsdp_shard_context(
+        model,
+        hsdp_config,
+        target_device=target_device,
+    )
+
+    hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)
+    if not hsdp_shard_conditions:
+        raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined")
+
+    # Apply HSDP sharding, this will automatically handle weight distribution
+    shard_model(
+        model,
+        hsdp_shard_conditions=hsdp_shard_conditions,
+        context=context,
+    )
+
+    logger.info("HSDP applied to model: %s", type(model).__name__)
+    return model
+
+
+def prepare_hsdp_shard_context(
+    model: nn.Module,
+    hsdp_config: HSDPInferenceConfig,
+    target_device: torch.device | None = None,
+) -> HSDPShardContext:
+    """Prepare mesh, precision, and ignored-parameter state for HSDP.
+
+    The returned context can be reused to shard completed child modules while
+    checkpoint weights are still streaming, then to shard the root once all
+    children have been loaded.
+    """
     if not hsdp_config.enabled:
         raise ValueError("HSDP is not enabled in config")
 
@@ -160,7 +200,7 @@ def apply_hsdp_to_model(
     )
 
     hsdp_shard_conditions = getattr(model, "_hsdp_shard_conditions", None)
-    if not hsdp_shard_conditions or len(hsdp_shard_conditions) == 0:
+    if not hsdp_shard_conditions:
         raise ValueError(f"Model {type(model).__name__} has no _hsdp_shard_conditions defined")
 
     # Collect parameters of any modules the model wants excluded from FSDP sharding.
@@ -188,7 +228,11 @@ def apply_hsdp_to_model(
         except AttributeError:
             logger.warning("_hsdp_ignored_modules entry %r not found on model", mod_name)
             continue
-        sub_mod.to(target_device)
+        has_meta_tensor = any(tensor.device.type == "meta" for tensor in (*sub_mod.parameters(), *sub_mod.buffers()))
+        if has_meta_tensor:
+            sub_mod.to_empty(device=target_device)
+        else:
+            sub_mod.to(target_device)
         ignored_params.update(sub_mod.parameters())
     if ignored_params:
         logger.info(
@@ -224,21 +268,36 @@ def apply_hsdp_to_model(
             target_device,
         )
 
-    # Apply HSDP sharding, this will automatically handle weight distribution
-    shard_model(
-        model,
-        reshard_after_forward=hsdp_config.reshard_after_forward,
-        mp_policy=mp_policy,
-        mesh=device_mesh,
-        hsdp_shard_conditions=hsdp_shard_conditions,
-        ignored_params=ignored_params if ignored_params else None,
+    return HSDPShardContext(
+        hsdp_kwargs={
+            "reshard_after_forward": hsdp_config.reshard_after_forward,
+            "mesh": device_mesh,
+            "mp_policy": mp_policy,
+        },
+        ignored_params=ignored_params or None,
     )
 
-    for param in model.parameters():
+
+def shard_hsdp_module(module: nn.Module, context: HSDPShardContext) -> None:
+    """Shard one completed child module with a prepared HSDP context."""
+    module_kwargs = dict(context.hsdp_kwargs)
+    if context.ignored_params:
+        module_ignored_params = context.ignored_params.intersection(module.parameters())
+        if module_ignored_params:
+            module_kwargs["ignored_params"] = module_ignored_params
+    fully_shard(module, **module_kwargs)
+    for param in module.parameters():
         param.requires_grad = False
 
-    logger.info("HSDP applied to model: %s", type(model).__name__)
-    return model
+
+def finalize_hsdp_root(model: nn.Module, context: HSDPShardContext) -> None:
+    """Shard the model root after all selected child modules are sharded."""
+    root_kwargs = dict(context.hsdp_kwargs)
+    if context.ignored_params:
+        root_kwargs["ignored_params"] = context.ignored_params
+    fully_shard(model, **root_kwargs)
+    for param in model.parameters():
+        param.requires_grad = False
 
 
 def shard_model(
@@ -249,6 +308,7 @@ def shard_model(
     mesh: DeviceMesh | None = None,
     hsdp_shard_conditions: list[Callable[[str, nn.Module], bool]],
     ignored_params: set[nn.Parameter] | None = None,
+    context: HSDPShardContext | None = None,
 ) -> None:
     """Apply HSDP sharding to model modules based on shard conditions.
 
@@ -259,32 +319,30 @@ def shard_model(
     This is required for packed integer parameters inside sharded transformer
     blocks; the root wrap receives the full set for all remaining parameters.
     """
-    hsdp_kwargs: dict[str, Any] = {
-        "reshard_after_forward": reshard_after_forward,
-        "mesh": mesh,
-        "mp_policy": mp_policy,
-    }
+    if context is None:
+        context = HSDPShardContext(
+            hsdp_kwargs={
+                "reshard_after_forward": reshard_after_forward,
+                "mesh": mesh,
+                "mp_policy": mp_policy,
+            },
+            ignored_params=ignored_params,
+        )
+    elif any(value is not None for value in (mp_policy, mesh, ignored_params)) or not reshard_after_forward:
+        raise ValueError("Pass either context or individual HSDP shard options, not both")
 
     num_sharded = 0
     for name, module in reversed(list(model.named_modules())):
         if any(cond(name, module) for cond in hsdp_shard_conditions):
-            module_kwargs = dict(hsdp_kwargs)
-            if ignored_params:
-                module_ignored_params = ignored_params.intersection(module.parameters())
-                if module_ignored_params:
-                    module_kwargs["ignored_params"] = module_ignored_params
-            fully_shard(module, **module_kwargs)
+            shard_hsdp_module(module, context)
             num_sharded += 1
 
     if num_sharded == 0:
         raise ValueError("No modules were sharded")
 
-    root_kwargs = dict(hsdp_kwargs)
-    if ignored_params:
-        root_kwargs["ignored_params"] = ignored_params
-    fully_shard(model, **root_kwargs)
+    finalize_hsdp_root(model, context)
     logger.info(
         "Sharded %d modules + root (ignored_params=%d)",
         num_sharded,
-        len(ignored_params) if ignored_params else 0,
+        len(context.ignored_params) if context.ignored_params else 0,
     )

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -16,6 +16,7 @@ import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
     download_weights_from_hf,
@@ -30,7 +31,14 @@ from vllm.utils.torch_utils import set_default_torch_dtype
 
 from vllm_omni.diffusion.config import set_current_diffusion_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig, apply_hsdp_to_model
+from vllm_omni.diffusion.distributed.hsdp import (
+    HSDPInferenceConfig,
+    HSDPShardContext,
+    apply_hsdp_to_model,
+    finalize_hsdp_root,
+    prepare_hsdp_shard_context,
+    shard_hsdp_module,
+)
 from vllm_omni.diffusion.lora.manager import LoRABackend
 from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
     get_checkpoint_adapter,
@@ -38,6 +46,8 @@ from vllm_omni.diffusion.model_loader.checkpoint_adapters import (
 from vllm_omni.diffusion.model_loader.host_weight_loader import HWRLoaderMixin, _HWRCommitError
 from vllm_omni.diffusion.model_loader.host_weight_plan import (
     HostWeightPlan,
+    TensorBinding,
+    build_checkpoint_binding_plan,
     build_checkpoint_mmap_plan,
     has_online_quantization,
 )
@@ -49,7 +59,7 @@ from vllm_omni.diffusion.offloader.config import (
     OffloadStrategy,
     resolve_offload,
 )
-from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery, PipelineModules
 from vllm_omni.diffusion.offloader.offload_plan import get_offload_plan
 from vllm_omni.diffusion.registry import initialize_model
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
@@ -138,6 +148,36 @@ def _resolve_custom_pipeline_cls(custom_pipeline_name: str | type | None) -> typ
     )
 
 
+def _is_distributed_rank_zero() -> bool:
+    return (
+        not torch.distributed.is_available()
+        or not torch.distributed.is_initialized()
+        or torch.distributed.get_rank() == 0
+    )
+
+
+@dataclasses.dataclass(frozen=True)
+class _HSDPLoadRoot:
+    name: str
+    module: nn.Module
+    context: HSDPShardContext
+
+
+@dataclasses.dataclass(frozen=True)
+class _HSDPModuleLoadGroup:
+    name: str
+    module: nn.Module
+    context: HSDPShardContext
+
+
+@dataclasses.dataclass(frozen=True)
+class _PreShardedHSDPLoadPlan:
+    roots: list[_HSDPLoadRoot]
+    groups: list[_HSDPModuleLoadGroup]
+    bindings: dict[str, TensorBinding]
+    binding_names: frozenset[str]
+
+
 class DiffusersPipelineLoader(HWRLoaderMixin):
     """Model loader that can load diffusers pipeline components from disk."""
 
@@ -1222,6 +1262,451 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
                     raise ValueError(f"Unknown load_format: {load_format}")
         return model
 
+    def _load_model_with_pre_sharded_hsdp(
+        self,
+        model: nn.Module,
+        discovered_modules: PipelineModules,
+        hsdp_config: HSDPInferenceConfig,
+        target_device: torch.device,
+    ) -> nn.Module:
+        """Shard meta DiT parameters before reading their rank-local checkpoint slices."""
+        from torch.distributed.checkpoint import HuggingFaceStorageReader
+        from torch.distributed.checkpoint import load as dcp_load
+
+        if self.counter_before_loading_weights == 0.0:
+            self.counter_before_loading_weights = time.perf_counter()
+
+        load_plan = self._prepare_pre_sharded_hsdp_load(
+            model,
+            discovered_modules,
+            hsdp_config,
+            target_device,
+        )
+        self._check_no_unplanned_meta_tensors(model, load_plan.binding_names)
+        self._validate_pre_sharded_hsdp_bindings(load_plan)
+
+        self._process_pre_sharded_hsdp_groups_on_meta(load_plan)
+
+        for group in load_plan.groups:
+            logger.debug("Applying meta-first HSDP shard to %s", group.name)
+            shard_hsdp_module(group.module, group.context)
+        for root in load_plan.roots:
+            logger.debug("Finalizing meta-first HSDP root %s", root.name)
+            finalize_hsdp_root(root.module, root.context)
+
+        self._materialize_pre_sharded_hsdp_state(model, load_plan, target_device)
+
+        targets = self._get_pre_sharded_hsdp_targets(model, load_plan)
+        checkpoint_states, local_payload_bytes = self._prepare_pre_sharded_checkpoint_states(
+            load_plan,
+            targets,
+        )
+
+        thread_count = int(getattr(self.od_config, "num_weight_load_threads", 4))
+        if _is_distributed_rank_zero():
+            directory_label = "directory" if len(checkpoint_states) == 1 else "directories"
+            logger.info(
+                "Pre-sharded HSDP loading %.2f GiB of rank-local tensor slices directly into "
+                "FSDP-owned storage from %d HF safetensors %s with %d reader thread(s) per rank",
+                local_payload_bytes / 1024**3,
+                len(checkpoint_states),
+                directory_label,
+                thread_count,
+            )
+
+        for checkpoint_dir in sorted(checkpoint_states, key=str):
+            reader = HuggingFaceStorageReader(str(checkpoint_dir), thread_count=thread_count)
+            # Each worker has the full checkpoint available and each DTensor
+            # already describes its rank-local destination slice. Avoid DCP's
+            # coordinator collectives, which can leave rank 0 with less GPU
+            # memory headroom than the other HSDP workers.
+            dcp_load(checkpoint_states[checkpoint_dir], storage_reader=reader, no_dist=True)
+
+        remaining_meta = self._pre_sharded_meta_target_names(model, load_plan)
+        if remaining_meta:
+            raise ValueError(
+                "Pre-sharded HSDP checkpoint loading left tensors on meta; "
+                f"first missing tensors: {remaining_meta[:16]}"
+            )
+
+        self._finalize_pre_sharded_hsdp_loaded_weights(load_plan)
+
+        self.counter_after_loading_weights = time.perf_counter()
+        logger.info_once(
+            "Loading weights took %.2f seconds",
+            self.counter_after_loading_weights - self.counter_before_loading_weights,
+        )
+        self._move_non_hsdp_modules(discovered_modules, target_device)
+        return model
+
+    @staticmethod
+    def _validate_pre_sharded_hsdp_bindings(load_plan: _PreShardedHSDPLoadPlan) -> None:
+        unsupported = [name for name, binding in load_plan.bindings.items() if binding.transform is not None]
+        if unsupported:
+            raise ValueError(
+                "Pre-sharded HSDP requires checkpoint tensors in runtime layout; "
+                f"tensor transforms are unsupported: {unsupported[:16]}"
+            )
+        non_safetensors = sorted(
+            {
+                binding.file_path
+                for binding in load_plan.bindings.values()
+                if not binding.file_path.endswith(".safetensors")
+            }
+        )
+        if non_safetensors:
+            raise ValueError(
+                f"Pre-sharded HSDP requires existing HF safetensors files; unsupported files: {non_safetensors[:8]}"
+            )
+
+    @staticmethod
+    def _get_pre_sharded_hsdp_targets(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> dict[str, torch.Tensor]:
+        targets: dict[str, torch.Tensor] = {
+            **dict(model.named_buffers()),
+            **dict(model.named_parameters()),
+        }
+        missing_targets = sorted(load_plan.binding_names - targets.keys())
+        if missing_targets:
+            raise ValueError(
+                f"Pre-sharded HSDP changed checkpoint target names; first missing targets: {missing_targets[:16]}"
+            )
+        return targets
+
+    @staticmethod
+    def _materialize_pre_sharded_hsdp_state(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+        target_device: torch.device,
+    ) -> None:
+        """Materialize the installed FSDP rank-local state before checkpoint I/O.
+
+        ``fully_shard`` supports meta initialization through ``Module.to_empty``:
+        its registered ``_apply`` hooks update the internal sharded storage to
+        point at the newly allocated rank-local tensors. Modules are materialized
+        in the same HSDP-group order and child-first traversal used by ordinary
+        ``fully_shard`` initialization instead of the forward traversal used by
+        a recursive ``to_empty`` call.
+        PyTorch's swap-on-conversion mode preserves Parameter identity so that
+        post-load kernel objects created before sharding cannot retain stale
+        references to the original meta Parameters.
+
+        A complete checkpoint plan covers every parameter and persistent buffer
+        in each DiT.  Non-persistent buffers are not checkpoint entries, so keep
+        their initialized values across ``to_empty``.
+        """
+        unplanned_parameters: list[str] = []
+        saved_buffers: dict[str, torch.Tensor] = {}
+        for root in load_plan.roots:
+            prefix = f"{root.name}."
+            for local_name, _parameter in root.module.named_parameters():
+                full_name = prefix + local_name
+                if full_name not in load_plan.binding_names:
+                    unplanned_parameters.append(full_name)
+            for local_name, buffer in root.module.named_buffers():
+                full_name = prefix + local_name
+                if full_name in load_plan.binding_names:
+                    continue
+                if buffer.device.type == "meta":
+                    raise ValueError(f"Pre-sharded HSDP cannot preserve an unplanned meta buffer: {full_name}")
+                saved_buffers[full_name] = buffer.detach().cpu().clone()
+
+        if unplanned_parameters:
+            raise ValueError(
+                "Pre-sharded HSDP requires checkpoint coverage for every DiT parameter; "
+                f"first unplanned parameters: {unplanned_parameters[:16]}"
+            )
+
+        previous_swap_mode = torch.__future__.get_swap_module_params_on_conversion()
+        torch.__future__.set_swap_module_params_on_conversion(True)
+        try:
+            sharded_group_roots = {group.module for group in load_plan.groups}
+            materialized_modules: set[nn.Module] = set()
+            for group in load_plan.groups:
+                for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+                    group.module,
+                    nested_hsdp_roots=sharded_group_roots - {group.module},
+                ):
+                    if module not in materialized_modules:
+                        module.to_empty(device=target_device, recurse=False)
+                        materialized_modules.add(module)
+            for root in load_plan.roots:
+                for module in DiffusersPipelineLoader._hsdp_managed_modules_post_order(
+                    root.module,
+                    nested_hsdp_roots=sharded_group_roots,
+                ):
+                    if module not in materialized_modules:
+                        module.to_empty(device=target_device, recurse=False)
+                        materialized_modules.add(module)
+        finally:
+            torch.__future__.set_swap_module_params_on_conversion(previous_swap_mode)
+
+        if saved_buffers:
+            materialized_buffers = dict(model.named_buffers())
+            with torch.no_grad():
+                for name, value in saved_buffers.items():
+                    target = materialized_buffers[name]
+                    target.copy_(value.to(device=target.device, dtype=target.dtype))
+
+    @staticmethod
+    def _hsdp_managed_modules_post_order(
+        root: nn.Module,
+        *,
+        nested_hsdp_roots: set[nn.Module],
+    ) -> list[nn.Module]:
+        """Mirror FSDP's child-first managed-module traversal.
+
+        Nested FSDP roots own their own parameters and are materialized in their
+        corresponding load group, so the enclosing traversal must not descend
+        into them.
+        """
+        modules: list[nn.Module] = []
+        visited: set[nn.Module] = set()
+
+        def visit(module: nn.Module) -> None:
+            if module in visited:
+                return
+            visited.add(module)
+            for child in module.children():
+                if child not in nested_hsdp_roots:
+                    visit(child)
+            modules.append(module)
+
+        visit(root)
+        return modules
+
+    @staticmethod
+    def _prepare_pre_sharded_checkpoint_states(
+        load_plan: _PreShardedHSDPLoadPlan,
+        targets: dict[str, torch.Tensor],
+    ) -> tuple[dict[Path, dict[str, torch.Tensor]], int]:
+        from torch.distributed.tensor import DTensor
+
+        checkpoint_states: dict[Path, dict[str, torch.Tensor]] = {}
+        checkpoint_owners: dict[tuple[Path, str], str] = {}
+        local_payload_bytes = 0
+        for name, binding in load_plan.bindings.items():
+            target = targets[name]
+
+            checkpoint_dir = Path(binding.file_path).parent
+            owner_key = (checkpoint_dir, binding.checkpoint_key)
+            if previous_owner := checkpoint_owners.get(owner_key):
+                raise ValueError(
+                    "Pre-sharded HSDP cannot bind one checkpoint tensor to multiple runtime tensors: "
+                    f"{binding.checkpoint_key!r} maps to both {previous_owner!r} and {name!r}"
+                )
+            checkpoint_owners[owner_key] = name
+
+            if target.device.type == "meta":
+                raise ValueError(f"Pre-sharded HSDP target {name!r} was not materialized")
+
+            local_target = target.to_local() if isinstance(target, DTensor) else target
+            local_payload_bytes += local_target.numel() * local_target.element_size()
+            checkpoint_states.setdefault(checkpoint_dir, {})[binding.checkpoint_key] = target
+        return checkpoint_states, local_payload_bytes
+
+    @staticmethod
+    def _pre_sharded_meta_target_names(
+        model: nn.Module,
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> list[str]:
+        from torch.distributed.tensor import DTensor
+
+        targets = {
+            **dict(model.named_buffers()),
+            **dict(model.named_parameters()),
+        }
+        return [
+            name
+            for name in load_plan.binding_names
+            if (target := targets[name]).device.type == "meta"
+            or (isinstance(target, DTensor) and target.to_local().device.type == "meta")
+        ]
+
+    @staticmethod
+    def _process_pre_sharded_hsdp_groups_on_meta(load_plan: _PreShardedHSDPLoadPlan) -> None:
+        modules_to_process: list[tuple[nn.Module, QuantizeMethodBase]] = []
+        for group in load_plan.groups:
+            for _, module in group.module.named_modules():
+                quant_method = getattr(module, "quant_method", None)
+                if quant_method is None or not isinstance(quant_method, QuantizeMethodBase):
+                    continue
+                modules_to_process.append((module, quant_method))
+
+        if not modules_to_process:
+            return
+        for root in load_plan.roots:
+            if not getattr(root.module, "_hsdp_pre_sharded_meta_post_load", False):
+                raise ValueError(
+                    f"Model {type(root.module).__name__} has not declared its post-load processing safe "
+                    "for pre-sharded meta initialization"
+                )
+        for module, quant_method in modules_to_process:
+            if type(quant_method) is not UnquantizedLinearMethod:
+                raise ValueError(
+                    "Pre-sharded HSDP does not support the selected post-load processing; "
+                    f"{type(module).__name__} uses {type(quant_method).__name__}"
+                )
+        for module, quant_method in modules_to_process:
+            quant_method.process_weights_after_loading(module)
+
+    def _prepare_pre_sharded_hsdp_load(
+        self,
+        model: nn.Module,
+        discovered_modules: PipelineModules,
+        hsdp_config: HSDPInferenceConfig,
+        target_device: torch.device,
+    ) -> _PreShardedHSDPLoadPlan:
+        outer_dit_names, outer_dits = discovered_modules.outermost_dits()
+        if not outer_dits:
+            raise ValueError("No DiT modules discovered for HSDP sharding")
+
+        sources = self._get_weight_sources(model)
+        plan_result = build_checkpoint_binding_plan(
+            model,
+            dit_modules=tuple(zip(outer_dit_names, outer_dits)),
+            sources=sources,
+            model_path=str(getattr(self.od_config, "model", "")) or None,
+            tensor_parallel_size=int(getattr(self.parallel_config, "tensor_parallel_size", 1)),
+            online_quantization=False,
+        )
+        if plan_result.plan is None:
+            raise ValueError(f"Pre-sharded HSDP checkpoint loading is incompatible: {plan_result.fallback_reason}")
+        plan = plan_result.plan
+        unsupported_sources = [source.prefix for source in sources if source.prefix not in plan.planned_source_prefixes]
+        if unsupported_sources:
+            raise ValueError(
+                "Pre-sharded HSDP checkpoint loading requires dedicated DiT sources; "
+                f"unsupported source prefixes: {unsupported_sources}"
+            )
+
+        roots: list[_HSDPLoadRoot] = []
+        groups: list[_HSDPModuleLoadGroup] = []
+        for outer_name, outer_dit in zip(outer_dit_names, outer_dits):
+            conditions = getattr(outer_dit, "_hsdp_shard_conditions", None)
+            if not conditions:
+                raise ValueError(f"Model {type(outer_dit).__name__} has no _hsdp_shard_conditions defined")
+
+            context = prepare_hsdp_shard_context(
+                outer_dit,
+                hsdp_config,
+                target_device=target_device,
+            )
+            roots.append(_HSDPLoadRoot(outer_name, outer_dit, context))
+            for local_name, module in reversed(list(outer_dit.named_modules())):
+                if not any(condition(local_name, module) for condition in conditions):
+                    continue
+                if not local_name:
+                    raise ValueError("_hsdp_shard_conditions must not select the DiT root")
+                full_name = f"{outer_name}.{local_name}" if local_name else outer_name
+                groups.append(_HSDPModuleLoadGroup(full_name, module, context))
+
+        parameters = dict(model.named_parameters())
+        targets: dict[str, torch.Tensor] = {**dict(model.named_buffers()), **parameters}
+        preserved_params = {param for root in roots for param in root.context.ignored_params or ()}
+        binding_names = frozenset(plan.bindings)
+        self._release_hsdp_checkpoint_targets_to_meta(
+            model,
+            targets,
+            parameters,
+            binding_names,
+            preserved_params,
+        )
+
+        return _PreShardedHSDPLoadPlan(
+            roots=roots,
+            groups=groups,
+            bindings=plan.bindings,
+            binding_names=binding_names,
+        )
+
+    @staticmethod
+    def _check_no_unplanned_meta_tensors(model: nn.Module, planned_names: frozenset[str]) -> None:
+        unplanned_meta = [
+            name
+            for name, tensor in (*model.named_parameters(), *model.named_buffers())
+            if tensor.device.type == "meta" and name not in planned_names
+        ]
+        if unplanned_meta:
+            raise ValueError(
+                "Pre-sharded HSDP initialized tensors on meta that are not covered by the "
+                f"checkpoint plan. First unplanned tensors: {unplanned_meta[:16]}"
+            )
+
+    @staticmethod
+    def _release_hsdp_checkpoint_targets_to_meta(
+        model: nn.Module,
+        targets: dict[str, torch.Tensor],
+        parameters: dict[str, nn.Parameter],
+        binding_names: frozenset[str],
+        preserved_params: set[nn.Parameter],
+    ) -> None:
+        released: dict[str, tuple[int, float]] = {}
+        for name in binding_names:
+            target = targets[name]
+            if target in preserved_params or target.device.type == "meta":
+                continue
+            logical_gib = target.numel() * target.element_size() / 1024**3
+            device_type = target.device.type
+            count, gib = released.get(device_type, (0, 0.0))
+            released[device_type] = (count + 1, gib + logical_gib)
+            if name in parameters:
+                replacement = DiffusersPipelineLoader._make_parameter_like(
+                    target,
+                    torch.empty_like(target, device="meta"),
+                )
+                torch.utils.swap_tensors(target, replacement)
+            else:
+                replacement = torch.empty_like(target, device="meta")
+                torch.utils.swap_tensors(target, replacement)
+        if released:
+            logger.info("Released HSDP checkpoint-covered tensors to meta before pre-sharded load: %s", released)
+
+    @staticmethod
+    def _make_parameter_like(template: nn.Parameter, tensor: torch.Tensor) -> nn.Parameter:
+        """Create a Parameter on new storage while preserving loader metadata."""
+        if type(template) is nn.Parameter:
+            replacement = nn.Parameter(tensor, requires_grad=template.requires_grad)
+            replacement.__dict__.update(getattr(template, "__dict__", {}))
+            return replacement
+
+        replacement = template.__class__.__new__(template.__class__, tensor)
+        replacement.__dict__.update(getattr(template, "__dict__", {}))
+        replacement.requires_grad_(template.requires_grad)
+        return replacement
+
+    @staticmethod
+    def _finalize_pre_sharded_hsdp_loaded_weights(
+        load_plan: _PreShardedHSDPLoadPlan,
+    ) -> None:
+        """Run model post-load hooks for tensors materialized by pre-sharded loading."""
+        for root in load_plan.roots:
+            post_load = getattr(root.module, "post_load_weights", None)
+            if callable(post_load):
+                post_load()
+            root.module.eval()
+            validate = getattr(root.module, "validate_loaded_weights", None)
+            if callable(validate):
+                prefix = f"{root.name}."
+                validate({name.removeprefix(prefix) for name in load_plan.binding_names if name.startswith(prefix)})
+
+    @staticmethod
+    def _move_non_hsdp_modules(
+        discovered_modules: PipelineModules,
+        target_device: torch.device,
+    ) -> None:
+        modules_to_move: list[nn.Module] = []
+        if discovered_modules.vaes is not None:
+            modules_to_move.extend(discovered_modules.vaes)
+        if discovered_modules.encoders is not None:
+            modules_to_move.extend(discovered_modules.encoders)
+        if discovered_modules.resident_modules is not None:
+            modules_to_move.extend(discovered_modules.resident_modules)
+        for module in modules_to_move:
+            module.to(target_device)
+
     def _load_model_with_hsdp(
         self,
         target_device: torch.device,
@@ -1251,6 +1736,29 @@ class DiffusersPipelineLoader(HWRLoaderMixin):
         # mapping (QKV fusion, etc.).
         if load_format == "diffusers":
             raise ValueError("HSDP is not supported with the diffusers adapter load format")
+        strategy = getattr(self.od_config, "hsdp_weight_load_strategy", "full")
+        if strategy == "pre_sharded":
+            if load_format != "default":
+                raise ValueError(
+                    f"hsdp_weight_load_strategy={strategy!r} currently supports only diffusion_load_format='default'"
+                )
+            if self.quant_config is not None:
+                raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support quantization yet")
+            if getattr(self.od_config, "lora_path", None):
+                raise ValueError(f"hsdp_weight_load_strategy={strategy!r} does not support LoRA yet")
+        elif strategy != "full":
+            raise ValueError(f"Unknown hsdp_weight_load_strategy: {strategy!r}")
+
+        if strategy == "pre_sharded":
+            model = self._init_from_load_format(load_format, target_device, custom_pipeline_name, is_hsdp=True)
+            discovered_modules = ModuleDiscovery.discover(model)
+            return self._load_model_with_pre_sharded_hsdp(
+                model,
+                discovered_modules,
+                hsdp_config,
+                target_device,
+            )
+
         # Same host-fallback bound as the ordinary path: with a quant config the
         # HSDP path initializes on the accelerator (hsdp_defer_to_cpu=False), so
         # over-wide NPU-unquantizable fallback weights would otherwise round-trip

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
-"""Loader-owned host-weight plans shared with diffusion offload backends."""
+"""Loader-owned host-weight plans shared by direct-checkpoint consumers."""
 
 from __future__ import annotations
 
@@ -205,7 +205,7 @@ def _build_source_map(
     if not model_to_ckpt:
         raise _PlanIncompatibleError("no compatible safetensors entries were found in the loader's model sources")
     logger.info(
-        "Indexed %d runtime tensor names from %d checkpoint keys for DLO preflight",
+        "Indexed %d runtime tensor names from %d checkpoint keys for checkpoint preflight",
         len(model_to_ckpt),
         indexed_keys,
     )
@@ -313,28 +313,28 @@ def _validate_source_metadata(
                     )
 
 
-def build_checkpoint_mmap_plan(
+def build_checkpoint_binding_plan(
     pipeline: nn.Module,
     *,
     dit_modules: Sequence[tuple[str, nn.Module]],
     sources: Sequence[object],
     model_path: str | None,
     tensor_parallel_size: int,
-    use_hsdp: bool,
     online_quantization: bool,
+    use_pipeline_key_remap: bool = True,
     has_distilled_lora: bool = False,
 ) -> HostWeightPlanResult:
-    """Build a complete direct-checkpoint plan or return a fallback reason."""
+    """Build a complete binding plan for the pre-sharded HSDP consumer."""
     if tensor_parallel_size != 1:
         return HostWeightPlanResult(None, f"TP={tensor_parallel_size} requires the ordinary loader")
-    if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
     if online_quantization:
         return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
     if has_distilled_lora:
         return HostWeightPlanResult(None, "distilled LoRA requires the ordinary loader")
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
+    if not callable(remap_fn) and use_pipeline_key_remap:
+        remap_fn = getattr(pipeline, "remap_checkpoint_key", None)
     if not callable(remap_fn):
         remap_fn = None
     adapter = get_direct_mmap_adapter(pipeline)
@@ -387,10 +387,37 @@ def build_checkpoint_mmap_plan(
     )
 
 
+def build_checkpoint_mmap_plan(
+    pipeline: nn.Module,
+    *,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    sources: Sequence[object],
+    model_path: str | None,
+    tensor_parallel_size: int,
+    use_hsdp: bool,
+    online_quantization: bool,
+    has_distilled_lora: bool = False,
+) -> HostWeightPlanResult:
+    """Build the DLO mmap plan without changing its HSDP eligibility."""
+    if use_hsdp:
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
+    return build_checkpoint_binding_plan(
+        pipeline,
+        dit_modules=dit_modules,
+        sources=sources,
+        model_path=model_path,
+        tensor_parallel_size=tensor_parallel_size,
+        online_quantization=online_quantization,
+        use_pipeline_key_remap=False,
+        has_distilled_lora=has_distilled_lora,
+    )
+
+
 __all__ = [
     "HostWeightPlan",
     "HostWeightPlanResult",
     "TensorBinding",
+    "build_checkpoint_binding_plan",
     "build_checkpoint_mmap_plan",
     "has_online_quantization",
 ]

vllm_omni/diffusion/models/auk/alias_free_triton.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Triton kernels for the AuK codec's alias-free SnakeBeta activation.
+
+The eager module oversamples with a grouped transposed convolution, applies
+SnakeBeta, and decimates with a grouped strided convolution. Both grouped
+convolutions carry one 12-tap filter shared by every channel, and cuDNN and
+ATen serve them with generic depthwise kernels that dominate the decode (more
+than half of a 12 s decode on H200). Here the oversampling is written as its
+polyphase form fused with the activation, and the decimation as a direct
+strided FIR, each one pass over memory per row. The arithmetic is the same
+as the eager module in fp32; only the summation order of the taps differs.
+
+The ops are registered with ``torch.library.triton_op`` so the compiled
+decode buckets trace them instead of breaking the graph.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+__all__ = ["alias_free_snake"]
+
+_BLOCK = 1024
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _upsample_snake_kernel(
+        x_ptr,
+        out_ptr,
+        taps_ptr,
+        exp_alpha_ptr,
+        inv_beta_ptr,
+        channels,
+        in_len,
+        out_len,
+        ratio: tl.constexpr,
+        taps: tl.constexpr,
+        pad: tl.constexpr,
+        crop: tl.constexpr,
+        block: tl.constexpr,
+    ):
+        # out[m] = snake(ratio * sum_k taps[k] * xpad[(m + crop - k) / ratio]) over the k
+        # with (m + crop - k) divisible by ratio, where xpad is x replicate-padded by pad
+        # on both sides. That is conv_transpose1d(stride=ratio) followed by the crop.
+        row = tl.program_id(0)
+        offs = tl.program_id(1) * block + tl.arange(0, block)
+        mask = offs < out_len
+        x_row = x_ptr + row.to(tl.int64) * in_len
+        acc = tl.zeros((block,), dtype=tl.float32)
+        for k in tl.static_range(taps):
+            pos = offs + (crop - k)
+            hit = (pos % ratio) == 0
+            src = pos // ratio - pad
+            src = tl.minimum(tl.maximum(src, 0), in_len - 1)
+            w = tl.load(taps_ptr + k)
+            val = tl.load(x_row + src, mask=mask & hit, other=0.0)
+            acc += tl.where(hit, w * val, 0.0)
+        u = acc * ratio
+        channel = row % channels
+        alpha = tl.load(exp_alpha_ptr + channel)
+        inv_beta = tl.load(inv_beta_ptr + channel)
+        s = tl.sin(u * alpha)
+        u = u + inv_beta * (s * s)
+        tl.store(out_ptr + row.to(tl.int64) * out_len + offs, u, mask=mask)
+
+    @triton.jit
+    def _lowpass_kernel(
+        x_ptr,
+        out_ptr,
+        taps_ptr,
+        in_len,
+        out_len,
+        stride: tl.constexpr,
+        taps: tl.constexpr,
+        pad_left: tl.constexpr,
+        block: tl.constexpr,
+    ):
+        # out[n] = sum_k taps[k] * x[clamp(stride * n + k - pad_left)], i.e. replicate
+        # padding followed by a strided conv1d.
+        row = tl.program_id(0)
+        offs = tl.program_id(1) * block + tl.arange(0, block)
+        mask = offs < out_len
+        x_row = x_ptr + row.to(tl.int64) * in_len
+        acc = tl.zeros((block,), dtype=tl.float32)
+        for k in tl.static_range(taps):
+            src = offs * stride + (k - pad_left)
+            src = tl.minimum(tl.maximum(src, 0), in_len - 1)
+            acc += tl.load(taps_ptr + k) * tl.load(x_row + src, mask=mask, other=0.0)
+        tl.store(out_ptr + row.to(tl.int64) * out_len + offs, acc, mask=mask)
+
+    @torch.library.triton_op("vllm_omni::auk_alias_free_snake", mutates_args=())
+    def alias_free_snake(
+        x: torch.Tensor,
+        up_taps: torch.Tensor,
+        down_taps: torch.Tensor,
+        exp_alpha: torch.Tensor,
+        inv_beta: torch.Tensor,
+        up_ratio: int,
+        up_pad: int,
+        up_crop: int,
+        down_stride: int,
+        down_pad_left: int,
+        down_pad_right: int,
+    ) -> torch.Tensor:
+        """Upsample by ``up_ratio``, SnakeBeta, low-pass and decimate: ``[B, C, T] -> [B, C, T_out]``."""
+        batch, channels, length = x.shape
+        x = x.contiguous()
+        up_len = length * up_ratio
+        oversampled = torch.empty((batch, channels, up_len), device=x.device, dtype=torch.float32)
+        rows = batch * channels
+        torch.library.wrap_triton(_upsample_snake_kernel)[(rows, triton.cdiv(up_len, _BLOCK))](
+            x,
+            oversampled,
+            up_taps,
+            exp_alpha,
+            inv_beta,
+            channels,
+            length,
+            up_len,
+            ratio=up_ratio,
+            taps=up_taps.numel(),
+            pad=up_pad,
+            crop=up_crop,
+            block=_BLOCK,
+        )
+        taps = down_taps.numel()
+        out_len = (up_len + down_pad_left + down_pad_right - taps) // down_stride + 1
+        out = torch.empty((batch, channels, out_len), device=x.device, dtype=torch.float32)
+        torch.library.wrap_triton(_lowpass_kernel)[(rows, triton.cdiv(out_len, _BLOCK))](
+            oversampled,
+            out,
+            down_taps,
+            up_len,
+            out_len,
+            stride=down_stride,
+            taps=taps,
+            pad_left=down_pad_left,
+            block=_BLOCK,
+        )
+        return out
+
+else:  # pragma: no cover
+
+    def alias_free_snake(*args, **kwargs):
+        raise RuntimeError("Triton is not available")

vllm_omni/diffusion/models/auk/auk_transformer.py

@@ -153,6 +153,33 @@ class TimeEmbedding(nn.Module):
         return self.time_mlp(hidden.to(self.time_mlp[0].weight.dtype))
 
 
+def _grouped_conv1d_btc(x: torch.Tensor, conv: nn.Conv1d) -> torch.Tensor:
+    """``conv`` applied to channels-last ``x`` ``[B, T, C]``, returning ``[B, T, C]``.
+
+    A same-padded grouped convolution written as one im2col copy and one
+    batched GEMM over the groups. cuDNN runs a grouped conv of this shape as
+    one kernel per group plus two layout transposes, which at AuK's sizes
+    (1536 channels in 16 groups, kernel 31, a few hundred frames) costs about
+    ten times the arithmetic.
+    """
+    batch, frames, channels = x.shape
+    groups = conv.groups
+    kernel = conv.kernel_size[0]
+    per_group = channels // groups
+    pad = conv.padding[0]
+    padded = F.pad(x, (0, 0, pad, pad))
+    # unfold gives a [B, T, C, K] view; one copy lays it out as [B, T, G, C/G * K].
+    cols = padded.unfold(1, kernel, 1).reshape(batch * frames, groups, per_group * kernel)
+    # conv.weight [C_out, C/G, K] -> [G, C/G * K, C_out/G]; the reduction order matches the columns.
+    weight = conv.weight.reshape(groups, channels // groups, per_group * kernel).transpose(1, 2).to(cols.dtype)
+    # [G, B*T, C/G*K] @ [G, C/G*K, C_out/G] -> [G, B*T, C_out/G]
+    out = torch.bmm(cols.transpose(0, 1), weight)
+    out = out.transpose(0, 1).reshape(batch, frames, channels)
+    if conv.bias is not None:
+        out = out + conv.bias.to(out.dtype)
+    return out
+
+
 class ConvPosEmbedding(nn.Module):
     """Depthwise-grouped conv stack that adds local position information."""
 
@@ -169,17 +196,19 @@ class ConvPosEmbedding(nn.Module):
         )
 
     def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
-        keep = None if mask is None else mask.unsqueeze(1)
-        x = x.transpose(1, 2)
+        keep = None if mask is None else mask.unsqueeze(-1)
         if keep is not None:
             x = x.masked_fill(~keep, 0.0)
         for layer in self.conv1d:
-            x = layer(x)
-            # Re-zero padding after each convolution so it cannot bleed into
-            # valid frames through the kernel window.
-            if keep is not None and isinstance(layer, nn.Conv1d):
-                x = x.masked_fill(~keep, 0.0)
-        return x.transpose(1, 2)
+            if isinstance(layer, nn.Conv1d):
+                x = _grouped_conv1d_btc(x, layer)
+                # Re-zero padding after each convolution so it cannot bleed into
+                # valid frames through the kernel window.
+                if keep is not None:
+                    x = x.masked_fill(~keep, 0.0)
+            else:
+                x = layer(x)
+        return x
 
 
 class AudioEmbedding(nn.Module):
@@ -205,9 +234,12 @@ class AdaLayerNorm(nn.Module):
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
 
     def forward(
-        self, x: torch.Tensor, emb: torch.Tensor
+        self, x: torch.Tensor, emb: torch.Tensor | None, mod: torch.Tensor | None = None
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
-        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.linear(self.silu(emb)).chunk(6, dim=1)
+        """``mod`` is the precomputed ``linear(silu(emb))``; ``emb`` is then unused."""
+        if mod is None:
+            mod = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
         x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
         return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
 
@@ -221,8 +253,11 @@ class AdaLayerNormFinal(nn.Module):
         self.linear = nn.Linear(dim, dim * 2)
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
 
-    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
-        scale, shift = self.linear(self.silu(emb)).chunk(2, dim=1)
+    def forward(self, x: torch.Tensor, emb: torch.Tensor | None, mod: torch.Tensor | None = None) -> torch.Tensor:
+        """``mod`` is the precomputed ``linear(silu(emb))``; ``emb`` is then unused."""
+        if mod is None:
+            mod = self.linear(self.silu(emb))
+        scale, shift = mod.chunk(2, dim=1)
         return self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
 
 
@@ -383,9 +418,11 @@ class DoubleBlock(nn.Module):
         c_rope: RopeCache,
         c_mask: torch.Tensor | None,
         bias: torch.Tensor | None = None,
+        mod_c: torch.Tensor | None = None,
+        mod_x: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t)
-        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t)
+        norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, t, mod_c)
+        norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, t, mod_x)
 
         x_attn, c_attn = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope, c_mask=c_mask, bias=bias)
 
@@ -416,8 +453,9 @@ class SingleBlock(nn.Module):
         mask: torch.Tensor | None,
         rope: RopeCache,
         bias: torch.Tensor | None = None,
+        mod: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t)
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(x, t, mod)
         x = x + gate_msa[:, None] * self.attn(x=norm, mask=mask, rope=rope, bias=bias)
         norm = self.ff_norm(x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
         return x + gate_mlp[:, None] * self.ff(norm)
@@ -491,6 +529,28 @@ class AuKTransformer(nn.Module):
 
         self.norm_out = AdaLayerNormFinal(dim)
         self.proj_out = nn.Linear(dim, latent_dim)
+        # Widths of the adaLN modulations in the order modulation_table() lays them out.
+        self._modulation_widths = [module.linear.out_features for module in self._adaln_modules()]
+
+    def _adaln_modules(self) -> list[nn.Module]:
+        """Every timestep-conditioned norm, in the order :meth:`step` consumes them."""
+        modules: list[nn.Module] = []
+        for block in self.transformer_blocks:
+            modules += [block.attn_norm_c, block.attn_norm_x]
+        modules += [block.attn_norm for block in self.single_transformer_blocks]
+        modules.append(self.norm_out)
+        return modules
+
+    def modulation_table(self, timesteps: torch.Tensor) -> torch.Tensor:
+        """adaLN modulations for every timestep, ``[len(timesteps), sum of widths]``.
+
+        The modulations depend only on the timestep, and a request's time grid
+        is known before its first step. Computing them for the whole grid at
+        once reads each modulation weight once per request instead of once per
+        step, which at batch 1 is what those matrix-vector products cost.
+        """
+        emb = F.silu(self.time_embed(timesteps.reshape(-1).float()))
+        return torch.cat([module.linear(emb) for module in self._adaln_modules()], dim=1)
 
     def project_text(self, text: torch.Tensor) -> torch.Tensor:
         """Project LLM hidden states ``[B, nt, text_hidden_dim]`` to model width."""
@@ -520,6 +580,7 @@ class AuKTransformer(nn.Module):
         drop_text: bool = False,
         cfg_infer: bool = False,
         cache: bool = False,
+        timesteps: torch.Tensor | None = None,
     ) -> "AuKStepContext":
         """Everything a denoise step needs that does not depend on ``x`` or the time.
 
@@ -528,6 +589,9 @@ class AuKTransformer(nn.Module):
         prompt embedding, the padding masks and biases and the rotary tables
         are not recomputed on every step. Arguments are those of
         :meth:`forward`; ``target_len`` is the target frame count ``n``.
+        ``timesteps`` are the times the steps will run at (the grid without
+        its end point); when given, their adaLN modulations are precomputed
+        and :meth:`step` selects them by ``step_index``.
         """
         batch = text.shape[0]
         if c_mask is None:
@@ -593,25 +657,44 @@ class AuKTransformer(nn.Module):
             rope_text=_rope_cos_sin(self.rotary_embed(text_len, c_mask)),
             rope_single=_rope_cos_sin(self.rotary_embed(text_len + audio_len, single_mask)),
             branches=branches,
+            modulation=None if timesteps is None else self.modulation_table(timesteps),
         )
 
-    def step(self, x: torch.Tensor, time: torch.Tensor, ctx: "AuKStepContext") -> torch.Tensor:
+    def step(
+        self,
+        x: torch.Tensor,
+        time: torch.Tensor,
+        ctx: "AuKStepContext",
+        step_index: torch.Tensor | int | None = None,
+    ) -> torch.Tensor:
         """Predict the velocity of ``x`` at ``time`` under a :meth:`prepare` context.
 
         The target embedding does not depend on the CFG branch, so it is
-        computed once and shared by both branches.
+        computed once and shared by both branches. When the context carries a
+        modulation table, ``step_index`` selects this step's row of it and
+        ``time`` is unused; the row broadcasts over the CFG branches.
         """
-        if time.ndim == 0:
-            time = time.repeat(x.shape[0])
-        t = self.time_embed(time)
         target = self.audio_embed(x, mask=ctx.target_mask)
+        if ctx.modulation is not None and step_index is not None:
+            if isinstance(step_index, int):
+                row = ctx.modulation[step_index : step_index + 1]
+            else:
+                row = ctx.modulation.index_select(0, step_index.reshape(1))
+            mods = list(torch.split(row, self._modulation_widths, dim=1))
+            t = None
+        else:
+            if time.ndim == 0:
+                time = time.repeat(x.shape[0])
+            t = self.time_embed(time)
+            if ctx.branches == 2:
+                t = torch.cat((t, t), dim=0)
+            mods = [None] * len(self._modulation_widths)
         if ctx.branches == 2:
-            t = torch.cat((t, t), dim=0)
             target = torch.cat((target, target), dim=0)
         audio = target if ctx.prompt is None else torch.cat([ctx.prompt, target], dim=1)
         c = ctx.c
 
-        for block in self.transformer_blocks:
+        for i, block in enumerate(self.transformer_blocks):
             c, audio = block(
                 audio,
                 c,
@@ -621,14 +704,24 @@ class AuKTransformer(nn.Module):
                 c_rope=ctx.rope_text,
                 c_mask=ctx.c_mask,
                 bias=ctx.joint_bias,
+                mod_c=mods[2 * i],
+                mod_x=mods[2 * i + 1],
             )
 
         h = torch.cat([c, audio], dim=1)
-        for block in self.single_transformer_blocks:
-            h = block(h, t, mask=ctx.single_mask, rope=ctx.rope_single, bias=ctx.single_bias)
+        offset = 2 * len(self.transformer_blocks)
+        for i, block in enumerate(self.single_transformer_blocks):
+            h = block(
+                h,
+                t,
+                mask=ctx.single_mask,
+                rope=ctx.rope_single,
+                bias=ctx.single_bias,
+                mod=mods[offset + i],
+            )
 
         h = h[:, h.shape[1] - x.shape[1] :]
-        return self.proj_out(self.norm_out(h, t))
+        return self.proj_out(self.norm_out(h, t, mods[-1]))
 
     def forward(
         self,
@@ -706,6 +799,8 @@ class AuKStepContext:
     rope_text: RopeCache
     rope_single: RopeCache
     branches: int
+    #: adaLN modulations per step, ``[steps, widths]``, when the time grid was known.
+    modulation: torch.Tensor | None = None
 
     def tensors(self) -> list[torch.Tensor | None]:
         """Every tensor field in a fixed order, rope pairs flattened."""
@@ -721,6 +816,7 @@ class AuKStepContext:
             *self.rope_audio,
             *self.rope_text,
             *self.rope_single,
+            self.modulation,
         ]
 
     def copy_(self, other: "AuKStepContext") -> None:
@@ -752,6 +848,7 @@ class AuKStepContext:
             rope_text=(self.rope_text[0].clone(), self.rope_text[1].clone()),
             rope_single=(self.rope_single[0].clone(), self.rope_single[1].clone()),
             branches=self.branches,
+            modulation=_clone(self.modulation),
         )
 
 
@@ -881,6 +978,8 @@ def sample_latents(
                     timestep=t[i],
                     cfg_strength=cfg_strength,
                     new_request=i == 0,
+                    timesteps=t[:-1],
+                    step_index=i,
                 )
                 x = x + (t[i + 1] - t[i]) * velocity
         finally:
@@ -922,9 +1021,10 @@ def _sample_latents(
             ref_mask=ref_mask,
             cfg_infer=guided,
             cache=guided,
+            timesteps=timesteps[:-1],
         )
         for i in range(timesteps.shape[0] - 1):
-            v = dit.step(x, timesteps[i], ctx)
+            v = dit.step(x, timesteps[i], ctx, step_index=i)
             if guided:
                 v_cond, v_uncond = v.chunk(2, dim=0)
                 v = v_cond + (v_cond - v_uncond) * cfg_strength

vllm_omni/diffusion/models/auk/auk_vae.py

@@ -32,7 +32,9 @@ from torch import nn
 from torch.nn.utils import remove_weight_norm as _fold_weight_norm
 from torch.nn.utils import weight_norm as _apply_weight_norm
 from vllm.logger import init_logger
+from vllm.triton_utils import HAS_TRITON
 
+from vllm_omni.diffusion.models.auk.alias_free_triton import alias_free_snake
 from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta as _SharedSnakeBeta
 
 logger = init_logger(__name__)
@@ -319,8 +321,45 @@ class AliasFreeActivation(nn.Module):
         self.downsample = Downsample(down_ratio, down_kernel_size, causal=causal)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self._fused_eligible(x):
+            return self._fused_forward(x)
         return self.downsample(self.act(self.upsample(x)))
 
+    def _fused_eligible(self, x: torch.Tensor) -> bool:
+        """Whether ``x`` can take the Triton kernels in alias_free_triton (CUDA fp32, AuK's filter geometry)."""
+        up = self.upsample
+        taps = up.filter.numel()
+        return (
+            x.is_cuda
+            and x.dtype == torch.float32
+            and x.dim() == 3
+            and isinstance(self.act, SnakeBeta)
+            and HAS_TRITON
+            # The kernel indexes the crop from the start of the transposed
+            # conv output, so the crop must cover the filter's reach, and the
+            # cropped length must be exactly ratio * T.
+            and up.pad_left >= taps - 1
+            and (2 * up.pad - 1) * up.ratio + taps == up.pad_left + up.pad_right
+        )
+
+    def _fused_forward(self, x: torch.Tensor) -> torch.Tensor:
+        up, low, act = self.upsample, self.downsample.lowpass, self.act
+        if not act._cached:
+            act.precompute_exp_cache()
+        return alias_free_snake(
+            x,
+            up.filter.reshape(-1),
+            low.filter.reshape(-1),
+            act._exp_alpha,
+            act._inv_beta,
+            up.ratio,
+            up.pad,
+            up.pad_left,
+            low.stride,
+            low.pad_left,
+            low.pad_right,
+        )
+
 
 class AmpBlock(nn.Module):
     """Anti-aliased multi-periodicity residual block: alternating alias-free snake and dilated convs."""

vllm_omni/diffusion/models/auk/cudagraph_wrapper.py

@@ -24,16 +24,19 @@ class _GraphEntry:
     static_ctx: AuKStepContext
     static_timestep: torch.Tensor
     static_cfg: torch.Tensor | None
+    static_step: torch.Tensor | None
     static_out: torch.Tensor
 
 
 class AuKCUDAGraphWrapper:
     """Replay one DiT denoise step and leave Euler scheduling to the caller.
 
-    The graph is keyed by the target, text and reference sequence lengths plus
-    whether the CFG branch is enabled. Timestep and the CFG strength are
+    The graph is keyed by the target, text and reference sequence lengths,
+    whether the CFG branch is enabled and, when the caller passes the time
+    grid, its length. Timestep (or the step index), and the CFG strength are
     mutable scalar buffers, so all Euler steps and CFG values within one path
-    reuse one graph.
+    reuse one graph. With the time grid, the adaLN modulations of every step
+    are part of the per-request context and the graph selects its row.
 
     Only :meth:`AuKTransformer.step` is captured. The per-request context
     (text projection, reference embedding, padding biases, rotary tables) is
@@ -85,8 +88,9 @@ class AuKCUDAGraphWrapper:
         text: torch.Tensor,
         ref: torch.Tensor,
         uses_cfg: bool,
-    ) -> tuple[int, int, int, bool]:
-        return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg)
+        steps: int = 0,
+    ) -> tuple[int, int, int, bool, int]:
+        return (x.shape[1], text.shape[1], ref.shape[1], uses_cfg, steps)
 
     def _retire_graph_generation_if_full(self) -> None:
         """Retire all graphs together so none outlive shared workspaces."""
@@ -103,6 +107,7 @@ class AuKCUDAGraphWrapper:
         ref: torch.Tensor,
         ref_mask: torch.Tensor,
         uses_cfg: bool,
+        timesteps: torch.Tensor | None = None,
     ) -> AuKStepContext:
         return self.dit.prepare(
             text,
@@ -112,6 +117,7 @@ class AuKCUDAGraphWrapper:
             ref=ref,
             ref_mask=ref_mask,
             cfg_infer=uses_cfg,
+            timesteps=timesteps,
         )
 
     def _step(
@@ -120,9 +126,10 @@ class AuKCUDAGraphWrapper:
         timestep: torch.Tensor,
         ctx: AuKStepContext,
         cfg_strength: torch.Tensor | float,
+        step_index: torch.Tensor | int | None = None,
     ) -> torch.Tensor:
         """One guided (or unguided) velocity; the CFG combine is part of the graph."""
-        velocity = self.dit.step(x, timestep, ctx)
+        velocity = self.dit.step(x, timestep, ctx, step_index=step_index)
         if ctx.branches == 1:
             return velocity
         conditional, unconditional = velocity.chunk(2, dim=0)
@@ -140,6 +147,8 @@ class AuKCUDAGraphWrapper:
         timestep: torch.Tensor,
         cfg_strength: float,
         new_request: bool = True,
+        timesteps: torch.Tensor | None = None,
+        step_index: int | None = None,
     ) -> torch.Tensor:
         """Velocity of ``x`` at ``timestep``.
 
@@ -147,26 +156,33 @@ class AuKCUDAGraphWrapper:
         (``text``, ``c_mask``, ``ref``, ``ref_mask``) is then prepared anew; the
         later steps of the request pass ``False`` and reuse it. The default
         prepares on every call, which is always correct.
+
+        ``timesteps`` (the times every step of the request runs at) and
+        ``step_index`` (this step's position in them) let the adaLN
+        modulations be computed once per request instead of once per step.
         """
         uses_cfg = cfg_strength >= 1e-5
+        if timesteps is None or step_index is None:
+            timesteps, step_index = None, None
         if not self.enabled or x.device.type != "cuda" or torch.cuda.is_current_stream_capturing():
             if new_request or self._eager_ctx is None:
-                self._eager_ctx = self._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg)
-            return self._step(x, timestep, self._eager_ctx, cfg_strength)
+                self._eager_ctx = self._prepare(x, None, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
+            return self._step(x, timestep, self._eager_ctx, cfg_strength, step_index)
 
         target_frames = x.shape[1]
         x, x_mask, text, c_mask, ref, ref_mask = self._bucket_inputs(x, text, c_mask, ref, ref_mask)
-        key = self._key(x, text, ref, uses_cfg)
+        steps = 0 if timesteps is None else int(timesteps.numel())
+        key = self._key(x, text, ref, uses_cfg, steps)
         entry = self._cache.get(key)
         ctx = None
         if entry is None or new_request or key != self._active_key:
-            ctx = self._prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg)
+            ctx = self._prepare(x, x_mask, text, c_mask, ref, ref_mask, uses_cfg, timesteps)
         if entry is None:
             # Same Retirement as #6587, details discussed in #7469
             self._retire_graph_generation_if_full()
             try:
                 assert ctx is not None
-                entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength)
+                entry = self._capture(x, ctx, timestep, cfg_strength=cfg_strength, step_index=step_index)
             except Exception:
                 self._cache.clear()
                 self._active_key = None
@@ -186,6 +202,8 @@ class AuKCUDAGraphWrapper:
 
         entry.static_x.copy_(x)
         entry.static_timestep.copy_(timestep)
+        if entry.static_step is not None:
+            entry.static_step.fill_(step_index)
         if entry.static_cfg is not None:
             entry.static_cfg.fill_(cfg_strength)
         entry.graph.replay()
@@ -198,10 +216,14 @@ class AuKCUDAGraphWrapper:
         timestep: torch.Tensor,
         *,
         cfg_strength: float,
+        step_index: int | None = None,
     ) -> _GraphEntry:
         static_x = x.clone()
         static_ctx = ctx.clone()
         static_timestep = timestep.clone()
+        static_step = None
+        if static_ctx.modulation is not None and step_index is not None:
+            static_step = torch.full((), step_index, device=static_x.device, dtype=torch.long)
         static_cfg = None
         if ctx.branches == 2:
             static_cfg = torch.empty((), device=static_x.device, dtype=torch.float32)
@@ -209,13 +231,13 @@ class AuKCUDAGraphWrapper:
         scale = static_cfg if static_cfg is not None else cfg_strength
         # Warm-up runs any lazy torch.compile outside the capture.
         for _ in range(3):
-            self._step(static_x, static_timestep, static_ctx, scale)
+            self._step(static_x, static_timestep, static_ctx, scale, static_step)
         if self._pool_handle is None:
             self._pool_handle = torch.cuda.graph_pool_handle()
         graph = torch.cuda.CUDAGraph()
         with torch.cuda.graph(graph, pool=self._pool_handle):
             # The CFG scalar stays a mutable buffer, so cfg=2 and cfg=3 share this graph.
-            static_out = self._step(static_x, static_timestep, static_ctx, scale)
+            static_out = self._step(static_x, static_timestep, static_ctx, scale, static_step)
 
         logger.info(
             "Captured AuK DiT single-step CUDA graph: target_frames=%d text_tokens=%d ref_frames=%d cfg=%s",
@@ -230,6 +252,7 @@ class AuKCUDAGraphWrapper:
             static_ctx=static_ctx,
             static_timestep=static_timestep,
             static_cfg=static_cfg,
+            static_step=static_step,
             static_out=static_out,
         )
 

vllm_omni/diffusion/models/auk/pipeline_auk.py

@@ -11,9 +11,11 @@ both, and the BigVGAN-flow decoder renders the waveform.
 
 from __future__ import annotations
 
+import hashlib
 import json
 import os
 import time
+from collections import OrderedDict
 from collections.abc import Iterable
 from typing import Any, ClassVar
 
@@ -27,7 +29,12 @@ from vllm.logger import init_logger
 from vllm_omni.diffusion.compile import regionally_compile
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
-from vllm_omni.diffusion.models.auk.auk_transformer import AuKTransformer, dit_state_dict, sample_latents
+from vllm_omni.diffusion.models.auk.auk_transformer import (
+    AuKTransformer,
+    build_time_grid,
+    dit_state_dict,
+    sample_latents,
+)
 from vllm_omni.diffusion.models.auk.auk_vae import AuKVAE
 from vllm_omni.diffusion.models.auk.cudagraph_wrapper import AuKCUDAGraphWrapper
 from vllm_omni.diffusion.models.auk.vae_cudagraph import AuKVAEDecodeGraph
@@ -59,6 +66,10 @@ _FLASH_CFG = 0.0
 _WARMUP_TARGET_FRAMES = (150, 300, 600)
 _WARMUP_REF_FRAMES = (0, 150)
 _WARMUP_TEXT_TOKENS = 96
+# Reference clips whose posterior-mean latents are kept. Voice cloning reuses a
+# small set of reference voices, so repeat requests skip the VAE encode.
+# ``model_config.auk_ref_cache_size`` overrides it; 0 disables the cache.
+_REF_CACHE_SIZE = 16
 # Decorrelates the VAE posterior draw from the initial latent for the same request seed.
 _VAE_SEED_OFFSET = 0x5EED_0A0C
 
@@ -152,6 +163,9 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         max_dit_graphs = od_config.model_config.get("max_dit_graphs", 32)
         if isinstance(max_dit_graphs, bool) or not isinstance(max_dit_graphs, int) or max_dit_graphs < 1:
             raise ValueError("AuK max_dit_graphs must be a positive integer")
+        ref_cache_size = od_config.model_config.get("auk_ref_cache_size", _REF_CACHE_SIZE)
+        if isinstance(ref_cache_size, bool) or not isinstance(ref_cache_size, int) or ref_cache_size < 0:
+            raise ValueError("AuK auk_ref_cache_size must be a non-negative integer (0 disables the cache)")
         self.device = get_local_device()
         self.dtype = getattr(od_config, "dtype", None) or torch.bfloat16
 
@@ -213,6 +227,8 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         if model_config.get("auk_vae_tile_frames") is not None:
             vae_decode_kwargs["tile_frames"] = int(model_config["auk_vae_tile_frames"])
         self.vae_decode = AuKVAEDecodeGraph(self.vae, enabled=not od_config.enforce_eager, **vae_decode_kwargs)
+        self._ref_cache_size = ref_cache_size
+        self._ref_cache: OrderedDict[str, torch.Tensor] = OrderedDict()
 
         logger.info(
             "AuK pipeline ready: variant=%s dtype=%s latent_dim=%d hop=%d sample_rate=%d",
@@ -267,6 +283,12 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
         text = torch.zeros(1, _WARMUP_TEXT_TOKENS, self.text_hidden_dim, device=self.device, dtype=self.dtype)
         c_mask = torch.ones(text.shape[:2], dtype=torch.bool, device=self.device)
         timestep = torch.zeros((), device=self.device, dtype=torch.float32)
+        # The graphs are keyed by the number of steps, so warm the variant's default grid.
+        if self.is_flash:
+            grid = torch.tensor(self.flash_t_grid, device=self.device, dtype=torch.float32)
+        else:
+            sway = None if self.default_sway is None else float(self.default_sway)
+            grid = build_time_grid(nfe=self.default_nfe, sway_sampling_coef=sway, t_grid=None, device=self.device)
         start = time.perf_counter()
         with torch.inference_mode(), self._dit_autocast():
             for ref_frames in _WARMUP_REF_FRAMES:
@@ -283,6 +305,8 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
                         timestep=timestep,
                         cfg_strength=cfg,
                         new_request=True,
+                        timesteps=grid[:-1],
+                        step_index=0,
                     )
         if self.device.type != "cpu":
             torch.accelerator.synchronize(self.device)
@@ -386,8 +410,24 @@ class AuKPipeline(nn.Module, SupportAudioInput, SupportAudioOutput, SupportsComp
 
         if audio is None:
             return torch.zeros(1, 0, self.latent_dim, device=self.device, dtype=torch.float32)
-        wav = self._prepare_waveform(_unwrap_single(audio)).to(self.device)
-        return self.vae.encode(wav, sample=sample, generator=generator)
+        wav = self._prepare_waveform(_unwrap_single(audio))
+        # The posterior mean is a pure function of the prepared waveform, so it
+        # is cached by content. A posterior draw depends on the generator and
+        # is always recomputed. Callers must not modify the returned tensor.
+        key = None
+        if not sample and self._ref_cache_size:
+            digest = hashlib.blake2b(wav.numpy().tobytes(), digest_size=16).hexdigest()
+            key = f"{wav.shape[-1]}:{digest}"
+            cached = self._ref_cache.get(key)
+            if cached is not None:
+                self._ref_cache.move_to_end(key)
+                return cached
+        latents = self.vae.encode(wav.to(self.device), sample=sample, generator=generator)
+        if key is not None:
+            self._ref_cache[key] = latents
+            while len(self._ref_cache) > self._ref_cache_size:
+                self._ref_cache.popitem(last=False)
+        return latents
 
     def _target_frames(self, gen_seconds: Any, ref_frames: int) -> int:
         """Target latent length, capped by the remaining context."""

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -38,6 +38,7 @@ from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.sp_plan import SequenceParallelInput, SequenceParallelOutput
 from vllm_omni.diffusion.forward_context import get_forward_context, is_forward_context_available
 from vllm_omni.diffusion.layers.norm import RMSNorm as _VllmRMSNorm
+from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
 from vllm_omni.diffusion.offloader.config import (
     OffloadStrategy,
     resolve_offload_strategy,
@@ -1084,6 +1085,7 @@ class Cosmos3LanguageModel(nn.Module):
         rope_theta: float,
         mrope_section: list[int],
         quant_config: QuantizationConfig | None = None,
+        release_completed_blocks_to_meta: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -1093,21 +1095,21 @@ class Cosmos3LanguageModel(nn.Module):
             rope_theta=rope_theta,
             mrope_section=mrope_section,
         )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3UndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
+        self.layers = nn.ModuleList()
+        for i in range(num_hidden_layers):
+            layer = Cosmos3UndDecoderLayer(
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=head_dim,
+                rms_norm_eps=rms_norm_eps,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.layers.append(layer)
         # TODO: Not used right now, will be used in the future for prompt upsampler.
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
@@ -1228,6 +1230,10 @@ class Cosmos3VFMTransformer(nn.Module):
 
     _hsdp_shard_conditions = [_is_transformer_block]
 
+    # Standard unquantized vLLM linears need no value-dependent post-load
+    # processing on accelerators. Edge inherits the same loading contract.
+    _hsdp_pre_sharded_meta_post_load = True
+
     # Modules whose parameters must NOT be FSDP-sharded at the root level.
     # time_embedder is cast to fp32 by post_load_weights for precision; if it
     # were swept into the root flat-parameter under MixedPrecisionPolicy(param_dtype=bf16),
@@ -1346,6 +1352,10 @@ class Cosmos3VFMTransformer(nn.Module):
         language_model_quant_config, gen_layers_quant_config = _resolve_cosmos3_quant_configs(
             getattr(od_config, "quantization_config", None)
         )
+        release_completed_blocks_to_meta = bool(
+            getattr(getattr(od_config, "parallel_config", None), "use_hsdp", False)
+            and getattr(od_config, "hsdp_weight_load_strategy", "full") == "pre_sharded"
+        )
         mixed_precision_config, mixed_precision_source = resolve_mixed_precision_config(od_config)
         if mixed_precision_config is None:
             if mixed_precision_source == "additional_config_disabled":
@@ -1372,6 +1382,7 @@ class Cosmos3VFMTransformer(nn.Module):
             rope_theta=self.rope_theta,
             mrope_section=self.mrope_section,
             quant_config=language_model_quant_config,
+            release_completed_blocks_to_meta=release_completed_blocks_to_meta,
             prefix="language_model",
             **self._language_model_kwargs(),
         )
@@ -1399,24 +1410,24 @@ class Cosmos3VFMTransformer(nn.Module):
             self.audio_proj_out = nn.Linear(self.hidden_size, self.sound_dim)
             self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size))
 
-        self.gen_layers = nn.ModuleList(
-            [
-                Cosmos3GenDecoderLayer(
-                    layer_idx=i,
-                    hidden_size=self.hidden_size,
-                    intermediate_size=self.intermediate_size,
-                    num_attention_heads=self.num_attention_heads,
-                    num_key_value_heads=self.num_key_value_heads,
-                    head_dim=self.head_dim,
-                    rms_norm_eps=self.rms_norm_eps,
-                    quant_config=gen_layers_quant_config,
-                    mlp_cls=self._gen_mlp_cls,
-                    qk_norm=self.qk_norm_for_diffusion,
-                    prefix=f"gen_layers.{i}",
-                )
-                for i in range(self.num_hidden_layers)
-            ]
-        )
+        self.gen_layers = nn.ModuleList()
+        for i in range(self.num_hidden_layers):
+            layer = Cosmos3GenDecoderLayer(
+                layer_idx=i,
+                hidden_size=self.hidden_size,
+                intermediate_size=self.intermediate_size,
+                num_attention_heads=self.num_attention_heads,
+                num_key_value_heads=self.num_key_value_heads,
+                head_dim=self.head_dim,
+                rms_norm_eps=self.rms_norm_eps,
+                quant_config=gen_layers_quant_config,
+                mlp_cls=self._gen_mlp_cls,
+                qk_norm=self.qk_norm_for_diffusion,
+                prefix=f"gen_layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.gen_layers.append(layer)
 
         self.mixed_precision_runtime: Cosmos3MixedPrecisionRuntime | None = None
         if mixed_precision_config is not None:

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3_edge.py

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelL
 from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
 from vllm_omni.diffusion.attention.layer import Attention as FrameworkAttention
+from vllm_omni.diffusion.models.utils import release_module_parameters_to_meta
 
 from .transformer_cosmos3 import (
     Cosmos3VFMTransformer,
@@ -258,6 +259,7 @@ class Cosmos3EdgeLanguageModel(nn.Module):
         mrope_section: list[int],
         use_und_k_norm_for_gen: bool,
         quant_config: QuantizationConfig | None = None,
+        release_completed_blocks_to_meta: bool = False,
         prefix: str = "",
     ) -> None:
         super().__init__()
@@ -267,22 +269,22 @@ class Cosmos3EdgeLanguageModel(nn.Module):
             rope_theta=rope_theta,
             mrope_section=mrope_section,
         )
-        self.layers = nn.ModuleList(
-            [
-                Cosmos3EdgeUndDecoderLayer(
-                    hidden_size=hidden_size,
-                    intermediate_size=intermediate_size,
-                    num_attention_heads=num_attention_heads,
-                    num_key_value_heads=num_key_value_heads,
-                    head_dim=head_dim,
-                    rms_norm_eps=rms_norm_eps,
-                    use_und_k_norm_for_gen=use_und_k_norm_for_gen,
-                    quant_config=quant_config,
-                    prefix=f"{prefix}.layers.{i}",
-                )
-                for i in range(num_hidden_layers)
-            ]
-        )
+        self.layers = nn.ModuleList()
+        for i in range(num_hidden_layers):
+            layer = Cosmos3EdgeUndDecoderLayer(
+                hidden_size=hidden_size,
+                intermediate_size=intermediate_size,
+                num_attention_heads=num_attention_heads,
+                num_key_value_heads=num_key_value_heads,
+                head_dim=head_dim,
+                rms_norm_eps=rms_norm_eps,
+                use_und_k_norm_for_gen=use_und_k_norm_for_gen,
+                quant_config=quant_config,
+                prefix=f"{prefix}.layers.{i}",
+            )
+            if release_completed_blocks_to_meta:
+                release_module_parameters_to_meta(layer)
+            self.layers.append(layer)
         self.norm = RMSNorm(hidden_size, eps=rms_norm_eps)
 
     def forward(

vllm_omni/diffusion/models/lingbot_video/pipeline_lingbot_video.py

@@ -288,13 +288,13 @@ def get_lingbot_video_post_process_func(od_config: OmniDiffusionConfig):
             output_key = output_keys[0]
             frames = frames[output_key]
         output_type = getattr(sampling_params, "output_type", None) or "pt"
-        # The image serving path currently accepts PIL images or NumPy arrays,
-        # while LingBot decodes T2I outputs to tensors. Keep that compatibility
-        # conversion without discarding the model's image/video payload key.
+        # Decoded LingBot pixels are already in [0, 1]. The video API treats
+        # floating tensors as [-1, 1], so use its NumPy path to avoid normalizing
+        # twice. Images also need NumPy for serving; keep latent tensors intact.
         if (
             isinstance(frames, torch.Tensor)
             and output_type != "latent"
-            and (output_type == "np" or output_key == "image")
+            and (output_type == "np" or output_key in {"image", "video"})
         ):
             frames = frames.float().cpu().numpy()
         return {output_key: frames} if output_key is not None else frames

vllm_omni/diffusion/models/utils.py

@@ -168,6 +168,28 @@ def init_parameters(
         init_parameters(child, dtype, device)
 
 
+def release_module_parameters_to_meta(module: nn.Module) -> None:
+    """Release a fully constructed module's parameter storage to meta.
+
+    The module must finish construction first so initialization code and vLLM
+    weight-loader metadata follow the ordinary CPU path. ``swap_tensors`` keeps
+    each registered Parameter object's identity while dropping only its storage.
+    Persistent buffers intentionally remain materialized.
+    """
+    for param in module.parameters():
+        if param.device.type == "meta":
+            continue
+
+        meta_tensor = torch.empty_like(param, device="meta")
+        if type(param) is nn.Parameter:
+            meta_param = nn.Parameter(meta_tensor, requires_grad=param.requires_grad)
+        else:
+            meta_param = param.__class__.__new__(param.__class__, meta_tensor)
+            meta_param.requires_grad_(param.requires_grad)
+        meta_param.__dict__.update(getattr(param, "__dict__", {}))
+        torch.utils.swap_tensors(param, meta_param)
+
+
 def init_parameters_preserving_vllm_linear(
     module: nn.Module,
     dtype: torch.dtype | None,