[Kimi K3] optimize: preprocess cpu-transport images on the vision owner (#33921)
This commit is contained in:
Mick
@@ -680,6 +680,138 @@ 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,
|
||||
)
|
||||
|
||||
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,
|
||||
prepare_deferred=Mock(
|
||||
return_value=(
|
||||
torch.tensor([[1, 99, 99, 2, 99, 3]]),
|
||||
[
|
||||
{
|
||||
"num_tokens": 2,
|
||||
"new_width": 4,
|
||||
"new_height": 2,
|
||||
"pad_width": 0,
|
||||
"pad_height": 2,
|
||||
},
|
||||
{
|
||||
"num_tokens": 1,
|
||||
"new_width": 2,
|
||||
"new_height": 2,
|
||||
"pad_width": 2,
|
||||
"pad_height": 2,
|
||||
},
|
||||
],
|
||||
{
|
||||
"image_mean": [0.5, 0.5, 0.5],
|
||||
"image_std": [0.5, 0.5, 0.5],
|
||||
"transparent_bg_config": None,
|
||||
},
|
||||
)
|
||||
),
|
||||
)
|
||||
images = [
|
||||
torch.arange(3 * 2 * 4, dtype=torch.uint8).reshape(3, 2, 4),
|
||||
torch.arange(3 * 2 * 2, dtype=torch.uint8).reshape(3, 2, 2),
|
||||
]
|
||||
base_output = SimpleNamespace(
|
||||
input_text="prompt", images=images, input_ids=[1, 99, 2, 99, 3]
|
||||
)
|
||||
|
||||
output = processor._build_deferred_output(base_output)
|
||||
|
||||
assert output.input_ids == [1, 99, 99, 2, 99, 3]
|
||||
assert [item.offsets for item in output.mm_items] == [[(1, 2)], [(4, 4)]]
|
||||
assert [item.feature.dtype for item in output.mm_items] == [
|
||||
torch.uint8,
|
||||
torch.uint8,
|
||||
]
|
||||
assert [item.feature.shape for item in output.mm_items] == [
|
||||
torch.Size([3, 2, 4]),
|
||||
torch.Size([3, 2, 2]),
|
||||
]
|
||||
assert [item.image_grid_thw.tolist() for item in output.mm_items] == [
|
||||
[[1, 2, 2]],
|
||||
[[1, 2, 2]],
|
||||
]
|
||||
assert all(item.hash is not None for item in output.mm_items)
|
||||
assert all(item.pad_value is not None for item in output.mm_items)
|
||||
assert all(
|
||||
DEFERRED_PREPROCESSING_KEY in item.model_specific_data
|
||||
for item in output.mm_items
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("image_shape", "in_patch_limit", "expected"),
|
||||
[((3, 32, 32), 65536, True), ((3, 1024, 1024), 1, False)],
|
||||
)
|
||||
def test_kimi_k3_defers_only_when_raw_transport_is_smaller(
|
||||
image_shape, in_patch_limit, expected
|
||||
):
|
||||
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,
|
||||
)
|
||||
image = torch.zeros(image_shape, dtype=torch.uint8)
|
||||
|
||||
with patch("sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True):
|
||||
assert processor._should_defer_gpu_preprocessing([image]) is expected
|
||||
|
||||
|
||||
def test_kimi_k3_does_not_defer_non_uint8_tensor_preprocessing():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_feature_transport = "cpu"
|
||||
|
||||
with patch("sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True):
|
||||
assert not processor._should_defer_gpu_preprocessing(
|
||||
[torch.zeros((3, 32, 32), dtype=torch.float32)]
|
||||
)
|
||||
|
||||
|
||||
def test_kimi_k3_does_not_defer_empty_image_batch():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_feature_transport = "cpu"
|
||||
|
||||
with patch("sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True):
|
||||
assert not processor._should_defer_gpu_preprocessing([])
|
||||
|
||||
|
||||
def test_kimi_k3_eager_preprocessing_preserves_float_tensor_support():
|
||||
from sglang.srt.multimodal.processors.kimi_k3 import _k3_to_cuda_chw
|
||||
|
||||
image = torch.zeros((1, 4, 4), dtype=torch.float32)
|
||||
with patch.object(torch.Tensor, "cuda", lambda self: self):
|
||||
output = _k3_to_cuda_chw(image)
|
||||
|
||||
assert output.dtype == torch.float32
|
||||
assert output.shape == (3, 4, 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("transport", ["cuda_ipc", "fabric"])
|
||||
def test_kimi_k3_keeps_gpu_transport_preprocessing_eager(transport):
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_feature_transport = transport
|
||||
|
||||
with patch("sglang.srt.multimodal.processors.kimi_k3.is_cuda", return_value=True):
|
||||
assert not processor._should_defer_gpu_preprocessing(
|
||||
[torch.zeros((3, 32, 32), dtype=torch.uint8)]
|
||||
)
|
||||
|
||||
|
||||
def test_kimi_k3_rejects_silently_dropped_images():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_tokens = Mock()
|
||||
@@ -699,6 +831,8 @@ def test_kimi_k3_uses_token_ids_to_preserve_media_boundaries():
|
||||
processor = object.__new__(KimiK3ImageProcessor)
|
||||
processor.mm_feature_transport = "cpu"
|
||||
processor.mm_tokens = SimpleNamespace(image_token_id=99)
|
||||
processor.mm_feature_transport = "cuda_ipc"
|
||||
processor.use_cuda_ipc = True
|
||||
processor.fast_load_mm_data = AsyncMock(
|
||||
return_value=SimpleNamespace(
|
||||
images=[object(), object()], input_ids=[1, 99, 2, 99, 3]
|
||||
|
||||
@@ -167,7 +167,7 @@ def test_kimi_k3_vision_tower_reuses_prepared_forward_metadata(monkeypatch):
|
||||
)
|
||||
|
||||
assert len(actual) == len(reference) == 1
|
||||
assert torch.equal(actual[0], reference[0])
|
||||
torch.testing.assert_close(actual[0], reference[0], rtol=0, atol=0, equal_nan=True)
|
||||
|
||||
|
||||
def test_kimi_k3_dp_helper_passes_host_grid_list_to_capable_tower():
|
||||
@@ -424,6 +424,7 @@ if __name__ == "__main__":
|
||||
class _K3TowerStub:
|
||||
device = torch.device("cpu")
|
||||
merge_kernel_size = (2, 2)
|
||||
patch_size = 2
|
||||
|
||||
def __init__(self):
|
||||
self.config = SimpleNamespace(hidden_size=2)
|
||||
@@ -467,14 +468,11 @@ def test_kimi_k3_encoder_dp_defers_feature_materialization(monkeypatch):
|
||||
"sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model",
|
||||
return_value=sharded_embeddings,
|
||||
) as run_dp, mock_patch(
|
||||
"sglang.srt.models.kimi_k3.get_server_args",
|
||||
return_value=SimpleNamespace(tp_size=1),
|
||||
), mock_patch(
|
||||
"sglang.srt.models.kimi_k3.get_parallel",
|
||||
return_value=SimpleNamespace(attn_tp_size=1),
|
||||
return_value=SimpleNamespace(tp_size=1, attn_tp_size=1),
|
||||
):
|
||||
output = model.get_image_feature(items)
|
||||
# exercise the loader inside the patch scope: it reads server args
|
||||
# Exercise the loader while the runtime topology is patched.
|
||||
loader_in_scope = run_dp.call_args.kwargs["load_local_pixel_values"]
|
||||
local = loader_in_scope([1])
|
||||
both = loader_in_scope([0, 1])
|
||||
@@ -502,6 +500,80 @@ def test_kimi_k3_encoder_dp_defers_feature_materialization(monkeypatch):
|
||||
)
|
||||
|
||||
|
||||
def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
|
||||
from unittest.mock import patch as mock_patch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
||||
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
|
||||
from sglang.srt.multimodal.kimi_k3_image_processing import (
|
||||
DEFERRED_PREPROCESSING_KEY,
|
||||
)
|
||||
|
||||
model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.use_data_parallel = True
|
||||
model.vision_tower = _K3TowerStub()
|
||||
model.mm_projector = lambda image_embeds: image_embeds
|
||||
|
||||
deferred_config = {
|
||||
"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,
|
||||
},
|
||||
}
|
||||
items = [
|
||||
MultimodalDataItem(
|
||||
modality=Modality.IMAGE,
|
||||
offsets=[(index, index)],
|
||||
feature=torch.full((3, 2, 2), index, dtype=torch.uint8),
|
||||
model_specific_data={
|
||||
"image_grid_thw": torch.tensor([[1, 1, 1]]),
|
||||
DEFERRED_PREPROCESSING_KEY: deferred_config,
|
||||
},
|
||||
)
|
||||
for index in range(2)
|
||||
]
|
||||
calls = []
|
||||
|
||||
def fake_preprocess(images, resize_configs, *args, **kwargs):
|
||||
calls.append([int(image[0, 0, 0]) for image in images])
|
||||
return torch.tensor([[float(calls[-1][0]), 0.0]]), torch.tensor([[1, 1, 1]])
|
||||
|
||||
with mock_patch(
|
||||
"sglang.srt.multimodal.mm_utils.run_dp_sharded_mrope_vision_model",
|
||||
return_value=torch.zeros(1, 2),
|
||||
) as run_dp, mock_patch(
|
||||
"sglang.srt.models.kimi_k3.get_parallel",
|
||||
return_value=SimpleNamespace(tp_size=1, attn_tp_size=1),
|
||||
), mock_patch(
|
||||
"sglang.srt.multimodal.processors.kimi_k25._gpu_preprocess_images",
|
||||
side_effect=fake_preprocess,
|
||||
):
|
||||
model.get_image_feature(items)
|
||||
loader = run_dp.call_args.kwargs["load_local_pixel_values"]
|
||||
one = loader([1])
|
||||
|
||||
assert calls == [[1]]
|
||||
assert one.dtype == torch.float32
|
||||
assert one.tolist() == [[1.0, 0.0]]
|
||||
|
||||
|
||||
def test_kimi_k3_scheduler_leaves_feature_placement_to_dp_owner():
|
||||
from sglang.srt.managers.mm_schedule import _can_skip_pre_embed_feature_move
|
||||
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
|
||||
|
||||
model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
|
||||
assert _can_skip_pre_embed_feature_move(model.get_image_feature)
|
||||
|
||||
|
||||
def test_kimi_k3_rejects_aggregated_items():
|
||||
"""One item must carry exactly one logical image: the DP owner
|
||||
assignment and the bounded CUDA-IPC lease accounting are per-item, so
|
||||
|
||||
Reference in New Issue
Block a user