cache-dit-model-integration

## GATE CHECK 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) 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 — Have you identified the new model's transformer architecture? NO → Read the model's diffusers source code. 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. ``` **Hard rules:** - **⚠️ 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: 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. - ALWAYS compute BOTH PSNR and SSIM when verifying correctness — PSNR alone cannot detect image corruption (garbled output). - 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. - **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| A["Read and analyze pipeline __call__ and transformer.forward()"] 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 S["Start TP for a new model"] --> DT{"Fused CUDA kernels? (flash_attn SwiGLU, triton RMSNorm)"} DT -->|Yes| DP["Patch to PyTorch equivalents FIRST - tp.md"] DT -->|No| FW DP --> FW{"Fused weight matrix? (QKV packed, out+down shared)"} FW -->|Yes| RW["Un-fuse / rearrange weights - tp.md"] FW -->|No| HEADS RW --> HEADS{"num_attention_heads % tp_size == 0?"} HEADS -->|No| E1["ERROR: MHA heads must be divisible. Check tp_size choice."] HEADS -->|Yes| GQA{"GQA? num_kv_heads < num_heads?"} GQA -->|"No (plain MHA)"| MHA["Standard: ColwiseParallel(to_q/k/v), RowwiseParallel(to_out.0), shard_div_attr 'heads'"] GQA -->|Yes| KVDIV{"num_kv_heads % tp_size == 0?"} KVDIV -->|Yes| STDGQA["Same MHA plan: K/V also shardable via ColwiseParallel"] KVDIV -->|No| BENCHQ{"Long sequence? (> ~10K tokens)"} BENCHQ -->|"Yes or unknown"| FFNO["Prefer FFN-only TP - tp.md GQA caveat: skip attention sharding entirely."] BENCHQ -->|"No, short seq"| REPL["Replicate - tp.md: ColwiseParallel(output_layouts=Replicate()). Skip shard_div_attr."] MHA --> FFN["Build FFN plan: ColwiseParallel(gate_proj/up_proj), RowwiseParallel(down_proj)"] STDGQA --> FFN FFNO --> FFN REPL --> BENCH2["⚠️ Benchmark FFN-only vs Replicate before shipping - tp.md"] BENCH2 --> FFN FFN --> MOD{"AdaLayerNormZero / norm.linear modulation?"} MOD -->|Yes| MODP["ColwiseParallel(output_layouts=Replicate()) for norm.linear"] MOD -->|No| VERIFY MODP --> VERIFY["Verify: PSNR > 35 dB, SSIM > 0.90. Visual inspection."] ``` ### 0.4 Decision Index — "If you're dealing with X, read this" | Situation | Section | |---|---| | Choosing the block I/O pattern | `./references/cache.md` and the ForwardPattern table below | | Block loop calls blocks with keyword args, or has extra ops inside the loop | `./references/cache.md` PatchFunctor pitfalls | | More complex structural patch (per-block forward, block-id injection, block-list merge) | `./references/cache.md` advanced PatchFunctor cases | | Third-party (non-diffusers) model | `./references/cache.md` third-party model section | | CP: choosing hook-based vs hybrid | `./references/cp.md` plus the CP chart above | | CP: a projection / rope / pos_embed output stays full-length | `./references/cp.md` output-split hook | | CP: uniform loop + attention-processor all-to-all | `./references/cp.md` single-point split | | CP: an inter-block op needs the full sequence | `./references/cp.md` sub-module hooks | | CP: a tensor is a method return or nested type that hooks cannot split | `./references/cp.md` hybrid CP | | CP: head count or sequence length not divisible by cp_size | `./references/cp.md` UAA | | CP: the model has an attention_mask | `./references/cp.md` attention mask reorder | | TP: choosing the overall TP strategy | `./references/tp.md` plus the TP chart above | | TP: attention output garbled | `./references/tp.md` shard_div_attr | | TP: fused CUDA kernels crash under TP | `./references/tp.md` DTensor-unsafe ops | | TP: a single Linear packs fused QKV or out+down weights | `./references/tp.md` fused-weight rearrangement | | TP: GQA with `num_kv_heads` not divisible by `tp_size` | `./references/tp.md` GQA performance caveat | | TP: when to use `output_layouts=Replicate()` | `./references/tp.md` Replicate strategy | | New text encoder / new VAE | `./references/tep.md` / `./references/vaep.md` | | Register the model in the CLI | `./references/cli.md` | | Verify correctness | `./references/testing.md` PSNR AND SSIM | ### 0.5 Top Pitfalls That Silently Corrupt Output These are the traps that pass without crashing but produce wrong images. Each is expanded in its section. 1. **Skipping SSIM.** PSNR alone cannot detect garbled output — a garbled image can still score PSNR > 25 dB. Always compute both. → `./references/testing.md` correctness verification 2. **CP + `attention_mask` misalignment.** Ulysses all-to-all reorders the sequence; a position-indexed mask no longer lines up → localized **top-left corruption**, PSNR stuck ~28–30. → `./references/cp.md` attention mask reorder 3. **Missing `shard_div_attr` in TP.** The attention processor reshapes with a stale head count → garbled output (the #1 TP bug). → `./references/tp.md` shard_div_attr 4. **PatchFunctor drops loop-body ops.** Extra operations inside the block loop are silently skipped after the cache wrapper takes over → stale `temb` / modulation. → `./references/cache.md` PatchFunctor Pitfall B 5. **GQA attention TP via `Replicate` is *correct* but often *slower* than a single GPU** (all-gather on Q dominates). Benchmark against FFN-only TP before shipping. → `./references/tp.md` GQA strategy 6. **DTensor-unsafe fused kernels** (flash_attn SwiGLU, triton RMSNorm) crash with illegal-memory-access under TP → patch them to PyTorch ops first. → `./references/tp.md` DTensor-unsafe ops 7. **Choosing pure patch-based CP when hybrid (hook + method patch) suffices.** If the ONLY obstacle is a method like `get_rotary_pos_embed` or `pos_embed` that returns a sequence-length-dependent tensor not reachable by hooks, patch that ONE method — do NOT patch the entire `forward()`. Pure patch-based CP is a maintenance burden that should never be green-lit without proving hybrid is truly impossible. → `./references/cp.md` implementation approaches and hybrid CP ### 0.6 Integration TODO List Template > **⚠️ MANDATORY: One feature at a time. Test thoroughly before starting the next.** Do NOT develop Cache + CP + TP together — a bug in one will contaminate all downstream results. Each phase depends on a verified-correct baseline from the previous phase. Copy this template into your plan, fill in model-specific details, and check off items as they pass. Phase 0 — Prerequisites ``` - [ ] Model files downloaded to local path - [ ] Test image prepared (for IE2I / I2V models) - [ ] Diffusers source code read: pipeline __call__, transformer.forward(), block class - [ ] ForwardPattern identified (Pattern_0 through Pattern_5) - [ ] Block ModuleList name(s) confirmed - [ ] GQA / attention_mask / modulation style documented - [ ] Planning complete and reviewed (use Review Plan agent) Phase 1 — Cache (BlockAdapter) - [ ] 1.1 Write BlockAdapter function in adapters.py - register name, _relaxed_assert, blocks, forward_pattern, check_forward_pattern=True - [ ] 1.2 Evaluate PatchFunctor need - Check: keyword args in block call? extra ops in loop body? - If needed: write PatchFunctor in patch_functors/ → wire in BlockAdapter - [ ] 1.3 Register in block_adapters/__init__.py - [ ] 1.4 Install editable cache-dit using `./references/install.md`