[Model] Support request-level batching for MammothModa2 DiT on shared runtime (#7476)

Signed-off-by: BeatSeat <wendavid552@gmail.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>

tests/diffusion/models/mammoth_moda2/test_pipeline_mammothmoda2_dit.py

@@ -4,6 +4,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import cache_dit
 import pytest
@@ -26,7 +27,10 @@ from vllm_omni.diffusion.models.mammoth_moda2 import pipeline_mammothmoda2_dit
 from vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit import (
     MammothModa2DiTPipeline,
     _build_mammoth_config,
+    _MammothRequest,
+    _pad_cond_sequence,
     _root_weight_source,
+    get_mammoth_moda2_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
@@ -62,9 +66,7 @@ def _od_config() -> OmniDiffusionConfig:
 
 def test_build_mammoth_config_uses_shared_transformer_projection() -> None:
     config = _build_mammoth_config(_od_config())
-    assert config.model_type == "mammothmoda2"
-    assert config.get_text_config().hidden_size == 8
-    assert config.gen_dit_config == {"hidden_size": 8, "in_channels": 4}
+    assert config.gen_dit_config["hidden_size"] == 8
 
 
 def test_build_mammoth_config_rejects_empty_shared_projection() -> None:
@@ -88,11 +90,11 @@ def test_root_weight_source_forwards_revision() -> None:
     assert _root_weight_source(config).revision == "rev-7"
 
 
-def test_pipeline_declares_native_components_and_single_request_mode_only() -> None:
+def test_pipeline_declares_native_components_and_batch_request_mode() -> None:
     assert MammothModa2DiTPipeline._dit_modules == ["gen_transformer"]
     assert MammothModa2DiTPipeline._encoder_modules == ["gen_image_condition_refiner"]
     assert MammothModa2DiTPipeline._vae_modules == ["gen_vae"]
-    assert MammothModa2DiTPipeline.supports_request_batch is False
+    assert MammothModa2DiTPipeline.supports_request_batch is True
     assert MammothModa2DiTPipeline.supports_step_execution is False
 
 
@@ -216,13 +218,17 @@ def test_parse_request_standard_fields_precede_request_level_fallbacks() -> None
     assert parsed.num_inference_steps == 7
 
 
-def test_parse_request_rejects_multiple_requests_and_outputs() -> None:
+def test_parse_request_parses_each_request_and_rejects_multi_output() -> None:
     pipeline = _pipeline_shell()
     batch = _batch()
     batch.requests.append(_batch(request_id="req-8").requests[0])
-    with pytest.raises(ValueError, match="exactly one request"):
-        pipeline._parse_request(batch)
-    with pytest.raises(ValueError, match="num_outputs_per_prompt=1"):
+    req0 = pipeline._parse_request(batch, index=0)
+    req1 = pipeline._parse_request(batch, index=1)
+    assert req0.request_id == "req-7"
+    assert req1.request_id == "req-8"
+    assert req0.index == 0
+    assert req1.index == 1
+    with pytest.raises(ValueError, match="num_outputs_per_prompt == 1"):
         pipeline._parse_request(_batch(sampling=OmniDiffusionSamplingParams(num_outputs_per_prompt=2)))
 
 
@@ -347,10 +353,14 @@ class _FakeVae(nn.Module):
         super().__init__()
         self.anchor = nn.Parameter(torch.zeros(1))
         self.config = _FakeVaeConfig()
+        self.last_decode_dtype: torch.dtype | None = None
 
     def decode(self, latents, return_dict=False):
         assert return_dict is False
-        return (torch.zeros(1, 3, 32, 48, dtype=latents.dtype),)
+        self.last_decode_dtype = latents.dtype
+        b = latents.shape[0]
+        base = latents[:, :1, :1, :1]
+        return (torch.zeros(b, 3, 32, 48, dtype=latents.dtype) + base,)
 
 
 class _FakeScheduler:
@@ -392,8 +402,9 @@ def test_forward_returns_diffusion_output_with_request_sampling(mocker) -> None:
         )
     )
 
-    assert isinstance(result, DiffusionOutput)
-    assert result.output.shape == (1, 3, 32, 48)
+    assert len(result) == 1
+    assert isinstance(result[0], DiffusionOutput)
+    assert result[0].output.shape == (1, 3, 32, 48)
     assert captured["shape"] == (1, 4, 4, 6)
     assert captured["seed"] == 42
     assert scheduler.requested_steps == 2
