# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. """Multimodal request data structures used across processors and model adapters.""" from __future__ import annotations import dataclasses import uuid from enum import Enum, auto from typing import Any import numpy as np import torch from tokenspeed.runtime.multimodal.hash import hash_feature from tokenspeed.runtime.multimodal.shm_transport import ShmTensorHandle from tokenspeed.runtime.utils.env import envs # Multimodal pad-value substitute IDs: a placeholder mm token's id is rewritten # to ``_MM_PAD_BASE + (hash & _MM_PAD_HASH_MASK)`` so duplicate features share # the same substitute and prefix-match in the text-only prefix cache. The base # sits well above any text vocab; the 30-bit mask keeps cross-hash collisions # rare enough for long-running servers (~10^9 slots). _MM_PAD_BASE = 1_000_000 _MM_PAD_HASH_MASK = (1 << 30) - 1 def is_mm_pad_value(token_ids: torch.Tensor) -> torch.Tensor: """Bool mask of positions rewritten to a hash-derived multimodal pad id.""" return (token_ids >= _MM_PAD_BASE) & (token_ids <= _MM_PAD_BASE + _MM_PAD_HASH_MASK) def maybe_substitute_mm_pad( input_ids: torch.Tensor, substitute_id: int | None ) -> torch.Tensor: """Replace hash mm-pad positions with ``substitute_id``; no-op if None.""" if substitute_id is None: return input_ids return input_ids.masked_fill(is_mm_pad_value(input_ids), substitute_id) class Modality(Enum): IMAGE = auto() VIDEO = auto() AUDIO = auto() # ``eq=False`` on every dataclass below: tensor-valued fields crash the # default element-wise ``__eq__`` and force ``__hash__`` to None. @dataclasses.dataclass(eq=False) class MultimodalDataItem: modality: Modality hash: int | None = None pad_value: int | None = None offsets: list | None = None feature: torch.Tensor | np.ndarray | ShmTensorHandle | None = None model_specific_data: dict[str, Any] = dataclasses.field(default_factory=dict) # Encoder output for this item, populated on first encoder pass and reused # across chunked-prefill iterations of the owning request. Lifetime is # tied to the request: when the request finishes the item is GC'd and # these tensors are released. ``encoded_deepstack`` is set only for # deepstack-enabled modalities. encoded: torch.Tensor | None = None encoded_deepstack: torch.Tensor | None = None # EPD (encode-prefill-decode): when set, this item's embedding is received # from an encode worker over Mooncake into ``encoded`` instead of running the # vision tower. A dict ``{bootstrap_room, bootstrap_host, bootstrap_port}`` # naming the encode worker's rendezvous for this item's image (one room per # item: the gateway splits the mm payload one item per image and the encode # worker row-splits the concatenated-subgrid embedding per item). None for # non-EPD items (left to the vision tower). encode_handshake: dict | None = None def __getattr__(self, name: str): if ( "model_specific_data" in self.__dict__ and name in self.__dict__["model_specific_data"] ): return self.__dict__["model_specific_data"][name] raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'") def ensure_hash(self): """Resolve ``self.hash`` to a concrete content id, lazily. The hash is resolved on demand rather than at construction because it is usually supplied by the caller, a SHM-backed feature cannot be hashed here without reading shared memory, and hashing inline bytes is only worth doing once the value is actually needed. Resolution order: * ``TOKENSPEED_MM_SKIP_COMPUTE_HASH`` -> a random id (dedup disabled); * an already-set hash (e.g. the gateway-provided ``content_hash`` for image/video) is kept as-is, no recompute; * inline features the gateway does not hash (e.g. audio) are hashed in-engine via ``hash_feature``; * SHM-backed features must carry a caller-provided hash, else raise -- we cannot hash a handle without reading shared memory. """ if envs.TOKENSPEED_MM_SKIP_COMPUTE_HASH.get(): self.hash = uuid.uuid4().int elif self.hash is None: if isinstance(self.feature, ShmTensorHandle): raise ValueError( "SHM-backed multimodal items must carry content hash or " "pad_value before TokenSpeed consumes them" ) self.hash = hash_feature(self.feature) if self.hash is None: raise RuntimeError("Failed to resolve multimodal item hash.") def set_pad_value(self): if self.pad_value is not None: return self.ensure_hash() self.pad_value = _MM_PAD_BASE + (self.hash & _MM_PAD_HASH_MASK) def is_modality(self, modality: Modality) -> bool: return self.modality == modality @dataclasses.dataclass(eq=False) class MultimodalInputs: mm_items: list[MultimodalDataItem] im_token_id: int | None = None video_token_id: int | None = None mrope_positions: torch.Tensor | None = None mrope_position_delta: torch.Tensor | None = None mrope_position_delta_scalar: int | None = None mrope_position_delta_repeated_cache: torch.Tensor | None = None def ensure_pad_values(self) -> None: for item in self.mm_items: item.set_pad_value() def publish_shm_features(self) -> None: for item in self.mm_items: if isinstance(item.feature, torch.Tensor): item.feature = ShmTensorHandle.publish(item.feature) def attach_shm_features(self) -> None: """Open every pending handle on this rank. Must run before the cross-rank barrier in ``request_handler.recv_reqs``. """ for item in self.mm_items: if isinstance(item.feature, ShmTensorHandle): item.feature.attach() def release_shm_features(self) -> None: for item in self.mm_items: if isinstance(item.feature, ShmTensorHandle): item.feature.release() item.feature = None def has_pending_shm_features(self) -> bool: return any(isinstance(item.feature, ShmTensorHandle) for item in self.mm_items) @dataclasses.dataclass(eq=False) class MultimodalForwardContext: """Per-forward multimodal metadata for prefill embedding replacement.""" mm_inputs: list[MultimodalInputs | None] extend_prefix_lens: list[int] extend_seq_lens: list[int] def has_inputs(self) -> bool: return bool(self.mm_inputs and any(x is not None for x in self.mm_inputs)) def has_extend_inputs(self) -> bool: return any( mm_input is not None and index < len(self.extend_seq_lens) for index, mm_input in enumerate(self.mm_inputs) )