[EPD][Perf] Async image preprocessing and cross-request ViT batching for encode_server (#25669)

This commit is contained in:
giang_ng_tr
2026-06-01 16:52:12 +08:00
committed by GitHub
parent bc36231d65
commit 2394dede0e
2 changed files with 23 additions and 3 deletions
@@ -2,6 +2,7 @@ import asyncio
import concurrent.futures
import contextlib
import ctypes
import functools
import logging
import multiprocessing as mp
import os
@@ -278,6 +279,11 @@ class MMEncoder:
self.context = zmq.asyncio.Context(2)
self.sync_context = zmq.Context() # Reuse sync context for thread pool
self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=10)
# Dedicated executor for image preprocessing (resize/normalize).
# Separate from self.executor (ZMQ sends) to avoid contention under high concurrency.
self.preproc_executor = concurrent.futures.ThreadPoolExecutor(
max_workers=envs.SGLANG_ENCODER_PREPROC_WORKERS.get()
)
embedding_cache_size = int(os.environ.get("SGLANG_VLM_CACHE_SIZE_MB", "4096"))
self.mm_cache = MultiModalStaticCache(embedding_cache_size * 1024 * 1024)
@@ -1358,7 +1364,10 @@ class MMEncoder:
image_config = self.vision_config.get("image", {})
if self.model_type in ["kimi_k25", "kimi_vl"]:
images = self._normalize_kimi_encoder_images(images)
return self.image_processor(images=images, **image_config)
return await asyncio.get_running_loop().run_in_executor(
self.preproc_executor,
functools.partial(self.image_processor, images=images, **image_config),
)
async def _process_video_items(self, mm_items, model_preprocessor):
if model_preprocessor:
@@ -1367,7 +1376,12 @@ class MMEncoder:
raise ValueError("No video processor available")
videos, video_processor_kwargs = await self._flatten_and_load_videos(mm_items)
processor_input = self.video_processor(videos=videos, **video_processor_kwargs)
processor_input = await asyncio.get_running_loop().run_in_executor(
self.preproc_executor,
functools.partial(
self.video_processor, videos=videos, **video_processor_kwargs
),
)
# Get additional video metadata
if (
@@ -1430,7 +1444,12 @@ class MMEncoder:
raise ValueError("No audio processor available")
audio_config = self.vision_config.get("audio", {})
processor_input = self.audio_processor.feature_extractor(audios, **audio_config)
processor_input = await asyncio.get_running_loop().run_in_executor(
self.preproc_executor,
functools.partial(
self.audio_processor.feature_extractor, audios, **audio_config
),
)
processor_input["feature_attention_mask"] = processor_input.pop(
"attention_mask"
)
+1
View File
@@ -735,6 +735,7 @@ class Envs:
SGLANG_ENCODER_DISPATCH_MIN_ITEMS = EnvInt(2)
SGLANG_ENCODER_IMAGE_PROCESSOR_USE_GPU = EnvBool(False)
SGLANG_ENCODER_MAX_BATCH_SIZE = EnvInt(8)
SGLANG_ENCODER_PREPROC_WORKERS = EnvInt(8)
# Persistent receiver-side GPU embedding pool size for mooncake EPD transport.
# 0 disables (per-request register/deregister). 4096 = 4GB default per TP
SGLANG_EMBEDDING_POOL_SIZE_MB = EnvInt(4096)