@@ -527,7 +538,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Reference run for phase 2: identical request while hooks were never
     # installed on this pipeline.
     assert not pipeline.is_cache_dit_enabled()
-    reference = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    reference = pipeline.forward(_batch(sampling=unguided_sampling))[0].output
     assert seq_lens == [cond_tokens] * 8
     assert torch.isfinite(reference.float()).all()
 
@@ -551,7 +562,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Phase 2: no-CFG request disables the hooks and must reproduce the
     # never-enabled reference bit-for-bit.
     seq_lens.clear()
-    unguided = pipeline.forward(_batch(sampling=unguided_sampling)).output
+    unguided = pipeline.forward(_batch(sampling=unguided_sampling))[0].output
     assert seq_lens == [cond_tokens] * 8
     assert not pipeline.is_cache_dit_enabled()
     assert not BlockAdapter.is_cached(pipeline.gen_transformer)
@@ -560,7 +571,7 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     # Phase 3: CFG with a new step count and image size re-installs the hooks
     # on a fresh context; output shape follows the new request geometry.
     seq_lens.clear()
-    regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64)).output
+    regenerated = pipeline.forward(_batch(sampling=changed_cfg_sampling, height=48, width=64))[0].output
     assert seq_lens == [cond_tokens, uncond_tokens] * 10
     assert pipeline.is_cache_dit_enabled()
     assert BlockAdapter.is_cached(pipeline.gen_transformer)
