EVS Framework: Support NemotronH_Nano_VL_V2 (#14051)
This commit is contained in:
@@ -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.
|
||||
|
||||
@@ -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 = }"
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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(
|
||||
|
||||
123
python/sglang/srt/multimodal/evs/README.md
Normal file
123
python/sglang/srt/multimodal/evs/README.md
Normal 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
|
||||
11
python/sglang/srt/multimodal/evs/__init__.py
Normal file
11
python/sglang/srt/multimodal/evs/__init__.py
Normal 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",
|
||||
]
|
||||
176
python/sglang/srt/multimodal/evs/evs_core.py
Normal file
176
python/sglang/srt/multimodal/evs/evs_core.py
Normal 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
|
||||
201
python/sglang/srt/multimodal/evs/evs_module.py
Normal file
201
python/sglang/srt/multimodal/evs/evs_module.py
Normal 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,
|
||||
)
|
||||
132
python/sglang/srt/multimodal/evs/evs_processor.py
Normal file
132
python/sglang/srt/multimodal/evs/evs_processor.py
Normal 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
|
||||
@@ -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,
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
54
test/srt/test_evs.py
Normal 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__])
|
||||
Reference in New Issue
Block a user