Skip to content

[Benchmark][Diffusion] MixFusion benchmarks - #8117

Draft
renyongxing wants to merge 25 commits into
vllm-project:mainfrom
renyongxing:feat/mixfusion-benchmark
Draft

renyongxing wants to merge 25 commits into
vllm-project:mainfrom
renyongxing:feat/mixfusion-benchmark

Conversation

@renyongxing

Copy link
Copy Markdown

Stacked on #6711 — review only the 4 files below; the rest of the diff is #6711's own content, which is already under review there. This PR will be rebased onto main once #6711 lands.

Benchmark tooling for MixFusion, plus the profiling knob it needs. Contents:

  • benchmarks/diffusion/mixfusion_kernel_benchmark.py — synthetic HunyuanImage-3.0-shaped DiT stack (RMSNorm, fused QKV, SwiGLU MLP). Times independent / padded / MixFusion layouts of the same sequences and reports correctness, latency, peak memory and relative work. Loads no checkpoint, so it isolates the layout effect and runs in seconds.
  • benchmarks/diffusion/mixfusion_benefit.py — real-pipeline A/B: one-request-at-a-time vs the same prompts admitted concurrently, parameterized over --family {qwen,hunyuan}. Replaces the two near-duplicate per-family scripts from [Perf][Diffusion][Qwen-Image] Replace per-bucket varlen loop with flat varlen attention for mixfusion #6711.
  • benchmarks/diffusion/README.md — a ## MixFusion benchmarks section covering all three entry points, including an online-serving recipe that reuses the repo's own diffusion_benchmark_serving.py via extra_args (no new serving code).
  • vllm_omni/profiler/omni_torch_profiler.py (+3) — OMNI_PROFILER_ACTIVITIES lets the pipeline profiler's activity list be narrowed from the default CPU,CUDA. Nothing in the tree sets it; unset, behaviour is byte-identical to main.

Size exemption request

This diff is 1071 non-test added lines, 71 over the 1000-line budget, so stating a concrete exemption as the review comment invited:

File Lines
mixfusion_benefit.py 597
mixfusion_kernel_benchmark.py 392
README.md 79
omni_torch_profiler.py 3

There is no production code path here. No model, scheduler, attention backend or request path is touched — the only non-benchmarks/ file is a 3-line profiler env override. The lines are measurement code and its documentation, which cannot be meaningfully shrunk further: the two per-family scripts from #6711 were already consolidated into one parameterized script (753 → 597 lines), and the synthetic benchmark was reviewed line by line and is dense by design. Removing documentation or the Qwen small-GCD pre-check to fit the budget would make the benchmark less trustworthy, not smaller in any useful sense.

Validation status

  • ruff check and ruff format --check pass; both scripts carry the correct SPDX header.
  • The scripts' pure helpers (argument parsing and per-family defaults, seed semantics, request construction, extra_args contents, engine config, output aggregation) are covered by 55 assertions.
  • Not run end-to-end: no checkpoint was loaded and no real generation was measured on the hardware this targets. The scripts are validated as code, not as reported numbers.

🤖 Generated with Claude Code