@@ -570,3 +581,442 @@ def test_forward_transitions_request_scoped_cache_dit_across_requests(request: p
     assert torch.isfinite(regenerated.float()).all()
 
     hook.remove()
+
+
+def test_pre_process_registers_batch_compatibility_key() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch().requests[0]
+    pre_process(req)
+    assert req.batch_compatibility_key == ("mammoth_moda2_dit", 32, 48, 7)
+
+
+def test_pre_process_rejects_missing_ar_conditions() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch(request_id="req-bad-ar", prompt={"prompt": "test"}).requests[0]
+    with pytest.raises(ValueError, match="Missing additional_information AR conditions.*req-bad-ar"):
+        pre_process(req)
+
+
+@pytest.mark.parametrize(
+    ("height", "width", "message"),
+    [
+        (0, 32, "Invalid image size.*req-dim"),
+        (32, -1, "Invalid image size.*req-dim"),
+        (30, 32, "multiples of 16.*req-dim"),
+        (32, 31, "multiples of 16.*req-dim"),
+    ],
+)
+def test_pre_process_rejects_invalid_dimensions(height, width, message) -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    batch = _batch(request_id="req-dim")
+    batch.prompts[0].update(height=height, width=width)
+    with pytest.raises(ValueError, match=message):
+        pre_process(batch.requests[0])
+
+
+def test_pre_process_rejects_non_positive_steps() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch(
+        request_id="req-zero-steps",
+        sampling=OmniDiffusionSamplingParams(num_inference_steps=0),
+    ).requests[0]
+    with pytest.raises(ValueError, match="num_inference_steps must be positive.*req-zero-steps"):
+        pre_process(req)
+
+
+def test_pre_process_allows_dummy_run() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+    req = _batch(
+        request_id="dummy_req_id",
+        prompt={"prompt": "dummy run"},
+        sampling=OmniDiffusionSamplingParams(height=512, width=512, num_inference_steps=20),
+    ).requests[0]
+    pre_process(req)
+    assert req.batch_compatibility_key == ("mammoth_moda2_dit", 512, 512, 20)
+
+
+def test_forward_casts_latents_to_vae_dtype() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    vae = _FakeVae()
+    vae.anchor = nn.Parameter(torch.zeros(1, dtype=torch.float16))
+    pipeline.gen_vae = vae
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+    ):
+        result = pipeline.forward(_batch())
+
+    assert len(result) == 1
+    assert vae.last_decode_dtype == torch.float16
+
+
+def test_group_requests_by_geometry_and_steps() -> None:
+    def make_req(idx: int, h: int, w: int, s: int) -> _MammothRequest:
+        return _MammothRequest(
+            index=idx,
+            request_id=f"r{idx}",
+            full_hidden_states=torch.zeros(1, 8),
+            full_token_ids=[1],
+            answer_start_index=0,
+            height=h,
+            width=w,
+            text_guidance_scale=1.0,
+            cfg_range=(0.0, 1.0),
+            num_inference_steps=s,
+            seed=None,
+            generator=None,
+            generator_device=None,
+        )
+
+    reqs = [make_req(0, 32, 32, 5), make_req(1, 32, 32, 5), make_req(2, 64, 64, 5)]
+    groups = MammothModa2DiTPipeline._group_requests(reqs)
+    assert sorted(groups) == [[0, 1], [2]]
+
+
+def test_pad_cond_sequence_right_pads() -> None:
+    embeds = [torch.ones(1, 2, 8), torch.ones(1, 4, 8) * 2]
+    masks = [torch.ones(1, 2, dtype=torch.bool), torch.ones(1, 4, dtype=torch.bool)]
+    padded, mask = _pad_cond_sequence(embeds, masks)
+    assert padded.shape == (2, 4, 8)
+    assert mask.tolist() == [[True, True, False, False], [True, True, True, True]]
+
+
+def test_forward_batched_matches_solo_output() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    def fake_randn_tensor(shape, *, generator, device, dtype):
+        if isinstance(generator, list):
+            tensors = [
+                torch.full((1, *shape[1:]), float(g.initial_seed() if g else 0), device=device, dtype=dtype)
+                for g in generator
+            ]
+            return torch.cat(tensors, dim=0)
+        val = float(generator.initial_seed() if generator else 0)
+        return torch.full(shape, val, device=device, dtype=dtype)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(f"{module}.randn_tensor", side_effect=fake_randn_tensor),
+    ):
+        req1 = _batch(
+            request_id="r1",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=10, guidance_scale=1.0, num_inference_steps=2
+            ),
+        ).requests[0]
+        req2 = _batch(
+            request_id="r2",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=20, guidance_scale=1.0, num_inference_steps=2
+            ),
+        ).requests[0]
+
+        batched_out = pipeline.forward(DiffusionRequestBatch([req1, req2]))
+        solo1 = pipeline.forward(DiffusionRequestBatch([req1]))
+        solo2 = pipeline.forward(DiffusionRequestBatch([req2]))
+
+    assert len(batched_out) == 2
+    torch.testing.assert_close(batched_out[0].output, solo1[0].output)
+    torch.testing.assert_close(batched_out[1].output, solo2[0].output)
+    assert not torch.equal(batched_out[0].output, batched_out[1].output)
+
+
+def test_refiner_output_length_preserved_when_input_len_differs_from_queries() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+
+    class _MockRefiner(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.anchor = nn.Parameter(torch.zeros(1))
+
+        def forward(self, features, mask=None):
+            b, _, d = features.shape
+            return torch.zeros(b, 8, d, device=features.device, dtype=features.dtype)
+
+    pipeline.gen_image_condition_refiner = _MockRefiner()
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    captured_cond_lens = []
+
+    def fake_forward(*, hidden_states, text_hidden_states, **kwargs):
+        captured_cond_lens.append(text_hidden_states.shape[1])
+        return torch.zeros_like(hidden_states)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+        patch.object(pipeline.gen_transformer, "forward", fake_forward),
+    ):
+        prompt = {
+            "prompt": "",
+            "additional_information": {
+                "full_hidden_states": torch.zeros(6, 8),
+                "full_token_ids": [10, 11, 100, 101, 102, 103],
+                "answer_start_index": 2,
+            },
+        }
+        batch = _batch(
+            prompt=prompt,
+            sampling=OmniDiffusionSamplingParams(height=32, width=48, guidance_scale=1.0, num_inference_steps=1),
+        )
+        pipeline.forward(batch)
+
+    # 2 text tokens + 8 refined queries = 10 total conditioning tokens seen by transformer
+    assert captured_cond_lens == [10]
+
+
+def test_nested_image_embedder_compacts_unequal_text_lengths() -> None:
+    from vllm_omni.diffusion.models.mammoth_moda2.mammothmoda2_dit_model import (
+        Transformer2DModel,
+    )
+
+    model = object.__new__(Transformer2DModel)
+    nn.Module.__init__(model)
+    object.__setattr__(model, "_internal_dict", SimpleNamespace(patch_size=2))
+
+    dim = 8
+    batch_size = 2
+    max_text_len = 4
+
+    class _MockTimeCaptionEmbed(nn.Module):
+        def forward(self, timestep, text_hidden_states, dtype):
+            return torch.zeros(batch_size, dim), text_hidden_states
+
+        def image_embedder(self, feats, mask=None):
+            return torch.full((batch_size, 3, dim), 99.0)
+
+    model.time_caption_embed = _MockTimeCaptionEmbed()
+    model.x_embedder = nn.Identity()
+
+    def mock_rope(freqs, mask, *args):
+        return None, None, None, None, mask.sum(dim=-1).tolist(), []
+
+    model.rope_embedder = mock_rope
+
+    hidden_states = torch.zeros(batch_size, 4, 4, 4)
+    timestep = torch.tensor([1.0, 1.0])
+    text_hidden_states = torch.zeros(batch_size, max_text_len, dim)
+    text_hidden_states[0, 0, :] = 1.0
+    text_hidden_states[0, 1, :] = 2.0
+    text_hidden_states[1, 0, :] = 10.0
+    text_hidden_states[1, 1, :] = 20.0
+    text_hidden_states[1, 2, :] = 30.0
+    text_hidden_states[1, 3, :] = 40.0
+
+    text_attention_mask = torch.tensor(
+        [
+            [True, True, False, False],
+            [True, True, True, True],
+        ]
+    )
+
+    ar_image_hidden_states = torch.zeros(batch_size, 5, dim)
+
+    res = model._prepare_embeddings(
+        hidden_states=hidden_states,
+        timestep=timestep,
+        text_hidden_states=text_hidden_states,
+        text_attention_mask=text_attention_mask,
+        freqs_cis=torch.zeros(1),
+        batch_size=batch_size,
+        height=4,
+        width=4,
+        ar_image_hidden_states=ar_image_hidden_states,
+    )
+
+    compacted_text = res[1]
+    compacted_mask = res[2]
+    encoder_seq_lengths = res[9]
+
+    assert encoder_seq_lengths == [5, 7]
+    assert compacted_text.shape == (2, 7, dim)
+    assert compacted_text[0, 0, 0].item() == 1.0
+    assert compacted_text[0, 1, 0].item() == 2.0
+    assert compacted_text[0, 2, 0].item() == 99.0
+    assert compacted_text[0, 3, 0].item() == 99.0
+    assert compacted_text[0, 4, 0].item() == 99.0
+    assert compacted_mask[0].tolist() == [True, True, True, True, True, False, False]
+    assert compacted_mask[1].tolist() == [True, True, True, True, True, True, True]
+
+
+def test_pre_process_rejects_missing_visual_tokens_and_multi_output() -> None:
+    pre_process = get_mammoth_moda2_pre_process_func(_od_config())
+
+    multi_out = _batch(
+        request_id="req-multi",
+        sampling=OmniDiffusionSamplingParams(num_outputs_per_prompt=2),
+    ).requests[0]
+    with pytest.raises(ValueError, match="num_outputs_per_prompt == 1.*req-multi"):
+        pre_process(multi_out)
+
+    no_vis_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(3, 8),
+            "full_token_ids": [10, 11, 12],
+            "answer_start_index": 1,
+        },
+    }
+    no_vis_req = _batch(request_id="req-no-vis", prompt=no_vis_prompt).requests[0]
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-no-vis"):
+        pre_process(no_vis_req)
+
+    valid_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(3, 8),
+            "full_token_ids": [10, 100, 101],
+            "answer_start_index": 1,
+        },
+    }
+    valid_req = _batch(request_id="req-valid", prompt=valid_prompt).requests[0]
+    pre_process(valid_req)
+    assert valid_req.batch_compatibility_key is not None
+
+
+def test_cfg_precomputation_uses_cpu_booleans_without_device_sync() -> None:
+    pipeline = _pipeline_shell()
+    pipeline.gen_transformer = _FakeTransformer()
+    pipeline.gen_image_condition_refiner = None
+    pipeline.gen_vae = _FakeVae()
+    pipeline.gen_freqs_cis = torch.zeros(1)
+    scheduler = _FakeScheduler()
+
+    raw_any = torch.Tensor.any
+    tensor_any_called = []
+
+    def wrapped_any(self, *args, **kwargs):
+        tensor_any_called.append(self)
+        return raw_any(self, *args, **kwargs)
+
+    module = "vllm_omni.diffusion.models.mammoth_moda2.pipeline_mammothmoda2_dit"
+    with (
+        patch(f"{module}.FlowMatchEulerDiscreteScheduler", return_value=scheduler),
+        patch(
+            f"{module}.randn_tensor",
+            side_effect=lambda s, **kw: torch.zeros(s, device=kw.get("device"), dtype=kw.get("dtype")),
+        ),
+        patch.object(torch.Tensor, "any", wrapped_any),
+    ):
+        req1 = _batch(
+            request_id="r1",
+            sampling=OmniDiffusionSamplingParams(
+                height=32,
+                width=48,
+                seed=1,
+                guidance_scale=4.0,
+                num_inference_steps=5,
+                extra_args={"cfg_range": [0.2, 0.8]},
+            ),
+        ).requests[0]
+        req2 = _batch(
+            request_id="r2",
+            sampling=OmniDiffusionSamplingParams(
+                height=32, width=48, seed=2, guidance_scale=1.0, num_inference_steps=5
+            ),
+        ).requests[0]
+
+        pipeline.forward(DiffusionRequestBatch([req1, req2]))
+
+    assert len(tensor_any_called) == 0
+
+
+def test_dummy_request_splits_valid_ar_conditions_without_error() -> None:
+    """Startup warmup dummy run must split into valid non-empty AR conditions."""
+    pipeline = _pipeline_shell()
+    batch = _batch(
+        request_id="dummy_req_id",
+        prompt={"prompt": "dummy run"},
+        sampling=OmniDiffusionSamplingParams(height=512, width=512, seed=1, guidance_scale=0.0, num_inference_steps=2),
+    )
+    parsed = pipeline._parse_request(batch)
+    text_cond, image_cond = pipeline._split_request_conditions(parsed)
+    assert text_cond.shape == (1, pipeline._llm_hidden_size)
+    assert image_cond.shape == (1, pipeline._llm_hidden_size)
+
+
+@pytest.mark.parametrize(
+    ("llm_cfg_patch", "expected_threshold"),
+    [
+        # Parent-level llm_config threshold
+        ({"gen_vocab_start_index": 200}, 200),
+        # Nested text_config threshold
+        ({"text_config": {"model_type": "mammothmoda2_qwen2_5_vl_text", "gen_vocab_start_index": 350}}, 350),
+        # Default-derived threshold (fallback to vocab_size / 152064)
+        ({}, 152064),
+    ],
+)
+def test_admission_and_inference_threshold_resolution_matches(llm_cfg_patch: dict, expected_threshold: int) -> None:
+    """Verify identical threshold resolution for parent-level, nested, and default configs."""
+    raw = _raw_config()
+    raw["llm_config"] = {"model_type": "mammothmoda2_qwen2_5_vl", **llm_cfg_patch}
+    od_cfg = OmniDiffusionConfig(
+        model="/models/MammothModa2-Preview",
+        model_class_name="MammothModa2DiTPipeline",
+        tf_model_config=TransformerConfig.from_dict(raw),
+    )
+    # 1. Check admission preprocessor resolution
+    pre_process = get_mammoth_moda2_pre_process_func(od_cfg)
+
+    # Token just below threshold -> rejected
+    below_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(2, 8),
+            "full_token_ids": [10, expected_threshold - 1],
+            "answer_start_index": 1,
+        },
+    }
+    below_req = _batch(request_id="req-below", prompt=below_prompt).requests[0]
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-below"):
+        pre_process(below_req)
+
+    # Token at or above threshold -> accepted
+    at_prompt = {
+        "prompt": "",
+        "additional_information": {
+            "full_hidden_states": torch.zeros(2, 8),
+            "full_token_ids": [10, expected_threshold],
+            "answer_start_index": 1,
+        },
+    }
+    at_req = _batch(request_id="req-at", prompt=at_prompt).requests[0]
+    pre_process(at_req)
+    assert at_req.batch_compatibility_key is not None
+
+    # 2. Check inference pipeline model config resolution
+    pipeline = _pipeline_shell()
+    pipeline.config = _build_mammoth_config(od_cfg)
+
+    # Token below threshold rejected by pipeline
+    parsed_below = pipeline._parse_request(_batch(request_id="req-inf-below", prompt=below_prompt))
+    with pytest.raises(ValueError, match="no visual-token hidden states.*req-inf-below"):
+        pipeline._split_request_conditions(parsed_below)
+
+    # Token at threshold accepted by pipeline
+    parsed_at = pipeline._parse_request(_batch(request_id="req-inf-at", prompt=at_prompt))
+    text_cond, image_cond = pipeline._split_request_conditions(parsed_at)
+    assert text_cond.shape == (1, 8)
+    assert image_cond.shape == (1, 8)