Production Add Diffusion Model

Productionizing a Diffusion Model

Scope

Use this skill to move a working model integration from Day-0 support to a measured, fail-closed production deployment. It is deliberately separate from add-diffusion-model; do not modify or copy that skill when productionizing a model. If registry, pipeline, or basic loader work is incomplete, read that skill first and finish the Day-0 vertical slice.

MiniMax-H3 PR #5691 is the merged Day-0 case study. Its open optimization issue #5700 is an actively maintained roadmap and evidence map, not a list of already supported features. It can lag merged implementation: read its current revision, cited PRs, target source, and maintained recipes, then leave every unverified row not tested until the target revision has scoped evidence.

Read these references before the corresponding gate:

• API and hardware recipes
• Production feature patterns
• Performance patterns and code examples
• Serving, soak, and CI validation

Non-negotiable evidence contract

Create a matrix whose row key is at least:

task x API mode x execution mode/capacity x shape/schedule x attention backend x
packed layout x cache policy x quantization x offload mode x topology x
hardware x dtype x output representation/transport

Use exactly these states:

Never infer compatibility from import success, server startup, another task, another card, or another topology. Record the model/checkpoint revision, vLLM-Omni commit, commands, request assets and hashes, seed/schedule/shape, raw output artifacts, environment, and result for every validated row.

Workflow

Follow the gates in order. Keep a dense BF16 single-device path as the correctness oracle until every advertised fast path has parity evidence.

Gate 0: Freeze the official contract

Pin the official implementation and checkpoint revision. Build a matrix for:

• constructor/call arguments, defaults, scheduler and sigma schedule;
• task-to-source modality/count/order rules;
• prompt, MIME, bytes, pixels, frames, duration, FPS, steps, seed, guidance, output count, and output format limits;
• offline API and every public sync/async serving endpoint;
• checkpoint partitions and task selection.

Normalize offline, sync, and async requests through one validation contract. Reject an invalid task/source combination before downloads, upload persistence, temporary files, async job creation, or engine submission. Use 400 for invalid semantics and 413 for payload limits; clean partial resources on every failure. See API and hardware recipes.

Gate 1: Re-prove correctness under production shapes

1. Make loading strict: track expected, loaded, unexpected, and intentionally ignored tensors. A fused destination is complete only after every source shard arrives. Missing retained parameters abort startup.
2. Fix reference revision, asset hashes, seed, schedule, and shape. Compare component outputs, intermediate denoise latents, and final artifacts.
3. Test unequal packed samples. Verify cu_seqlens, padding exclusion, modality/CFG ownership, split/gather boundaries, and one output per request.
4. Test partial process groups with non-member ranks and group size smaller than world size. Do not construct a private group per component or request.
5. Run negative API cases before positive E2E tests to prove side-effect-free rejection and cleanup.

Gate 2: Wire shared fast primitives

Prefer explicit, testable vLLM-Omni/vLLM operators with native fallbacks. torch.compile remains a separately validated regional optimization; it is not a substitute for correct operator selection or lifecycle-safe boundaries.

RMSNorm + RoPE baseline

Use the shared cross-platform operators, preserving the official Q/K norm and RoPE order:

from vllm_omni.diffusion.layers.norm import RMSNorm
from vllm_omni.diffusion.layers.rope import (
    RotaryEmbedding,
    apply_rope_to_qk,
)

self.norm_q = RMSNorm(head_dim, eps=eps)
self.norm_k = RMSNorm(head_dim, eps=eps)
self.rope = RotaryEmbedding(is_neox_style=False)

query = self.norm_q(query)  # [B, S, Hq, D] or packed [T, Hq, D]
key = self.norm_k(key)      # [B, S, Hkv, D] or packed [T, Hkv, D]
query, key = apply_rope_to_qk(
    self.rope,
    query,
    key,
    (cos, sin),
)

Verify is_neox_style, half_head_dim, partial rotary dimensions, cos/sin layout, and packed row ownership against the official model. False means interleaved/GPT-J style; do not assume the default is NeoX. Shared RMSNorm and RoPE are two dispatched operators, not automatically one fused kernel.

