From 773faf992d1160c69b234dce0dd56e36c74f1dde Mon Sep 17 00:00:00 2001 From: Oguz Ulgen Date: Wed, 12 Aug 2026 11:04:07 -0700 Subject: [PATCH] Reserve multimodal runtime allocations and keep padded inputs aligned (#34141) Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com> Co-authored-by: Lianmin Zheng Co-authored-by: wangwenchen0407 Co-authored-by: Hanming Lu --- .../srt/mem_cache/kv_cache_configurator.py | 26 ++++++++++++++++++- .../srt/model_executor/forward_batch_info.py | 3 +++ .../kv_pool_runtime.py | 7 ++++- .../vlm/test_vision_openai_server_a.py | 6 +++++ 4 files changed, 40 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 8aa875482..4e3e4df65 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -114,6 +114,26 @@ def _should_enable_lazy_compaction() -> bool: return not envs.SGLANG_DISABLE_LAZY_COMPACTION.get() +def mm_runtime_reservation_gb( + *, is_multimodal: bool, mm_feature_transport: Optional[str] +) -> float: + """Multimodal GPU memory allocated only after the KV pool is sized + (mm embedding cache + GPU feature-transport pools); reserve it out of + the KV budget so it doesn't have to fit in the runtime slack.""" + if not is_multimodal: + return 0.0 + reserved_mb = envs.SGLANG_VLM_CACHE_SIZE_MB.get() + if mm_feature_transport in ("cuda_ipc", "cuda_vmm"): + reserved_mb += envs.SGLANG_MM_FEATURE_CACHE_MB.get() + if reserved_mb > 0: + logger.info( + "Reserving %.2f GB of the KV budget for post-sizing multimodal " + "allocations (feature-transport pools + embedding cache).", + reserved_mb / 1024, + ) + return reserved_mb / 1024 + + # base ratio of mamba pool size to max_running_requests. Under # SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK the decode-time skip frees one resident slot # per running request, so the base drops by 1 (overlap 5->4, lazy 4->3). no_buffer @@ -1732,7 +1752,11 @@ class KVCacheConfigurator: ) / 1024, ) - rest_memory = available_gpu_memory - slack_gb + mm_reservation_gb = mm_runtime_reservation_gb( + is_multimodal=self.model_config.is_multimodal, + mm_feature_transport=self.server_args.mm_feature_transport, + ) + rest_memory = available_gpu_memory - slack_gb - mm_reservation_gb if self.mambaish_config is not None: rest_memory = self._handle_max_mamba_cache(rest_memory) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 2d01ef258..d71be5ae9 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -1439,6 +1439,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # padding self._original_num_tokens = self.positions.shape[0] self.input_ids = self._pad_tensor_to_size(self.input_ids, num_tokens) + if self.input_embeds is not None: + # Keep token-aligned inputs consistent after padding. + self.input_embeds = self._pad_tensor_to_size(self.input_embeds, num_tokens) self.req_pool_indices = self._pad_tensor_to_size(self.req_pool_indices, bs) if self.lora_ids is not None: self.lora_ids.extend((bs - len(self.lora_ids)) * [None]) diff --git a/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py b/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py index c0a16bc56..58c315149 100644 --- a/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py +++ b/python/sglang/srt/model_executor/model_runner_components/kv_pool_runtime.py @@ -8,6 +8,7 @@ import torch from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.distributed import get_world_group +from sglang.srt.mem_cache.kv_cache_configurator import mm_runtime_reservation_gb from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.platforms import current_platform from sglang.srt.utils.common import get_available_gpu_memory, get_device_memory_capacity @@ -76,8 +77,12 @@ def compute_post_capture_kv_resize( ) / 1024, ) + mm_reservation_gb = mm_runtime_reservation_gb( + is_multimodal=model_runner.model_config.is_multimodal, + mm_feature_transport=model_runner.server_args.mm_feature_transport, + ) budget_bytes = ( - int(max(0.0, free_gb - headroom_gb) * (1 << 30)) + int(max(0.0, free_gb - headroom_gb - mm_reservation_gb) * (1 << 30)) + pool.post_capture_backed_bytes ) config = model_runner.kv_cache_configurator.config_from_budget( diff --git a/test/registered/vlm/test_vision_openai_server_a.py b/test/registered/vlm/test_vision_openai_server_a.py index 3624e35b0..e96250395 100644 --- a/test/registered/vlm/test_vision_openai_server_a.py +++ b/test/registered/vlm/test_vision_openai_server_a.py @@ -8,6 +8,7 @@ import unittest import openai +from sglang.srt.environ import envs from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -46,6 +47,11 @@ class TestQwen3VLServer(ImageOpenAITestMixin, VideoOpenAITestMixin): model = "Qwen/Qwen3-VL-30B-A3B-Instruct" extra_args = ["--cuda-graph-max-bs-decode=4"] + @classmethod + def setUpClass(cls): + with envs.SGLANG_MM_FEATURE_CACHE_MB.override(512): + super().setUpClass() + class TestQwen2VLContextLengthServer(CustomTestCase): # --context-length 300 is calibrated to this model's mm-token expansion: