[codex] Optimize Z-Image packed QKV

[BBuf]

Summary

This PR enables the existing packed QKV projection path for Z-Image when loading the standard BF16 checkpoint, not only the Nunchaku quantized path.

• Map checkpoint attention.to_q/to_k/to_v.weight tensors into attention.to_qkv.weight during load.
• Always instantiate ZImageAttention.to_qkv so the non-quantized path also uses one merged QKV projection.
• No benchmark images or media files are committed; the generated outputs are attached below as GitHub user-attachments.

Benchmark

Environment:

• Host/container: radixark03, sglang-diffusion-bbuf
• GPU: H200, CUDA_VISIBLE_DEVICES=7
• Baseline: origin/main at 2d2be5d7b247626c3259fe33c145d0281db6ac4f
• Tuned: this PR at 75d923c25
• Native backend check: benchmark logs contain no Falling back to diffusers backend, Using diffusers backend, or Loaded diffusers pipeline lines.

Benchmark-only note: current origin/main hits the installed H200 FA3 ABI mismatch (flash_attn_varlen_func() got an unexpected keyword argument 'out'). I applied the same local FA3 compatibility shim to both benchmark worktrees only; that shim is intentionally not included in this PR.

Skill preset command used for both runs, with only the worktree and --label differing:

CUDA_VISIBLE_DEVICES=7 FLASHINFER_DISABLE_VERSION_CHECK=1 PYTHONPATH=python \
python3 /root/.codex/skills/sglang-diffusion-benchmark-profile/scripts/bench_diffusion_denoise.py \
  --model zimage \
  --label <zimage_qkv_pr_base_fa3shim_03|zimage_qkv_pr_tuned_fa3shim_03> \
  --output-dir outputs/diffusion_benchmarks

Expanded preset command:

sglang generate \
  --model-path=Tongyi-MAI/Z-Image-Turbo \
  --prompt="A futuristic cyberpunk city at night, neon lights reflecting on wet streets" \
  --log-level=info \
  --seed=42 \
  --width=1024 \
  --height=1024 \
  --num-inference-steps=9 \
  --guidance-scale=4.0 \
  --save-output \
  --warmup \
  --enable-torch-compile \
  --perf-dump-path <json>

Run	Denoise latency	E2E latency	Peak memory	Delta
Baseline	1.2818s	1.5609s	20.8GB	-
Tuned	0.8288s	1.0739s	20.6GB	-35.3% denoise / -31.2% E2E

compare_perf.py reports:

• E2E: 1560.94 ms -> 1073.93 ms, -487.01 ms (-31.2%)
• DenoisingStage: 1281.80 ms -> 828.80 ms, -453.00 ms (-35.3%)

Raw metric files on the H200 run:

• Baseline: /tmp/sglang_zimage_pr_base/outputs/diffusion_benchmarks/zimage_zimage_qkv_pr_base_fa3shim_03.json
• Tuned: /tmp/sglang_zimage_pr_tuned/outputs/diffusion_benchmarks/zimage_zimage_qkv_pr_tuned_fa3shim_03.json

Image Comparison

Baseline, generated by the same preset command:

Tuned, generated by the same preset command:

Fallback links:

• Baseline: https://github.com/user-attachments/assets/9b7763e6-f62d-4dfd-b4ba-994337601bf0
• Tuned: https://github.com/user-attachments/assets/195b2d35-0071-40f2-a80d-eff9b3c17189

Validation

• git diff --check
• python3 -m py_compile python/sglang/multimodal_gen/configs/models/dits/zimage.py python/sglang/multimodal_gen/runtime/models/dits/zimage.py
• bench_diffusion_denoise.py --model zimage baseline + tuned on H200
• native backend log gate via grep for diffusers fallback strings
• compare_perf.py on the two benchmark JSON files

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci