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,