[Diffusion] Cache-DiT 1.5.1: DMD Calibrator, SVDQuant DQ, etc.

[DefTruth]

[Diffusion] Cache-DiT 1.5.1: DMD Calibrator, Per-Request Knobs, Agent Skill

Port of upstream PR #29269 ("[Diffusion] Cache-DiT 1.5.0: DMD Calibrator, SVDQuant DQ, etc.") onto local main, upgraded to cache-dit 1.5.1, plus a reusable agent skill that captures the whole integration workflow.

Branch: cache-dit-1.5.1 (5 commits, based on origin/main)

Motivation

Upstream #29269 brings DMD (Dynamic Mode Decomposition) calibrator support into SGLang Diffusion, but it was written against an older main. Our main has since been refactored to a per-request knob system (_build_cache_dit_config() + _cache_dit_knob() in denoising.py, request-level sampling_params.cache_dit_params > env priority, primary/secondary transformer fallback). This PR ports the upstream feature while preserving that refactor — every new DMD parameter is injected as a knob, never as a direct envs.XXX read.

cache-dit 1.5.1 release: https://github.com/vipshop/cache-dit

Modifications

1. Upgrade cache-dit to 1.5.1 (4 commits, cherry-picked from #29269 with conflict resolution)

• upgrade cache-dit -> 1.5.0 (b10d33323, cherry-pick of 87b13c81d): python/pyproject.toml cache-dit==1.5.0 (1.3.0 → 1.5.0) + top-level import re-organization.
• upgrade cache-dit -> 1.5.0 (c1884b77a, cherry-pick of d4fa42e17): patched_similarity type annotation fix.
• update (2c3ac64da, cherry-pick of 48b711d59, the DMD body, conflicts resolved by hand):
  • runtime/cache/cache_dit_integration.py: import DMDCalibratorConfig; 5 new CacheDitConfig fields — enable_dmd, dmd_history (default 6), dmd_rank (default 0 = auto), dmd_ridge (default 1e-8), dmd_svd_precision (default "medium"); _assert_calibrator_exclusive() raises ValueError when both DMD and TaylorSeer are enabled (single calibrator_config slot per transformer); both enable_cache_on_transformer / enable_cache_on_dual_transformer build DMDCalibratorConfig when enable_dmd is set (DMD takes precedence) with the config logged.
  • multimodal_gen/envs.py: 5 primary env vars (SGLANG_CACHE_DIT_DMD, _DMD_HISTORY, _DMD_RANK, _DMD_RIDGE, _DMD_SVD_PRECISION) + secondary (Wan2.2 low-noise expert) support via _CACHE_DIT_SECONDARY_CONFIGS tuples and a bool _secondary_dmd_getter falling back to primary.
  • runtime/pipelines_core/stages/denoising.py: DMD fields injected via the knob pattern, e.g. enable_dmd=knob("enable_dmd", envs.SGLANG_CACHE_DIT_DMD, envs.SGLANG_CACHE_DIT_SECONDARY_DMD, secondary=secondary) — one change covers primary/secondary call sites and the minimax_h3 subclass (its override calls super()).
• upgrade cache-dit -> 1.5.1 (1890bace9): version bump in pyproject.toml.

Conflict-resolution principle: keep main's knob refactoring; re-inject upstream's direct-write logic following the new pattern (5 DMD keys added to CACHE_DIT_REQUEST_KNOB_KEYS, symmetric with taylorseer). Net diff on denoising.py is +30 lines of knob() entries.

2. Tests (c0e38b75b)

• test/unit/test_cache_dit_integration.py: stub extended with the new top-level symbols; new TestCalibratorSelection — DMD occupies the calibrator slot, DMD × TaylorSeer mutual exclusion raises on both single- and dual-transformer paths.
• test/unit/test_cache_dit_per_request.py: request-level DMD knobs reach CacheDitConfig; secondary inherits request-level primary DMD values.
• All 32 unit tests pass (python test_file.py, unittest style — the env has no pytest).

3. New agent skill: sglang-diffusion-cache-dit (.claude/skills/)

A reusable workflow skill (in-repo, no runtime impact) so future cache-dit upgrades/knob additions follow the same validated path:

• Integration map of the 4 sglang-side files and the knob data flow (env → request knob → CacheDitConfig → guards → cache-dit Config).
• Standard six steps to add a cache knob (DMD as the worked example) + upstream-PR conflict forecast.
• 9-case CLI test matrix (SGLD × diffusers × SVDQ/compile, env-driven and yaml-driven), battle-tested full commands, failure-diagnosis order, cache-dit-metrics psnr ssim accuracy workflow with the PSNR-vs-SSIM garbled-image guard, PSNR reference baselines.
• Environment troubleshooting: wheel-ABI fallback to a source SVDQuant build, stale flashinfer-cubin, stub-test sync.
• references/block_adapter.md: custom BlockAdapter reference (ForwardPattern contracts, has_separate_cfg, third-party/non-diffusers rules, interception pitfalls) — ideas only: sglang adapter code stays in runtime/cache/cache_dit_integration.py, never registered in the cache-dit repo, and PatchFunctor is explicitly not recommended for sglang.

Usage

# DBCache + DMD (SGLD backend, env-driven)
SGLANG_CACHE_DIT_ENABLED=true SGLANG_CACHE_DIT_DMD=true \
  sglang generate --model-path=$FLUX_DIR --log-level=info \
  --prompt='A fantasy landscape with mountains and a river, detailed, vibrant colors' \
  --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 outputs --output-file-name flux_cache_dmd_sgld.png

# DBCache + DMD (diffusers backend, yaml-driven; stacks with SVDQuant DQ + compile)
sglang generate --model-path=$FLUX_DIR --backend diffusers --log-level=info \
  --prompt='A fantasy landscape with mountains and a river, detailed, vibrant colors' \
  --width=1024 --height=1024 --num-inference-steps=28 \
  --warmup-mode request --warmup-steps 28 \
  --dit-cpu-offload false --text-encoder-cpu-offload false \
  --enable-torch-compile \
  --cache-dit-config <cache_dit_dir>/examples/configs/blackwell/cache_dmd_svdq.yaml \
  --save-output --output-path outputs --output-file-name flux_cache_dmd_svdq_nvfp4_compile_diffusers.png

Note: the legacy --warmup flag is gone — use --warmup-mode request --warmup-steps N (compile cases use --warmup-steps 28).

Per-request overrides work through sampling_params.cache_dit_params (all 5 DMD keys are in CACHE_DIT_REQUEST_KNOB_KEYS), with request > env > secondary-fallback priority.

Accuracy Tests

Verified on FLUX.1-dev, single NVIDIA PRO 5000 (sm_120a), 1024×1024, 28 steps, prompt = "A fantasy landscape with mountains and a river, detailed, vibrant colors". PSNR computed against each backend's own baseline (SGLD and diffusers are scored separately; sglang's default DBCache R=0.24 is aggressive, so magnitudes sit below the cache-dit-side PSNR>30 standard).

