cache-dit-model-integration

## ROUTE FIRST — what is the user actually asking for? Pick one route before reading anything else; the rest of this file is only route 2. 1. **Benchmarking / acceptance of an ALREADY-supported model** — the model runs (`cdit.generate` has an example, cache/CP/TP already work) and the user wants speed + PSNR/SSIM numbers: "跑一遍全量测试矩阵", "4 卡验收", "量化掉多少精度", "和 bf16 比加速比" → read **`./references/perf_matrix.md`**. It owns the row list, command templates, the layout mapping / preflight, and Tier A (1/2 GPUs, default) vs the on-request Tier B (4 GPUs) / Tier C (8 GPUs) extensions (hybrid CP+TP included). **Do NOT open Tier B/C unless the user explicitly asks to extend to 4/8 cards.** The GATE CHECK, the Hard rules and the §0.6 checklist below do not apply to this route. 量化类型与配置语义（fp8/int8/int4/nvfp4、svdq PTQ vs DQ、precision plan、regional、 `per_tensor_fallback`、observer）以 `cache-dit-quantization` 为准；本 route 只负责跑矩阵、 出加速比与 PSNR/SSIM，不负责解释量化后端行为。 2. **Integrating a NEW model** — no `BlockAdapter`, no CLI example, new pipeline/transformer code → continue with the GATE CHECK and §0 below. 3. **Second-pass optimization of a model that is ALREADY in diffusers** — the model runs, cache/CP/TP work, and the task is "wrap it as `_<Model>Pipeline` and add fused ops / kernels without vendoring" (e.g. `_QwenImage21`) → read **`./references/custom_models.md` §15** (and §1.3 for the entry-point decision). It owns the subclass-in-place runbook: `from_pretrained` class swap, the `_NAME` registry checklist, seam selection (why a module-level QKV helper beats a processor subclass when CP replaces `__call__`), the fast-math numerical contract, the "fold per-sample constants first" rule, the compile traps, and the validation protocol. The GATE CHECK below does **not** apply — no new model is being integrated. 4. **A measured speedup looks wrong, or choosing which optimizations can pay off before benchmarking** → read `./references/perf_analysis.md` (FLOPs budget: attention vs GEMM share). 5. **The model is audio-only** (TTS / text-to-audio: it returns a waveform, not frames) → read `./references/audio.md` **before** the Cache chapter. It covers the latent layout and duration semantics, the CP obstacles specific to audio (dataclass return, per-rank RoPE, replicated text, duration mask), the wav output path in the CLI, and the audio metrics + their noise floor — the image PSNR/SSIM gates below do **not** transfer. The rest of this file (GATE CHECK, phases) still applies — only route 1 skips them. ## Related skills（跨 skill 路由） 模型接完只是第一步；下面这些场景不要在本文里找答案： | 触发场景 | 去读 | |---|---| | 量化相关的模型测试与精度归因：`--quantize-type` 选哪个、precision plan / regional / `per_tensor_fallback`、svdq 选 PTQ 还是 DQ、observer 怎么定位掉精度的层、"量化掉多少精度"该归因到哪条路径 | `cache-dit-quantization` | | 模型里要新增或改动 CUDA 算子、算子注册与构建、从 nunchaku 等上游移植 kernel | `cache-dit-custom-kernel`（迁移方法论在其 `references/operator-migration.md`） | | 发版本、构建 cu13 wheel、索引与安装验收 | `cache-dit-build-release` | | 只要 CUDA/PTX 细节、CUTLASS/CuTe C++ 模板、CuTe DSL Python、NCU 驱动的 kernel 优化 | `cuda-cpp-kernel` / `cutlass-cpp-kernel` / `cute-dsl-kernel` / `cuda-auto-tune` | 边界：本 skill 管"模型怎么接、怎么验收"；量化配置与后端语义归 `cache-dit-quantization`，算子实现归 `cache-dit-custom-kernel`，发布流程归 `cache-dit-build-release`。 ## GATE CHECK > Scope: route 2 (integrating a new model) only. Route 1 (benchmarking an already-supported model) goes > to `./references/perf_matrix.md`, and route 3 (second-pass optimization of a diffusers model, `_NAME`) > goes to `./references/custom_models.md` §15 — neither needs transformer source paths or an > integration plan. Before writing any code, confirm the following: ``` STOP — Has the user provided BOTH of the following? 1. A local model path (e.g., /workspace/dev/vipdev/hf_models/Krea-2-Turbo) (if the weights are not downloaded yet, see `./references/model_download.md`) 2. The model's pipeline/transformer code info: - Pipeline class name (e.g., Krea2Pipeline, or third-party like BooguImagePipeline) - Transformer class name (e.g., Krea2Transformer2DModel) - File paths to the pipeline and transformer source code (diffusers or third-party) NO → **MUST ask the user to specify these before proceeding.** Do NOT guess, search blindly, or assume defaults. YES → Continue. STOP — Does this task either (a) add a model that is NOT in the diffusers library (ships its own `pipeline_<model>.py` / transformer code, `model_index.json` points at an unimportable module, or the code lives only in a third-party package / `trust_remote_code`), or (b) move a diffusers model's source code INTO cache-dit so we own the model instead of patching diffusers (the path for custom-kernel / deep-optimization work)? YES → **Read `./references/custom_models.md` FIRST, before the Cache chapter.** (a) Decide the ownership model (vendor into `src/cache_dit/models/pipelines/` vs. depend on the third-party package), then keep every dispatch contract it lists. (b) Follow the migration runbook in that file (§3): audit the diffusers-detection flips (`check_forward_pattern`, `_relaxed_assert`, leftover `PatchFunctor` silently skipped), keep class / attribute / state-dict names, fold the patches into the owned code and delete them. Both: vendored `DiffusionPipeline` / `ModelMixin`, diffusers VAE and transformers text encoder, top-level block `ModuleList`, one block call per loop iteration, `hidden_states` first, `_repeated_blocks`, and skip/step-varying state moved INTO the block. Then continue with the phases below. NO → Continue. STOP — Have you identified the new model's transformer architecture? NO → Read the model's source code (diffusers, or the vendored / third-party copy). Identify: - The ModuleList name(s) containing transformer blocks (e.g., transformer_blocks) - The block forward() signature (inputs and outputs) YES → Proceed to Cache (`./references/cache.md`) and CLI (`./references/cli.md`) in parallel. STOP — Is the model only available as a diffusers ModularPipeline (e.g. FLUX.3)? (pipe is NOT a DiffusionPipeline / nn.Module; its steps live in ModularPipelineBlocks) YES → Read `./references/modular_pipeline.md` FIRST (core plumbing, which block list to cache, CLI template, CP/TP status), then continue with the phases below. STOP — Is the model audio-only (TTS / text-to-audio: the pipeline returns a waveform, no frames)? YES → Read `./references/audio.md` BEFORE the Cache chapter. It owns the request shape (`duration` -> latent frames, per-call seq len), the CP obstacles (dataclass return, per-rank RoPE table, replicated text cross-attn, all-True duration mask), the T2A wav path in the CLI, and the audio metrics with their noise floor. The image PSNR/SSIM thresholds elsewhere in this skill do NOT apply to a chaotic TTS trajectory. ``` **Hard rules:** > Route 2 only. A route-1 acceptance run is governed by `./references/perf_matrix.md` §2/§7 plus the > gates in `./references/testing.md` §7.6 — it needs neither model source paths nor an approved plan. - **⚠️ MANDATORY: Local model path and code info.** If the user has not explicitly provided (a) the local model path and (b) the pipeline/transformer class names with source file paths (diffusers or third-party), you **MUST** ask the user to specify them via ``vscode_askQuestions`` before any code changes. Do NOT search the codebase or assume default paths — the user knows their setup best. - **⚠️ MANDATORY: `./references/custom_models.md` for BOTH ownership cases.** Read it **before** writing the plan whenever either holds: **(a)** the model is **not part of the diffusers library** (own `pipeline_<model>.py` / transformer code, a `model_index.json` module diffusers cannot import, or code from a third-party package); **(b)** the task **moves a diffusers model's source into cache-dit** (owning the model instead of patching it — the required path when an optimization touches attention/kernel math or the block loop itself, where no hook point is left). Follow its ownership decision (vendor / absorb / depend), its contract checklist and, for (b), its migration runbook: the diffusers-detection branches flip once `__module__` is `cache_dit.*` — `check_forward_pattern` defaults to False, `_relaxed_assert` stops checking, and a leftover `PatchFunctor` is **silently skipped** — while class, block-attribute and state-dict names must stay unchanged. Reference integrations: MOSS-Image (vendored, case a); Boogu-Image (third-party package, case a, adapter only). No case-(b) precedent exists yet — treat its runbook as authoritative. - **⚠️ MANDATORY: Plan before code.** Before writing ANY implementation code, you **MUST**: (1) thoroughly analyze the model's pipeline and transformer source code, (2) create a detailed integration plan following this skill's workflow and read the relevant references (Cache → CP → TP → TE-P → VAE-P → CLI → Testing), (3) present the plan to the user for review and approval. **Do NOT start implementing until the user explicitly approves the plan.** This prevents wasted effort from incorrect assumptions about the model architecture. - ALWAYS set up local model paths via environment variables BEFORE testing — downloading from HuggingFace Hub is extremely slow. Download flows (ModelScope first, HF fallback): [`./references/model_download.md`](./references/model_download.md). - ALWAYS compute BOTH PSNR and SSIM when verifying correctness — PSNR alone cannot detect image corruption (garbled output). For **audio** models the equivalent pair is a spectrogram metric (`mel_psnr` / `mel_ssim` via `cache_dit.metrics`) plus `si_sdr` on the waveform, and the gate must be calibrated against a **same-configuration rerun** (a diffusion TTS run is chaotic; measured floor: mel PSNR ≈ 21 dB / SSIM ≈ 0.75 / SI-SDR ≈ −5 dB). See [`./references/audio.md`](./references/audio.md) §6. - For Python-only changes, `pip install -e "." --no-build-isolation` is sufficient; SVDQuant C++ compilation is NOT required. - **Do NOT alter core dependency versions** (torch, torchvision, transformers, diffusers, cache-dit, triton) in the `cdit` conda environment. Other dependencies may be installed only if they do not conflict with these. - **Do NOT modify any code in the diffusers library.** If a model requires patches (e.g., monkey-patching `forward()`, attention processors, etc.), write all patch code inside the cache-dit repository. Diffusers is a third-party dependency and must not be altered. - **⚠️ MANDATORY: Use cache-dit's attention dispatcher.** Every new model-specific attention processor or CP attention call in cache-dit must use `_dispatch_attention_fn` from `cache_dit.attention` (normally `from ...attention import _dispatch_attention_fn`), not diffusers' `dispatch_attention_fn` or a backend function directly. When the processor already owns CP communication, pass `parallel_config=None` to prevent a second collective. Read [`./references/attention.md`](./references/attention.md) before writing attention code. - **All examples in this skill are references, not templates to copy.** Every model has unique architecture details (block signatures, tensor layouts, shared vs per-block modulation, attention mask requirements, etc.). Before following any example, first analyze whether the referenced model's architecture is actually comparable to the target model. Blindly copying an example that was designed for a different architecture will produce incorrect or broken code. - ControlNet parallelism is a special case and is NOT covered in this guide. - **⚠️ GQA attention backend pitfall**: When a model uses GQA (e.g., ``num_heads=48, num_kv_heads=12``), the ``dispatch_attention_fn(..., enable_gqa=True)`` path may cause PyTorch SDPA to fall back to a **slow backend** (``math`` or an inefficient ``mem_efficient`` kernel) because flash-attention / cuDNN SDPA backends have limited GQA support. **Always benchmark ``enable_gqa=True`` vs. manually repeating K/V heads to match Q heads and passing ``enable_gqa=False``** (MHA path). On NVIDIA L20, the MHA repeat gave a **~2.2× single-GPU speedup** for Krea-2-Turbo (48 Q / 12 KV heads, 128 head_dim, 4608 seq). This is not CP-specific — any model with GQA should evaluate whether the repeat→MHA path is faster. If confirmed, apply the repeat unconditionally in the attention processor patch, not just in the CP path. Document the finding in the planner's docstring as well (see ``krea2.py`` for an example). --- ## 0. Model Integration Practical Navigation > Start here. This map orients you before you dive into any chapter. The integration order is **Cache → CP → TP → TE-P → VAE-P → CLI → Testing**; each feature is verified against a single-GPU baseline before moving on. ### 0.1 End-to-End Workflow > The integration order is **Cache → CP → TP → TE-P → VAE-P**. Each feature is tested against a single-GPU baseline before moving to the next. Use the detailed checklist in §0.6 for planning and tracking. ```mermaid flowchart TD G{"GATE CHECK: local model path AND pipeline/transformer source both provided?"} G -->|No| GA["STOP - ask the user via vscode_askQuestions"] G -->|Yes| CM{"Own the model source? (absent from diffusers, or migrating one in)"} CM -->|"Yes: read custom_models.md (vendor / depend / migrate)"| A["Read and analyze pipeline __call__ and transformer.forward()"] CM -->|No| A A --> PL["Draft an integration plan with §0.6 TODO list, get user approval"] PL --> C1["Cache: BlockAdapter + ForwardPattern - cache.md"] C1 -->|"Parallel: CLI"| CLI["Register generate CLI + env mapping - cli.md"] C1 --> CQ{"Block loop: keyword/positional mismatch, or extra ops in loop body?"} CQ -->|Yes| CF["Write a PatchFunctor - cache.md"] CQ -->|No| CA["Pick ForwardPattern, build BlockAdapter - cache.md"] CF --> CA CA --> TST1["TEST: baseline + cache. PSNR > 30, SSIM > 0.90"] TST1 -->|FAIL| CA TST1 -->|PASS| C2["CP: hook-based first, hybrid second - cp.md"] C2 --> TST2["TEST: CP Ulysses vs baseline. PSNR > 35, SSIM > 0.90"] TST2 -->|FAIL| C2 TST2 -->|PASS| T3["TP: shard_div_attr, GQA, DTensor-unsafe ops - tp.md"] T3 --> TST3["TEST: TP vs baseline. PSNR > 35, SSIM > 0.90"] TST3 -->|FAIL| T3 TST3 -->|PASS| T4["TE-P: check support, add planner if needed - tep.md"] T4 --> TST4["TEST: TP+TE-P vs baseline. PSNR > 35, SSIM > 0.90"] TST4 -->|FAIL| T4 TST4 -->|PASS| T5["VAE-P: check support, add planner if needed - vaep.md"] T5 --> TST5["TEST: VAE-P vs baseline. PSNR > 35, SSIM > 0.90"] TST5 -->|FAIL| T5 TST5 -->|PASS| T7["Final: hybrid CP+TP, results table, visual check - testing.md"] ``` ### 0.2 Context Parallelism Decision Chart (the highest-risk step) ```mermaid flowchart TD S["Start CP for a new model"] --> MK{"Model uses an attention_mask? Check inside joint_attention_kwargs / attention_kwargs too"} MK -->|"Yes"| MKN["Flag: the mask must be permuted/reordered under CP later - cp.md"] MK -->|"No"| HK MKN --> HK{"Can every sequence tensor be split at a module-call boundary?"} HK -->|"Yes: plain tensors or 1-D lists"| HT{"What shape is transformer.forward()?"} HK -->|"No: a FEW tensors are methods / nested types / local vars"| HB{"Can the gap be closed by a MINIMAL patch (one method, not entire forward)?"} HB -->|"Yes: patch get_rotary_pos_embed / pos_embed / one sub-module"| HY["Hybrid CP - cp.md: hook plan + minimal method patch"] HB -->|"No: tensors are deep local intermediates unreachable by any hook"| STOP["STOP: pure patch-based is REJECTED. Re-analyze the model structure."] HT -->|"Simple preprocess to loop to postprocess"| RT["Root split with key '' - cp.md"] HT -->|"One sub-module output stays full-length: caption_projection / rope / pos_embed"| OS["Add output-split hook split_output=True - cp.md"] HT -->|"Uniform loop, processor does the all-to-all"| B0["Single-point split at transformer_blocks.0 - cp.md"] HT -->|"Inter-block op needs the FULL sequence"| SM["Sub-module-level hooks - cp.md"] RT --> UA OS --> UA B0 --> UA SM --> UA HY --> UA UA{"num_kv_heads or sequence length not divisible by cp_size?"} UA -->|"Yes"| UAA["Enable UAA: --ulysses-anything - cp.md"] UA -->|"No"| VF UAA --> VF["Verify: top-left corner clean, PSNR > 35, SSIM > 0.90"] ``` ### 0.3 Tensor Parallelism Decision Chart (second-highest-risk step) ```mermaid flowchart TD