[Feature]: support Flux.2-dev tea_cache

[nuclearwu]

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.

Purpose

Ref #1806, support Flux.2-dev tea_cache

Test Plan

python examples/offline_inference/text_to_image/text_to_image.py \
  --model /workspace/cache/ymttest/johnjan/models/black-forest-labs/FLUX___2-dev/ \
  --prompt "a lovely bunny holding a sign that says 'vllm-omni'" \
  --seed 42 \
  --tensor-parallel-size 4 \
  --num-images-per-prompt 1 \
  --num-inference-steps 50 \
  --guidance-scale 4.0 \
  --height 1024 \
  --width 1024 \
  --cache-backend tea_cache \
  --output outputs/flux2-dev-teacache.png

Test Result

vLLM-Omni:
Reproduced with 4xA800.

Config	Flux.2-dev	Baseline	With Tea-Cache	Speedup
TP=1 & CPU offload		90.5477s	49.3883s	1.83x
TP=2 & CPU offload		62.3353s	35.4450s	1.76x
TP=4		28.8728s	7.9040s	3.65x

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
• The test results. Please paste the results comparison before and after, or the e2e results.
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
• (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

[nuclearwu]

cc @yuanheng-zhao @hsliuustc0106 @wtomin

[Gaohan123]

@wtomin @SamitHuang @ZJY0516 PTAL

[Gaohan123]

Please supplement UT for this feature. Thanks

[nuclearwu]

Please supplement UT for this feature. Thanks

@Gaohan123 Done

[nuclearwu]

@Gaohan123 @hsliuustc0106 please review again.

[nuclearwu]

cc @Gaohan123 @hsliuustc0106

[nuclearwu]

@hsliuustc0106 The CI error (Docker container timeout with exit code 124) has no direct logical relationship with the PR code changes

[nuclearwu]

cc @hsliuustc0106 @Gaohan123

[hsliuustc0106]

resolve conflicts please

[nuclearwu]

resolve conflicts please

@hsliuustc0106 Done

[nuclearwu]

cc @hsliuustc0106

[wtomin]

There is a known bug when both teacache and cpu-offload are enabled. #1868. I think it will be fixed soon by #2689. Please test with your implementation again after this bugfix.

[nuclearwu]

There is a known bug when both teacache and cpu-offload are enabled. #1868. I think it will be fixed soon by #2689. Please test with your implementation again after this bugfix.

@wtomin @hsliuustc0106 OK, Done. The verification results are shown in the above table and the bugfix is effective.

[wtomin]

LGTM.