For packed non-interleaved RoPE, use the shared fused boundary when its contract matches the model:

from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope

query, key = fused_qk_norm_rope(
    query,                    # [T, Hq, D]
    key,                      # [T, Hkv, D]
    self.norm_q.weight,       # [D]
    self.norm_k.weight,       # [D]
    rope_table,               # [T, rotary_dim] = [cos | sin]
    self.norm_q.variance_epsilon,
)

The current CUDA fast path specializes BF16 head_dim=128, rotary_dim=96; the public function keeps an eager fallback for unsupported inputs. Verify the packed row/frequency contract and trace the actual fast path. Do not reshape an incompatible official RoPE layout merely to enter this kernel.

MiniMax-H3-style partial NeoX RoPE rotates a prefix and passes the suffix through:

self.q_norm = RMSNorm(head_dim, eps=eps, dtype=torch.bfloat16)
self.k_norm = RMSNorm(head_dim, eps=eps, dtype=torch.bfloat16)
self.rope = RotaryEmbedding(is_neox_style=True, half_head_dim=False)

def apply_partial_rope(self, x, freqs, rot_dim):
    x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:]
    cos, sin = torch.cos(freqs).to(x.dtype), torch.sin(freqs).to(x.dtype)
    return torch.cat((self.rope(x_rot, cos, sin), x_pass), dim=-1)

The partial-RoPE snippet assumes import torch and lives on the owning attention module; adapt names and validated frequency shapes rather than copying it at module scope.

Call the modules normally so CustomOp selects CUDA/HIP/NPU/XPU/native implementations. Do not call forward_cuda() directly. Compare the shared path with native BF16 at operator, block, denoise-trajectory, and artifact levels before claiming either accuracy or speed.

Faster BF16 attention and fused projections

Use local TP head counts and role-aware shared attention:

from vllm.model_executor.layers.linear import QKVParallelLinear
from vllm_omni.diffusion.attention.layer import Attention

self.qkv_proj = QKVParallelLinear(
    hidden_size=hidden_size,
    head_size=head_dim,
    total_num_heads=num_heads,
    total_num_kv_heads=num_kv_heads,
    bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.qkv_proj",
)
self.attn = Attention(
    num_heads=self.qkv_proj.num_heads,
    num_kv_heads=self.qkv_proj.num_kv_heads,
    head_size=head_dim,
    softmax_scale=head_dim**-0.5,
    causal=False,
    qkv_layout="BSND",
    role="self",
    prefix=prefix,
)

Prove the selected backend from runtime logs/trace and compare it with dense SDPA. Packed metadata must express real sample boundaries; do not build masks that the selected backend ignores. A backend that is faster for one shape or platform is not a global default. See performance patterns for QKV, SwiGLU, AdaLN, sparse attention, redundancy, and fusion examples.

Gate 3: Add online FP8 without breaking the loader

Route the active config to each component, then pass a stable runtime prefix and config into every quantizable vLLM linear:

def resolve_component_quant_config(quant_config, component):
    return quant_config.resolve(component) if hasattr(quant_config, "resolve") else quant_config

transformer_quant_config = resolve_component_quant_config(
    od_config.quantization_config,
    "transformer",
)
self.transformer = MyDiT(
    od_config,
    quant_config=transformer_quant_config,
    prefix="transformer",
)

self.qkv_proj = QKVParallelLinear(
    hidden_size=hidden_size,
    head_size=head_dim,
    total_num_heads=num_heads,
    total_num_kv_heads=num_kv_heads,
    quant_config=quant_config,
    prefix=f"{prefix}.qkv_proj",
)

Use build_quant_config() for direct construction/testing:

from vllm_omni.quantization import build_quant_config

quant_config = build_quant_config({
    "transformer": {"method": "fp8"},
    "text_encoder": None,
    "vae": None,
    "default": None,
})

Preserve each parameter's vLLM weight_loader; perform checkpoint layout conversion before invoking it. Account for every QKV/gate-up source shard. Keep precision-sensitive norm/modulation/embedders in BF16 unless separately proven, using the current safe_quant_config pattern where applicable.

