Sglang Diffusion Cache Dit

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

#!/bin/bash
set -u
BASE=<sglang_dir>/.tmp/{task}
FLUX_DIR=<flux_model_dir>   # model actually under test — ask the user if not specified
CFG=<cache_dit_dir>/examples/configs
OUT=$BASE/outputs; LOGS=$BASE/logs; mkdir -p "$OUT" "$LOGS"
export CUDA_VISIBLE_DEVICES=<gpu_id>
PROMPT='A fantasy landscape with mountains and a river, detailed, vibrant colors'
COMMON=(
  --model-path="$FLUX_DIR" --log-level=info --prompt="$PROMPT"
  --width=1024 --height=1024 --num-inference-steps=28
  --warmup-mode request --warmup-steps 1          # old --warmup is gone; verify via --help first
  --dit-cpu-offload false --text-encoder-cpu-offload false
  --save-output --output-path "$OUT"
)
run_case() {  # $1=name $2=timeout_s; remaining args are case-specific
  local name=$1; shift; local timeout_s=$1; shift
  echo "[$(date '+%H:%M:%S')] START $name" >> "$LOGS/summary.log"
  timeout "$timeout_s" sglang generate "${COMMON[@]}" "$@" \
    --output-file-name "$name.png" > "$LOGS/$name.log" 2>&1
  local rc=$?
  [[ $rc -eq 0 && -f "$OUT/$name.png" ]] && st=PASS || st=FAIL
  echo "[$(date '+%H:%M:%S')] $st $name (rc=$rc, png=$([[ -f $OUT/$name.png ]] && echo yes || echo no))" >> "$LOGS/summary.log"
}
# SGLD triple (env-driven)
run_case flux_sgld 1800
SGLANG_CACHE_DIT_ENABLED=true run_case flux_cache_sgld 1800
SGLANG_CACHE_DIT_ENABLED=true SGLANG_CACHE_DIT_DMD=true run_case flux_cache_dmd_sgld 1800
# diffusers triple (yaml-driven)
run_case flux_diffusers 1800 --backend diffusers
run_case flux_cache_diffusers 1800 --backend diffusers --cache-dit-config "$CFG/cache.yaml"
run_case flux_cache_dmd_diffusers 1800 --backend diffusers --cache-dit-config "$CFG/cache_dmd.yaml"
# SVDQ nvfp4 triple (requires the svdq extension)
run_case flux_svdq_nvfp4_diffusers 2400 --backend diffusers --cache-dit-config "$CFG/blackwell/quantize_svdq.yaml"
run_case flux_svdq_nvfp4_compile_diffusers 3600 --backend diffusers --warmup-steps 28 \
  --enable-torch-compile --cache-dit-config "$CFG/blackwell/quantize_svdq.yaml"
run_case flux_cache_dmd_svdq_nvfp4_compile_diffusers 3600 --backend diffusers --warmup-steps 28 \
  --enable-torch-compile --cache-dit-config "$CFG/blackwell/cache_dmd_svdq.yaml"
echo "[$(date '+%H:%M:%S')] MATRIX DONE" >> "$LOGS/summary.log"

sleep 60 && tail -3 logs/flux_sgld.log | cut -c1-160 && nvidia-smi --query-gpu=index,memory.used --format=csv,noheader -i <gpu_id>
# Expect: model loading / inferring + tens of GB of VRAM in use;
# `ambiguous option` / Traceback → kill immediately and fix the args

grep -E 'DMD=True|Calibrator Config: DMD' logs/flux_cache_dmd_sgld.log          # DMD active
grep -E 'Match Blocks|Collected Context Config' logs/flux_cache_sgld.log        # DBCache active
grep -E 'SVDQuant.*Type: svdq_nvfp4_r128_dq' logs/flux_svdq_nvfp4_diffusers.log # quantization active

