EVS Framework: Support NemotronH_Nano_VL_V2 (#14051)

This commit is contained in:
Netanel Haber
2026-01-05 10:18:07 +02:00
committed by GitHub
parent b12258bfaa
commit bebd625ba1
14 changed files with 821 additions and 56 deletions

View File

@@ -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.

View File

@@ -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 = }"

View File

@@ -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,

View File

@@ -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):

View File

@@ -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(

View File

@@ -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

View File

@@ -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",
]

View File

@@ -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

View File

@@ -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,
)

View File

@@ -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

View File

@@ -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,

View File

@@ -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.

View File

@@ -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),

54
test/srt/test_evs.py Normal file
View File

@@ -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__])