vlm: cache kimi-k3 per-image processor artifacts (#34404)
This commit is contained in:
@@ -19,6 +19,7 @@ from sglang.srt.disaggregation.encode_receiver import (
|
||||
EmbeddingData,
|
||||
MMReceiverHTTP,
|
||||
MultiModalEmbeddingData,
|
||||
_encoder_media_item,
|
||||
_select_mm_processor_prompt,
|
||||
)
|
||||
from sglang.srt.disaggregation.encode_server import MMEncoder, _get_mm_grid_dim
|
||||
@@ -27,8 +28,10 @@ from sglang.srt.managers.tokenizer_manager import (
|
||||
_reject_missing_dispatched_encoder_embedding,
|
||||
)
|
||||
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
|
||||
from sglang.srt.multimodal.cache import snapshot_media
|
||||
from sglang.srt.multimodal.encoder_preprocessing import (
|
||||
LOCAL_PREPROCESSED_KEY,
|
||||
EncoderMediaProcessorConfig,
|
||||
EncoderPreprocessOutput,
|
||||
get_encoder_preprocessed_items,
|
||||
hash_raw_encoder_item,
|
||||
@@ -41,6 +44,7 @@ from sglang.srt.multimodal.kimi_k3_image_processing import (
|
||||
)
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.srt.server_args import resolve_encoder_transfer_backend
|
||||
from sglang.srt.utils import ImageData
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||
@@ -208,6 +212,11 @@ def _encoder(model_type="kimi_k3"):
|
||||
vision_config=SimpleNamespace(merge_kernel_size=(2, 2))
|
||||
)
|
||||
)
|
||||
encoder.encoder_media_processor_config = (
|
||||
KimiK3ForConditionalGeneration.encoder_media_processor_config
|
||||
if model_type == "kimi_k3"
|
||||
else EncoderMediaProcessorConfig()
|
||||
)
|
||||
return encoder
|
||||
|
||||
|
||||
@@ -296,6 +305,19 @@ def test_kimi_k3_epd_preprocess_preserves_raw_per_image_items():
|
||||
assert deferred.image_std == [0.5, 0.5, 0.5]
|
||||
|
||||
|
||||
def test_kimi_k3_epd_preserves_verified_content_identity():
|
||||
image = Image.new("RGB", (8, 6), color=(1, 2, 3))
|
||||
digest = "sha256:" + "ab" * 32
|
||||
|
||||
output = prepare_kimi_k3_encoder_inputs(
|
||||
[{"type": "image", "image": image, "content_hash": digest}],
|
||||
_kimi_k3_image_processor(),
|
||||
)
|
||||
|
||||
item = get_encoder_preprocessed_items(output)[0]
|
||||
assert item.model_specific_data["content_digest"] == digest
|
||||
|
||||
|
||||
def test_kimi_k3_epd_model_preprocessor_receives_image_processor():
|
||||
image = Image.new("RGB", (8, 6), color=(1, 2, 3))
|
||||
image_processor = _kimi_k3_image_processor()
|
||||
@@ -458,6 +480,67 @@ def test_kimi_k3_epd_selects_matching_jpeg_decode_mode(
|
||||
load.assert_called_once_with(b"jpeg", expected_decode_mode)
|
||||
|
||||
|
||||
def test_kimi_k3_epd_verifies_content_hash_before_decode():
|
||||
payload = b"jpeg"
|
||||
digest = snapshot_media(payload).content_digest
|
||||
expected = torch.zeros((3, 2, 3), dtype=torch.uint8)
|
||||
encoder = _encoder()
|
||||
encoder.use_image_processor_gpu = False
|
||||
|
||||
with patch(
|
||||
"sglang.srt.disaggregation.encode_server.load_image",
|
||||
return_value=(expected, None),
|
||||
) as load:
|
||||
output = encoder._load_single_item(
|
||||
{"url": payload, "content_hash": digest}, Modality.IMAGE
|
||||
)
|
||||
|
||||
assert output == {
|
||||
"type": "image",
|
||||
"image": expected,
|
||||
"content_hash": digest,
|
||||
}
|
||||
load.assert_called_once_with(payload, False)
|
||||
|
||||
|
||||
def test_epd_receiver_keeps_content_hash_aligned_with_image():
|
||||
digest = "sha256:" + "cd" * 32
|
||||
receiver = MMReceiverHTTP.__new__(MMReceiverHTTP)
|
||||
request = SimpleNamespace(
|
||||
image_data=[
|
||||
ImageData(
|
||||
url="image",
|
||||
detail="high",
|
||||
max_dynamic_patch=12,
|
||||
preprocess_kwargs={"crop": False},
|
||||
content_hash=digest,
|
||||
)
|
||||
],
|
||||
video_data=None,
|
||||
audio_data=None,
|
||||
mm_content_hashes=[digest],
|
||||
)
|
||||
|
||||
assert receiver._extract_url_data(request) == [
|
||||
{
|
||||
"url": "image",
|
||||
"modality": Modality.IMAGE,
|
||||
"detail": "high",
|
||||
"max_dynamic_patch": 12,
|
||||
"preprocess_kwargs": {"crop": False},
|
||||
"content_hash": digest,
|
||||
}
|
||||
]
|
||||
|
||||
assert _encoder_media_item(receiver._extract_url_data(request)[0]) == {
|
||||
"url": "image",
|
||||
"detail": "high",
|
||||
"max_dynamic_patch": 12,
|
||||
"preprocess_kwargs": {"crop": False},
|
||||
"content_hash": digest,
|
||||
}
|
||||
|
||||
|
||||
def test_kimi_k3_epd_aggregates_original_image_sizes_in_part_order():
|
||||
first = EmbeddingData(
|
||||
req_id="request",
|
||||
|
||||
@@ -127,12 +127,14 @@ class TestBaseProcessorConfigExtraction(CustomTestCase):
|
||||
|
||||
with patch.dict(os.environ, {}, clear=False):
|
||||
os.environ.pop("SGLANG_IO_WORKERS", None)
|
||||
with patch.object(
|
||||
BaseMultimodalProcessor, "auto_mm_processor_worker_num", 4
|
||||
), patch.object(
|
||||
BaseMultimodalProcessor, "auto_mm_io_worker_num", 16
|
||||
), patch.object(
|
||||
BaseMultimodalProcessor, "supports_mm_processor_concurrency", True
|
||||
with (
|
||||
patch.object(
|
||||
BaseMultimodalProcessor, "auto_mm_processor_worker_num", 4
|
||||
),
|
||||
patch.object(BaseMultimodalProcessor, "auto_mm_io_worker_num", 16),
|
||||
patch.object(
|
||||
BaseMultimodalProcessor, "supports_mm_processor_concurrency", True
|
||||
),
|
||||
):
|
||||
proc = self._make_processor({})
|
||||
try:
|
||||
@@ -172,14 +174,18 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
|
||||
def _server_args(mm_feature_transport):
|
||||
return SimpleNamespace(
|
||||
mm_feature_transport=mm_feature_transport,
|
||||
image_processor_backend="auto",
|
||||
disable_fast_image_processor=False,
|
||||
skip_tokenizer_init=False,
|
||||
mm_process_config={},
|
||||
mm_preprocess_cache_size_mb=0,
|
||||
trust_mm_content_hashes=False,
|
||||
mm_processor_worker_num=0,
|
||||
mm_io_worker_num=0,
|
||||
tokenizer_worker_num=1,
|
||||
base_gpu_id=2,
|
||||
tp_size=8,
|
||||
rl_on_policy_target=None,
|
||||
allowed_media_domains=[],
|
||||
media_url_max_file_size_mb=64,
|
||||
)
|
||||
@@ -195,9 +201,13 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
|
||||
# transport policy must still resolve from the instance's ServerArgs.
|
||||
from sglang.srt.multimodal.processors import base_processor
|
||||
|
||||
with envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(True), patch.object(
|
||||
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
), patch.object(base_processor, "MmItemMemoryPool") as memory_pool:
|
||||
with (
|
||||
envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(True),
|
||||
patch.object(
|
||||
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
),
|
||||
patch.object(base_processor, "MmItemMemoryPool") as memory_pool,
|
||||
):
|
||||
processor = base_processor.BaseMultimodalProcessor(
|
||||
hf_config=MagicMock(),
|
||||
server_args=self._server_args("cuda_ipc"),
|
||||
@@ -213,9 +223,13 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
|
||||
def test_cuda_ipc_pool_handle_cache_can_be_disabled(self):
|
||||
from sglang.srt.multimodal.processors import base_processor
|
||||
|
||||
with envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(False), patch.object(
|
||||
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
), patch.object(base_processor, "MmItemMemoryPool") as memory_pool:
|
||||
with (
|
||||
envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(False),
|
||||
patch.object(
|
||||
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
),
|
||||
patch.object(base_processor, "MmItemMemoryPool") as memory_pool,
|
||||
):
|
||||
processor = base_processor.BaseMultimodalProcessor(
|
||||
hf_config=MagicMock(),
|
||||
server_args=self._server_args("cuda_ipc"),
|
||||
@@ -230,9 +244,13 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
|
||||
def test_cpu_transport_does_not_allocate_ipc_pool(self):
|
||||
from sglang.srt.multimodal.processors import base_processor
|
||||
|
||||
with envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(True), patch.object(
|
||||
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
), patch.object(base_processor, "MmItemMemoryPool") as memory_pool:
|
||||
with (
|
||||
envs.SGLANG_USE_IPC_POOL_HANDLE_CACHE.override(True),
|
||||
patch.object(
|
||||
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
),
|
||||
patch.object(base_processor, "MmItemMemoryPool") as memory_pool,
|
||||
):
|
||||
processor = base_processor.BaseMultimodalProcessor(
|
||||
hf_config=MagicMock(),
|
||||
server_args=self._server_args("cpu"),
|
||||
@@ -251,9 +269,12 @@ class TestMultimodalFeatureTransportRuntime(CustomTestCase):
|
||||
hf_processor = self._processor()
|
||||
feature = torch.empty(1, device="meta")
|
||||
hf_processor.return_value = {"pixel_values": feature}
|
||||
with patch.object(
|
||||
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
), patch.object(base_processor, "MmItemMemoryPool") as memory_pool:
|
||||
with (
|
||||
patch.object(
|
||||
base_processor.BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
),
|
||||
patch.object(base_processor, "MmItemMemoryPool") as memory_pool,
|
||||
):
|
||||
processor = base_processor.BaseMultimodalProcessor(
|
||||
hf_config=MagicMock(),
|
||||
server_args=self._server_args("cuda_vmm"),
|
||||
@@ -365,9 +386,10 @@ class TestPrecomputeHashBeforeCpuTransfer(CustomTestCase):
|
||||
BaseMultimodalProcessor,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
), patch.object(BaseMultimodalProcessor, "__init__", lambda self: None):
|
||||
with (
|
||||
patch.object(BaseMultimodalProcessor, "__abstractmethods__", set()),
|
||||
patch.object(BaseMultimodalProcessor, "__init__", lambda self: None),
|
||||
):
|
||||
processor = BaseMultimodalProcessor()
|
||||
processor.precompute_hash_before_cpu_transfer = enabled
|
||||
processor.use_cuda_ipc = False
|
||||
@@ -409,9 +431,10 @@ class TestMultimodalProcessorConcurrency(unittest.IsolatedAsyncioTestCase):
|
||||
MultimodalProcessorExecutor,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
), patch.object(BaseMultimodalProcessor, "__init__", lambda self: None):
|
||||
with (
|
||||
patch.object(BaseMultimodalProcessor, "__abstractmethods__", set()),
|
||||
patch.object(BaseMultimodalProcessor, "__init__", lambda self: None),
|
||||
):
|
||||
processor = BaseMultimodalProcessor()
|
||||
|
||||
processor.mm_processor_executor = MultimodalProcessorExecutor(
|
||||
@@ -438,9 +461,10 @@ class TestMultimodalProcessorConcurrency(unittest.IsolatedAsyncioTestCase):
|
||||
BaseMultimodalProcessor,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
BaseMultimodalProcessor, "__abstractmethods__", set()
|
||||
), patch.object(BaseMultimodalProcessor, "__init__", lambda self: None):
|
||||
with (
|
||||
patch.object(BaseMultimodalProcessor, "__abstractmethods__", set()),
|
||||
patch.object(BaseMultimodalProcessor, "__init__", lambda self: None),
|
||||
):
|
||||
processor = BaseMultimodalProcessor()
|
||||
|
||||
processor.mm_processor_executor = None
|
||||
|
||||
@@ -1,7 +1,14 @@
|
||||
"""CPU coverage for Kimi-K2.5/K2.7 encoder-DP wiring."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import functools
|
||||
import io
|
||||
import pickle
|
||||
import tempfile
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
@@ -18,11 +25,26 @@ from sglang.srt.managers.schedule_batch import (
|
||||
MultimodalInputs,
|
||||
MultimodalProcessorOutput,
|
||||
)
|
||||
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
|
||||
from sglang.srt.models.kimi_k25 import (
|
||||
KimiK25ForConditionalGeneration,
|
||||
mm_projection_auto,
|
||||
)
|
||||
from sglang.srt.models.kimi_vl_moonvit import tpool_patch_merger
|
||||
from sglang.srt.multimodal.cache import (
|
||||
MultimodalPreprocessCache,
|
||||
resolve_multimodal_item_hash,
|
||||
snapshot_media,
|
||||
)
|
||||
from sglang.srt.multimodal.kimi_k3_image_processing import (
|
||||
DEFERRED_PREPROCESSING_KEY,
|
||||
KimiK3DeferredPreprocessing,
|
||||
)
|
||||
from sglang.srt.multimodal.media_artifacts.kimi_k3 import (
|
||||
KimiK3ImagePreprocessArtifact,
|
||||
KimiK3PreprocessConfig,
|
||||
KimiK3ResizeConfig,
|
||||
)
|
||||
from sglang.srt.multimodal.mm_utils import run_dp_sharded_mrope_vision_model
|
||||
from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor
|
||||
from sglang.srt.multimodal.processors.kimi_common import KimiGridMMDataMixin
|
||||
@@ -45,6 +67,8 @@ from sglang.srt.multimodal.transport.cuda_ipc import (
|
||||
CudaIpcTensorTransportProxy,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_context, get_parallel
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils import ImageData
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
@@ -320,8 +344,9 @@ def test_dp_helper_supports_moonvit3d_packed_embeddings_on_tp1():
|
||||
|
||||
# The IPC consumer count asks for the *configured* TP size (matching
|
||||
# MmItemMemoryPool.try_to_recycle), so the double publishes one too.
|
||||
with get_context().override_server_args(tp_size=1), get_parallel().override(
|
||||
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
|
||||
with (
|
||||
get_context().override_server_args(tp_size=1),
|
||||
get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0),
|
||||
):
|
||||
output = run_dp_sharded_mrope_vision_model(
|
||||
tower, pixel_values, [[1, 2, 2]], rope_type="rope_2d_packed"
|
||||
@@ -338,8 +363,9 @@ def test_dp_helper_can_lazily_load_kimi_features_on_tp1():
|
||||
|
||||
# The IPC consumer count asks for the *configured* TP size (matching
|
||||
# MmItemMemoryPool.try_to_recycle), so the double publishes one too.
|
||||
with get_context().override_server_args(tp_size=1), get_parallel().override(
|
||||
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
|
||||
with (
|
||||
get_context().override_server_args(tp_size=1),
|
||||
get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0),
|
||||
):
|
||||
output = run_dp_sharded_mrope_vision_model(
|
||||
tower,
|
||||
@@ -490,8 +516,9 @@ def test_kimi_non_dp_keeps_grid_thws_on_the_host():
|
||||
|
||||
# The IPC consumer count asks for the *configured* TP size (matching
|
||||
# MmItemMemoryPool.try_to_recycle), so the double publishes one too.
|
||||
with get_context().override_server_args(tp_size=1), get_parallel().override(
|
||||
tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0
|
||||
with (
|
||||
get_context().override_server_args(tp_size=1),
|
||||
get_parallel().override(tp_size=1, tp_rank=0, attn_tp_size=1, attn_tp_rank=0),
|
||||
):
|
||||
model.get_image_feature(items)
|
||||
|
||||
@@ -582,6 +609,30 @@ class _HFProcessor:
|
||||
)
|
||||
|
||||
|
||||
class _AnySizeTokenizer:
|
||||
def encode(self, text, allowed_special=None):
|
||||
if text.startswith("<|media_begin|>image "):
|
||||
return [10, 11]
|
||||
if text == "<|media_end|>":
|
||||
return [14]
|
||||
return []
|
||||
|
||||
|
||||
def _k3_preprocess_config(
|
||||
*, patch_size=14, in_patch_limit=16384
|
||||
) -> KimiK3PreprocessConfig:
|
||||
return KimiK3PreprocessConfig(
|
||||
patch_size=patch_size,
|
||||
merge_kernel_size=2,
|
||||
in_patch_limit=in_patch_limit,
|
||||
patch_limit_on_one_side=512,
|
||||
fixed_output_tokens=None,
|
||||
image_mean=(0.5, 0.5, 0.5),
|
||||
image_std=(0.5, 0.5, 0.5),
|
||||
transparent_bg_config=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("processor_cls", "wrapper_cls"),
|
||||
[
|
||||
@@ -592,13 +643,17 @@ class _HFProcessor:
|
||||
def test_kimi_processor_workers_clone_the_gpu_wrapper(processor_cls, wrapper_cls):
|
||||
server_args = SimpleNamespace(
|
||||
mm_feature_transport="cpu",
|
||||
image_processor_backend="auto",
|
||||
disable_fast_image_processor=False,
|
||||
skip_tokenizer_init=False,
|
||||
mm_process_config={},
|
||||
mm_io_worker_num=0,
|
||||
mm_processor_worker_num=0,
|
||||
tokenizer_worker_num=1,
|
||||
mm_preprocess_cache_size_mb=0,
|
||||
trust_mm_content_hashes=False,
|
||||
base_gpu_id=0,
|
||||
rl_on_policy_target=None,
|
||||
allowed_media_domains=[],
|
||||
media_url_max_file_size_mb=64,
|
||||
)
|
||||
@@ -615,6 +670,12 @@ def test_kimi_processor_workers_clone_the_gpu_wrapper(processor_cls, wrapper_cls
|
||||
assert isinstance(processor._processor, wrapper_cls)
|
||||
assert isinstance(worker_processor, wrapper_cls)
|
||||
assert worker_processor is not processor._processor
|
||||
if processor_cls is KimiK3ImageProcessor:
|
||||
fingerprint_config = processor.preprocess_fingerprint_payload()[
|
||||
"wrapped_processor"
|
||||
]
|
||||
assert isinstance(fingerprint_config, KimiK3PreprocessConfig)
|
||||
assert fingerprint_config.patch_size == 14
|
||||
finally:
|
||||
processor.mm_processor_executor.shutdown()
|
||||
processor.io_executor.shutdown()
|
||||
@@ -690,18 +751,498 @@ def test_kimi_k3_epd_rebuild_uses_the_same_media_contract():
|
||||
)
|
||||
|
||||
|
||||
def test_kimi_k3_cpu_transport_defers_gpu_preprocessing():
|
||||
from sglang.srt.multimodal.kimi_k3_image_processing import (
|
||||
DEFERRED_PREPROCESSING_KEY,
|
||||
KimiK3DeferredPreprocessing,
|
||||
def _cached_k3_artifact(content_digest, artifact_key, value=1):
|
||||
return KimiK3ImagePreprocessArtifact(
|
||||
content_digest=content_digest,
|
||||
artifact_key=artifact_key,
|
||||
feature_hash=123,
|
||||
original_size=(1536, 1024),
|
||||
resize_config=KimiK3ResizeConfig(
|
||||
num_tokens=3,
|
||||
new_width=6,
|
||||
new_height=2,
|
||||
pad_width=0,
|
||||
pad_height=0,
|
||||
),
|
||||
grid_thw=(1, 2, 6),
|
||||
feature=torch.full((12, 2), value, dtype=torch.float32),
|
||||
)
|
||||
|
||||
|
||||
def test_kimi_k3_cached_artifact_is_composed_per_prompt():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_tokens = SimpleNamespace(image_token_id=99)
|
||||
processor._tokenizer = _Tokenizer()
|
||||
processor.mm_feature_transport = "cpu"
|
||||
processor.use_cuda_ipc = False
|
||||
artifact = _cached_k3_artifact("sha256:" + "ab" * 32, "artifact")
|
||||
|
||||
first = processor.compose_request([1, 99, 2], [artifact])
|
||||
second = processor.compose_request([3, 4, 99, 5], [artifact])
|
||||
|
||||
assert first.input_ids != second.input_ids
|
||||
assert first.mm_items[0].offsets == [(3, 5)]
|
||||
assert second.mm_items[0].offsets == [(4, 6)]
|
||||
assert first.mm_items[0].hash == second.mm_items[0].hash == 123
|
||||
torch.testing.assert_close(first.mm_items[0].feature, second.mm_items[0].feature)
|
||||
|
||||
|
||||
def test_kimi_k3_cached_deferred_artifact_has_model_contract():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_feature_transport = "cpu"
|
||||
feature = torch.zeros((3, 2, 2), dtype=torch.uint8)
|
||||
|
||||
artifact = processor._make_artifact(
|
||||
content_digest="sha256:" + "ab" * 32,
|
||||
artifact_key="sha256:" + "cd" * 32,
|
||||
original_size=(2, 2),
|
||||
resize_config={
|
||||
"num_tokens": 1,
|
||||
"new_width": 2,
|
||||
"new_height": 2,
|
||||
"pad_width": 0,
|
||||
"pad_height": 0,
|
||||
},
|
||||
grid_thw=(1, 1, 1),
|
||||
feature=feature,
|
||||
deferred=KimiK3DeferredPreprocessing(
|
||||
backend="gpu",
|
||||
image_mean=[0.5, 0.5, 0.5],
|
||||
image_std=[0.5, 0.5, 0.5],
|
||||
transparent_bg_config=None,
|
||||
resize_config={
|
||||
"num_tokens": 1,
|
||||
"new_width": 2,
|
||||
"new_height": 2,
|
||||
"pad_width": 0,
|
||||
"pad_height": 0,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
config = artifact.deferred
|
||||
assert config.backend == "gpu"
|
||||
assert config.resize_config["new_width"] == 2
|
||||
|
||||
|
||||
def test_kimi_k3_normal_cache_path_connects_real_producer_to_model_consumer():
|
||||
hf_processor = _HFProcessor()
|
||||
hf_processor.tokenizer = _AnySizeTokenizer()
|
||||
hf_config = SimpleNamespace(
|
||||
media_placeholder_token_id=42,
|
||||
to_dict=lambda: {
|
||||
"model_type": "kimi_k3",
|
||||
"architectures": ["KimiK3ForConditionalGeneration"],
|
||||
},
|
||||
)
|
||||
server_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
mm_feature_transport="cpu",
|
||||
mm_process_config={},
|
||||
mm_io_worker_num=1,
|
||||
mm_processor_worker_num=0,
|
||||
tokenizer_worker_num=1,
|
||||
mm_preprocess_cache_size_mb=1,
|
||||
)
|
||||
processor = KimiK3ImageProcessor(
|
||||
hf_config=hf_config,
|
||||
server_args=server_args,
|
||||
_processor=hf_processor,
|
||||
transport_mode=None,
|
||||
)
|
||||
image = Image.new("RGB", (28, 28), color=(1, 2, 3))
|
||||
encoded_image = io.BytesIO()
|
||||
image.save(encoded_image, format="PNG")
|
||||
image_data = ImageData(
|
||||
url="data:image/png;base64,"
|
||||
+ base64.b64encode(encoded_image.getvalue()).decode()
|
||||
)
|
||||
request = SimpleNamespace(video_data=None, mm_content_hashes=None)
|
||||
|
||||
class _Tower(nn.Module):
|
||||
device = torch.device("cpu")
|
||||
patch_size = 14
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.patch_embed = SimpleNamespace(
|
||||
proj=SimpleNamespace(weight=torch.empty(1, dtype=torch.float32))
|
||||
)
|
||||
|
||||
def forward(self, pixel_values, _grid_thws):
|
||||
return pixel_values
|
||||
|
||||
model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
|
||||
nn.Module.__init__(model)
|
||||
model.use_data_parallel = False
|
||||
model.vision_tower = _Tower()
|
||||
model.mm_projector = _Projector()
|
||||
|
||||
try:
|
||||
with (
|
||||
patch(
|
||||
"sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True
|
||||
),
|
||||
patch.object(
|
||||
processor,
|
||||
"prepare_artifact_batch",
|
||||
wraps=processor.prepare_artifact_batch,
|
||||
) as prepare_artifacts,
|
||||
):
|
||||
cold = asyncio.run(
|
||||
processor.process_mm_data_async([image_data], [1, 42, 2], request)
|
||||
)
|
||||
hot = asyncio.run(
|
||||
processor.process_mm_data_async([image_data], [3, 42, 4], request)
|
||||
)
|
||||
cold_items = pickle.loads(pickle.dumps(cold.mm_items))
|
||||
hot_items = pickle.loads(pickle.dumps(hot.mm_items))
|
||||
|
||||
with (
|
||||
patch("sglang.srt.models.kimi_k3.configured_tp_size", return_value=1),
|
||||
patch(
|
||||
"sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images",
|
||||
return_value=(
|
||||
torch.ones((4, 3), dtype=torch.float32),
|
||||
torch.tensor([[1, 2, 2]], dtype=torch.int64),
|
||||
),
|
||||
),
|
||||
):
|
||||
cold_features = model.get_image_feature(cold_items)
|
||||
hot_features = model.get_image_feature(hot_items)
|
||||
finally:
|
||||
processor.shutdown()
|
||||
|
||||
assert prepare_artifacts.call_count == 1
|
||||
assert cold.mm_items[0].hash == hot.mm_items[0].hash
|
||||
assert cold.mm_items[0].offsets == hot.mm_items[0].offsets == [(3, 3)]
|
||||
assert (
|
||||
cold_items[0].model_specific_data[DEFERRED_PREPROCESSING_KEY].backend == "gpu"
|
||||
)
|
||||
torch.testing.assert_close(cold_features, hot_features)
|
||||
|
||||
|
||||
def test_kimi_k3_model_accepts_mixed_cached_eager_and_deferred_artifacts():
|
||||
class _Tower(nn.Module):
|
||||
device = torch.device("cpu")
|
||||
patch_size = 2
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.patch_embed = SimpleNamespace(
|
||||
proj=SimpleNamespace(weight=torch.empty(1, dtype=torch.float32))
|
||||
)
|
||||
|
||||
def forward(self, pixel_values, _grid_thws):
|
||||
return pixel_values
|
||||
|
||||
model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
|
||||
nn.Module.__init__(model)
|
||||
model.use_data_parallel = False
|
||||
model.vision_tower = _Tower()
|
||||
model.mm_projector = _Projector()
|
||||
eager = _image_item(torch.ones((1, 3)), [[1, 1, 1]])
|
||||
deferred = _image_item(torch.zeros((3, 2, 2), dtype=torch.uint8), [[1, 1, 1]])
|
||||
deferred.model_specific_data[DEFERRED_PREPROCESSING_KEY] = (
|
||||
KimiK3DeferredPreprocessing(
|
||||
backend="gpu",
|
||||
image_mean=[0.5, 0.5, 0.5],
|
||||
image_std=[0.5, 0.5, 0.5],
|
||||
transparent_bg_config=None,
|
||||
resize_config={
|
||||
"num_tokens": 1,
|
||||
"new_width": 2,
|
||||
"new_height": 2,
|
||||
"pad_width": 0,
|
||||
"pad_height": 0,
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
with (
|
||||
patch("sglang.srt.models.kimi_k3.configured_tp_size", return_value=1),
|
||||
patch(
|
||||
"sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images",
|
||||
return_value=(torch.full((1, 3), 2.0), torch.tensor([[1, 1, 1]])),
|
||||
),
|
||||
):
|
||||
output = model.get_image_feature([eager, deferred])
|
||||
|
||||
torch.testing.assert_close(output, torch.tensor([[1.0, 1.0, 1.0], [2.0, 2.0, 2.0]]))
|
||||
|
||||
|
||||
def test_kimi_k3_trusted_hot_hit_skips_media_read():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.processor_fingerprint = "processor"
|
||||
processor.trust_mm_content_hashes = True
|
||||
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
|
||||
processor.io_executor = ThreadPoolExecutor(max_workers=1)
|
||||
digest = "sha256:" + "ab" * 32
|
||||
key = processor._artifact_key(digest, "unread-source")
|
||||
artifact = _cached_k3_artifact(digest, key)
|
||||
processor.mm_preprocess_cache.put(key, artifact)
|
||||
|
||||
try:
|
||||
with patch(
|
||||
"sglang.srt.multimodal.media_artifacts.base.snapshot_media",
|
||||
side_effect=AssertionError("trusted cache hit must not read media"),
|
||||
):
|
||||
result = asyncio.run(
|
||||
processor.prepare_media_artifacts(
|
||||
["unread-source"],
|
||||
content_hashes=[digest],
|
||||
)
|
||||
)
|
||||
finally:
|
||||
processor.io_executor.shutdown()
|
||||
|
||||
assert result == [artifact]
|
||||
|
||||
|
||||
def test_kimi_k3_default_media_options_share_one_artifact_key():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.processor_fingerprint = "processor"
|
||||
digest = "sha256:" + "ab" * 32
|
||||
|
||||
keys = {
|
||||
processor._artifact_key(digest, "image.png"),
|
||||
processor._artifact_key(digest, ImageData(url="image.png")),
|
||||
processor._artifact_key(digest, {"url": "image.png", "detail": "auto"}),
|
||||
}
|
||||
|
||||
assert len(keys) == 1
|
||||
|
||||
|
||||
def test_kimi_k3_output_affecting_media_options_do_not_share_artifacts():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.processor_fingerprint = "processor"
|
||||
digest = "sha256:" + "ab" * 32
|
||||
base = processor._artifact_key(digest, ImageData(url="image.png"))
|
||||
|
||||
assert base != processor._artifact_key(
|
||||
digest, ImageData(url="image.png", detail="low")
|
||||
)
|
||||
assert base != processor._artifact_key(
|
||||
digest, ImageData(url="image.png", max_dynamic_patch=4)
|
||||
)
|
||||
assert base != processor._artifact_key(
|
||||
digest,
|
||||
ImageData(url="image.png", preprocess_kwargs={"max_pixels": 1024}),
|
||||
)
|
||||
assert base != processor._artifact_key(
|
||||
digest,
|
||||
{"url": "image.png", "future_model_option": "new-behavior"},
|
||||
)
|
||||
|
||||
|
||||
def test_kimi_k3_rejects_changed_feature_hash_for_same_artifact():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.processor_fingerprint = "processor"
|
||||
processor.trust_mm_content_hashes = False
|
||||
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
|
||||
processor.mm_processor_executor = None
|
||||
processor.io_executor = ThreadPoolExecutor(max_workers=2)
|
||||
image = Image.new("RGB", (2, 2), color=(1, 2, 3))
|
||||
digest = snapshot_media(image).content_digest
|
||||
key = processor._artifact_key(digest, image)
|
||||
old = replace(_cached_k3_artifact(digest, key), feature=None)
|
||||
new = replace(_cached_k3_artifact(digest, key), feature_hash=old.feature_hash + 1)
|
||||
processor.mm_preprocess_cache.put(key, old)
|
||||
|
||||
async def prepare(_entries):
|
||||
return [new]
|
||||
|
||||
processor._run_preprocess_and_build_artifact_batch = prepare
|
||||
try:
|
||||
with pytest.raises(ValueError, match="feature hash changed"):
|
||||
asyncio.run(processor.prepare_media_artifacts([image]))
|
||||
finally:
|
||||
processor.io_executor.shutdown()
|
||||
|
||||
assert key not in processor.mm_preprocess_cache
|
||||
|
||||
|
||||
def test_kimi_k3_artifact_and_data_item_share_hash_resolution():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_feature_transport = "cpu"
|
||||
processor.mm_tokens = SimpleNamespace(image_token_id=99)
|
||||
processor._tokenizer = _Tokenizer()
|
||||
processor.use_cuda_ipc = False
|
||||
feature = torch.zeros((4, 3), dtype=torch.float32)
|
||||
digest = "sha256:" + "ab" * 32
|
||||
|
||||
artifact = processor._make_artifact(
|
||||
content_digest=digest,
|
||||
artifact_key="sha256:" + "01" * 32,
|
||||
original_size=(2, 2),
|
||||
resize_config={
|
||||
"num_tokens": 1,
|
||||
"new_width": 2,
|
||||
"new_height": 2,
|
||||
"pad_width": 0,
|
||||
"pad_height": 0,
|
||||
},
|
||||
grid_thw=(1, 1, 1),
|
||||
feature=feature,
|
||||
)
|
||||
direct_item = MultimodalDataItem(modality=Modality.IMAGE, feature=feature)
|
||||
direct_item.set_pad_value()
|
||||
composed_item = processor.compose_request([1, 99, 2], [artifact]).mm_items[0]
|
||||
|
||||
expected_hash = resolve_multimodal_item_hash(
|
||||
existing_hash=direct_item.hash,
|
||||
namespace=artifact.artifact_key,
|
||||
)
|
||||
expected_item = MultimodalDataItem(modality=Modality.IMAGE, hash=expected_hash)
|
||||
expected_item.set_pad_value()
|
||||
assert artifact.feature_hash == composed_item.hash == expected_hash
|
||||
assert composed_item.pad_value == expected_item.pad_value
|
||||
|
||||
|
||||
def test_kimi_k3_untrusted_path_change_is_a_cache_miss():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.processor_fingerprint = "processor"
|
||||
processor.trust_mm_content_hashes = False
|
||||
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
|
||||
processor.mm_processor_executor = None
|
||||
processor.io_executor = ThreadPoolExecutor(max_workers=2)
|
||||
|
||||
async def prepare(entries):
|
||||
return [
|
||||
_cached_k3_artifact(
|
||||
entry.content_digest,
|
||||
entry.artifact_key,
|
||||
entry.media.getpixel((0, 0))[0],
|
||||
)
|
||||
for entry in entries
|
||||
]
|
||||
|
||||
processor._run_preprocess_and_build_artifact_batch = prepare
|
||||
try:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "mutable.png"
|
||||
Image.new("RGB", (2, 2), color=(1, 0, 0)).save(path)
|
||||
first = asyncio.run(processor.prepare_media_artifacts([str(path)]))[0]
|
||||
Image.new("RGB", (2, 2), color=(2, 0, 0)).save(path)
|
||||
second = asyncio.run(processor.prepare_media_artifacts([str(path)]))[0]
|
||||
finally:
|
||||
processor.io_executor.shutdown()
|
||||
|
||||
assert first.content_digest != second.content_digest
|
||||
assert first.artifact_key != second.artifact_key
|
||||
assert first.feature[0, 0].item() == 1
|
||||
assert second.feature[0, 0].item() == 2
|
||||
|
||||
|
||||
def test_kimi_k3_partial_hits_deduplicate_misses_and_preserve_order():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.processor_fingerprint = "processor"
|
||||
processor.trust_mm_content_hashes = False
|
||||
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
|
||||
processor.mm_processor_executor = None
|
||||
processor.io_executor = ThreadPoolExecutor(max_workers=4)
|
||||
cached_image = Image.new("RGB", (2, 2), color=(1, 0, 0))
|
||||
missed_image = Image.new("RGB", (2, 2), color=(2, 0, 0))
|
||||
cached_digest = snapshot_media(cached_image).content_digest
|
||||
missed_digest = snapshot_media(missed_image).content_digest
|
||||
cached_key = processor._artifact_key(cached_digest, cached_image)
|
||||
cached = _cached_k3_artifact(cached_digest, cached_key, value=1)
|
||||
processor.mm_preprocess_cache.put(cached_key, cached)
|
||||
batches = []
|
||||
|
||||
async def prepare(entries):
|
||||
batches.append(entries)
|
||||
return [
|
||||
replace(
|
||||
_cached_k3_artifact(
|
||||
entry.content_digest,
|
||||
entry.artifact_key,
|
||||
entry.media.getpixel((0, 0))[0],
|
||||
),
|
||||
feature_hash=456,
|
||||
)
|
||||
for entry in entries
|
||||
]
|
||||
|
||||
processor._run_preprocess_and_build_artifact_batch = prepare
|
||||
try:
|
||||
artifacts = asyncio.run(
|
||||
processor.prepare_media_artifacts(
|
||||
[cached_image, missed_image, missed_image, cached_image],
|
||||
)
|
||||
)
|
||||
finally:
|
||||
processor.io_executor.shutdown()
|
||||
|
||||
assert len(batches) == 1
|
||||
assert len(batches[0]) == 1
|
||||
assert batches[0][0].content_digest == missed_digest
|
||||
assert batches[0][0].artifact_key == processor._artifact_key(
|
||||
missed_digest, missed_image
|
||||
)
|
||||
assert [artifact.content_digest for artifact in artifacts] == [
|
||||
cached_digest,
|
||||
missed_digest,
|
||||
missed_digest,
|
||||
cached_digest,
|
||||
]
|
||||
assert artifacts[0] is artifacts[3]
|
||||
assert artifacts[1] is artifacts[2]
|
||||
|
||||
|
||||
def test_kimi_k3_cancelled_artifact_owner_does_not_fail_joiner():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.processor_fingerprint = "processor"
|
||||
processor.trust_mm_content_hashes = False
|
||||
processor.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
|
||||
processor.mm_processor_executor = None
|
||||
processor.io_executor = ThreadPoolExecutor(max_workers=2)
|
||||
processor._preprocess_metrics_callback = None
|
||||
image = Image.new("RGB", (2, 2), color=(1, 2, 3))
|
||||
started = asyncio.Event()
|
||||
release = asyncio.Event()
|
||||
|
||||
async def prepare(entries):
|
||||
started.set()
|
||||
await release.wait()
|
||||
return [
|
||||
_cached_k3_artifact(
|
||||
entry.content_digest,
|
||||
entry.artifact_key,
|
||||
entry.media.getpixel((0, 0))[0],
|
||||
)
|
||||
for entry in entries
|
||||
]
|
||||
|
||||
processor._run_preprocess_and_build_artifact_batch = prepare
|
||||
|
||||
async def run():
|
||||
owner = asyncio.create_task(processor.prepare_media_artifacts([image]))
|
||||
await started.wait()
|
||||
joiner = asyncio.create_task(processor.prepare_media_artifacts([image]))
|
||||
await asyncio.sleep(0)
|
||||
owner.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await owner
|
||||
|
||||
release.set()
|
||||
artifacts = await joiner
|
||||
assert len(artifacts) == 1
|
||||
assert artifacts[0].feature[0, 0].item() == 1
|
||||
|
||||
try:
|
||||
asyncio.run(run())
|
||||
finally:
|
||||
processor.io_executor.shutdown()
|
||||
|
||||
|
||||
def test_kimi_k3_cpu_transport_defers_gpu_preprocessing():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_tokens = SimpleNamespace(image_token_id=99)
|
||||
processor.mm_feature_transport = "cpu"
|
||||
processor.use_cuda_ipc = False
|
||||
processor._processor = SimpleNamespace(
|
||||
_patch_size=2,
|
||||
preprocess_config=_k3_preprocess_config(patch_size=2),
|
||||
prepare_deferred=Mock(
|
||||
return_value=(
|
||||
torch.tensor([[1, 99, 99, 2, 99, 3]]),
|
||||
@@ -776,11 +1317,7 @@ def test_kimi_k3_defers_only_when_raw_transport_is_smaller(
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_feature_transport = "cpu"
|
||||
processor._processor = SimpleNamespace(
|
||||
_patch_size=14,
|
||||
_merge_kernel_size=2,
|
||||
_in_patch_limit=in_patch_limit,
|
||||
_patch_limit_on_one_side=512,
|
||||
_fixed_output_tokens=None,
|
||||
preprocess_config=_k3_preprocess_config(in_patch_limit=in_patch_limit),
|
||||
)
|
||||
image = torch.zeros(image_shape, dtype=torch.uint8)
|
||||
|
||||
@@ -817,7 +1354,7 @@ def test_kimi_k3_eager_preprocessing_preserves_float_tensor_support():
|
||||
assert output.shape == (3, 4, 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("transport", ["cuda_ipc", "fabric"])
|
||||
@pytest.mark.parametrize("transport", ["cuda_ipc", "cuda_vmm"])
|
||||
def test_kimi_k3_keeps_gpu_transport_preprocessing_eager(transport):
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_feature_transport = transport
|
||||
@@ -831,6 +1368,7 @@ def test_kimi_k3_keeps_gpu_transport_preprocessing_eager(transport):
|
||||
def test_kimi_k3_rejects_silently_dropped_images():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_tokens = Mock()
|
||||
processor.mm_preprocess_cache = MultimodalPreprocessCache(0)
|
||||
processor.load_mm_data = AsyncMock(return_value=SimpleNamespace(images=[object()]))
|
||||
|
||||
with pytest.raises(ValueError, match="expected 2, loaded 1"):
|
||||
@@ -838,13 +1376,14 @@ def test_kimi_k3_rejects_silently_dropped_images():
|
||||
processor.process_mm_data_async(
|
||||
image_data=["image-1", "image-2"],
|
||||
input_text="<|media_pad|><|media_pad|>",
|
||||
request_obj=SimpleNamespace(video_data=None),
|
||||
request_obj=SimpleNamespace(video_data=None, mm_content_hashes=None),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_kimi_k3_uses_token_ids_to_preserve_media_boundaries():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_preprocess_cache = MultimodalPreprocessCache(0)
|
||||
processor.mm_feature_transport = "cpu"
|
||||
processor.mm_tokens = SimpleNamespace(image_token_id=99)
|
||||
processor.mm_feature_transport = "cuda_ipc"
|
||||
@@ -863,7 +1402,7 @@ def test_kimi_k3_uses_token_ids_to_preserve_media_boundaries():
|
||||
processor.process_mm_data_async(
|
||||
image_data=["image-1", "image-2"],
|
||||
input_text=[1, 99, 2, 99, 3],
|
||||
request_obj=SimpleNamespace(video_data=None),
|
||||
request_obj=SimpleNamespace(video_data=None, mm_content_hashes=None),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -69,8 +69,11 @@ def make_processor(config, image_processor_cls=None):
|
||||
model_impl="sglang",
|
||||
keep_mm_feature_on_device=False,
|
||||
mm_feature_transport="cpu",
|
||||
image_processor_backend="auto",
|
||||
disable_fast_image_processor=True,
|
||||
skip_tokenizer_init=False,
|
||||
mm_preprocess_cache_size_mb=0,
|
||||
trust_mm_content_hashes=False,
|
||||
# Read by NativeMmHost._use_feature_shm (single-rank fixture → the
|
||||
# inline zero-copy transport, like the 1-GPU e2e).
|
||||
tp_size=1,
|
||||
@@ -80,6 +83,7 @@ def make_processor(config, image_processor_cls=None):
|
||||
mm_processor_worker_num=1,
|
||||
tokenizer_worker_num=1,
|
||||
base_gpu_id=0,
|
||||
rl_on_policy_target=None,
|
||||
allowed_media_domains=[],
|
||||
media_url_max_file_size_mb=64,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
import asyncio
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Optional
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Modality
|
||||
from sglang.srt.multimodal.cache import MultimodalPreprocessCache, snapshot_media
|
||||
from sglang.srt.multimodal.media_artifacts import (
|
||||
MediaArtifactCacheMixin,
|
||||
MediaArtifactInput,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Artifact:
|
||||
content_digest: str
|
||||
artifact_key: str
|
||||
feature_hash: int
|
||||
feature: Optional[bytes]
|
||||
|
||||
@property
|
||||
def has_feature(self) -> bool:
|
||||
return self.feature is not None
|
||||
|
||||
def cache_value(self):
|
||||
return self
|
||||
|
||||
def cache_size_items(self):
|
||||
return (
|
||||
self.content_digest,
|
||||
self.artifact_key,
|
||||
self.feature_hash,
|
||||
self.feature,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _FutureMediaInput:
|
||||
url: str
|
||||
content_hash: Optional[str] = None
|
||||
frame_sampling: int = 2
|
||||
|
||||
|
||||
class _Processor(MediaArtifactCacheMixin):
|
||||
artifact_modality = Modality.IMAGE
|
||||
artifact_option_defaults = {"detail": "auto", "frame_sampling": 2}
|
||||
|
||||
def __init__(self):
|
||||
self.processor_fingerprint = "processor"
|
||||
self.trust_mm_content_hashes = False
|
||||
self.mm_preprocess_cache = MultimodalPreprocessCache(1024 * 1024)
|
||||
self.mm_processor_executor = None
|
||||
self.io_executor = ThreadPoolExecutor(max_workers=4)
|
||||
self.batches = []
|
||||
|
||||
def decode_media_snapshot(self, snapshot, modality):
|
||||
self.assert_artifact_modality(modality)
|
||||
return snapshot.data
|
||||
|
||||
@staticmethod
|
||||
def assert_artifact_modality(modality):
|
||||
if modality != Modality.IMAGE:
|
||||
raise AssertionError(f"unexpected modality: {modality}")
|
||||
|
||||
def prepare_artifact_batch(
|
||||
self, entries: list[MediaArtifactInput]
|
||||
) -> list[_Artifact]:
|
||||
self.batches.append(entries)
|
||||
return [
|
||||
_Artifact(
|
||||
content_digest=entry.content_digest,
|
||||
artifact_key=entry.artifact_key,
|
||||
feature_hash=int(entry.content_digest[-16:], 16),
|
||||
feature=entry.media,
|
||||
)
|
||||
for entry in entries
|
||||
]
|
||||
|
||||
def close(self):
|
||||
self.io_executor.shutdown()
|
||||
|
||||
|
||||
class TestMediaArtifactProcessor(unittest.TestCase):
|
||||
def test_unknown_model_option_is_part_of_artifact_identity(self):
|
||||
processor = _Processor()
|
||||
digest = snapshot_media(b"image").content_digest
|
||||
try:
|
||||
base = processor._artifact_key(digest, _FutureMediaInput(url="image.png"))
|
||||
self.assertEqual(
|
||||
base,
|
||||
processor._artifact_key(
|
||||
digest,
|
||||
{
|
||||
"url": "image.png",
|
||||
"content_hash": digest,
|
||||
"frame_sampling": 2,
|
||||
},
|
||||
),
|
||||
)
|
||||
self.assertNotEqual(
|
||||
base,
|
||||
processor._artifact_key(
|
||||
digest,
|
||||
{"url": "image.png", "future_model_knob": "different"},
|
||||
),
|
||||
)
|
||||
self.assertNotEqual(
|
||||
base,
|
||||
processor._artifact_key(
|
||||
digest,
|
||||
_FutureMediaInput(url="image.png"),
|
||||
modality=Modality.VIDEO,
|
||||
),
|
||||
)
|
||||
finally:
|
||||
processor.close()
|
||||
|
||||
def test_non_image_models_can_override_identity_and_decode_hooks(self):
|
||||
class _VideoProcessor(_Processor):
|
||||
artifact_modality = Modality.VIDEO
|
||||
|
||||
def snapshot_media_source(self, source, modality):
|
||||
self.assert_video_modality(modality)
|
||||
return snapshot_media(source.encode())
|
||||
|
||||
def decode_media_snapshot(self, snapshot, modality):
|
||||
self.assert_video_modality(modality)
|
||||
return snapshot.data
|
||||
|
||||
@staticmethod
|
||||
def assert_video_modality(modality):
|
||||
if modality != Modality.VIDEO:
|
||||
raise AssertionError(f"unexpected modality: {modality}")
|
||||
|
||||
processor = _VideoProcessor()
|
||||
try:
|
||||
artifacts = asyncio.run(processor.prepare_media_artifacts(["clip.mp4"]))
|
||||
finally:
|
||||
processor.close()
|
||||
|
||||
self.assertEqual(len(artifacts), 1)
|
||||
self.assertEqual(
|
||||
artifacts[0].content_digest, snapshot_media(b"clip.mp4").content_digest
|
||||
)
|
||||
|
||||
def test_partial_hits_and_duplicate_misses_are_shared_by_contract(self):
|
||||
processor = _Processor()
|
||||
first_digest = snapshot_media(b"first").content_digest
|
||||
first_key = processor._artifact_key(first_digest, b"first")
|
||||
first = _Artifact(first_digest, first_key, 1, b"first")
|
||||
processor.mm_preprocess_cache.put(first_key, first)
|
||||
|
||||
try:
|
||||
artifacts = asyncio.run(
|
||||
processor.prepare_media_artifacts(
|
||||
[b"first", b"second", b"second", b"first"]
|
||||
)
|
||||
)
|
||||
finally:
|
||||
processor.close()
|
||||
|
||||
self.assertEqual(len(processor.batches), 1)
|
||||
self.assertEqual(len(processor.batches[0]), 1)
|
||||
self.assertEqual(
|
||||
[artifact.content_digest for artifact in artifacts],
|
||||
[
|
||||
first_digest,
|
||||
snapshot_media(b"second").content_digest,
|
||||
snapshot_media(b"second").content_digest,
|
||||
first_digest,
|
||||
],
|
||||
)
|
||||
self.assertIs(artifacts[0], artifacts[3])
|
||||
self.assertIs(artifacts[1], artifacts[2])
|
||||
|
||||
def test_adapter_cannot_change_validated_artifact_identity(self):
|
||||
processor = _Processor()
|
||||
|
||||
async def wrong_identity(entries):
|
||||
artifact = processor.prepare_artifact_batch(entries)[0]
|
||||
return [replace(artifact, artifact_key="sha256:" + "0" * 64)]
|
||||
|
||||
processor._run_preprocess_and_build_artifact_batch = wrong_identity
|
||||
try:
|
||||
with self.assertRaisesRegex(ValueError, "changed the media artifact key"):
|
||||
asyncio.run(processor.prepare_media_artifacts([b"image"]))
|
||||
finally:
|
||||
processor.close()
|
||||
|
||||
def test_trusted_hit_uses_identity_without_reading_source(self):
|
||||
processor = _Processor()
|
||||
processor.trust_mm_content_hashes = True
|
||||
digest = snapshot_media(b"cached").content_digest
|
||||
key = processor._artifact_key(digest, "unread-source")
|
||||
artifact = _Artifact(digest, key, 1, b"cached")
|
||||
processor.mm_preprocess_cache.put(key, artifact)
|
||||
|
||||
try:
|
||||
artifacts = asyncio.run(
|
||||
processor.prepare_media_artifacts(
|
||||
["unread-source"], content_hashes=[digest]
|
||||
)
|
||||
)
|
||||
finally:
|
||||
processor.close()
|
||||
|
||||
self.assertEqual(artifacts, [artifact])
|
||||
self.assertEqual(processor.batches, [])
|
||||
|
||||
def test_cached_artifact_must_match_content_identity(self):
|
||||
processor = _Processor()
|
||||
digest = snapshot_media(b"fresh").content_digest
|
||||
key = processor._artifact_key(digest, b"fresh")
|
||||
processor.mm_preprocess_cache.put(
|
||||
key,
|
||||
_Artifact(
|
||||
snapshot_media(b"stale").content_digest,
|
||||
key,
|
||||
1,
|
||||
b"stale",
|
||||
),
|
||||
)
|
||||
try:
|
||||
artifacts = asyncio.run(processor.prepare_media_artifacts([b"fresh"]))
|
||||
finally:
|
||||
processor.close()
|
||||
|
||||
self.assertEqual(artifacts[0].content_digest, digest)
|
||||
self.assertEqual(artifacts[0].feature, b"fresh")
|
||||
self.assertEqual(len(processor.batches), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -4,22 +4,24 @@ import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
||||
from sglang.srt.multimodal.cache import (
|
||||
CacheMiss,
|
||||
MultimodalPreprocessCache,
|
||||
build_artifact_key,
|
||||
build_feature_hash,
|
||||
build_processor_fingerprint,
|
||||
estimate_cache_size_bytes,
|
||||
parse_content_hash,
|
||||
resolve_multimodal_item_hash,
|
||||
snapshot_media,
|
||||
)
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
@@ -209,21 +211,25 @@ class TestMediaIdentity(unittest.TestCase):
|
||||
def preprocess_fingerprint_payload(self):
|
||||
return {"backend": self.backend, "antialias": True}
|
||||
|
||||
config = SimpleNamespace(model_type="vlm", architectures=["VLM"])
|
||||
args = SimpleNamespace(
|
||||
class Config:
|
||||
def to_dict(self):
|
||||
return {"model_type": "vlm", "architectures": ["VLM"]}
|
||||
|
||||
config = Config()
|
||||
args = ServerArgs(
|
||||
model_path="dummy",
|
||||
revision="model-revision",
|
||||
tokenizer_revision="tokenizer-revision",
|
||||
disable_fast_image_processor=False,
|
||||
mm_process_config={"image": {"max_pixels": 1024}},
|
||||
)
|
||||
base = build_processor_fingerprint(Processor("gpu"), config, args)
|
||||
|
||||
changed_backend = build_processor_fingerprint(Processor("cpu"), config, args)
|
||||
changed_args = SimpleNamespace(
|
||||
**{
|
||||
**vars(args),
|
||||
"mm_process_config": {"image": {"max_pixels": 2048}},
|
||||
}
|
||||
changed_args = ServerArgs(
|
||||
model_path="dummy",
|
||||
revision="model-revision",
|
||||
disable_fast_image_processor=False,
|
||||
mm_process_config={"image": {"max_pixels": 2048}},
|
||||
)
|
||||
changed_config = build_processor_fingerprint(
|
||||
Processor("gpu"), config, changed_args
|
||||
@@ -231,7 +237,7 @@ class TestMediaIdentity(unittest.TestCase):
|
||||
self.assertNotEqual(base, changed_backend)
|
||||
self.assertNotEqual(base, changed_config)
|
||||
|
||||
def test_feature_hash_includes_artifact_and_processor_output(self):
|
||||
def test_item_hash_namespace_covers_identity_and_processor_output(self):
|
||||
digest = snapshot_media(b"image").content_digest
|
||||
first = build_artifact_key(
|
||||
digest,
|
||||
@@ -243,11 +249,25 @@ class TestMediaIdentity(unittest.TestCase):
|
||||
modality="image",
|
||||
processor_fingerprint="processor-b",
|
||||
)
|
||||
self.assertNotEqual(build_feature_hash(first, 1), build_feature_hash(second, 1))
|
||||
self.assertNotEqual(build_feature_hash(first, 1), build_feature_hash(first, 2))
|
||||
self.assertIsInstance(build_feature_hash(first, 1 << 128), int)
|
||||
self.assertNotEqual(
|
||||
resolve_multimodal_item_hash(existing_hash=1, namespace=first),
|
||||
resolve_multimodal_item_hash(existing_hash=1, namespace=second),
|
||||
)
|
||||
self.assertNotEqual(
|
||||
resolve_multimodal_item_hash(existing_hash=1, namespace=first),
|
||||
resolve_multimodal_item_hash(existing_hash=2, namespace=first),
|
||||
)
|
||||
with self.assertRaises(ValueError):
|
||||
build_feature_hash(first, -1)
|
||||
resolve_multimodal_item_hash(existing_hash=-1, namespace=first)
|
||||
|
||||
def test_multimodal_data_item_uses_shared_feature_hash(self):
|
||||
feature = torch.arange(12, dtype=torch.float32).reshape(4, 3)
|
||||
expected = resolve_multimodal_item_hash(feature=feature)
|
||||
item = MultimodalDataItem(modality=Modality.IMAGE, feature=feature)
|
||||
|
||||
item.set_pad_value()
|
||||
|
||||
self.assertEqual(item.hash, expected)
|
||||
|
||||
|
||||
class TestMultimodalPreprocessCache(unittest.TestCase):
|
||||
@@ -262,6 +282,31 @@ class TestMultimodalPreprocessCache(unittest.TestCase):
|
||||
self.assertIn("c", cache)
|
||||
self.assertEqual(cache.current_size_bytes, 6)
|
||||
|
||||
def test_compatible_lookup_is_atomic_and_does_not_count_bypass_as_miss(self):
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
cache.put("key", b"metadata-only")
|
||||
|
||||
self.assertIsNone(cache.get_if_present("key", lambda value: False))
|
||||
self.assertEqual((cache.hits, cache.misses), (0, 0))
|
||||
self.assertEqual(
|
||||
cache.get_if_present("key", lambda value: value.startswith(b"metadata")),
|
||||
b"metadata-only",
|
||||
)
|
||||
self.assertEqual((cache.hits, cache.misses), (1, 0))
|
||||
|
||||
def test_claimed_miss_rejects_an_incompatible_racing_entry(self):
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
cache.put("key", b"metadata-only")
|
||||
|
||||
miss = cache.lookup_or_claim_many(
|
||||
["key"], predicate=lambda key, value: value == b"full-feature"
|
||||
)[0]
|
||||
|
||||
self.assertIsInstance(miss, CacheMiss)
|
||||
self.assertTrue(miss.should_compute)
|
||||
self.assertNotIn("key", cache)
|
||||
self.assertEqual(cache.current_size_bytes, 0)
|
||||
|
||||
def test_gpu_backed_values_are_not_implicitly_copied(self):
|
||||
if not torch.cuda.is_available():
|
||||
self.skipTest("CUDA is not available")
|
||||
@@ -370,6 +415,61 @@ class TestMultimodalPreprocessCache(unittest.TestCase):
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_lookup_or_claim_many_batches_owned_and_joined_misses(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
results = cache.lookup_or_claim_many(["a", "b", "a"])
|
||||
misses_to_compute = [
|
||||
item
|
||||
for item in results
|
||||
if isinstance(item, CacheMiss) and item.should_compute
|
||||
]
|
||||
self.assertEqual([item.key for item in misses_to_compute], ["a", "b"])
|
||||
|
||||
cache.complete_miss(misses_to_compute[0], b"value-a")
|
||||
cache.complete_miss(misses_to_compute[1], b"value-b")
|
||||
self.assertEqual(await cache.wait_for_miss(results[2]), b"value-a")
|
||||
self.assertEqual(cache.get("b"), b"value-b")
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_cancelled_miss_waiter_does_not_cancel_computing_caller(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
computing_miss = cache.lookup_or_claim_many(["key"])[0]
|
||||
waiting_miss = cache.lookup_or_claim_many(["key"])[0]
|
||||
self.assertTrue(computing_miss.should_compute)
|
||||
self.assertFalse(waiting_miss.should_compute)
|
||||
|
||||
waiter = asyncio.create_task(cache.wait_for_miss(waiting_miss))
|
||||
await asyncio.sleep(0)
|
||||
waiter.cancel()
|
||||
with self.assertRaises(asyncio.CancelledError):
|
||||
await waiter
|
||||
|
||||
cache.complete_miss(computing_miss, b"artifact")
|
||||
self.assertEqual(computing_miss.future.result(), b"artifact")
|
||||
self.assertEqual(cache.get("key"), b"artifact")
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_disabled_cache_does_not_join_or_retain(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=0)
|
||||
misses = cache.lookup_or_claim_many(["a", "a"])
|
||||
self.assertTrue(
|
||||
all(
|
||||
isinstance(item, CacheMiss) and item.should_compute
|
||||
for item in misses
|
||||
)
|
||||
)
|
||||
for item in misses:
|
||||
cache.complete_miss(item, b"value")
|
||||
self.assertEqual(len(cache), 0)
|
||||
self.assertEqual(cache.stats()["singleflight_joins"], 0)
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_clear_starts_a_new_singleflight_generation(self):
|
||||
async def run():
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
@@ -397,6 +497,20 @@ class TestMultimodalPreprocessCache(unittest.TestCase):
|
||||
|
||||
asyncio.run(run())
|
||||
|
||||
def test_clear_starts_a_new_cache_miss_generation(self):
|
||||
cache = MultimodalPreprocessCache[str, bytes](max_size_bytes=1024)
|
||||
old = cache.lookup_or_claim_many(["key"])[0]
|
||||
cache.clear()
|
||||
new = cache.lookup_or_claim_many(["key"])[0]
|
||||
|
||||
self.assertTrue(old.should_compute)
|
||||
self.assertTrue(new.should_compute)
|
||||
self.assertIsNot(old.future, new.future)
|
||||
cache.complete_miss(old, b"old")
|
||||
self.assertNotIn("key", cache)
|
||||
cache.complete_miss(new, b"new")
|
||||
self.assertEqual(cache.get("key"), b"new")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user