desenSunUBW and others added 20 commits September 10, 2026 09:10
…Image-3.0 and Qwen-Image (vllm-project#3638)

Signed-off-by: renyongxing <renyongxing2@huawei.com>
…en-Image mixfusion

Signed-off-by: renyongxing <renyongxing2@huawei.com>
…ache pop key

- qwen_image_mixfusion_benefit: OmniDiffusionRequest now takes prompt/request_id
  (single), not prompts/request_ids lists; construct per-request and gather.
- hunyuan_image3_real_mixfusion_benefit: same API migration (one request per
  prompt, asyncio.gather for the batched path).
- diffusion_model_runner: state_cache pop key uses state.request_id instead of
  the non-existent state.req_id in the interrupted/denoise_completed path.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
MixFusion scheduling and pipeline code (base_scheduler, qwen/hunyuan
pipelines) still iterates request.prompts after the single-request API
migration. Add a backwards-compatible prompts property returning [self.prompt]
so the mixfusion batch path works on the current single-prompt request API.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
Replace the native apply_rotary_emb_torch path in QwenImageCrossAttention
with the fused Triton rotary_kernel path. MixFusion lays out per-token
(chunk-position, resolution) rotary frequencies in chunk order; flatten the
3D image_freq_chunks to [B*chunk_size, D/2] so every image token carries its
own frequencies and feed them alongside text freqs through the fused kernel.
Also adapt the mixfusion benchmark scripts to the single-request API and
enable torch profiler output for the batched path.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
diffusion_worker.py referenced get_fs_group/get_hsdp_replicate_group
without definitions in parallel_state.py after the upstream revert,
causing ImportError at worker startup. Cherry-pick upstream vllm-project#6722 which
defines both group accessors and initializes FS/HSDP-replicate groups
when use_hsdp=True, restoring topology-aware DiffusionWorker process
titles.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
Rebase the mixfusion-varlen feature branch on the latest upstream main
(279 commits ahead of the previous merge base 13173d6) and resolve
conflicts in:
- diffusion/executor/multiproc_executor.py: adopt upstream lock-protected
  shutdown cleaner (drop wake_events from cleaner)
- diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py: keep
  MixFusion attention helpers, adopt upstream _forward_paged split and
  inject MixFusion branch into _forward_dense_legacy
- diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py: merge
  MixFusion helpers with upstream num_special_tokens / guidance /
  timesteps_r support
- diffusion/sched/base_scheduler.py: keep MixFusion shape relaxation and
  add upstream use_step_execution to sampling params key

Signed-off-by: renyongxing <renyongxing2@huawei.com>
Run ruff check --fix and ruff format on the two files touched by the
upstream merge to satisfy pre-commit:
- hunyuan_image3_transformer.py: reorder imports, drop stray blank line
- pipeline_hunyuan_image3.py: reorder imports, drop stray blank lines

Signed-off-by: renyongxing <renyongxing2@huawei.com>
…ate.extra

Propagate mixfusion_min_chunk_tokens / mixfusion_max_chunks from
sampling.extra_args into state.extra during prepare_encode so the
request-specific values reach _denoise_step_mixfusion instead of the
hard-coded 256/128 defaults.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
…st attention

Address lishunyang12 review comments on PR vllm-project#6711 for the Qwen-Image
MixFusion flat varlen path:

* Keep RoPE frequencies in FP32 and cast only the rotated outputs back to
  the query/key dtype. The vLLM fused kernel requires query/key and cos/sin
  to share a dtype, so promote the rotated tensors to FP32 around the fused
  call and cast the outputs back.

* Compute real_txt_lens once per batch in _forward_mixfusion instead of
  calling mask.sum().tolist() inside every transformer block, and thread
  the value through the block into the attention layer.

* Fall back to per-request dense attention when the selected backend cannot
  handle flat varlen inputs (no _forward_varlen_flat), so MixFusion still
  runs correctly on backends like TORCH_SDPA.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
Resolve conflicts in diffusion_model_runner.py stepwise scheduler loop:
keep the MixFusion list-noise_pred handling from feat/mixfusion-varlen while
adopting main's local 'pipeline' variable rename and latents None check.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
…mage

Support num_outputs_per_prompt > 1 in Qwen MixFusion: encoder_hidden_states,
mask, timesteps, and guidance are already per-image rows upstream, but the
_forward_mixfusion pack loop built chunk mappings per request. Expand each
state's latents into one row per image so every image gets its own prompt
row and chunk mapping; reassemble outputs back to [num_images, seq, C] per
request to preserve the external contract.

Also fix _qwen_mixfusion_candidate chunk_count to multiply in the per-request
image count, and slice per-image prompt rows by cumulative offset in the
serial-ragged fallback.

Verified against real Qwen-Image weights: a [N_img=2, N_img=1] batch produces
shapes [2, seq, C] / [1, seq, C] and is bit-identical to the same images run
as three single-image states.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
Signed-off-by: renyongxing <renyongxing2@huawei.com>
- Slice timesteps by per-image rows (offset:offset+num_images) in
  _denoise_step_serial_ragged instead of by request index, so requests
  after one with num_outputs_per_prompt>1 no longer inherit a duplicate
  timestep from the previous request (wrong sigma under mixed
  resolutions).
- Gate the flat varlen attention path on actual platform dispatch
  (CUDA/ROCm/MUSA, XPU only when flash_attn_varlen_func is available)
  instead of method presence: forward_npu ignores is_varlen and would
  otherwise receive packed multi-document Q/K/V without cu_seqlens.

Co-Authored-By: Claude Code <noreply@anthropic.com>
Signed-off-by: renyongxing <renyongxing2@huawei.com>
Resolve conflicts with HunyuanImage3 paged prefix caching (vllm-project#7719):
- pipeline: keep both new staticmethods (_build_mixfusion_attention_mask_buckets from this PR, _slice_cached_prefix_inputs from main)
- transformer: take main's defensive getattr for use_meanflow

Co-Authored-By: Claude Code <noreply@anthropic.com>
Signed-off-by: renyongxing <renyongxing2@huawei.com>
Co-Authored-By: Claude Code <noreply@anthropic.com>
Signed-off-by: renyongxing <renyongxing2@huawei.com>
@hsliuustc0106 hsliuustc0106 added benchmark/profiler/metrics/logger Codes related to benchmarks, profiler, metrics and logger system diffusion codes related to diffusion models labels Sep 24, 2026
@hsliuustc0106

Copy link
Copy Markdown
Collaborator

This PR touches vllm_omni/diffusion/, benchmarks/diffusion/, tests/diffusion/, .gitignore, vllm_omni/profiler/ (17 files). Based on CODEOWNERS coverage of the changed files, the most-related reviewers appear to be:

@Isotr0py @princepride @SamitHuang

Could one of you take a look when you get a chance? Thanks!

renyongxing and others added 4 commits September 28, 2026 02:31
Co-Authored-By: Claude Code <noreply@anthropic.com>
Signed-off-by: renyongxing <renyongxing2@huawei.com>
Co-Authored-By: Claude Code <noreply@anthropic.com>
Signed-off-by: renyongxing <renyongxing2@huawei.com>
Add CPU contract tests for the two behaviours this PR introduces:

- ``tests/diffusion/attention/test_varlen_flat_attn.py`` pins the metadata
  fields the flat packed varlen path consumes, the validation that keeps
  dense or padded callers off it, the packed boundaries handed to the
  FlashAttention kernel, and the ``forward_cuda``/``forward_xpu`` dispatch.
  The kernel is stubbed so the arguments can be asserted on CPU.
- ``tests/diffusion/sched/test_mixfusion_shape_relaxation.py`` pins the
  batch-key shape relaxation and the ``None`` key that per-prompt
  ``additional_information`` requests get, for both ``BaseScheduler`` and
  ``RequestScheduler``.

Co-Authored-By: Claude Code <noreply@anthropic.com>
Signed-off-by: renyongxing <renyongxing2@huawei.com>
Revert the earlier benchmark removal and replace the four ad-hoc entry
points with two focused tools plus documentation:

- ``mixfusion_benefit.py`` merges ``qwen_image_mixfusion_benefit.py`` and
  ``hunyuan_image3_real_mixfusion_benefit.py`` into one ``--family``
  parameterised script that loads the real pipeline through
  ``DiffusionEngine`` and compares sequential serving against the same
  prompts admitted concurrently. It keeps the Qwen small-GCD pre-check so
  a run that silently falls back to independent execution is not reported
  as a MixFusion speedup.
- ``mixfusion_kernel_benchmark.py`` (renamed from
  ``hunyuan_image3_mixfusion_benefit.py``) times independent, padded and
  MixFusion layouts on a synthetic HunyuanImage-3.0-shaped stack, so the
  layout effect can be swept without an 80B checkpoint.
- ``benchmarks/diffusion/README.md`` documents both, plus the online
  serving recipe that attaches ``enable_mixfusion`` through
  ``--extra-body``.
- ``run_qwen_image_mixfusion_benefit.sh`` and
  ``run_hunyuan_image3_mixfusion_benefit.sh`` are dropped; the README
  carries the invocation recipes instead.

Restores the ``OMNI_PROFILER_ACTIVITIES`` override in
``OmniTorchProfilerWrapper`` that the removal had reverted.

Benchmark-tooling only: no production code path changes.

Co-Authored-By: Claude Code <noreply@anthropic.com>
Signed-off-by: renyongxing <renyongxing2@huawei.com>
@renyongxing
renyongxing force-pushed the feat/mixfusion-benchmark branch from 60321d7 to 4d9b9d0 Compare September 28, 2026 03:00
@hsliuustc0106 hsliuustc0106 added the high priority high priority issue, needs to be done asap label Sep 30, 2026 — with ChatGPT Codex Connector
Add the full-matrix A/B harness used to endorse the MixFusion PR against
main: per-case timing driver, latent-dump parity matrix, noise-floor
controls, and a summarizer that aggregates timing, flat-varlen evidence,
latent parity and an image-level PSNR/SSIM layer into one markdown report.

mixfusion_benefit.py gains the pieces the endorsement run needs on both
sides of the A/B: quantization passthrough, --output-type for latent
dumps, --dump-outputs-dir / --profile-dir / --json-output, and a
VLLM_OMNI_ROOT sys.path isolation shim so sitecustomized environments
import the intended checkout.

All paths are parameterized (A_WT/B_WT/RESULTS/PARITY/GPU/MODEL); runs
are resumable and PILOT/ONLY/FORCE-controlled.

Signed-off-by: renyongxing <renyongxing2@huawei.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

benchmark/profiler/metrics/logger Codes related to benchmarks, profiler, metrics and logger system diffusion codes related to diffusion models high priority high priority issue, needs to be done asap

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants