mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 14:57:09 +08:00
190 lines
8.9 KiB
Python
190 lines
8.9 KiB
Python
"""Pure-CPU unit tests for _chunk_mm_metadata: lock down the boundary semantics of bisecting image-range intersections, byte-range slicing, and patch-dimension slicing.
|
|
|
|
No real ckpt / GPU compute is needed: only CPU-tensor metadata dicts are built and passed to pure functions; importing deepinfer still needs _C (a GPU machine).
|
|
"""
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from deepinfer.ve.prepare_tensor_input_vl import _chunk_mm_metadata, _prepare_raw_image_transform_inputs
|
|
|
|
# The raw byte count / patch count per image is fixed, which makes the rebased byte and patch ranges easy to assert.
|
|
RAW_BYTES_PER_IMG = 12
|
|
PATCHES_PER_IMG = 3
|
|
|
|
|
|
def _md(spans: list[tuple[int, int]], patches_per_img: int = PATCHES_PER_IMG) -> dict[str, Any]:
|
|
"""Build a metadata dict in the mapper's output shape from per-image (start, count).
|
|
|
|
spans: [(start, count)] is where image_starts/image_lens come from; each image has RAW_BYTES_PER_IMG raw RGB bytes and
|
|
patches_per_img patches. The contents of images_raw are global byte indices (so the sliced ranges can be asserted directly).
|
|
"""
|
|
starts = torch.tensor([s for s, _ in spans], dtype=torch.int32)
|
|
lens = torch.tensor([c for _, c in spans], dtype=torch.int32)
|
|
n = len(spans)
|
|
cu_patches = torch.zeros(n + 1, dtype=torch.int64)
|
|
cu_raw = torch.zeros(n + 1, dtype=torch.int64)
|
|
for i in range(n):
|
|
cu_patches[i + 1] = cu_patches[i] + patches_per_img
|
|
cu_raw[i + 1] = cu_raw[i] + RAW_BYTES_PER_IMG
|
|
total_patch = int(cu_patches[-1])
|
|
total_raw = int(cu_raw[-1])
|
|
return {
|
|
"images_raw": torch.arange(total_raw, dtype=torch.uint8) if total_raw else torch.zeros(0, dtype=torch.uint8),
|
|
"images_input_size": torch.full((n, 2), 4, dtype=torch.int64) if n else torch.zeros(0, 2, dtype=torch.int64),
|
|
"cu_raw_bytes": cu_raw,
|
|
"image_starts": starts,
|
|
"image_lens": lens,
|
|
"cu_image_patch_lens": cu_patches,
|
|
"images_spatial_crop": torch.arange(n * 2).reshape(n, 2) if n else torch.zeros(0, 2, dtype=torch.int64),
|
|
"images_global_attn_pos_id": torch.arange(total_patch * 2).reshape(total_patch, 2),
|
|
"images_seq_mask": None,
|
|
"mm_token_type": None,
|
|
}
|
|
|
|
|
|
def _hit_range(spans: list[tuple[int, int]], start_idx: int, end_idx: int) -> tuple[int, int] | None:
|
|
"""Reference for the old min/max semantics: images intersecting the chunk are [first, last+1); None when there are none."""
|
|
hit = [i for i, (s, c) in enumerate(spans) if s < end_idx and s + c > start_idx]
|
|
return (hit[0], hit[-1] + 1) if hit else None
|
|
|
|
|
|
def _assert_chunk(spans: list[tuple[int, int]], start_idx: int, end_idx: int, patches_per_img: int = PATCHES_PER_IMG) -> None:
|
|
md = _chunk_mm_metadata(_md(spans, patches_per_img), start_idx, end_idx)
|
|
rng = _hit_range(spans, start_idx, end_idx)
|
|
if rng is None:
|
|
# image-free chunk: all image-related keys are set to None (downstream skips image encoding as "no images")
|
|
for k in (
|
|
"images_raw",
|
|
"images_input_size",
|
|
"cu_raw_bytes",
|
|
"images_spatial_crop",
|
|
"images_global_attn_pos_id",
|
|
"cu_image_patch_lens",
|
|
):
|
|
assert md[k] is None, k
|
|
return
|
|
si, ei = rng
|
|
# images_raw is sliced by byte range: this chunk's raw bytes = [si*RAW, ei*RAW)
|
|
assert md["images_raw"].tolist() == list(range(si * RAW_BYTES_PER_IMG, ei * RAW_BYTES_PER_IMG))
|
|
assert md["images_input_size"].shape[0] == ei - si
|
|
# cu_raw_bytes is rebased to this chunk's start and accumulated per image
|
|
assert md["cu_raw_bytes"].tolist() == [i * RAW_BYTES_PER_IMG for i in range(ei - si + 1)]
|
|
# pos_id is sliced by patch range; cu_image_patch_lens is rebased to start at 0
|
|
assert md["images_global_attn_pos_id"].shape[0] == (ei - si) * patches_per_img
|
|
assert md["cu_image_patch_lens"][0].item() == 0
|
|
assert md["cu_image_patch_lens"][-1].item() == (ei - si) * patches_per_img
|
|
assert md["images_spatial_crop"].shape[0] == ei - si
|
|
|
|
|
|
def test_chunk_at_image_end_excludes_image() -> None:
|
|
# image0=[0,3), image1=[3,6); the chunk ends exactly at image0's end, end_idx==3 => only image0 is included
|
|
spans = [(0, 3), (3, 3)]
|
|
_assert_chunk(spans, 0, 3)
|
|
md = _chunk_mm_metadata(_md(spans), 0, 3)
|
|
assert md["images_raw"].shape[0] == RAW_BYTES_PER_IMG # only image0's bytes
|
|
assert md["cu_image_patch_lens"][-1].item() == PATCHES_PER_IMG # only image0's patches
|
|
|
|
|
|
def test_chunk_at_image_start_includes_image() -> None:
|
|
# the chunk starts at image1's start, start_idx==3 => image1 is included (image0 is not)
|
|
spans = [(0, 3), (3, 3)]
|
|
md = _chunk_mm_metadata(_md(spans), 3, 6)
|
|
assert md["images_raw"].tolist() == list(range(RAW_BYTES_PER_IMG, 2 * RAW_BYTES_PER_IMG))
|
|
assert md["images_spatial_crop"].shape[0] == 1
|
|
|
|
|
|
def test_empty_chunk_has_no_image() -> None:
|
|
spans = [(0, 3), (3, 3)]
|
|
md = _chunk_mm_metadata(_md(spans), 3, 3) # start==end
|
|
assert md["images_raw"] is None
|
|
assert md["cu_image_patch_lens"] is None
|
|
|
|
|
|
def test_multi_image_partial_overlap() -> None:
|
|
# images [2,4), [5,7), [8,11); chunk [3,9) intersects all three
|
|
spans = [(2, 2), (5, 2), (8, 3)]
|
|
for s in range(12):
|
|
for e in range(s + 1, 13):
|
|
_assert_chunk(spans, s, e, patches_per_img=4)
|
|
|
|
|
|
def test_no_image_returns_early() -> None:
|
|
md = _chunk_mm_metadata(_md([]), 0, 4)
|
|
# no images: image_starts is an empty tensor and intersect is empty -> all image-related keys are set to None
|
|
assert md["images_raw"] is None
|
|
assert md["images_spatial_crop"] is None
|
|
|
|
|
|
def test_no_raw_bytes_drops_image_metadata() -> None:
|
|
"""When this seq has no raw bytes (every image was cache-hit and pre-split), intersecting a zero-placeholder image must not crash:
|
|
images_input_size/cu_raw_bytes are None -> the whole image key group is set to None and downstream treats it as image-free."""
|
|
metadata = _md([(0, 3)])
|
|
metadata["images_input_size"] = None
|
|
metadata["images_raw"] = None
|
|
metadata["cu_raw_bytes"] = None
|
|
md = _chunk_mm_metadata(metadata, 0, 3)
|
|
for k in ("images_raw", "images_input_size", "cu_raw_bytes", "images_spatial_crop", "cu_image_patch_lens"):
|
|
assert md[k] is None, k
|
|
|
|
|
|
def test_does_not_mutate_input() -> None:
|
|
"""Shallow-copy semantics: tensors produced by slicing must not write back into the original metadata."""
|
|
metadata = _md([(0, 3), (3, 3)])
|
|
orig_raw = metadata["images_raw"]
|
|
orig_cu_raw = metadata["cu_raw_bytes"]
|
|
orig_cu_patch = metadata["cu_image_patch_lens"]
|
|
_chunk_mm_metadata(metadata, 3, 6)
|
|
assert metadata["images_raw"] is orig_raw
|
|
assert metadata["cu_raw_bytes"] is orig_cu_raw
|
|
assert metadata["cu_image_patch_lens"] is orig_cu_patch
|
|
assert metadata["cu_raw_bytes"].tolist() == [0, RAW_BYTES_PER_IMG, 2 * RAW_BYTES_PER_IMG] # unaffected by rebase
|
|
|
|
|
|
@pytest.mark.parametrize("sizes", [[(1302, 1302)] * 424, [(65536, 32768), (14, 14)]])
|
|
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
|
|
def test_prepare_raw_image_inputs_int64(sizes: list[tuple[int, int]], dtype: torch.dtype) -> None:
|
|
"""Realistic multi-image overflow plus a single-image multiply-before-cast regression.
|
|
|
|
Only the raw buffer is meta: preparation never reads pixels, so this exercises
|
|
the real CPU arithmetic without allocating several GiB of image data.
|
|
"""
|
|
byte_lens = [w * h * 3 for w, h in sizes]
|
|
raw = torch.empty(sum(byte_lens), dtype=torch.uint8, device="meta")
|
|
metadata = {"images_raw": raw, "images_input_size": torch.tensor(sizes, dtype=dtype)}
|
|
_prepare_raw_image_transform_inputs(metadata, total_image_patches=123)
|
|
expected = [0]
|
|
for n in byte_lens:
|
|
expected.append(expected[-1] + n)
|
|
assert expected[-1] > 2**31
|
|
assert metadata["cu_raw_bytes"].dtype == torch.int64
|
|
assert metadata["cu_raw_bytes"].tolist() == expected
|
|
assert metadata["image_raw_sizes"].dtype == torch.int32
|
|
assert metadata["image_raw_sizes"].tolist() == [list(s) for s in sizes]
|
|
assert metadata["images_raw"] is raw
|
|
assert metadata["total_image_patches"] == 123
|
|
assert "images_input_size" not in metadata
|
|
|
|
|
|
def test_chunk_raw_bytes_above_int32() -> None:
|
|
metadata = _md([(0, 3), (3, 6)])
|
|
raw_start = 2**31 + 12
|
|
metadata["images_raw"] = torch.empty(raw_start + 12, dtype=torch.uint8, device="meta")
|
|
metadata["cu_raw_bytes"] = torch.tensor([0, raw_start, raw_start + 12], dtype=torch.int64)
|
|
chunk = _chunk_mm_metadata(metadata, 3, 9)
|
|
assert chunk["images_raw"].numel() == 12
|
|
assert chunk["images_raw"].storage_offset() == raw_start
|
|
assert chunk["cu_raw_bytes"].dtype == torch.int64
|
|
assert chunk["cu_raw_bytes"].tolist() == [0, 12]
|
|
|
|
|
|
def test_prepare_empty_raw_image_inputs() -> None:
|
|
metadata = {"images_raw": None, "images_input_size": None}
|
|
_prepare_raw_image_transform_inputs(metadata, total_image_patches=0)
|
|
assert metadata["images_raw"] is None
|
|
assert metadata["image_raw_sizes"] is None
|
|
assert metadata["cu_raw_bytes"] is None
|
|
assert metadata["total_image_patches"] == -1
|