From 11a951b9626ba5e947fe30e44dc83cf1d7c70687 Mon Sep 17 00:00:00 2001 From: NumberWan Date: Tue, 26 May 2026 11:01:10 +0800 Subject: [PATCH] [BugFix]: CUDA `device-side assert` failures on single-stage BAGEL i2i requests (#3680) Signed-off-by: NumberWan --- .../bagel/test_forward_cache_update_vae.py | 169 ++++++++++++++++++ .../models/bagel/bagel_transformer.py | 14 +- 2 files changed, 182 insertions(+), 1 deletion(-) create mode 100644 tests/diffusion/models/bagel/test_forward_cache_update_vae.py diff --git a/tests/diffusion/models/bagel/test_forward_cache_update_vae.py b/tests/diffusion/models/bagel/test_forward_cache_update_vae.py new file mode 100644 index 000000000..e49b408f4 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/bagel_transformer.py b/vllm_omni/diffusion/models/bagel/bagel_transformer.py index d33b26f60..ecaba8e08 100644 --- a/vllm_omni/diffusion/models/bagel/bagel_transformer.py +++ b/vllm_omni/diffusion/models/bagel/bagel_transformer.py @@ -1389,6 +1389,18 @@ class Bagel(nn.Module): packed_timestep_embeds = self.time_embedder(packed_timesteps) packed_latent = self.vae2llm(packed_latent) + packed_timestep_embeds + packed_pos_embed + # Mirror forward_cache_update_vit: build a full packed_sequence so MoE gen-mode + # indexes (packed_text_indexes / packed_vae_token_indexes) match tensor length. + packed_text_embedding = self.language_model.forward( + packed_text_ids=packed_text_ids, + return_embeddings_only=True, + ).packed_query_sequence + packed_sequence = packed_text_embedding.new_zeros((sum(packed_seqlens), self.hidden_size)) + packed_sequence[packed_text_indexes] = packed_text_embedding + if packed_latent.dtype != packed_sequence.dtype: + packed_latent = packed_latent.to(packed_sequence.dtype) + packed_sequence[packed_vae_token_indexes] = packed_latent + extra_inputs = {} if self.use_moe: extra_inputs = { @@ -1398,7 +1410,7 @@ class Bagel(nn.Module): } output = self.language_model.forward( - packed_text_ids=packed_text_ids, + packed_query_sequence=packed_sequence, query_lens=packed_seqlens, packed_query_position_ids=packed_position_ids, packed_query_indexes=packed_indexes,