[EPD][VLM] support video/audio input (#17824)

Co-authored-by: siyu <liusy58@linux.alibaba.com>

python/sglang/srt/disaggregation/encode_grpc_server.py

@@ -26,6 +26,7 @@ from sglang.srt.disaggregation.encode_server import (
     handle_scheduler_receive_url_request,
     launch_encoder,
 )
+from sglang.srt.managers.schedule_batch import Modality
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import random_uuid
 from sglang.srt.utils.network import get_zmq_socket
@@ -96,6 +97,7 @@ class SGLangEncoderServer(SGLangEncoderServicer):
             for socket in self.send_sockets:
                 await socket.send_pyobj(request_dict)
 
+            # gRPC encode is image-only; encoder.encode() requires modality
             (
                 nbytes,
                 embedding_len,
@@ -104,6 +106,7 @@ class SGLangEncoderServer(SGLangEncoderServicer):
                 error_code,
             ) = await self.encoder.encode(
                 mm_items=list(request.mm_items),
+                modality=Modality.IMAGE,
                 req_id=request.req_id,
                 num_parts=request.num_parts,
                 part_idx=request.part_idx,

python/sglang/srt/disaggregation/encode_receiver.py

@@ -1,4 +1,5 @@
 import asyncio
+import itertools
 import logging
 import pickle
 import random
@@ -6,9 +7,10 @@ import threading
 import time
 import uuid
 from abc import ABC, abstractmethod
+from collections import OrderedDict, defaultdict
 from enum import IntEnum
 from http import HTTPStatus
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional
 
 import aiohttp
 import torch
@@ -23,7 +25,7 @@ from sglang.srt.distributed.parallel_state import (
 from sglang.srt.environ import envs
 from sglang.srt.managers.io_struct import GenerateReqInput, TokenizedGenerateReqInput
 from sglang.srt.managers.multimodal_processor import get_mm_processor, import_processors
-from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.managers.schedule_batch import Modality, Req
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import ImageData
 from sglang.srt.utils.hf_transformers_utils import get_processor
@@ -125,54 +127,38 @@ class EmbeddingData:
         req_id,
         num_parts,
         part_idx,
-        image_grid_dim,
+        grid_dim,
+        modality,
         embedding=None,
+        embedding_shape=None,
         error_msg=None,
         error_code=None,
+        **kwargs,
     ):
         self.req_id = req_id
         self.num_parts = num_parts
         self.part_idx = part_idx
-        self.image_grid_dim = image_grid_dim
+        self.grid_dim = grid_dim
+        self.modality = modality
         self.embedding = embedding
         self.send_time = None
         self.dtype = embedding.dtype if embedding is not None else None
-        self.shape = list(embedding.shape) if embedding is not None else None
+        if embedding_shape is not None:
-        # aggregated data
+            self.shape = embedding_shape
-        self.ready_list = [i == self.part_idx for i in range(self.num_parts)]
+        else:
-        self.embedding_list = [
+            self.shape = list(embedding.shape) if embedding is not None else None
-            embedding if i == self.part_idx else None for i in range(self.num_parts)
-        ]
-        self.image_grid_dim_list = [
-            self.image_grid_dim if i == self.part_idx else None
-            for i in range(self.num_parts)
-        ]
         self.error_msg = error_msg
         self.error_code = error_code
+        # Store additional metadata (e.g., video_timestamps for qwen3_vl)
+        for key, value in kwargs.items():
+            setattr(self, key, value)
 
-    def add(self, embedding_data):
+    def get_grid(self):
-        assert self.req_id == embedding_data.req_id
+        """Get the grid dimension of the embedding, used for image/video/audio."""
-        assert not self.ready_list[embedding_data.part_idx]
+        return self.grid_dim
-        self.ready_list[embedding_data.part_idx] = True
-        self.image_grid_dim_list[embedding_data.part_idx] = (
-            embedding_data.image_grid_dim
-        )
-        self.embedding_list[embedding_data.part_idx] = embedding_data.embedding
 
-    def get_embedding(self, is_concat=False):
+    def get_embedding(self):
-        if is_concat:
+        return self.embedding
-            return torch.concat(
-                [embedding.cuda() for embedding in self.embedding_list]
-            ).to("cpu", non_blocking=True)
-        else:
-            return self.embedding_list
-
-    def get_img_grid(self):
-        return torch.concatenate(self.image_grid_dim_list)
-
-    @property
-    def ready(self):
-        return sum(self.ready_list) == self.num_parts
 
     def __repr__(self):
         return f"EmbeddingData(req_id={self.req_id}, num_parts={self.num_parts}, part_idx={self.part_idx}) error_msg={self.error_msg}"
@@ -182,16 +168,173 @@ class EmbeddingData:
             req_id=self.req_id,
             num_parts=self.num_parts,
             part_idx=self.part_idx,
-            image_grid_dim=self.image_grid_dim,
+            grid_dim=self.grid_dim,
+            modality=self.modality,
+            embedding=None,
+            embedding_shape=self.shape,
             error_msg=self.error_msg,
             error_code=self.error_code,
         )
-        new_data.send_time = self.send_time
+        for key, value in self.__dict__.items():
-        new_data.dtype = self.dtype
+            if key.startswith("_") or key == "embedding":
-        new_data.shape = self.shape
+                continue
+            setattr(new_data, key, value)
         return new_data
 
 
+# Modality -> (list attr name, whether to flatten grid for that list)
+_MODALITY_GRID_ATTRS = {
+    Modality.IMAGE: ("img_grid_thw", False),
+    Modality.VIDEO: ("video_grid_thw", False),
+    Modality.AUDIO: ("audio_feature_lens", True),
+}
+_VIDEO_META_ATTRS = ("video_timestamps", "second_per_grid_ts")
+
+
+def _cat_grid(dims, flatten_items=False):
+    """Concatenate non-None tensors from a list; optionally flatten each before cat."""
+    valid = (
+        [g.flatten() for g in dims if g is not None]
+        if flatten_items
+        else [g for g in dims if g is not None]
+    )
+    return torch.cat(valid, dim=0) if valid else None
+
+
+class MultiModalEmbeddingData(EmbeddingData):
+    def __init__(
+        self,
+        part_idx,
+        num_parts,
+        req_id,
+        grid_dim,
+        modality,
+        embedding,
+        embedding_shape,
+        **kwargs,
+    ):
+        super().__init__(
+            req_id,
+            num_parts,
+            part_idx,
+            grid_dim,
+            modality,
+            embedding,
+            embedding_shape,
+            **kwargs,
+        )
+        self.img_grid_thw = [None] * num_parts
+        self.video_grid_thw = [None] * num_parts
+        self.audio_feature_lens = [None] * num_parts
+        self.modality_list = [
+            modality if part_idx == i else None for i in range(num_parts)
+        ]
+        self.ready_list = [i == part_idx for i in range(num_parts)]
+        self.embedding_list = [
+            embedding if i == part_idx else None for i in range(num_parts)
+        ]
+        self.embedding_shape_list = [
+            embedding_shape if i == part_idx else None for i in range(num_parts)
+        ]
+        self.video_timestamps = [None] * num_parts
+        self.second_per_grid_ts = [None] * num_parts
+
+        self._set_part_grid(part_idx, modality, self.get_grid())
+        if modality == Modality.VIDEO:
+            self._set_video_meta_for_part(part_idx, kwargs)
+
+    def _set_part_grid(self, part_idx, modality, grid):
+        """Set the grid for one part according to modality (IMAGE/VIDEO/AUDIO)."""
+        spec = _MODALITY_GRID_ATTRS.get(modality)
+        if spec is None:
+            raise ValueError(f"Invalid modality: {modality}")
+        attr_name, flatten = spec
+        value = grid.flatten() if flatten else grid
+        getattr(self, attr_name)[part_idx] = value
+
+    def _set_video_meta_for_part(self, part_idx, source):
+        """Copy video_timestamps and second_per_grid_ts from source (dict or object)."""
+        for attr_name in _VIDEO_META_ATTRS:
+            val = (
+                source.get(attr_name)
+                if isinstance(source, dict)
+                else getattr(source, attr_name, None)
+            )
+            if val is not None:
+                getattr(self, attr_name)[part_idx] = val
+
+    @classmethod
+    def from_embedding_data(cls, embedding_data: EmbeddingData):
+        """Create MultiModalEmbeddingData from an EmbeddingData instance."""
+        # Only forward known optional attrs (e.g. video metadata) so they land on the instance
+        extra = {}
+        for attr in _VIDEO_META_ATTRS:
+            val = getattr(embedding_data, attr, None)
+            if val is not None:
+                extra[attr] = val
+        mm_data = cls(
+            part_idx=embedding_data.part_idx,
+            num_parts=embedding_data.num_parts,
+            req_id=embedding_data.req_id,
+            grid_dim=embedding_data.grid_dim,
+            modality=embedding_data.modality,
+            embedding=embedding_data.embedding,
+            embedding_shape=embedding_data.shape,
+            **extra,
+        )
+        mm_data.send_time = embedding_data.send_time
+        return mm_data
+
+    def __repr__(self):
+        return f"MultiModalEmbeddingData(req_id={self.req_id}, num_parts={self.num_parts}, part_idx={self.part_idx}, modality={self.modality})"
+
+    def get_embedding(self, is_concat=False):
+        if is_concat:
+            groups = defaultdict(list)
+            for i, e in enumerate(self.embedding_list):
+                if e is not None:
+                    groups[self.modality_list[i]].append(e.cuda())
+            return {
+                mod: torch.concat(tensors).to("cpu", non_blocking=True)
+                for mod, tensors in groups.items()
+            }
+        return self.embedding_list
+
+    @property
+    def ready(self):
+        return sum(self.ready_list) == self.num_parts
+
+    def get_mm_extra_meta(self):
+        """Build kwargs for mm_processor.get_mm_data() from grid and optional video meta."""
+        kwargs = {
+            "img_grid_thw": _cat_grid(self.img_grid_thw),
+            "video_grid_thw": _cat_grid(self.video_grid_thw),
+            "audio_feature_lens": _cat_grid(
+                self.audio_feature_lens, flatten_items=True
+            ),
+        }
+        for attr in _VIDEO_META_ATTRS:
+            lst = getattr(self, attr, None)
+            if not lst:
+                continue
+            valid = [a for a in lst if a is not None]
+            if valid:
+                kwargs[attr] = list(itertools.chain(*valid))
+        return kwargs
+
+    def add(self, embedding_data: EmbeddingData):
+        assert self.req_id == embedding_data.req_id
+        assert not self.ready_list[embedding_data.part_idx]
+        pid = embedding_data.part_idx
+        self.ready_list[pid] = True
+        self.modality_list[pid] = embedding_data.modality
+        self.embedding_list[pid] = embedding_data.get_embedding()
+        self.embedding_shape_list[pid] = embedding_data.shape
+        self._set_part_grid(pid, embedding_data.modality, embedding_data.get_grid())
+        if embedding_data.modality == Modality.VIDEO:
+            self._set_video_meta_for_part(pid, embedding_data)
+
+
 class WaitingImageRequestStatus(IntEnum):
     FAIL = -1
     PENDING = 0
@@ -199,6 +342,41 @@ class WaitingImageRequestStatus(IntEnum):
     TIMEOUT = -2
 
 
+def create_part_req_id(original_req_id: str, part_idx: int) -> str:
+    """Create a unique part request ID by appending part index suffix."""
+    return f"{original_req_id}_local_part_{part_idx}"
+
+
+def extract_original_req_id(part_req_id: str) -> str:
+    """Extract the original request ID from a part request ID."""
+    if "_local_part_" in part_req_id:
+        return part_req_id.rsplit("_local_part_", 1)[0]
+    return part_req_id
+
+
+def calculate_modality_num_parts(modalities, num_items_assigned):
+    """
+    Calculate total number of parts and number of parts per modality.
+
+    Args:
+        modalities: List of modalities in order
+        num_items_assigned: Dictionary mapping modality to list of assignment counts per encoder
+
+    Returns:
+        Tuple of (total_num_parts, modality_num_parts_dict)
+        - total_num_parts: Total number of parts across all modalities
+        - modality_num_parts: Dictionary mapping modality to number of parts for that modality
+    """
+    total_num_parts = 0
+    modality_num_parts = {}
+    for modality in modalities:
+        num_items_assigned_modality = num_items_assigned.get(modality)
+        num_parts = sum(1 for x in num_items_assigned_modality if x != 0)
+        modality_num_parts[modality] = num_parts
+        total_num_parts += num_parts
+    return total_num_parts, modality_num_parts
+
+
 # For zmq_to_scheduler
 class WaitingImageRequest:
     def __init__(
@@ -247,21 +425,38 @@ class WaitingImageRequest:
             ) as session:
                 tasks = []
                 logger.info(f"{self.num_items_assigned = } ")
-                for idx, assigned_num in enumerate(self.num_items_assigned):
-                    if assigned_num == 0:
-                        continue
-                    encoder_url = self.encoder_urls[idx]
-                    target_url = f"{encoder_url}/scheduler_receive_url"
-                    payload = {
-                        "req_id": req_id,
-                        "receive_count": receive_count,
-                        "receive_url": f"{host_name}:{embedding_port}",
-                    }
 
-                    logger.info(f"Preparing to send  to {target_url}")
+                # Calculate part_idx_offset similar to encode() method
+                modalities = list(self.num_items_assigned.keys())
+                _, modality_num_parts = calculate_modality_num_parts(
+                    modalities, self.num_items_assigned
+                )
 
-                    task = _send_single_request(session, target_url, payload)
+                part_idx_offset = 0
-                    tasks.append(task)
+                for modality in modalities:
+                    assigned_nums = self.num_items_assigned[modality]
+                    num_parts = modality_num_parts[modality]
+                    cum_idx = 0
+                    for idx, assigned_num in enumerate(assigned_nums):
+                        if assigned_num == 0:
+                            continue
+                        part_idx = part_idx_offset + cum_idx
+                        part_req_id = create_part_req_id(req_id, part_idx)
+                        encoder_url = self.encoder_urls[idx]
+                        target_url = f"{encoder_url}/scheduler_receive_url"
+                        payload = {
+                            "req_id": part_req_id,  # use part_req_id to match encode request
+                            "receive_count": receive_count,
+                            "receive_url": f"{host_name}:{embedding_port}",
+                            "modality": modality.name,
+                        }
+                        logger.info(
+                            f"Preparing to send to {target_url} with part_req_id={part_req_id}"
+                        )
+                        task = _send_single_request(session, target_url, payload)
+                        tasks.append(task)
+                        cum_idx += 1
+                    part_idx_offset += num_parts
 
                 if not tasks:
                     logger.info("No tasks to send.")
@@ -305,20 +500,30 @@ class WaitingImageRequest:
                 return
 
             buffer = parts[1].buffer if hasattr(parts[1], "buffer") else parts[1]
-            recv_obj.embedding = torch.frombuffer(buffer, dtype=recv_obj.dtype).reshape(
+            recv_obj.embedding = (
-                recv_obj.shape
+                torch.frombuffer(buffer, dtype=recv_obj.dtype)
+                .reshape(recv_obj.shape)
+                .clone()
             )
-            recv_obj.embedding_list[recv_obj.part_idx] = recv_obj.embedding
+
+            # Extract original req_id from part_req_id
+            part_req_id = recv_obj.req_id
+            original_req_id = extract_original_req_id(part_req_id)
+            # Update recv_obj.req_id to original for aggregation
+            recv_obj.req_id = original_req_id
+
             if self.recv_embedding_data is None:
-                self.recv_embedding_data = recv_obj
+                self.recv_embedding_data = MultiModalEmbeddingData.from_embedding_data(
+                    recv_obj
+                )
             else:
                 self.recv_embedding_data.add(recv_obj)
 
         recv_embedding = self.recv_embedding_data.get_embedding(is_concat=True)
-        img_grid_thw = self.recv_embedding_data.get_img_grid()
-
         mm_inputs = self.mm_processor.get_mm_data(
-            self.recv_req.input_text, recv_embedding, img_grid_thw
+            self.recv_req.input_text,
+            recv_embedding,
+            **self.recv_embedding_data.get_mm_extra_meta(),
         )
         self.recv_req.mm_inputs = mm_inputs
         self.recv_req.input_ids = mm_inputs["input_ids"]
@@ -330,8 +535,11 @@ class WaitingImageRequestGrpc(WaitingImageRequest):
     def send_encode_request(self):
         async def send_embedding_port(req_id, receive_count, host_name, embedding_port):
             tasks = []
-            logger.info(f"{self.num_items_assigned = } ")
+            # gRPC image-only: flatten modality dict to flat list
-            for idx, assigned_num in enumerate(self.num_items_assigned):
+            assigned = list(self.num_items_assigned.values())[0]
+            logger.info(f"num_items_assigned={assigned}")
+
+            for idx, assigned_num in enumerate(assigned):
                 if assigned_num == 0:
                     continue
                 encoder_url = self.encoder_urls[idx]
@@ -394,11 +602,22 @@ class MMReceiverBase(ABC):
         self.context = zmq.asyncio.Context(20)
         self.encoder_transfer_backend = server_args.encoder_transfer_backend
         self.encode_urls = server_args.encoder_urls
-        self.encode_idx = list(range(len(self.encode_urls)))
         self.host = get_local_ip_auto(server_args.host)
         if self.encoder_transfer_backend == "mooncake":
             self.dtype = dtype
             self.embeddings_engine = get_mooncake_transfer_engine()
+            if self.embeddings_engine is None:
+                from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+                    init_mooncake_transfer_engine,
+                )
+
+                self.embeddings_engine = init_mooncake_transfer_engine(
+                    hostname=self.host,
+                    ib_device=(
+                        server_args.disaggregation_ib_device
+                        or server_args.mooncake_ib_device
+                    ),
+                )
             self.embeddings_buffer = dict()
         elif self.encoder_transfer_backend == "zmq_to_scheduler":
             self.pp_rank = pp_rank
@@ -455,20 +674,17 @@ class MMReceiverBase(ABC):
         pass
 
     async def recv_mm_data(
-        self, img_data, mm_processor, prompt, need_wait_for_image=True
+        self, request_obj, mm_processor, prompt, need_wait_for_mm_inputs=True
     ):
         req_id = None
         try:
-            if len(self.encode_urls) == 0 or not need_wait_for_image:
+            if len(self.encode_urls) == 0 or not need_wait_for_mm_inputs:
                 return None
             req_id = uuid.uuid4().hex
             embedding_port, recv_socket = get_zmq_socket_on_host(self.context, zmq.PULL)
-            if not isinstance(img_data, list):
+            mm_data = self._extract_url_data(request_obj)
-                img_data = [img_data.url]
-            else:
-                img_data = [img.url for img in img_data]
             asyncio.create_task(
-                self.encode(req_id, img_data, embedding_port, "encode", "send")
+                self.encode(req_id, mm_data, embedding_port, "encode", "send")
             )
             return await asyncio.wait_for(
                 self._recv_mm_data(req_id, recv_socket, mm_processor, prompt),
@@ -501,7 +717,7 @@ class MMReceiverBase(ABC):
 
         recv_embedding = None
 
-        recv_embedding_data: EmbeddingData = None
+        recv_embedding_data: MultiModalEmbeddingData = None
 
         try:
             while recv_embedding_data is None or not recv_embedding_data.ready:
@@ -517,6 +733,11 @@ class MMReceiverBase(ABC):
                     self._cleanup_mooncake_buffer(req_id)
                     return None
                 logger.debug("recv_obj=%s", recv_obj)
+                # Extract original req_id from part_req_id
+                part_req_id = recv_obj.req_id
+                original_req_id = extract_original_req_id(part_req_id)
+                # Update recv_obj.req_id to original for aggregation
+                recv_obj.req_id = original_req_id
                 if self.encoder_transfer_backend == "zmq_to_tokenizer":
                     if len(parts) < 2:
                         logger.error(
@@ -534,8 +755,9 @@ class MMReceiverBase(ABC):
                         .clone()
                     )
                 if recv_embedding_data is None:
-                    recv_obj.embedding_list[recv_obj.part_idx] = recv_obj.embedding
+                    recv_embedding_data = MultiModalEmbeddingData.from_embedding_data(
-                    recv_embedding_data = recv_obj
+                        recv_obj
+                    )
                 else:
                     recv_embedding_data.add(recv_obj)
 
@@ -545,14 +767,32 @@ class MMReceiverBase(ABC):
                         "mooncake: embeddings_buffer missing req_id=%s", req_id
                     )
                     return None
-                recv_embedding = self.embeddings_buffer[req_id]
+                raw_buffer = self.embeddings_buffer.pop(req_id)
-                del self.embeddings_buffer[req_id]
+                self.embeddings_engine.deregister(raw_buffer.data_ptr())
-                self.embeddings_engine.deregister(recv_embedding.data_ptr())
+                byte_offset = 0
-            elif self.encoder_transfer_backend == "zmq_to_tokenizer":
+                for i in range(recv_embedding_data.num_parts):
-                recv_embedding = recv_embedding_data.get_embedding(is_concat=True)
+                    shape = recv_embedding_data.embedding_shape_list[i]
+                    if shape is None:
+                        continue
+                    part_bytes = (
+                        shape[0]
+                        * shape[1]
+                        * torch.tensor([], dtype=self.dtype).element_size()
+                    )
+                    recv_embedding_data.embedding_list[i] = (
+                        raw_buffer[byte_offset : byte_offset + part_bytes]
+                        .view(self.dtype)
+                        .reshape(shape)
+                    )
+                    byte_offset += part_bytes
 
-            img_grid_thw = recv_embedding_data.get_img_grid()
+            recv_embedding = recv_embedding_data.get_embedding(is_concat=True)
-            mm_inputs = mm_processor.get_mm_data(prompt, recv_embedding, img_grid_thw)
+
+            mm_inputs = mm_processor.get_mm_data(
+                prompt,
+                recv_embedding,
+                **recv_embedding_data.get_mm_extra_meta(),
+            )
             return mm_inputs
         finally:
             recv_socket.close()
@@ -561,30 +801,24 @@ class MMReceiverBase(ABC):
         self._send_encode_request(obj)
 
     def _send_encode_request(self, obj):
-        if obj.image_data is None:
+        mm_data = self._extract_url_data(obj)
-            image_urls = []
-        elif not isinstance(obj.image_data, list):
-            image_urls = [obj.image_data.url]
-        else:
-            image_urls = [img.url for img in obj.image_data]
         if obj.rid is None:
             obj.rid = uuid.uuid4().hex
-        if image_urls and self.encode_urls:
+        if mm_data and self.encode_urls:
-            logger.info(f"Processing {len(image_urls)} images for request {obj.rid}")
+            logger.info(f"Processing {len(mm_data)} mm items for request {obj.rid}")
-            obj.need_wait_for_image = True
+            obj.need_wait_for_mm_inputs = True
 
-            encode_idx = list(range(len(self.encode_urls)))
+            num_items_assigned = self._assign_items_by_modality(
-            random.shuffle(encode_idx)
+                mm_data, len(self.encode_urls)
-            obj.num_items_assigned = [
+            )
-                (idx + len(image_urls)) // len(self.encode_urls) for idx in encode_idx
+            obj.num_items_assigned = num_items_assigned
-            ]
             encode_thread = threading.Thread(
                 target=self._run_encode_in_thread,
                 args=(
                     obj.rid,
-                    image_urls,
+                    mm_data,
                     "encode",
-                    obj.num_items_assigned,
+                    num_items_assigned,
                     None,
                 ),
                 daemon=True,
@@ -597,7 +831,7 @@ class MMReceiverBase(ABC):
         for recv_req in recv_reqs:
             if (
                 isinstance(recv_req, TokenizedGenerateReqInput)
-                and recv_req.need_wait_for_image is True
+                and recv_req.need_wait_for_mm_inputs is True
             ):
                 waiting_req = waiting_cls(
                     rid=recv_req.rid,
@@ -666,13 +900,13 @@ class MMReceiverBase(ABC):
         return new_recv_reqs, abort_reqs
 
     def _run_encode_in_thread(
-        self, req_id, img_data, endpoint_encode, num_items_assigned, embedding_port
+        self, req_id, mm_data, endpoint_encode, num_items_assigned, embedding_port
     ):
         try:
             asyncio.run(
                 self.encode(
                     req_id=req_id,
-                    img_data=img_data,
+                    mm_data=mm_data,
                     embedding_port=embedding_port,
                     endpoint_encode=endpoint_encode,
                     endpoint_send=None,
@@ -718,11 +952,8 @@ class MMReceiverBase(ABC):
         req.tokenizer = self.scheduler.tokenizer
         return req
 
-    async def allocate_embedding_buffer(self, req_id, embedding_length, embedding_dim):
+    async def allocate_embedding_buffer(self, req_id, total_bytes):
-        embeddings = torch.zeros(
+        embeddings = torch.empty(total_bytes, dtype=torch.uint8)
-            (embedding_length, embedding_dim),
-            dtype=self.dtype,
-        )
         self.embeddings_engine.register(
             embeddings.data_ptr(),
             embeddings.nbytes,
@@ -730,6 +961,76 @@ class MMReceiverBase(ABC):
         self.embeddings_buffer[req_id] = embeddings
         return embeddings.data_ptr()
 
+    def _assign_items_by_modality(
+        self, mm_data, encoder_num, random_shuffle=True
+    ) -> Dict:
+        """
+        Assign multimodal items across encoders by modality with cross-modality load balancing.
+
+        Args:
+            mm_data: List of multimodal data items, each with a "modality" key
+            encoder_num: Number of encoders
+            random_shuffle: Whether to shuffle the encoder indices
+
+        Returns:
+            Dictionary mapping modality to list of assignment counts per encoder
+            Format: {modality: [count_for_encoder_0, count_for_encoder_1, ...]}
+        """
+        encode_idx = list(range(encoder_num))
+        if random_shuffle:
+            random.shuffle(encode_idx)
+        # Get unique modalities with order preserved
+        modalities = list(dict.fromkeys(mm_item.get("modality") for mm_item in mm_data))
+        # Use OrderedDict to explicitly maintain modality order
+        num_items_assigned = OrderedDict()
+        current_offset = 0
+
+        for modality in modalities:
+            mm_data_modality = [
+                mm_item for mm_item in mm_data if mm_item.get("modality") == modality
+            ]
+            num_items = len(mm_data_modality)
+            if num_items == 0:
+                continue
+
+            base = num_items // len(encode_idx)
+            remainder = num_items % len(encode_idx)
+            # Rotate assignments based on current_offset to balance load across modalities
+            assignments = [0] * len(encode_idx)
+            for i in range(len(encode_idx)):
+                # keep shuffle order when assigning items to encoders
+                pos_in_shuffled = (current_offset + i) % len(encode_idx)
+                actual_encoder_idx = encode_idx[pos_in_shuffled]
+                assignments[actual_encoder_idx] = base + (1 if i < remainder else 0)
+            num_items_assigned[modality] = assignments
+            current_offset = (current_offset + remainder) % len(encode_idx)
+
+        return num_items_assigned
+
+    def _extract_url_data(self, request_obj) -> List[Dict]:
+        mm_data = []
+        for attr, modality in [
+            ("image_data", Modality.IMAGE),
+            ("video_data", Modality.VIDEO),
+            ("audio_data", Modality.AUDIO),
+        ]:
+            mm_items = getattr(request_obj, attr, None)
+            if mm_items:
+                if not isinstance(mm_items, list):
+                    mm_items = [mm_items]
+                for mm_item in mm_items:
+                    mm_data.append(
+                        {
+                            "url": (
+                                mm_item.url
+                                if isinstance(mm_item, ImageData)
+                                else mm_item
+                            ),
+                            "modality": modality,
+                        }
+                    )
+        return mm_data
+
 
 class MMReceiverHTTP(MMReceiverBase):
     def __init__(
@@ -759,42 +1060,65 @@ class MMReceiverHTTP(MMReceiverBase):
     async def encode(
         self,
         req_id,
-        img_data,
+        mm_data,
         embedding_port,
         endpoint_encode,
         endpoint_send,
         num_items_assigned=None,
     ):
-        if len(img_data) == 0:
+        if len(mm_data) == 0:
             return
 
-        # Split mm_items
+        # get unique modalities with order preserved
+        modalities = [mm_item.get("modality") for mm_item in mm_data]
+        modalities = list(dict.fromkeys(modalities))
         encode_requests = []
+
         if num_items_assigned is None:
-            random.shuffle(self.encode_idx)
+            num_items_assigned = self._assign_items_by_modality(
-            num_items_assigned = [
+                mm_data, len(self.encode_urls)
-                (idx + len(img_data)) // len(self.encode_urls)
-                for idx in self.encode_idx
-            ]
-        num_parts = sum(1 for x in num_items_assigned if x != 0)
-        cum_num_items = 0
-        cum_idx = 0
-        for idx, assigned_num in enumerate(num_items_assigned):
-            if assigned_num == 0:
-                continue
-            encode_requests.append(
-                {
-                    "encoder_idx": idx,
-                    "mm_items": img_data[cum_num_items : cum_num_items + assigned_num],
-                    "num_parts": num_parts,
-                    "part_idx": cum_idx,
-                    "req_id": req_id,
-                    "prefill_host": self.host,
-                    "embedding_port": embedding_port,
-                }
             )
-            cum_idx += 1
+
-            cum_num_items += assigned_num
+        # Calculate total num_parts across all modalities
+        total_num_parts, modality_num_parts = calculate_modality_num_parts(
+            modalities, num_items_assigned
+        )
+
+        part_idx_offset = 0
+        for modality in modalities:
+            num_items_assigned_modality = num_items_assigned.get(modality)
+            mm_data_modality = [
+                mm_item for mm_item in mm_data if mm_item.get("modality") == modality
+            ]
+
+            num_parts = modality_num_parts[modality]
+            cum_num_items = 0
+            cum_idx = 0
+            for idx, assigned_num in enumerate(num_items_assigned_modality):
+                if assigned_num == 0:
+                    continue
+                part_idx = part_idx_offset + cum_idx
+                part_req_id = create_part_req_id(req_id, part_idx)
+                encode_requests.append(
+                    {
+                        "encoder_idx": idx,
+                        "mm_items": [
+                            mm_item.get("url")
+                            for mm_item in mm_data_modality[
+                                cum_num_items : cum_num_items + assigned_num
+                            ]
+                        ],
+                        "num_parts": total_num_parts,
+                        "part_idx": part_idx,
+                        "req_id": part_req_id,  # use part_req_id to avoid key collision
+                        "modality": modality.name,  # convert enum to string for json serialization
+                        "prefill_host": self.host,
+                        "embedding_port": embedding_port,
+                    }
+                )
+                cum_idx += 1
+                cum_num_items += assigned_num
+            part_idx_offset += num_parts
 
         async with aiohttp.ClientSession(
             timeout=aiohttp.ClientTimeout(
@@ -832,21 +1156,21 @@ class MMReceiverHTTP(MMReceiverBase):
 
             # mooncake backend: send bootstrap info
 
-            embedding_size_list_sort = [None for _ in range(num_parts)]
+            embedding_size_list_sort = [None for _ in range(total_num_parts)]
-            embedding_length_tot = 0
+            response_json_list_sort = [None for _ in range(total_num_parts)]
-            response_json_list_sort = [None for _ in range(num_parts)]
             for response_json in response_json_list_unsort:
                 idx = response_json["part_idx"]
                 embedding_size_list_sort[idx] = response_json["embedding_size"]
-                embedding_length_tot += response_json["embedding_len"]
                 response_json_list_sort[idx] = response_json
 
+            total_embedding_bytes = sum(
+                s for s in embedding_size_list_sort if s is not None
+            )
             offset = 0
             metadata_tasks = []
             buffer_address = await self.allocate_embedding_buffer(
                 req_id,
-                embedding_length_tot,
+                total_embedding_bytes,
-                response_json_list_sort[0]["embedding_dim"],
             )
             for idx in range(len(tasks)):
                 response_json = response_json_list_sort[idx]
@@ -903,21 +1227,38 @@ class MMReceiverGrpc(MMReceiverBase):
     async def encode(
         self,
         req_id,
-        img_data,
+        mm_data,
         embedding_port,
         endpoint_encode,
         endpoint_send,
         num_items_assigned=None,
     ):
-        if not img_data:
+        if not mm_data:
             return
 
+        # gRPC currently only supports image; flatten new dict formats to simple lists
+        if mm_data and isinstance(mm_data[0], dict):
+            non_image = [
+                item.get("modality")
+                for item in mm_data
+                if item.get("modality") != Modality.IMAGE
+            ]
+            if non_image:
+                raise NotImplementedError(
+                    f"gRPC encode only supports IMAGE modality, got: {non_image}"
+                )
+            img_data = [item.get("url") for item in mm_data]
+        else:
+            img_data = mm_data
+        if isinstance(num_items_assigned, dict):
+            num_items_assigned = list(num_items_assigned.values())[0]
+
         encode_requests = []
         if num_items_assigned is None:
-            random.shuffle(self.encode_idx)
+            encode_idx = list(range(len(self.encode_urls)))
+            random.shuffle(encode_idx)
             num_items_assigned = [
-                (idx + len(img_data)) // len(self.encode_urls)
+                (idx + len(img_data)) // len(self.encode_urls) for idx in encode_idx
-                for idx in self.encode_idx
             ]
         num_parts = sum(1 for x in num_items_assigned if x != 0)
         cum_num_items = 0
@@ -972,19 +1313,17 @@ class MMReceiverGrpc(MMReceiverBase):
             return
 
         embedding_size_by_part = [None for _ in range(num_parts)]
-        embedding_length_tot = 0
         response_json_sorted = [None for _ in range(num_parts)]
         for response_json in response_json_unsorted:
             idx = response_json["part_idx"]
             embedding_size_by_part[idx] = response_json["embedding_size"]
-            embedding_length_tot += response_json["embedding_len"]
             response_json_sorted[idx] = response_json
 
+        total_embedding_bytes = sum(s for s in embedding_size_by_part if s is not None)
         offset = 0
         buffer_address = await self.allocate_embedding_buffer(
             req_id,
-            embedding_length_tot,
+            total_embedding_bytes,
-            response_json_sorted[0]["embedding_dim"],
         )
         grpc_metadata_tasks = []
         for response_json in response_json_sorted:

python/sglang/srt/disaggregation/encode_server.py

@@ -18,7 +18,7 @@ import zmq
 import zmq.asyncio
 from fastapi import FastAPI
 from fastapi.responses import ORJSONResponse, Response
-from transformers import AutoImageProcessor
+from transformers import AutoProcessor
 
 from sglang.srt.configs.device_config import DeviceConfig
 from sglang.srt.configs.load_config import LoadConfig
@@ -37,6 +37,7 @@ from sglang.srt.managers.io_struct import ProfileReq, ProfileReqInput, ProfileRe
 from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
 from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
 from sglang.srt.model_loader import get_model
+from sglang.srt.multimodal.processors.qwen_vl import preprocess_video
 from sglang.srt.server_args import (
     PortArgs,
     ServerArgs,
@@ -117,18 +118,46 @@ def _convert(data):
         return data
 
 
-_image_grid_attrs = ["image_grid_thw", "image_grid_hws"]
+_mm_grid_attrs = {
+    Modality.IMAGE: ["image_grid_thw", "image_grid_hws"],
+    Modality.VIDEO: ["video_grid_thw"],
+    Modality.AUDIO: ["audio_feature_lens_raw"],
+}
+
+_mm_feature_attrs = {
+    Modality.IMAGE: ["pixel_values"],
+    Modality.VIDEO: ["pixel_values_videos"],
+    Modality.AUDIO: ["input_features"],
+}
 
 
-def _get_image_grid_dim(images_input):
+def _get_mm_grid_dim(mm_inputs, modality):
-    for attr in _image_grid_attrs:
+    for attr in _mm_grid_attrs[modality]:
-        if attr in images_input:
+        if attr in mm_inputs:
-            return images_input[attr]
+            return mm_inputs[attr]
+    raise ValueError(f"Grid dim ({_mm_grid_attrs[modality]}) not found in {mm_inputs}")
+
+
+def _get_mm_feature(mm_inputs, modality):
+    for attr in _mm_feature_attrs[modality]:
+        if attr in mm_inputs:
+            return mm_inputs[attr]
     raise ValueError(
-        f"Image grid dim ({_image_grid_attrs}) not found in {images_input}"
+        f"Feature attrs ({_mm_feature_attrs[modality]}) not found in {mm_inputs}"
     )
 
 
+def _build_mm_aux_data(mm_inputs):
+    """
+    Build auxiliary data for video modality.
+    """
+    aux_data = {
+        "video_timestamps": mm_inputs.get("video_timestamps", None),
+        "second_per_grid_ts": mm_inputs.get("second_per_grid_ts", None),
+    }
+    return aux_data
+
+
 class MMEncoder:
     def __init__(
         self,
@@ -142,17 +171,11 @@ class MMEncoder:
         set_global_server_args_for_scheduler(server_args)
         self.rank = rank
         self.profiler = EncoderProfiler(rank)
-
+        self._load_mm_processor(server_args)
-        self.image_processor = AutoImageProcessor.from_pretrained(
-            server_args.model_path,
-            trust_remote_code=server_args.trust_remote_code,
-            use_fast=True,
-        )
 
         self.model_config = ModelConfig.from_server_args(
             server_args,
         )
-
         self.load_config = LoadConfig(
             load_format=server_args.load_format,
             download_dir=server_args.download_dir,
@@ -161,6 +184,9 @@ class MMEncoder:
             remote_instance_weight_loader_seed_instance_service_port=server_args.remote_instance_weight_loader_seed_instance_service_port,
             remote_instance_weight_loader_send_weights_group_ports=server_args.remote_instance_weight_loader_send_weights_group_ports,
         )
+        self.model_type = getattr(
+            self.model_config.hf_config, "model_type", "unknown"
+        ).lower()
 
         self.device = server_args.device
         self.gpu_id = server_args.base_gpu_id + rank
@@ -172,7 +198,10 @@ class MMEncoder:
 
         torch.get_device_module(self.device).set_device(self.gpu_id)
 
-        self.use_image_processor_gpu = use_image_processor_gpu
+        self.use_image_processor_gpu = (
+            use_image_processor_gpu and not server_args.disable_fast_image_processor
+        )
+        self._build_vision_config(server_args.mm_process_config)
 
         init_distributed_environment(
             backend=get_default_distributed_backend(self.device),
@@ -214,10 +243,11 @@ class MMEncoder:
                 EmbeddingCacheController,
             )
 
+            hidden_dims = self._infer_embedding_dims()
             self.mm_global_cache = EmbeddingCacheController(
                 rank,
                 server_args.tp_size,
-                hidden_dim=self.model_config.hidden_size,
+                hidden_dims=hidden_dims,
                 tp_group=get_tp_group().cpu_group,
                 all_rank_get=False,
             )
@@ -233,11 +263,157 @@ class MMEncoder:
                 self.local_ip = get_local_ip_auto()
 
                 self.engine = get_mooncake_transfer_engine()
+                if self.engine is None:
+                    from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+                        init_mooncake_transfer_engine,
+                    )
+
+                    self.engine = init_mooncake_transfer_engine(
+                        hostname=self.local_ip,
+                        gpu_id=self.gpu_id,
+                        ib_device=(
+                            self.server_args.disaggregation_ib_device
+                            or self.server_args.mooncake_ib_device
+                        ),
+                    )
 
             self.embedding_to_send = dict()
 
         logger.info(f"rank {rank} init finish ")
 
+    def _infer_embedding_dims(self) -> dict:
+        """Infer per-modality embedding dimensions from hf_config at init time."""
+        default = self.model_config.hidden_size
+        hf_cfg = self.model_config.hf_config
+        thinker_cfg = getattr(hf_cfg, "thinker_config", None)
+        dims = {
+            Modality.IMAGE: default,
+            Modality.VIDEO: default,
+            Modality.AUDIO: default,
+        }
+
+        vision_cfg = getattr(thinker_cfg, "vision_config", None) or getattr(
+            hf_cfg, "vision_config", None
+        )
+        if vision_cfg is not None:
+            out_hs = getattr(vision_cfg, "out_hidden_size", None)
+            if out_hs is not None:
+                ds = getattr(vision_cfg, "deepstack_visual_indexes", None)
+                vis_dim = (
+                    out_hs * (1 + len(ds))
+                    if isinstance(ds, (list, tuple)) and ds
+                    else out_hs
+                )
+                dims[Modality.IMAGE] = vis_dim
+                dims[Modality.VIDEO] = vis_dim
+
+        audio_cfg = getattr(thinker_cfg, "audio_config", None) or getattr(
+            hf_cfg, "audio_config", None
+        )
+        if audio_cfg is not None:
+            for attr in ("output_dim", "d_model"):
+                val = getattr(audio_cfg, attr, None)
+                if val and int(val) > 0:
+                    dims[Modality.AUDIO] = int(val)
+                    break
+
+        logger.info(f"Global cache embedding dims: {dims}")
+        return dims
+
+    def _build_vision_config(self, mm_process_config):
+        """
+        Validate vision config, used for image/video/audio.
+        If not provided, keep default values.
+        """
+        self.vision_config = (
+            mm_process_config.get("vision_config", {})
+            if mm_process_config is not None
+            else {}
+        )
+        for modality_str in ["image", "video", "audio"]:
+            if not self.vision_config.get(modality_str, None):
+                self.vision_config[modality_str] = {}
+            if self.use_image_processor_gpu:
+                self.vision_config[modality_str]["device"] = self.device
+
+            if modality_str == "video":
+                video_defaults = {"fps": 2.0, "max_frames": 768, "min_frames": 4}
+                for k, v in video_defaults.items():
+                    self.vision_config["video"].setdefault(k, v)
+
+            if modality_str == "audio":
+                if "return_attention_mask" not in self.vision_config["audio"]:
+                    self.vision_config["audio"]["return_attention_mask"] = True
+                if "padding" not in self.vision_config["audio"]:
+                    if self.model_type == "qwen2_audio":
+                        # For Qwen2Audio, use padding="max_length"
+                        # (same as https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_audio/processing_qwen2_audio.py#L93)
+                        self.vision_config["audio"]["padding"] = "max_length"
+                    else:
+                        self.vision_config["audio"]["padding"] = True
+                if "truncation" not in self.vision_config["audio"]:
+                    # keep same logic as base_processor.py
+                    if (
+                        hasattr(self, "audio_processor")
+                        and self.audio_processor is not None
+                    ):
+                        if self.audio_processor.__class__.__name__ in {
+                            "Gemma3nProcessor",
+                            "GlmAsrProcessor",
+                            "Qwen2AudioProcessor",
+                            "Qwen3OmniMoeProcessor",
+                        }:
+                            self.vision_config["audio"]["truncation"] = False
+
+    def _load_mm_processor(self, server_args: ServerArgs):
+        """
+        Load image/video/audio processor separately,
+        avoid issues with AutoProcessor not recognizing certain models
+        """
+        from transformers import AutoImageProcessor, AutoVideoProcessor
+
+        try:
+            self.image_processor = AutoImageProcessor.from_pretrained(
+                server_args.tokenizer_path or server_args.model_path,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+                use_fast=not server_args.disable_fast_image_processor,
+            )
+        except Exception as e:
+            logger.warning(f"Failed to load image processor: {e}")
+            self.image_processor = None
+
+        try:
+            self.video_processor = AutoVideoProcessor.from_pretrained(
+                server_args.tokenizer_path or server_args.model_path,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+                use_fast=not server_args.disable_fast_image_processor,
+            )
+        except Exception as e:
+            logger.warning(f"Failed to load video processor: {e}")
+            self.video_processor = None
+
+        try:
+            # Note: AutoProcessor is used for audio processor
+            _audio_proc = AutoProcessor.from_pretrained(
+                server_args.tokenizer_path or server_args.model_path,
+                trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
+                use_fast=not server_args.disable_fast_image_processor,
+            )
+            if not hasattr(_audio_proc, "feature_extractor"):
+                logger.warning(
+                    "Loaded AutoProcessor has no feature_extractor attribute, "
+                    "audio processing will be unavailable."
+                )
+                self.audio_processor = None
+            else:
+                self.audio_processor = _audio_proc
+        except Exception as e:
+            logger.warning(f"Failed to load audio processor: {e}")
+            self.audio_processor = None
+
     def _load_single_item(
         self,
         data,
@@ -282,16 +458,69 @@ class MMEncoder:
                 task_info.append((modality, data))
         return futures, task_info
 
-    async def _flatten_and_load_images(self, mm_items):
+    def _get_feat_extract_output_lengths(self, feature_lens):
         """
-        Flatten mm_items structure, load images concurrently, and restore original structure.
+        Computes the output length of the convolutional layers and the output length of the audio encoder
+        """
+        # qwen2_audio/qwen2.5_omni
+        if self.model_type in ["qwen2_audio", "qwen2_5_omni"]:
+            input_length = (feature_lens - 1) // 2 + 1
+            return (input_length - 2) // 2 + 1
+        # qwen3_omni_moe
+        elif self.model_type == "qwen3_omni_moe":
+            input_lengths_leave = feature_lens % 100
+            feat_lengths = (input_lengths_leave - 1) // 2 + 1
+            output_lengths = (
+                ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (feature_lens // 100) * 13
+            )
+            return output_lengths
+        else:
+            # fallback to original HF audio sample logic for other models
+            logger.warning(
+                f"Fallback to original HF audio sample logic for {self.model_type}"
+            )
+            input_length = (feature_lens - 1) // 2 + 1
+            return (input_length - 2) // 2 + 1
+
+    async def _flatten_and_load_videos(self, mm_items):
+        if not isinstance(mm_items, (list, tuple)):
+            mm_items = [mm_items]
+
+        futures, _ = self.submit_data_loading_tasks(
+            mm_items, [Modality.VIDEO] * len(mm_items)
+        )
+        async_futures = [asyncio.wrap_future(f) for f in futures]
+        video_items = await asyncio.gather(*async_futures)
+
+        video_processor_kwargs = {}
+        if "qwen" in self.model_type:
+            # for qwen-series model, do sample frames before preprocess
+            video_processed = [
+                await preprocess_video(
+                    video, video_config=self.vision_config.get("video", {})
+                )
+                for video in video_items
+            ]
+            videos, video_metadata = map(list, zip(*video_processed))
+            video_processor_kwargs["do_sample_frames"] = False
+            if video_metadata:
+                video_processor_kwargs["video_metadata"] = video_metadata
+            return videos, video_processor_kwargs
+        else:
+            raise NotImplementedError(
+                f"Video processing is not supported for {self.model_type} model."
+            )
+
+    async def _flatten_and_load_data_by_modality(self, mm_items, modality):
+        """
+        Flatten mm_items structure, load multimodal data concurrently, and restore original structure.
 
         Returns:
-            Same structure as load_images would return
+            Same structure as load_mm_items would return, support for image/audio
         """
-        # Handle single image (not a list)
+        # Handle single mm_item (not a list)
         if not isinstance(mm_items, (list, tuple)):
-            futures, _ = self.submit_data_loading_tasks([mm_items], [Modality.IMAGE])
+            futures, _ = self.submit_data_loading_tasks([mm_items], [modality])
             return await asyncio.wrap_future(futures[0])
 
         # Handle nested list (list of lists)
@@ -299,14 +528,14 @@ class MMEncoder:
             # Flatten nested structure
             flat_data = []
             flat_indices = []  # Track which group each item belongs to
-            for group_idx, image_group in enumerate(mm_items):
+            for group_idx, item_group in enumerate(mm_items):
-                for item in image_group:
+                for item in item_group:
                     flat_data.append(item)
                     flat_indices.append(group_idx)
 
             # Submit all tasks concurrently
             futures, _ = self.submit_data_loading_tasks(
-                flat_data, [Modality.IMAGE] * len(flat_data)
+                flat_data, [modality] * len(flat_data)
             )
 
             # Wait for all tasks to complete asynchronously
@@ -323,123 +552,137 @@ class MMEncoder:
         # Handle simple list
         else:
             futures, _ = self.submit_data_loading_tasks(
-                mm_items, [Modality.IMAGE] * len(mm_items)
+                mm_items, [modality] * len(mm_items)
             )
             # Wait for all tasks to complete asynchronously
             async_futures = [asyncio.wrap_future(f) for f in futures]
             return await asyncio.gather(*async_futures)
 
-    def get_num_patches(self, grid: Union[torch.Tensor, List[int]]) -> int:
+    def get_num_patches(
-        """Calculate number of raw patches (before 2x2 merge). Used for pixel_values slicing."""
+        self, grid: Union[torch.Tensor, List[int]], modality: Modality
-        return int(grid[0] * grid[1] * grid[2])
+    ) -> int:
+        """Calculate number of raw patches (before merge/sampling). Used for pixel_values slicing."""
+        if modality == Modality.AUDIO:
+            return int(grid.item())
+        else:
+            return int(grid[0] * grid[1] * grid[2])
 
-    def get_num_tokens(self, grid: Union[torch.Tensor, List[int]]) -> int:
+    def get_num_tokens(
+        self, grid: Union[torch.Tensor, List[int]], modality: Modality
+    ) -> int:
         """Calculate number of tokens (after 2x2 merge). Used for mm_embedding slicing."""
-        merge_size = getattr(self.image_processor, "merge_size", 2)
+        if modality == Modality.AUDIO:
-        return self.get_num_patches(grid) // (merge_size**2)
+            input_length = self.get_num_patches(grid, modality)
+            return self._get_feat_extract_output_lengths(input_length)
+        else:
+            merge_size = getattr(self.image_processor, "merge_size", 2)
+            return self.get_num_patches(grid, modality) // (merge_size**2)
 
     def slice_embedding(
-        self, mm_embedding: torch.Tensor, grid_thw: List
+        self, mm_embedding: torch.Tensor, grid_thw: List, modality: Modality
     ) -> List[torch.Tensor]:
         """Slice a concatenated embedding tensor into individual image embeddings."""
         slices, offset = [], 0
         for grid in grid_thw:
-            count = self.get_num_tokens(grid)
+            count = self.get_num_tokens(grid, modality)
             slices.append(mm_embedding[offset : offset + count])
             offset += count
         return slices
 
     def _calculate_hashes_from_features(
-        self, pixel_values: torch.Tensor, grid_thw: List
+        self, mm_feature: torch.Tensor, grid_thw: List, modality: Modality
     ) -> List[str]:
-        """CPU Task: Compute hashes based on processed feature patches (pixel_values)."""
+        """CPU Task: Compute hashes based on processed feature patches."""
         hashes, offset = [], 0
+        logger.info(f"{mm_feature.shape=} with {modality=}")
         for grid in grid_thw:
-            num_patches = self.get_num_patches(grid)
+            num_patches = self.get_num_patches(grid, modality)
-            feature_slice = pixel_values[offset : offset + num_patches]
+            feature_slice = mm_feature[offset : offset + num_patches]
-            tmp_item = MultimodalDataItem(
+            tmp_item = MultimodalDataItem(modality=modality, feature=feature_slice)
-                modality=Modality.IMAGE, feature=feature_slice
-            )
             tmp_item.set_pad_value()
             hashes.append(tmp_item.hash)
             offset += num_patches
         return hashes
 
     async def _encode_missing(
-        self, pixel_values: torch.Tensor, images_input: dict, indices: List[int]
+        self,
+        mm_feature: torch.Tensor,
+        mm_inputs: dict,
+        indices: List[int],
+        modality: Modality = Modality.IMAGE,
+        get_feature_fn=None,
     ) -> List[torch.Tensor]:
         """
-        GPU Task: Run ViT inference ONLY on the subset of images missing from the cache.
+        GPU Task: Run ViT inference ONLY on the subset of mm items missing from the cache.
         """
-        grid_thw = images_input["image_grid_thw"]
+        grid_thw = _get_mm_grid_dim(mm_inputs, modality)
 
-        # 1. Slice pixel_values to get only the patches for missing images
+        # 1. Slice mm_feature to get only the patches for missing mm items
-        sub_pixel_list = []
+        sub_feature_list = []
         offsets = [0]
         curr = 0
         for g in grid_thw:
-            curr += self.get_num_patches(g)
+            curr += self.get_num_patches(g, modality)
             offsets.append(curr)
 
         for idx in indices:
-            sub_pixel_list.append(pixel_values[offsets[idx] : offsets[idx + 1]])
+            sub_feature_list.append(mm_feature[offsets[idx] : offsets[idx + 1]])
 
-        sub_feature = torch.cat(sub_pixel_list, dim=0)
+        sub_feature = torch.cat(sub_feature_list, dim=0)
 
         mm_item = MultimodalDataItem.from_dict(
             {
-                "modality": Modality.IMAGE,
+                "modality": modality,
                 "feature": _convert(sub_feature),
             }
         )
 
-        for k, v in images_input.items():
+        for k, v in mm_inputs.items():
-            if k == "pixel_values":
+            if k in _mm_feature_attrs.get(modality, []):
                 continue
             val = _convert(v)
-            if k in _image_grid_attrs:
+            if k in _mm_grid_attrs.get(modality, []):
                 mm_item.set(k, val[indices])
             else:
                 mm_item.set(k, val)
 
         with torch.inference_mode():
-            new_embeddings = self.model.get_image_feature([mm_item]).cpu()
+            new_embeddings = get_feature_fn([mm_item]).cpu()
             if new_embeddings.ndim != 2:
                 new_embeddings = new_embeddings.reshape(-1, new_embeddings.shape[-1])
 
         sub_grids = [grid_thw[i] for i in indices]
-        return self.slice_embedding(new_embeddings, sub_grids)
+        return self.slice_embedding(new_embeddings, sub_grids, modality)
 
     async def encode_with_global_cache(
         self,
         mm_items,
+        modality: Modality,
         req_id: str,
         num_parts: int,
         part_idx: int,
         hashes: Optional[List[str]] = None,
     ) -> torch.Tensor:
-        images = await self._flatten_and_load_images(mm_items)
+        mm_inputs, get_feature_fn = await self._process_mm_items(mm_items, modality)
-        kwargs = {"device": self.device} if self.use_image_processor_gpu else {}
+        grid_thw = _get_mm_grid_dim(mm_inputs, modality)
-        images_input = self.image_processor(images=images, **kwargs)
+        mm_feature = _convert(_get_mm_feature(mm_inputs, modality))
-        pixel_values = images_input["pixel_values"]
+        num_items = len(grid_thw)
-        grid_thw = images_input["image_grid_thw"]
-        num_images = len(grid_thw)
 
         # Step 1: Rank 0 checks global cache and broadcasts hit/miss mask to all ranks.
         if self.rank == 0:
             if hashes is None:
-                image_hashes = self._calculate_hashes_from_features(
+                mm_hashes = self._calculate_hashes_from_features(
-                    pixel_values, grid_thw
+                    mm_feature, grid_thw, modality
                 )
             else:
-                image_hashes = hashes
+                mm_hashes = hashes
-            exist_mask = await self.mm_global_cache.batch_is_exist(image_hashes)
+            exist_mask = await self.mm_global_cache.batch_is_exist(mm_hashes)
             mask_tensor = torch.tensor(
                 [1 if e else 0 for e in exist_mask], dtype=torch.int32
             )
         else:
-            image_hashes = None
+            mm_hashes = None
-            mask_tensor = torch.zeros(num_images, dtype=torch.int32)
+            mask_tensor = torch.zeros(num_items, dtype=torch.int32)
 
         if self.server_args.tp_size > 1:
             torch.distributed.broadcast(
@@ -456,7 +699,7 @@ class MMEncoder:
         new_slices = []
         if missing_indices:
             new_slices = await self._encode_missing(
-                pixel_values, images_input, missing_indices
+                mm_feature, mm_inputs, missing_indices, modality, get_feature_fn
             )
 
         # Step 3: Rank 0 prefetches cache-hit embeddings from global cache.
@@ -464,9 +707,11 @@ class MMEncoder:
 
         if self.rank == 0:
             if hit_indices:
-                hit_hashes = [image_hashes[i] for i in hit_indices]
+                hit_hashes = [mm_hashes[i] for i in hit_indices]
-                hit_tokens = [self.get_num_tokens(grid_thw[i]) for i in hit_indices]
+                hit_tokens = [
-                self.mm_global_cache.prefetch(req_id, hit_hashes, hit_tokens)
+                    self.get_num_tokens(grid_thw[i], modality) for i in hit_indices
+                ]
+                self.mm_global_cache.prefetch(req_id, hit_hashes, hit_tokens, modality)
 
                 try:
 
@@ -478,7 +723,7 @@ class MMEncoder:
                 except (asyncio.TimeoutError, Exception) as e:
                     logger.error(
                         f"Prefetch failed for req {req_id}: {e}. "
-                        f"Falling back to ViT for {len(hit_indices)} hit images."
+                        f"Falling back to ViT for {len(hit_indices)} hit items."
                     )
                     prefetch_status[0] = 0
 
@@ -490,21 +735,21 @@ class MMEncoder:
                 group=self.mm_global_cache.prefetch_tp_group,
             )
 
-        # Step 5: If prefetch failed, all ranks fallback to ViT for the hit images.
+        # Step 5: If prefetch failed, all ranks fallback to ViT for the hit mm items.
         if prefetch_status.item() == 0 and hit_indices:
             logger.info(
                 f"Req {req_id}: Prefetch failed, all ranks running ViT fallback "
-                f"for {len(hit_indices)} images."
+                f"for {len(hit_indices)} mm items."
             )
             fallback_slices = await self._encode_missing(
-                pixel_values, images_input, hit_indices
+                mm_feature, mm_inputs, hit_indices, modality, get_feature_fn
             )
         else:
             fallback_slices = None
 
         # Step 6: Rank 0 assembles final embedding and prepares for sending.
         if self.rank == 0:
-            final_slices = [None] * num_images
+            final_slices = [None] * num_items
 
             for i, idx in enumerate(missing_indices):
                 final_slices[idx] = new_slices[i]
@@ -512,7 +757,7 @@ class MMEncoder:
             # Fill in cache-hit embeddings (from prefetch or fallback)
             if prefetch_status.item() == 1 and hit_indices:
                 cached_slices = self.mm_global_cache.get_embeddings(
-                    [image_hashes[i] for i in hit_indices]
+                    [mm_hashes[i] for i in hit_indices]
                 )
                 for i, idx in enumerate(hit_indices):
                     final_slices[idx] = cached_slices[i]
@@ -524,10 +769,10 @@ class MMEncoder:
 
             # Background insert: store newly computed embeddings into global cache.
             # Includes both original misses and fallback-recomputed hits.
-            all_new_hashes = [image_hashes[i] for i in missing_indices]
+            all_new_hashes = [mm_hashes[i] for i in missing_indices]
             all_new_slices = list(new_slices)
             if fallback_slices is not None:
-                all_new_hashes += [image_hashes[i] for i in hit_indices]
+                all_new_hashes += [mm_hashes[i] for i in hit_indices]
                 all_new_slices += list(fallback_slices)
 
             if all_new_hashes:
@@ -543,8 +788,15 @@ class MMEncoder:
                 self.background_tasks.add(task)
                 task.add_done_callback(self.background_tasks.discard)
 
+            aux_data = _build_mm_aux_data(mm_inputs)
             self.embedding_to_send[req_id] = EmbeddingData(
-                req_id, num_parts, part_idx, grid_thw, mm_embedding
+                req_id,
+                num_parts,
+                part_idx,
+                grid_thw,
+                modality,
+                mm_embedding,
+                **aux_data,
             )
             return (
                 mm_embedding.nbytes,
@@ -556,31 +808,152 @@ class MMEncoder:
         else:
             return (0, 0, 0, None, None)
 
-    async def _encode(self, mm_items) -> torch.Tensor:
+    async def _flatten_and_load_audios(self, mm_items):
-        try:
+        """
-            images = await self._flatten_and_load_images(mm_items)
+        Flatten mm_items structure, load audios concurrently, and restore original structure.
-        except Exception as e:
+        """
-            raise BadRequestError(f"Failed to load images from input: {str(e)}")
+        return await self._flatten_and_load_data_by_modality(mm_items, Modality.AUDIO)
 
-        try:
+    async def _flatten_and_load_images(self, mm_items):
-            kwargs = {"device": self.device} if self.use_image_processor_gpu else {}
+        """
-            images_input = self.image_processor(images=images, **kwargs)
+        Flatten mm_items structure, load images concurrently, and restore original structure.
-            feature = images_input["pixel_values"]
+        """
-            mm_item = MultimodalDataItem.from_dict(
+        return await self._flatten_and_load_data_by_modality(mm_items, Modality.IMAGE)
-                {
+
-                    "modality": Modality.IMAGE,
+    def _calculate_timestamps(self, indices, video_fps: float, merge_size: int = 2):
-                    "feature": _convert(feature),
+        """Calculate timestamps for video frames, used for qwen3_vl models."""
-                }
+        # refer to https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen3_vl/processing_qwen3_vl.py#L255
+        if not isinstance(indices, list):
+            indices = indices.tolist()
+        if len(indices) % merge_size != 0:
+            indices.extend(
+                indices[-1] for _ in range(merge_size - len(indices) % merge_size)
             )
-            for k, v in images_input.items():
+        timestamps = [idx / video_fps for idx in indices]
-                if k == "pixel_values":
+        # Frames are merged by merge_size, so we need to average the timestamps
-                    continue
+        # between the first/last frame within the temporal patch
-                mm_item.set(k, _convert(v))
+        timestamps = [
+            (timestamps[i] + timestamps[i + merge_size - 1]) / 2
+            for i in range(0, len(timestamps), merge_size)
+        ]
+        return timestamps
 
+    async def _process_mm_items(self, mm_items, modality):
+        if modality == Modality.IMAGE and self.image_processor:
+            images = await self._flatten_and_load_images(mm_items)
+            image_config = self.vision_config.get("image", {})
+            processor_input = self.image_processor(images=images, **image_config)
+            feature = processor_input["pixel_values"]
+            if hasattr(self.model, "thinker"):  # for omni models
+                get_feature_method = self.model.thinker.get_image_feature
+            else:
+                get_feature_method = self.model.get_image_feature
+        elif modality == Modality.VIDEO and self.video_processor:
+            videos, video_processor_kwargs = await self._flatten_and_load_videos(
+                mm_items
+            )
+            processor_input = self.video_processor(
+                videos=videos, **video_processor_kwargs
+            )
+            # Get additional video metadata
+            if (
+                self.model_type in ["qwen3_vl", "qwen3_vl_moe"]
+                and video_processor_kwargs.get("video_metadata", None) is not None
+            ):
+                # For qwen3-vl models, we need to store the video timestamps
+                video_metadata = video_processor_kwargs["video_metadata"]
+                try:
+                    merge_size = (
+                        self.model_config.hf_config.vision_config.spatial_merge_size
+                    )
+                except (AttributeError, KeyError):
+                    merge_size = 2  # Default merge_size
+
+                video_timestamps = []
+                for metadata in video_metadata:
+                    video_fps = metadata.get("fps", None) or 24  # original video fps
+                    frames_indices = metadata.get("frames_indices", None)
+                    timestamps = self._calculate_timestamps(
+                        frames_indices, video_fps, merge_size
+                    )
+                    video_timestamps.append(timestamps)
+                processor_input["video_timestamps"] = video_timestamps
+            elif (
+                self.model_type in ["qwen2_5_vl", "qwen2_5_omni", "qwen3_omni_moe"]
+                and processor_input.get("video_grid_thw", None) is not None
+            ):
+                # For omni/qwen2_5_vl models, calculate second_per_grid_ts for rotary embedding
+                video_grid_thw = processor_input["video_grid_thw"]
+                try:
+                    temporal_patch_size = self.video_processor.temporal_patch_size
+                except AttributeError:
+                    temporal_patch_size = 2  # Default temporal_patch_size
+                # get sampled fps, default: 2
+                fps_list = [
+                    self.vision_config.get("video", {}).get("fps", None) or 2
+                ] * len(video_grid_thw)
+                second_per_grid_ts = [(temporal_patch_size / fps) for fps in fps_list]
+                second_per_grid_ts_tensor = torch.tensor(
+                    second_per_grid_ts, dtype=torch.float32
+                )
+                processor_input["second_per_grid_ts"] = second_per_grid_ts_tensor
+
+            feature = processor_input["pixel_values_videos"]
+            if hasattr(self.model, "thinker"):  # for omni models
+                get_feature_method = self.model.thinker.get_video_feature
+            else:
+                get_feature_method = self.model.get_video_feature
+        elif modality == Modality.AUDIO and self.audio_processor:
+            audios = await self._flatten_and_load_audios(mm_items)
+            audio_config = self.vision_config.get("audio", {})
+            processor_input = self.audio_processor.feature_extractor(
+                audios, **audio_config
+            )
+            processor_input["feature_attention_mask"] = processor_input.pop(
+                "attention_mask"
+            )
+            # convert to same format as image/video
+            input_lengths = torch.tensor(
+                processor_input["feature_attention_mask"].sum(-1), dtype=torch.long
+            )
+            processor_input["audio_feature_lens_raw"] = input_lengths
+            output_lengths = self._get_feat_extract_output_lengths(input_lengths)
+            processor_input["audio_feature_lens"] = output_lengths
+            feature = processor_input["input_features"]
+            if hasattr(self.model, "thinker"):  # for omni models
+                get_feature_method = self.model.thinker.get_audio_feature
+            else:
+                get_feature_method = self.model.get_audio_feature
+        else:
+            raise ValueError(
+                f"Currently only support image, video and audio modalities, {modality} modality has no processor available."
+            )
+
+        return processor_input, get_feature_method
+
+    async def _encode(self, mm_items, modality: Modality) -> torch.Tensor:
+        try:
+            mm_inputs, get_feature_fn = await self._process_mm_items(mm_items, modality)
+        except NotImplementedError as e:
+            raise InternalError(f"Not implemented error: {str(e)}")
+        except Exception as e:
+            raise BadRequestError(f"Failed to process mm items: {str(e)}")
+        try:
             # support mm_cache
             mm_embedding = None
             mm_hash = None
 
+            mm_item = MultimodalDataItem.from_dict(
+                {
+                    "modality": modality,
+                    "feature": _convert(_get_mm_feature(mm_inputs, modality)),
+                }
+            )
+            for k, v in mm_inputs.items():
+                if k in _mm_feature_attrs[modality]:
+                    continue
+                mm_item.set(k, _convert(v))
+
             if self.server_args.enable_prefix_mm_cache:
                 mm_item.set_pad_value()
                 mm_hash = MultiModalStaticCache.combine_hashes([mm_item.hash])
@@ -591,7 +964,7 @@ class MMEncoder:
 
             if mm_embedding is None:
                 with torch.inference_mode():
-                    mm_embedding: torch.Tensor = self.model.get_image_feature([mm_item])
+                    mm_embedding: torch.Tensor = get_feature_fn([mm_item])
                     mm_embedding = mm_embedding.cpu()
                 if len(mm_embedding.shape) != 2:
                     mm_embedding = mm_embedding.reshape(-1, mm_embedding.shape[-1])
@@ -602,7 +975,8 @@ class MMEncoder:
             if self.profiler is not None:
                 self.profiler.step()
 
-            return _get_image_grid_dim(images_input), mm_embedding
+            aux_data = _build_mm_aux_data(mm_inputs)
+            return _get_mm_grid_dim(mm_inputs, modality), mm_embedding, aux_data
         except BadRequestError as e:
             raise BadRequestError(f"Bad request error: {str(e)}")
         except Exception as e:
@@ -626,7 +1000,6 @@ class MMEncoder:
             self.engine.deregister(embedding.data_ptr())
 
             mm_data.embedding = None
-            mm_data.embedding_list[mm_data.part_idx] = None
 
         # Send ack/data
         endpoint = (
@@ -665,16 +1038,19 @@ class MMEncoder:
 
         await asyncio.get_event_loop().run_in_executor(self.executor, send_with_socket)
 
-    async def encode_with_hash(self, mm_items, req_id, num_parts, part_idx, hashes):
+    async def encode(self, mm_items, modality: Modality, req_id, num_parts, part_idx):
-        images = await self._flatten_and_load_images(mm_items)
-
-    async def encode(self, mm_items, req_id, num_parts, part_idx):
         try:
-            image_grid_dim, mm_embedding = await self._encode(mm_items)
+            grid_dim, mm_embedding, aux_data = await self._encode(mm_items, modality)
 
             if self.rank == 0:
                 mm_data = EmbeddingData(
-                    req_id, num_parts, part_idx, image_grid_dim, mm_embedding
+                    req_id,
+                    num_parts,
+                    part_idx,
+                    grid_dim,
+                    modality,
+                    mm_embedding,
+                    **aux_data,
                 )
                 self.embedding_to_send[req_id] = mm_data
             return (
@@ -694,6 +1070,7 @@ class MMEncoder:
                     num_parts,
                     part_idx,
                     None,
+                    modality,
                     error_msg=error_msg,
                     error_code=error_code,
                 )
@@ -894,6 +1271,7 @@ async def run_encoder(
             if encoder.mm_global_cache is not None:
                 await encoder.encode_with_global_cache(
                     mm_items=request["mm_items"],
+                    modality=Modality.from_str(request["modality"]),
                     req_id=request["req_id"],
                     num_parts=request["num_parts"],
                     part_idx=request["part_idx"],
@@ -902,6 +1280,7 @@ async def run_encoder(
             else:
                 await encoder.encode(
                     mm_items=request["mm_items"],
+                    modality=Modality.from_str(request["modality"]),
                     req_id=request["req_id"],
                     num_parts=request["num_parts"],
                     part_idx=request["part_idx"],
@@ -966,6 +1345,7 @@ async def handle_encode_request(request: dict):
             nbytes, embedding_len, embedding_dim, error_msg, error_code = (
                 await encoder.encode_with_global_cache(
                     mm_items=request["mm_items"],
+                    modality=Modality.from_str(request["modality"]),
                     req_id=request["req_id"],
                     num_parts=request["num_parts"],
                     part_idx=request["part_idx"],
@@ -976,6 +1356,7 @@ async def handle_encode_request(request: dict):
             nbytes, embedding_len, embedding_dim, error_msg, error_code = (
                 await encoder.encode(
                     mm_items=request["mm_items"],
+                    modality=Modality.from_str(request["modality"]),
                     req_id=request["req_id"],
                     num_parts=request["num_parts"],
                     part_idx=request["part_idx"],

python/sglang/srt/managers/io_struct.py

@@ -28,7 +28,7 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
 import torch
 
 from sglang.srt.lora.lora_registry import LoRARef
-from sglang.srt.managers.schedule_batch import BaseFinishReason
+from sglang.srt.managers.schedule_batch import BaseFinishReason, Modality
 from sglang.srt.multimodal.mm_utils import has_valid_data
 from sglang.srt.observability.req_time_stats import (
     APIServerReqTimeStats,
@@ -226,8 +226,8 @@ class GenerateReqInput(BaseReq):
     received_time: Optional[float] = None
 
     # For EPD-disaggregated inference
-    need_wait_for_image: Optional[bool] = None
+    need_wait_for_mm_inputs: Optional[bool] = None
-    num_items_assigned: Optional[List] = None
+    num_items_assigned: Optional[Dict[Modality, List[int]]] = None
 
     # Multimodal tiling controls (extensions)
     max_dynamic_patch: Optional[int] = None
@@ -731,8 +731,8 @@ class TokenizedGenerateReqInput(BaseReq):
     # Whether to return entropy
     return_entropy: bool = False
 
-    need_wait_for_image: bool = False
+    need_wait_for_mm_inputs: bool = False
-    num_items_assigned: Optional[List] = None
+    num_items_assigned: Optional[Dict[Modality, List[int]]] = None
 
     # For observability
     time_stats: Optional[Union[APIServerReqTimeStats, DPControllerReqTimeStats]] = None

python/sglang/srt/managers/mm_utils.py

@@ -427,6 +427,7 @@ def get_embedding_chunk(
 
 def _get_precomputed_embedding(
     items: List[MultimodalDataItem],
+    items_size: List[int],
     prefix_length: List[int],
     extend_length: List[int],
     items_offset_list: List[List[Tuple[int, int]]],
@@ -437,38 +438,32 @@ def _get_precomputed_embedding(
     If none have precomputed_embeddings, return None.
     """
     precomputed_embeddings = []
-    for idx, item in enumerate(items):
+    max_iterations = min(len(items_size) - 1, len(prefix_length))
-        if item.precomputed_embeddings is None:
+
-            precomputed_embeddings.append(None)
+    for i in range(max_iterations):
+        if items_size[i] == items_size[i + 1]:
             continue
-        seq_start_idx = prefix_length[idx]
+
-        seq_end_idx = seq_start_idx + extend_length[idx] - 1
+        items_per_req = items[items_size[i] : items_size[i + 1]]
-        prefix_embedding_length = []
+        extend_len = extend_length[i] if i < len(extend_length) else 0
-        extend_embedding_length = []
+        items_offset = items_offset_list[i]
-        for mm_start_idx, mm_end_idx in items_offset_list[idx]:
+
-            if mm_start_idx > seq_end_idx:
+        if any(item.precomputed_embeddings is None for item in items_per_req):
-                break
+            chunk = None
-            if seq_start_idx > mm_start_idx:
+        else:
-                prefix_embedding_length.append(
+            req_embeddings = torch.concat(
-                    min(seq_start_idx - mm_start_idx, mm_end_idx - mm_start_idx + 1)
+                [item.precomputed_embeddings for item in items_per_req]
-                )
+            )
-            if mm_end_idx >= seq_start_idx:
+            chunk, _, _ = get_embedding_chunk(
-                extend_embedding_length.append(
+                embedding=req_embeddings,
-                    min(
+                extend_prefix_len=prefix_length[i],
-                        mm_end_idx - seq_start_idx + 1,
+                extend_seq_len=extend_len,
-                        seq_end_idx - mm_start_idx + 1,
+                items_offset=items_offset,
-                        mm_end_idx - mm_start_idx + 1,
+            )
-                        seq_end_idx - seq_start_idx + 1,
+
-                    )
+        if chunk is None and len(items_per_req) > 1:
-                )
+            return None
-        prefix_embedding_length = int(np.sum(prefix_embedding_length))
+        precomputed_embeddings.append(chunk)
-        extend_embedding_length = int(np.sum(extend_embedding_length))
-        precomputed_embeddings.append(
-            item.precomputed_embeddings[
-                prefix_embedding_length : prefix_embedding_length
-                + extend_embedding_length
-            ]
-        )
 
     if any(feature is not None for feature in precomputed_embeddings):
         if not all(feature is not None for feature in precomputed_embeddings):
@@ -886,7 +881,7 @@ def get_embedding_and_mask(
     """
     # 1. Get embedding
     embedding = _get_precomputed_embedding(
-        embedding_items, prefix_length, extend_length, items_offset_list
+        embedding_items, items_size, prefix_length, extend_length, items_offset_list
     )
     if embedding is None:
         embedding, input_ids = _get_chunked_prefill_embedding(

python/sglang/srt/managers/tokenizer_manager.py

@@ -724,10 +724,10 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
             ):
                 if self.server_args.language_only:
                     mm_inputs = await self.mm_receiver.recv_mm_data(
-                        img_data=obj.image_data,
+                        request_obj=obj,
                         mm_processor=self.mm_processor,
                         prompt=(input_text or input_ids),
-                        need_wait_for_image=obj.need_wait_for_image,
+                        need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs,
                     )
                 if mm_inputs is None:
                     mm_inputs: Dict = await self.mm_data_processor.process(
@@ -740,7 +740,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
             elif (
                 self.server_args.language_only
                 and self.server_args.encoder_transfer_backend == "zmq_to_scheduler"
-                and not obj.need_wait_for_image
+                and not obj.need_wait_for_mm_inputs
             ):
                 # In language_only mode with zmq_to_scheduler, if we didn't dispatch
                 # to encoder (e.g., only one image), process locally like non-language_only mode
@@ -982,7 +982,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
                 priority=obj.priority,
                 extra_key=obj.extra_key,
                 routing_key=obj.routing_key,
-                need_wait_for_image=obj.need_wait_for_image,
+                need_wait_for_mm_inputs=obj.need_wait_for_mm_inputs,
                 num_items_assigned=obj.num_items_assigned,
             )
         elif isinstance(obj, EmbeddingReqInput):
@@ -2362,14 +2362,14 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
             if self.server_args.enable_adaptive_dispatch_to_encoder:
                 should_dispatch = self._should_dispatch_to_encoder(obj)
 
-            # Set need_wait_for_image flag based on whether we dispatch to encoder
+            # Set need_wait_for_mm_inputs flag based on whether we dispatch to encoder
             # This flag will be used in _tokenize_one_request to determine processing path
             if should_dispatch:
-                obj.need_wait_for_image = True
+                obj.need_wait_for_mm_inputs = True
                 if self.server_args.encoder_transfer_backend == "zmq_to_scheduler":
                     self.mm_receiver.send_encode_request(obj)
             else:
-                obj.need_wait_for_image = False
+                obj.need_wait_for_mm_inputs = False
 
     def convert_to_span_attrs(
         self,

python/sglang/srt/mem_cache/storage/mooncake_store/embedding_cache_controller.py

@@ -95,14 +95,14 @@ class EmbeddingCacheController:
         tp_rank,
         tp_size,
         max_pool_size_gb=4.0,
-        hidden_dim=1024,
+        hidden_dims: dict = None,
         tp_group=None,
         all_rank_get=False,
     ):
         self.tp_world_size = tp_size
         self.tp_group = tp_group
         self.all_rank_get = all_rank_get
-        self.hidden_dim = hidden_dim
+        self.hidden_dims = hidden_dims or {}
         self.element_size = torch.float32.itemsize
 
         # 1. Mooncake Backend & Pinned Buffer
@@ -115,7 +115,8 @@ class EmbeddingCacheController:
 
         # 2. Variable Size Memory Management
         self.allocator = ContiguousMemoryAllocator(self.total_pool_size_bytes)
-        self.hash_to_metadata = {}  # {image_hash: (offset, num_tokens, size_bytes)}
+        # {hash: (offset, num_tokens, embedding_dim, size_bytes)}
+        self.hash_to_metadata = {}
 
         # 3. Task Tracking
         self.ongoing_prefetch = {}  # {req_id: EmbeddingPrefetchOperation}
@@ -142,9 +143,19 @@ class EmbeddingCacheController:
             self.prefetch_tp_group = None
 
     def prefetch(
-        self, req_id: str, image_hashes: List[str], expected_tokens: List[int]
+        self,
+        req_id: str,
+        image_hashes: List[str],
+        expected_tokens: List[int],
+        modality=None,
     ):
         """Issues ONE batch GET for all missing images in the request."""
+        dim = self.hidden_dims.get(modality) if modality is not None else None
+        if not dim:
+            logger.warning(
+                f"Req {req_id}: Unknown dim for modality={modality}, skipping prefetch (will fallback to ViT)."
+            )
+            return
         keys, ptrs, sizes = [], [], []
 
         with self.lock:
@@ -155,12 +166,12 @@ class EmbeddingCacheController:
                     )
                     continue
 
-                size_bytes = num_tokens * self.hidden_dim * self.element_size
+                size_bytes = num_tokens * dim * self.element_size
                 offset = self.allocator.allocate(size_bytes)
                 if offset is None:
                     continue
 
-                self.hash_to_metadata[h] = (offset, num_tokens, size_bytes)
+                self.hash_to_metadata[h] = (offset, num_tokens, dim, size_bytes)
                 keys.append(h)
                 ptrs.append(self.cpu_pool.data_ptr() + offset)
                 sizes.append(size_bytes)
@@ -187,20 +198,20 @@ class EmbeddingCacheController:
                 if h in self.hash_to_metadata:
                     continue
 
-                num_tokens = tensor.shape[0]
+                num_tokens, dim = tensor.shape[0], tensor.shape[1]
-                size_bytes = num_tokens * self.hidden_dim * self.element_size
+                size_bytes = num_tokens * dim * self.element_size
                 offset = self.allocator.allocate(size_bytes)
                 if offset is None:
                     continue
 
                 # Copy to pinned pool for RDMA
-                self.hash_to_metadata[h] = (offset, num_tokens, size_bytes)
                 target_view = (
                     self.cpu_pool[offset : offset + size_bytes]
                     .view(torch.float32)
-                    .view(num_tokens, self.hidden_dim)
+                    .view(num_tokens, dim)
                 )
                 target_view.copy_(tensor.cpu())
+                self.hash_to_metadata[h] = (offset, num_tokens, dim, size_bytes)
 
                 keys.append(h)
                 ptrs.append(self.cpu_pool.data_ptr() + offset)
@@ -277,11 +288,11 @@ class EmbeddingCacheController:
         with self.lock:
             tensors = []
             for h in image_hashes:
-                offset, num_tokens, size_bytes = self.hash_to_metadata[h]
+                offset, num_tokens, dim, size_bytes = self.hash_to_metadata[h]
                 tensors.append(
                     self.cpu_pool[offset : offset + size_bytes]
                     .view(torch.float32)
-                    .view(num_tokens, self.hidden_dim)
+                    .view(num_tokens, dim)
                 )
             return tensors
 

python/sglang/srt/multimodal/processors/base_processor.py

@@ -256,15 +256,20 @@ class BaseMultimodalProcessor(ABC):
     def spatial_merge_size(self):
         return self.hf_config.vision_config.spatial_merge_size
 
-    def build_input_ids(self, prompt, img_grid_thw):
+    def build_input_ids(
+        self, prompt, img_grid_thw=None, video_grid_thw=None, audio_seq_lens=None
+    ):
         """
-        Use prompt and img_grid_thw to build input_ids
+        Use prompt, img_grid_thw, video_grid_thw, and audio_seq_lens to build input_ids.
+        Supports image, video, and audio tokens.
         """
         if not isinstance(prompt, list):
             prompt = self._tokenizer.encode(prompt)
 
-        img_token_id = self.IM_TOKEN_ID
+        img_token_id = getattr(self, "IM_TOKEN_ID", None)
-        spatial_merge_size = self.spatial_merge_size
+        video_token_id = getattr(self, "VIDEO_TOKEN_ID", None)
+        audio_token_id = getattr(self, "audio_token_id", None)
+        spatial_merge_size = getattr(self, "spatial_merge_size", 1)
 
         input_ids = []
         offsets = []
@@ -273,34 +278,81 @@ class BaseMultimodalProcessor(ABC):
 
         # Use img_token_id instead of im_start_id, because a dummy im_start_id
         # may be generated by the tokenizer.
-        img_start_indices = list(
+        vision_start_indices = []
-            filter(lambda i: prompt[i + 1] == img_token_id, range(len(prompt) - 1))
+        for i in range(len(prompt) - 1):
-        )
+            if img_token_id is not None and prompt[i + 1] == img_token_id:
+                vision_start_indices.append((i, Modality.IMAGE))
+            elif video_token_id is not None and prompt[i + 1] == video_token_id:
+                vision_start_indices.append((i, Modality.VIDEO))
+            elif audio_token_id is not None and prompt[i + 1] == audio_token_id:
+                vision_start_indices.append((i, Modality.AUDIO))
+        # get modality list with order preserved
+        modality_list = [modality for _, modality in vision_start_indices]
 
-        for cur_img_idx, img_start_idx in enumerate(img_start_indices):
+        img_idx = 0
-            assert cur_idx <= img_start_idx
+        video_idx = 0
-            # include img_start_id
+        audio_idx = 0
-            input_ids.extend(prompt[cur_idx : img_start_idx + 1])
+        for mm_start_idx, modality in vision_start_indices:
-            img_offset_start = len(input_ids)
+            if modality == Modality.IMAGE:
-            img_token_num = img_grid_thw[cur_img_idx].prod() // (spatial_merge_size**2)
+                mm_token_num = img_grid_thw[img_idx].prod() // (spatial_merge_size**2)
-            input_ids.extend([img_token_id] * img_token_num)
+                mm_token_id = img_token_id
-            # jump to img_end_id
+                img_idx += 1
-            cur_idx = img_start_idx + 2
+            elif modality == Modality.VIDEO:
-            offsets.append((img_offset_start, len(input_ids) - 1))
+                mm_token_num = video_grid_thw[video_idx].prod() // (
+                    spatial_merge_size**2
+                )
+                mm_token_id = video_token_id
+                video_idx += 1
+            elif modality == Modality.AUDIO:
+                mm_token_num = int(audio_seq_lens[audio_idx].item())
+                mm_token_id = audio_token_id
+                audio_idx += 1
+            else:
+                raise ValueError(f"Invalid modality: {modality}")
+            assert cur_idx <= mm_start_idx
+
+            input_ids.extend(prompt[cur_idx : mm_start_idx + 1])
+            mm_offset_start = len(input_ids)
+            input_ids.extend([mm_token_id] * mm_token_num)
+            cur_idx = (
+                mm_start_idx + 2
+            )  # jump to img_end_id, video_end_id, or audio_end_id
+            offsets.append((mm_offset_start, len(input_ids) - 1))
         else:
             input_ids.extend(prompt[cur_idx:])
 
-        return input_ids, offsets
+        return input_ids, offsets, modality_list
 
-    def get_mm_data(self, prompt, embeddings, img_grid_thw):
+    def get_mm_data(self, prompt, embeddings, **kwargs):
-        input_ids, offsets = self.build_input_ids(prompt, img_grid_thw)
+        img_grid_thw = kwargs.get("img_grid_thw", None)
-        mm_items = [
+        video_grid_thw = kwargs.get("video_grid_thw", None)
-            MultimodalDataItem(
+        audio_feature_lens = kwargs.get("audio_feature_lens", None)
-                modality=Modality.IMAGE,
+
-                offsets=offsets,
+        input_ids, offsets, modality_list = self.build_input_ids(
-                precomputed_embeddings=embeddings,
+            prompt,
+            img_grid_thw=img_grid_thw,
+            video_grid_thw=video_grid_thw,
+            audio_seq_lens=audio_feature_lens,
+        )
+        assert all(isinstance(modality, Modality) for modality in modality_list)
+
+        mm_items = []
+        consumed_per_modality = {}
+
+        for modality, offset in zip(modality_list, offsets):
+            num_tokens = offset[1] - offset[0] + 1
+            embedding_start = consumed_per_modality.get(modality, 0)
+            embedding_slice = embeddings[modality][
+                embedding_start : embedding_start + num_tokens
+            ]
+            consumed_per_modality[modality] = embedding_start + num_tokens
+            mm_items.append(
+                MultimodalDataItem(
+                    modality=modality,
+                    offsets=offset,
+                    precomputed_embeddings=embedding_slice,
+                )
             )
-        ]
 
         return {
             "input_ids": input_ids,
@@ -308,6 +360,7 @@ class BaseMultimodalProcessor(ABC):
             "im_start_id": self.IM_START_TOKEN_ID,
             "im_end_id": self.IM_END_TOKEN_ID,
             "im_token_id": self.IM_TOKEN_ID,
+            "video_token_id": getattr(self, "VIDEO_TOKEN_ID", None),
         }
 
     def process_mm_data(

python/sglang/srt/multimodal/processors/qwen_audio.py

@@ -1,6 +1,6 @@
 import re
 
-from sglang.srt.managers.schedule_batch import Modality
+from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem
 from sglang.srt.models.qwen2_audio import Qwen2AudioForConditionalGeneration
 from sglang.srt.multimodal.processors.base_processor import (
     BaseMultimodalProcessor,
@@ -31,6 +31,52 @@ class Qwen2AudioMultimodalProcessor(BaseMultimodalProcessor):
 
         self.ATTR_NAME_TO_MODALITY.update({"feature_attention_mask": Modality.AUDIO})
 
+    def get_mm_data(self, prompt, embeddings, **kwargs):
+        audio_feature_lens = kwargs.get("audio_feature_lens", None)
+
+        # Convert audio_feature_lens to token counts for build_input_ids
+        output_lengths = None
+        input_lengths = None
+        if audio_feature_lens is not None:
+            if audio_feature_lens.dim() > 1:
+                audio_feature_lens = audio_feature_lens.flatten()
+            input_lengths = (audio_feature_lens - 1) // 2 + 1
+            output_lengths = (input_lengths - 2) // 2 + 1
+
+        input_ids, offsets, modality_list = self.build_input_ids(
+            prompt,
+            audio_seq_lens=output_lengths,
+        )
+
+        mm_items = []
+        consumed_per_modality = {}
+
+        for modality, offset in zip(modality_list, offsets):
+            num_tokens = offset[1] - offset[0] + 1
+            embedding_start = consumed_per_modality.get(modality, 0)
+            embedding_slice = embeddings[modality][
+                embedding_start : embedding_start + num_tokens
+            ]
+            consumed_per_modality[modality] = embedding_start + num_tokens
+            mm_items.append(
+                MultimodalDataItem(
+                    modality=modality,
+                    offsets=offset,
+                    precomputed_embeddings=embedding_slice,
+                )
+            )
+
+        if mm_items:
+            mm_items[0].audio_feature_lens = output_lengths
+
+        return {
+            "mm_items": mm_items,
+            "input_ids": input_ids,
+            "audio_start_id": self.audio_start_id,
+            "audio_token_id": self.audio_token_id,
+            "audio_end_id": self.audio_end_id,
+        }
+
     async def process_mm_data_async(
         self,
         audio_data,

python/sglang/srt/multimodal/processors/qwen_vl.py

@@ -254,6 +254,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
         self.IM_START_TOKEN_ID = hf_config.vision_start_token_id
         self.IM_END_TOKEN_ID = hf_config.vision_end_token_id
         self.IM_TOKEN_ID = hf_config.image_token_id
+        self.VIDEO_TOKEN_ID = hf_config.video_token_id
 
         self.vision_start_token_id = hf_config.vision_start_token_id
         self.vision_end_token_id = getattr(hf_config, "vision_end_token_id", None)
@@ -271,12 +272,138 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
             image_token_regex=re.compile(
                 r"<\|vision_start\|>(?:<\|image_pad\|>)+<\|vision_end\|>"
             ),
-            video_token_id=hf_config.video_token_id,
+            video_token_id=self.VIDEO_TOKEN_ID,
             audio_token_id=self.audio_token_id,
         ).build(_processor)
 
-    def get_mm_data(self, prompt, embeddings, img_grid_thw):
+    def build_input_ids_with_timestamps(
-        input_ids, offsets = self.build_input_ids(prompt, img_grid_thw)
+        self, prompt, embeddings, img_grid_thw, video_grid_thw, video_timestamps
+    ):
+        """
+        Build input_ids with timestamps for qwen3_vl models.
+        """
+        if not isinstance(prompt, list):
+            prompt = self._processor.tokenizer.encode(prompt)
+
+        img_token_id = getattr(self, "IM_TOKEN_ID", None)
+        video_token_id = getattr(self, "VIDEO_TOKEN_ID", None)
+        audio_token_id = getattr(self, "audio_token_id", None)
+        spatial_merge_size = getattr(self, "spatial_merge_size", 1)
+        vision_start_token_id = getattr(self, "vision_start_token_id", None)
+        vision_end_token_id = getattr(self, "vision_end_token_id", None)
+
+        input_ids = []
+        offsets = []
+        modality_list = []
+        cur_idx = 0
+
+        vision_start_indices = []
+        for i in range(len(prompt) - 1):
+            if img_token_id is not None and prompt[i + 1] == img_token_id:
+                vision_start_indices.append((i, Modality.IMAGE))
+            elif video_token_id is not None and prompt[i + 1] == video_token_id:
+                vision_start_indices.append((i, Modality.VIDEO))
+
+        img_idx = 0
+        video_idx = 0
+        model_type = getattr(self, "model_type", None)
+        for mm_start_idx, modality in vision_start_indices:
+            modality_list.append(modality)
+            video_tokens = None
+            if modality == Modality.IMAGE:
+                mm_token_num = img_grid_thw[img_idx].prod() // (spatial_merge_size**2)
+                mm_token_id = img_token_id
+                img_idx += 1
+            elif modality == Modality.VIDEO:
+                curr_timestamps = video_timestamps[video_idx]
+                num_frames = video_grid_thw[video_idx][0]
+                frame_seqlen = video_grid_thw[video_idx][1:].prod().item() // (
+                    spatial_merge_size**2
+                )
+                video_tokens = []
+                _current_offset = len(input_ids) + mm_start_idx + 1 - cur_idx
+                # take single frame as one mm_item
+                for frame_idx in range(num_frames):
+                    if frame_idx > 0:
+                        modality_list.append(Modality.VIDEO)
+                    curr_time = curr_timestamps[frame_idx]
+                    timestamp_text = f"<{curr_time:.1f} seconds>"
+                    timestamp_tokens = self._processor.tokenizer.encode(
+                        timestamp_text, add_special_tokens=False
+                    )
+                    video_tokens.extend(timestamp_tokens)
+                    _current_offset += len(timestamp_tokens)
+                    if vision_start_token_id is not None:
+                        video_tokens.append(vision_start_token_id)
+                        _current_offset += 1
+                    video_tokens.extend([video_token_id] * frame_seqlen)
+                    if vision_end_token_id is not None:
+                        video_tokens.append(vision_end_token_id)
+                    offsets.append(
+                        (_current_offset, _current_offset + frame_seqlen - 1)
+                    )
+                    _current_offset += (
+                        frame_seqlen + 1
+                        if vision_end_token_id is not None
+                        else frame_seqlen
+                    )  # for vision_end_token_id
+                mm_token_num = len(video_tokens)
+                mm_token_id = None
+                video_idx += 1
+            else:
+                logger.warning(
+                    f"{modality} modality is not supported for qwen3_vl models with timestamps."
+                )
+                continue
+            assert cur_idx <= mm_start_idx
+            input_ids.extend(prompt[cur_idx : mm_start_idx + 1])
+            if modality == Modality.VIDEO:
+                input_ids.extend(video_tokens)
+            else:
+                mm_offset_start = len(input_ids)
+                input_ids.extend([mm_token_id] * mm_token_num)
+                offsets.append((mm_offset_start, len(input_ids) - 1))
+            cur_idx = mm_start_idx + 2  # jump to vision_end_id
+        else:
+            input_ids.extend(prompt[cur_idx:])
+
+        return input_ids, offsets, modality_list
+
+    def get_mm_data(self, prompt, embeddings, **kwargs):
+        img_grid_thw = kwargs.get("img_grid_thw", None)
+        video_grid_thw = kwargs.get("video_grid_thw", None)
+        audio_feature_lens = kwargs.get("audio_feature_lens", None)
+        video_timestamps = kwargs.get("video_timestamps", None)
+        second_per_grid_ts = kwargs.get("second_per_grid_ts", None)
+
+        audio_seq_lens = None
+        if audio_feature_lens is not None:
+            if self.model_type == "qwen3_omni_moe":
+                # apply _get_feat_extract_lengths to get seq_lens
+                input_lengths_leave = audio_feature_lens % 100
+                feat_lengths = (input_lengths_leave - 1) // 2 + 1
+                audio_seq_lens = (
+                    ((feat_lengths - 1) // 2 + 1 - 1) // 2
+                    + 1
+                    + (audio_feature_lens // 100) * 13
+                )
+            elif self.model_type == "qwen2_5_omni":
+                audio_seq_lens = (audio_feature_lens - 1) // 2 + 1
+                audio_seq_lens = (audio_seq_lens - 2) // 2 + 1
+
+        if (
+            self.model_type in ["qwen3_vl", "qwen3_vl_moe"]
+            and video_timestamps is not None
+        ):
+            input_ids, offsets, modality_list = self.build_input_ids_with_timestamps(
+                prompt, embeddings, img_grid_thw, video_grid_thw, video_timestamps
+            )
+        else:
+            input_ids, offsets, modality_list = self.build_input_ids(
+                prompt, img_grid_thw, video_grid_thw, audio_seq_lens=audio_seq_lens
+            )
+        assert all(isinstance(modality, Modality) for modality in modality_list)
+
         mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index(
             spatial_merge_size=self.hf_config.vision_config.spatial_merge_size,
             image_token_id=self.mm_tokens.image_token_id,
@@ -285,19 +412,41 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
             model_type=self.model_type,
             input_ids=torch.tensor(input_ids, dtype=torch.long).unsqueeze(0),
             image_grid_thw=img_grid_thw,
+            video_grid_thw=video_grid_thw,
+            second_per_grid_ts=second_per_grid_ts,
+            use_audio_in_video=False,
+            audio_seqlens=(
+                audio_feature_lens if self.model_type == "qwen3_omni_moe" else None
+            ),
+            audio_token_id=getattr(self.hf_config, "audio_token_id", None),
+            audio_start_token_id=self.audio_start_token_id,
+            position_id_per_seconds=getattr(
+                self.hf_config, "position_id_per_seconds", None
+            ),
             tokens_per_second=getattr(
                 self.hf_config.vision_config, "tokens_per_second", None
             ),
         )
         mrope_positions = mrope_positions.squeeze(1)
 
-        mm_items = [
+        mm_items = []
-            MultimodalDataItem(
+        consumed_per_modality = {}
-                modality=Modality.IMAGE,
+
-                offsets=offsets,
+        for modality, offset in zip(modality_list, offsets):
-                precomputed_embeddings=embeddings,
+            num_tokens = offset[1] - offset[0] + 1
+            embedding_start = consumed_per_modality.get(modality, 0)
+            embedding_slice = embeddings[modality][
+                embedding_start : embedding_start + num_tokens
+            ]
+            consumed_per_modality[modality] = embedding_start + num_tokens
+            logger.info(f"Get embedding slice for {modality}, num_tokens={num_tokens}")
+            mm_items.append(
+                MultimodalDataItem(
+                    modality=modality,
+                    offsets=offset,
+                    precomputed_embeddings=embedding_slice,
+                )
             )
-        ]
 
         return {
             "input_ids": input_ids,

python/sglang/srt/server_args.py

@@ -2984,6 +2984,22 @@ class ServerArgs:
                 self.disaggregation_ib_device
             )
 
+        # Validate model type: only support Qwen models for now
+        hf_config = self.get_model_config().hf_config
+        model_arch = hf_config.architectures[0]
+        if (self.encoder_only or self.language_only) and model_arch not in [
+            "Qwen2VLForConditionalGeneration",
+            "Qwen3VLForConditionalGeneration",
+            "Qwen2_5_VLForConditionalGeneration",
+            "Qwen3VLMoeForConditionalGeneration",
+            "Qwen3OmniMoeForConditionalGeneration",
+            "Qwen2AudioForConditionalGeneration",
+            "Qwen2_5OmniForConditionalGeneration",
+        ]:
+            raise ValueError(
+                f"Model type {model_arch} is not supported for encoder disaggregation, only Qwen models are supported for now."
+            )
+
     def _validate_ib_devices(self, device_str: str) -> Optional[str]:
         """
         Validate IB devices before passing to mooncake.

test/registered/distributed/test_epd_disaggregation.py

@@ -1,10 +1,13 @@
+import io
 import os
+import re
 import subprocess
 import threading
 import time
 import unittest
 
 import grpc
+import openai
 import zmq
 from grpc_health.v1 import health_pb2, health_pb2_grpc
 
@@ -21,10 +24,595 @@ from sglang.test.test_utils import (
     is_in_ci,
     popen_launch_server,
 )
+from sglang.test.vlm_utils import (
+    AUDIO_TRUMP_SPEECH_URL,
+    IMAGE_MAN_IRONING_URL,
+    IMAGE_SGL_LOGO_URL,
+    VIDEO_JOBS_URL,
+)
+
+# Omni model for local testing; override via env var EPD_OMNI_MODEL
+DEFAULT_OMNI_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+
 
 register_cuda_ci(est_time=150, suite="stage-c-test-4-gpu-h100")
 
 
+@unittest.skipIf(
+    is_in_ci(),
+    "Omni model EPD test with image, video, and audio modalities, running locally only",
+)
+class TestEPDDisaggregationOmni(PDDisaggregationServerBase):
+    """
+    EPD disaggregation test for omni models (e.g. Qwen3-Omni). Covers image, video,
+    and audio when server_type=http (encoder_transfer_backend: mooncake/zmq_to_scheduler/zmq_to_tokenizer).
+    When server_type=grpc, only image is tested (gRPC encode is image-only).
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        cls.model = os.environ.get("EPD_OMNI_MODEL", DEFAULT_OMNI_MODEL)
+        cls.server_type = os.environ.get("EPD_ENCODE_SERVER_TYPE", "http")
+        assert cls.server_type in (
+            "grpc",
+            "http",
+        ), f"Invalid EPD_ENCODE_SERVER_TYPE: {cls.server_type}"
+        cls.encoder_transfer_backend = os.environ.get(
+            "EPD_ENCODER_TRANSFER_BACKEND", "zmq_to_scheduler"
+        )
+        assert cls.encoder_transfer_backend in (
+            "mooncake",
+            "zmq_to_scheduler",
+            "zmq_to_tokenizer",
+        ), f"Invalid EPD_ENCODER_TRANSFER_BACKEND: {cls.encoder_transfer_backend}"
+        cls.enable_global_cache = (
+            os.environ.get("MOONCAKE_MASTER") is not None
+            or os.environ.get("MOONCAKE_CLIENT") is not None
+        )
+        if cls.server_type == "grpc":
+            cls.encode_port = f"{int(cls.lb_port) + 305}"
+            cls.encode_url = f"grpc://{cls.base_host}:{cls.encode_port}"
+        else:
+            cls.encode_port = f"{int(cls.lb_port) + 300}"
+            cls.encode_url = f"http://{cls.base_host}:{cls.encode_port}"
+
+        cls.image_man_ironing = IMAGE_MAN_IRONING_URL
+        cls.image_sgl_logo = IMAGE_SGL_LOGO_URL
+        cls.video_jobs = VIDEO_JOBS_URL
+        cls.audio_trump = AUDIO_TRUMP_SPEECH_URL
+
+        print(
+            f"Setting up EPD Omni: model={cls.model}, encode={cls.encode_port}, "
+            f"prefill={cls.prefill_port}, decode={cls.decode_port}, "
+            f"server_type={cls.server_type}, backend={cls.encoder_transfer_backend}, "
+            f"global_cache={cls.enable_global_cache}"
+        )
+        print(f"Data URLs: image={cls.image_man_ironing}, audio={cls.audio_trump}")
+
+        cls.start_encode()
+        prefill_thread = threading.Thread(target=cls.start_prefill)
+        decode_thread = threading.Thread(target=cls.start_decode)
+        prefill_thread.start()
+        decode_thread.start()
+        prefill_thread.join()
+        decode_thread.join()
+
+        if cls.server_type == "grpc":
+            cls._wait_grpc_ready(cls.base_host, cls.encode_port, cls.process_encode)
+        else:
+            cls.wait_server_ready(
+                cls.encode_url + "/health", process=cls.process_encode
+            )
+        cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill)
+        cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode)
+
+        cls.launch_lb()
+
+        cls.api_key = "sk-123456"
+        os.environ["OPENAI_API_KEY"] = cls.api_key
+        os.environ["OPENAI_API_BASE"] = f"{cls.lb_url}/v1"
+
+    @classmethod
+    def start_encode(cls):
+        if cls.server_type == "grpc":
+            cls.encode_stdout = io.StringIO()
+            cls.encode_stderr = io.StringIO()
+            cls.process_encode = subprocess.Popen(
+                [
+                    "python3",
+                    "-m",
+                    "sglang.launch_server",
+                    "--model-path",
+                    cls.model,
+                    "--host",
+                    cls.base_host,
+                    "--port",
+                    cls.encode_port,
+                    "--trust-remote-code",
+                    "--encoder-only",
+                    "--grpc-mode",
+                    "--encoder-transfer-backend",
+                    "zmq_to_scheduler",
+                    "--tp",
+                    "1",
+                ]
+            )
+        else:
+            encode_args = [
+                "--trust-remote-code",
+                "--encoder-only",
+                "--encoder-transfer-backend",
+                cls.encoder_transfer_backend,
+                "--tp",
+                "1",
+                "--port",
+                cls.encode_port,
+            ]
+            if cls.enable_global_cache:
+                encode_args.append("--enable-mm-global-cache")
+            cls.encode_stdout = io.StringIO()
+            cls.encode_stderr = io.StringIO()
+            cls.process_encode = popen_launch_server(
+                cls.model,
+                base_url=cls.encode_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=encode_args,
+                return_stdout_stderr=(cls.encode_stdout, cls.encode_stderr),
+            )
+
+    @classmethod
+    def start_prefill(cls):
+        prefill_args = [
+            "--trust-remote-code",
+            "--language-only",
+            "--encoder-urls",
+            cls.encode_url,
+            "--encoder-transfer-backend",
+            (
+                "zmq_to_scheduler"
+                if cls.server_type == "grpc"
+                else cls.encoder_transfer_backend
+            ),
+            "--disaggregation-mode",
+            "prefill",
+            "--tp",
+            "1",
+            "--base-gpu-id",
+            "1",
+            "--port",
+            cls.prefill_port,
+        ]
+        prefill_args += cls.transfer_backend + cls.rdma_devices
+        prefill_env = os.environ.copy()
+        if cls.server_type == "grpc":
+            prefill_env["SGLANG_ENCODER_MM_RECEIVER_MODE"] = "grpc"
+        cls.process_prefill = popen_launch_server(
+            cls.model,
+            base_url=cls.prefill_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=prefill_args,
+            env=prefill_env,
+        )
+
+    @classmethod
+    def start_decode(cls):
+        decode_args = [
+            "--trust-remote-code",
+            "--disaggregation-mode",
+            "decode",
+            "--tp",
+            "1",
+            "--base-gpu-id",
+            "2",
+            "--port",
+            cls.decode_port,
+        ]
+        decode_args += cls.transfer_backend + cls.rdma_devices
+        cls.process_decode = popen_launch_server(
+            cls.model,
+            base_url=cls.decode_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=decode_args,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        for process in [
+            cls.process_lb,
+            cls.process_decode,
+            cls.process_prefill,
+            cls.process_encode,
+        ]:
+            if process:
+                try:
+                    kill_process_tree(process.pid)
+                except Exception as e:
+                    print(f"Error killing process: {e}")
+
+    @staticmethod
+    def _wait_grpc_ready(
+        host, port, process, timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH
+    ):
+        deadline = time.time() + timeout
+        channel = grpc.insecure_channel(f"{host}:{port}")
+        stub = health_pb2_grpc.HealthStub(channel)
+        try:
+            while time.time() < deadline:
+                if process.poll() is not None:
+                    raise RuntimeError(
+                        f"gRPC encoder exited with code {process.returncode}"
+                    )
+                try:
+                    response = stub.Check(
+                        health_pb2.HealthCheckRequest(service=""), timeout=2
+                    )
+                    if response.status == health_pb2.HealthCheckResponse.SERVING:
+                        return
+                except grpc.RpcError:
+                    pass
+                time.sleep(1)
+        finally:
+            channel.close()
+        raise RuntimeError(f"gRPC encoder not ready at {host}:{port} within {timeout}s")
+
+    # ---- helpers ----
+
+    def _client(self):
+        return openai.Client(api_key=self.api_key, base_url=f"{self.lb_url}/v1")
+
+    def _skip_if_grpc(self, msg="gRPC encode is image-only"):
+        """Skip this test when encode server is gRPC (image-only)."""
+        if self.server_type == "grpc":
+            self.skipTest(msg)
+
+    def _parse_cache_log(self):
+        """Parse encode server logs and return list of (local_hits, global_hits, misses)
+        tuples from '=== Multi-Level Cache Check ===' lines."""
+        log = self.encode_stdout.getvalue() + self.encode_stderr.getvalue()
+        pattern = re.compile(
+            r"Multi-Level Cache Check.*?"
+            r"Local Hits:\s*(\d+).*?"
+            r"Global Hits:\s*(\d+).*?"
+            r"Misses.*?:\s*(\d+)"
+        )
+        return [(int(m[1]), int(m[2]), int(m[3])) for m in pattern.finditer(log)]
+
+    # ---- image ----
+    def test_image(self):
+        client = self._client()
+        response = client.chat.completions.create(
+            model="default",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": self.image_man_ironing},
+                        },
+                        {
+                            "type": "text",
+                            "text": "Describe this image in a sentence.",
+                        },
+                    ],
+                },
+            ],
+            temperature=0,
+            max_tokens=256,
+        )
+        text = response.choices[0].message.content
+        print(f"[Omni EPD] Image response:\n{text}")
+        self.assertIsNotNone(text)
+        self.assertGreater(len(text), 0)
+
+        text_lower = text.lower()
+        self.assertTrue(
+            any(w in text_lower for w in ("man", "person", "driver")),
+            f"Image response should mention a person: {text}",
+        )
+        self.assertTrue(
+            any(w in text_lower for w in ("iron", "cloth", "hang", "holding")),
+            f"Image response should mention ironing/clothes: {text}",
+        )
+
+    def test_image_cache_hit(self):
+        """Send the same image twice; the second request should hit the global-mm-cache."""
+        self._skip_if_grpc("gRPC encode is image-only; cache test uses HTTP path")
+        if not self.enable_global_cache:
+            self.skipTest("global-mm-cache not enabled (MOONCAKE_MASTER not set)")
+        client = self._client()
+        baseline = len(self._parse_cache_log())
+        for i in range(2):
+            response = client.chat.completions.create(
+                model="default",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "image_url",
+                                "image_url": {"url": self.image_sgl_logo},
+                            },
+                            {
+                                "type": "text",
+                                "text": "What is shown in this image?",
+                            },
+                        ],
+                    },
+                ],
+                temperature=0,
+                max_tokens=128,
+            )
+            text = response.choices[0].message.content
+            print(f"[Omni EPD] Image cache-hit round {i}: {text}")
+            self.assertIsNotNone(text)
+            self.assertGreater(len(text), 0)
+            time.sleep(1)
+
+        entries = self._parse_cache_log()[baseline:]
+        print(f"[Omni EPD] Image cache log entries: {entries}")
+        self.assertGreaterEqual(
+            len(entries), 2, "Expected at least 2 cache-check log entries"
+        )
+        local_hits, global_hits, _ = entries[-1]
+        self.assertGreater(
+            local_hits + global_hits,
+            0,
+            f"Second image request should have cache hits, got: {entries[-1]}",
+        )
+
+    # ---- video ----
+    def test_video(self):
+        self._skip_if_grpc()
+        client = self._client()
+        response = client.chat.completions.create(
+            model="default",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": "Describe the video."},
+                        {
+                            "type": "video_url",
+                            "video_url": {"url": self.video_jobs},
+                        },
+                    ],
+                },
+            ],
+            max_tokens=8192,
+            stream=False,
+        )
+        text = response.choices[0].message.content
+        print(f"[Omni EPD] Video response:\n{text}")
+        self.assertIsNotNone(text)
+        self.assertGreater(len(text), 0)
+
+        text_lower = text.lower()
+        self.assertTrue(
+            any(
+                w in text_lower
+                for w in ("ipod", "device", "microphone", "smartphone", "phone")
+            ),
+            f"Video response should mention a device: {text}",
+        )
+        self.assertTrue(
+            any(
+                w in text_lower
+                for w in (
+                    "man",
+                    "person",
+                    "individual",
+                    "speaker",
+                    "presenter",
+                    "steve",
+                    "hand",
+                    "hands",
+                )
+            ),
+            f"Video response should mention a person: {text}",
+        )
+        self.assertTrue(
+            any(
+                w in text_lower
+                for w in (
+                    "present",
+                    "presenting",
+                    "examine",
+                    "examining",
+                    "display",
+                    "displaying",
+                    "hold",
+                    "holding",
+                    "gestur",
+                    "speak",
+                    "speaking",
+                )
+            ),
+            f"Video response should mention an action: {text}",
+        )
+
+    def test_video_cache_hit(self):
+        """Send the same video twice; the second request should hit the global-mm-cache."""
+        self._skip_if_grpc()
+        if not self.enable_global_cache:
+            self.skipTest("global-mm-cache not enabled (MOONCAKE_MASTER not set)")
+        client = self._client()
+        baseline = len(self._parse_cache_log())
+        for i in range(2):
+            response = client.chat.completions.create(
+                model="default",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "text", "text": "Describe the video."},
+                            {
+                                "type": "video_url",
+                                "video_url": {"url": self.video_jobs},
+                            },
+                        ],
+                    },
+                ],
+                max_tokens=256,
+                stream=False,
+            )
+            text = response.choices[0].message.content
+            print(f"[Omni EPD] Video cache-hit round {i}: {text}")
+            self.assertIsNotNone(text)
+            self.assertGreater(len(text), 0)
+            time.sleep(1)
+
+        entries = self._parse_cache_log()[baseline:]
+        print(f"[Omni EPD] Video cache log entries: {entries}")
+        self.assertGreaterEqual(
+            len(entries), 2, "Expected at least 2 cache-check log entries"
+        )
+        local_hits, global_hits, _ = entries[-1]
+        self.assertGreater(
+            local_hits + global_hits,
+            0,
+            f"Second video request should have cache hits, got: {entries[-1]}",
+        )
+
+    # ---- audio ----
+
+    def test_audio(self):
+        self._skip_if_grpc()
+        client = self._client()
+        response = client.chat.completions.create(
+            model="default",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "audio_url",
+                            "audio_url": {"url": self.audio_trump},
+                        },
+                        {
+                            "type": "text",
+                            "text": "Listen to this audio and write down the audio transcription in English.",
+                        },
+                    ],
+                },
+            ],
+            temperature=0,
+            max_tokens=256,
+            stream=False,
+        )
+        text = response.choices[0].message.content
+        print(f"[Omni EPD] Audio response:\n{text}")
+        self.assertIsNotNone(text)
+        self.assertGreater(len(text), 0)
+
+        text_lower = text.lower()
+        for keyword in ("thank you", "leader"):
+            self.assertIn(
+                keyword,
+                text_lower,
+                f"Audio response should contain '{keyword}': {text}",
+            )
+
+    def test_audio_cache_hit(self):
+        """Send the same audio twice; the second request should hit the global-mm-cache."""
+        self._skip_if_grpc()
+        if not self.enable_global_cache:
+            self.skipTest("global-mm-cache not enabled (MOONCAKE_MASTER not set)")
+        client = self._client()
+        baseline = len(self._parse_cache_log())
+        for i in range(2):
+            response = client.chat.completions.create(
+                model="default",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": [
+                            {
+                                "type": "audio_url",
+                                "audio_url": {"url": self.audio_trump},
+                            },
+                            {
+                                "type": "text",
+                                "text": "What is this audio about?",
+                            },
+                        ],
+                    },
+                ],
+                temperature=0,
+                max_tokens=128,
+                stream=False,
+            )
+            text = response.choices[0].message.content
+            print(f"[Omni EPD] Audio cache-hit round {i}: {text}")
+            self.assertIsNotNone(text)
+            self.assertGreater(len(text), 0)
+            time.sleep(1)
+
+        entries = self._parse_cache_log()[baseline:]
+        print(f"[Omni EPD] Audio cache log entries: {entries}")
+        self.assertGreaterEqual(
+            len(entries), 2, "Expected at least 2 cache-check log entries"
+        )
+        local_hits, global_hits, _ = entries[-1]
+        self.assertGreater(
+            local_hits + global_hits,
+            0,
+            f"Second audio request should have cache hits, got: {entries[-1]}",
+        )
+
+    # ---- mixed modality ----
+
+    def test_mixed_image_audio_video(self):
+        """Image + audio + video in one request to test multi-modal routing."""
+        self._skip_if_grpc()
+        client = self._client()
+        response = client.chat.completions.create(
+            model="default",
+            messages=[
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": self.image_man_ironing},
+                        },
+                        {
+                            "type": "audio_url",
+                            "audio_url": {"url": self.audio_trump},
+                        },
+                        {
+                            "type": "video_url",
+                            "video_url": {"url": self.video_jobs},
+                        },
+                        {
+                            "type": "text",
+                            "text": (
+                                "I have an image, an audio clip, and a video, which are not related at all. "
+                                "Please: 1. Describe the image in a sentence, "
+                                "2. Summarize the audio content briefly, "
+                                "3. Describe what happens in the video."
+                            ),
+                        },
+                    ],
+                },
+            ],
+            temperature=0,
+            max_tokens=512,
+            stream=False,
+        )
+        text = response.choices[0].message.content
+        print(f"[Omni EPD] Mixed image+audio+video response:\n{text}")
+        self.assertIsNotNone(text)
+        self.assertGreater(len(text), 0)
+
+        text_lower = text.lower()
+        self.assertTrue(
+            any(w in text_lower for w in ("man", "person", "iron", "cloth")),
+            f"Mixed response should describe the image: {text}",
+        )
+
+
 @unittest.skipIf(is_in_ci(), "Skipping in CI to reduce multi-GPU runtime")
 class TestEPDDisaggregationOneEncoder(PDDisaggregationServerBase):
     """Test EPD disaggregation with single encode server"""