[Bugfix][Qwen3.5] Give image prompts embedding spans so prefix hits past an image don't crash

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

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

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]