Skip to content

perf(vllm-omni): batch BAGEL grouped t2i rollout into one packed generate_image - #203

Merged
CjhHa1 merged 10 commits into
Tencent-Hunyuan:mainfrom
zzhuoxin1508:bagel-dit-batch-main
Jul 21, 2026
Merged

perf(vllm-omni): batch BAGEL grouped t2i rollout into one packed generate_image#203
CjhHa1 merged 10 commits into
Tencent-Hunyuan:mainfrom
zzhuoxin1508:bagel-dit-batch-main

Conversation

@zzhuoxin1508

@zzhuoxin1508 zzhuoxin1508 commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Group BAGEL vLLM-Omni t2i rollout requests by prompt and use
    num_outputs_per_prompt=samples_per_prompt (spp) instead of sending one diffusion
    request per image sample. The shared prompt KV is built once and replicated ×spp;
    all spp images are denoised in one packed generate_image per step, then the
    engine splits them back per request.
  • Keep downstream UniRL semantics sample-level by preserving sample-level x_T
    recipes (grouped-span slice of resolve_request_noise), splitting log-prob +
    trajectory per image ([spp, T+1, seq, C]), and validating the grouped prompt
    layout via grouped_texts_from_req.
  • **Bugfix ** feat(trainer): add periodic eval to PE / ReFL / UnifiedModel #202 renamed the
    pack_initial_noise_extra_args() keyword n_prompts= -> n_samples= and updated the
    shared dit.py caller, but BagelInputAdapter (merged from a side branch) still
    passes n_prompts=. Since that rename, every bagel_t2i rollout that packs a
    driver-authored x_T crashes on its first generate with `TypeError: pack_initial_noise_extra_args() got an
    unexpected keyword argument 'n_prompts'

Compatibility / Risk

  • Packed path is scoped to pure bagel_t2i + cfg≤1 (_is_packable_t2i /
    _is_batchable_t2i). cfg>1 / text-output / i2i keep the sample-level layout
    (num_outputs_per_prompt=1); a mismatched num_outputs>1 on the worker raises
    instead of a broken silent fallback. spp=1 is byte-identical to before.
  • Missing CFG keys are treated as non-batchable (upstream Bagel defaults absent
    keys to CFG-ON @4.0/1.5).
  • Packed images are isolated by varlen block-diagonal attention, so a packed forward
    == spp separate forwards. Only per-step SDE RNG differs (<0.1%, no collapse, ratio 1.0).

Test Plan

  • 8×H20 BAGEL vLLM-Omni E2E, 3-way A/B (spp=16):

    variant e2e s/rollout ratio
    trainside ~187 1.0000
    unmodified vllm_omni ~150 1.0000
    batched (this PR) ~139 1.0000
  • Confirm bagel_t2i with driver x_T (recipe / initial_noise_batch) no longer raises
    on pack_initial_noise_extra_args.

  • Smoke cfg>1 + spp>1: stays on sample-level path (no pack).

Reviewer Notes

@github-actions github-actions Bot added the wip Draft / work in progress label Jul 13, 2026
@zzhuoxin1508
zzhuoxin1508 force-pushed the bagel-dit-batch-main branch from 8c534f8 to d685cce Compare July 14, 2026 06:21
@zzhuoxin1508
zzhuoxin1508 marked this pull request as ready for review July 15, 2026 07:21
@github-actions github-actions Bot added need review Ready and waiting for review and removed wip Draft / work in progress labels Jul 15, 2026
@zzhuoxin1508
zzhuoxin1508 requested a review from CjhHa1 July 15, 2026 07:30
CjhHa1 and others added 4 commits July 20, 2026 22:14
Only collapse samples_per_prompt when cfg<=1 so cfg>1 keeps the real
sample-level path; refuse num_outputs>1 when not batchable. Treat missing
CFG keys as non-batchable (upstream defaults CFG-ON). Rename/clarify the
generate_image tap and reuse grouped_texts_from_req.
…decode

Simplify spp access, rename packability helper, shorten comments, reuse
upstream's first VAE decode, contiguous log-probs for IPC, and reject
text-output modalities in the batchable guard.

@CjhHa1 CjhHa1 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@CjhHa1
CjhHa1 merged commit c96fce7 into Tencent-Hunyuan:main Jul 21, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

need review Ready and waiting for review

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants