diff --git a/docs/supported_models/multimodal_language_models.md b/docs/supported_models/multimodal_language_models.md index 8aa642d53..792c67e30 100644 --- a/docs/supported_models/multimodal_language_models.md +++ b/docs/supported_models/multimodal_language_models.md @@ -58,7 +58,7 @@ SGLang supports video input for Vision-Language Models (VLMs), enabling temporal | **GLM-4v** (4.5V, 4.1V, MOE) | `zai-org/GLM-4.5V` | Video clips are read with Decord, converted to tensors, and passed to the model alongside metadata for rotary-position handling. | | **NVILA** (Full & Lite) | `Efficient-Large-Model/NVILA-8B` | The runtime samples eight frames per clip and attaches them to the multimodal request when `video_data` is present. | | **LLaVA video variants** (LLaVA-NeXT-Video, LLaVA-OneVision) | `lmms-lab/LLaVA-NeXT-Video-7B` | The processor routes video prompts to the LlavaVid video-enabled architecture, and the provided example shows how to query it with `sgl.video(...)` clips. | -| **NVIDIA Nemotron Nano 2.0 VL** | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` | For video, the processor is configured to sample at 2 FPS, at a max of 128 frames, as per model training. | +| **NVIDIA Nemotron Nano 2.0 VL** | `nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16` | The processor samples at 2 FPS, at a max of 128 frames, as per model training. The model uses [EVS](../../python/sglang/srt/multimodal/evs/README.md), a pruning method that removes redundant tokens from video embeddings. By default `video_pruning_rate=0.7`. Change this by providing: `--json-model-override-args '{"video_pruning_rate": 0.0}'` to disable EVS, for example. | | **JetVLM** | | The runtime samples eight frames per clip and attaches them to the multimodal request when `video_data` is present. | Use `sgl.video(path, num_frames)` when building prompts to attach clips from your SGLang programs. diff --git a/python/sglang/srt/disaggregation/encode_server.py b/python/sglang/srt/disaggregation/encode_server.py index aa8e43be7..bde59f96b 100644 --- a/python/sglang/srt/disaggregation/encode_server.py +++ b/python/sglang/srt/disaggregation/encode_server.py @@ -31,7 +31,7 @@ from sglang.srt.distributed.parallel_state import ( from sglang.srt.layers.dp_attention import initialize_dp_attention from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileReqType from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem -from sglang.srt.mem_cache.multimodal_cache import MultiModalStaticCache +from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.model_loader import get_model from sglang.srt.server_args import ( PortArgs, @@ -209,7 +209,7 @@ class MMEncoder: async with self.mm_cache_lock: mm_cache = self.mm_cache.get([mm_item.hash]) if mm_cache is not None: - mm_embedding = mm_cache + mm_embedding = mm_cache.embedding if mm_embedding is None: with torch.inference_mode(): @@ -220,7 +220,7 @@ class MMEncoder: if self.server_args.enable_prefix_mm_cache: async with self.mm_cache_lock: - self.mm_cache.set(mm_hash, mm_embedding) + self.mm_cache.set(mm_hash, EmbeddingResult(embedding=mm_embedding)) end_time = time.perf_counter() logger.info( f"Vit time : {(end_time - start_time)*1000:.2f} ms {mm_embedding.shape = }" diff --git a/python/sglang/srt/managers/mm_utils.py b/python/sglang/srt/managers/mm_utils.py index 6bbbf51cc..1e4d09036 100644 --- a/python/sglang/srt/managers/mm_utils.py +++ b/python/sglang/srt/managers/mm_utils.py @@ -21,8 +21,9 @@ from sglang.srt.managers.schedule_batch import ( MultimodalDataItem, MultimodalInputs, ) -from sglang.srt.mem_cache.multimodal_cache import MultiModalStaticCache +from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.multimodal.evs import EVSEmbeddingResult from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import flatten_nested_list, is_npu, print_warning_once from sglang.utils import logger @@ -478,6 +479,11 @@ def _get_precomputed_embedding( return None +DataEmbeddingFunc = Callable[ + [List[MultimodalDataItem]], torch.Tensor | EVSEmbeddingResult +] + + def get_embedding_items_per_chunk_with_extra_padding( embedding_items_per_req: List["MultimodalDataItem"], extend_prefix_len: int, @@ -540,13 +546,14 @@ def get_embedding_items_per_chunk_with_extra_padding( # TODO: To be obsoleted. def _get_chunked_prefill_embedding( - data_embedding_func: Callable[[List[MultimodalDataItem]], torch.Tensor], + data_embedding_func: DataEmbeddingFunc, embedding_items: List[MultimodalDataItem], items_size: List[int], prefix_length: List[int], extend_length: List[int], items_offset_list: List[List[Tuple[int, int]]], -) -> Optional[torch.Tensor]: + input_ids: torch.Tensor, +) -> tuple[torch.Tensor | None, torch.Tensor]: # Calculate embedding for each request, try to get it from cache to avoid repeated calculation embedding_list = [] # FIXME(Xinyuan): temporary workaround for eagle3, which may have len(items_size) > len(prefix_length) @@ -564,7 +571,12 @@ def _get_chunked_prefill_embedding( embedding_items_hash = MultiModalStaticCache.combine_hashes(item_hashes) embedding_per_req = embedding_cache.get(item_hashes) if embedding_per_req is None: - embedding_per_req = data_embedding_func(embedding_items_per_req) + embedding = data_embedding_func(embedding_items_per_req) + embedding_per_req = ( + EmbeddingResult(embedding=embedding) + if isinstance(embedding, torch.Tensor) + else embedding + ) if not embedding_cache.set(embedding_items_hash, embedding_per_req): print_warning_once( "Multimodal embedding cache is full. This typically occurs when a single " @@ -573,16 +585,31 @@ def _get_chunked_prefill_embedding( "embedding size." ) + extend_prefix_len = prefix_length[i] + extend_seq_len = extend_length[i] if i < len(extend_length) else 0 + + if isinstance(embedding_per_req, EVSEmbeddingResult): + item = embedding_items_per_req[0] + input_ids, items_offset = ( + embedding_per_req.redistribute_pruned_frames_placeholders( + input_ids, + items_offset, + item=item, + extend_prefix_len=extend_prefix_len, + extend_seq_len=extend_seq_len, + ) + ) + embedding_per_req_chunk, _, _ = get_embedding_chunk( - embedding=embedding_per_req, - extend_prefix_len=prefix_length[i], - extend_seq_len=extend_length[i] if i < len(extend_length) else 0, + embedding=embedding_per_req.embedding, + extend_prefix_len=extend_prefix_len, + extend_seq_len=extend_seq_len, items_offset=items_offset, ) embedding_list.append(embedding_per_req_chunk) if len(embedding_list) == 0: - return None - return torch.concat(embedding_list, dim=0) + return None, input_ids + return torch.concat(embedding_list, dim=0), input_ids def get_embedding_chunk_remove_extra_padding( @@ -826,7 +853,7 @@ def _adjust_embedding_length( def get_embedding_and_mask( - data_embedding_func: Callable[[List[MultimodalDataItem]], torch.Tensor], + data_embedding_func: DataEmbeddingFunc, embedding_items: List[MultimodalDataItem], placeholder_tensor: torch.Tensor, input_ids: torch.Tensor, @@ -834,7 +861,7 @@ def get_embedding_and_mask( prefix_length: List[int], extend_length: List[int], items_offset_list: List[List[Tuple[int, int]]], -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor]: """ Generate multimodal embeddings and create a mask for identifying their positions in the input sequence. @@ -852,29 +879,31 @@ def get_embedding_and_mask( A tuple containing: - The generated embeddings tensor - A boolean mask tensor indicating where these embeddings should be placed + - If EVS is used, the pruned input ids tensor; otherwise, the original input ids tensor """ # 1. Get embedding embedding = _get_precomputed_embedding( embedding_items, prefix_length, extend_length, items_offset_list ) if embedding is None: - embedding = _get_chunked_prefill_embedding( + embedding, input_ids = _get_chunked_prefill_embedding( data_embedding_func, embedding_items, items_size, prefix_length, extend_length, items_offset_list, + input_ids, ) if embedding is None: - return None, None + return None, None, input_ids # 2. Get mask if _is_npu: torch.npu.current_stream().synchronize() special_multimodal_mask = _get_multimodal_mask(input_ids, placeholder_tensor) # 3. Adjust embedding length if needed embedding = _adjust_embedding_length(embedding, special_multimodal_mask, logger) - return embedding, special_multimodal_mask + return embedding, special_multimodal_mask, input_ids def embed_mm_inputs( @@ -884,9 +913,7 @@ def embed_mm_inputs( input_ids: torch.Tensor, input_embedding: nn.Embedding, multimodal_model: nn.Module = None, - data_embedding_func_mapping: Dict[ - Modality, Callable[[List[MultimodalDataItem]], torch.Tensor] - ] = None, + data_embedding_func_mapping: Dict[Modality, DataEmbeddingFunc] = None, placeholder_tokens: dict[Modality, List[int]] = None, use_deepstack: Dict[Modality, bool] = {}, ) -> Optional[torch.Tensor]: @@ -953,7 +980,7 @@ def embed_mm_inputs( ) items_size = torch.cumsum(items_size, dim=0).tolist() - embedding, mask = get_embedding_and_mask( + embedding, mask, input_ids = get_embedding_and_mask( data_embedding_func=embedder, embedding_items=items, placeholder_tensor=placeholder_tensor, @@ -1020,9 +1047,7 @@ def general_mm_embed_routine( forward_batch: ForwardBatch, language_model: nn.Module, multimodal_model: Optional[nn.Module] = None, - data_embedding_funcs: Dict[ - Modality, Callable[[List[MultimodalDataItem]], torch.Tensor] - ] = None, + data_embedding_funcs: Dict[Modality, DataEmbeddingFunc] = None, placeholder_tokens: Optional[dict[Modality, List[int]]] = None, use_deepstack: Dict[Modality, bool] = {}, **kwargs, diff --git a/python/sglang/srt/mem_cache/multimodal_cache.py b/python/sglang/srt/mem_cache/multimodal_cache.py index 604048700..ac1cb93de 100644 --- a/python/sglang/srt/mem_cache/multimodal_cache.py +++ b/python/sglang/srt/mem_cache/multimodal_cache.py @@ -1,5 +1,6 @@ import abc from collections import OrderedDict +from dataclasses import dataclass from typing import List, Optional import torch @@ -67,6 +68,11 @@ def _get_tensor_size(embedding: torch.Tensor): return embedding.element_size() * embedding.numel() +@dataclass(kw_only=True) +class EmbeddingResult: + embedding: torch.Tensor + + class MultiModalStaticCache(MultimodalCache): """ A server-level cache for multimodal embedding. @@ -79,12 +85,12 @@ class MultiModalStaticCache(MultimodalCache): ): super().__init__() self.max_size = max_size - self.mm_cache: OrderedDict[int, torch.Tensor] = OrderedDict() + self.mm_cache: OrderedDict[int, EmbeddingResult] = OrderedDict() self.current_size = 0 def get( self, mm_hashes: List[int], combined_hash: Optional[int] = None - ) -> Optional[torch.Tensor]: + ) -> Optional[EmbeddingResult]: combined_hash = self.combine_hashes(mm_hashes) # MultiModalStaticCache does not fallback to individual item lookup @@ -94,17 +100,21 @@ class MultiModalStaticCache(MultimodalCache): return embedding def set( - self, mm_hash: int, embedding: torch.Tensor, loc: Optional[torch.Tensor] = None + self, + mm_hash: int, + embedding: EmbeddingResult, + loc: Optional[torch.Tensor] = None, ) -> bool: + assert isinstance(embedding, EmbeddingResult), embedding if mm_hash in self.mm_cache: self.mm_cache.move_to_end(mm_hash) return True - data_size = _get_tensor_size(embedding) + data_size = _get_tensor_size(embedding.embedding) while self.current_size + data_size > self.max_size: if not self.mm_cache: return False lru_hash, lru_embedding = self.mm_cache.popitem(last=False) - self.current_size -= _get_tensor_size(lru_embedding) + self.current_size -= _get_tensor_size(lru_embedding.embedding) self.mm_cache[mm_hash] = embedding self.current_size += data_size @@ -119,7 +129,7 @@ class MultiModalStaticCache(MultimodalCache): if mm_hash not in self.mm_cache: return False old_embedding = self.mm_cache.pop(mm_hash) - self.current_size -= _get_tensor_size(old_embedding) + self.current_size -= _get_tensor_size(old_embedding.embedding) return True def clear(self): diff --git a/python/sglang/srt/models/nano_nemotron_vl.py b/python/sglang/srt/models/nano_nemotron_vl.py index e33767217..cc140a333 100644 --- a/python/sglang/srt/models/nano_nemotron_vl.py +++ b/python/sglang/srt/models/nano_nemotron_vl.py @@ -36,19 +36,24 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.nemotron_h import NemotronHForCausalLM from sglang.srt.models.radio import RadioModel +from sglang.srt.multimodal.evs import EVS, EVSConfig from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) -class NemotronH_Nano_VL_V2(nn.Module): +class NemotronH_Nano_VL_V2(EVS): + @staticmethod + def create_evs_config(config: NemotronH_Nano_VL_V2_Config): + return EVSConfig(video_pruning_rate=config.video_pruning_rate) + def __init__( self, config: NemotronH_Nano_VL_V2_Config, quant_config: QuantizationConfig | None = None, prefix: str = "", ) -> None: - super().__init__() + super().__init__(config) self.downsample_ratio = config.downsample_ratio self.language_model = NemotronHForCausalLM( diff --git a/python/sglang/srt/multimodal/evs/README.md b/python/sglang/srt/multimodal/evs/README.md new file mode 100644 index 000000000..3bbf87996 --- /dev/null +++ b/python/sglang/srt/multimodal/evs/README.md @@ -0,0 +1,123 @@ +# Efficient Video Sampling (EVS) + +Implementation of [Efficient Video Sampling: Pruning Temporally Redundant Tokens for Faster VLM Inference](https://arxiv.org/abs/2510.14624). + +## Overview + +> NOTE: The current implementation in sglang is cannot work with VLMs that use positional embeddings [Such as Qwen2.5VL]. Further work is warranted. + +Video frames often contain redundant information, as consecutive frames may be nearly identical. EVS exploits this in the latent space [=embedding space] by computing similarity between adjacent frame token embeddings and pruning tokens that are highly similar to the previous frames. This reduces the token count while preserving informative content. + +Key properties: +- The first frame is always fully retained (provides complete initial context) +- Configurable via `video_pruning_rate` in model config.json (0 = disabled, 0.7 = ~70% reduction; ~30% retained.) + + +## Performance Characteristics VS. Accuracy - Example + +> NOTE: Actual retained accuracy post-EVS may depend on how dynamic the input videos are, how high the pruning rate is, whether or not the model was trained with EVS on or not, etc. +> To learn more, read the paper above. It is incumbent on the user to evaluate as per their use case and benchmarks. + +A cursory example of a performance boost due to EVS: + +```bash +export SGLANG_VLM_CACHE_SIZE_MB=0 +sglang serve --model-path nvidia/Nemotron-Nano-12B-v2-VL-BF16 --trust-remote-code --mem-fraction-static 0.8 --max-mamba-cache-size 128 --chunked-prefill-size 8192 +``` +Example Request: +```json +{ "model": "nvidia/Nemotron-Nano-12B-v2-VL-BF16", "stream": true, "temperature": 0.0, "max_completion_tokens": 3, "messages": [{ "role": "user", "content": [{ "type": "video_url", "video_url": { "url": "file:///tmp/01.mp4" } }]}]} +``` + +- `1XH100 95GiB` +- `BS=1` +- All 30 videos of `https://huggingface.co/datasets/lmms-lab/Video-MME/blob/main/videos_chunked_01.zip` +- Default [for this model] pruning rate of `--json-model-override-args '{"video_pruning_rate": 0.7}'` [i.e., 30% of tokens are preserved] VS. `--json-model-override-args '{"video_pruning_rate": 0.0}'` [EVS off] + +| Scenario\ Metric | Online TTFT (Seconds) stderr: ±0.38 | VideoMME Accuracy | +|--------------------------------- |------------------------------------- |------------------------- | +| EVS Off [q=0.0] | 11.96 [100%] | Between 0.665 and 0.668 | +| EVS Off [q=0.4] | 09.97 [ 83%] | | +| EVS On [q=0.7] (default value) | 08.79 [ 73%] | | +| EVS Off [q=0.9] | 08.39 [ 70%] | 0.644 | + + +## Architecture + +### Request Flow + +1. Prompt Construction (EVSProcessor) + * Calculates estimated tokens per frame based on pruning rate, so the emitted input_ids tensor's length will by definition match the final sequence length post pruning. This is necessary for 3. +2. Embedding Generation (EVS) + * Calls original model `get_video_feature()` for full embeddings + * Retains top-k dissimilar tokens + * Returns EVSEmbeddingResult in addition to pruned token counts *per frame* +3. Token Redistribution (mm_utils) + * Adjusts input_ids so each frame's placeholder tokens matches the pruned count from 2. + + +## Integration Guide + +### Step 1: Model [See `NemotronH_Nano_VL_V2`] + +Make your model inherit from `EVS` and implement `create_evs_config`: + +```python +from sglang.srt.multimodal.evs import EVSConfig, EVS + +class MyEVSVideoModel(EVS): + @staticmethod + def create_evs_config(config): + return EVSConfig( + video_pruning_rate=config.video_pruning_rate + ) + + def __init__(self, config, ...): + super().__init__(config) # EVS wraps get_video_feature + ... + + def get_video_feature(self, items): + # Your existing implementation + # Returns: (total_frames, tokens_per_frame, hidden_dim) + ... +``` + +### Step 2: Processor [See `NanoNemotronVLImageProcessor`] + +Create an `EVSProcessor` as a member of your VLImageProcessor: + +```python +from sglang.srt.multimodal.evs import EVSProcessor + +class MyProcessor: + models = [MyEVSVideoModel, MyNonEVSModel] # You may mix evs and non evs models in a processor + + def __init__(hf_config): + self.evs = EVSProcessor(hf_config, config_to_evs_model={MyEVSVideoModelConfig: MyEVSVideoModel}) + + def process_video(self, ...): + for video in videos: + tokens_per_frame = self.tokens_per_frame() + mm_items = create_data_items( + image=image_feature, + image_offsets=img_offsets, + video=video_feature, + video_offsets=video_offsets, + ) +``` + +### Step 3: Config [See `NemotronH_Nano_VL_V2_Config`] + +Add `video_pruning_rate` to your model config: + +```python +class MyModelConfig(PretrainedConfig): + def __init__(self, ..., video_pruning_rate=0.0, ...): + self.video_pruning_rate = video_pruning_rate +``` + +## Files + +- `evs_core.py`: Core algorithms (retention mask computation, token redistribution) +- `evs_module.py`: EVS, configs) +- `evs_processor.py`: EVSProcessor diff --git a/python/sglang/srt/multimodal/evs/__init__.py b/python/sglang/srt/multimodal/evs/__init__.py new file mode 100644 index 000000000..bb6462bd1 --- /dev/null +++ b/python/sglang/srt/multimodal/evs/__init__.py @@ -0,0 +1,11 @@ +"""https://arxiv.org/abs/2510.14624: Efficient Video Sampling: Pruning Temporally Redundant Tokens for Faster VLM Inference""" + +from .evs_module import EVS, EVSConfig, EVSEmbeddingResult +from .evs_processor import EVSProcessor + +__all__ = [ + "EVS", + "EVSConfig", + "EVSEmbeddingResult", + "EVSProcessor", +] diff --git a/python/sglang/srt/multimodal/evs/evs_core.py b/python/sglang/srt/multimodal/evs/evs_core.py new file mode 100644 index 000000000..6ab53362b --- /dev/null +++ b/python/sglang/srt/multimodal/evs/evs_core.py @@ -0,0 +1,176 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/multimodal/evs.py + +import torch + + +def compute_retained_tokens_count( + tokens_per_frame: int, num_frames: int, q: float +) -> int: + """ + Compute the number of retained tokens for a given video. + Method ensures that we retain all the tokens from the first frame + regardless of the pruning rate. + + Args: + tokens_per_frame: The number of tokens per frame. + num_frames: The total number of frames. + q: The pruning rate. + + Returns: + The number of retained tokens. + """ + total_tokens = tokens_per_frame * num_frames + evs_num_tokens = int(total_tokens * (1 - q)) + min_num_tokens = tokens_per_frame + return max(min_num_tokens, evs_num_tokens) + + +def compute_retention_mask( + video_embeds: torch.Tensor, + video_size_thw: torch.LongTensor | tuple[int, int, int], + spatial_merge_size: int, + q: float, +) -> torch.Tensor: + """ + Computes the retention mask for input video embeddings. + + Args: + video_embeds (`torch.Tensor`): The input video embeddings + of shape `(T * H * W // spatial_merge_size ^ 2, hidden_size)` + video_size_thw (`torch.LongTensor` of shape `(3)`): + The temporal, height and width of video. + spatial_merge_size: Size reduction for rows & cols dimensions. + q: (`float`): Pruning rate factor [0,1) + + Returns: + `torch.Tensor`: The retention mask for the video embeddings of + `(T * H * W // spatial_merge_size ^ 2)` shape. + """ + T, H, W = map(int, video_size_thw) + + # Use reshape instead of einops to avoid graph breaks + video_embeds = video_embeds.reshape( + T, + H // spatial_merge_size, + W // spatial_merge_size, + video_embeds.size(-1), + ) + tokens_per_frame = (H // spatial_merge_size) * (W // spatial_merge_size) + # Core EVS + similarity = torch.nn.functional.cosine_similarity( + video_embeds[1:, ...], video_embeds[:-1, ...], dim=-1 + ) + dissimilarity = 1 - similarity + + # Always ensure we include all tokens from the first frame + dissimilarity = torch.cat( + [255 * torch.ones_like(video_embeds[:1, :, :, 0]), dissimilarity], dim=0 + ) + + dissimilarity_flat = dissimilarity.view(-1) + order = torch.argsort(dissimilarity_flat, dim=-1, descending=True, stable=True) + retain_num_tokens = compute_retained_tokens_count( + tokens_per_frame=tokens_per_frame, num_frames=T, q=q + ) + topk_indices = order[:retain_num_tokens] + + retention_mask = torch.zeros_like(dissimilarity_flat, dtype=torch.bool) + retention_mask[topk_indices] = True + retention_mask = retention_mask.reshape(dissimilarity.size()) + + mask = retention_mask.view(-1) # "T H W -> (T H W)" + return mask + + +# ▲ End of VLLM code + + +def tokens_per_frame( + *, + q: float, + num_frames: int, + frame_num_tokens: int, +) -> list[int]: + """ + Before EVS pruning, we want to pre-reduce input_ids to be the same length that will be retained of embeddings due to EVS pruning, so the forward batch metadata will be correct post EVS. + We don't know the exact number of tokens per frame after EVS pruning, but we know the *total* number of tokens that will be retained. + So, we create a bogus tokens_per_frame list that sums to the total number of tokens that will be retained, and use it for placeholder spans, later to replaced, see `replace_offsets_with_tokens_per_frame` below. + """ + retained = compute_retained_tokens_count( + tokens_per_frame=frame_num_tokens, num_frames=num_frames, q=q + ) + base = retained // num_frames + rem = retained % num_frames + tpf = [base] * (num_frames - 1) + [base + rem] + assert sum(tpf) == retained + return tpf + + +def replace_offsets_with_tokens_per_frame( + *, + pre_chunked_input_ids: list[int], + num_tokens_per_frame: list[int], + frame_offsets_inclusive: list[tuple[int, int]], + filler_token_id: int, +) -> list[int]: + """ + Given a single video, after EVS pruning of redundant tokens, we have a new `num_tokens_per_frame`, therefore the existing input_ids and offsets are stale. + We need to replace all stale offsets with new offsets that reflect the new `num_tokens_per_frame`, respectively. + + Returns: + Modified input_ids with offsets replaced with new offsets. + + Examples: + >>> assert replace_offsets_with_tokens_per_frame( + ... pre_chunked_input_ids=[1, 0, 0, 4, 5, 0, 0, 0, 9, 10, 0, 0, 12, 13], + ... frame_offsets_inclusive=[(1, 2), (5, 7), (10, 11)], + ... num_tokens_per_frame=[1, 4, 2], + ... filler_token_id=0, + ... ) == [1, 0, 4, 5, 0, 0, 0, 0, 9, 10, 0, 0, 12, 13] + + >>> assert replace_offsets_with_tokens_per_frame( + ... pre_chunked_input_ids=[1, 0, 0, 4, 5, 9, 10, 0, 0, 0], + ... frame_offsets_inclusive=[(1, 2), (7, 9)], + ... num_tokens_per_frame=[1, 4], + ... filler_token_id=0, + ... ) == [1, 0, 4, 5, 9, 10, 0, 0, 0, 0] + + >>> assert replace_offsets_with_tokens_per_frame( + ... pre_chunked_input_ids=[0, 0, 1, 4, 0, 0, 0, 5, 9, 10], + ... frame_offsets_inclusive=[(0, 1), (4, 6)], + ... num_tokens_per_frame=[1, 4], + ... filler_token_id=0, + ... ) == [0, 1, 4, 0, 0, 0, 0, 5, 9, 10] + """ + assert isinstance(pre_chunked_input_ids, list) + ids = pre_chunked_input_ids + + if len(frame_offsets_inclusive) == 1: + """There might be no frame separators, in which case there will be one contiguous span of tokens""" + final = ids[0 : frame_offsets_inclusive[0][0]] + frames = [filler_token_id] * sum(num_tokens_per_frame) + final.extend(frames) + else: + cursor = 0 + final = [] + for (start, end), num_tokens in zip( + frame_offsets_inclusive, num_tokens_per_frame, strict=True + ): + final.extend(ids[cursor:start]) + final.extend([filler_token_id] * num_tokens) + cursor = end + 1 + final.extend(ids[frame_offsets_inclusive[-1][1] + 1 :]) + return final diff --git a/python/sglang/srt/multimodal/evs/evs_module.py b/python/sglang/srt/multimodal/evs/evs_module.py new file mode 100644 index 000000000..a7ab6296e --- /dev/null +++ b/python/sglang/srt/multimodal/evs/evs_module.py @@ -0,0 +1,201 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + + +import dataclasses +import typing +from abc import ABC, abstractmethod +from dataclasses import dataclass + +import torch +from transformers import PretrainedConfig + +from sglang.srt.managers.schedule_batch import MultimodalDataItem +from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult +from sglang.srt.multimodal.processors.base_processor import BaseMultimodalProcessor +from sglang.utils import logger + +from .evs_core import compute_retention_mask, replace_offsets_with_tokens_per_frame + + +@dataclasses.dataclass(kw_only=True) +class EVSDataItem(MultimodalDataItem): + thw_grids: list[tuple[int, int, int]] + + +@dataclasses.dataclass(kw_only=True) +class VideoEVSDataItem(EVSDataItem): + pre_chunked_input_ids: torch.Tensor + + def __post_init__(self): + assert self.is_video() + + +@dataclass(kw_only=True) +class EVSEmbeddingResult(EmbeddingResult): + """ + Embedding result that includes per-frame token counts after EVS pruning. + + After pruning, each frame retains a different number of tokens based on its + dissimilarity to the previous frame. This metadata is needed downstream to + adjust the input_ids placeholder spans to match the actual embedding sizes. + + Attributes: + embedding: The pruned video embeddings tensor. + num_tokens_per_frame: Actual retained token count for each frame. + For example, [256, 180, 195, 256] means frame 0 kept all 256 tokens + (first frame is never pruned), while frames 1-2 were pruned. + """ + + num_tokens_per_frame: list[int] + + def redistribute_pruned_frames_placeholders( + self, + input_ids: torch.Tensor, + offsets: list[tuple[int, int]], + *, + item: VideoEVSDataItem, + extend_prefix_len: int, + extend_seq_len: int, + ) -> tuple[torch.Tensor, list[tuple[int, int]]]: + assert len(input_ids) == extend_seq_len + assert isinstance( + item, VideoEVSDataItem + ), f"Expected VideoEVSDataItem, got {type(item)}" + pre_chunked_input_ids = item.pre_chunked_input_ids + filler_token_id = item.pad_value + input_ids_list = replace_offsets_with_tokens_per_frame( + pre_chunked_input_ids=pre_chunked_input_ids, + num_tokens_per_frame=self.num_tokens_per_frame, + frame_offsets_inclusive=offsets, + filler_token_id=filler_token_id, + ) + input_ids = torch.tensor( + input_ids_list, dtype=input_ids.dtype, device=input_ids.device + ) + offsets = BaseMultimodalProcessor.get_mm_items_offset( + input_ids, filler_token_id + ) + input_ids = input_ids[extend_prefix_len : extend_prefix_len + extend_seq_len] + assert ( + len(input_ids) == extend_seq_len + ), f"Input ids length changed after redistribution, got {len(input_ids)} != {extend_seq_len}" + return input_ids, offsets + + +@dataclass(frozen=True, kw_only=True) +class EVSConfig: + video_pruning_rate: float + spatial_merge_size: int = 1 + + def __post_init__(self): + assert ( + self.video_pruning_rate >= 0.0 and self.video_pruning_rate < 1.0 + ), f"Video pruning rate must be between 0.0 and 1.0, got {self.video_pruning_rate=}" + + +class EVS(torch.nn.Module, ABC): + """ + Base class for video models that support EVS pruning. + + Subclass this alongside your model class and implement the static `create_evs_config`. + On initialization, if video_pruning_rate > 0, this mixin replaces the model's + get_video_feature() method with a wrapper that applies EVS pruning. + + Example: See `NemotronH_Nano_VL_V2` + """ + + @staticmethod + @abstractmethod + def create_evs_config(config: PretrainedConfig) -> EVSConfig: + """Extract EVS parameters from model config. Must be implemented by subclass.""" + raise NotImplementedError + + @abstractmethod + def get_video_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: + """Extract EVS parameters from model config. Must be implemented by subclass.""" + raise NotImplementedError + + def __init__( + self, + config: PretrainedConfig, + *args: typing.Any, + **kwargs: typing.Any, + ) -> None: + super().__init__() + model_name = self.__class__.__name__ + self.original_get_video_feature = self.get_video_feature + self.evs_config = self.create_evs_config(config) + self.evs_enabled = self.evs_config.video_pruning_rate > 0.0 + if self.evs_enabled: + logger.info(f"[EVS] enabled for {model_name} [{self.evs_config}]") + self.get_video_feature = self.evs_video + else: + logger.info( + f"[EVS] requested on model {model_name} but is disabled for pruning_rate == 0.0." + ) + + def evs_video(self, items: list[MultimodalDataItem]) -> EVSEmbeddingResult: + """ + Apply EVS pruning to video embeddings. + + Args: + items: List containing a single VideoEVSDataItem with video features. + + Returns: + EVSEmbeddingResult with pruned embeddings and actual token counts per frame. + """ + logger.debug( + f"[EVS] beginning for model {self.__class__.__name__} [evs_config={self.evs_config=}]" + ) + assert len(items) == 1, f"Expected 1 item, got {len(items)}" + item = items[0] + assert isinstance( + item, VideoEVSDataItem + ), f"Expected VideoEVSDataItem with modality VIDEO, got {item}" + + q = self.evs_config.video_pruning_rate + merge = self.evs_config.spatial_merge_size + videos_features = self.original_get_video_feature([item]) + if videos_features.ndim == 3: + videos_features = videos_features.flatten(0, 1) + assert videos_features.ndim == 2, videos_features.ndim + + final_embeddings: list[torch.Tensor] = [] + num_tokens_per_frame: list[int] = [] + + sizes = [(t * h * w // merge**2) for t, h, w in item.thw_grids] + for single_video, video_size_thw in zip( + videos_features.split(sizes), + item.thw_grids, + strict=True, + ): + retention_mask = compute_retention_mask( + single_video, + video_size_thw=video_size_thw, + spatial_merge_size=merge, + q=q, + ) + preserved = single_video[retention_mask] + final_embeddings.append(preserved) + num_frames = video_size_thw[0] + tokens_per_frame = ( + retention_mask.reshape(num_frames, -1).sum(dim=-1).tolist() + ) + num_tokens_per_frame.extend(tokens_per_frame) + final_embeddings_tensor = torch.cat(final_embeddings) + return EVSEmbeddingResult( + embedding=final_embeddings_tensor, + num_tokens_per_frame=num_tokens_per_frame, + ) diff --git a/python/sglang/srt/multimodal/evs/evs_processor.py b/python/sglang/srt/multimodal/evs/evs_processor.py new file mode 100644 index 000000000..251f17bca --- /dev/null +++ b/python/sglang/srt/multimodal/evs/evs_processor.py @@ -0,0 +1,132 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + + +import torch +from transformers import PretrainedConfig + +from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem +from sglang.utils import logger + +from .evs_core import tokens_per_frame +from .evs_module import EVS, EVSConfig, EVSDataItem, VideoEVSDataItem + + +def _non_evs_data_items( + *, + image: torch.Tensor | None, + image_offsets: list[tuple[int, int]], + video: torch.Tensor | None, + video_offsets: list[tuple[int, int]], + input_ids_list: list[int], +): + items: list[MultimodalDataItem] = [] + if image is not None: + item = MultimodalDataItem( + modality=Modality.IMAGE, feature=image, offsets=image_offsets + ) + items.append(item) + if video is not None: + item = MultimodalDataItem( + modality=Modality.VIDEO, feature=video, offsets=video_offsets + ) + items.append(item) + return items + + +class EVSProcessor: + """ + This processor handles prompt construction with the correct number of + placeholder tokens per frame. When EVS is active, it allocates fewer + placeholders based on the pruning rate. When inactive, it uses the full + token count. + """ + + def __init__( + self, + hf_config: PretrainedConfig, + config_to_evs_model: dict[type[PretrainedConfig], type[EVS]], + ): + assert len(config_to_evs_model) > 0 + assert all(issubclass(model, EVS) for model in config_to_evs_model.values()) + + self.evs_config: EVSConfig | None = None + + config_name = hf_config.__class__.__name__ + evs_model = config_to_evs_model.get(hf_config.__class__) + if evs_model is None: + logger.info( + f"[EVS] no model matches {config_name} in {config_to_evs_model}" + ) + return + evs_config = evs_model.create_evs_config(hf_config) + logger.info( + f"""[EVS] {evs_config} {'enabled' if evs_config.video_pruning_rate > 0.0 else 'disabled'} for model={evs_model.__name__}; model_config={config_name}""" + ) + if evs_config.video_pruning_rate > 0.0: + self.evs_config = evs_config + + def static_size_data_items( + self, *, frames_per_video: list[int], num_images: int, rows: int, cols: int + ): + """helper function to create data items for models with static image and video tokens per frame""" + + frame_num_tokens = rows * cols + + if self.evs_config is None: + tpf = [[frame_num_tokens] * num_frames for num_frames in frames_per_video] + return _non_evs_data_items, tpf + + def create_evs_data_items( + *, + input_ids_list: list[int], + image: torch.Tensor | None, + image_offsets: list[tuple[int, int]], + video: torch.Tensor | None, + video_offsets: list[tuple[int, int]], + ) -> list[MultimodalDataItem]: + items = [] + if image is not None: + image_thw_grids = [(1, rows, cols)] * num_images + item = EVSDataItem( + modality=Modality.IMAGE, + feature=image, + offsets=image_offsets, + thw_grids=image_thw_grids, + ) + items.append(item) + if video is not None: + video_thw_grids = [ + (num_frames, rows, cols) for num_frames in frames_per_video + ] + item = VideoEVSDataItem( + modality=Modality.VIDEO, + feature=video, + offsets=video_offsets, + thw_grids=video_thw_grids, + pre_chunked_input_ids=input_ids_list, + ) + items.append(item) + return items + + tpf = [ + tokens_per_frame( + q=self.evs_config.video_pruning_rate, + num_frames=num_frames, + frame_num_tokens=frame_num_tokens, + ) + for num_frames in frames_per_video + ] + + return create_evs_data_items, tpf diff --git a/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py b/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py index cb0ccad67..7464bd341 100644 --- a/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py +++ b/python/sglang/srt/multimodal/processors/nano_nemotron_vl.py @@ -11,14 +11,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from math import sqrt from typing import TYPE_CHECKING import numpy as np import torch from PIL import Image -from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem +from sglang.srt.configs.nano_nemotron_vl import NemotronH_Nano_VL_V2_Config from sglang.srt.models.nano_nemotron_vl import NemotronH_Nano_VL_V2 +from sglang.srt.multimodal.evs import EVSProcessor from sglang.srt.multimodal.internvl_utils import image_to_pixel_values from sglang.srt.multimodal.processors.base_processor import ( BaseMultimodalProcessor, @@ -40,6 +42,9 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): def __init__(self, hf_config, server_args, _image_processor, *args, **kwargs): super().__init__(hf_config, server_args, _image_processor, *args, **kwargs) + self.evs = EVSProcessor( + hf_config, {NemotronH_Nano_VL_V2_Config: NemotronH_Nano_VL_V2} + ) Image.MAX_IMAGE_PIXELS = None self.image_size = hf_config.image_size self.VIDEO_CONTEXT_TOKEN = hf_config.video_context_token @@ -90,10 +95,8 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): def render_image(self, *, num_tiles: int): return f"{self.IMG_START_TOKEN}{self.IMG_CONTEXT_TOKEN * self.num_image_token * num_tiles}{self.IMG_END_TOKEN}" - def render_frame( - self, frame_index: int, *, timestamp: float, start_placeholder_token: str - ): - return f"Frame {frame_index + 1} sampled at {timestamp:.2f} seconds: {start_placeholder_token}{self.IMG_CONTEXT_TOKEN * self.num_image_token}{self.IMG_END_TOKEN}" + def render_frame(self, frame_index: int, *, timestamp: float, num_tokens: int): + return f"Frame {frame_index + 1} sampled at {timestamp:.2f} seconds: {self.PLACEHOLDER}{self.IMG_CONTEXT_TOKEN * num_tokens}{self.IMG_END_TOKEN}" @staticmethod def parse_video(video: "VideoReader") -> tuple[np.ndarray, list[float]]: @@ -117,8 +120,17 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): discard_alpha_channel=True, ) - prompt = input_text + videos = [self.parse_video(video) for video in base_output.videos] + rows = cols = int(sqrt(self.num_image_token)) + create_data_items, tokens_per_frame = self.evs.static_size_data_items( + frames_per_video=[len(frames) for frames, _ in videos], + num_images=len(base_output.images), + rows=rows, + cols=cols, + ) + + prompt = input_text image_feature = None if base_output.images: preprocessed_images = [ @@ -134,8 +146,9 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): video_feature = None if base_output.videos: preprocessed_videos = [] - for video in base_output.videos: - video_array, timestamps = self.parse_video(video) + for (video_array, timestamps), tpf in zip( + videos, tokens_per_frame, strict=True + ): frames_tensors = [ self.preprocess_image( Image.fromarray(frame, mode="RGB"), @@ -149,9 +162,11 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): self.render_frame( i, timestamp=timestamp, - start_placeholder_token=self.PLACEHOLDER, + num_tokens=num_tokens, + ) + for i, (timestamp, num_tokens) in enumerate( + zip(timestamps, tpf, strict=True) ) - for i, timestamp in enumerate(timestamps) ] prompt = prompt.replace( self.VIDEO_CONTEXT_TOKEN, "".join(rendered_frames), 1 @@ -175,20 +190,18 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor): # Cleanup: prompt_ids[prompt_ids == self.PLACEHOLDER_ID] = self.img_start_token_id - items = [] - if image_feature is not None: - item = MultimodalDataItem( - Modality.IMAGE, feature=image_feature, offsets=img_offsets - ) - items.append(item) - if video_feature is not None: - item = MultimodalDataItem( - Modality.VIDEO, feature=video_feature, offsets=video_offsets - ) - items.append(item) + prompt_ids_list = prompt_ids.tolist() + + items = create_data_items( + image=image_feature, + image_offsets=img_offsets, + video=video_feature, + video_offsets=video_offsets, + input_ids_list=prompt_ids_list, + ) return { - "input_ids": prompt_ids.tolist(), + "input_ids": prompt_ids_list, "mm_items": items, "im_start_id": self.img_start_token_id, "im_end_id": self.img_end_token_id, diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index 8dc8094ac..387edfa7c 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -3,6 +3,8 @@ import argparse import asyncio import copy +import doctest +import inspect import json import logging import os @@ -19,7 +21,7 @@ from datetime import datetime from functools import partial, wraps from io import BytesIO from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace from typing import Any, Awaitable, Callable, List, Optional, Tuple import aiohttp @@ -1954,6 +1956,18 @@ def intel_amx_benchmark(extra_args=None, min_throughput=None): return decorator +def run_doctests(obj: Callable[..., Any] | ModuleType): + mod = inspect.getmodule(obj) + globals = dict(mod.__dict__) + finder = doctest.DocTestFinder() + runner = doctest.DocTestRunner(verbose=True) + tests = finder.find(obj, obj.__name__, globs=globals) + assert len(tests) >= 1, f"No tests found for {obj.__name__}" + for test in tests: + result = runner.run(test) + assert result.failed == 0, f"Test {test.name} failed" + + def dump_metric(metric_name: str, value: Any, labels: Optional[dict] = None): """ Output test metric to JSONL and stdout for CI performance tracking. diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 70fb9cdf6..d46ed0be0 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -43,6 +43,7 @@ suites = { TestFile("test_deterministic.py", 228), TestFile("test_constrained_decoding.py", 111), TestFile("test_eval_fp8_accuracy.py", 250), + TestFile("test_evs.py", 20), TestFile("test_external_models.py", 30), TestFile("test_fp8_utils.py", 9), TestFile("rotary_embedding/test_mrope.py", 10), diff --git a/test/srt/test_evs.py b/test/srt/test_evs.py new file mode 100644 index 000000000..917098a9b --- /dev/null +++ b/test/srt/test_evs.py @@ -0,0 +1,54 @@ +from dataclasses import asdict, dataclass +from types import SimpleNamespace + +import pytest + +from sglang.test.test_utils import run_doctests + + +def test_resolve_evs_config(): + from sglang.srt.multimodal.evs import EVS, EVSConfig, EVSProcessor + + @dataclass(frozen=True, kw_only=True) + class EVSModelConfig: + video_pruning_rate: float = 0.1 + spatial_merge_size: int = 2 + + class EVSModel(EVS): + @staticmethod + def create_evs_config(hf_config: EVSModelConfig) -> EVSConfig: + return EVSConfig( + video_pruning_rate=hf_config.video_pruning_rate, + spatial_merge_size=hf_config.spatial_merge_size, + ) + + processor = EVSProcessor( + hf_config=EVSModelConfig(spatial_merge_size=3), + config_to_evs_model={EVSModelConfig: EVSModel}, + ) + expected = EVSConfig(video_pruning_rate=0.1, spatial_merge_size=3) + assert asdict(processor.evs_config) == asdict(expected) + + # No EVS for pruning rate 0.0 + processor = EVSProcessor( + hf_config=EVSModelConfig(video_pruning_rate=0.0), + config_to_evs_model={EVSModelConfig: EVSModel}, + ) + assert processor.evs_config is None + + # No EVS for non-EVS config + processor = EVSProcessor( + hf_config=SimpleNamespace(), + config_to_evs_model={EVSModelConfig: EVSModel}, + ) + assert processor.evs_config is None + + +def test_replace_offsets_with_tokens_per_frame(): + from sglang.srt.multimodal.evs.evs_core import replace_offsets_with_tokens_per_frame + + run_doctests(replace_offsets_with_tokens_per_frame) + + +if __name__ == "__main__": + pytest.main([__file__])