Validate resident DiT FP8 first: prove which modules are quantized, HBM saved, BF16-vs-FP8 trajectory/artifact quality, latency, and throughput. Text encoder, VAE, other hardware, TP/cache/DLO combinations are separate rows. Pre-quantized, pruned, rotated, or adapter-modified checkpoints are separate model/loading lanes; compare them with the released BF16 model without calling the result a pure quantization ablation. Read ../quantization/SKILL.md if present and prefer the target revision's implementation/docs when sibling guidance differs.

Gate 4: Make every component offloadable

Declare topology instead of relying on heuristic discovery:

from typing import ClassVar

import torch.nn as nn

from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.offloader import OffloadPlan

class MyPipeline(nn.Module, SupportsComponentDiscovery):
    _dit_modules: ClassVar[list[str]] = ["transformer"]
    _encoder_modules: ClassVar[list[str]] = ["text_encoder"]
    _vae_modules: ClassVar[list[str]] = ["vae"]
    _resident_modules: ClassVar[list[str]] = []
    _offload_plan: ClassVar[OffloadPlan] = OffloadPlan(
        on_demand_component_paths=frozenset({"text_encoder", "vae"}),
        block_attrs={"transformer": ("blocks",)},
        offload_submodules={"token_refiner": "blocks"},
        resident_dit_paths=frozenset({"transformer"}),
        encoder_block_attrs={"text_encoder": ("encoder.layers",)},
    )

Add a test that resolves every dotted component/block path and fails if it is missing or is not the expected module/container; discovery may otherwise warn and skip. Validate ordinary layerwise offload and distributed layerwise offload (DLO) independently.

Keep component lifecycle and host-weight storage separate. OffloadPlan declares what can be streamed; the diffusion loader owns any HostWeightPlan and hands the exact prevalidated plan to DLO. Do not make the offloader rescan checkpoint files or duplicate loader name/shape/dtype decisions. Direct mmap is a fail-closed optimization: current preflight requires TP1, no HSDP, no online quantization, complete bindings, and represented transforms. Otherwise the ordinary loader remains authoritative.

If the target revision supports Host Weight Runtime (HWR), treat it as a separate, exact-identity startup/cache path rather than generic offload or zero-copy execution. Qualify preferred population/fallback and required consume-only failure independently; verify immutable manifests, final-layout restore, lease cleanup, corruption/quarantine, and node-local sharing. Current model contracts may be BF16 no-AllGather-only, so reject AllGather, online quantization, HSDP, LoRA/adapted weights, and unrecognized load formats unless the exact target revision explicitly adds them.

On small-HBM cards, record load/materialization peak, resident peak, encode, each denoise wave, decode, transient per-rank HBM, host PSS, H2D, and collective time. Test DLO AllGather and --dlo-no-use-allgather as different deployments. Sweep the supported resident-block policies and publish the non-dominated latency-HBM frontier; a single minimum-memory or fastest point is not the DLO trade-off. Keep request-specific preparation replica-local. Only weight materialization may enter the declared DLO collective group; never use WORLD for per-request work in a DP deployment. Direct checkpoint mmap preflight is limited to TP1 without HSDP or online quantization; TP>1 falls back to the ordinary TP-aware loader and can still feed DLO AllGather or no-AllGather after scoped E2E. HSDP remains incompatible with DLO AllGather. A target revision may allow finalized per-tensor online FP8 through the ordinary loader and AllGather; keep every other online quantizer fail-closed, and require model/card/topology E2E before promotion. Use deploy YAML for DP under --omni; vLLM DP CLI flags are rejected.

Gate 5: Add request-scoped acceleration safely

Server-wide Cache-DiT support does not imply a per-request quality API. Only add request policy after defining calibrated quality tiers and a safe hook lifecycle using SupportsRequestScopedCacheDiT, CacheDiTRequestSpec, and RequestScopedCacheDiTRuntime. Validate alternating and concurrent quality tiers, real cache hits, refresh/reset, success, error, disconnect, abort, and the next uncached request. Do not batch different quality policies unless the batch compatibility key and hook ownership make that safe.

Then test Cache-DiT against FP8, DLO, TP/SP, compile, and batching one combination at a time. See production feature patterns for the lifecycle skeleton and compatibility matrix.