[Doc] Add production diffusion model skill (#6097)
Signed-off-by: David Chen <530634352@qq.com>
This commit is contained in:
WeiQing Chen authored and GitHub committed
1 parent
1dae537313
commit
dff28e6867
7 files changed
+2188
No files matched your search
@@ -0,0 +1,458 @@
|
||||
---
|
||||
name: production-add-diffusion-model
|
||||
description: Productionize a vLLM-Omni diffusion model after its Day-0 vertical slice works. Use when work requires official API parity and input limits, feature-combination evidence, online FP8, distributed layerwise offload, per-request Cache-DiT, CUDA/ROCm/NPU/XPU recipes, faster or sparse attention, fused operators, USP or disaggregation analysis, continuous batching and abort handling, long-running RPS stability, or accuracy/performance/reliability CI. For initial architecture porting, registry wiring, and basic weight loading, use add-diffusion-model first.
|
||||
---
|
||||
|
||||
# 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`](../add-diffusion-model/SKILL.md); 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](https://github.com/vllm-project/vllm-omni/pull/5691)
|
||||
is the merged Day-0 case study. Its open optimization
|
||||
[issue #5700](https://github.com/vllm-project/vllm-omni/issues/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](references/api-and-recipes.md)
|
||||
- [Production feature patterns](references/feature-patterns.md)
|
||||
- [Performance patterns and code examples](references/performance-patterns.md)
|
||||
- [Serving, soak, and CI validation](references/production-validation.md)
|
||||
|
||||
## Non-negotiable evidence contract
|
||||
|
||||
Create a matrix whose row key is at least:
|
||||
|
||||
```text
|
||||
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:
|
||||
|
||||
| State | Meaning |
|
||||
|---|---|
|
||||
| `validated` | The exact row passed correctness and measured deployment tests, with reproducible artifacts. |
|
||||
| `limited` | A bounded subset passed; the limitation and rejection/fallback behavior are explicit. |
|
||||
| `unsupported` | A stable architectural/platform restriction is proven and fails early with an actionable error. |
|
||||
| `not tested` | No adequate evidence exists. This is the default. |
|
||||
|
||||
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](references/api-and-recipes.md).
|
||||
|
||||
### 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:
|
||||
|
||||
```python
|
||||
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:
|
||||
|
||||
```python
|
||||
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:
|
||||
|
||||
```python
|
||||
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:
|
||||
|
||||
```python
|
||||
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](references/performance-patterns.md)
|
||||
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:
|
||||
|
||||
```python
|
||||
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:
|
||||
|
||||
```python
|
||||
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`](../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:
|
||||
|
||||
```python
|
||||
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](references/feature-patterns.md)
|
||||
for the lifecycle skeleton and compatibility matrix.
|
||||
|
||||
### Gate 6: Optimize only from profiles
|
||||
|
||||
Freeze one canonical workload manifest before profiling. Keep lossless runtime
|
||||
A/Bs, accelerated paths with numerical/quality trade-offs, and production
|
||||
topology studies in separate result lanes. Reconcile the client boundary as
|
||||
queue + encode + denoise + video/audio decode + output transport/codec +
|
||||
residual; do not hide a material residual in another stage.
|
||||
|
||||
Profile stage time, GPU kernels, CPU/GPU synchronization, allocations, H2D,
|
||||
output payload/copies, and collectives before choosing work. Apply and A/B one
|
||||
change at a time:
|
||||
|
||||
1. Remove synchronization and redundant denoise-loop work: hoist static prompt
|
||||
conditioning/token refinement, RoPE tables, masks/metadata, bounded AdaLN
|
||||
schedule projections, and reference scans; remove per-layer `.item()` and
|
||||
unused masks/allocations.
|
||||
2. Select the fastest correct BF16 attention backend for each role and shape.
|
||||
3. Prefer explicit fused QKV, gate-up + `SiluAndMul`, RMSNorm/RoPE, AdaLN, and
|
||||
layout/residual operators with guarded native fallbacks. Preserve every
|
||||
materialized dtype/rounding boundary consumed by the unfused graph.
|
||||
4. Enable sparse attention only when platform, shape, topology, realized
|
||||
sparsity, quality, fallback, and end-to-end speedup are all proven.
|
||||
5. Profile USP all-to-all count/bytes, packing, contiguous copies, overlap, and
|
||||
scaling efficiency. Audit replicated cross-attention before using
|
||||
`skip_sequence_parallel`. Qualify regular and accelerated Ulysses transport
|
||||
independently; activation logs, JIT/readiness time, workspace growth,
|
||||
stream ownership, maximum-shape warmup, and numerical drift are gates.
|
||||
Preserve the checkpoint's Q:KV head ratio when padding GQA for Ulysses, and
|
||||
validate strided QKV staging rather than assuming packed contiguity.
|
||||
6. Evaluate text-encoder or VAE disaggregation only if stage share, transfer
|
||||
volume, reuse/concurrency, and independent scaling justify it. Tiling,
|
||||
offload, or patch parallelism is not disaggregation.
|
||||
|
||||
Do not combine independent optimizations into one performance claim. Preserve
|
||||
raw before/after runs with fixed workload and warmup policy.
|
||||
|
||||
### Gate 7: Make serving interruptible and load-stable
|
||||
|
||||
Treat request-mode batching and step-wise execution as separate capabilities:
|
||||
|
||||
- request batching requires `supports_request_batch = True` and
|
||||
`DiffusionRequestBatch -> list[DiffusionOutput]`, exactly one output per
|
||||
request;
|
||||
- step execution requires the complete `SupportsStepExecution` contract:
|
||||
`prepare_encode`, `denoise_step`, `step_scheduler`, and `post_decode`.
|
||||
|
||||
Keep all mutable scheduler, RNG, latent, mask, cache, and step state request
|
||||
scoped. First validate `--step-execution --max-num-seqs 1` and step-boundary
|
||||
abort; only then test heterogeneous multi-request waves. Step continuous
|
||||
batching is experimental and is not automatically compatible with Cache-DiT.
|
||||
Before recommending it, predeclare a useful case and success threshold, then
|
||||
A/B it against request mode under the same arrival process. Structural support
|
||||
or a generic bridge is not latency, throughput, or cancellation evidence for a
|
||||
model.
|
||||
|
||||
Test queued and in-flight abort, client disconnect, timeout, OOM, worker error,
|
||||
and restart. Require one terminal result, idempotent cleanup, no post-decode
|
||||
after abort, no cache/temp/VRAM/request-ID leak, and a successful next request.
|
||||
Inject late worker results/exceptions after cancellation and prove the result
|
||||
pump stays alive instead of raising `InvalidStateError`. For large batched or
|
||||
long-window video outputs, freeze dtype, range, layout, contiguity, ownership,
|
||||
payload size, serializer/codec limits, and offline-versus-HTTP semantics.
|
||||
Validate device-side output preparation, D2H/IPC or shared handles, remote
|
||||
encoding, event-loop responsiveness, and client materialization as separate
|
||||
stages, including abort and consumer failure cleanup.
|
||||
A model/VAE callback that publishes ordered chunks is only a producer contract;
|
||||
it is not transport, backpressure, cancellation, public streaming, or
|
||||
time-to-first-frame evidence. Distinguish complete-response faster-than-playback
|
||||
(`client E2E / output duration <= 1`) from streaming and report first-chunk,
|
||||
steady cadence, finalization, and complete-artifact latency separately.
|
||||
Run below-, near-, and above-saturation mixed-RPS soaks and report success/error
|
||||
rate, throughput, queue time, memory slope, temp growth, and worker health.
|
||||
Report p50/p95/p99 only from a declared arrival model with enough measured
|
||||
requests for the claimed percentile. See
|
||||
[serving and CI validation](references/production-validation.md).
|
||||
|
||||
### Gate 8: Publish recipes and four independent CI tracks
|
||||
|
||||
For every target vendor/card/topology, publish exact environment, serve/deploy
|
||||
command, one complete JSON or multipart curl per supported task, expected
|
||||
output checks, benchmark command, warmup/repeat/concurrency policy, raw metrics,
|
||||
per-rank HBM/host PSS, accuracy artifact, and limitations. Mark each
|
||||
CUDA/ROCm/NPU/XPU row `validated`, `unsupported`, or `not tested`; never inherit
|
||||
support from the platform abstraction or another card. Re-run DLO on each
|
||||
small-HBM recipe.
|
||||
|
||||
Maintain separate CI gates:
|
||||
|
||||
1. Function: strict load, API/offline positive and negative contracts.
|
||||
2. Accuracy: fixed reference revision/assets/seed/schedule and scoped tolerance.
|
||||
3. Performance: fixed best deployment plus memory-constrained/DLO row, raw JSON,
|
||||
explicit warmup, enough samples for each claimed statistic, and owned
|
||||
regression thresholds.
|
||||
4. Reliability: long mixed-RPS soak, abort/disconnect/fault cleanup, memory and
|
||||
temporary-resource trend.
|
||||
|
||||
## Production Definition of Done
|
||||
|
||||
A model is production-ready only when:
|
||||
|
||||
- every official advertised task passes API parity, offline/online parity,
|
||||
strict loading, fixed-reference quality, and negative validation;
|
||||
- the recommended deployment row passes correctness, isolation, abort/error
|
||||
cleanup, benchmark, and soak gates on named hardware;
|
||||
- online FP8 and DLO are each implemented and either validated in scoped rows
|
||||
or explicitly limited/unsupported with fail-fast behavior;
|
||||
- per-request cache quality tiers have measured speed/quality frontiers and
|
||||
request isolation, if advertised;
|
||||
- CUDA/ROCm/NPU/XPU recipes describe evidence rather than inferred support;
|
||||
- Function, Accuracy, Performance, and Reliability CI have owners, artifacts,
|
||||
and actionable failure output;
|
||||
- accelerated paths identify every precision/quality trade-off, and the raw
|
||||
offline plus online encoded output contracts are versioned and tested;
|
||||
- all unverified task/backend/cache/quant/offload/topology/hardware combinations
|
||||
remain `not tested` and are not presented as supported.
|
||||
|
||||
Day-0 support is still a valid milestone, but title, docs, and matrix must say
|
||||
which representative task passed and which production gates remain.
|
||||
@@ -0,0 +1,4 @@
|
||||
interface:
|
||||
display_name: "Production Add Diffusion Model"
|
||||
short_description: "Productionize vLLM-Omni diffusion models"
|
||||
default_prompt: "Use $production-add-diffusion-model to take this integrated diffusion model from Day-0 support to evidence-backed production readiness."
|
||||
@@ -0,0 +1,315 @@
|
||||
# API Contract and Hardware Recipes
|
||||
|
||||
## Contents
|
||||
|
||||
1. [Freeze the upstream contract](#freeze-the-upstream-contract)
|
||||
2. [Validate before side effects](#validate-before-side-effects)
|
||||
3. [Task request templates](#task-request-templates)
|
||||
4. [Serving benchmark template](#serving-benchmark-template)
|
||||
5. [Hardware recipe contract](#hardware-recipe-contract)
|
||||
6. [Best deployment selection](#best-deployment-selection)
|
||||
|
||||
## Freeze the upstream contract
|
||||
|
||||
Pin three revisions independently:
|
||||
|
||||
- official source implementation commit or package version;
|
||||
- checkpoint/model-card revision and partition;
|
||||
- vLLM-Omni target commit.
|
||||
|
||||
Build a table before editing code. One row is one official task/partition.
|
||||
|
||||
| Field | Required content |
|
||||
|---|---|
|
||||
| Task and partition | Official task name, checkpoint subdirectory, task selection rules |
|
||||
| Call contract | Constructor and call arguments, types, required/optional, official defaults |
|
||||
| Prompt/input | Prompt count/length, modality, MIME, order, count, decoded shape |
|
||||
| Size/time | Width/height alignment, pixels, frames, duration, FPS, start offset |
|
||||
| Denoise | Scheduler, sigma schedule, steps, CFG/guidance, seed behavior |
|
||||
| Output | Count, format, codec/sample rate, return type |
|
||||
| API | Offline mapping, sync endpoint, async endpoint/job lifecycle |
|
||||
| Rejection | Status code and stable error class/message for every invalid boundary |
|
||||
|
||||
Official API parity includes limits and defaults, not just names. Do not add a
|
||||
generic default that changes a distilled checkpoint's fixed schedule, or accept
|
||||
an input combination the official task rejects.
|
||||
|
||||
For every accepted boundary, add the neighboring rejected values. Examples:
|
||||
`max_images` and `max_images + 1`, smallest/largest legal duration, aligned and
|
||||
misaligned dimensions, omitted and illegal steps, one and too many outputs.
|
||||
|
||||
## Validate before side effects
|
||||
|
||||
Use one normalization/validation path for offline, sync, and async requests.
|
||||
The ordering must be:
|
||||
|
||||
```text
|
||||
parse -> normalize aliases -> validate task/source matrix -> validate declared
|
||||
limits -> stream with byte caps -> inspect/decode with frame/pixel/duration caps
|
||||
-> persist only if needed -> submit to engine
|
||||
```
|
||||
|
||||
Fail before download, persistence, job creation, or engine submit whenever the
|
||||
available metadata is sufficient. For streamed uploads, enforce:
|
||||
|
||||
- total reference count and per-modality count before reading bodies;
|
||||
- per-file bytes while streaming;
|
||||
- aggregate bytes across all files;
|
||||
- decoded frames, duration, pixels, and media type before engine submission.
|
||||
|
||||
On rejection, close streams and remove every partial temporary resource. Use
|
||||
HTTP 400 for invalid task/field combinations and 413 for payload limits. Test
|
||||
both sync and async paths; an invalid async request must not create a job ID.
|
||||
|
||||
Do not trust client MIME alone. Sniff/validate content and keep ordered repeated
|
||||
multipart fields ordered. Avoid logging prompts, raw references, signed URLs,
|
||||
or base64 payloads.
|
||||
|
||||
## Task request templates
|
||||
|
||||
These are endpoint-shape templates, not model recipes. Replace placeholders
|
||||
with the exact official defaults and validated assets for the target model.
|
||||
Every production recipe must include all advertised tasks and an expected
|
||||
output check.
|
||||
|
||||
### Text to image: JSON
|
||||
|
||||
```bash
|
||||
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/images/generations \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "<model>",
|
||||
"prompt": "<validated prompt>",
|
||||
"size": "1024x1024",
|
||||
"num_inference_steps": 30,
|
||||
"seed": 42,
|
||||
"response_format": "b64_json"
|
||||
}' | jq -er '.data[0].b64_json' | base64 -d > output.png
|
||||
file output.png
|
||||
```
|
||||
|
||||
### Image edit: multipart
|
||||
|
||||
```bash
|
||||
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/images/edits \
|
||||
-F 'model=<model>' \
|
||||
-F 'prompt=<validated edit prompt>' \
|
||||
-F 'image=@/path/to/reference.png;type=image/png' \
|
||||
-F 'size=1024x1024' \
|
||||
-F 'num_inference_steps=30' \
|
||||
-F 'seed=42' \
|
||||
-F 'output_format=png' \
|
||||
| jq -er '.data[0].b64_json' | base64 -d > edited.png
|
||||
file edited.png
|
||||
```
|
||||
|
||||
Use repeated `image` fields only if the model's official input matrix allows
|
||||
multiple ordered sources. Never describe a file upload as JSON-only.
|
||||
|
||||
### Public request fields versus internal sampling fields
|
||||
|
||||
Use the field exposed by the target revision's endpoint schema; do not rename
|
||||
it from the pipeline's internal attribute. For example, the current video JSON
|
||||
and multipart contracts expose `extra_params`, while serving translates that
|
||||
object into `DiffusionSamplingParams.extra_args`. Other request paths may
|
||||
expose canonical `extra_args` or nest it under `extra_body`. Inspect the exact
|
||||
Pydantic model/FastAPI form dependency and exercise the generated curl before
|
||||
publishing a recipe. A deprecation warning in a different endpoint is not
|
||||
evidence that an unrecognized multipart form field works here.
|
||||
|
||||
Current endpoint map (re-check it at the target commit):
|
||||
|
||||
| Endpoint | Public model-specific field |
|
||||
|---|---|
|
||||
| `/v1/videos`, `/v1/videos/sync` | Multipart `extra_params` JSON string |
|
||||
| `/v1/images/generations` | JSON `extra_params` object |
|
||||
| `/v1/chat/completions` diffusion | `extra_body.extra_args` |
|
||||
| `/v1/images/edits` | Declared typed form fields only |
|
||||
| `/v1/audio/generate` | Declared typed JSON fields only |
|
||||
|
||||
Prefer declared top-level fields. Use a generic model-specific envelope only
|
||||
when the endpoint actually exposes it.
|
||||
|
||||
### Text to video: multipart
|
||||
|
||||
```bash
|
||||
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
|
||||
-H 'Accept: video/mp4' \
|
||||
-F 'model=<model>' \
|
||||
-F 'prompt=<validated prompt>' \
|
||||
-F 'size=1280x720' \
|
||||
-F 'num_frames=81' \
|
||||
-F 'fps=16' \
|
||||
-F 'num_inference_steps=30' \
|
||||
-F 'seed=42' \
|
||||
-F 'extra_params={"task":"<official-task>"}' \
|
||||
-o output.mp4
|
||||
ffprobe -v error -show_entries stream=codec_name,width,height,r_frame_rate \
|
||||
-of json output.mp4
|
||||
```
|
||||
|
||||
### Image/video to video: multipart
|
||||
|
||||
```bash
|
||||
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
|
||||
-H 'Accept: video/mp4' \
|
||||
-F 'model=<model>' \
|
||||
-F 'prompt=<validated prompt>' \
|
||||
-F 'num_frames=81' \
|
||||
-F 'fps=16' \
|
||||
-F 'num_inference_steps=30' \
|
||||
-F 'seed=42' \
|
||||
-F 'extra_params={"task":"<official-task>","start_time_seconds":0}' \
|
||||
-F 'input_reference=@/path/to/reference.png;type=image/png' \
|
||||
-o i2v.mp4
|
||||
|
||||
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
|
||||
-H 'Accept: video/mp4' \
|
||||
-F 'model=<model>' \
|
||||
-F 'prompt=<validated prompt>' \
|
||||
-F 'num_frames=81' \
|
||||
-F 'num_inference_steps=30' \
|
||||
-F 'seed=43' \
|
||||
-F 'extra_params={"task":"<official-task>","start_time_seconds":0}' \
|
||||
-F 'input_reference=@/path/to/reference.mp4;type=video/mp4' \
|
||||
-o v2v.mp4
|
||||
```
|
||||
|
||||
For multiple references, repeat the model's accepted multipart field in
|
||||
official order. Include complete downloadable or repository fixtures plus
|
||||
SHA256 hashes. A placeholder-only recipe is not runnable evidence.
|
||||
|
||||
### Text to audio: JSON
|
||||
|
||||
```bash
|
||||
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/audio/generate \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{
|
||||
"model": "<model>",
|
||||
"input": "<validated audio prompt>",
|
||||
"audio_length": 10.0,
|
||||
"num_inference_steps": 100,
|
||||
"seed": 42,
|
||||
"response_format": "wav"
|
||||
}' -o output.wav
|
||||
ffprobe -v error -show_entries stream=codec_name,sample_rate,channels \
|
||||
-of json output.wav
|
||||
```
|
||||
|
||||
Use `/v1/audio/generate` for general diffusion audio. Do not substitute the
|
||||
speech API unless the model is actually served through the speech contract.
|
||||
|
||||
### Async video lifecycle
|
||||
|
||||
Alongside `/v1/videos/sync`, include `POST /v1/videos`, status polling,
|
||||
content retrieval, cancellation/abort if public, and cleanup assertions. The
|
||||
same request must normalize to the same engine sampling parameters in both
|
||||
modes.
|
||||
|
||||
## Serving benchmark template
|
||||
|
||||
Run a smoke first, then fixed-concurrency and arrival-rate tests. Video task
|
||||
names for the harness are `t2v`, `i2v`, `ti2v`, or `v2v`; model-specific task
|
||||
selection belongs in `--extra-body` when required.
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://127.0.0.1:8091 \
|
||||
--endpoint /v1/videos \
|
||||
--model '<model>' \
|
||||
--dataset random \
|
||||
--task t2v \
|
||||
--num-prompts 100 \
|
||||
--width 1280 --height 720 \
|
||||
--num-frames 81 --fps 16 \
|
||||
--num-inference-steps 30 \
|
||||
--request-rate 0.2 \
|
||||
--max-concurrency 4 \
|
||||
--warmup-requests 4 \
|
||||
--warmup-concurrency 4 \
|
||||
--extra-body '{"extra_params":"{\"task\":\"<official-task>\"}"}' \
|
||||
--output-file result.json
|
||||
```
|
||||
|
||||
The video backend currently stringifies each multipart value. Pre-serialize
|
||||
the nested `extra_params` object as above; passing it as a nested object would
|
||||
produce Python dict syntax rather than valid JSON. If the target revision
|
||||
changes the backend to JSON-encode dict/list values, update and smoke-test this
|
||||
command against that revision.
|
||||
|
||||
For i2v/v2v, use the harness's validated dataset/trace path or provide the
|
||||
documented media source; do not silently benchmark t2v and label it i2v. Keep
|
||||
prompt, assets, seed, dimensions, frames, steps, guidance, warmup, concurrency,
|
||||
and arrival process identical across A/B runs.
|
||||
|
||||
Keep three benchmark lanes distinct:
|
||||
|
||||
1. lossless runtime A/B with the released schedule/precision/attention policy;
|
||||
2. accelerated-path A/B that changes one adapter, quantization, cache, or
|
||||
sparsity policy and adds same-seed quality evidence;
|
||||
3. production-topology studies that vary placement or parallelism and report a
|
||||
latency/throughput/memory frontier rather than a kernel speedup.
|
||||
|
||||
For a fixed single-request A/B, retain every measured run and report median or
|
||||
mean with range after an explicit warmup; do not infer p95/p99 from a few
|
||||
repetitions. Tail latency requires a declared arrival process, request count,
|
||||
concurrency, and enough successful samples. Reconcile client E2E with queue,
|
||||
encoder, denoise, video/audio decode, D2H/IPC, codec, and residual timings.
|
||||
|
||||
For generated media, define `real-time` rather than relying on the label. A
|
||||
complete-response claim uses `client E2E / validated output duration <= 1` and
|
||||
does not imply streaming or low time-to-first-frame. Report first chunk/fragment
|
||||
and cadence separately when a public streaming route exists. Never add gains
|
||||
across benchmark lanes with different schedules, prompts, seeds, artifacts, or
|
||||
topologies.
|
||||
|
||||
The diffusion serving harness currently covers image/video task families, not
|
||||
`/v1/audio/generate`. For text-to-audio, use the target model's task-specific
|
||||
audio benchmark if one exists; otherwise add a serving harness that records
|
||||
request rate, concurrency, latency, generated-audio seconds per second, error
|
||||
rate, and output validation. Do not relabel the TTS harness or a video workload
|
||||
as diffusion-audio evidence.
|
||||
|
||||
## Hardware recipe contract
|
||||
|
||||
Create a separate row and runnable section for each vendor/card/topology:
|
||||
|
||||
| Evidence | Required detail |
|
||||
|---|---|
|
||||
| Identity | Vendor, exact SKU, device count, HBM, interconnect |
|
||||
| Software | Driver, CUDA/ROCm/CANN/oneAPI, framework/container, dependency versions |
|
||||
| Source | vLLM-Omni commit, official/checkpoint revision and hashes |
|
||||
| Deployment | dtype, attention backend per role, TP/SP/CFG/HSDP/DP, cache, quant, offload |
|
||||
| Commands | Complete deploy YAML/serve command, every task curl, benchmark command |
|
||||
| Method | Warmup exclusion, repetitions or request count, concurrency, arrival distribution, duration, statistic definitions |
|
||||
| Results | Success rate, p50/p95/p99, RPS/throughput per device, stage latency |
|
||||
| Memory | Load/materialize, resident, encode, denoise, decode, transient per-rank HBM and host PSS |
|
||||
| Quality | Metric/tolerance, artifact links/hashes, temporal/audio checks |
|
||||
| Output | Raw/offline dtype, range, layout and ownership; online codec, payload bytes, transport route and client boundary |
|
||||
| Limits | `validated`/`limited`/`unsupported`/`not tested`, fallback/rejection behavior |
|
||||
|
||||
Use current paths such as `recipes/<vendor>/`, `recipes/README.md`, supported
|
||||
models, and feature compatibility tables. Top-level documentation must link to
|
||||
the scoped matrix; it must not use an unqualified checkmark that hides a ROCm,
|
||||
NPU, XPU, task, topology, or dtype limitation.
|
||||
|
||||
DLO evidence is per card, especially on small-HBM devices. A DLO result on one
|
||||
CUDA card does not establish another CUDA SKU, ROCm, NPU, or XPU support.
|
||||
|
||||
## Best deployment selection
|
||||
|
||||
Benchmark a correctness-approved candidate matrix, not every theoretical
|
||||
combination. Choose the recommended row by the production objective—latency,
|
||||
throughput, cost, or small-HBM fit—and document that objective.
|
||||
|
||||
Do not recommend a row until it passes:
|
||||
|
||||
- all advertised task/API contracts;
|
||||
- fixed-reference quality;
|
||||
- request isolation and abort/error cleanup;
|
||||
- sustained arrival-rate tests;
|
||||
- per-rank HBM and host PSS limits;
|
||||
- hardware-specific accuracy and performance CI.
|
||||
|
||||
Keep the most memory-constrained DLO row even when it is not the fastest; it is
|
||||
a separate production recipe, not a failed latency candidate.
|
||||
@@ -0,0 +1,462 @@
|
||||
# Production Feature Patterns
|
||||
|
||||
## Contents
|
||||
|
||||
1. [Compatibility matrix](#compatibility-matrix)
|
||||
2. [Online FP8](#online-fp8)
|
||||
3. [Strict fused-weight loading](#strict-fused-weight-loading)
|
||||
4. [Distributed layerwise offload](#distributed-layerwise-offload)
|
||||
5. [Per-request Cache-DiT](#per-request-cache-dit)
|
||||
6. [Combination gates](#combination-gates)
|
||||
|
||||
## Compatibility matrix
|
||||
|
||||
Create rows at the granularity below. Add evidence links and a reason for every
|
||||
non-validated state.
|
||||
|
||||
| Task/shape | Execution mode/capacity | Backend/layout | Quant | Cache policy | Offload | Topology | Hardware/dtype | Output contract/transport | State | Evidence |
|
||||
|---|---|---|---|---|---|---|---|---|---|---|
|
||||
| `<task>/<shape>/<schedule>` | serial request / 1 | SDPA, padded | BF16 | none | resident | 1 device | `<card>` BF16 | raw float / local | `not tested` | — |
|
||||
| `<task>/<shape>/<schedule>` | step / N | `<fast>`, packed | BF16 | none | resident | `<TP/SP>` | `<card>` BF16 | encoded / subprocess | `not tested` | — |
|
||||
| `<task>/<shape>/<schedule>` | step / N | `<fast>`, packed | FP8 | high | DLO no-AG | `<TP/SP>` | `<card>` | device uint8 / SHM | `not tested` | — |
|
||||
|
||||
Initialize every row as `not tested`; the table is a work queue, not evidence.
|
||||
Promote the dense BF16 row to the oracle only after its scoped parity passes.
|
||||
Add one axis at a time. If the runtime cannot
|
||||
enforce a known-incompatible combination, add construction/admission
|
||||
validation before writing a recipe.
|
||||
|
||||
## Online FP8
|
||||
|
||||
### Component routing
|
||||
|
||||
The public builder accepts one method or a longest-prefix component map:
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
config = build_quant_config({
|
||||
"transformer": {
|
||||
"method": "fp8",
|
||||
"ignored_layers": ["transformer.final_layer"],
|
||||
},
|
||||
"text_encoder": None,
|
||||
"vae": None,
|
||||
"default": None,
|
||||
})
|
||||
|
||||
assert config.resolve("transformer.blocks.0.attn.qkv_proj") is not None
|
||||
assert config.resolve("text_encoder.layers.0.mlp") is None
|
||||
assert config.resolve("vae.decoder.conv_out") is None
|
||||
```
|
||||
|
||||
Use actual runtime prefixes. Multi-DiT or nested pipelines may not use the
|
||||
generic names `transformer` and `vae`. A prefix that falls through to a default
|
||||
is not component validation. Add tests for every included and ignored prefix.
|
||||
|
||||
For CLI serving, first validate resident DiT-only FP8:
|
||||
|
||||
```bash
|
||||
vllm serve '<model>' --omni --quantization fp8 --max-num-seqs 1
|
||||
```
|
||||
|
||||
For scoped routing:
|
||||
|
||||
```bash
|
||||
vllm serve '<model>' --omni \
|
||||
--diffusion-quantization-config \
|
||||
'{"<dit-runtime-prefix>":{"method":"fp8","ignored_layers":["<exact-linear-prefix>"]},"<vae-runtime-prefix>":null,"default":null}' \
|
||||
--max-num-seqs 1
|
||||
```
|
||||
|
||||
Do not publish these placeholders. A recipe must contain the target model's
|
||||
verified prefixes.
|
||||
|
||||
### Quantizable module wiring
|
||||
|
||||
Pass `quant_config` and a stable prefix through every vLLM linear. Keep
|
||||
precision-sensitive modulation/norm layers safe unless proven otherwise.
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.linear import (
|
||||
MergedColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
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}.attn.qkv_proj",
|
||||
)
|
||||
self.gate_up_proj = MergedColumnParallelLinear(
|
||||
hidden_size,
|
||||
[intermediate_size, intermediate_size],
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.mlp.gate_up_proj",
|
||||
)
|
||||
self.act_fn = SiluAndMul()
|
||||
self.down_proj = RowParallelLinear(
|
||||
intermediate_size,
|
||||
hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.mlp.down_proj",
|
||||
)
|
||||
```
|
||||
|
||||
Set required checkpoint parameters to fail rather than silently initialize if
|
||||
the current loader supports `missing_param_init = "error"`.
|
||||
|
||||
### FP8 evidence
|
||||
|
||||
For each task/hardware row, capture:
|
||||
|
||||
- startup proves the FP8 quant method is attached to intended linear modules;
|
||||
- no unexpected BF16 fallback and no meta/uninitialized parameters;
|
||||
- same-seed component, trajectory, and final-artifact comparison with BF16;
|
||||
- HBM at load, resident, encode, denoise, decode, and peak;
|
||||
- cold load time; fixed-work warm-latency raw runs; and serving p50/p95/p99 plus
|
||||
throughput only from a declared arrival load with enough samples;
|
||||
- named ignored layers and why their BF16 retention is required.
|
||||
|
||||
DiT FP8 evidence does not cover the text encoder or VAE. CUDA evidence does not
|
||||
cover ROCm/NPU/XPU. A pre-quantized checkpoint is a different loading path from
|
||||
online FP8 and needs its own state. If it is also pruned, rotated, distilled, or
|
||||
adapter-modified, report that model delta explicitly; BF16-versus-artifact
|
||||
quality and speed are deployment comparisons, not pure quantization ablations.
|
||||
|
||||
## Strict fused-weight loading
|
||||
|
||||
Never use `param.data.copy_()` for quantizable vLLM parameters. Transform the
|
||||
checkpoint layout before invoking the parameter loader so the online quant
|
||||
loader remains outermost:
|
||||
|
||||
```python
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
param = params[target_name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
|
||||
if target_name.endswith(".qkv_proj.weight"):
|
||||
packed_qkv = reorder_checkpoint_qkv(checkpoint_weight)
|
||||
weight_loader(param, packed_qkv)
|
||||
elif target_name.endswith(".gate_up_proj.weight"):
|
||||
gate, up = checkpoint_weight.chunk(2, dim=0)
|
||||
weight_loader(param, gate, 0)
|
||||
weight_loader(param, up, 1)
|
||||
else:
|
||||
weight_loader(param, checkpoint_weight)
|
||||
```
|
||||
|
||||
Track source shards per fused destination:
|
||||
|
||||
```python
|
||||
expected = {
|
||||
"blocks.0.attn.qkv_proj.weight": {"q", "k", "v"},
|
||||
"blocks.0.mlp.gate_up_proj.weight": {0, 1},
|
||||
}
|
||||
seen = {name: set() for name in expected}
|
||||
|
||||
# After mapping/loading each source:
|
||||
seen[target_name].add(shard_id)
|
||||
|
||||
incomplete = {
|
||||
name: sorted(shards - seen[name], key=str)
|
||||
for name, shards in expected.items()
|
||||
if shards - seen[name]
|
||||
}
|
||||
if incomplete:
|
||||
raise RuntimeError(f"incomplete fused checkpoint parameters: {incomplete}")
|
||||
```
|
||||
|
||||
Also compare retained model parameters with the returned loaded set. Explicit
|
||||
checkpoint ignores need a reason and test; an unknown or missing weight is not
|
||||
an ignore policy.
|
||||
|
||||
Load-time adapter fusion creates a dedicated model, not a request-switchable
|
||||
LoRA. Pin the base and adapter revisions, validate every low-rank and full-rank
|
||||
delta target, freeze the adapter-declared schedule/task, and fuse before
|
||||
sharding only when the loader contract proves identical rank-local weights.
|
||||
Reject offload or host-weight paths that bypass that fusion rather than serving
|
||||
the unfused base model silently.
|
||||
|
||||
## Distributed layerwise offload
|
||||
|
||||
### Declare component topology
|
||||
|
||||
Use `SupportsComponentDiscovery` and `OffloadPlan`. Validate all dotted paths
|
||||
at construction/test time because discovery may otherwise warn and skip.
|
||||
|
||||
Treat a future/requested `--layerwise-offload-components` selector as a public
|
||||
API only if it exists in the target revision. Today, express component topology
|
||||
with `OffloadPlan` (`on_demand_component_paths`, `block_attrs`,
|
||||
`encoder_block_attrs`, `resident_dit_paths`, and `offload_submodules`). Do not
|
||||
invent a CLI flag or claim component-selective DLO from discovery declarations
|
||||
alone; validate the selector separately when it lands.
|
||||
|
||||
```python
|
||||
from operator import attrgetter
|
||||
|
||||
import torch
|
||||
|
||||
def assert_offload_path(root, path):
|
||||
value = attrgetter(path)(root)
|
||||
if not isinstance(value, torch.nn.Module):
|
||||
raise TypeError(f"offload path {path!r} is not an nn.Module")
|
||||
return value
|
||||
|
||||
for path in (
|
||||
*pipeline._dit_modules,
|
||||
*pipeline._encoder_modules,
|
||||
*pipeline._vae_modules,
|
||||
*pipeline._offload_plan.on_demand_component_paths,
|
||||
):
|
||||
assert_offload_path(pipeline, path)
|
||||
```
|
||||
|
||||
For `block_attrs`, resolve the DiT first and then each declared block attribute.
|
||||
Require an indexable block container with at least one `nn.Module`. Validate
|
||||
`offload_submodules`, `resident_dit_paths`, and `encoder_block_attrs` similarly.
|
||||
|
||||
### Keep host storage loader-owned
|
||||
|
||||
Current DLO accepts a loader-produced `HostWeightPlan`. Pipeline authors should
|
||||
not instantiate this plan or add ad hoc mmap capability flags. The loader must
|
||||
preflight complete runtime/checkpoint bindings, names, shapes, dtypes, persistent
|
||||
buffers, and represented transforms before it skips ordinary materialization.
|
||||
The exact accepted plan is consumed once by DLO; a rejected plan records an
|
||||
observable fallback reason and continues through the ordinary loader.
|
||||
|
||||
Direct checkpoint mmap currently requires TP1 without HSDP or online
|
||||
quantization. TP>1 falls back to the ordinary TP-aware loader and may feed DLO
|
||||
AllGather or no-AllGather; the current design records a bounded DP2xTP2
|
||||
AllGather smoke, but that path does not receive direct-mmap host-page sharing.
|
||||
Per-tensor online FP8 must use the ordinary loader. Current revisions can admit
|
||||
its finalized weight/scale layout to DLO AllGather, including transposed runtime
|
||||
weights, while unsupported online quantizers still fail closed. Direct mmap
|
||||
remains unavailable, so capture the temporary full-model host materialization
|
||||
before rank sharding.
|
||||
|
||||
Host Weight Runtime is a distinct opt-in path. Current final-layout consumers
|
||||
use no-AllGather DLO and exact model-declared BF16 representations. In
|
||||
`preferred` mode, consume an exact local hit or canonically load and publish for
|
||||
a future startup; in `required` mode, consume only and fail on a miss or unusable
|
||||
artifact. HWR must not interact with DLO AllGather, online quantization, HSDP,
|
||||
LoRA/adapted weights, or an unrecognized load format unless the target revision
|
||||
adds an exact producer/restore contract. It is immutable node-local CPU backing,
|
||||
not zero-copy GPU execution; the DLO transport still owns staging or direct H2D.
|
||||
Qualify identity inputs, cold population, warm reuse, atomic concurrent
|
||||
publication, source digests, corruption/quarantine, lease cleanup, capacity,
|
||||
node-local PSS, H2D, and E2E latency. A host-memory saving that materially
|
||||
regresses transfer or request latency is a separate memory-first deployment.
|
||||
|
||||
### DLO deployment modes
|
||||
|
||||
Treat these paths separately:
|
||||
|
||||
| Path | Weight/runtime behavior | Required boundaries |
|
||||
|---|---|---|
|
||||
| DLO AllGather | Rank-sharded pinned host tensors; reconstruct layer through collective | TP>1 uses ordinary TP-aware loader output; HSDP and unsupported online quantizers are rejected; finalized per-tensor online FP8 is eligible only when recognized by the target revision; concurrent ranks must participate consistently |
|
||||
| DLO no-AllGather | Loader-approved checkpoint mmap or ordinary runtime tensors; each rank streams a complete block | Direct mmap is TP1/non-HSDP/non-online-quant only; TP/HSDP/online quant use ordinary runtime tensors and require scoped E2E |
|
||||
| SP + DLO | SP group can supply the DLO collective group | Packed boundaries and collective order require E2E |
|
||||
|
||||
Single-stage examples:
|
||||
|
||||
```bash
|
||||
# SP + DLO AllGather candidate; validate exact model/card before publishing.
|
||||
vllm serve '<model>' --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4 \
|
||||
--max-num-seqs 1
|
||||
|
||||
# TP + DLO rank-local candidate; not a support statement.
|
||||
vllm serve '<model>' --omni \
|
||||
--tensor-parallel-size 2 \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--dlo-no-use-allgather \
|
||||
--max-num-seqs 1
|
||||
```
|
||||
|
||||
Under `--omni`, configure DP in a deploy YAML rather than passing vLLM DP CLI
|
||||
flags:
|
||||
|
||||
```yaml
|
||||
pipeline: <registered-pipeline>
|
||||
async_chunk: false
|
||||
data_parallel_size: 4
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0,1,2,3"
|
||||
max_num_seqs: 1
|
||||
enable_distributed_layerwise_offload: true
|
||||
dlo_use_allgather: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 1
|
||||
sequence_parallel_size: 1
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve '<model>' --omni --deploy-config /path/to/dlo_dp4.yaml
|
||||
```
|
||||
|
||||
Keep the YAML next to the validated recipe. Match its schema to the target
|
||||
revision; do not assume a model accepts another pipeline's deploy config.
|
||||
|
||||
### DLO validation sequence
|
||||
|
||||
1. Resident BF16 reference.
|
||||
2. Ordinary layerwise offload parity.
|
||||
3. DLO enable, warmup, same-seed generation, disable/restart.
|
||||
4. AllGather and no-AllGather separately.
|
||||
5. Concurrent requests with different prompts and identical collective-affecting
|
||||
parameters; then reject or schedule incompatible steps/shapes safely.
|
||||
6. Queued/in-flight abort, error, idle rank, worker exit, and next request.
|
||||
7. TP, SP, cache, compile, online FP8 one axis at a time.
|
||||
8. Loader-plan selection/fallback reason, cold/warm storage lifecycle, and
|
||||
corruption/stale-entry recovery when mmap/cache backing is advertised.
|
||||
9. Replica-local request preparation plus DLO-group-only weight collectives in
|
||||
DP; exercise concurrent replicas to catch accidental WORLD rendezvous.
|
||||
10. Sweep resident block counts; publish raw latency/HBM pairs, dominated
|
||||
points, and the non-dominated Pareto frontier.
|
||||
11. Per-rank HBM, host PSS for the full process tree, H2D and collective trace.
|
||||
|
||||
On a small-HBM card, measure transient encode/VAE peaks as well as resident DiT
|
||||
weights. If a non-DiT component still exceeds HBM, combine only independently
|
||||
validated on-demand component staging, VAE tiling/patching, or another topology;
|
||||
do not hide the remaining peak behind the phrase “DLO enabled.”
|
||||
|
||||
## Per-request Cache-DiT
|
||||
|
||||
### Policy and lifecycle skeleton
|
||||
|
||||
First verify that the target revision contains the shared protocol/runtime.
|
||||
If not, follow the nearest current model's lifecycle without inventing a new
|
||||
public API.
|
||||
|
||||
The following is a **serial/exclusive request-mode sketch only**. Pipeline-wide
|
||||
hooks are mutable; concurrent requests need an admission/batching key and
|
||||
runner-owned transition serialization before this shape is safe.
|
||||
|
||||
```python
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.cache.cachedit import (
|
||||
CacheDiTBackend,
|
||||
CacheDiTRequestSpec,
|
||||
RequestScopedCacheDiTRuntime,
|
||||
)
|
||||
|
||||
class MyPipeline(nn.Module):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__()
|
||||
# ... components ...
|
||||
self._cache_dit_runtime = RequestScopedCacheDiTRuntime(self)
|
||||
|
||||
def adopt_cache_dit_backend(self, backend: CacheDiTBackend) -> None:
|
||||
self._cache_dit_runtime.adopt(
|
||||
backend,
|
||||
installation_key="my_model.generic",
|
||||
)
|
||||
|
||||
def is_cache_dit_enabled(self) -> bool:
|
||||
return self._cache_dit_runtime.is_enabled
|
||||
|
||||
def _cache_spec(self, quality, steps):
|
||||
if quality == "lossless":
|
||||
return None
|
||||
if quality == "high":
|
||||
return CacheDiTRequestSpec(
|
||||
installation_key="my_model.high.v1",
|
||||
cache_config=self._validated_high_config,
|
||||
num_inference_steps=steps,
|
||||
)
|
||||
raise ValueError(f"unsupported quality tier: {quality!r}")
|
||||
|
||||
def forward(self, req):
|
||||
quality, steps = validate_and_resolve_request(req)
|
||||
self._cache_dit_runtime.prepare(self._cache_spec(quality, steps))
|
||||
return self._run_validated_request(req)
|
||||
```
|
||||
|
||||
The snippet shows the model-owned transition for exclusive execution. Put
|
||||
disable/restore handling in the runtime/runner lifecycle that owns success,
|
||||
exception, disconnect, and asynchronous abort; a `try/finally` inside
|
||||
`forward()` alone cannot observe every server cancellation boundary.
|
||||
|
||||
The `installation_key` must include every policy dimension that changes hook
|
||||
installation; changes in step count can use runtime refresh. `None` disables
|
||||
installed hooks. Do not mutate the global startup cache config for a request.
|
||||
|
||||
Centralize cleanup in the owning request/runner boundary. An exception, abort,
|
||||
disconnect, or admission failure must not leave a changed hook policy visible
|
||||
to the next request. The exact owner depends on the target revision; add a
|
||||
failure-injection test rather than relying on `forward()` alone to observe
|
||||
asynchronous abort.
|
||||
|
||||
Example model-specific request form after calibration:
|
||||
|
||||
```bash
|
||||
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
|
||||
-F 'model=<model>' \
|
||||
-F 'prompt=<prompt>' \
|
||||
-F 'quality=lossless' \
|
||||
-F 'num_inference_steps=<validated-steps>' \
|
||||
-o lossless.mp4
|
||||
|
||||
curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \
|
||||
-F 'model=<model>' \
|
||||
-F 'prompt=<prompt>' \
|
||||
-F 'quality=high' \
|
||||
-F 'num_inference_steps=<validated-steps>' \
|
||||
-o high.mp4
|
||||
```
|
||||
|
||||
Do not expose these tier names merely because another model uses them.
|
||||
|
||||
### Cache-DiT evidence
|
||||
|
||||
For every tier/task/shape/schedule:
|
||||
|
||||
- show real cache hits/skipped blocks, not only “backend enabled” logs;
|
||||
- measure fixed-work speed and quality against lossless/native;
|
||||
- publish a speed/quality frontier and chosen tier mapping;
|
||||
- alternate lossless/high/lossless and vary steps to prove restore/refresh;
|
||||
- interleave tasks and shapes; run concurrent requests only when scheduling
|
||||
guarantees exclusive compatible hook state;
|
||||
- inject validation error, denoise error, disconnect, and abort, then prove the
|
||||
next request starts with its requested policy;
|
||||
- bound any cached tensors and clear request-local state.
|
||||
|
||||
## Combination gates
|
||||
|
||||
Start every combination as `not tested`.
|
||||
|
||||
| Combination | Default production posture |
|
||||
|---|---|
|
||||
| Direct checkpoint mmap + TP>1/HSDP/online quant | Preflight falls back to the ordinary loader; do not claim direct-mmap savings |
|
||||
| Per-tensor online FP8 + DLO AllGather | Generic compatibility exists for finalized weights/scales; each model/card/topology remains a candidate until E2E, and other online quantizers fail closed |
|
||||
| Online FP8 + DLO no-AllGather | Candidate only after model/card/topology E2E |
|
||||
| TP>1 + DLO AllGather | Ordinary TP-aware loader only; bounded DP2xTP2 smoke exists, but each new model/card/topology still needs E2E |
|
||||
| TP>1 + DLO no-AllGather | Ordinary TP-aware loader only; candidate with limited generic coverage |
|
||||
| HSDP + DLO AllGather | Reject; it would double-shard HSDP-managed parameters |
|
||||
| HSDP + DLO no-AllGather | Configuration-compatible candidate with limited E2E coverage |
|
||||
| Cache-DiT + online FP8 | Separate trajectory, hit-rate, HBM, latency test |
|
||||
| Cache-DiT + DLO | Verify hooks, streamed blocks, memory peaks, abort cleanup |
|
||||
| Cache-DiT + step execution | Do not advertise until lifecycle/batching support is explicit |
|
||||
| Sparse attention + Ring | Verify backend execution; reject silently ignored sparse settings |
|
||||
| Load-time fused adapter + offload/HWR | Reject unless the selected loader path applies and validates the same fusion before sharding/storage |
|
||||
| Packed attention + any topology | Unequal-sample boundary parity is mandatory |
|
||||
| Device-side uint8 preparation + offline caller | Version the dtype/range/layout change; HTTP MP4 parity does not preserve the raw tensor contract |
|
||||
| Device-side preparation + remote codec | Validate route selection, strided/contiguous layouts, payload bytes, ordering, fallback, and encoded-media parity |
|
||||
|
||||
A startup-only test never upgrades a combination. Require request output,
|
||||
parity/quality, actual fast-path evidence, resource cleanup, and a benchmark on
|
||||
the named hardware.
|
||||
@@ -0,0 +1,481 @@
|
||||
# Performance Patterns and Code Examples
|
||||
|
||||
## Contents
|
||||
|
||||
1. [Measurement loop](#measurement-loop)
|
||||
2. [Faster BF16 attention](#faster-bf16-attention)
|
||||
3. [Sparse attention](#sparse-attention)
|
||||
4. [Fusion-first operator patterns](#fusion-first-operator-patterns)
|
||||
5. [Remove redundant computation](#remove-redundant-computation)
|
||||
6. [USP communication optimization](#usp-communication-optimization)
|
||||
7. [Text encoder and VAE disaggregation](#text-encoder-and-vae-disaggregation)
|
||||
8. [Performance acceptance](#performance-acceptance)
|
||||
|
||||
## Measurement loop
|
||||
|
||||
Do not start with a list of optimizations. Start with a fixed correctness and
|
||||
performance workload:
|
||||
|
||||
```text
|
||||
model/checkpoint revision + prompt/media hashes + seed + dimensions/frames +
|
||||
scheduler/sigma/steps/guidance + task + dtype + backend + topology + hardware
|
||||
```
|
||||
|
||||
Capture:
|
||||
|
||||
- stage time: load, encode, conditioning, denoise, VAE decode/postprocess;
|
||||
- GPU trace: kernel duration/count, launch gaps, allocations, H2D, collectives;
|
||||
- CPU trace: preprocessing, per-layer sync, Python dispatch, temporary files;
|
||||
- memory: per-rank HBM peaks and process-tree host PSS;
|
||||
- output: device preparation, D2H/IPC, payload bytes, codec wall/process CPU,
|
||||
event-loop wait, client materialization, and signed residual;
|
||||
- serving: queue time, throughput, errors, request mix, and distribution samples.
|
||||
|
||||
For a fixed-work single-request A/B, run one explicit warmup policy and at
|
||||
least three measured repetitions, keep every raw run, and report median or mean
|
||||
with range. Do not derive p95/p99 from three samples. For tail latency, declare
|
||||
the arrival model and collect enough successful requests for the percentile and
|
||||
confidence claimed. Implement one optimization at a time, re-run parity, and
|
||||
report both local kernel and end-to-end deltas. A kernel win that moves no E2E
|
||||
metric is not a production speed claim.
|
||||
|
||||
Use the pipeline profiler already exposed by the target revision when
|
||||
available, and use PyTorch/CUDA/ROCm/NPU/XPU profiling tools appropriate to the
|
||||
platform. Do not force every platform through a CUDA-only workflow.
|
||||
|
||||
## Faster BF16 attention
|
||||
|
||||
### Shared attention wiring
|
||||
|
||||
Use the shared role-aware attention layer and give it the true local head
|
||||
counts and QKV layout:
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.linear import (
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
|
||||
class SelfAttention(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size,
|
||||
head_dim,
|
||||
num_heads,
|
||||
num_kv_heads,
|
||||
quant_config=None,
|
||||
prefix="",
|
||||
):
|
||||
super().__init__()
|
||||
self.head_dim = head_dim
|
||||
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.attention = 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,
|
||||
)
|
||||
self.out_proj = RowParallelLinear(
|
||||
num_heads * head_dim,
|
||||
hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.out_proj",
|
||||
)
|
||||
|
||||
def forward(self, x, attn_metadata=None):
|
||||
qkv, _ = self.qkv_proj(x)
|
||||
q_size = self.qkv_proj.num_heads * self.head_dim
|
||||
kv_size = self.qkv_proj.num_kv_heads * self.head_dim
|
||||
q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
|
||||
q = q.unflatten(-1, (self.qkv_proj.num_heads, self.head_dim))
|
||||
k = k.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim))
|
||||
v = v.unflatten(-1, (self.qkv_proj.num_kv_heads, self.head_dim))
|
||||
out = self.attention(q, k, v, attn_metadata)
|
||||
out = out.flatten(-2)
|
||||
out, _ = self.out_proj(out)
|
||||
return out
|
||||
```
|
||||
|
||||
Add Q/K norm and RoPE in the official order before the `Attention` call. Keep
|
||||
checkpoint mapping for the fused QKV and output projection complete.
|
||||
|
||||
### Backend qualification
|
||||
|
||||
For each role (`self`, cross-attention category, token refiner, etc.), record:
|
||||
|
||||
- actual selected backend from logs/trace;
|
||||
- supported dtype, head size, QKV layout, mask/metadata, packed varlen behavior;
|
||||
- platform and compute capability;
|
||||
- TP/SP/Ring/AllGather-KV restrictions;
|
||||
- fallback behavior and whether fallback changes the claimed performance.
|
||||
|
||||
Use unequal packed samples and compare with dense SDPA. Verify `cu_seqlens`,
|
||||
max sequence length, padding exclusion, modality boundaries, CFG ownership, and
|
||||
output slicing. Do not pass an `attn_mask` and assume it is consumed. Inspect
|
||||
the backend contract and trace.
|
||||
|
||||
TRTLLM, Flash Attention, cuDNN, NPU varlen, and other backends are independent
|
||||
qualification rows. The shared `Attention` class does not make every backend
|
||||
packed-safe or faster automatically. Ring and hybrid paths require their own
|
||||
boundary parity. AllGather-KV and specific backend combinations may fail fast;
|
||||
keep that behavior visible in docs.
|
||||
|
||||
## Sparse attention
|
||||
|
||||
Sparse attention is a backend/model/platform optimization, not a synonym for
|
||||
low-precision dense attention. Enable it only after dense BF16 parity.
|
||||
|
||||
For every candidate task/shape/topology:
|
||||
|
||||
1. State the sparsity policy, selected backend/kernel, block/head dimensions,
|
||||
supported hardware, external kernel/package revision, and fallback.
|
||||
2. Prove the sparse kernel executed and report realized sparsity; an enabled
|
||||
flag is insufficient.
|
||||
3. Compare same-seed per-block outputs, denoise trajectory, and final artifact
|
||||
with dense attention.
|
||||
4. Add modality-specific metrics: temporal consistency and audio/video sync for
|
||||
video, spectral/audio quality for audio, and prompt/image metrics as relevant.
|
||||
5. Report attention time and full E2E time. Include preprocessing/index-build
|
||||
overhead, peak memory, and the exact dense control.
|
||||
6. Validate each task and shape. Do not extrapolate a T2V result to I2V/Ref2V,
|
||||
or NPU sparse support to CUDA/ROCm/XPU.
|
||||
|
||||
If Ring bypasses the selected backend or a sparse setting would be silently
|
||||
ignored, reject the combination. Prefer a visible dense fallback only when the
|
||||
recipe labels the row `limited` and does not claim sparse speed.
|
||||
|
||||
## Fusion-first operator patterns
|
||||
|
||||
Prefer current shared/custom operators with explicit guards and native
|
||||
fallbacks. Prove parity at operator, block, trajectory, and artifact levels,
|
||||
then use a trace to prove fewer launches/bytes or lower latency.
|
||||
|
||||
### Q/K RMSNorm + RoPE
|
||||
|
||||
```python
|
||||
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)
|
||||
# Example only: set these booleans from the official checkpoint contract.
|
||||
use_neox_layout = False
|
||||
cos_sin_are_half_dim = True
|
||||
self.rope = RotaryEmbedding(
|
||||
is_neox_style=use_neox_layout,
|
||||
half_head_dim=cos_sin_are_half_dim,
|
||||
)
|
||||
|
||||
q = self.norm_q(q)
|
||||
k = self.norm_k(k)
|
||||
q, k = apply_rope_to_qk(self.rope, q, k, (cos, sin))
|
||||
```
|
||||
|
||||
Replace the illustrative booleans with the official layout in real code.
|
||||
Preserve partial rotary dimensions by splitting the rotated prefix and concatenating
|
||||
the pass-through suffix. If 3D cos/sin includes a batch axis, the current
|
||||
shared path assumes the values are shared across the batch; per-sample
|
||||
different positions require a proven representation/path.
|
||||
|
||||
For packed `[tokens, heads, head_dim]` tensors with a non-interleaved
|
||||
`[cos | sin]` table, the current shared fused boundary is:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.layers.fused_qk_norm_rope import fused_qk_norm_rope
|
||||
|
||||
q, k = fused_qk_norm_rope(
|
||||
q,
|
||||
k,
|
||||
self.norm_q.weight,
|
||||
self.norm_k.weight,
|
||||
rope_table,
|
||||
self.norm_q.variance_epsilon,
|
||||
)
|
||||
```
|
||||
|
||||
Its CUDA fast path currently specializes BF16 `head_dim=128` and
|
||||
`rotary_dim=96`; other valid inputs take the eager implementation. The separate
|
||||
shared RMSNorm and RotaryEmbedding still provide the general cross-platform
|
||||
baseline. Treat NPU/MUSA and other vendor fusions as independent qualification
|
||||
rows, and do not generalize a model-specific Triton kernel without compatible
|
||||
shape, layout, frequency partition, dtype, and platform evidence.
|
||||
|
||||
### Fused QKV projection
|
||||
|
||||
Use `QKVParallelLinear` as shown above. Qualification includes:
|
||||
|
||||
- `num_heads` and `num_kv_heads` divisibility by TP;
|
||||
- local split sizes from the layer, never reusing global counts after TP;
|
||||
- Q/K/V checkpoint order conversion before the vLLM loader;
|
||||
- all Q, K, and V shards present;
|
||||
- quantization scale/loader mapping when FP8 or another method is enabled.
|
||||
|
||||
### Gate-up projection + fused SwiGLU
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.linear import (
|
||||
MergedColumnParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
self.gate_up_proj = MergedColumnParallelLinear(
|
||||
input_size=hidden_size,
|
||||
output_sizes=[intermediate_size, intermediate_size],
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.gate_up_proj",
|
||||
)
|
||||
self.act_fn = SiluAndMul()
|
||||
self.down_proj = RowParallelLinear(
|
||||
input_size=intermediate_size,
|
||||
output_size=hidden_size,
|
||||
bias=False,
|
||||
input_is_parallel=True,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.down_proj",
|
||||
)
|
||||
|
||||
gate_up, _ = self.gate_up_proj(x)
|
||||
x = self.act_fn(gate_up)
|
||||
x, _ = self.down_proj(x)
|
||||
```
|
||||
|
||||
Map both checkpoint `gate_proj` and `up_proj` shards into the merged parameter
|
||||
and require both. A merged projection followed by `F.silu(gate) * up` is not
|
||||
the same claim as the `SiluAndMul` fused activation kernel; identify what the
|
||||
trace actually uses.
|
||||
|
||||
### AdaLN
|
||||
|
||||
Use a shared layer only when its mathematical/return contract matches the
|
||||
official model:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
|
||||
|
||||
self.adaln = AdaLayerNorm(
|
||||
hidden_size,
|
||||
elementwise_affine=False,
|
||||
eps=1e-6,
|
||||
)
|
||||
x = self.adaln(x, scale[:, None, :], shift[:, None, :])
|
||||
```
|
||||
|
||||
`AdaLayerNormZero`, `AdaLayerNormZeroSingle`, and continuous variants include
|
||||
different projections, broadcasts, and tuple outputs; they are not drop-in
|
||||
replacements. A shared AdaLN may have an NPU fast path while using native ops on
|
||||
CUDA/HIP/XPU. Do not call it a universal fused AdaLN kernel.
|
||||
|
||||
For sigma schedules with a small finite set of timesteps, profile whether
|
||||
precomputing the timestep/AdaLN projection once per request is worthwhile.
|
||||
Cache keys must include the exact schedule/config/device/dtype and stay bounded.
|
||||
|
||||
### Layout, residual, and other candidates
|
||||
|
||||
After profiling, inspect:
|
||||
|
||||
- bias/dropout/add or gated residual chains;
|
||||
- norm + linear or norm + modulation boundaries;
|
||||
- repeated transpose/reshape/contiguous around attention;
|
||||
- VAE normalization/activation/convolution boundaries;
|
||||
- MoE routing and fused experts for models that actually use MoE.
|
||||
|
||||
For fused residual-plus-norm or modulation chains, preserve the unfused
|
||||
materialization boundary. If the residual was stored in BF16 before RMSNorm,
|
||||
round to BF16 and promote again inside the fused kernel; using the unrounded
|
||||
FP32 temporary changes the denoising trajectory even when the stored residual
|
||||
looks correct.
|
||||
|
||||
Keep DLO materialize/release hooks, process-group collectives, request cache
|
||||
transitions, and cleanup outside compiled/fused regions unless lifecycle safety
|
||||
is explicitly tested. Compile only stable regions and retain eager fallback.
|
||||
|
||||
## Remove redundant computation
|
||||
|
||||
### Hoist request invariants
|
||||
|
||||
Audit the denoise loop for work that depends only on the request rather than
|
||||
the current timestep:
|
||||
|
||||
```python
|
||||
# Before the denoise loop.
|
||||
condition = self.condition_proj(text_embeddings)
|
||||
refined_condition = self.token_refiner(condition, packed_metadata)
|
||||
rope_tables = self.build_rope_tables(position_ids)
|
||||
attn_metadata = self.build_attention_metadata(layout)
|
||||
|
||||
for step, sigma in enumerate(sigmas):
|
||||
timestep_state = self.prepare_timestep_state(sigma)
|
||||
noise = self.transformer(
|
||||
latents,
|
||||
refined_condition,
|
||||
rope_tables=rope_tables,
|
||||
attn_metadata=attn_metadata,
|
||||
timestep_state=timestep_state,
|
||||
)
|
||||
latents = self.scheduler_step(noise, sigma, latents)
|
||||
```
|
||||
|
||||
Candidates include prompt conditioning, token-refiner output, reference
|
||||
encoding/scans, position IDs/RoPE tables, packed masks/metadata, static CFG
|
||||
branches, and finite-schedule AdaLN projections. First prove they are invariant
|
||||
for all official tasks, LoRA/adapters, prompt-update features, shapes, and
|
||||
schedules. In continuous batching, store request-specific results on request
|
||||
state rather than the pipeline singleton.
|
||||
|
||||
### Eliminate per-layer synchronization
|
||||
|
||||
Avoid device-to-host scalar extraction such as `.item()` inside every layer:
|
||||
|
||||
```python
|
||||
# Bad: forces a GPU/CPU synchronization in every attention layer.
|
||||
max_len = int(cu_seqlens[-1].item())
|
||||
|
||||
# Better: compute validated host metadata once during request packing and pass it.
|
||||
packed = PackedLayout(
|
||||
cu_seqlens=cu_seqlens,
|
||||
max_seqlen=max(sample_lengths),
|
||||
)
|
||||
for block in self.blocks:
|
||||
hidden = block(hidden, packed_layout=packed)
|
||||
```
|
||||
|
||||
Trace to verify the synchronization disappears. Do not replace a correct
|
||||
dynamic value with a stale constant merely to remove `.item()`.
|
||||
|
||||
### Remove unused masks and allocations
|
||||
|
||||
If packed attention consumes `cu_seqlens` and ignores a dense mask, do not
|
||||
construct the dense mask. Reuse immutable request metadata and preallocated
|
||||
buffers when safe. Prefer direct indexed scatter/gather over constructing large
|
||||
temporary tensors only when duplicate-index semantics and gradients (normally
|
||||
disabled for inference) remain correct.
|
||||
|
||||
Any cross-request cache needs a bounded lifetime and a complete key containing
|
||||
shape, layout, device, dtype, schedule, task, adapter/LoRA, and any value that
|
||||
changes the tensor. Otherwise keep it request-scoped.
|
||||
|
||||
### VAE and reference work
|
||||
|
||||
Profile duplicate VAE decode/postprocess, repeated reference decoding/scans,
|
||||
unnecessary host-device round trips, and full-frame materialization. Apply VAE
|
||||
tiling, slicing, patch parallelism, or on-demand staging only with seam,
|
||||
ordering, color/range, and memory parity tests.
|
||||
|
||||
## USP communication optimization
|
||||
|
||||
“USP works” is not evidence that USP communication is optimized. For Ulysses,
|
||||
Ring, hybrid, and AllGather-KV separately, profile:
|
||||
|
||||
- collective count, bytes, duration, stream, and overlap with compute;
|
||||
- Q/K/V and output all-to-all boundaries;
|
||||
- pack/unpack, padding, metadata broadcast, transpose, and contiguous copies;
|
||||
- head divisibility or advanced uneven-head support;
|
||||
- scaling efficiency from 1 to N devices on the same workload/interconnect;
|
||||
- load balance for unequal packed samples.
|
||||
|
||||
Declare `_sp_plan` split/gather boundaries around meaningful module outputs,
|
||||
not arbitrary Python statements. Keep packed metadata sharded/gathered in the
|
||||
same ownership model as tokens.
|
||||
|
||||
When the target revision offers accelerated Ulysses transport such as
|
||||
`--ulysses-a2a-permute`, qualify it separately from regular Ulysses. Prove the
|
||||
strict scatter-heads/gather-sequence layout selected the fast path, record
|
||||
one-time JIT/readiness cost, prewarm the maximum workspace shape before CUDA
|
||||
graph capture, keep the grow-only workspace on one stream, and exercise cleanup
|
||||
before process-group destruction. Compare collective and E2E deltas after the
|
||||
dense backend and topology are frozen; keep unsupported layouts on an explicit
|
||||
regular-Ulysses fallback.
|
||||
|
||||
For GQA, validate `Q_heads % KV_heads == 0`, pad KV heads first, and derive Q
|
||||
padding from the original ratio. Exercise both contiguous and row-strided QKV
|
||||
layouts; a fast copy path must preserve bitwise values and packed boundaries.
|
||||
|
||||
Cross-attention sometimes uses replicated text K/V while only video/audio Q is
|
||||
sequence-sharded. In that specific case, `skip_sequence_parallel=True` can
|
||||
avoid incorrect or wasteful Q/K/V redistribution:
|
||||
|
||||
```python
|
||||
self.cross_attention = Attention(
|
||||
num_heads=local_heads,
|
||||
num_kv_heads=local_kv_heads,
|
||||
head_size=head_dim,
|
||||
softmax_scale=head_dim**-0.5,
|
||||
causal=False,
|
||||
role="cross",
|
||||
skip_sequence_parallel=True,
|
||||
prefix=f"{prefix}.cross_attn",
|
||||
)
|
||||
```
|
||||
|
||||
Use it only after proving K/V are replicated and Q/output ownership is correct.
|
||||
Otherwise it disables needed communication. Report whether the optimization
|
||||
reduces collective bytes/time and improves E2E scaling.
|
||||
|
||||
## Text encoder and VAE disaggregation
|
||||
|
||||
Disaggregation is a measured architecture decision. First collect:
|
||||
|
||||
| Question | Required measurement |
|
||||
|---|---|
|
||||
| Is the stage large enough? | Encode/decode share at target concurrency and shape |
|
||||
| Is transfer economical? | Tensor/media bytes, serialization, connector latency and bandwidth |
|
||||
| Is there reuse? | Prompt/reference reuse rate, cacheability, fan-out |
|
||||
| Can it scale independently? | Queue saturation and resource utilization per stage |
|
||||
| Is the contract stable? | Shape, dtype, mask/tag/position metadata, ordering |
|
||||
|
||||
### Text encoder candidate
|
||||
|
||||
Define an explicit tensor contract containing hidden states, masks, tags or
|
||||
position IDs, dtype, shape, task/checkpoint revision, and request ID. Support
|
||||
disabling the local encoder so the DiT cannot accidentally recompute it. Test
|
||||
all tasks, prompt lengths, partial TP group membership, retries, backpressure,
|
||||
abort, connector cleanup, and independent replica scaling.
|
||||
|
||||
### VAE candidate
|
||||
|
||||
Separate VAE encode and decode contracts if both exist. Define latent shape,
|
||||
scale/shift, dtype/range, tiling/chunk order, output media ordering, and stage
|
||||
metadata. Measure latent/media transfer against local decode cost. Test tiling
|
||||
seams, RNG/determinism where relevant, chunk ordering, codec/postprocess,
|
||||
backpressure, abort, and worker failure.
|
||||
|
||||
Do not label VAE tiling, patch parallelism, compile, or offload as VAE
|
||||
disaggregation. Likewise, a separate text-encoder TP group is not a separate
|
||||
stage. If no reusable generic stage/connector exists in the target revision,
|
||||
produce a feasibility result or RFC and keep the capability `not tested` rather
|
||||
than inventing a model-local production architecture.
|
||||
|
||||
## Performance acceptance
|
||||
|
||||
For each accepted optimization, require:
|
||||
|
||||
- same fixed workload and correctness artifacts before/after;
|
||||
- actual fast-path execution proven by logs/trace/counters;
|
||||
- local metric plus E2E latency/throughput and memory;
|
||||
- at least three raw fixed-work repetitions after explicit warmup, or a
|
||||
declared arrival-load sample set large enough for every tail percentile;
|
||||
- named hardware, backend, topology, dtype, and task scope;
|
||||
- fallback/rejection tests for unsupported shapes/platforms;
|
||||
- no regression in abort/error cleanup or long-run memory trend.
|
||||
|
||||
Update the feature matrix narrowly. Do not convert a profile-driven direction
|
||||
into a generic “supported” feature merely because a model-specific PR exists.
|
||||
@@ -0,0 +1,464 @@
|
||||
# Production Serving and Validation
|
||||
|
||||
## Contents
|
||||
|
||||
1. [Execution-mode decision](#execution-mode-decision)
|
||||
2. [Request-level batching](#request-level-batching)
|
||||
3. [Step execution and continuous batching](#step-execution-and-continuous-batching)
|
||||
4. [Abort and cleanup](#abort-and-cleanup)
|
||||
5. [Mixed-RPS soak](#mixed-rps-soak)
|
||||
6. [Four CI tracks](#four-ci-tracks)
|
||||
7. [PR evidence checklist](#pr-evidence-checklist)
|
||||
|
||||
## Execution-mode decision
|
||||
|
||||
Request batching and step-wise execution solve different problems:
|
||||
|
||||
| Mode | Pipeline contract | Scheduling/abort boundary |
|
||||
|---|---|---|
|
||||
| Serial request | `forward(DiffusionRequestBatch)` for one request | Whole request; in-denoise preemption is unavailable |
|
||||
| Request batching | `supports_request_batch = True`; one forward consumes compatible independent requests | Whole batch/request |
|
||||
| Single-request step | Full `SupportsStepExecution`; `max_num_seqs=1` | Between denoise steps |
|
||||
| Step continuous batching | Full step contract plus heterogeneous batched-step evidence | Between denoise waves/steps |
|
||||
|
||||
Do not expose four stub methods as an opt-out. The current protocol is
|
||||
structural; if a native pipeline is not step-capable, do not define the four
|
||||
step methods. If it is capable, implement and test all four.
|
||||
|
||||
Treat structural support, a generic request-batch bridge, and model benefit as
|
||||
three separate claims. Before recommending step execution, name a hypothesis
|
||||
that request mode cannot satisfy, define a success threshold, and run a matched
|
||||
request-mode control. If throughput or latency regresses, retain step execution
|
||||
as functional/limited and keep request mode as the production recommendation.
|
||||
|
||||
Choose compatibility keys from every value that changes tensor shapes,
|
||||
semantics, hooks, or scheduler behavior: task, output shape/count, frames,
|
||||
steps/schedule, guidance/CFG, dtype/backend, LoRA/adapter, cache quality policy,
|
||||
and model-specific conditions. Reject or schedule separately rather than
|
||||
silently using the first request's settings.
|
||||
|
||||
## Request-level batching
|
||||
|
||||
A request-batched pipeline returns exactly one `DiffusionOutput` per logical
|
||||
request. This is model-owned pseudocode: `validate_batch_compatibility()` and
|
||||
`_generate_batch()` are placeholders to implement from the official contract,
|
||||
not shared APIs to import.
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.worker.request_batch import (
|
||||
DiffusionRequestBatch,
|
||||
split_diffusion_output_by_request,
|
||||
)
|
||||
|
||||
class MyPipeline(torch.nn.Module):
|
||||
supports_request_batch = True
|
||||
|
||||
def forward(self, req: DiffusionRequestBatch) -> list[DiffusionOutput]:
|
||||
validate_batch_compatibility(req.sampling_params_list, req.prompts)
|
||||
common = req.sampling_params_list[0]
|
||||
|
||||
text_prompts: list[str] | None = [
|
||||
prompt
|
||||
if isinstance(prompt, str)
|
||||
else (prompt.get("prompt") or "")
|
||||
for prompt in req.prompts
|
||||
]
|
||||
|
||||
prompt_fields = DiffusionRequestBatch.collate_prompt_field_map(
|
||||
req.prompts,
|
||||
{"prompt_embeds": None, "prompt_mask": None},
|
||||
)
|
||||
if prompt_fields["prompt_embeds"] is not None:
|
||||
text_prompts = None
|
||||
generators = req.collate_request_generators(
|
||||
common.num_outputs_per_prompt,
|
||||
default_generator=None,
|
||||
)
|
||||
output = self._generate_batch(
|
||||
prompt=text_prompts,
|
||||
prompt_fields=prompt_fields,
|
||||
generators=generators,
|
||||
sampling=common,
|
||||
)
|
||||
return split_diffusion_output_by_request(
|
||||
DiffusionOutput(output=output),
|
||||
req,
|
||||
num_outputs_per_prompt=common.num_outputs_per_prompt,
|
||||
)
|
||||
```
|
||||
|
||||
The prompt-field collator only handles tensor fields; extract raw text prompts
|
||||
and model-specific negative/alias fields separately, as current in-tree batched
|
||||
pipelines do. Set a raw prompt field to `None` when the matching precomputed
|
||||
embedding is supplied. The exact aliases and generator device depend on the
|
||||
model. Use current collators rather than manual `torch.cat` when they cover the
|
||||
contract; they reject mixed present/missing tensors and incompatible
|
||||
shape/dtype/device.
|
||||
|
||||
Before reading `sampling_params_list[0]`, validate every field that must be
|
||||
common. Tests must cover:
|
||||
|
||||
- two and N requests, independent seeds/generators, one output each;
|
||||
- mixed prompt lengths and packed/padded boundaries;
|
||||
- multiple outputs per prompt and correct result slicing;
|
||||
- compatible versus incompatible shapes, steps, guidance, LoRA, and quality;
|
||||
- one request failure without cross-request output or state corruption;
|
||||
- output ordering and request IDs under different completion times.
|
||||
|
||||
After the model advertises support, candidate serving is:
|
||||
|
||||
```bash
|
||||
vllm serve '<model>' --omni \
|
||||
--max-num-seqs 4 \
|
||||
--request-batch-max-wait-ms 20
|
||||
```
|
||||
|
||||
The wait window trades admission latency for batch formation. Benchmark at 0
|
||||
and candidate nonzero values under the target arrival process.
|
||||
|
||||
## Step execution and continuous batching
|
||||
|
||||
### Full protocol skeleton
|
||||
|
||||
Use the current shared state types and make every mutable value request-local:
|
||||
|
||||
The helpers beginning with `_` and `validate_step_batch()` below are
|
||||
model-owned placeholders; scheduler initialization is scheduler-specific.
|
||||
|
||||
```python
|
||||
import copy
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, ClassVar
|
||||
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.models.interface import SupportsStepExecution
|
||||
from vllm_omni.diffusion.worker.input_batch import InputBatch
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
|
||||
class MyPipeline(torch.nn.Module, SupportsStepExecution):
|
||||
supports_step_execution: ClassVar[bool] = True
|
||||
|
||||
def prepare_encode(
|
||||
self,
|
||||
state: StepRequestState,
|
||||
**kwargs: Any,
|
||||
) -> StepRequestState:
|
||||
prepared = self._prepare_request(state.prompt, state.sampling)
|
||||
state.prompt_embeds = prepared.prompt_embeds
|
||||
state.latents = prepared.latents
|
||||
state.timesteps = prepared.timesteps
|
||||
state.scheduler = copy.deepcopy(self.scheduler)
|
||||
state.scheduler.set_begin_index(0)
|
||||
state.step_index = 0
|
||||
state.extra["task_state"] = prepared.task_state
|
||||
return state
|
||||
|
||||
def denoise_step(
|
||||
self,
|
||||
input_batch: InputBatch,
|
||||
*,
|
||||
states: Sequence[StepRequestState] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor | None:
|
||||
validate_step_batch(input_batch, states)
|
||||
return self.transformer(**build_batched_denoise_inputs(input_batch))
|
||||
|
||||
def step_scheduler(
|
||||
self,
|
||||
state: StepRequestState,
|
||||
noise_pred: torch.Tensor,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
timestep = state.current_timestep
|
||||
if timestep is None:
|
||||
raise RuntimeError("step_scheduler called without a current timestep")
|
||||
state.latents = state.scheduler.step(
|
||||
noise_pred,
|
||||
timestep,
|
||||
state.latents,
|
||||
return_dict=False,
|
||||
)[0].to(state.latents.dtype)
|
||||
state.step_index += 1
|
||||
|
||||
def post_decode(
|
||||
self,
|
||||
state: StepRequestState,
|
||||
**kwargs: Any,
|
||||
) -> DiffusionOutput:
|
||||
return DiffusionOutput(output=self.decode(state.latents))
|
||||
```
|
||||
|
||||
This is a shape guide, not drop-in code: verify the target revision's
|
||||
`StepRequestState.extra`, scheduler signature, CFG handling, state fields, and
|
||||
decode contract. Follow a current native pipeline such as Qwen-Image for the
|
||||
exact revision.
|
||||
|
||||
Never use the shared `self.scheduler` to advance a request. Store scheduler,
|
||||
latents, embeddings, timesteps, masks, RNG/generator, task metadata, cache
|
||||
policy, adapter, and output assembly state on `StepRequestState` or its
|
||||
request-local extension. Pipeline fields may hold immutable weights/config and
|
||||
bounded immutable caches only.
|
||||
|
||||
### Qualification sequence
|
||||
|
||||
1. Compare request-mode and step-mode at `max_num_seqs=1` for every task using
|
||||
identical seed/schedule/shape and intermediate latents.
|
||||
2. Abort queued and in-flight requests at several step boundaries.
|
||||
3. Increase to two compatible requests; verify independent scheduler/RNG/state
|
||||
and output order.
|
||||
4. Increase to the target capacity and mix prompt lengths.
|
||||
5. Exercise mixed shape, frames, steps/schedule, task, guidance, and quality.
|
||||
Batch only compatible rows; otherwise split/reject predictably.
|
||||
6. Inject a failure in encode, denoise, scheduler, and decode independently.
|
||||
7. Run the production soak and fault matrix.
|
||||
|
||||
Candidate commands only after capability validation:
|
||||
|
||||
```bash
|
||||
# Conservative step execution and step-boundary abort.
|
||||
vllm serve '<model>' --omni \
|
||||
--step-execution \
|
||||
--max-num-seqs 1
|
||||
|
||||
# Experimental batched-step candidate after multi-request E2E.
|
||||
vllm serve '<model>' --omni \
|
||||
--step-execution \
|
||||
--max-num-seqs 4
|
||||
```
|
||||
|
||||
If an attention backend is only correct for single-request step mode, fail
|
||||
early or keep the recipe at `max_num_seqs=1`. Do not describe generic runtime
|
||||
support as model support.
|
||||
|
||||
## Abort and cleanup
|
||||
|
||||
### Required semantics
|
||||
|
||||
For queued and in-flight abort/disconnect:
|
||||
|
||||
- emit exactly one terminal aborted/error result;
|
||||
- do not run `post_decode` after abort unless a documented partial-output mode
|
||||
explicitly requires it;
|
||||
- do not increment `step_index` or apply a scheduler result after cancellation;
|
||||
- remove scheduler/state cache entries and request IDs;
|
||||
- release latents, prompt/reference tensors, cache hooks, DLO buffers, connector
|
||||
messages, uploaded/temp files, and output buffers;
|
||||
- keep other requests alive;
|
||||
- make cleanup idempotent when disconnect, timeout, and worker error race;
|
||||
- drop late results/exceptions for already-resolved or cancelled futures
|
||||
without terminating the shared result-pump thread;
|
||||
- complete the next valid request with its requested cache/offload policy.
|
||||
|
||||
Test cancellation before engine admission, while queued, after encode, during
|
||||
early/middle/final denoise steps, before decode, and while returning/persisting
|
||||
output. Include async job cancellation and client socket disconnect.
|
||||
|
||||
### Failure injection matrix
|
||||
|
||||
| Location | Injection | Assertions |
|
||||
|---|---|---|
|
||||
| Validation/upload | bad MIME, count/bytes overflow, truncated media | 400/413, no job/engine submit, no temp leak |
|
||||
| Load/start | missing fused shard, unsupported topology | startup fails, actionable error, no silent init |
|
||||
| Encode | exception/OOM/timeout | one error, request state/temp released |
|
||||
| Denoise | exception/OOM/worker exit | peers handled, hooks/buffers released, health behavior explicit |
|
||||
| Scheduler | invalid state/shape | no latent/step partial commit |
|
||||
| Decode/post | exception/codec error | no corrupt successful result, partial output removed |
|
||||
| Result/IPC | late result after cancel, large split tensor, orphan SHM | pump remains alive, one terminal result, exact ownership cleanup |
|
||||
| Connector/disagg | timeout/retry/duplicate | idempotency, backpressure, no orphan payload |
|
||||
|
||||
After each injection, check server health and issue a known-good request. If
|
||||
the process must restart, document and test the restart policy; do not claim
|
||||
transparent recovery.
|
||||
|
||||
### Large-output transport
|
||||
|
||||
Treat worker-to-engine IPC and engine-to-client transport as separate hops.
|
||||
For batched video, prove that each logical output does not accidentally retain
|
||||
or serialize the complete batch storage. Exercise the target revision's shared
|
||||
memory or artifact-handle path above and below its threshold, including
|
||||
non-contiguous views, multiple outputs, consumer failure, timeout, abort, and
|
||||
worker exit. Verify every shared segment/file is released exactly once.
|
||||
|
||||
Freeze the media contract at each hop: dtype, value range, shape/layout,
|
||||
contiguity, request ownership, encoding/color model, payload bytes, and pending
|
||||
consumers. Device-side float-to-uint8 preparation, D2H/IPC, planar/interleaved
|
||||
conversion, CPU encode/mux, and delivery off the asyncio event loop are
|
||||
different optimizations. Validate their route selection and fallbacks
|
||||
independently. Byte-identical HTTP MP4 output does not prove that a raw offline
|
||||
tensor kept its dtype or range.
|
||||
|
||||
For long-window output, validate segment ordering and audio/video continuity,
|
||||
then exercise payloads around serializer, shared-memory, HTTP, and codec size
|
||||
limits. Chunking must reassemble exactly, reject malformed/missing chunks, and
|
||||
release partial state after abort, timeout, or peer failure.
|
||||
|
||||
Qualify chunk production, transport, and public streaming as separate layers.
|
||||
A VAE callback must define owned representation, frame offsets, final marker,
|
||||
rank ownership, source/config gate, and collective-safe failure, but it does not
|
||||
prove SHM/ZMQ delivery, bounded backpressure, cancellation, codec reuse, or
|
||||
client streaming. Measure time to first owned chunk, first encoded fragment,
|
||||
steady cadence, mux finalization, and complete client artifact independently.
|
||||
|
||||
Measure remote MP4 encoding, serialization/copies, network transfer, and client
|
||||
materialization separately from denoise/decode. An inference-kernel speedup does
|
||||
not establish serving throughput if output transport dominates or leaks.
|
||||
|
||||
## Mixed-RPS soak
|
||||
|
||||
Use a realistic request mix and arrival distribution. Include small/large
|
||||
shapes, short/long prompts, all advertised tasks, cache tiers, and a small rate
|
||||
of abort/disconnect/error injections. Run phases below, near, and above measured
|
||||
saturation so queue/backpressure behavior is visible.
|
||||
|
||||
Example serving benchmark phase:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://127.0.0.1:8091 \
|
||||
--endpoint /v1/videos \
|
||||
--model '<model>' \
|
||||
--dataset trace \
|
||||
--dataset-path /path/to/versioned-production-trace.jsonl \
|
||||
--task t2v \
|
||||
--num-prompts 500 \
|
||||
--request-rate 0.5 \
|
||||
--max-concurrency 8 \
|
||||
--warmup-requests 8 \
|
||||
--warmup-concurrency 8 \
|
||||
--output-file soak-near-saturation.json
|
||||
```
|
||||
|
||||
Use a task-appropriate trace/media dataset. Run the same phase at rates selected
|
||||
from an initial saturation curve; a literal `0.5` is illustrative only.
|
||||
|
||||
On supported GPU hosts, wrap a test or soak with the repository monitor:
|
||||
|
||||
```bash
|
||||
bash tests/dfx/stability/scripts/resource_monitor.sh run --backend gpu -- \
|
||||
pytest -s -v tests/dfx/stability/scripts/test_stability_<model>.py -m slow
|
||||
```
|
||||
|
||||
Check the target revision's monitor backends. Do not claim CPU/NPU monitoring
|
||||
from a reserved CLI choice unless collection is implemented and verified.
|
||||
|
||||
Report:
|
||||
|
||||
- offered and achieved RPS, success/error/abort rates;
|
||||
- p50/p95/p99 and queue/service/stage times, with the successful sample count
|
||||
and arrival process used for each percentile;
|
||||
- throughput per device and batch/wave utilization;
|
||||
- per-rank HBM and process-tree PSS trend/slope;
|
||||
- cache size/hit rate, temp-file/disk growth, open FDs/handles;
|
||||
- outstanding shared-memory/artifact handles and encoded/network bytes;
|
||||
- worker/engine health, restarts, OOM/timeouts, backpressure response;
|
||||
- total duration, request count/mix, raw logs/JSON/monitor artifact hashes.
|
||||
|
||||
A short burst or repeated serial curl loop is not a stability test.
|
||||
|
||||
## Four CI tracks
|
||||
|
||||
Keep the tracks independent so a failure answers a specific question. These
|
||||
four readiness tracks are not the repository's numbered CI levels. Map them
|
||||
onto the current L1-L5 taxonomy:
|
||||
|
||||
| Readiness track | Repository CI placement |
|
||||
|---|---|
|
||||
| Function | L1 unit/logic plus L2 basic online/offline E2E; extend real-model scenarios in L3/L4 |
|
||||
| Accuracy | L3 real-model gates for important cases; deeper and broader nightly coverage in L4 |
|
||||
| Performance | Important thresholds in L3 where they fit the time budget; full baselines/regression jobs in nightly L4 |
|
||||
| Reliability | Cheap rejection/recovery assertions can run earlier, but long stability, fault injection, and reliability suites belong to weekly L5 |
|
||||
|
||||
Distributed topology is a test case within the appropriate level, not a
|
||||
replacement definition for a level.
|
||||
|
||||
### 1. Function CI
|
||||
|
||||
Cover:
|
||||
|
||||
- strict unquantized and advertised quantized loading;
|
||||
- fused source-shard completeness and prefix routing;
|
||||
- official task positive boundaries and neighboring negative cases;
|
||||
- offline, sync, async normalization/parity;
|
||||
- single-device reference plus every advertised topology smoke;
|
||||
- output type/count/order and request ID isolation;
|
||||
- feature rejection paths that must fail early.
|
||||
|
||||
### 2. Accuracy CI
|
||||
|
||||
Pin official implementation/checkpoint/vLLM-Omni revisions, prompts and media
|
||||
hashes, seed, scheduler/sigma/steps, dimensions/frames/FPS, guidance, dtype,
|
||||
backend, and topology. Compare at suitable levels:
|
||||
|
||||
- component outputs or selected blocks;
|
||||
- one or more intermediate denoise latents;
|
||||
- final image/video/audio artifact metrics;
|
||||
- temporal consistency and audio/video synchronization where relevant.
|
||||
|
||||
Set tolerance per modality/hardware/backend with rationale. Keep failure
|
||||
artifacts and metadata. A single visual inspection does not define CI.
|
||||
|
||||
For quantization trajectory tooling, inspect the target revision's current
|
||||
`vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` and
|
||||
reuse it when compatible instead of creating an ad hoc comparator.
|
||||
|
||||
### 3. Performance CI
|
||||
|
||||
Pin workload and best deployment. Include a resident recommended row and a
|
||||
small-HBM/DLO row when both are production targets. Declare:
|
||||
|
||||
- one explicit warmup exclusion and at least three measured fixed-work
|
||||
repetitions;
|
||||
- raw result JSON and environment identity;
|
||||
- median/mean plus range for small fixed-work A/Bs; p50/p95/p99 only for an
|
||||
arrival-load run with enough samples, plus throughput/RPS, stage time,
|
||||
HBM/PSS;
|
||||
- regression metric, threshold, baseline owner, and update procedure;
|
||||
- correctness/accuracy check linked to the same configuration.
|
||||
|
||||
Use the current harness config style, for example:
|
||||
|
||||
```bash
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py \
|
||||
--test-config-file tests/dfx/perf/tests/test_<model>_vllm_omni.json
|
||||
```
|
||||
|
||||
Verify parser options on the target revision. Do not copy a JSON schema from a
|
||||
different model without matching its endpoint, task, dataset, and deploy path.
|
||||
|
||||
### 4. Reliability CI
|
||||
|
||||
Run a bounded CI soak plus a longer scheduled soak. Cover:
|
||||
|
||||
- below/near/above-saturation arrival phases;
|
||||
- queued/in-flight abort and disconnect;
|
||||
- OOM/timeout and worker/process failure appropriate to the platform;
|
||||
- DLO/cache/connector cleanup;
|
||||
- health, fast-fail/backpressure, restart/recovery policy;
|
||||
- memory/temp/FD trend and successful known-good request after each fault.
|
||||
|
||||
Keep reliability thresholds independent from latency thresholds. A fast server
|
||||
that leaks memory or hangs after cancellation fails production readiness.
|
||||
|
||||
## PR evidence checklist
|
||||
|
||||
The PR description must include or link:
|
||||
|
||||
- official source/checkpoint revisions and API/limit matrix;
|
||||
- scoped feature compatibility matrix with the four allowed states;
|
||||
- implementation notes for FP8, DLO, request cache, attention/fusion, batching;
|
||||
- exact environment, deploy/serve, every task curl, benchmark and soak commands;
|
||||
- fixed accuracy inputs and artifacts;
|
||||
- raw before/after performance runs and memory traces;
|
||||
- output-contract manifests and stage accounting from device preparation
|
||||
through the complete client artifact;
|
||||
- abort/disconnect/fault results and cleanup evidence;
|
||||
- Function/Accuracy/Performance/Reliability CI files and ownership;
|
||||
- limitations/TODOs that remain `not tested` rather than implied support.
|
||||
|
||||
Separate the Day-0 milestone from production readiness. If the PR only closes
|
||||
some gates, title and description should say which scoped gates it closes.
|
||||
@@ -37,6 +37,10 @@ self-check.
|
||||
- [`precheck-pr`](precheck-pr/SKILL.md): self-checks a branch before creating a
|
||||
PR by validating title format, dead code, simplification opportunities,
|
||||
accuracy and performance claims, and merge readiness
|
||||
- [`production-add-diffusion-model`](production-add-diffusion-model/SKILL.md):
|
||||
takes a working Day-0 diffusion integration through API/limit parity,
|
||||
feature-combination evidence, quantization/offload, performance, production
|
||||
serving, hardware recipes, and CI readiness
|
||||
- [`quantization`](quantization/SKILL.md): guides quantization method selection,
|
||||
model integration, checkpoint loading, and quality/performance validation
|
||||
for vLLM-Omni
|
||||
|
||||
Reference in new issue
Block a user