backend	variant	PSNR (dB)
SGLD	+ DBCache	23.5
SGLD	+ DBCache + DMD	24.5
diffusers	+ DBCache	31.0
diffusers	+ DBCache + DMD	29.0
diffusers	+ SVDQuant W4A4 NVFP4 (DQ)	23.4
svdq stack	+ DBCache + DMD (vs plain svdq)	27.8

DMD beats pure DBCache on the SGLD path (24.5 vs 23.5 dB) and stacks on top of SVDQuant without collapsing quality (27.8 dB vs the svdq-only output). Every accelerated case also verified by log evidence (Calibrator Config: DMD_H(6, medium), Match Blocks, SVDQuant Type: svdq_nvfp4_r128_dq). All 9 CLI matrix cases PASS; 32/32 unit tests PASS.

Speed Tests and Profiling

Same setup; diffusers timing covers text encode + denoise + VAE decode.

configuration	time (s)	speedup
baseline fp16 diffusers	17.15	1.00×
+ SVDQuant W4A4 NVFP4 (DQ)	8.52	2.01×
+ SVDQuant + torch.compile	6.50	2.64×
+ SVDQuant + compile + DBCache + DMD	4.01	4.28×

The four-way stack (SVDQuant DQ + compile + DBCache + DMD) is the fastest configuration at ~4.3× speedup, confirming the four orthogonal acceleration axes (quantization, kernel fusion, block caching, residual forecasting) compose.

Baseline, PRO 5000	SVDQuant DQ R=128, NVFP4	SVDQ DQ + DMD + Compile
17.15	8.52	4.01

Environment Notes

• The PyPI cache-dit-cu13 wheel's SVDQuant extension can be ABI-incompatible with a newer torch (undefined symbol: materialize_cow_storage); fall back to a source build: CUDA_HOME=<cuda_home> CACHE_DIT_BUILD_SVDQUANT=1 pip install ".[quantization]" --no-build-isolation from the cache-dit checkout. Self-check: python -c "from cache_dit.quantization.svdquant import svdq_is_available, svdq_get_load_error as e; print(svdq_is_available(), e())".
• A stale flashinfer-cubin package blocks all sglang imports with a version mismatch; uninstall it when the installed flashinfer has no matching cubin release.

@mickqian @BBuf

---

CI States

Latest PR Test (Base): ❌ Run #34826830765
Latest PR Test (Extra): ❌ Run #34826830290
Latest PR Test (AMD ROCm 10): ❌ Run #34826830508

[mickqian]

/tag-and-rerun-ci

