minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Kimi K3] optimize: preprocess cpu-transport images on the vision owner (#33921)

Browse Source
This commit is contained in:
Mick
2026-08-09 16:13:20 +08:00
committed by GitHub
parent d0aa37b49b
commit 22e003580b
6 changed files with 470 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/mm_schedule.py
+1
View File
@@ -178,6 +178,7 @@ def _can_skip_pre_embed_feature_move(data_embedding_func: DataEmbeddingFunc) ->
"Qwen3_5ForConditionalGeneration",
"Qwen3_5MoeForConditionalGeneration",
"KimiK25ForConditionalGeneration",
"KimiK3ForConditionalGeneration",
}
python/sglang/srt/models/kimi_k3.py
+46 -21
View File
@@ -104,6 +104,12 @@ from sglang.srt.models.kimi_k3_vl import (
)
from sglang.srt.models.transformers import maybe_prefix
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.multimodal.kimi_k3_image_processing import (
DEFERRED_PREPROCESSING_KEY,
fill_transparent_bg,
normalization_tensors,
to_chw_uint8,
)
from sglang.srt.multimodal.mm_utils import materialize_multimodal_features
from sglang.srt.runtime_context import get_exec, get_parallel, get_server_args
from sglang.srt.utils import is_blackwell_supported, is_hip, make_layers
@@ -3075,40 +3081,59 @@ class KimiK3ForConditionalGeneration(nn.Module):
grid_thw_list = grid_thws_host.tolist()
def materialize_item_features(image_indices: List[int]) -> torch.Tensor:
"""Materialize features for the images assigned to this rank.
K3 vision is image-wise data-parallel, so each image is consumed
by exactly one TP rank. Deferred CUDA-IPC proxies are
reconstructed here, after the assignment is known, so an image
crosses the tokenizer/scheduler boundary once instead of once
per rank; CPU-transport features likewise only pay their H2D
copy on the owner rank. The consumer count matches
MmItemMemoryPool.try_to_recycle(), which waits for the server TP
size rather than the attention subgroup size.
"""
parallel = get_parallel()
server_args = get_server_args()
ipc_consumer_count = max(
getattr(server_args, "tp_size", parallel.attn_tp_size), 1
)
"""Materialize only the images assigned to this vision-DP rank."""
ipc_consumer_count = max(get_parallel().tp_size, 1)
device_index = device.index
if device.type == "cuda" and device_index is None:
device_index = torch.cuda.current_device()
features = []
selected_items = []
for image_index in image_indices:
item = items[image_index]
if device.type == "cuda":
item.reconstruct(
device_index, ipc_consumer_count=ipc_consumer_count
)
feature = item.feature
if not isinstance(feature, torch.Tensor):
selected_items.append(item)
deferred = [
item.model_specific_data.get(DEFERRED_PREPROCESSING_KEY)
for item in selected_items
]
if any(config is not None for config in deferred):
if not all(config is not None for config in deferred):
raise ValueError(
"Kimi-K3 cannot mix deferred and preprocessed image features"
)
from sglang.srt.multimodal.processors.kimi_k25 import (
_gpu_preprocess_images,
)
first_config = deferred[0]
image_scale, image_bias = normalization_tensors(
first_config["image_mean"], first_config["image_std"], device
)
pixel_values, _ = _gpu_preprocess_images(
[item.feature for item in selected_items],
[config["resize_config"] for config in deferred],
image_scale,
image_bias,
self.vision_tower.patch_size,
to_chw=lambda image: to_chw_uint8(image, device=device),
post_resize=lambda x: fill_transparent_bg(
x, first_config["transparent_bg_config"]
),
)
return pixel_values.to(dtype=target_dtype)
features = []
for item in selected_items:
if not isinstance(item.feature, torch.Tensor):
raise TypeError(
"Kimi-K3 image feature must be a torch.Tensor, "
f"got {type(feature)}"
f"got {type(item.feature)}"
)
features.append(feature)
features.append(item.feature)
return materialize_multimodal_features(
features, device=device, dtype=target_dtype
)
python/sglang/srt/multimodal/kimi_k3_image_processing.py
+81
View File
@@ -0,0 +1,81 @@
from typing import Union
import numpy as np
import torch
from PIL import Image
DEFERRED_PREPROCESSING_KEY = "kimi_k3_deferred_preprocessing"
def to_chw_uint8(
image: Union[torch.Tensor, Image.Image],
device: torch.device | str | None = None,
) -> torch.Tensor:
if isinstance(image, Image.Image):
has_alpha = image.mode != "RGB" and (
"A" in image.getbands() or "transparency" in image.info
)
array = np.array(image.convert("RGBA" if has_alpha else "RGB"), copy=True)
image = torch.from_numpy(array).permute(2, 0, 1)
if image.dtype != torch.uint8:
raise ValueError(
f"Kimi-K3 preprocessing expects raw uint8 pixels, got {image.dtype}"
)
if image.dim() == 2:
image = image.unsqueeze(0)
if image.shape[0] == 1:
image = image.repeat(3, 1, 1)
if device is not None:
image = image.to(device)
return image
def fill_transparent_bg(x: torch.Tensor, bg_config: Union[dict, None]) -> torch.Tensor:
if x.shape[1] == 3:
return x
rgb = x[:, :3]
if bg_config is None:
return rgb
_, _, height, width = x.shape
pattern = bg_config.get("pattern", "black")
if pattern == "chessboard":
square = bg_config.get("chessboard_square_size", 16)
white = float(bg_config.get("chessboard_white_value", 255))
gray = float(bg_config.get("chessboard_gray_value", 200))
top_left = bg_config.get("chessboard_square_on_top_left", True)
ys = torch.arange(height, device=x.device) // square
xs = torch.arange(width, device=x.device) // square
parity = (ys.unsqueeze(1) + xs.unsqueeze(0)) % 2
background = torch.where(parity == (1 if top_left else 0), gray, white)
background = background.unsqueeze(0).expand(3, height, width)
elif pattern == "white":
background = torch.full((3, height, width), 255.0, device=x.device)
elif pattern == "black":
background = torch.zeros(3, height, width, device=x.device)
elif pattern == "gray":
background = torch.full((3, height, width), 128.0, device=x.device)
else:
raise ValueError(f"Invalid background pattern: {pattern}")
alpha = (x[:, 3:4] / 255.0).clamp(0.0, 1.0)
return (alpha * rgb + (1.0 - alpha) * background).clamp(0.0, 255.0).floor_()
def normalization_tensors(
image_mean: list[float],
image_std: list[float],
device: torch.device | str,
) -> tuple[torch.Tensor, torch.Tensor]:
scale = torch.tensor(
[1.0 / (255.0 * std) for std in image_std],
device=device,
dtype=torch.float32,
).view(1, 3, 1, 1)
bias = torch.tensor(
[-mean / std for mean, std in zip(image_mean, image_std)],
device=device,
dtype=torch.float32,
).view(1, 3, 1, 1)
return scale, bias
python/sglang/srt/multimodal/processors/kimi_k3.py
+130 -55
View File
@@ -15,8 +15,21 @@ import numpy as np
import torch
from PIL import Image
from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput
from sglang.srt.managers.schedule_batch import (
Modality,
MultimodalDataItem,
MultimodalProcessorOutput,
)
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
from sglang.srt.multimodal.kimi_k3_image_processing import (
DEFERRED_PREPROCESSING_KEY,
)
from sglang.srt.multimodal.kimi_k3_image_processing import (
fill_transparent_bg as _fill_transparent_bg,
)
from sglang.srt.multimodal.kimi_k3_image_processing import (
to_chw_uint8,
)
from sglang.srt.multimodal.processors.base_processor import (
BaseMultimodalProcessor as SGLangBaseProcessor,
)
@@ -28,8 +41,10 @@ from sglang.srt.multimodal.processors.kimi_k25 import (
KimiGPUProcessorWrapper,
_get_image_dimensions,
_gpu_preprocess_images,
_grid_thw_from_resize_config,
navit_resize_config,
)
from sglang.srt.utils import is_cuda
from sglang.srt.utils.cuda_ipc_transport_utils import (
DEFER_CUDA_IPC_FEATURE_RECONSTRUCTION_KEY,
)
@@ -124,15 +139,7 @@ def _expand_k3_image_prompt_text(
def _k3_to_cuda_chw(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor:
if isinstance(image, Image.Image):
# The checkpoint's fill_transparent_bg_with() returns RGB-mode images
# untouched before it ever inspects the alpha bands, so an RGB image
# carrying a stray "transparency" info key must NOT be promoted to
# RGBA here.
has_alpha = image.mode != "RGB" and (
"A" in image.getbands() or "transparency" in image.info
)
arr = np.asarray(image.convert("RGBA" if has_alpha else "RGB"))
return torch.from_numpy(arr).permute(2, 0, 1).cuda()
return to_chw_uint8(image, device="cuda")
image = image.cuda()
if image.dim() == 2:
@@ -142,51 +149,6 @@ def _k3_to_cuda_chw(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor:
return image
def _chessboard_background(
height: int, width: int, cfg: dict, device: torch.device
) -> torch.Tensor:
square = cfg.get("chessboard_square_size", 16)
white = float(cfg.get("chessboard_white_value", 255))
gray = float(cfg.get("chessboard_gray_value", 200))
on_top_left = cfg.get("chessboard_square_on_top_left", True)
ys = torch.arange(height, device=device) // square
xs = torch.arange(width, device=device) // square
parity = (ys.unsqueeze(1) + xs.unsqueeze(0)) % 2
gray_parity = 1 if on_top_left else 0
bg = torch.where(parity == gray_parity, gray, white)
return bg.unsqueeze(0).expand(3, height, width)
def _fill_transparent_bg(x: torch.Tensor, bg_cfg: Union[dict, None]) -> torch.Tensor:
"""Composite a resized (1, 4, H, W) float image in [0, 255] onto the
configured background; 3-channel input passes through."""
if x.shape[1] == 3:
return x
rgb = x[:, :3]
if bg_cfg is None:
return rgb
_, _, height, width = x.shape
pattern = bg_cfg.get("pattern", "black")
if pattern == "chessboard":
bg = _chessboard_background(height, width, bg_cfg, x.device)
elif pattern == "white":
bg = torch.full((3, height, width), 255.0, device=x.device)
elif pattern == "black":
bg = torch.zeros(3, height, width, device=x.device)
elif pattern == "gray":
bg = torch.full((3, height, width), 128.0, device=x.device)
else:
raise ValueError(f"Invalid background pattern: {pattern}")
alpha = (x[:, 3:4] / 255.0).clamp(0.0, 1.0)
# The checkpoint processor casts the composited float result back with
# numpy's astype(np.uint8), which truncates; floor matches that exactly
# (a composite of [0, 255] inputs is always non-negative).
return (alpha * rgb + (1.0 - alpha) * bg).clamp(0.0, 255.0).floor_()
class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper):
def __init__(self, *args, transparent_bg_config=None, **kwargs):
super().__init__(*args, **kwargs)
@@ -292,6 +254,31 @@ class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper):
out["image_grid_thw"] = grid_thws
return out
def prepare_deferred(self, text, images, original_input_ids=None):
input_text = text[0] if isinstance(text, list) else text
image_sizes = [_get_image_dimensions(image) for image in images]
resize_configs = [
navit_resize_config(
width,
height,
self._patch_size,
self._merge_kernel_size,
self._in_patch_limit,
self._patch_limit_on_one_side,
self._fixed_output_tokens,
)
for width, height in image_sizes
]
input_ids = self._prepare_input_ids(
input_text, resize_configs, original_input_ids, image_sizes
)
deferred_config = {
"image_mean": list(self._image_mean),
"image_std": list(self._image_std),
"transparent_bg_config": self._transparent_bg_config,
}
return input_ids, resize_configs, deferred_config
class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
models = [KimiK3ForConditionalGeneration]
@@ -328,6 +315,91 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
super().__init__(hf_config, server_args, processor, *args, **kwargs)
self.mm_tokens = mm_tokens
def _should_defer_gpu_preprocessing(self, images) -> bool:
if (
not images
or self.mm_feature_transport != "cpu"
or not is_cuda()
or not all(
isinstance(image, Image.Image)
or (isinstance(image, torch.Tensor) and image.dtype == torch.uint8)
for image in images
)
):
return False
raw_bytes = 0
processed_bytes = 0
patch_size = self._processor._patch_size
for image in images:
width, height = _get_image_dimensions(image)
resize_config = navit_resize_config(
width,
height,
patch_size,
self._processor._merge_kernel_size,
self._processor._in_patch_limit,
self._processor._patch_limit_on_one_side,
self._processor._fixed_output_tokens,
)
if isinstance(image, torch.Tensor):
channels = (
3 if image.dim() == 2 or image.shape[0] == 1 else image.shape[0]
)
else:
channels = (
4
if image.mode != "RGB"
and ("A" in image.getbands() or "transparency" in image.info)
else 3
)
raw_bytes += channels * width * height
padded_width = resize_config["new_width"] + resize_config["pad_width"]
padded_height = resize_config["new_height"] + resize_config["pad_height"]
processed_bytes += 3 * padded_width * padded_height * torch.float32.itemsize
return raw_bytes <= processed_bytes
def _build_deferred_output(self, base_output):
input_ids, resize_configs, deferred_config = self._processor.prepare_deferred(
base_output.input_text,
base_output.images,
base_output.input_ids,
)
offsets = self.get_mm_items_offset(
input_ids.flatten(), self.mm_tokens.image_token_id
)
if len(offsets) != len(base_output.images):
raise ValueError("Expected one Kimi-K3 image span for each image")
items = []
for image, resize_config, offset in zip(
base_output.images, resize_configs, offsets
):
grid_thw = _grid_thw_from_resize_config(
resize_config, self._processor._patch_size
)
item = MultimodalDataItem(
modality=Modality.IMAGE,
feature=to_chw_uint8(image),
offsets=[offset],
model_specific_data={
"image_grid_thw": torch.tensor([grid_thw], dtype=torch.int64),
DEFERRED_PREPROCESSING_KEY: {
**deferred_config,
"resize_config": resize_config,
},
},
)
items.append(item)
self._precompute_hashes_before_cpu_transfer(items)
return MultimodalProcessorOutput(
input_ids=input_ids.flatten().tolist(),
mm_items=items,
im_token_id=self.mm_tokens.image_token_id,
)
async def process_mm_data_async(
self,
image_data: List[Union[str, bytes, Dict]],
@@ -378,6 +450,9 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
f"expected {expected_image_count}, loaded {len(base_output.images)}"
)
if self._should_defer_gpu_preprocessing(base_output.images):
return self._build_deferred_output(base_output)
mm_items, input_ids, _ = await self.process_and_combine_mm_data_async(
base_output,
self.mm_tokens,
test/registered/unit/models/test_kimi_k25.py
+134
View File
@@ -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]
test/registered/unit/models/test_kimi_k3_vision.py
+78 -6
View File
@@ -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