[Perf][Bagel] Avoid per-step device syncs in Bagel img2img

[natureofnature]

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

Purpose

AR

The Bagel img2img path forces implicit GPU→CPU synchronizations in two hot loops, leaving the GPU idle while the CPU stalls:

DIT

• AR stage (model_executor/models/bagel/bagel.py) — the MoT routing calls vae_mask.any() / ~vae_mask.any() (each is an implicit .item() device sync), and _adjust_positions_for_img2img indexes the CUDA positions tensor element-by-element in a Python loop (one sync per token).
• DiT stage (diffusion/models/bagel/bagel_transformer.py) — the denoise loops iterate for t in timesteps where timesteps is a CUDA tensor (so t is a 0-d tensor → building the timestep tensor and the cfg_interval comparison sync eachstep), and attention uses Python max()/sum() over length tensors (one sync per element).

Fix

• Cache the VAE-mask occupancy once per request as plain bools; the per-layer routing reads the bools instead of calling .any().
• Copy positions to the host once before the boundary loop.
• Iterate timesteps.tolist() so t is a Python float (timesteps[i] tensor is still used for the scheduler).
• Use tensor reductions (query_lens.max(), query_lens.sum()) instead of Python max()/sum() over tensors.

Test Plan

• Correctness: ran i2i before/after on the same input + seed at 256² and 1024²; generated PNGs are md5-identical between baseline and patched (changes are numerically identical, not approximations).
• Performance: split serving (stage-0 AR / stage-1 DiT on separate GPUs, RDMA connector), 1 warmup + 5 measured i2i requests per config; comparedmetrics.stage_durations (stage_0_gen_ms / stage_1_gen_ms) baseline vs patched.
• Smoke: service boots and serves t2i + i2i (HTTP 200) with the patched code.
• Prompt: Change the color to blue

Test Result

Split serving (RDMA connector, 15 steps, i2i), stage_0 = AR, stage_1 = DiT, All changes are numerically identical — outputs are byte-for-byte unchanged (md5 of generated images matches before/after at both 256² and 1024²).
mean of 5 measured runs (H800, AR/DIT disaggregation):

resolution	stage	before	after	delta
256×256	AR	1064 ms	890 ms	−16.4%
256×256	DiT	755 ms	724 ms	−4.1%
1024×1024	AR	2265 ms	1969 ms	−13.1%
1024×1024	DiT	7714 ms	7362 ms	−4.6%

---

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)

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

[natureofnature]

@princepride PTAL

[Gaohan123]

LGTM. Thanks