From dff28e686711add0e929ce0c5660b09ef26c8c88 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Thu, 3 Sep 2026 20:09:28 +0800 Subject: [PATCH] [Doc] Add production diffusion model skill (#6097) Signed-off-by: David Chen <530634352@qq.com> --- .../production-add-diffusion-model/SKILL.md | 458 +++++++++++++++++ .../agents/openai.yaml | 4 + .../references/api-and-recipes.md | 315 ++++++++++++ .../references/feature-patterns.md | 462 +++++++++++++++++ .../references/performance-patterns.md | 481 ++++++++++++++++++ .../references/production-validation.md | 464 +++++++++++++++++ .claude/skills/readme.md | 4 + 7 files changed, 2188 insertions(+) create mode 100644 .claude/skills/production-add-diffusion-model/SKILL.md create mode 100644 .claude/skills/production-add-diffusion-model/agents/openai.yaml create mode 100644 .claude/skills/production-add-diffusion-model/references/api-and-recipes.md create mode 100644 .claude/skills/production-add-diffusion-model/references/feature-patterns.md create mode 100644 .claude/skills/production-add-diffusion-model/references/performance-patterns.md create mode 100644 .claude/skills/production-add-diffusion-model/references/production-validation.md diff --git a/.claude/skills/production-add-diffusion-model/SKILL.md b/.claude/skills/production-add-diffusion-model/SKILL.md new file mode 100644 index 000000000..a54e72823 --- /dev/null +++ b/.claude/skills/production-add-diffusion-model/SKILL.md @@ -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. diff --git a/.claude/skills/production-add-diffusion-model/agents/openai.yaml b/.claude/skills/production-add-diffusion-model/agents/openai.yaml new file mode 100644 index 000000000..38fc52586 --- /dev/null +++ b/.claude/skills/production-add-diffusion-model/agents/openai.yaml @@ -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." diff --git a/.claude/skills/production-add-diffusion-model/references/api-and-recipes.md b/.claude/skills/production-add-diffusion-model/references/api-and-recipes.md new file mode 100644 index 000000000..b73ceb165 --- /dev/null +++ b/.claude/skills/production-add-diffusion-model/references/api-and-recipes.md @@ -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": "", + "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=' \ + -F '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=' \ + -F 'prompt=' \ + -F 'size=1280x720' \ + -F 'num_frames=81' \ + -F 'fps=16' \ + -F 'num_inference_steps=30' \ + -F 'seed=42' \ + -F 'extra_params={"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=' \ + -F 'prompt=' \ + -F 'num_frames=81' \ + -F 'fps=16' \ + -F 'num_inference_steps=30' \ + -F 'seed=42' \ + -F 'extra_params={"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=' \ + -F 'prompt=' \ + -F 'num_frames=81' \ + -F 'num_inference_steps=30' \ + -F 'seed=43' \ + -F 'extra_params={"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": "", + "input": "", + "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 '' \ + --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\":\"\"}"}' \ + --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//`, `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. diff --git a/.claude/skills/production-add-diffusion-model/references/feature-patterns.md b/.claude/skills/production-add-diffusion-model/references/feature-patterns.md new file mode 100644 index 000000000..1b4dabb5a --- /dev/null +++ b/.claude/skills/production-add-diffusion-model/references/feature-patterns.md @@ -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 | +|---|---|---|---|---|---|---|---|---|---|---| +| `//` | serial request / 1 | SDPA, padded | BF16 | none | resident | 1 device | `` BF16 | raw float / local | `not tested` | — | +| `//` | step / N | ``, packed | BF16 | none | resident | `` | `` BF16 | encoded / subprocess | `not tested` | — | +| `//` | step / N | ``, packed | FP8 | high | DLO no-AG | `` | `` | 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 '' --omni --quantization fp8 --max-num-seqs 1 +``` + +For scoped routing: + +```bash +vllm serve '' --omni \ + --diffusion-quantization-config \ + '{"":{"method":"fp8","ignored_layers":[""]},"":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 '' --omni \ + --enable-distributed-layerwise-offload \ + --usp 4 \ + --max-num-seqs 1 + +# TP + DLO rank-local candidate; not a support statement. +vllm serve '' --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: +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 '' --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=' \ + -F 'prompt=' \ + -F 'quality=lossless' \ + -F 'num_inference_steps=' \ + -o lossless.mp4 + +curl --fail-with-body -sS -X POST http://127.0.0.1:8091/v1/videos/sync \ + -F 'model=' \ + -F 'prompt=' \ + -F 'quality=high' \ + -F 'num_inference_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. diff --git a/.claude/skills/production-add-diffusion-model/references/performance-patterns.md b/.claude/skills/production-add-diffusion-model/references/performance-patterns.md new file mode 100644 index 000000000..6840d2d5f --- /dev/null +++ b/.claude/skills/production-add-diffusion-model/references/performance-patterns.md @@ -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. diff --git a/.claude/skills/production-add-diffusion-model/references/production-validation.md b/.claude/skills/production-add-diffusion-model/references/production-validation.md new file mode 100644 index 000000000..eed568415 --- /dev/null +++ b/.claude/skills/production-add-diffusion-model/references/production-validation.md @@ -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 '' --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 '' --omni \ + --step-execution \ + --max-num-seqs 1 + +# Experimental batched-step candidate after multi-request E2E. +vllm serve '' --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 '' \ + --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_.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__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. diff --git a/.claude/skills/readme.md b/.claude/skills/readme.md index 95edc5021..54008b82c 100644 --- a/.claude/skills/readme.md +++ b/.claude/skills/readme.md @@ -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