for f in logs/*.log; do echo "$f: $(grep -oE 'finished in [0-9.]+ seconds' $f | head -1) $(grep -oE '[0-9.]+it/s' $f | tail -1)"; done

# Compare each accelerated result against the same-backend baseline
cache-dit-metrics psnr ssim -i1 outputs/flux_sgld.png -i2 outputs/flux_cache_sgld.png
cache-dit-metrics psnr ssim -i1 outputs/flux_sgld.png -i2 outputs/flux_cache_dmd_sgld.png
cache-dit-metrics psnr ssim -i1 outputs/flux_diffusers.png -i2 outputs/flux_svdq_nvfp4_diffusers.png
cache-dit-metrics psnr ssim -i1 outputs/flux_svdq_nvfp4_compile_diffusers.png \
  -i2 outputs/flux_cache_dmd_svdq_nvfp4_compile_diffusers.png

import numpy as np, torch
from PIL import Image
def load(p): return torch.from_numpy(np.array(Image.open(p))).float() / 255.0
def psnr(a, b):
    mse = ((a - b) ** 2).mean().item()
    return float('inf') if mse == 0 else 10 * np.log10(1.0 / mse)
pairs = [  # (label, baseline, accelerated output) — fill per actual cases
    ('sgld: cache vs base', 'flux_sgld.png', 'flux_cache_sgld.png'),
    ('sgld: cache+dmd vs base', 'flux_sgld.png', 'flux_cache_dmd_sgld.png'),
    ('diff: svdq vs base', 'flux_diffusers.png', 'flux_svdq_nvfp4_diffusers.png'),
    ('svdq: cache+dmd+compile vs svdq', 'flux_svdq_nvfp4_compile_diffusers.png',
     'flux_cache_dmd_svdq_nvfp4_compile_diffusers.png'),
]
import os; os.chdir(os.path.dirname(__file__) + '/outputs')
for name, a, b in pairs: print(f'{name:38s} PSNR = {psnr(load(a), load(b)):6.2f} dB')

conda activate sgl
cd <sglang_dir>
export FLUX_DIR=<flux_model_dir>   # model actually under test — ask the user if not specified
export CUDA_VISIBLE_DEVICES=<gpu_id>
mkdir -p .tmp/{task}/outputs
# SVDQuant cases require: pip install cache-dit-cu13==<ver> --no-deps;
# fall back to a source build if the wheel is torch-ABI-incompatible (see §3);
# uninstall flashinfer-cubin if a stale copy reports a version mismatch

# baseline
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 .tmp/{task}/outputs --output-file-name flux_sgld.png

# DBCache
SGLANG_CACHE_DIT_ENABLED=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 .tmp/{task}/outputs --output-file-name flux_cache_sgld.png

# DBCache + DMD Calibrator
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 .tmp/{task}/outputs --output-file-name flux_cache_dmd_sgld.png

# baseline
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 1 \
  --dit-cpu-offload false --text-encoder-cpu-offload false \
  --save-output --output-path .tmp/{task}/outputs --output-file-name flux_diffusers.png

# DBCache
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 1 \
  --dit-cpu-offload false --text-encoder-cpu-offload false \
  --cache-dit-config $CFG/cache.yaml \
  --save-output --output-path .tmp/{task}/outputs --output-file-name flux_cache_diffusers.png

# DBCache + DMD Calibrator
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 1 \
  --dit-cpu-offload false --text-encoder-cpu-offload false \
  --cache-dit-config $CFG/cache_dmd.yaml \
  --save-output --output-path .tmp/{task}/outputs --output-file-name flux_cache_dmd_diffusers.png

# SVDQuant W4A4 NVFP4
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 1 \
  --dit-cpu-offload false --text-encoder-cpu-offload false \
  --cache-dit-config $CFG/blackwell/quantize_svdq.yaml \
  --save-output --output-path .tmp/{task}/outputs --output-file-name flux_svdq_nvfp4_diffusers.png

# + 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 $CFG/blackwell/quantize_svdq.yaml \
  --save-output --output-path .tmp/{task}/outputs --output-file-name flux_svdq_nvfp4_compile_diffusers.png

# + compile + DBCache + DMD Calibrator
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 $CFG/blackwell/cache_dmd_svdq.yaml \
  --save-output --output-path .tmp/{task}/outputs --output-file-name flux_cache_dmd_svdq_nvfp4_compile_diffusers.png

Enabling cache-dit ... DMD=True (history=6, rank=0, svd=medium), TaylorSeer=False ..., steps=28
[Cache-DiT] Collected Context Config: DBCache_F1B0_W4I1M0MC3_R0.24_N28_CFG0, Calibrator Config: DMD_H(6, medium)
[Cache-DiT] Match Blocks: CachedBlocks_Pattern_0_1_2, for transformer_blocks ...
[Cache-DiT] SVDQuant Type: svdq_nvfp4_r128_dq, Rank: 128