[Diffusion] Enable request-level batching for Boogu-Image (TI2I) (#7227)

Signed-off-by: Shenglei Fu <sfu@confluent.io>
Signed-off-by: Shenglei Fu <117230642+ShengleiFu@users.noreply.github.com>
Co-authored-by: Claude <noreply@anthropic.com>

benchmarks/diffusion/backends.py

@@ -107,6 +107,8 @@ async def async_request_image_edits(
 
     if extra_body.get("guidance_scale") is not None:
         form.add_field("guidance_scale", str(extra_body["guidance_scale"]))
+    if extra_body.get("guidance_scale_2") is not None:
+        form.add_field("guidance_scale_2", str(extra_body["guidance_scale_2"]))
     if extra_body.get("negative_prompt") is not None:
         form.add_field("negative_prompt", str(extra_body["negative_prompt"]))
     if extra_body.get("true_cfg_scale") is not None:

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -1107,6 +1107,11 @@ def calculate_metrics(
     error_outputs = [o for o in outputs if not o.success]
 
     num_success = len(success_outputs)
+    error_log_path = os.environ.get("BENCH_ERROR_LOG")
+    if error_outputs and error_log_path:
+        with open(error_log_path, "a") as f:
+            for output in error_outputs:
+                f.write(f"ERROR latency={output.latency:.3f} detail={output.error!r}\n")
     latencies = [o.latency for o in success_outputs]
     peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
 

recipes/Boogu/Boogu-Image.md

@@ -286,7 +286,6 @@ curl -X POST http://localhost:8091/v1/images/generations \
     "size": "1024x1024",
     "seed": 42
   }' | jq -r '.data[0].b64_json' | base64 -d > output-turbo.png
-```
 
 ### Notes
 
@@ -418,6 +417,45 @@ curl -s http://localhost:8091/v1/chat/completions \
   supports Boogu-Image-Edit directly
   (`--model Boogu/Boogu-Image-0.1-Edit --guidance-scale 5.0`, optional
   `--guidance-scale-2`).
+- **Request batching (`--max-num-seqs > 1`) is not bit-exact vs. serial.**
+  Compatible TI2I requests (same task type, same `guidance_scale_2`
+  explicitness) can share one forward pass. The instruction encoder batches
+  with `padding="longest"`, so a request's *padded batch shape* depends on
+  whichever partner request(s) happen to co-arrive, even though production
+  RoPE derives each row's own effective position count from that row's own
+  `attention_mask.sum()` rather than from the batch's padded width. In the
+  measurements below, `--max-num-seqs 1` and `--max-num-seqs > 1` do not
+  produce pixel-identical output for the same seed; the exact mechanism
+  (padding-shape-dependent kernel/reduction-order numerics vs. some other
+  batch-shape effect) has not been isolated, so treat this as an observed
+  drift under the specific configs below, not a diagnosed root cause, and
+  do not extrapolate the numbers below as a general upper bound on other
+  configs. Two distinct regimes have been measured so far, and they are
+  **not** interchangeable tolerances:
+    - **Fixed-shape co-batch** (partner instructions encode to the same
+      length): `Boogu-Image-0.1-Edit`, 512x512, 20 steps, guidance 5/2, one
+      fixed reference image and instruction pair, seeds 42-73, B=2 vs. serial
+      -- mean pixel MAE **1.70/255**, max per-image MAE **2.68/255**, min PSNR
+      **34.35 dB** (2/32 image pairs were pixel-exact).
+    - **Heterogeneous partner length** (co-batched instructions encode to
+      different lengths): MAD **2.6-6.0** in the fixed-length partner-mixing
+      check on this PR's test head -- larger than the fixed-shape case above
+      in that check. A targeted current-head A/B run observed a comparable
+      range: `Boogu-Image-0.1-Edit`, 512x512, 20 steps, guidance 5/2, one
+      fixed A request (5-token instruction) co-batched in turn with a
+      2-token and a 49-token partner instruction (both also differing in
+      reference image and content, not length alone; `max_num_seqs=2`,
+      co-batching verified via request tracing, not assumed) --
+      solo vs. `A`+2-token partner MAD **4.06/255** (PSNR 29.9 dB),
+      solo vs. `A`+49-token partner MAD **3.87/255** (PSNR 29.7 dB),
+      2-token vs. 49-token partner MAD **4.05/255** (PSNR 30.2 dB). A
+      controlled ablation (same partner content/reference, varying only
+      instruction length) would be needed to isolate a pure length effect,
+      and has not been run.
+  Pin `--max-num-seqs 1` for any workload that requires strict
+  serial-reproducibility (e.g. golden-image regression tests); use
+  `--max-num-seqs > 1` only when this drift is acceptable for the throughput
+  gain.
 
 ## Fast image editing (Boogu-Image-0.1-Edit-Turbo)
 

tests/diffusion/models/boogu_image/test_boogu_image_transformer.py

@@ -699,6 +699,116 @@ def test_transformer_forward_ti2i_shape():
     assert torch.isfinite(out).all()
 
 
+def test_transformer_forward_batch_invariant_to_partner_padding():
+    """Row A must not change when only row B's padded length changes.
+
+    Both rows carry a non-empty reference image (the TI2I/editing path), so
+    this exercises the reference-image refiner and the joint
+    instruction/reference/noise sequence concatenation and position
+    offsets -- not just the T2I path, which skips the reference refiner
+    entirely and would leave that concatenation logic uncovered. Real
+    attention masking, RoPE, and sequence concatenation are exercised
+    end-to-end through the actual (tiny) weights -- the mechanism the CPU
+    pipeline-level fakes in ``test_pipeline_boogu_image.py`` cannot reach,
+    since those stop at a fake encoder/transformer boundary. A's instruction
+    content, reference latent, and attention mask are held byte-identical
+    across runs; only B's padded region grows.
+    """
+    from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
+        BooguImageDoubleStreamRotaryPosEmbed,
+        BooguImageTransformer2DModel,
+    )
+
+    torch.manual_seed(0)
+    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+    _randomize_parameters(model)
+    # _randomize_parameters' uniform(-0.02, 0.02) scale (tuned for the other
+    # shape/finite-value tests in this file) makes the instruction stream's
+    # contribution to the final output smaller than float32 rounding noise
+    # after 4 layers of untrained mixing with the image latents -- both the
+    # invariance check and the negative control below were passing/failing
+    # vacuously at that scale (diffs ~1e-6) until confirmed via a scaled-up
+    # diagnostic that masking is in fact respected; amplify locally so a real
+    # masking regression is distinguishable from noise.
+    with torch.no_grad():
+        for param in model.parameters():
+            param.mul_(20.0)
+    model.eval()
+
+    batch_size = 2
+    in_channels = 4
+    latent_h = latent_w = 8
+    instruction_feat_dim = 32
+    a_len = 5
+
+    torch.manual_seed(0)
+    latents = torch.randn(batch_size, in_channels, latent_h, latent_w)
+    timestep = torch.full((batch_size,), 0.5)
+    freqs_real = BooguImageDoubleStreamRotaryPosEmbed.get_freqs_real(model.axes_dim_rope, model.axes_lens, theta=10000)
+    a_content = torch.randn(a_len, instruction_feat_dim)
+    # Non-empty, fixed per-row reference latents route both rows through the
+    # TI2I reference-image refiner and the joint instruction/reference/noise
+    # concatenation, rather than the T2I path (which skips both entirely).
+    a_ref = torch.randn(in_channels, 6, 10)
+    b_ref = torch.randn(in_channels, 6, 10)
+    ref_image_hidden_states = [[a_ref], [b_ref]]
+
+    def run(b_len, a_mask_len=a_len, a_tail_value=0.0):
+        """``a_mask_len`` marks how many of row A's leading positions are valid.
+
+        Real usage always leaves it at ``a_len``; the negative control below
+        widens it to include A's own padding tail so that tail's content
+        starts to matter, proving the invariance check exercises real
+        masking rather than an architecture that ignores trailing positions
+        regardless of the mask.
+        """
+        width = max(a_len, b_len)
+        instruction_hidden_states = torch.zeros(batch_size, width, instruction_feat_dim)
+        instruction_hidden_states[0, :a_len] = a_content
+        if width > a_len:
+            instruction_hidden_states[0, a_len:] = a_tail_value
+        torch.manual_seed(100 + b_len)
+        instruction_hidden_states[1, :b_len] = torch.randn(b_len, instruction_feat_dim)
+
+        instruction_attention_mask = torch.zeros(batch_size, width, dtype=torch.bool)
+        instruction_attention_mask[0, :a_mask_len] = True
+        instruction_attention_mask[1, :b_len] = True
+
+        with torch.no_grad():
+            return model(
+                latents,
+                timestep,
+                instruction_hidden_states,
+                freqs_real,
+                instruction_attention_mask,
+                ref_image_hidden_states=ref_image_hidden_states,
+            )
+
+    out_short_b = run(b_len=3)
+    out_long_b = run(b_len=8)  # B now sits in a much wider padded batch than A.
+
+    torch.testing.assert_close(
+        out_short_b[0],
+        out_long_b[0],
+        rtol=1e-4,
+        atol=1e-4,
+        msg="row A changed when only B's padded length grew",
+    )
+
+    # Negative control: mark A's own padding tail "valid" (mask length ==
+    # full width) and show that what sits there now measurably changes A's
+    # output -- a magnitude-based outlier there did not (SDPA masking makes
+    # softmax weight on a correctly-excluded position exactly zero regardless
+    # of its magnitude), so this instead varies the tail's *content* under a
+    # broken mask and checks the two runs diverge.
+    width = max(a_len, 8)
+    tail_x = torch.randn(width - a_len, instruction_feat_dim)
+    tail_y = torch.randn(width - a_len, instruction_feat_dim)
+    out_broken_x = run(b_len=8, a_mask_len=width, a_tail_value=tail_x)
+    out_broken_y = run(b_len=8, a_mask_len=width, a_tail_value=tail_y)
+    assert not torch.allclose(out_broken_x[0], out_broken_y[0], rtol=1e-4, atol=1e-4)
+
+
 # ---------------------------------------------------------------------------
 # The three tests below are for the fused qk-norm+RoPE kernel's wiring into
 # this transformer: the packed-table layout conversion, the fused path against

tests/diffusion/models/boogu_image/test_pipeline_boogu_image.py

@@ -14,6 +14,7 @@ Two groups:
 """
 
 import os
+import zlib
 from types import SimpleNamespace
 
 import pytest
@@ -520,8 +521,17 @@ def test_constructor_rejects_non_cuda_sequence_parallelism(
 # ---------------------------------------------------------------------------
 
 
+def _content_id(text: str) -> int:
+    """Bounded text IDs, stable across processes unlike builtin ``hash()``.
+
+    Seed/reference checks use separate transformers so BF16 rounding cannot
+    hide their smaller signals beneath the content term.
+    """
+    return zlib.crc32(text.encode()) % 9973
+
+
 class _RecordingProcessor:
-    """Fake Qwen3VLProcessor: deterministic token ids derived from the text."""
+    """Fake Qwen3VLProcessor with content-derived token IDs."""
 
     def __init__(self):
         self.calls = []
@@ -533,8 +543,8 @@ class _RecordingProcessor:
         for i, messages in enumerate(prompts):
             system_text = messages[0]["content"][0]["text"]
             user_text = messages[1]["content"][0]["text"]
-            input_ids[i, 0] = len(system_text) % 997
-            input_ids[i, 1] = len(user_text) % 997
+            input_ids[i, 0] = _content_id(system_text)
+            input_ids[i, 1] = _content_id(user_text)
             input_ids[i, 2:] = torch.arange(2, _SEQ_LEN) + i
         attention_mask = torch.ones(batch, _SEQ_LEN, dtype=torch.long)
         attention_mask[:, -1] = 0  # fake right-padding
@@ -1283,7 +1293,10 @@ def test_apply_chat_template_ti2i_places_image_before_text():
 
 
 class _ImageAwareRecordingProcessor:
-    """Records whether reference images reached the processor."""
+    """Record reference-image presence and encode text with stable IDs.
+
+    Find text by type because TI2I puts the image first.
+    """
 
     def __init__(self):
         self.calls = []
@@ -1295,7 +1308,13 @@ class _ImageAwareRecordingProcessor:
             has_image.append(any(c.get("type") == "image" for c in user_content))
         self.calls.append({"prompts": prompts, "kwargs": kwargs, "has_image": has_image})
         batch = len(prompts)
-        input_ids = torch.arange(batch * _SEQ_LEN, dtype=torch.long).view(batch, _SEQ_LEN)
+        input_ids = torch.zeros(batch, _SEQ_LEN, dtype=torch.long)
+        for i, messages in enumerate(prompts):
+            system_text = messages[0]["content"][0]["text"]
+            user_text = next(c["text"] for c in messages[1]["content"] if c.get("type") == "text")
+            input_ids[i, 0] = _content_id(system_text)
+            input_ids[i, 1] = _content_id(user_text)
+            input_ids[i, 2:] = torch.arange(2, _SEQ_LEN) + i
         attention_mask = torch.ones(batch, _SEQ_LEN, dtype=torch.long)
         return {"input_ids": input_ids, "attention_mask": attention_mask}
 
@@ -1558,55 +1577,10 @@ def test_forward_image_guidance_ignored_without_reference():
 
 
 # ---------------------------------------------------------------------------
-# Request-batch: compatibility key, generator routing, output split
+# Request-batch: generator routing, output split
 # ---------------------------------------------------------------------------
 
 
-def test_boogu_batch_compatibility_key_t2i_stable_ti2i_unique():
-    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import _boogu_batch_compatibility_key
-
-    # t2i: request_id does not enter the key -> t2i requests batch together.
-    assert _boogu_batch_compatibility_key(False, "req-a") == _boogu_batch_compatibility_key(False, "req-b")
-    assert _boogu_batch_compatibility_key(False, "req-a")[1] == "t2i"
-
-    # ti2i: request_id in the key -> each edit gets a unique key, never co-batched.
-    assert _boogu_batch_compatibility_key(True, "req-a") != _boogu_batch_compatibility_key(True, "req-b")
-    assert _boogu_batch_compatibility_key(True, "req-a")[1] == "ti2i"
-
-    # t2i and ti2i never share a key.
-    assert _boogu_batch_compatibility_key(False, "req-a") != _boogu_batch_compatibility_key(True, "req-a")
-
-
-def test_pre_process_key_wiring_t2i_batches_ti2i_isolated(tmp_path):
-    # End-to-end: real pre-process sets request.batch_compatibility_key, and the
-    # scheduler's key builder reads it into condition_key. Two t2i requests share
-    # a key (co-batchable); two edit requests get distinct keys (batch=1).
-    import PIL.Image
-
-    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import get_boogu_image_pre_process_func
-    from vllm_omni.diffusion.request import OmniDiffusionRequest
-    from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
-    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-    pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
-
-    def condition_key(prompt, rid):
-        req = OmniDiffusionRequest(
-            prompt=prompt, sampling_params=OmniDiffusionSamplingParams(height=512, width=512), request_id=rid
-        )
-        pre(req)
-        return build_request_batch_sampling_params_key(req).condition_key
-
-    t2i_a = condition_key({"prompt": "a cat"}, "t-a")
-    t2i_b = condition_key({"prompt": "a dog"}, "t-b")
-    assert t2i_a == t2i_b and t2i_a[1] == "t2i"
-
-    img = PIL.Image.new("RGB", (64, 64))
-    ti2i_a = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-a")
-    ti2i_b = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-b")
-    assert ti2i_a != ti2i_b and ti2i_a[1] == "ti2i"
-
-
 class _GeneratorRecordingVAE:
     """Fake VAE that records the generator passed to each latent ``sample()``."""
 
@@ -1736,78 +1710,6 @@ def test_forward_request_batch_num_outputs_slices_and_generators():
     assert float(outs[1].output[0, 0, 0, 0]) == 2.0 and float(outs[1].output[1, 0, 0, 0]) == 3.0
 
 
-def test_forward_batch_isolation_partner_content_and_seed():
-    """CFG-on, B=2: request A's output must not change when only the
-    co-batched partner's prompt content, negative prompt, or seed changes.
-
-    ``_FakeTransformer``/``_FakeScheduler`` are content-blind (always-zero
-    velocity, latents passed through unchanged), so they cannot catch a
-    cross-request value leak. This test swaps in a transformer whose output
-    depends on both ``instruction_embeds`` (content, positive or negative
-    depending on which CFG branch called it) and ``latents`` (seed), and a
-    scheduler that actually applies the predicted velocity, so a batching bug
-    that mixes rows in either the cond or uncond prediction would change A's
-    result. A negative-prompt-only perturbation is required to cover the
-    uncond branch: varying only the positive prompt never touches
-    ``negative_instruction_embeds``, so an earlier version of this test
-    passed even with a synthetic row-mixing bug injected into the uncond
-    predict() call (verified via a RED-arm check before this fix).
-    """
-
-    class _ContentAwareTransformer(_FakeTransformer):
-        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
-            content = instruction_embeds.mean(dim=(1, 2)).view(-1, 1, 1, 1)
-            return latents + content
-
-    class _ApplyingScheduler(_FakeScheduler):
-        def step(self, model_output, t, latents, return_dict=False):
-            return (model_output,)
-
-    def run(prompt_a, seed_a, neg_a, prompt_b, seed_b, neg_b):
-        pipeline = _make_forward_pipeline()
-        pipeline.transformer = _ContentAwareTransformer()
-        pipeline.scheduler = _ApplyingScheduler()
-        kw = dict(height=64, width=64, num_inference_steps=2, guidance_scale=4.0, output_type="latent")
-        req = _wrap_request_batch(
-            [
-                (
-                    {"prompt": prompt_a, "negative_prompt": neg_a},
-                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)),
-                ),
-                (
-                    {"prompt": prompt_b, "negative_prompt": neg_b},
-                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)),
-                ),
-            ]
-        )
-        return pipeline.forward(req)[0].output
-
-    baseline = run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "blurry")
-    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a totally different scene", 2, "blurry"))
-    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 999, "blurry"))
-    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "watermark"))
-
-
-def test_forward_batched_ti2i_fails_closed():
-    # A batched ti2i must fail closed (it is gated to batch=1).
-    pipeline = _make_forward_pipeline()
-
-    def edit_prompt():
-        return {
-            "prompt": "make it winter",
-            "additional_information": {"preprocessed_image": torch.zeros(1, 3, 64, 64), "prompt_image": None},
-        }
-
-    req = _wrap_request_batch(
-        [
-            (edit_prompt(), _sampling(num_inference_steps=1, guidance_scale=1.0)),
-            (edit_prompt(), _sampling(num_inference_steps=1, guidance_scale=1.0)),
-        ]
-    )
-    with pytest.raises(RuntimeError, match="gated to batch=1"):
-        pipeline.forward(req)
-
-
 def test_supports_request_batch_enabled():
     from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline
 

tests/diffusion/models/boogu_image/test_pipeline_boogu_image_batching.py

@@ -0,0 +1,746 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Boogu-Image request-batching tests: compatibility keys, cross-request
+isolation, CFG-parallel dispatch, and variable-length padding.
+
+Split out from ``test_pipeline_boogu_image.py`` to keep that file's
+constructor/prompt-encoding focus separate from these request-batch-specific
+cases. Shared pipeline/fake builders are imported from there rather than
+duplicated.
+"""
+
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+from .test_pipeline_boogu_image import (
+    _content_id,
+    _EditForwardVAE,
+    _FakeScheduler,
+    _FakeTransformer,
+    _make_edit_forward_pipeline,
+    _make_edit_od_config,
+    _make_forward_pipeline,
+    _sampling,
+    _wrap_request_batch,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class _RefContentVAE(_EditForwardVAE):
+    """Fake VAE whose ``encode`` reflects the input image's mean into the latent."""
+
+    def encode(self, img):
+        val = float(img.float().mean().item())
+        dist = SimpleNamespace(sample=lambda generator=None: torch.full((1, 4, 8, 8), val))
+        return SimpleNamespace(latent_dist=dist)
+
+
+def _edit_prompt(text, neg, ref_fill):
+    import PIL.Image
+
+    return {
+        "prompt": text,
+        "negative_prompt": neg,
+        "additional_information": {
+            "prompt_image": PIL.Image.new("RGB", (64, 64)),
+            "preprocessed_image": torch.full((1, 3, 64, 64), ref_fill),
+        },
+    }
+
+
+# ---------------------------------------------------------------------------
+# Request-batch: compatibility key
+# ---------------------------------------------------------------------------
+
+
+def test_boogu_batch_compatibility_key_t2i_and_ti2i_stable_and_isolated():
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import _boogu_batch_compatibility_key
+
+    # Each task shares a stable key; guidance compatibility is checked upstream.
+    assert _boogu_batch_compatibility_key(False) == ("boogu_image", "t2i")
+
+    assert _boogu_batch_compatibility_key(True) == ("boogu_image", "ti2i")
+
+    # T2I and TI2I use different denoise paths.
+    assert _boogu_batch_compatibility_key(False) != _boogu_batch_compatibility_key(True)
+
+
+def test_pre_process_key_wiring_t2i_and_ti2i_batch(tmp_path):
+    # Verify preprocessing keys reach the scheduler and keep T2I/TI2I apart.
+    import PIL.Image
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import get_boogu_image_pre_process_func
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
+
+    def condition_key(prompt, rid):
+        req = OmniDiffusionRequest(
+            prompt=prompt, sampling_params=OmniDiffusionSamplingParams(height=512, width=512), request_id=rid
+        )
+        pre(req)
+        return build_request_batch_sampling_params_key(req).condition_key
+
+    t2i_a = condition_key({"prompt": "a cat"}, "t-a")
+    t2i_b = condition_key({"prompt": "a dog"}, "t-b")
+    assert t2i_a == t2i_b and t2i_a[1] == "t2i"
+
+    img = PIL.Image.new("RGB", (64, 64))
+    ti2i_a = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-a")
+    ti2i_b = condition_key({"prompt": "edit", "multi_modal_data": {"image": img}}, "e-b")
+    assert ti2i_a == ti2i_b and ti2i_a[1] == "ti2i"
+    assert t2i_a != ti2i_a
+
+
+def test_ti2i_batch_key_distinguishes_explicit_image_guidance(tmp_path):
+    import PIL.Image
+
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import get_boogu_image_pre_process_func
+    from vllm_omni.diffusion.request import OmniDiffusionRequest
+    from vllm_omni.diffusion.sched.request_scheduler import build_request_batch_sampling_params_key
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+    pre = get_boogu_image_pre_process_func(_make_edit_od_config(tmp_path))
+    image = PIL.Image.new("RGB", (64, 64))
+
+    def make_request(guidance_scale_2: float | None) -> OmniDiffusionRequest:
+        request = OmniDiffusionRequest(
+            prompt={"prompt": "edit", "multi_modal_data": {"image": image}},
+            sampling_params=OmniDiffusionSamplingParams(
+                height=512,
+                width=512,
+                seed=123,
+                guidance_scale=2.0,
+                guidance_scale_2=guidance_scale_2,
+            ),
+            request_id="ti2i",
+        )
+        pre(request)
+        return request
+
+    omitted = build_request_batch_sampling_params_key(make_request(None))
+    explicit = build_request_batch_sampling_params_key(make_request(2.0))
+
+    # Omission auto-fills the same value but must not enable Boogu's image-guidance branch.
+    assert omitted.condition_key == explicit.condition_key == ("boogu_image", "ti2i")
+    assert omitted.guidance_scale_2 == explicit.guidance_scale_2 == 2.0
+    assert omitted.guidance_scale_2_provided is False
+    assert explicit.guidance_scale_2_provided is True
+    assert omitted != explicit
+
+
+# ---------------------------------------------------------------------------
+# Request-batch: cross-request isolation
+# ---------------------------------------------------------------------------
+
+
+def test_forward_batch_isolation_partner_content_and_seed(monkeypatch):
+    """Keep A unchanged under B's prompt, negative-prompt, or seed changes.
+
+    Exercise the real CFG/encoding path at B=2 with signal-sensitive fakes.
+    """
+
+    class _ContentAwareTransformer(_FakeTransformer):
+        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+            content = instruction_embeds.mean(dim=(1, 2)).view(-1, 1, 1, 1)
+            return latents + content
+
+    class _ApplyingScheduler(_FakeScheduler):
+        def step(self, model_output, t, latents, return_dict=False):
+            return (model_output,)
+
+    def run(prompt_a, seed_a, neg_a, prompt_b, seed_b, neg_b):
+        pipeline = _make_forward_pipeline()
+        pipeline.transformer = _ContentAwareTransformer()
+        pipeline.scheduler = _ApplyingScheduler()
+        kw = dict(height=64, width=64, num_inference_steps=2, guidance_scale=4.0, output_type="latent")
+        req = _wrap_request_batch(
+            [
+                (
+                    {"prompt": prompt_a, "negative_prompt": neg_a},
+                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)),
+                ),
+                (
+                    {"prompt": prompt_b, "negative_prompt": neg_b},
+                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)),
+                ),
+            ]
+        )
+        return pipeline.forward(req)[0].output
+
+    baseline = run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "blurry")
+    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a totally different scene", 2, "blurry"))
+    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 999, "blurry"))
+    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "watermark"))
+    # Same-length partner rewrites exercise content rather than length.
+    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a fox in a cave", 2, "blurry"))
+    assert torch.equal(baseline, run("a cat on a mat", 1, "ugly", "a dog in a park", 2, "grainy"))
+
+    # Sensitivity control: A's own prompt and negative prompt must affect A.
+    assert not torch.equal(baseline, run("a totally different scene", 1, "ugly", "a dog in a park", 2, "blurry"))
+    assert not torch.equal(baseline, run("a cat on a mat", 1, "watermark", "a dog in a park", 2, "blurry"))
+
+    # Test seed separately: BF16 content magnitudes can hide O(1) noise.
+    class _LatentsOnlyTransformer(_FakeTransformer):
+        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+            return latents
+
+    def run_seed(seed_a, seed_b):
+        pipeline = _make_forward_pipeline()
+        pipeline.transformer = _LatentsOnlyTransformer()
+        pipeline.scheduler = _ApplyingScheduler()
+        kw = dict(height=64, width=64, num_inference_steps=2, guidance_scale=4.0, output_type="latent")
+        req = _wrap_request_batch(
+            [
+                ("a cat on a mat", _sampling(**kw, generator=torch.Generator().manual_seed(seed_a))),
+                ("a dog in a park", _sampling(**kw, generator=torch.Generator().manual_seed(seed_b))),
+            ]
+        )
+        return pipeline.forward(req)[0].output
+
+    seed_baseline = run_seed(1, 2)
+    assert torch.equal(seed_baseline, run_seed(1, 999))  # B's seed must not affect A
+    assert not torch.equal(seed_baseline, run_seed(999, 2))  # A's own seed must affect A
+
+    # RED: compare partner seeds under the same reversed-generator mutation.
+    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
+    real_collate = DiffusionRequestBatch.collate_request_generators
+
+    def reversed_collate(self, num_outputs_per_prompt, default_generator):
+        result = real_collate(self, num_outputs_per_prompt, default_generator)
+        return list(reversed(result)) if isinstance(result, list) and len(result) > 1 else result
+
+    monkeypatch.setattr(DiffusionRequestBatch, "collate_request_generators", reversed_collate)
+    buggy_seed_baseline = run_seed(1, 2)
+    assert not torch.equal(buggy_seed_baseline, run_seed(1, 999))
+
+
+def test_forward_batched_ti2i_partner_reference_isolation(monkeypatch):
+    """Keep A independent of B's content, seed, and reference at B=2.
+
+    Double guidance exercises all three prediction branches; separate fakes
+    keep seed/reference sensitivity independent of content magnitude.
+    """
+
+    class _ContentRefTransformer(_FakeTransformer):
+        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+            out = latents + instruction_embeds.mean(dim=(1, 2)).view(-1, 1, 1, 1)
+            # Reference latents are request-major: one optional list per row.
+            ref = kwargs.get("ref_image_hidden_states")
+            if ref is not None:
+                ref_means = torch.tensor(
+                    [float(r[0].float().mean()) if r is not None else 0.0 for r in ref],
+                    dtype=out.dtype,
+                    device=out.device,
+                ).view(-1, 1, 1, 1)
+                out = out + ref_means
+            return out
+
+    class _ApplyingScheduler(_FakeScheduler):
+        def step(self, model_output, t, latents, return_dict=False):
+            return (model_output,)
+
+    def run(prompt_a, neg_a, seed_a, ref_fill_a, prompt_b, neg_b, seed_b, ref_fill_b):
+        pipeline = _make_edit_forward_pipeline()
+        pipeline.transformer = _ContentRefTransformer()
+        pipeline.scheduler = _ApplyingScheduler()
+        pipeline.vae = _RefContentVAE()
+        kw = dict(
+            height=64,
+            width=64,
+            num_inference_steps=2,
+            guidance_scale=5.0,
+            guidance_scale_2=2.0,
+            output_type="latent",
+        )
+        req = _wrap_request_batch(
+            [
+                (
+                    _edit_prompt(prompt_a, neg_a, ref_fill_a),
+                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)),
+                ),
+                (
+                    _edit_prompt(prompt_b, neg_b, ref_fill_b),
+                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)),
+                ),
+            ]
+        )
+        return pipeline.forward(req)[0].output
+
+    A = ("make it winter", "ugly", 1, 0.25)
+    baseline = run(*A, "a dog in a park", "blurry", 2, 0.75)
+    assert torch.equal(baseline, run(*A, "a totally different scene", "blurry", 2, 0.75))
+    assert torch.equal(baseline, run(*A, "a dog in a park", "watermark", 2, 0.75))
+    assert torch.equal(baseline, run(*A, "a dog in a park", "blurry", 999, 0.75))
+    assert torch.equal(baseline, run(*A, "a dog in a park", "blurry", 2, 0.5))
+    # Same-length partner rewrites exercise content rather than length.
+    assert torch.equal(baseline, run(*A, "a fox in a cave", "blurry", 2, 0.75))
+    assert torch.equal(baseline, run(*A, "a dog in a park", "grainy", 2, 0.75))
+
+    # Sensitivity control: A's own prompt and negative prompt must affect A.
+    B = ("a dog in a park", "blurry", 2, 0.75)
+    assert not torch.equal(baseline, run("a totally different scene", "ugly", 1, 0.25, *B))
+    assert not torch.equal(baseline, run("make it winter", "watermark", 1, 0.25, *B))
+
+    # Keep reference and seed signals separate from the larger BF16 content term.
+    class _RefOnlyTransformer(_FakeTransformer):
+        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+            ref = kwargs.get("ref_image_hidden_states")
+            if ref is None:
+                # Branches [r, r, 0] give 2r at image guidance 2.
+                # Returning latents would cancel r after two scheduler steps.
+                return torch.zeros_like(latents)
+            ref_means = torch.tensor(
+                [float(r[0].float().mean()) if r is not None else 0.0 for r in ref],
+                dtype=latents.dtype,
+                device=latents.device,
+            ).view(-1, 1, 1, 1)
+            return torch.zeros_like(latents) + ref_means
+
+    class _LatentsOnlyTransformer(_FakeTransformer):
+        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+            return latents
+
+    def run_isolated(transformer_cls, ref_fill_a, seed_a, ref_fill_b, seed_b):
+        pipeline = _make_edit_forward_pipeline()
+        pipeline.transformer = transformer_cls()
+        pipeline.scheduler = _ApplyingScheduler()
+        pipeline.vae = _RefContentVAE()
+        kw = dict(
+            height=64, width=64, num_inference_steps=2, guidance_scale=5.0, guidance_scale_2=2.0, output_type="latent"
+        )
+        req = _wrap_request_batch(
+            [
+                (
+                    _edit_prompt("make it winter", "ugly", ref_fill_a),
+                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_a)),
+                ),
+                (
+                    _edit_prompt("a dog in a park", "blurry", ref_fill_b),
+                    _sampling(**kw, generator=torch.Generator().manual_seed(seed_b)),
+                ),
+            ]
+        )
+        return pipeline.forward(req)[0].output
+
+    ref_baseline = run_isolated(_RefOnlyTransformer, 0.25, 1, 0.75, 2)
+    assert torch.equal(ref_baseline, run_isolated(_RefOnlyTransformer, 0.25, 1, 0.5, 2))  # B's ref must not affect A
+    assert not torch.equal(
+        ref_baseline, run_isolated(_RefOnlyTransformer, 0.75, 1, 0.75, 2)
+    )  # A's own ref must affect A
+
+    seed_baseline = run_isolated(_LatentsOnlyTransformer, 0.25, 1, 0.75, 2)
+    # Only A's seed changes A.
+    assert torch.equal(seed_baseline, run_isolated(_LatentsOnlyTransformer, 0.25, 1, 0.75, 999))
+    assert not torch.equal(seed_baseline, run_isolated(_LatentsOnlyTransformer, 0.25, 999, 0.75, 2))
+
+    # RED: compare baseline and partner changes under each active mutation.
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+    from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+
+    real_build_ref_latents = BooguImagePipeline._build_ref_latents
+
+    def reversed_build_ref_latents(self, preprocessed_images, num_images_per_prompt, device, generators=None):
+        swapped = list(reversed(preprocessed_images))
+        return real_build_ref_latents(self, swapped, num_images_per_prompt, device, generators)
+
+    monkeypatch.setattr(BooguImagePipeline, "_build_ref_latents", reversed_build_ref_latents)
+    buggy_ref_baseline = run_isolated(_RefOnlyTransformer, 0.25, 1, 0.75, 2)
+    assert not torch.equal(buggy_ref_baseline, run_isolated(_RefOnlyTransformer, 0.25, 1, 0.5, 2))
+    monkeypatch.undo()
+
+    real_collate = DiffusionRequestBatch.collate_request_generators
+
+    def reversed_collate(self, num_outputs_per_prompt, default_generator):
+        result = real_collate(self, num_outputs_per_prompt, default_generator)
+        return list(reversed(result)) if isinstance(result, list) and len(result) > 1 else result
+
+    monkeypatch.setattr(DiffusionRequestBatch, "collate_request_generators", reversed_collate)
+    buggy_seed_baseline = run_isolated(_LatentsOnlyTransformer, 0.25, 1, 0.75, 2)
+    assert not torch.equal(buggy_seed_baseline, run_isolated(_LatentsOnlyTransformer, 0.25, 1, 0.75, 999))
+
+
+# ---------------------------------------------------------------------------
+# Request-batch: CFG-parallel dispatch
+# ---------------------------------------------------------------------------
+
+
+def _patch_cfg_parallel_state(monkeypatch, *, world_size, rank, fake_group):
+    """Make ``CFGParallelMixin`` see an active M-rank CFG group.
+
+    ``is_cfg_group_initialized()`` reads ``parallel_state._CFG`` through that
+    function's own closure, so patching ``_CFG`` flips it regardless of which
+    module imported the function. ``get_classifier_free_guidance_world_size``,
+    ``get_classifier_free_guidance_rank``, and ``get_cfg_group`` are each
+    rebound by name inside ``cfg_parallel.py`` (``from ...parallel_state
+    import ...``), so those three must be patched on both modules or
+    ``cfg_parallel``'s copy keeps resolving to the real (single-rank)
+    implementation.
+    """
+    from vllm_omni.diffusion.distributed import cfg_parallel as cfg_parallel_mod
+    from vllm_omni.diffusion.distributed import parallel_state
+
+    monkeypatch.setattr(parallel_state, "_CFG", fake_group)
+    for module in (parallel_state, cfg_parallel_mod):
+        monkeypatch.setattr(module, "get_classifier_free_guidance_world_size", lambda: world_size)
+        monkeypatch.setattr(module, "get_classifier_free_guidance_rank", lambda: rank)
+        monkeypatch.setattr(module, "get_cfg_group", lambda: fake_group)
+
+
+class _SlotTableCfgGroup:
+    """Fake CFG process group backed by a precomputed per-rank slot table.
+
+    Mirrors the real ``all_gather`` contract: called once per (slot, tensor
+    element) in slot-major order. Validates that the calling rank's local
+    prediction (``input_``) matches what that rank was supposed to compute
+    for this slot, and uses the *actual* ``input_`` (not the precomputed
+    value) for its own rank in the returned list -- a dispatch bug (wrong
+    branch id, wrong rank assignment) fails the assertion below instead of
+    being silently overwritten by the precomputed "correct" table.
+    """
+
+    def __init__(self, world_size, rank_in_group, slots_per_rank):
+        self.world_size = world_size
+        self.rank_in_group = rank_in_group
+        self._slots_per_rank = slots_per_rank
+        self.gather_calls = 0
+
+    def all_gather(self, input_, dim=0, separate_tensors=False):
+        del dim
+        assert separate_tensors
+        slot_idx = self.gather_calls
+        self.gather_calls += 1
+        expected_local = self._slots_per_rank[self.rank_in_group][slot_idx]
+        torch.testing.assert_close(
+            input_,
+            expected_local,
+            rtol=0,
+            atol=0,
+            msg=(
+                f"rank {self.rank_in_group} computed the wrong local branch prediction "
+                f"for slot {slot_idx} (dispatch bug)"
+            ),
+        )
+        return [
+            input_ if rank == self.rank_in_group else self._slots_per_rank[rank][slot_idx]
+            for rank in range(self.world_size)
+        ]
+
+
+def _make_counting_predict_noise(real_predict_noise, call_log):
+    def counting_predict_noise(**kwargs):
+        call_log.append(kwargs)
+        return real_predict_noise(**kwargs)
+
+    return counting_predict_noise
+
+
+def test_forward_batched_ti2i_double_guidance_cfg_parallel_matches_sequential(monkeypatch):
+    """CFG-parallel dispatch/gather must reproduce the sequential result.
+
+    Builds a real TI2I B=2 double-guidance request batch (distinct prompts,
+    negatives, seeds, and references per row, via the same signal-visible
+    fakes as ``test_forward_batched_ti2i_partner_reference_isolation``), then
+    replays the exact same branch kwargs through the real
+    ``_dispatch_branches``/``all_gather`` path for 2- and 3-rank CFG-parallel
+    layouts. A wrongly dispatched branch or a row swap during gather would change
+    the result, since request A and B carry distinct content/reference
+    signals into every branch.
+    """
+    from vllm_omni.diffusion.distributed.cfg_parallel import _dispatch_branches
+    from vllm_omni.diffusion.models.boogu_image.pipeline_boogu_image import BooguImagePipeline
+
+    class _ContentRefTransformer(_FakeTransformer):
+        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+            # FP32 throughout: the encode path is BF16, and the content term
+            # (derived from crc32-scaled token ids) reaches several hundred in
+            # magnitude, which would round away the 0.25/0.75 reference delta
+            # in BF16 and make the with-ref/without-ref branches indistinguishable.
+            out = latents.float() + instruction_embeds.float().mean(dim=(1, 2)).view(-1, 1, 1, 1)
+            ref = kwargs.get("ref_image_hidden_states")
+            if ref is not None:
+                ref_means = torch.tensor(
+                    [float(r[0].float().mean()) if r is not None else 0.0 for r in ref],
+                    dtype=out.dtype,
+                    device=out.device,
+                ).view(-1, 1, 1, 1)
+                out = out + ref_means
+            return out
+
+    pipeline = _make_edit_forward_pipeline()
+    pipeline.transformer = _ContentRefTransformer()
+    pipeline.vae = _RefContentVAE()
+    kw = dict(
+        height=64,
+        width=64,
+        num_inference_steps=1,
+        guidance_scale=5.0,
+        guidance_scale_2=2.0,
+        output_type="latent",
+    )
+    req = _wrap_request_batch(
+        [
+            (
+                _edit_prompt("make it winter", "ugly", 0.25),
+                _sampling(**kw, generator=torch.Generator().manual_seed(1)),
+            ),
+            (
+                _edit_prompt("a dog in a park", "blurry", 0.75),
+                _sampling(**kw, generator=torch.Generator().manual_seed(2)),
+            ),
+        ]
+    )
+
+    # Step 1: run the real sequential (single-rank) forward once, capturing
+    # the exact double-guidance branch kwargs and the resulting model_pred
+    # for this request batch's one denoising step.
+    captured: dict[str, Any] = {}
+    real_multi_branch = BooguImagePipeline.predict_noise_with_multi_branch_cfg
+
+    def capturing_multi_branch(
+        self, do_true_cfg, true_cfg_scale, branches_kwargs, cfg_normalize=False, output_slice=None
+    ):
+        result = real_multi_branch(
+            self, do_true_cfg, true_cfg_scale, branches_kwargs, cfg_normalize=cfg_normalize, output_slice=output_slice
+        )
+        captured.update(
+            branches_kwargs=branches_kwargs,
+            true_cfg_scale=true_cfg_scale,
+            cfg_normalize=cfg_normalize,
+            output_slice=output_slice,
+            sequential_pred=result,
+        )
+        return result
+
+    monkeypatch.setattr(BooguImagePipeline, "predict_noise_with_multi_branch_cfg", capturing_multi_branch)
+    pipeline.forward(req)
+    monkeypatch.undo()
+
+    branches_kwargs = captured["branches_kwargs"]
+    assert len(branches_kwargs) == 3  # cond+ref, neg+ref, neg+no-ref
+    sequential_pred = captured["sequential_pred"]
+
+    # Step 2: replay the identical branches through the real dispatch/gather
+    # path for 2- and 3-rank layouts (rank 0: [0, 2] / [0]; rank 1: [1] / [1];
+    # rank 2 only exists for the 3-rank layout: [2]).
+    n_branches = len(branches_kwargs)
+    # Branch predictions are a deterministic function of (pipeline, branches_kwargs)
+    # alone; compute once and reuse across both rank-layout replays below.
+    with torch.no_grad():
+        branch_preds = [pipeline.predict_noise(**branch_kw) for branch_kw in branches_kwargs]
+
+    # If any two branches ever collapsed to the same tensor, a wrongly
+    # dispatched or wrongly gathered branch below could go unnoticed even with exact-equality
+    # checks -- confirm the fakes are actually signal-visible.
+    for i in range(n_branches):
+        for j in range(i + 1, n_branches):
+            assert not torch.equal(branch_preds[i], branch_preds[j]), (
+                f"branches {i} and {j} are numerically identical; the fakes cannot "
+                "distinguish a branch mix-up from a correct dispatch"
+            )
+
+    for cfg_world_size in (2, 3):
+        assignments = _dispatch_branches(n_branches, cfg_world_size)
+        max_per_rank = max(len(a) for a in assignments)
+
+        slots_per_rank = []
+        for branch_ids in assignments:
+            slots = [branch_preds[bid] for bid in branch_ids]
+            while len(slots) < max_per_rank:
+                slots.append(torch.zeros_like(branch_preds[0]))
+            slots_per_rank.append(slots)
+
+        rank_outputs = []
+        for cfg_rank in range(cfg_world_size):
+            fake_group = _SlotTableCfgGroup(cfg_world_size, cfg_rank, slots_per_rank)
+            _patch_cfg_parallel_state(monkeypatch, world_size=cfg_world_size, rank=cfg_rank, fake_group=fake_group)
+
+            call_log: list[dict] = []
+            pipeline.predict_noise = _make_counting_predict_noise(pipeline.predict_noise, call_log)
+            try:
+                with torch.no_grad():
+                    rank_pred = pipeline.predict_noise_with_multi_branch_cfg(
+                        do_true_cfg=True,
+                        true_cfg_scale=captured["true_cfg_scale"],
+                        branches_kwargs=branches_kwargs,
+                        cfg_normalize=captured["cfg_normalize"],
+                        output_slice=captured["output_slice"],
+                    )
+            finally:
+                del pipeline.predict_noise  # drop the instance override, restore the class method
+                monkeypatch.undo()
+
+            # A silent fallback to the sequential path (e.g. the CFG-group
+            # patch not taking effect) would compute every branch locally and
+            # never call all_gather; catch that instead of letting it pass
+            # only because sequential also happens to match the baseline.
+            expected_calls = len(assignments[cfg_rank])
+            assert len(call_log) == expected_calls, (
+                f"cfg_world_size={cfg_world_size} rank={cfg_rank} called predict_noise "
+                f"{len(call_log)}x, expected {expected_calls}x (silent sequential fallback?)"
+            )
+            assert fake_group.gather_calls == max_per_rank, (
+                f"cfg_world_size={cfg_world_size} rank={cfg_rank} issued {fake_group.gather_calls} "
+                f"all_gather call(s), expected {max_per_rank} (silent sequential fallback?)"
+            )
+            rank_outputs.append(rank_pred)
+
+        for cfg_rank, rank_pred in enumerate(rank_outputs):
+            torch.testing.assert_close(
+                rank_pred,
+                sequential_pred,
+                rtol=0,
+                atol=0,
+                msg=(
+                    f"cfg_world_size={cfg_world_size} rank={cfg_rank} diverged from the "
+                    "sequential (single-rank) baseline"
+                ),
+            )
+
+
+# ---------------------------------------------------------------------------
+# Request-batch: variable-length padding
+# ---------------------------------------------------------------------------
+
+
+def test_forward_batched_ti2i_variable_length_partner_padding(monkeypatch):
+    """A must be invariant to B's real padded length, not just same-length content.
+
+    ``test_forward_batched_ti2i_partner_reference_isolation`` pins every fake
+    encoder call to a fixed ``_SEQ_LEN`` with an all-ones mask, so it never
+    exercises production's ``padding="longest"`` batching or a real 0/1
+    ``instruction_attention_mask``. Here the fake processor tokenizes by
+    content length and right-pads to the batch max, and the fake transformer
+    does a masked reduction over ``instruction_attention_mask`` instead of a
+    blind ``.mean()`` -- so A's row can only stay correct if the real
+    encode/reshape/CFG-duplication plumbing threads the real per-row mask
+    through unchanged when B's length (or B's negative-prompt length) grows.
+    """
+
+    class _VariableLengthImageAwareProcessor:
+        """Tokenizes by word count and right-pads to the batch max length."""
+
+        def __init__(self, force_all_ones_mask=False):
+            self.calls = []
+            self.force_all_ones_mask = force_all_ones_mask
+
+        def apply_chat_template(self, prompts, **kwargs):
+            has_image = []
+            rows = []
+            for messages in prompts:
+                user_content = messages[1]["content"]
+                has_image.append(any(c.get("type") == "image" for c in user_content))
+                system_text = messages[0]["content"][0]["text"]
+                user_text = next(c["text"] for c in messages[1]["content"] if c.get("type") == "text")
+                length = 2 + len(user_text.split())
+                ids = torch.zeros(length, dtype=torch.long)
+                ids[0] = _content_id(system_text)
+                ids[1] = _content_id(user_text)
+                if length > 2:
+                    ids[2:] = torch.arange(2, length)
+                rows.append(ids)
+
+            batch = len(prompts)
+            max_len = max(len(r) for r in rows)
+            input_ids = torch.zeros(batch, max_len, dtype=torch.long)
+            attention_mask = torch.zeros(batch, max_len, dtype=torch.long)
+            for i, ids in enumerate(rows):
+                input_ids[i, : len(ids)] = ids
+                attention_mask[i, : len(ids)] = 1
+            if self.force_all_ones_mask:
+                attention_mask = torch.ones_like(attention_mask)
+
+            self.calls.append(
+                {"prompts": prompts, "kwargs": kwargs, "has_image": has_image, "lengths": [len(r) for r in rows]}
+            )
+            return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+    class _MaskedContentRefTransformer(_FakeTransformer):
+        """Masked-mean over real padding, unlike the fixed-``_SEQ_LEN`` fakes."""
+
+        def __call__(self, latents, timestep, instruction_embeds, freqs_real, instruction_attention_mask, **kwargs):
+            mask = instruction_attention_mask.float().unsqueeze(-1)
+            valid_counts = mask.sum(dim=1).clamp(min=1.0)
+            content = ((instruction_embeds.float() * mask).sum(dim=1) / valid_counts).mean(dim=-1).view(-1, 1, 1, 1)
+            out = latents.float() + content
+            ref = kwargs.get("ref_image_hidden_states")
+            if ref is not None:
+                ref_means = torch.tensor(
+                    [float(r[0].float().mean()) if r is not None else 0.0 for r in ref],
+                    dtype=out.dtype,
+                    device=out.device,
+                ).view(-1, 1, 1, 1)
+                out = out + ref_means
+            return out
+
+    class _ApplyingScheduler(_FakeScheduler):
+        def step(self, model_output, t, latents, return_dict=False):
+            return (model_output,)
+
+    A = _edit_prompt("make it winter", "ugly", 0.25)
+
+    def run(partner=None, force_all_ones_mask=False):
+        pipeline = _make_edit_forward_pipeline()
+        pipeline.processor = _VariableLengthImageAwareProcessor(force_all_ones_mask=force_all_ones_mask)
+        pipeline.transformer = _MaskedContentRefTransformer()
+        pipeline.scheduler = _ApplyingScheduler()
+        pipeline.vae = _RefContentVAE()
+        kw = dict(
+            height=64,
+            width=64,
+            num_inference_steps=1,
+            guidance_scale=5.0,
+            guidance_scale_2=2.0,
+            output_type="latent",
+        )
+        items = [(A, _sampling(**kw, generator=torch.Generator().manual_seed(1)))]
+        if partner is not None:
+            items.append((partner, _sampling(**kw, generator=torch.Generator().manual_seed(2))))
+        req = _wrap_request_batch(items)
+        outs = pipeline.forward(req)
+        return outs[0].output, pipeline.processor.calls
+
+    a_alone, _ = run(partner=None)
+
+    short_b = _edit_prompt("a dog", "blurry", 0.75)
+    long_b = _edit_prompt("a dog running through a large green park full of tall autumn trees", "blurry", 0.75)
+    long_neg_b = _edit_prompt(
+        "a dog in a park", "extremely blurry ugly low quality watermark grainy artifact noisy", 0.75
+    )
+
+    a_with_short_b, calls_short = run(partner=short_b)
+    a_with_long_b, calls_long = run(partner=long_b)
+    a_with_long_neg_b, calls_long_neg = run(partner=long_neg_b)
+
+    # calls[0] is the positive-instruction batch call, calls[1] the negative.
+    width_short = max(calls_short[0]["lengths"])
+    width_long = max(calls_long[0]["lengths"])
+    assert width_long > width_short, "B's longer instruction must actually widen the padded batch"
+
+    a_len_in_short = calls_short[0]["lengths"][0]
+    a_len_in_long = calls_long[0]["lengths"][0]
+    assert a_len_in_short == a_len_in_long  # A's own content length never changes...
+    assert width_long > a_len_in_long  # ...but A now sits inside a wider, padded row.
+
+    neg_width_short = max(calls_short[1]["lengths"])
+    neg_width_long_neg = max(calls_long_neg[1]["lengths"])
+    assert neg_width_long_neg > neg_width_short, "B's longer negative prompt must widen the negative-branch batch too"
+
+    assert torch.equal(a_alone, a_with_short_b)
+    assert torch.equal(a_alone, a_with_long_b)
+    assert torch.equal(a_alone, a_with_long_neg_b)
+
+    # Negative control: force an all-ones mask (ignore real padding). A's row
+    # then sees B's padding-region positions as "real" content, so this must
+    # break invariance -- otherwise the checks above could pass vacuously
+    # even with no masking at all, exactly like the fixed-_SEQ_LEN fakes.
+    a_with_long_b_broken_mask, _ = run(partner=long_b, force_all_ones_mask=True)
+    assert not torch.equal(a_alone, a_with_long_b_broken_mask)

vllm_omni/diffusion/models/boogu_image/pipeline_boogu_image.py

@@ -112,18 +112,15 @@ def get_boogu_image_post_process_func(od_config: OmniDiffusionConfig):
     return post_process_func
 
 
-def _boogu_batch_compatibility_key(has_reference: bool, request_id: str) -> tuple:
-    """Request-batch isolation key. ``forward`` reads shared guidance/shape fields
-    from the batch's first request, so t2i and ti2i must not share a key.
+def _boogu_batch_compatibility_key(has_reference: bool) -> tuple:
+    """Separate T2I and TI2I denoise paths.
 
-    ti2i is held at batch=1 (request-unique key) because
-    ``guidance_scale_2_provided`` is absent from ``RequestBatchSamplingParamsKey``
-    while ``forward`` reads it from the first request, so mixed-``_provided`` edits
-    with equal numeric ``guidance_scale_2`` would co-batch into the wrong mode.
+    Shared shape/guidance fields, including ``guidance_scale_2_provided``,
+    are checked by ``RequestBatchSamplingParamsKey``.
     """
     if not has_reference:
         return ("boogu_image", "t2i")
-    return ("boogu_image", "ti2i", request_id)
+    return ("boogu_image", "ti2i")
 
 
 def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig):
@@ -157,14 +154,14 @@ def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig):
         prompt = request.prompt
         if isinstance(prompt, str):
             # Plain-text prompt cannot carry an image -> text-to-image.
-            request.batch_compatibility_key = _boogu_batch_compatibility_key(False, request.request_id)
+            request.batch_compatibility_key = _boogu_batch_compatibility_key(False)
             return request
 
         multi_modal_data = prompt.get("multi_modal_data") or {}
         raw_image = multi_modal_data.get("image")
         if not raw_image:
             # No reference image -> text-to-image (Base checkpoint).
-            request.batch_compatibility_key = _boogu_batch_compatibility_key(False, request.request_id)
+            request.batch_compatibility_key = _boogu_batch_compatibility_key(False)
             return request
 
         if isinstance(raw_image, list):
@@ -199,7 +196,7 @@ def get_boogu_image_pre_process_func(od_config: OmniDiffusionConfig):
         prompt["additional_information"]["preprocessed_image"] = preprocessed_image
         prompt["additional_information"]["prompt_image"] = prompt_image
         request.prompt = prompt
-        request.batch_compatibility_key = _boogu_batch_compatibility_key(True, request.request_id)
+        request.batch_compatibility_key = _boogu_batch_compatibility_key(True)
         return request
 
     return pre_process_func
@@ -961,14 +958,6 @@ class BooguImagePipeline(CFGParallelMixin, nn.Module, ProgressBarMixin, Supports
         if self._is_turbo and req.num_reqs > 1:
             raise RuntimeError("BooguImageTurboPipeline does not support request batching.")
 
-        # Fail-closed: a batched ti2i must never reach here (it is gated to batch=1).
-        if has_reference and req.num_reqs > 1:
-            raise RuntimeError(
-                f"BooguImagePipeline received a batched TI2I (edit) request "
-                f"(num_reqs={req.num_reqs}); TI2I batching is gated to batch=1 "
-                "pending guidance-mode / compatibility-key validation."
-            )
-
         sampling_params_list = req.sampling_params_list
         # Shared shape/step/guidance fields are guaranteed identical across the
         # batch by RequestBatchSamplingParamsKey; request-local values (seeds,