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

[Fix] Carry the backend on Kimi-K3 deferred preprocessing configs (#34766)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-13 13:30:33 -07:00
committed by GitHub
parent 8ad04a9bee
commit 8554d9a5bc
6 changed files with 80 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/kimi_k3.py
+5 -5
View File
@@ -3276,8 +3276,8 @@ class KimiK3ForConditionalGeneration(nn.Module):
"Kimi-K3 cannot mix deferred and preprocessed image features"
)
first_config = deferred[0]
backend = first_config["backend"]
if any(config["backend"] != backend for config in deferred):
backend = first_config.backend
if any(config.backend != backend for config in deferred):
raise ValueError(
"Kimi-K3 cannot mix deferred preprocessing backends"
)
@@ -3287,17 +3287,17 @@ class KimiK3ForConditionalGeneration(nn.Module):
)
image_scale, image_bias = normalization_tensors(
first_config["image_mean"], first_config["image_std"], device
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],
[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"]
x, first_config.transparent_bg_config
),
)
elif backend == "cpu":
python/sglang/srt/multimodal/kimi_k3_image_processing.py
+28 -15
View File
@@ -1,6 +1,7 @@
import functools
import math
from typing import Union
from dataclasses import dataclass
from typing import Literal, Optional, Union
import numpy as np
import torch
@@ -9,6 +10,22 @@ from PIL import Image
DEFERRED_PREPROCESSING_KEY = "kimi_k3_deferred_preprocessing"
@dataclass(frozen=True)
class KimiK3DeferredPreprocessing:
"""Parameters the vision-DP owner needs to finish one deferred image.
``backend`` is the producer's decision and cannot be recovered from the
item; the feature's own layout stays observable on ``item.feature``, so it
is not mirrored here.
"""
backend: Literal["gpu", "cpu"]
image_mean: list[float]
image_std: list[float]
transparent_bg_config: Optional[dict]
resize_config: dict
def prepare_kimi_k3_encoder_inputs(
images, image_processor, *, use_gpu_preprocessing=False
):
@@ -61,12 +78,13 @@ def prepare_kimi_k3_encoder_inputs(
patch_size = int(media_proc_cfg["patch_size"])
merge_kernel_size = int(media_proc_cfg["merge_kernel_size"])
common_deferred_config = {
"backend": "gpu" if use_gpu_preprocessing else "cpu",
"image_mean": list(media_proc_cfg["image_mean"]),
"image_std": list(media_proc_cfg["image_std"]),
"transparent_bg_config": media_proc_cfg.get("transparent_bg_config"),
}
deferred_preprocessing = functools.partial(
KimiK3DeferredPreprocessing,
backend="gpu" if use_gpu_preprocessing else "cpu",
image_mean=list(media_proc_cfg["image_mean"]),
image_std=list(media_proc_cfg["image_std"]),
transparent_bg_config=media_proc_cfg.get("transparent_bg_config"),
)
items = []
grids = []
@@ -93,11 +111,9 @@ def prepare_kimi_k3_encoder_inputs(
feature=to_chw_uint8(image) if use_gpu_preprocessing else image,
model_specific_data={
"grid_thws": grid_tensor,
DEFERRED_PREPROCESSING_KEY: {
**common_deferred_config,
"feature_layout": "chw" if use_gpu_preprocessing else "raw",
"resize_config": resize_config,
},
DEFERRED_PREPROCESSING_KEY: deferred_preprocessing(
resize_config=resize_config
),
},
)
if not use_gpu_preprocessing:
@@ -133,9 +149,6 @@ def materialize_kimi_k3_cpu_features(items, image_processor) -> torch.Tensor:
medias = []
for item in items:
image = item.feature
config = item.model_specific_data[DEFERRED_PREPROCESSING_KEY]
if config["feature_layout"] != "raw":
raise ValueError("Kimi-K3 deferred CPU preprocessing expects raw inputs")
if not isinstance(image, Image.Image):
if not isinstance(image, torch.Tensor) or image.dtype != torch.uint8:
raise TypeError(
python/sglang/srt/multimodal/processors/kimi_k3.py
+20 -11
View File
@@ -8,6 +8,7 @@ images onto the checkpoint-configured background
at load time.
"""
import functools
import re
from typing import Dict, List, Union
@@ -23,6 +24,7 @@ from sglang.srt.managers.schedule_batch import (
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
from sglang.srt.multimodal.kimi_k3_image_processing import (
DEFERRED_PREPROCESSING_KEY,
KimiK3DeferredPreprocessing,
)
from sglang.srt.multimodal.kimi_k3_image_processing import (
fill_transparent_bg as _fill_transparent_bg,
@@ -272,12 +274,16 @@ class KimiK3GPUProcessorWrapper(KimiGPUProcessorWrapper):
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
# This path only ever defers GPU preprocessing: the caller gates on
# `_should_defer_gpu_preprocessing` and stages CHW uint8 features.
deferred_preprocessing = functools.partial(
KimiK3DeferredPreprocessing,
backend="gpu",
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_preprocessing
class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
@@ -365,7 +371,11 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
return raw_bytes <= processed_bytes
def _build_deferred_output(self, base_output):
input_ids, resize_configs, deferred_config = self._processor.prepare_deferred(
(
input_ids,
resize_configs,
deferred_preprocessing,
) = self._processor.prepare_deferred(
base_output.input_text,
base_output.images,
base_output.input_ids,
@@ -389,10 +399,9 @@ class KimiK3ImageProcessor(KimiGridMMDataMixin, SGLangBaseProcessor):
offsets=[offset],
model_specific_data={
"image_grid_thw": torch.tensor([grid_thw], dtype=torch.int64),
DEFERRED_PREPROCESSING_KEY: {
**deferred_config,
"resize_config": resize_config,
},
DEFERRED_PREPROCESSING_KEY: deferred_preprocessing(
resize_config=resize_config
),
},
)
items.append(item)
test/registered/unit/disaggregation/test_kimi_k3_encoder_mode.py
+3 -3
View File
@@ -188,8 +188,8 @@ def test_kimi_k3_epd_preprocess_preserves_raw_per_image_items():
assert item.hash is not None
assert item.pad_value is not None
deferred = item.model_specific_data[DEFERRED_PREPROCESSING_KEY]
assert deferred["image_mean"] == [0.5, 0.5, 0.5]
assert deferred["image_std"] == [0.5, 0.5, 0.5]
assert deferred.image_mean == [0.5, 0.5, 0.5]
assert deferred.image_std == [0.5, 0.5, 0.5]
def test_kimi_k3_epd_model_preprocessor_receives_image_processor():
@@ -289,7 +289,7 @@ def test_kimi_k3_epd_default_cpu_materialization_is_owner_only_and_exact():
assert len(processor.calls[0]) == 1
assert processor.calls[0][0]["image"].getpixel((0, 0)) == (11, 0, 0)
assert torch.all(materialized == 11)
assert items[0].model_specific_data[DEFERRED_PREPROCESSING_KEY]["backend"] == "cpu"
assert items[0].model_specific_data[DEFERRED_PREPROCESSING_KEY].backend == "cpu"
def test_encoder_preprocess_materializes_only_local_size_balanced_items():
test/registered/unit/models/test_kimi_k25.py
+16 -9
View File
@@ -1,6 +1,7 @@
"""CPU coverage for Kimi-K2.5/K2.7 encoder-DP wiring."""
import asyncio
import functools
from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock, patch
@@ -690,6 +691,7 @@ 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,
)
processor = object.__new__(KimiK3ImageProcessor)
@@ -717,11 +719,13 @@ def test_kimi_k3_cpu_transport_defers_gpu_preprocessing():
"pad_height": 2,
},
],
{
"image_mean": [0.5, 0.5, 0.5],
"image_std": [0.5, 0.5, 0.5],
"transparent_bg_config": None,
},
functools.partial(
KimiK3DeferredPreprocessing,
backend="gpu",
image_mean=[0.5, 0.5, 0.5],
image_std=[0.5, 0.5, 0.5],
transparent_bg_config=None,
),
)
),
)
@@ -751,10 +755,13 @@ def test_kimi_k3_cpu_transport_defers_gpu_preprocessing():
]
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
)
deferred = [
item.model_specific_data[DEFERRED_PREPROCESSING_KEY] for item in output.mm_items
]
# The staged features are CHW uint8, so the config has to route them to the
# GPU arm of `materialize_item_features`.
assert [config.backend for config in deferred] == ["gpu", "gpu"]
assert [config.resize_config["new_width"] for config in deferred] == [4, 2]
@pytest.mark.parametrize(
test/registered/unit/models/test_kimi_k3_vision.py
+8 -8
View File
@@ -509,6 +509,7 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
from sglang.srt.models.kimi_k3 import KimiK3ForConditionalGeneration
from sglang.srt.multimodal.kimi_k3_image_processing import (
DEFERRED_PREPROCESSING_KEY,
KimiK3DeferredPreprocessing,
)
model = KimiK3ForConditionalGeneration.__new__(KimiK3ForConditionalGeneration)
@@ -517,20 +518,19 @@ def test_kimi_k3_preprocesses_only_dp_owner_images(monkeypatch):
model.vision_tower = _K3TowerStub()
model.mm_projector = lambda image_embeds: image_embeds
deferred_config = {
"backend": "gpu",
"feature_layout": "chw",
"image_mean": [0.5, 0.5, 0.5],
"image_std": [0.5, 0.5, 0.5],
"transparent_bg_config": None,
"resize_config": {
deferred_config = 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,
},
}
)
items = [
MultimodalDataItem(
modality=Modality.IMAGE,