[mickqian]

brilliant! are there any documentation that needs updated too?

[mickqian]

/tag-and-rerun-ci

[DefTruth]

brilliant! are there any documentation that needs updated too?

Thanks for the reminder, the documentation really needs to be updated.

[mickqian]

Lint is currently blocked by Markdown formatting in python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-cache-dit/references/block_adapter.md: ## More references has a trailing space and the file has no final newline. I reproduced and fixed both locally; pre-commit passes on the file, but this fork has maintainer edits disabled so I cannot push the fix. Please remove the trailing space and add the EOF newline.

[mickqian]

/tag-and-rerun-ci

[DefTruth]

brilliant! are there any documentation that needs updated too?

Thanks for the reminder, the documentation really needs to be updated.

done

[mickqian]

/tag-and-rerun-ci

[DefTruth]

Can we merge this PR?

[mickqian]

The current head has a dependency mismatch in Base NVIDIA CI, separate from the ExpertPack timeout. Please address the installation/import boundary described below before rerunning the affected shard.

[mickqian]

[P1] Apply the cache-dit upgrade to Base jobs that import diffusion modules. This pin is under the diffusion extra, but base-b-test-1-gpu-small (1) invokes ci_install_dependency.sh without that extra and installs python[dev,runai,tracing]. On this exact head, its package list still reports cache-dit 1.3.0; importing test/registered/unit/models/test_qwen_image_fp8_norm_quant.py then reaches the new top-level imports in cache_dit_integration.py and fails with ImportError: cannot import name 'BlockAdapterRegister' from 'cache_dit' (failed job). Please make the Base installation path install the required Cache-DiT version before running these tests, and cover the real package/import boundary with a smoke test. The Cache-DiT unit tests add these exports to a stub, so they cannot detect the installed-package mismatch. Rerunning alone does not ensure that this dependency is upgraded.

[DefTruth]

Thanks for pointing out this error. This really needs to be fixed. We do need to consider compatibility with other cases where cache-dit hasn't been upgraded to 1.5.1. I'm working on this issue.

[DefTruth]

Do you prefer to upgrade all scenarios to cache-dit 1.5.1, or only upgrade python/pyproject.toml first?

[DefTruth]

@mickqian fixed.

python -m  pytest test/registered/unit/models/test_qwen_image_fp8_norm_quant.py python/sglang/multimodal_gen/test/unit/test_cache_dit_integration.py -v
======================================================= test session starts =======================================================
platform linux -- Python 3.12.13, pytest-9.1.1, pluggy-1.6.0 -- /workspace/dev/miniconda3/envs/sgl/bin/python
cachedir: .pytest_cache
rootdir: /workspace/dev/vipshop/sglang/test
configfile: pytest.ini
plugins: anyio-4.13.0
collected 19 items

test/registered/unit/models/test_qwen_image_fp8_norm_quant.py::TestQwenImageFp8NormQuantGate::test_merged_qkv_uses_its_materialized_input_scale PASSED [  5%]
test/registered/unit/models/test_qwen_image_fp8_norm_quant.py::TestQwenImageFp8NormQuantGate::test_nonpositive_scale_keeps_fusion_disabled PASSED [ 10%]
test/registered/unit/models/test_qwen_image_fp8_norm_quant.py::TestQwenImageFp8NormQuantGate::test_separate_qkv_requires_identical_input_scales PASSED [ 15%]
test::TestCacheDitRefreshContext::test_dual_refresh_without_scm_preset_skips_steps_mask PASSED                              [ 21%]
test::TestCacheDitRefreshContext::test_refresh_context_with_scm_preset_uses_steps_mask PASSED                               [ 26%]
test::TestCacheDitRefreshContext::test_refresh_context_without_scm_preset_skips_steps_mask PASSED                           [ 31%]
test::TestBuildCustomBlockAdapter::test_builds_adapter_for_registered_class PASSED                                          [ 36%]
test::TestBuildCustomBlockAdapter::test_custom_adapter_is_retained_until_disable PASSED                                     [ 42%]
test::TestBuildCustomBlockAdapter::test_has_separate_cfg_follows_runtime PASSED                                             [ 47%]
test::TestBuildCustomBlockAdapter::test_minimax_h3_uses_main_blocks_with_hidden_state_pattern PASSED                        [ 52%]
test::TestBuildCustomBlockAdapter::test_raises_when_blocks_attr_missing PASSED                                              [ 57%]
test::TestBuildCustomBlockAdapter::test_returns_none_for_unknown_class PASSED                                               [ 63%]
test::TestCalibratorSelection::test_both_calibrators_raise_on_dual_transformer PASSED                                       [ 68%]
test::TestCalibratorSelection::test_both_calibrators_raise_on_transformer PASSED                                            [ 73%]
test::TestCalibratorSelection::test_dmd_takes_calibrator_slot PASSED                                                        [ 78%]
test::TestCacheDitLegacyFallback::test_enable_dmd_raises_clear_error PASSED                                                 [ 84%]
test::TestCacheDitLegacyFallback::test_fallback_import_binds_registry_and_nulls_dmd PASSED                                  [ 89%]
test::TestCacheDitLegacyFallback::test_taylorseer_path_still_enables_cache PASSED                                           [ 94%]
test::TestCacheDitRealPackageBoundary::test_import_chain_matches_installed_package PASSED                                   [100%]

