From b2fd9d72724a7dcb0cb50b1af75d9ed1aaf0ab65 Mon Sep 17 00:00:00 2001 From: syang-genai <191082996+syang-genai@users.noreply.github.com> Date: Wed, 4 Nov 2026 07:55:44 -0800 Subject: [PATCH] fix(qwen3_5): give image prompts embedding spans so prefix hits past an image don't crash Qwen3.5 image prompts reached the runner without `embedding_spans`, so they took the whole-image merge: every vision row is scattered into the image tokens of the forward. After a prefix-cache hit that ends inside or past the image, the forward holds fewer image tokens than the vision tower returns rows, and `merge_multimodal_embeddings` dies with a shape mismatch. `build_qwen_vl_inputs` now returns each image's run of image tokens as an embedding span for the `qwen3_5` and `qwen3_5_moe` model types (the same gate as the Qwen3.5 MRoPE path), plus the media `cache_seed`. The runner then scatters only the slice of each image a forward contains, and image prompts can be chunked instead of prefilled in one forward. Image tokens that do not match the grids are rejected when the request is built. Qwen3.8 (`qwen4_exp`) is left on the old path for now. Fixes #2865 --- atom/models/qwen3_5_vl.py | 73 +++++++++++++++++++++- tests/test_multimodal_inputs.py | 76 +++++++++++++++++++++++ tests/test_multimodal_runtime.py | 100 +++++++++++++++++++++++++++++++ 3 files changed, 247 insertions(+), 2 deletions(-) diff --git a/atom/models/qwen3_5_vl.py b/atom/models/qwen3_5_vl.py index dd43bc2f3d..ee6ca7760e 100644 --- a/atom/models/qwen3_5_vl.py +++ b/atom/models/qwen3_5_vl.py @@ -413,6 +413,53 @@ def _images_before_text(messages: list[dict]) -> list[dict]: return reordered +# Qwen3.5 checkpoints whose image prompts carry explicit embedding spans. Same +# model types as the Qwen3.5 MRoPE path (`_get_qwen3_5_mrope_input_positions`). +_QWEN3_5_SPAN_MODEL_TYPES = frozenset({"qwen3_5", "qwen3_5_moe"}) + + +def qwen3_5_image_spans( + input_ids: list[int], + image_grid_thw, + image_token_id: int, + spatial_merge_size: int, +) -> tuple[tuple[int, int], ...]: + """Each image's ``(start, count)`` run of image tokens, in image order. + + The HF Qwen processor has already expanded every ``<|image_pad|>`` into + ``t * h * w / spatial_merge_size**2`` image tokens, one per row the vision + tower returns for that grid. Walk the prompt and claim each image's run in + turn, so adjacent images without separators split correctly too. + """ + grids = ( + image_grid_thw.tolist() if hasattr(image_grid_thw, "tolist") else image_grid_thw + ) + merge_unit = int(spatial_merge_size) ** 2 + spans: list[tuple[int, int]] = [] + position = 0 + for index, (t, h, w) in enumerate(grids): + count = int(t) * int(h) * int(w) // merge_unit + while position < len(input_ids) and input_ids[position] != image_token_id: + position += 1 + run = input_ids[position : position + count] + found = next( + (i for i, token in enumerate(run) if token != image_token_id), len(run) + ) + if found != count: + raise ValueError( + f"image {index} needs {count} consecutive image tokens (grid " + f"{t}x{h}x{w}, merge {spatial_merge_size}) but only {found} start " + f"at position {position}" + ) + spans.append((position, count)) + position += count + if any(token == image_token_id for token in input_ids[position:]): + raise ValueError( + f"prompt has more image tokens than its {len(grids)} image grids produce" + ) + return tuple(spans) + + def build_qwen_vl_inputs( atom_config, processor, @@ -421,7 +468,15 @@ def build_qwen_vl_inputs( chat_template_kwargs: dict, tools=None, ) -> tuple[list[int], dict]: - """Apply the existing template/processor convention and validate image slots.""" + """Apply the existing template/processor convention and validate image slots. + + For Qwen3.5 the image token runs are also returned as ``embedding_spans``, + so the runner scatters only the slice of each image a forward contains. + Without spans the prompt goes through the whole-image merge, which assumes + every image token is in the forward; a prefix-cache hit that ends inside or + past an image breaks that assumption (ROCm/ATOM#2865). Spans also let an + image prompt be chunked instead of prefilled in one forward. + """ if isinstance(prompt, str): text = prompt else: @@ -440,7 +495,21 @@ def build_qwen_vl_inputs( f"Prompt has {num_placeholders} image placeholders but {len(images)} images" ) processed = processor(text=[text], images=images, return_tensors="pt") - return processed["input_ids"][0].tolist(), { + input_ids = processed["input_ids"][0].tolist() + multimodal_data = { "pixel_values": processed["pixel_values"], "image_grid_thw": processed["image_grid_thw"], } + multimodal_config = getattr(atom_config, "multimodal_config", None) + if getattr(multimodal_config, "model_type", None) in _QWEN3_5_SPAN_MODEL_TYPES: + from atom.model_engine.multimodal_runtime import multimodal_cache_seed + + vision_config = getattr(multimodal_config, "vision_config", None) + multimodal_data["embedding_spans"] = qwen3_5_image_spans( + input_ids, + multimodal_data["image_grid_thw"], + image_token_id=getattr(multimodal_config, "image_token_id", 248056), + spatial_merge_size=getattr(vision_config, "spatial_merge_size", 2), + ) + multimodal_data["cache_seed"] = multimodal_cache_seed(multimodal_data) + return input_ids, multimodal_data diff --git a/tests/test_multimodal_inputs.py b/tests/test_multimodal_inputs.py index eb21f164ea..d82e00f298 100644 --- a/tests/test_multimodal_inputs.py +++ b/tests/test_multimodal_inputs.py @@ -372,3 +372,79 @@ def test_kimi_rendered_prompt_reports_required_conversation_before_loading(): "prompt", {"image": "unused.png"}, ) + + +IMG = 248056 # Qwen3.5 <|image_pad|> + + +def qwen3_5_config(model_type="qwen3_5_moe"): + return SimpleNamespace( + hf_config=SimpleNamespace(architectures=["Qwen3_5MoeForConditionalGeneration"]), + multimodal_config=SimpleNamespace( + model_type=model_type, + image_token_id=IMG, + vision_config=SimpleNamespace(spatial_merge_size=2), + ), + ) + + +def prepare_qwen3_5(image_module, input_ids, image_grid_thw, model_type="qwen3_5_moe"): + hf_output = { + "input_ids": np.array([input_ids]), + "pixel_values": np.ones((sum(t * h * w for t, h, w in image_grid_thw), 3)), + "image_grid_thw": np.array(image_grid_thw), + } + + class Processor: + def apply_chat_template(self, messages, **kwargs): + return "<|image_pad|>" * len(image_grid_thw) + " prompt" + + def __call__(self, **kwargs): + return hf_output + + images = [image_module.new("RGB", (2, 2)) for _ in image_grid_thw] + conversation = chat(*({"type": "image"} for _ in images)) + return prepare_multimodal_inputs( + qwen3_5_config(model_type), Processor(), conversation, {"image": images} + ) + + +@pytest.mark.parametrize("model_type", ["qwen3_5", "qwen3_5_moe"]) +def test_qwen3_5_image_prompts_carry_one_span_per_image(image_module, model_type): + # Two images: a 4x4 grid (4 tokens after the 2x2 merge) and a 2x4 grid + # (2 tokens), the second directly after the first with no separator. + input_ids = [1, 2, IMG, IMG, IMG, IMG, IMG, IMG, 3] + token_ids, data = prepare_qwen3_5( + image_module, input_ids, [[1, 4, 4], [1, 2, 4]], model_type + ) + assert token_ids == input_ids + assert data["embedding_spans"] == ((2, 4), (6, 2)) + assert isinstance(data["cache_seed"], int) + + +def test_qwen3_5_cache_seed_tells_images_apart(image_module): + _, first = prepare_qwen3_5(image_module, [IMG] * 4, [[1, 4, 4]]) + _, same = prepare_qwen3_5(image_module, [IMG] * 4, [[1, 4, 4]]) + _, other_grid = prepare_qwen3_5(image_module, [IMG] * 4, [[1, 2, 8]]) + assert first["cache_seed"] == same["cache_seed"] + assert first["cache_seed"] != other_grid["cache_seed"] + + +@pytest.mark.parametrize( + "input_ids", + [ + [1, IMG, IMG, IMG, 2], # one token short of the grid's 4 + [1, IMG, IMG, 2, IMG, IMG], # the run is split + [1, IMG, IMG, IMG, IMG, IMG], # one image token too many + ], +) +def test_qwen3_5_image_tokens_that_disagree_with_the_grid_are_rejected( + image_module, input_ids +): + with pytest.raises(ValueError, match="image tokens"): + prepare_qwen3_5(image_module, input_ids, [[1, 4, 4]]) + + +def test_other_qwen_model_types_keep_the_whole_image_merge(image_module): + _, data = prepare_qwen3_5(image_module, [IMG] * 4, [[1, 4, 4]], "qwen4_exp") + assert "embedding_spans" not in data and "cache_seed" not in data diff --git a/tests/test_multimodal_runtime.py b/tests/test_multimodal_runtime.py index 4a09615e45..0c1ff4b7ee 100644 --- a/tests/test_multimodal_runtime.py +++ b/tests/test_multimodal_runtime.py @@ -257,3 +257,103 @@ def processor(**kwargs): batch=batch, ) assert chunks == [6, 6, 6] + + +def qwen3_5_payload(token_ids, image_grid_thw, seed=21): + from atom.models.qwen3_5_vl import qwen3_5_image_spans + + spans = qwen3_5_image_spans(token_ids, image_grid_thw, 42, 2) + rows = sum(count for _, count in spans) + return { + "cache_seed": seed, + "embedding_spans": spans, + "pixel_values": torch.arange(2 * rows).view(rows, 2).float(), + "image_grid_thw": torch.tensor(image_grid_thw), + } + + +# 3 text tokens, one 1x4x8 image (8 tokens after the 2x2 merge), 5 text tokens. +QWEN3_5_IDS = [7, 8, 9] + [42] * 8 + [5, 6, 7, 8, 9] + + +@pytest.mark.parametrize("cached", [0, 3, 5, 10, 11, 13]) +def test_qwen3_5_prefix_hit_scatters_only_the_image_rows_left_in_the_forward(cached): + # ROCm/ATOM#2865: the whole-image merge scattered all 8 image rows into the + # forward even when a prefix hit had already covered some (cached=5, 10) or + # all (cached=11, 13) of them, and crashed on the shape mismatch. + model = VisionModel() + data = qwen3_5_payload(QWEN3_5_IDS, [[1, 4, 8]]) + full_ids = torch.tensor(QWEN3_5_IDS) + expected = model.embed_input_ids(full_ids) + expected[3:11] = data["pixel_values"] + 100 + + length = len(QWEN3_5_IDS) - cached + batch = SimpleNamespace( + req_ids=[4], + num_scheduled_tokens=[length], + context_lens=[len(QWEN3_5_IDS)], + multimodal_data={4: data}, + ) + actual = embed_multimodal_batch( + model, VisionEmbeddingCache(), full_ids[cached:], batch, "cpu", torch.float32 + ) + torch.testing.assert_close(actual, expected[cached:], rtol=0, atol=0) + + +def test_qwen3_5_image_prompt_is_chunked_and_resumes_after_a_prefix_hit(): + from atom.models.qwen3_5_vl import build_qwen_vl_inputs + + def processor(**kwargs): + return { + "input_ids": torch.tensor([QWEN3_5_IDS]), + "pixel_values": torch.ones(32, 3), + "image_grid_thw": torch.tensor([[1, 4, 8]]), + } + + atom_config = SimpleNamespace( + multimodal_config=SimpleNamespace( + model_type="qwen3_5", + image_token_id=42, + vision_config=SimpleNamespace(spatial_merge_size=2), + ) + ) + token_ids, data = build_qwen_vl_inputs( + atom_config, processor, "<|image_pad|> prompt", [object()], {} + ) + assert data["embedding_spans"] == ((3, 8),) + + scheduler = Scheduler( + MockConfig( + max_num_batched_tokens=6, + max_model_len=128, + num_kvcache_blocks=100, + enable_chunked_prefill=True, + enable_prefix_caching=True, + ) + ) + + def prefill(seq): + """Schedule `seq` to the end of its prompt; return each chunk's size.""" + assert scheduler._unschedulable_reason(seq) is None + assert not scheduler._requires_atomic_prefill(seq) + scheduler.add(seq) + chunks = [] + while seq.num_cached_tokens < len(token_ids): + batch, _ = scheduler.schedule() + assert batch.multimodal_data[seq.id]["embedding_spans"] == ((3, 8),) + chunks.append(batch.total_tokens_num_prefill) + scheduler.postprocess( + list(scheduler.running), + ScheduledBatchOutput([], [], None, None, None), + batch=batch, + ) + return chunks + + params = SamplingParams(temperature=0, max_tokens=2) + # The 16-token prompt no longer has to fit one forward. + assert prefill(Sequence(token_ids, 4, params, multimodal_data=data)) == [6, 6, 4] + # The same image again hits the 12 published tokens, past the image's end + # at 11: the forward holds 4 text tokens and none of the image (the + # `[0, 2048]` crash in ROCm/ATOM#2865). + again = Sequence(token_ids, 4, params, multimodal_data=dict(data)) + assert prefill(again) == [4]