[BugFix]: CUDA device-side assert failures on single-stage BAGEL i2i requests (#3680)

Signed-off-by: NumberWan <wantszkin2003@gmail.com>

tests/diffusion/models/bagel/test_forward_cache_update_vae.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for Bagel.forward_cache_update_vae packed-sequence assembly."""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+from transformers import PretrainedConfig
+from vllm.transformers_utils.configs.bagel import BagelConfig
+
+from vllm_omni.diffusion.models.bagel.bagel_transformer import (
+    Bagel,
+    BaseNavitOutputWithPast,
+    NaiveCache,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+HIDDEN_SIZE = 32
+NUM_LAYERS = 2
+
+
+def _make_bagel_config() -> BagelConfig:
+    llm_config = PretrainedConfig()
+    llm_config.hidden_size = HIDDEN_SIZE
+    llm_config.layer_module = "Qwen2MoTDecoderLayer"
+    llm_config.num_hidden_layers = NUM_LAYERS
+    llm_config.num_attention_heads = 4
+
+    vae_config = PretrainedConfig()
+    vae_config.downsample = 8
+    vae_config.z_channels = 4
+
+    return BagelConfig(
+        llm_config=llm_config,
+        vae_config=vae_config,
+        visual_gen=True,
+        visual_und=False,
+        latent_patch_size=2,
+        max_latent_size=32,
+    )
+
+
+def _make_tracking_language_model() -> MagicMock:
+    """Return a mock LM that records forward() kwargs."""
+    calls: list[dict] = []
+    hidden_size = HIDDEN_SIZE
+
+    def forward(*_args, **kwargs):
+        calls.append(dict(kwargs))
+        if kwargs.get("return_embeddings_only"):
+            text_ids = kwargs["packed_text_ids"]
+            emb = torch.randn(text_ids.shape[0], hidden_size)
+            return BaseNavitOutputWithPast(packed_query_sequence=emb)
+
+        packed_query_sequence = kwargs.get("packed_query_sequence")
+        if packed_query_sequence is None:
+            text_ids = kwargs["packed_text_ids"]
+            packed_query_sequence = torch.randn(text_ids.shape[0], hidden_size)
+        return BaseNavitOutputWithPast(
+            packed_query_sequence=packed_query_sequence,
+            past_key_values=kwargs.get("past_key_values"),
+        )
+
+    lm = MagicMock()
+    lm.forward = forward
+    lm._calls = calls
+    return lm
+
+
+def _vae_transform(image: Image.Image) -> torch.Tensor:
+    arr = torch.from_numpy(np.array(image.convert("RGB"), dtype=np.float32))
+    arr = arr / 127.5 - 1.0
+    return arr.permute(2, 0, 1)
+
+
+@pytest.fixture
+def bagel_and_vae_input():
+    language_model = _make_tracking_language_model()
+    bagel = Bagel(language_model=language_model, vit_model=None, config=_make_bagel_config())
+
+    image = Image.new("RGB", (64, 64), color=(128, 64, 32))
+    new_token_ids = {"start_of_image": 100, "end_of_image": 101}
+
+    gen_input, _newlens, _new_rope = bagel.prepare_vae_images(
+        curr_kvlens=[0],
+        curr_rope=[0],
+        images=[image],
+        transforms=_vae_transform,
+        new_token_ids=new_token_ids,
+    )
+
+    vae_model = MagicMock()
+    patch_shapes = gen_input["patchified_vae_latent_shapes"]
+    p = bagel.latent_patch_size
+    c = bagel.latent_channel
+
+    def encode(padded_images: torch.Tensor) -> torch.Tensor:
+        # Match AutoEncoder.encode layout: [B, z_channels, h*p, w*p] per image.
+        batch = padded_images.shape[0]
+        assert batch == len(patch_shapes)
+        latents = [torch.randn(c, h * p, w * p) for h, w in patch_shapes]
+        return torch.stack(latents, dim=0)
+
+    vae_model.encode = encode
+
+    expected_seq_len = int(gen_input["packed_seqlens"].sum().item())
+    num_text_tokens = int(gen_input["packed_text_ids"].numel())
+
+    return SimpleNamespace(
+        bagel=bagel,
+        language_model=language_model,
+        vae_model=vae_model,
+        gen_input=gen_input,
+        past_key_values=NaiveCache(NUM_LAYERS),
+        expected_seq_len=expected_seq_len,
+        num_text_tokens=num_text_tokens,
+    )
+
+
+class TestForwardCacheUpdateVae:
+    """Regression tests for single-stage img2img VAE cache prefill."""
+
+    def test_forwards_full_packed_query_sequence(self, bagel_and_vae_input):
+        """Update forward must use packed_query_sequence spanning text + VAE tokens."""
+        ctx = bagel_and_vae_input
+        assert ctx.expected_seq_len > ctx.num_text_tokens, "test setup should include many VAE tokens"
+
+        ctx.bagel.forward_cache_update_vae(
+            ctx.vae_model,
+            ctx.past_key_values,
+            **ctx.gen_input,
+        )
+
+        update_calls = [c for c in ctx.language_model._calls if not c.get("return_embeddings_only")]
+        assert len(update_calls) == 1
+
+        call = update_calls[0]
+        assert "packed_query_sequence" in call
+        assert "packed_text_ids" not in call
+
+        packed = call["packed_query_sequence"]
+        assert packed.shape[0] == ctx.expected_seq_len
+        assert call.get("mode") == "gen"
+
+    def test_packed_sequence_covers_text_and_vae_index_ranges(self, bagel_and_vae_input):
+        """Text and VAE index ranges must lie within packed_query_sequence."""
+        ctx = bagel_and_vae_input
+        ctx.bagel.forward_cache_update_vae(
+            ctx.vae_model,
+            ctx.past_key_values,
+            **ctx.gen_input,
+        )
+
+        update_call = next(c for c in ctx.language_model._calls if not c.get("return_embeddings_only"))
+        packed = update_call["packed_query_sequence"]
+
+        text_idx = ctx.gen_input["packed_text_indexes"]
+        vae_idx = ctx.gen_input["packed_vae_token_indexes"]
+        assert int(text_idx.max()) < packed.shape[0]
+        assert int(vae_idx.max()) < packed.shape[0]
+        assert int(text_idx.min()) >= 0
+        assert int(vae_idx.min()) >= 0