diff --git a/deepinfer/ve/prepare_tensor_input_vl.py b/deepinfer/ve/prepare_tensor_input_vl.py index 24a0b3f..f79cf2b 100644 --- a/deepinfer/ve/prepare_tensor_input_vl.py +++ b/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) diff --git a/tests/ve/test_mixed_token_metadata.py b/tests/ve/test_mixed_token_metadata.py new file mode 100644 index 0000000..630ae07 --- /dev/null +++ b/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))