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).