guqiong96/lsglang/sglang-diffusion-cache-dit

SGLang Diffusion × Cache-DiT Integration

Path placeholders used throughout: <sglang_dir> = sglang repo root, <cache_dit_dir> = local cache-dit repo root (if not present locally, clone it: git clone https://github.com/vipshop/cache-dit), <flux_model_dir> = the DiT model weights actually under test (if the user has not specified the model, ask for the model name and checkpoint path first — the workflow is not FLUX-specific; FLUX.1-dev is only the reference run), <cuda_home> = CUDA toolkit root (typically /usr/local/cuda), <gpu_id> = a free GPU index.

GATE CHECK (confirm before starting)

STOP — are all of the following confirmed?
  1. All work happens in a dedicated env for sglang diffusion testing (e.g. `conda activate sgl`),
     fully isolated from cache-dit/ffpa dev envs: never touch other envs, and never let sglang
     dependencies leak into them. If the env does not exist, create a new dedicated one first
     (conda/venv both fine, see §3) — never reuse an existing env just to save effort. Install
     missing dependencies directly into the dedicated env.
  2. GPU: run sglang jobs on the GPUs allocated for them (e.g. CUDA_VISIBLE_DEVICES=<gpu_id>);
     other GPUs may be busy with other jobs.
  3. The test artifact directory <sglang_dir>/.tmp/{task}/ exists (.tmp/ is gitignored; never
     write test outputs into the repo root).
  4. CLI args have been verified once via `sglang generate --help | grep -- <every arg you use>`
     (the old --warmup from earlier docs/PRs is gone; it is now
     --warmup-mode {off,request,server} + --warmup-steps N).
  5. Plan-time alignment — align ALL of the following with the user while drafting the plan
     (before any run), not afterwards:
     - Model under test: model name + weights path. Do NOT default to FLUX.1-dev; adapt the
       case-name prefix and the PSNR baseline table to the actual model.
     - Generation settings: resolution, step count, prompt/seed (reference run used
       1024x1024 / 28 steps — follow the user's actual setup instead).
     - Local cache-dit checkout: does `<cache_dit_dir>` exist? The yaml configs in §4 come from
       `<cache_dit_dir>/examples/configs/`; if absent, agree with the user whether to
       `git clone https://github.com/vipshop/cache-dit` or obtain the configs another way.
     Anything unknown → ASK the user first.
  NO → fix these before touching anything.

Hard rules

• Conflict resolution: keep the target branch's refactored structure (the knob system in §2). If the upstream PR's direct-write style targets code that has since been refactored, re-inject it following the new pattern; never revert the target branch's refactor.
• After touching cache-related modules, run all test/unit/test_cache_dit*.py (not just one file — a skipped stub test once left a 9/9 ImportError that only surfaced at CLI stage).

1. Integration map (4 files on the sglang side)

Data flow: env (with secondary fallback) → knob (request > env) → CacheDitConfig fields → mutual-exclusion guard → cache-dit Config (DBCacheConfig / DMDCalibratorConfig / TaylorSeerCalibratorConfig)

Key facts:

• Knob priority: request override (sampling_params.cache_dit_params) > env; a secondary knob first inherits the request-level primary value, then falls back to SGLANG_CACHE_DIT_SECONDARY_*, then to the primary env default.
• A single change in _build_cache_dit_config automatically covers both the primary/secondary call sites and the minimax_h3 subclass (its override calls super()).
• Dual-transformer models (wan2.2 etc.) can reuse the primary config wholesale for the secondary transformer (_cache_dit_secondary_uses_primary_config()); new fields are inherited for free.
• The calibrator is a single slot: DMD × TaylorSeer are mutually exclusive; _assert_calibrator_exclusive(config, label) guards both enable entry points (enabling both raises ValueError).

2. Standard six steps to add a cache knob (DMD as the example)

1. envs.py annotation section: add the primary + secondary variables (e.g. SGLANG_CACHE_DIT_DMD: bool = False).
2. envs.py getter: non-bool → add a (SUFFIX, type, default) tuple to _CACHE_DIT_SECONDARY_CONFIGS (auto-generates the secondary fallback); bool → write a dedicated _secondary_xxx_getter (get_bool_env_var(SECONDARY_X, default=os.getenv(PRIMARY_X, "false"))) and register it.
3. cache_dit_integration.py: add the field to CacheDitConfig (with docstring and default).
4. cache_dit_integration.py: add the key to CACHE_DIT_REQUEST_KNOB_KEYS (symmetric with taylorseer; automatically joins the request/secondary validation sets and the remount detection cache_dit_overrides_key).
5. denoising.py: append a knob() entry after taylorseer_order inside _build_cache_dit_config:

enable_dmd=knob("enable_dmd", envs.SGLANG_CACHE_DIT_DMD, envs.SGLANG_CACHE_DIT_SECONDARY_DMD, secondary=secondary)

1. Calibrator/special logic: construct the corresponding cache-dit Config inside enable_cache_on_transformer and enable_cache_on_dual_transformer (DMD: if enable_dmd: DMDCalibratorConfig(...) elif enable_taylorseer: ...), and add the field to the log format string (placeholder and argument counts must stay aligned).

Conflict forecast when porting an upstream PR: denoising.py always conflicts (the PR's direct envs.XXX → Config construction was refactored away on the target branch) → apply only step 5; cache_dit_integration.py conflicts locally (the per-request knob system added on the target branch sits right next to the PR's insertion points) → keep the target branch's frozenset and append the new keys.

3. Environment setup (dedicated env)

Environment principle (mandatory, top priority): everything — pip installs/uninstalls, source builds, unit tests and CLI runs — happens inside the dedicated env (e.g. conda activate sgl). It is fully isolated from the cache-dit/ffpa dev envs: never install/uninstall packages in other envs, and never let sglang dependency changes (torch, cache-dit, flashinfer, ...) leak into them. If the env does not exist yet, create a dedicated one before continuing — do not reuse any existing env:

conda create -n sgl python=3.12
conda activate sgl
cd <sglang_dir>
pip install -e ".[diffusion]" --no-build-isolation

(Any other virtualenv mechanism works as well; what matters is dedicated and isolated. After switching envs, confirm versions with pip show sglang cache-dit before testing.)

conda activate sgl
# Pure Python (no SVDQuant):
pip install cache-dit==<ver>            # or cache-dit-cu13==<ver> --no-deps
# SVDQuant (PTQ nvfp4) requires a working CUDA extension. The PyPI wheel can be ABI-incompatible
# with a newer torch (undefined symbol: materialize_cow_storage → the wheel was built against an
# older torch). In that case, build from the local cache-dit source tree
# (no local checkout? git clone https://github.com/vipshop/cache-dit first):
cd <cache_dit_dir>
export CUDA_HOME=<cuda_home>
pip install setuptools-scm              # missing → metadata-generation-failed
CACHE_DIT_BUILD_SVDQUANT=1 pip install ".[quantization]" --no-build-isolation
# Extension self-check (pinpoints the load error in one step):
python -c "from cache_dit.quantization.svdquant import svdq_is_available, svdq_get_load_error as e; print(svdq_is_available(), e())"

Other environment pitfalls:

• A stale flashinfer-cubin that mismatches the flashinfer main package → blocks every sglang import; if the installed flashinfer has no matching cubin release (e.g. 0.6.18), pip uninstall flashinfer-cubin directly.
• The dedicated env may lack pytest → the test files use unittest.main() style; run python <test_file.py> directly.
• test_cache_dit_integration.py replaces the cache_dit module with a stub (no real install required): after changing top-level imports in cache_dit_integration.py, sync the stub's top-level symbols in _install_cache_dit_stub() (missing BlockAdapterRegister/Parallelism*/DMDCalibratorConfig once caused 9/9 ImportError).

4. CLI test matrix (reference run: PRO 5000, FLUX.1-dev, 1024×1024, 28 steps)

Nine-case design (backend × acceleration feature combos):

Common args: --model-path=$FLUX_DIR --log-level=info --prompt='...' --width=1024 --height=1024 --num-inference-steps=28 --warmup-mode request --warmup-steps 1 --dit-cpu-offload false --text-encoder-cpu-offload false --save-output --output-path .tmp/{task}/outputs; compile cases add --warmup-steps 28; the blackwell yaml is already svdq_nvfp4_r128_dq (nvfp4, not int4) — grep to confirm before running.

run_case script pattern (saved as .tmp/{task}/run_matrix.sh):

• timeout <1800~3600> guard + stdout redirect to logs/{name}.log + summary.log records PASS/FAIL (dual criteria: rc + png existence) + a failure does not abort the remaining cases.
• When several cases fail with a common root cause, fix it and re-run only the failed segment via a small sub-script, not the whole matrix.
• Submit in the background; then do a one-shot health check (sleep 45-60 && tail -3 logs/<first_case>.log && nvidia-smi -i <gpu_id>) to confirm the model is loading without CLI errors, then stop and wait for the completion notification — never poll.
• Save long verification scripts (PSNR / perf extraction) as .py files; avoid long python -c one-liners (nested f-strings are error-prone).

Fixed failure-diagnosis order:

1. Check summary.log and whether the output files were produced (EOFError / worker did not terminate gracefully, forcing / leaked-semaphore lines in the tail are multiprocess shutdown noise, not the failure itself);
2. grep -iE 'error|assert|exception|raise' logs/{case}.log | head to find the first traceback (the real cause is often mid-log, e.g. AssertionError: Quantization backend ... not supported);
3. Use package-level diagnostics when available (svdq_is_available()/svdq_get_load_error()).

4.1 Full command reference (battle-tested during the 1.5.1 upgrade; adapt and reuse)

Unit tests + negatives (mandatory after touching cache modules):

cd <sglang_dir>/.tmp/{task}
conda activate sgl
# Run ALL cache-dit-related tests via glob (unittest style, direct run; the env may lack pytest)
for t in ../../python/sglang/multimodal_gen/test/unit/test_cache_dit*.py; do echo "== $t"; python $t 2>&1 | tail -3; done
# Mutual-exclusion negatives already exist as unit tests
# (test_both_calibrators_raise_on_{,dual_}transformer) — no ad-hoc script needed