======================================================== warnings summary =========================================================
<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyPacked has no __module__ attribute

<frozen importlib._bootstrap>:488
  <frozen importlib._bootstrap>:488: DeprecationWarning: builtin type SwigPyObject has no __module__ attribute

../../miniconda3/envs/sgl/lib/python3.12/site-packages/torch/jit/_script.py:365: 14 warnings
  /workspace/dev/miniconda3/envs/sgl/lib/python3.12/site-packages/torch/jit/_script.py:365: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

../../miniconda3/envs/sgl/lib/python3.12/site-packages/_pytest/config/__init__.py:1464
  /workspace/dev/miniconda3/envs/sgl/lib/python3.12/site-packages/_pytest/config/__init__.py:1464: PytestConfigWarning: Unknown config option: asyncio_mode

    self._warn_or_fail_if_strict(f"Unknown config option: {key}\n")

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================================================ 19 passed, 17 warnings in 10.24s =================================================

[nvpohanh]

Seeing these CI failures caused by this:

https://github.com/sgl-project/sglang/actions/runs/34321708659/job/102392704016?pr=38426
https://github.com/sgl-project/sglang/actions/runs/34334438529/job/102430646325?pr=35599

Root cause:

PR #37774’s multimodal CI ran on 5090-e-runner-3, upgraded the shared system environment from cache-dit 1.3.0 to 1.5.1, and skipped virtualenv cleanup. PR #38426 later reused that runner without diffusion extras, inheriting the incompatible package.

[mickqian]

CI follow-up on c9b22eb:

The current Base run still has 1200-second server-startup timeouts in the 2-GPU Wan cases: shard 0, shard 1, and shard 2. Upstream #39034, merged as 358c163, fixes an interrupted IPC JIT build leaving subsequent initialization stuck; that recovery path and Wan warmup passed H100 CI. This head predates that fix. Please merge the latest main into this branch and rerun Base CI to remove that known startup blocker. The timeout logs alone do not prove that every stalled case has the same cause.

Shard 2 also reports a separate LTX2AVDecodingStage performance failure for LTX-2.3; the IPC fix does not establish a fix for that performance check. Recheck it separately on the updated head rather than relaxing the baseline.

I cannot push this update: maintainer edits are disabled and I do not have write access to the head repository.

[DefTruth]

CI follow-up on c9b22eb:

The current Base run still has 1200-second server-startup timeouts in the 2-GPU Wan cases: shard 0, shard 1, and shard 2. Upstream #39034, merged as 358c163, fixes an interrupted IPC JIT build leaving subsequent initialization stuck; that recovery path and Wan warmup passed H100 CI. This head predates that fix. Please merge the latest main into this branch and rerun Base CI to remove that known startup blocker. The timeout logs alone do not prove that every stalled case has the same cause.

Shard 2 also reports a separate LTX2AVDecodingStage performance failure for LTX-2.3; the IPC fix does not establish a fix for that performance check. Recheck it separately on the updated head rather than relaxing the baseline.

I cannot push this update: maintainer edits are disabled and I do not have write access to the head repository.

This issue happens because the fork belongs to the xlite-dev group. Pull requests created from organization-owned forks do not support the "Allow edits from maintainers" option on GitHub, so maintainers cannot push commits or update the PR branch directly.

[mickqian]

CI follow-up on 9fefa7f:

The current blocker is base-a-test-cpu (1): TestDSV4TopKDispatch.test_v2_raw_output_uses_sparse_prefill_buffer_with_capture raises AttributeError: PagedIndexerMetadata has no attribute compressed_seq_lens, including its internal retry. This is an outdated DSV4 test fixture, unrelated to this PR's diffusion changes. The other failed Base jobs report health-check fast-fail downstream of this CPU job.

Main already contains the exact fixture fix in #39101 (4309c7c): rename c4_seq_lens to compressed_seq_lens and set compressed_page_size. Please merge the latest main into this branch with a normal merge commit, then run CI on the updated head; rerunning this unchanged head will not fix the deterministic failure.

Maintainer edits are disabled and I have no write access to the fork, so I cannot push this update.

[DefTruth]

If CI continues to be blocked, I will resubmit this PR through my personal GitHub account. Thank you for your patient reply.