Fix token metadata alignment in mixed text and image batches

deepinfer/ve/prepare_tensor_input_vl.py

@@ -90,6 +90,11 @@ def prepare_input_metadata_vl(
         if multi_modal_data is None:
             multi_modal_data = dict.fromkeys(MM_METADATA_KEYS)
         chunked_multi_modal_data = _chunk_mm_metadata(multi_modal_data, start_idx, end_idx)
+        if seq_metadata.multi_modal_data is None:
+            # Text/profile/padding sequences still occupy token slots in a mixed VL batch.
+            # Fill after slicing: seqlen_k can include an already processed prefix.
+            chunked_multi_modal_data["mm_token_type"] = torch.zeros(seq_metadata.seqlen_q, dtype=torch.int32)
+            chunked_multi_modal_data["images_seq_mask"] = torch.zeros(seq_metadata.seqlen_q, dtype=torch.bool)
         # chunk 切片专用键不进 batch（模型层无消费者）；"handle" 留在原 dict 管理所有权
         for key in (*MM_IMAGE_SPAN_KEYS, "cu_image_patch_lens"):
             chunked_multi_modal_data.pop(key)

tests/ve/test_mixed_token_metadata.py

@@ -0,0 +1,74 @@
+"""CPU regression for token positions in mixed text/image prefill batches."""
+
+import importlib.util
+import sys
+from collections.abc import Iterator
+from pathlib import Path
+from types import ModuleType, SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
+
+
+@pytest.fixture
+def vl(monkeypatch: pytest.MonkeyPatch) -> Iterator[ModuleType]:
+    """Load the real preparation and schema without starting an accelerator runtime."""
+    root = Path(__file__).resolve().parents[2] / "deepinfer"
+
+    def module(name: str, **values: Any) -> ModuleType:
+        value = ModuleType(name)
+        value.__dict__.update(values)
+        monkeypatch.setitem(sys.modules, name, value)
+        return value
+
+    module("deepinfer", PLATFORM=SimpleNamespace(is_npu=False))
+    module("deepinfer.timer", cpu_timer=SimpleNamespace(record=lambda _: lambda function: function))
+
+    def concat(rows: list[dict]) -> dict:
+        return {key: torch.cat(tensors) if (tensors := [row[key] for row in rows if row[key] is not None]) else None for key in rows[0]}
+
+    module("deepinfer.utils", batch_tensor_dict_list=concat)
+    module("deepinfer.ve.processor.image_processor", downsampled_seqlens=lambda h, w, ratio: h * w)
+    module("deepinfer.shared_memory", SharedTorchMemoryPool=object, get_multi_modal_pool=lambda: None)
+    monkeypatch.setattr(torch.cuda, "Event", SimpleNamespace)
+    schema_spec = importlib.util.spec_from_file_location("deepinfer.ve.schema", root / "ve/schema.py")
+    schema = importlib.util.module_from_spec(schema_spec)
+    monkeypatch.setitem(sys.modules, schema_spec.name, schema)
+    schema_spec.loader.exec_module(schema)
+    spec = importlib.util.spec_from_file_location("mixed_vl_target", root / "ve/prepare_tensor_input_vl.py")
+    target = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(target)
+    yield target
+
+
+def sequence(length: int, metadata: dict | None = None, offset: int = 0) -> SimpleNamespace:
+    return SimpleNamespace(seqlen_q=length, seqlen_k=length + offset, multi_modal_data=metadata)
+
+
+@pytest.mark.parametrize("position", [0, 1, 2])
+@pytest.mark.parametrize("text_length", [0, 1, 7])
+def test_mixed_batch_keeps_text_and_padding_positions(vl: ModuleType, position: int, text_length: int) -> None:
+    metadata = dict.fromkeys(vl.MM_METADATA_KEYS)
+    metadata["mm_token_type"] = torch.full((3,), 5, dtype=torch.int32)
+    metadata["images_seq_mask"] = torch.ones(3, dtype=torch.bool)
+    rows = [sequence(3, metadata), sequence(3, metadata)]
+    rows.insert(position, sequence(text_length))
+    result = vl.prepare_input_metadata_vl(rows)
+    expected = torch.cat(
+        [torch.zeros(text_length, dtype=torch.int32) if row.multi_modal_data is None else metadata["mm_token_type"] for row in rows]
+    )
+    assert torch.equal(result["mm_token_type"], expected)
+    assert torch.equal(result["images_seq_mask"], expected != 0)
+    assert result["mm_token_type"].dtype == torch.int32
+    assert result["images_seq_mask"].dtype == torch.bool
+    assert metadata["mm_token_type"].tolist() == [5, 5, 5]
+
+
+@pytest.mark.parametrize("offset", [0, 3, 4096])
+def test_text_chunk_uses_query_length(vl: ModuleType, offset: int) -> None:
+    result = vl.prepare_input_metadata_vl([sequence(5, offset=offset)])
+    assert torch.equal(result["mm_token_type"], torch.zeros(5, dtype=torch.int32))
+    assert torch.equal(result["images_seq_mask"], torch.zeros(5, dtype=torch.bool))