[EPD] feat: pipeline owner-only multimodal preprocessing (#34206)
This commit is contained in:
Mick
@@ -58,6 +58,11 @@ from sglang.srt.model_executor.model_runner_components.load_model_utils import (
|
||||
maybe_precompile_model_kernels_after_loading,
|
||||
)
|
||||
from sglang.srt.model_loader import get_model
|
||||
from sglang.srt.multimodal.encoder_preprocessing import (
|
||||
EncoderPreprocessOutput,
|
||||
get_encoder_preprocessed_items,
|
||||
invoke_encoder_preprocessor,
|
||||
)
|
||||
from sglang.srt.multimodal.processors.qwen_vl import preprocess_video
|
||||
from sglang.srt.observability.metrics_collector import EncoderMetricsCollector
|
||||
from sglang.srt.observability.req_time_stats import EncoderReqTimeStats
|
||||
@@ -65,7 +70,13 @@ from sglang.srt.observability.trace import (
|
||||
process_tracing_init,
|
||||
trace_set_thread_info,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_disagg, get_exec, get_mm, publish
|
||||
from sglang.srt.runtime_context import (
|
||||
get_disagg,
|
||||
get_exec,
|
||||
get_mm,
|
||||
get_parallel,
|
||||
publish,
|
||||
)
|
||||
from sglang.srt.server_args import (
|
||||
PortArgs,
|
||||
ServerArgs,
|
||||
@@ -862,9 +873,24 @@ class MMEncoder:
|
||||
return slices
|
||||
|
||||
def _calculate_hashes_from_features(
|
||||
self, mm_feature, grid_thw: List, modality: Modality
|
||||
self, mm_feature, grid_thw: List, modality: Modality, mm_inputs=None
|
||||
) -> List[int]:
|
||||
"""CPU Task: Compute hashes based on processed feature patches."""
|
||||
preprocessed_items = (
|
||||
get_encoder_preprocessed_items(mm_inputs) if mm_inputs is not None else None
|
||||
)
|
||||
if preprocessed_items is not None:
|
||||
if len(preprocessed_items) != len(grid_thw):
|
||||
raise ValueError(
|
||||
"Encoder preprocess item/grid mismatch: "
|
||||
f"{len(preprocessed_items)} items != {len(grid_thw)} grids"
|
||||
)
|
||||
hashes = []
|
||||
for item in preprocessed_items:
|
||||
item.set_pad_value()
|
||||
hashes.append(item.hash)
|
||||
return hashes
|
||||
|
||||
hashes = []
|
||||
if modality == Modality.AUDIO and isinstance(mm_feature, list):
|
||||
for feature in mm_feature:
|
||||
@@ -884,21 +910,37 @@ class MMEncoder:
|
||||
offset += num_patches
|
||||
return hashes
|
||||
|
||||
def _encode_missing(
|
||||
def _build_mm_data_items(
|
||||
self,
|
||||
mm_feature,
|
||||
mm_inputs: dict,
|
||||
indices: List[int],
|
||||
modality: Modality = Modality.IMAGE,
|
||||
get_feature_fn=None,
|
||||
modality: Modality,
|
||||
grid_thw: Optional[List] = None,
|
||||
keep_on_gpu: bool = False,
|
||||
) -> List[torch.Tensor]:
|
||||
"""
|
||||
GPU Task: Run ViT inference ONLY on the subset of mm items missing from the cache.
|
||||
) -> List[MultimodalDataItem]:
|
||||
"""Build the model-facing items selected for one encoder forward.
|
||||
|
||||
A model preprocessor can preserve an item-wise representation with
|
||||
``EncoderPreprocessOutput``. This path avoids concatenating and then
|
||||
re-slicing features before encoder-DP knows which rank owns each item.
|
||||
Legacy Hugging Face processor outputs retain their existing aggregate
|
||||
tensor behavior.
|
||||
"""
|
||||
if grid_thw is None:
|
||||
grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
|
||||
|
||||
preprocessed_items = get_encoder_preprocessed_items(mm_inputs)
|
||||
if preprocessed_items is not None:
|
||||
if len(preprocessed_items) != len(grid_thw):
|
||||
raise ValueError(
|
||||
"Encoder preprocess item/grid mismatch: "
|
||||
f"{len(preprocessed_items)} items != {len(grid_thw)} grids"
|
||||
)
|
||||
selected = [preprocessed_items[index] for index in indices]
|
||||
if any(item.modality != modality for item in selected):
|
||||
raise ValueError("Encoder preprocess output contains wrong modality")
|
||||
return selected
|
||||
|
||||
split_kimi_k3_images = (
|
||||
self.model_type == "kimi_k3" and modality == Modality.IMAGE
|
||||
)
|
||||
@@ -916,11 +958,11 @@ class MMEncoder:
|
||||
sub_feature_list = []
|
||||
offsets = [0]
|
||||
curr = 0
|
||||
for g in grid_thw:
|
||||
curr += self.get_num_patches(g, modality)
|
||||
for grid in grid_thw:
|
||||
curr += self.get_num_patches(grid, modality)
|
||||
offsets.append(curr)
|
||||
for idx in indices:
|
||||
sub_feature_list.append(mm_feature[offsets[idx] : offsets[idx + 1]])
|
||||
for index in indices:
|
||||
sub_feature_list.append(mm_feature[offsets[index] : offsets[index + 1]])
|
||||
if not split_kimi_k3_images:
|
||||
sub_feature = torch.cat(sub_feature_list, dim=0)
|
||||
|
||||
@@ -948,19 +990,39 @@ class MMEncoder:
|
||||
)
|
||||
]
|
||||
|
||||
for k, v in mm_inputs.items():
|
||||
if k in _mm_feature_attrs.get(modality, []):
|
||||
for key, value in mm_inputs.items():
|
||||
if key in _mm_feature_attrs.get(modality, []):
|
||||
continue
|
||||
val = _convert(v)
|
||||
if k in _mm_grid_attrs.get(modality, []):
|
||||
value = _convert(value)
|
||||
if key in _mm_grid_attrs.get(modality, []):
|
||||
if split_kimi_k3_images:
|
||||
for mm_item, idx in zip(mm_items, indices):
|
||||
mm_item.set(k, val[idx : idx + 1])
|
||||
for mm_item, index in zip(mm_items, indices):
|
||||
mm_item.set(key, value[index : index + 1])
|
||||
else:
|
||||
mm_items[0].set(k, val[indices])
|
||||
mm_items[0].set(key, value[indices])
|
||||
else:
|
||||
for mm_item in mm_items:
|
||||
mm_item.set(k, val)
|
||||
mm_item.set(key, value)
|
||||
return mm_items
|
||||
|
||||
def _encode_missing(
|
||||
self,
|
||||
mm_feature,
|
||||
mm_inputs: dict,
|
||||
indices: List[int],
|
||||
modality: Modality = Modality.IMAGE,
|
||||
get_feature_fn=None,
|
||||
grid_thw: Optional[List] = None,
|
||||
keep_on_gpu: bool = False,
|
||||
) -> List[torch.Tensor]:
|
||||
"""
|
||||
GPU Task: Run ViT inference ONLY on the subset of mm items missing from the cache.
|
||||
"""
|
||||
if grid_thw is None:
|
||||
grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
|
||||
mm_items = self._build_mm_data_items(
|
||||
mm_feature, mm_inputs, indices, modality, grid_thw
|
||||
)
|
||||
|
||||
forward_start = time.perf_counter()
|
||||
with torch.inference_mode():
|
||||
@@ -1002,7 +1064,7 @@ class MMEncoder:
|
||||
if self.rank == 0:
|
||||
if hashes is None:
|
||||
mm_hashes = self._calculate_hashes_from_features(
|
||||
mm_feature, grid_thw, modality
|
||||
mm_feature, grid_thw, modality, mm_inputs
|
||||
)
|
||||
else:
|
||||
mm_hashes = hashes
|
||||
@@ -1651,12 +1713,31 @@ class MMEncoder:
|
||||
if not (self.image_processor or model_preprocessor):
|
||||
raise ValueError("No image processor available")
|
||||
images = await self._flatten_and_load_images(mm_items)
|
||||
if model_preprocessor:
|
||||
return model_preprocessor(images, Modality.IMAGE, self.vision_config)
|
||||
image_config = self.vision_config.get("image", {})
|
||||
original_image_sizes = [_get_original_image_size(item) for item in images]
|
||||
if self.model_type in ["kimi_k25", "kimi_k3", "kimi_vl"]:
|
||||
images = self._normalize_kimi_encoder_images(images)
|
||||
original_image_sizes = [_get_original_image_size(item) for item in images]
|
||||
if model_preprocessor:
|
||||
processor_output = invoke_encoder_preprocessor(
|
||||
model_preprocessor,
|
||||
images,
|
||||
Modality.IMAGE,
|
||||
self.vision_config,
|
||||
image_processor=self.image_processor,
|
||||
use_gpu_preprocessing=self.use_image_processor_gpu,
|
||||
)
|
||||
if (
|
||||
isinstance(processor_output, EncoderPreprocessOutput)
|
||||
and processor_output.materialize_local_items is not None
|
||||
):
|
||||
parallel = get_parallel()
|
||||
await asyncio.get_running_loop().run_in_executor(
|
||||
self.preproc_executor,
|
||||
processor_output.materialize_for_rank,
|
||||
parallel.attn_tp_rank,
|
||||
parallel.attn_tp_size,
|
||||
)
|
||||
return processor_output
|
||||
image_config = self.vision_config.get("image", {})
|
||||
processor_input = await asyncio.get_running_loop().run_in_executor(
|
||||
self.preproc_executor,
|
||||
functools.partial(self.image_processor, images=images, **image_config),
|
||||
@@ -1775,25 +1856,25 @@ class MMEncoder:
|
||||
# support mm_cache
|
||||
mm_embedding = None
|
||||
mm_hash = None
|
||||
|
||||
mm_item = MultimodalDataItem.from_dict(
|
||||
{
|
||||
"modality": modality,
|
||||
"feature": _convert(_get_mm_feature(mm_inputs, modality)),
|
||||
}
|
||||
mm_feature = _convert(_get_mm_feature(mm_inputs, modality))
|
||||
grid_thw = _get_mm_grid_dim(mm_inputs, modality, self.model_type)
|
||||
model_mm_items = self._build_mm_data_items(
|
||||
mm_feature,
|
||||
mm_inputs,
|
||||
list(range(len(grid_thw))),
|
||||
modality,
|
||||
grid_thw,
|
||||
)
|
||||
for k, v in mm_inputs.items():
|
||||
if k in _mm_feature_attrs[modality]:
|
||||
continue
|
||||
mm_item.set(k, _convert(v))
|
||||
|
||||
cache_hit = False
|
||||
use_mm_cache = get_mm().enable_prefix_mm_cache and log_metrics
|
||||
if use_mm_cache:
|
||||
mm_item.set_pad_value()
|
||||
mm_hash = MultiModalStaticCache.combine_hashes([mm_item.hash])
|
||||
for item in model_mm_items:
|
||||
item.set_pad_value()
|
||||
item_hashes = [item.hash for item in model_mm_items]
|
||||
mm_hash = MultiModalStaticCache.combine_hashes(item_hashes)
|
||||
async with self.mm_cache_lock:
|
||||
mm_cache = self.mm_cache.get([mm_item.hash])
|
||||
mm_cache = self.mm_cache.get(item_hashes)
|
||||
if mm_cache is not None:
|
||||
mm_embedding = mm_cache.embedding
|
||||
cache_hit = True
|
||||
@@ -1801,7 +1882,7 @@ class MMEncoder:
|
||||
if mm_embedding is None:
|
||||
forward_start = time.perf_counter()
|
||||
with torch.inference_mode():
|
||||
mm_embedding: torch.Tensor = get_feature_fn([mm_item])
|
||||
mm_embedding: torch.Tensor = get_feature_fn(model_mm_items)
|
||||
mm_embedding = mm_embedding.cpu()
|
||||
if len(mm_embedding.shape) != 2:
|
||||
mm_embedding = mm_embedding.reshape(-1, mm_embedding.shape[-1])
|
||||
@@ -1810,7 +1891,8 @@ class MMEncoder:
|
||||
time.perf_counter() - forward_start, modality=modality_str
|
||||
)
|
||||
|
||||
# Per-request cache hit metrics: tokens = embedding rows, files = 1 item.
|
||||
# Per-request cache hit metrics: tokens = embedding rows, files =
|
||||
# logical multimodal items (not the legacy aggregate tensor count).
|
||||
if use_mm_cache and encoder_metrics_collector is not None:
|
||||
total_tokens = int(mm_embedding.shape[0])
|
||||
hit_tokens = total_tokens if cache_hit else 0
|
||||
@@ -1818,7 +1900,9 @@ class MMEncoder:
|
||||
hit_tokens, total_tokens, modality=modality_str
|
||||
)
|
||||
encoder_metrics_collector.record_cache_files(
|
||||
1 if cache_hit else 0, 1, modality=modality_str
|
||||
len(model_mm_items) if cache_hit else 0,
|
||||
len(model_mm_items),
|
||||
modality=modality_str,
|
||||
)
|
||||
|
||||
if use_mm_cache:
|
||||
@@ -1865,7 +1949,7 @@ class MMEncoder:
|
||||
)
|
||||
|
||||
return (
|
||||
_get_mm_grid_dim(mm_inputs, modality, self.model_type),
|
||||
grid_thw,
|
||||
mm_embedding,
|
||||
aux_data,
|
||||
)
|
||||
@@ -2223,23 +2307,21 @@ class MMEncoder:
|
||||
)
|
||||
)
|
||||
|
||||
# Build mm_item (all ranks)
|
||||
mm_item = MultimodalDataItem.from_dict(
|
||||
{
|
||||
"modality": modality,
|
||||
"feature": _convert(_get_mm_feature(mm_inputs, modality)),
|
||||
}
|
||||
# Build model-facing items on all ranks. Owner-deferred processor
|
||||
# outputs stay per-item until encoder-DP assigns them.
|
||||
mm_feature = _convert(_get_mm_feature(mm_inputs, modality))
|
||||
model_mm_items = self._build_mm_data_items(
|
||||
mm_feature,
|
||||
mm_inputs,
|
||||
list(range(len(grid_thw))),
|
||||
modality,
|
||||
grid_thw,
|
||||
)
|
||||
for k, v in mm_inputs.items():
|
||||
if k in _mm_feature_attrs.get(modality, []):
|
||||
continue
|
||||
val = _convert(v)
|
||||
mm_item.set(k, val)
|
||||
|
||||
async def _run_forward():
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
emb = get_feature_fn([mm_item])
|
||||
emb = get_feature_fn(model_mm_items)
|
||||
if len(emb.shape) != 2:
|
||||
emb = emb.reshape(-1, emb.shape[-1])
|
||||
# mooncake's transfer_sync is a host-side
|
||||
|
||||
@@ -3064,6 +3064,32 @@ class KimiK3ForConditionalGeneration(nn.Module):
|
||||
)
|
||||
self.language_model.set_dspark_layers_to_capture(layer_ids)
|
||||
|
||||
def preprocess_mm_for_encoder(
|
||||
self,
|
||||
mm_data,
|
||||
modality,
|
||||
config,
|
||||
*,
|
||||
image_processor=None,
|
||||
use_gpu_preprocessing=False,
|
||||
):
|
||||
"""Prepare per-image raw inputs for owner-side EPD preprocessing."""
|
||||
if modality != Modality.IMAGE:
|
||||
raise ValueError("Kimi-K3 encoder mode supports image input only")
|
||||
if image_processor is None:
|
||||
raise ValueError("Kimi-K3 encoder preprocessing needs an image processor")
|
||||
|
||||
from sglang.srt.multimodal.kimi_k3_image_processing import (
|
||||
prepare_kimi_k3_encoder_inputs,
|
||||
)
|
||||
|
||||
self._encoder_image_processor = image_processor
|
||||
return prepare_kimi_k3_encoder_inputs(
|
||||
mm_data,
|
||||
image_processor,
|
||||
use_gpu_preprocessing=use_gpu_preprocessing,
|
||||
)
|
||||
|
||||
def get_image_feature(self, items: List[MultimodalDataItem]) -> torch.Tensor:
|
||||
device = self.vision_tower.device
|
||||
target_dtype = self.vision_tower.patch_embed.proj.weight.dtype
|
||||
@@ -3087,6 +3113,10 @@ class KimiK3ForConditionalGeneration(nn.Module):
|
||||
|
||||
def materialize_item_features(image_indices: List[int]) -> torch.Tensor:
|
||||
"""Materialize only the images assigned to this vision-DP rank."""
|
||||
from sglang.srt.multimodal.encoder_preprocessing import (
|
||||
LOCAL_PREPROCESSED_KEY,
|
||||
)
|
||||
|
||||
# Match the configured TP consumer count captured when the
|
||||
# tokenizer creates MmItemMemoryPool. A live attention subgroup
|
||||
# size could leave acknowledgements missing and strand the lease.
|
||||
@@ -3104,6 +3134,21 @@ class KimiK3ForConditionalGeneration(nn.Module):
|
||||
)
|
||||
selected_items.append(item)
|
||||
|
||||
locally_preprocessed = [
|
||||
item.model_specific_data.get(LOCAL_PREPROCESSED_KEY, False)
|
||||
for item in selected_items
|
||||
]
|
||||
if any(locally_preprocessed):
|
||||
if not all(locally_preprocessed):
|
||||
raise ValueError(
|
||||
"Kimi-K3 cannot mix local preprocessed and deferred images"
|
||||
)
|
||||
return materialize_multimodal_features(
|
||||
[item.feature for item in selected_items],
|
||||
device=device,
|
||||
dtype=target_dtype,
|
||||
)
|
||||
|
||||
deferred = [
|
||||
item.model_specific_data.get(DEFERRED_PREPROCESSING_KEY)
|
||||
for item in selected_items
|
||||
@@ -3113,25 +3158,44 @@ class KimiK3ForConditionalGeneration(nn.Module):
|
||||
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"]
|
||||
),
|
||||
)
|
||||
backend = first_config["backend"]
|
||||
if any(config["backend"] != backend for config in deferred):
|
||||
raise ValueError(
|
||||
"Kimi-K3 cannot mix deferred preprocessing backends"
|
||||
)
|
||||
if backend == "gpu":
|
||||
from sglang.srt.multimodal.processors.kimi_k25 import (
|
||||
_gpu_preprocess_images,
|
||||
)
|
||||
|
||||
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"]
|
||||
),
|
||||
)
|
||||
elif backend == "cpu":
|
||||
from sglang.srt.multimodal.kimi_k3_image_processing import (
|
||||
materialize_kimi_k3_cpu_features,
|
||||
)
|
||||
|
||||
pixel_values = materialize_kimi_k3_cpu_features(
|
||||
selected_items, self._encoder_image_processor
|
||||
)
|
||||
pixel_values = pixel_values.to(device, non_blocking=True)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported Kimi-K3 deferred preprocessing backend: {backend}"
|
||||
)
|
||||
return pixel_values.to(dtype=target_dtype)
|
||||
|
||||
features = []
|
||||
|
||||
@@ -0,0 +1,140 @@
|
||||
import hashlib
|
||||
import inspect
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Callable, Iterable, Sequence
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.managers.schedule_batch import MultimodalDataItem
|
||||
|
||||
LOCAL_PREPROCESSED_KEY = "encoder_local_preprocessed"
|
||||
|
||||
|
||||
def hash_raw_encoder_item(value: Any) -> int:
|
||||
"""Hash raw CPU media including layout metadata, before owner materialization."""
|
||||
if isinstance(value, torch.Tensor):
|
||||
value = value.detach().cpu().contiguous().numpy()
|
||||
elif not isinstance(value, np.ndarray):
|
||||
from PIL import Image
|
||||
|
||||
if not isinstance(value, Image.Image):
|
||||
raise TypeError(f"Unsupported raw encoder item: {type(value)}")
|
||||
value = np.asarray(value)
|
||||
|
||||
value = np.ascontiguousarray(value)
|
||||
hasher = hashlib.sha256()
|
||||
hasher.update(value.dtype.str.encode())
|
||||
hasher.update(repr(value.shape).encode())
|
||||
hasher.update(memoryview(value))
|
||||
return int.from_bytes(hasher.digest()[:8], byteorder="big", signed=False)
|
||||
|
||||
|
||||
class EncoderPreprocessOutput(dict):
|
||||
"""Processor output that preserves one feature object per multimodal item.
|
||||
|
||||
Most Hugging Face processors concatenate every image into one tensor before
|
||||
the vision model decides its encoder-DP assignment. A model-specific
|
||||
``preprocess_mm_for_encoder`` hook can return this mapping instead, allowing
|
||||
the encoder to carry raw or partially processed items to the model. The
|
||||
model can then materialize only the items owned by its local vision rank.
|
||||
|
||||
The mapping remains compatible with existing encoder metadata helpers;
|
||||
``mm_items`` is an out-of-band, per-item representation used only by the
|
||||
encoder forward/cache paths.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
values: Mapping[str, Any] | None = None,
|
||||
*,
|
||||
mm_items: Iterable[MultimodalDataItem],
|
||||
item_sizes: Sequence[int] | None = None,
|
||||
materialize_local_items: (
|
||||
Callable[[list[MultimodalDataItem]], Sequence[torch.Tensor]] | None
|
||||
) = None,
|
||||
) -> None:
|
||||
super().__init__(values or {})
|
||||
self.mm_items = list(mm_items)
|
||||
if not self.mm_items:
|
||||
raise ValueError("EncoderPreprocessOutput requires at least one item")
|
||||
self.item_sizes = list(item_sizes) if item_sizes is not None else None
|
||||
if self.item_sizes is not None and len(self.item_sizes) != len(self.mm_items):
|
||||
raise ValueError("Encoder preprocess item_sizes must match mm_items")
|
||||
self.materialize_local_items = materialize_local_items
|
||||
|
||||
def local_item_indices(self, rank: int, world_size: int) -> list[int]:
|
||||
"""Return the same size-balanced owner assignment used by vision DP."""
|
||||
if self.materialize_local_items is None:
|
||||
return []
|
||||
if world_size < 1 or not 0 <= rank < world_size:
|
||||
raise ValueError(
|
||||
f"Invalid encoder preprocess rank {rank} for world size {world_size}"
|
||||
)
|
||||
if world_size == 1:
|
||||
return list(range(len(self.mm_items)))
|
||||
if self.item_sizes is None:
|
||||
raise ValueError(
|
||||
"Owner-side encoder preprocessing requires per-item load sizes"
|
||||
)
|
||||
|
||||
from sglang.srt.multimodal.mm_utils import get_dp_encoder_lb_assignment
|
||||
|
||||
shuffled, counts, _ = get_dp_encoder_lb_assignment(self.item_sizes, world_size)
|
||||
start = sum(counts[:rank])
|
||||
return shuffled[start : start + counts[rank]]
|
||||
|
||||
def materialize_for_rank(self, rank: int, world_size: int) -> None:
|
||||
"""Materialize only this vision-DP rank's items in-place."""
|
||||
indices = self.local_item_indices(rank, world_size)
|
||||
if not indices:
|
||||
return
|
||||
items = [self.mm_items[index] for index in indices]
|
||||
materialize = self.materialize_local_items
|
||||
assert materialize is not None
|
||||
features = list(materialize(items))
|
||||
if len(features) != len(items):
|
||||
raise ValueError(
|
||||
"Encoder local materializer must return one feature per item"
|
||||
)
|
||||
for item, feature in zip(items, features):
|
||||
item.feature = feature
|
||||
item.model_specific_data[LOCAL_PREPROCESSED_KEY] = True
|
||||
|
||||
|
||||
def get_encoder_preprocessed_items(
|
||||
processor_output: Mapping[str, Any],
|
||||
) -> list[MultimodalDataItem] | None:
|
||||
if isinstance(processor_output, EncoderPreprocessOutput):
|
||||
return processor_output.mm_items
|
||||
return None
|
||||
|
||||
|
||||
def invoke_encoder_preprocessor(
|
||||
preprocessor,
|
||||
mm_data,
|
||||
modality,
|
||||
config,
|
||||
**available_context,
|
||||
):
|
||||
"""Call a model hook with only the optional context it declares.
|
||||
|
||||
Existing hooks keep their three-argument contract. New model integrations
|
||||
can request shared processor state or backend policy by adding named
|
||||
keyword-only parameters, without requiring encode-server model branches.
|
||||
"""
|
||||
parameters = inspect.signature(preprocessor).parameters
|
||||
accepts_kwargs = any(
|
||||
parameter.kind == inspect.Parameter.VAR_KEYWORD
|
||||
for parameter in parameters.values()
|
||||
)
|
||||
context = (
|
||||
available_context
|
||||
if accepts_kwargs
|
||||
else {
|
||||
name: value
|
||||
for name, value in available_context.items()
|
||||
if name in parameters
|
||||
}
|
||||
)
|
||||
return preprocessor(mm_data, modality, config, **context)
|
||||
@@ -1,3 +1,5 @@
|
||||
import functools
|
||||
import math
|
||||
from typing import Union
|
||||
|
||||
import numpy as np
|
||||
@@ -7,6 +9,195 @@ from PIL import Image
|
||||
DEFERRED_PREPROCESSING_KEY = "kimi_k3_deferred_preprocessing"
|
||||
|
||||
|
||||
def prepare_kimi_k3_encoder_inputs(
|
||||
images, image_processor, *, use_gpu_preprocessing=False
|
||||
):
|
||||
"""Keep K3 EPD images raw until the vision-DP owner is known.
|
||||
|
||||
The lightweight NaViT shape calculation runs on every encoder rank. The
|
||||
expensive resize, normalization, and patchification remain deferred in each
|
||||
item and are executed by ``KimiK3ForConditionalGeneration`` only for images
|
||||
assigned to the local vision rank.
|
||||
"""
|
||||
from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
|
||||
from sglang.srt.multimodal.encoder_preprocessing import (
|
||||
EncoderPreprocessOutput,
|
||||
hash_raw_encoder_item,
|
||||
)
|
||||
from sglang.srt.multimodal.processors.kimi_k25 import (
|
||||
_grid_thw_from_resize_config,
|
||||
navit_resize_config,
|
||||
)
|
||||
|
||||
media_proc_cfg = getattr(image_processor, "media_proc_cfg", None)
|
||||
if not isinstance(media_proc_cfg, dict):
|
||||
raise ValueError(
|
||||
"Kimi-K3 EPD owner-side preprocessing requires "
|
||||
"image_processor.media_proc_cfg"
|
||||
)
|
||||
|
||||
required = (
|
||||
"patch_size",
|
||||
"merge_kernel_size",
|
||||
"in_patch_limit",
|
||||
"patch_limit_on_one_side",
|
||||
"image_mean",
|
||||
"image_std",
|
||||
)
|
||||
missing = [name for name in required if name not in media_proc_cfg]
|
||||
if missing:
|
||||
raise ValueError(
|
||||
"Kimi-K3 image processor is missing deferred-preprocessing config: "
|
||||
+ ", ".join(missing)
|
||||
)
|
||||
|
||||
concrete_images = []
|
||||
for image in images:
|
||||
if isinstance(image, dict):
|
||||
if image.get("type") != "image" or "image" not in image:
|
||||
raise ValueError(f"Unsupported Kimi-K3 encoder media item: {image}")
|
||||
image = image["image"]
|
||||
concrete_images.append(image)
|
||||
|
||||
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"),
|
||||
}
|
||||
|
||||
items = []
|
||||
grids = []
|
||||
original_image_sizes = []
|
||||
for image in concrete_images:
|
||||
width, height = (
|
||||
(int(image.shape[-1]), int(image.shape[-2]))
|
||||
if isinstance(image, torch.Tensor)
|
||||
else image.size
|
||||
)
|
||||
resize_config = navit_resize_config(
|
||||
width,
|
||||
height,
|
||||
patch_size,
|
||||
merge_kernel_size,
|
||||
int(media_proc_cfg["in_patch_limit"]),
|
||||
int(media_proc_cfg["patch_limit_on_one_side"]),
|
||||
media_proc_cfg.get("fixed_output_tokens"),
|
||||
)
|
||||
grid_thw = _grid_thw_from_resize_config(resize_config, patch_size)
|
||||
grid_tensor = torch.tensor([grid_thw], dtype=torch.int64)
|
||||
item = MultimodalDataItem(
|
||||
modality=Modality.IMAGE,
|
||||
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,
|
||||
},
|
||||
},
|
||||
)
|
||||
if not use_gpu_preprocessing:
|
||||
item.set_hash(hash_raw_encoder_item(image))
|
||||
items.append(item)
|
||||
grids.append(grid_thw)
|
||||
original_image_sizes.append([width, height])
|
||||
|
||||
grid_thws = torch.tensor(grids, dtype=torch.int64)
|
||||
return EncoderPreprocessOutput(
|
||||
{
|
||||
# Preserve the conventional feature key for cache/accounting code;
|
||||
# EncoderPreprocessOutput.mm_items is the authoritative per-item view.
|
||||
"pixel_values": [item.feature for item in items],
|
||||
"grid_thws": grid_thws,
|
||||
"original_image_sizes": original_image_sizes,
|
||||
},
|
||||
mm_items=items,
|
||||
item_sizes=[math.prod(grid) for grid in grids],
|
||||
materialize_local_items=(
|
||||
None
|
||||
if use_gpu_preprocessing
|
||||
else functools.partial(
|
||||
materialize_kimi_k3_cpu_item_features,
|
||||
image_processor=image_processor,
|
||||
)
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def materialize_kimi_k3_cpu_features(items, image_processor) -> torch.Tensor:
|
||||
"""Run the checkpoint's exact processor only on locally owned images."""
|
||||
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(
|
||||
"Kimi-K3 deferred CPU preprocessing expects PIL or uint8 tensors"
|
||||
)
|
||||
image = to_hwc_uint8(image).numpy()
|
||||
channels = image.shape[-1]
|
||||
if channels == 1:
|
||||
image = Image.fromarray(image[..., 0], mode="L")
|
||||
elif channels == 3:
|
||||
image = Image.fromarray(image, mode="RGB")
|
||||
elif channels == 4:
|
||||
image = Image.fromarray(image, mode="RGBA")
|
||||
else:
|
||||
raise ValueError(f"Unsupported Kimi-K3 image channel count: {channels}")
|
||||
medias.append({"type": "image", "image": image})
|
||||
|
||||
output = image_processor.preprocess(medias, return_tensors="pt")
|
||||
expected_grids = torch.cat(
|
||||
[item.model_specific_data["grid_thws"] for item in items], dim=0
|
||||
)
|
||||
if not torch.equal(output["grid_thws"].cpu(), expected_grids.cpu()):
|
||||
raise ValueError("Kimi-K3 deferred CPU preprocessing produced wrong grids")
|
||||
return output["pixel_values"]
|
||||
|
||||
|
||||
def materialize_kimi_k3_cpu_item_features(items, image_processor) -> list[torch.Tensor]:
|
||||
"""Return exact checkpoint-processor features split by logical image."""
|
||||
pixel_values = materialize_kimi_k3_cpu_features(items, image_processor)
|
||||
patch_counts = [
|
||||
math.prod(item.model_specific_data["grid_thws"][0].tolist()) for item in items
|
||||
]
|
||||
if sum(patch_counts) != pixel_values.shape[0]:
|
||||
raise ValueError(
|
||||
"Kimi-K3 processor feature length does not match image grids: "
|
||||
f"{pixel_values.shape[0]} != {sum(patch_counts)}"
|
||||
)
|
||||
return list(pixel_values.split(patch_counts))
|
||||
|
||||
|
||||
def to_hwc_uint8(image: Union[torch.Tensor, Image.Image]) -> torch.Tensor:
|
||||
"""Stage an exact CPU image without doing resize/normalize work."""
|
||||
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)
|
||||
return torch.from_numpy(array)
|
||||
|
||||
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(-1)
|
||||
elif image.dim() == 3 and image.shape[0] in (1, 3, 4):
|
||||
image = image.permute(1, 2, 0)
|
||||
if image.shape[-1] == 1:
|
||||
image = image.repeat(1, 1, 3)
|
||||
return image.cpu().contiguous()
|
||||
|
||||
|
||||
def to_chw_uint8(
|
||||
image: Union[torch.Tensor, Image.Image],
|
||||
device: torch.device | str | None = None,
|
||||
|
||||
Reference in New Issue
Block a user