SGLang 源码完全拆解（十四）：多模态输入——图像如何在文本流里占位、编码与回填

1
class MultimodalDataItem(msgspec.Struct, kw_only=True, dict=True, array_like=True):
2
    """
3
    One MultimodalDataItem represents a single multimodal input
4
    (one image, one video, or one audio).
5
    For example, if there are 3 images and 1 audio, there will be 4
6
    MultimodalDataItems.
7

8
    Each item has its own hash and pad_value, enabling per-image
9
    RadixAttention caching.
10
    """
11

12
    modality: Modality
13
    hash: Optional[int] = None
14
    pad_value: Optional[int] = None
15
    offsets: Optional[List[Tuple[int, int]]] = None
16

17
    format: MultimodalInputFormat = MultimodalInputFormat.NORMAL
18

19
    # the raw features returned by processor, e.g. pixel_values or audio_features
20
    feature: Optional[MultimodalDataValue] = None
21
    # the precomputed embeddings, passed as final encoder embeddings
22
    precomputed_embeddings: Optional[MultimodalDataValue] = None
23
    keep_device_embedding: bool = False
24

25
    model_specific_data: Dict[str, MultimodalDataValue] = msgspec.field(
26
        default_factory=dict
27
    )
28

29
    identity: Optional[bytes] = None
30
    pad_values: Optional[Tuple[int, ...]] = None

1
# 读：未知字段从 model_specific_data 里取，取不到才抛 AttributeError
2
def __getattr__(self, name: str) -> MultimodalDataValue:
3
    if name in self.model_specific_data:
4
        return self.model_specific_data[name]
5
    raise AttributeError(
6
        f"'{self.__class__.__name__}' object has no attribute '{name}'"
7
    )
8

9
# 写：声明过的字段走 msgspec 正常赋值，其余一律落进 model_specific_data
10
def __setattr__(self, name: str, value: MultimodalDataValue) -> None:
11
    if name in self.__struct_fields__:
12
        if name == "hash" and isinstance(value, int):
13
            value &= _MM_HASH_MASK
14
        msgspec.Struct.__setattr__(self, name, value)
15
    else:
16
        self.model_specific_data[name] = value

1
# Constant used as the base offset for MM (multimodal) pad values.
2
# This ensures pad_values don't overlap with valid text token IDs.
3
MM_PAD_SHIFT_VALUE = 1_000_000
4
_MM_HASH_MASK = (1 << 64) - 1
5
_MM_IDENTITY_BYTES = 32
6
_MM_WIDE_PAD_TAG = b"sglang-mm-wide-pad-v1"
7

8
# Keep MM_PAD_SHIFT_VALUE + sentinel below 2**63 so it fits an int64 token ID.
9
_MM_WIDE_PAD_MASK = (1 << 62) - 1
10

11

12
def _compute_pad_values(identity: bytes) -> Tuple[int, ...]:
13
    """derive shared wide sentinels from an item hash or full identity digest"""
14
    digest = hashlib.sha256(_MM_WIDE_PAD_TAG + identity).digest()
15
    # prefix matching can stop after one sentinel, so each must retain wide bits
16
    return tuple(
17
        MM_PAD_SHIFT_VALUE
18
        + (int.from_bytes(digest[i : i + 8], byteorder="big") & _MM_WIDE_PAD_MASK)
19
        for i in range(0, _MM_IDENTITY_BYTES, 8)
20
    )

1
def set_identity(self, identity: bytes) -> None:
2
    """Adopt a full 32-byte digest and derive four 62-bit pad sentinels.
3

4
    ``hash`` keeps the leading 64 bits for scalar consumers such as the
5
    embedding cache; the radix key uses the full sentinel pattern.
6
    """
7
    if len(identity) != _MM_IDENTITY_BYTES:
8
        raise ValueError(
9
            f"multimodal identity must be {_MM_IDENTITY_BYTES} bytes, "
10
            f"got {len(identity)}"
11
        )
12
    self.identity = bytes(identity)
13
    self.hash = int.from_bytes(identity[:8], byteorder="big", signed=False)
14
    self.pad_values = _compute_pad_values(self.identity)
15
    self.pad_value = self.pad_values[0]

1
@staticmethod
2
def build_padded_input_ids(input_ids, mm_items: List[MultimodalDataItem]):
3
    """pad the input_ids with mm_items if it's not already padded"""
4
    if input_ids is None or not mm_items:
5
        return None
6

7
    for item in mm_items:
8
        if item.pad_value is None or item.offsets is None:
9
            return None
10

11
    if isinstance(input_ids, torch.Tensor):
12
        padded_input_ids = input_ids.flatten().tolist()
13
    else:
14
        padded_input_ids = list(input_ids)
15

16
    for item in mm_items:
17
        for start, end in item.offsets:
18
            padded_input_ids[start : end + 1] = item.padding_sequence(
19
                end - start + 1
20
            )
21
    return padded_input_ids

1
def pad_mm_token_pairs(
2
    input_ids: array,
3
    mm_inputs: MultimodalInputs,
4
    *,
5
    data_token_pairs: Optional[List[Tuple[int, int]]],
6
    data_start_token_ids: Optional[List[int]] = None,
7
) -> array:
8
    """Pad enclosed tokens without mutating input_ids; record item start offsets.
9

10
    data_start_token_ids identifies new items, so slices can reuse an item's
11
    pad value. By default, every start token begins a new item.
12
    """
13
    assert_int64_array(input_ids, "input_ids")
14
    pad_items = mm_inputs.mm_items
15
    mm_inputs.data_offsets = []
16
    if data_token_pairs is None:
17
        data_token_pairs = [mm_inputs.im_start_id, mm_inputs.im_end_id]
18
    if data_token_pairs is None:
19
        print_warning_once(
20
            "No data_token_pairs provided, RadixAttention might be influenced."
21
        )
22
        return input_ids
23
    start_token_ids = {s for s, _e in data_token_pairs}
24
    end_tokens_ids = {e for _s, e in data_token_pairs}
25

26
    padded_ids = array("q")
27
    last_idx = 0
28
    data_idx = -1
29

30
    start_indices = [i for i, x in enumerate(input_ids) if x in start_token_ids]
31
    end_indices = [i for i, x in enumerate(input_ids) if x in end_tokens_ids]
32

33
    if len(start_indices) != len(end_indices):
34
        return input_ids
35

36
    for start_idx, end_idx in zip(start_indices, end_indices):
37
        padded_ids.extend(input_ids[last_idx : start_idx + 1])
38

39
        if input_ids[start_idx] in data_start_token_ids:
40
            data_idx += 1
41
            mm_inputs.data_offsets += [start_idx]
42

43
        if data_idx >= len(pad_items):
44
            data_idx = len(pad_items) - 1
45

46
        num_tokens = end_idx - start_idx - 1
47
        if num_tokens > 0:
48
            padded_ids.extend(pad_items[data_idx].padding_sequence(num_tokens))
49

50
        last_idx = end_idx
51

52
    padded_ids.extend(input_ids[last_idx:])
53

54
    assert len(input_ids) == len(padded_ids), "Length validation fails"
55
    return padded_ids

1
def data_hash(data) -> int:
2
    hash_bytes = hashlib.sha256(data).digest()[:8]
3
    return int.from_bytes(hash_bytes, byteorder="big", signed=False)
4

5

6
def tensor_hash(tensor_list) -> int:
7
    """
8
    hash a tensor or a tensor list
9
    """
10
    tensor = tensor_list
11
    if isinstance(tensor_list, list):
12
        tensor_list = flatten_nested_list(tensor_list)
13
        tensors = [
14
            x.flatten() if isinstance(x, torch.Tensor) else x for x in tensor_list
15
        ]
16
        # GPU path: concat + triton hash (unchanged)
17
        if any(isinstance(t, torch.Tensor) and t.is_cuda for t in tensors):
18
            tensor = torch.concat(tensors)
19
            return gpu_tensor_hash(tensor.cuda())
20
        # CPU path: hash each tensor incrementally without concat
21
        hasher = hashlib.sha256()
22
        for t in tensors:
23
            t = t.detach().cpu().contiguous()
24
            hasher.update(memoryview(t.reshape(-1).view(torch.uint8).numpy()))
25
        hash_bytes = hasher.digest()[:8]
26
        return int.from_bytes(hash_bytes, byteorder="big", signed=False)

1
caller_mm_hashes = getattr(obj, "mm_hashes", None)
2
if caller_mm_hashes and mm_inputs and mm_inputs.mm_items:
3
    if len(caller_mm_hashes) != len(mm_inputs.mm_items):
4
        logger.warning(
5
            "mm_hashes length (%d) != mm_items length (%d); "
6
            "ignoring caller hashes for this request.",
7
            len(caller_mm_hashes),
8
            len(mm_inputs.mm_items),
9
        )
10
    else:
11
        for item, hex_hash in zip(mm_inputs.mm_items, caller_mm_hashes):
12
            if not isinstance(item, MultimodalDataItem):
13
                continue
14
            try:
15
                item.set_hash(int(hex_hash, 16))
16
            except (TypeError, ValueError):
17
                logger.warning(
18
                    "Ignoring malformed mm_hashes entry %r; "
19
                    "this item will fall back to hash_feature().",
20
                    hex_hash,
21
                )

1
@dataclasses.dataclass
2
class MultimodalInputs:
3
    """The multimodal data related inputs."""
4

5
    # items of data
6
    mm_items: List[MultimodalDataItem]
7
    padded_input_ids: Optional[List[int]] = None
8
    image_pad_len: Optional[list] = None
9
    num_image_tokens: Optional[int] = None
10
    token_modalities: Optional[List[int]] = None
11

12
    # image
13
    im_token_id: Optional[int] = None
14
    im_start_id: Optional[int] = None
15
    im_end_id: Optional[int] = None
16
    slice_start_id: Optional[int] = None
17
    slice_end_id: Optional[int] = None
18

19
    # video
20
    video_token_id: Optional[int] = None
21

22
    # audio
23
    audio_token_id: Optional[int] = None
24
    audio_start_id: Optional[int] = None
25
    audio_end_id: Optional[int] = None
26

27
    # QWen2-VL related
28
    mrope_positions: Optional[torch.Tensor] = None
29
    mrope_position_delta: Optional[torch.Tensor] = None
30
    mrope_position_delta_repeated_cache: Optional[torch.Tensor] = None

1
token_modalities = [0] * num_tokens
2
for item in mm_items:
3
    if not item.offsets:
4
        continue
5
    modality = item.modality.value
6
    for start, end in item.offsets:
7
        if start < 0 or end < start or end >= num_tokens:
8
            raise ValueError(
9
                "Invalid multimodal token offsets: "
10
                f"offset=({start}, {end}), num_tokens={num_tokens}"
11
            )
12
        if any(token_modalities[index] for index in range(start, end + 1)):
13
            raise ValueError(
14
                f"Overlapping multimodal token offsets at ({start}, {end})"
15
            )
16
        token_modalities[start : end + 1] = [modality] * (end - start + 1)

1
def _maybe_compute_mrope_positions(self, req) -> None:
2
    """Compute M-RoPE positions when they are missing
3
    (e.g. gRPC preprocessed path)."""
4
    if self._mm_processor is None:
5
        return
6
    mm = req.multimodal_inputs
7
    if mm is None or mm.mrope_positions is not None:
8
        return
9

10
    mrope_positions, mrope_position_delta = (
11
        self._mm_processor.compute_mrope_positions(
12
            req.origin_input_ids, mm.mm_items
13
        )
14
    )
15
    if mrope_positions is not None:
16
        mm.mrope_positions = mrope_positions
17
        mm.mrope_position_delta = mrope_position_delta

1
PROCESSOR_MAPPING = {}
2

3

4
def import_processors(package_name: str, overwrite: bool = False):
5
    package = importlib.import_module(package_name)
6
    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
7
        if not ispkg:
8
            try:
9
                module = importlib.import_module(name)
10
            except Exception as e:
11
                logger.warning(f"Ignore import error when loading {name}: {e}")
12
                continue
13
            all_members = inspect.getmembers(module, inspect.isclass)
14
            classes = [
15
                member
16
                for name, member in all_members
17
                if member.__module__ == module.__name__
18
            ]
19
            for cls in (
20
                cls for cls in classes if issubclass(cls, BaseMultimodalProcessor)
21
            ):
22
                assert hasattr(cls, "models")
23
                for arch in getattr(cls, "models"):
24
                    if overwrite:
25
                        for model_cls, processor_cls in PROCESSOR_MAPPING.items():
26
                            if model_cls.__name__ == arch.__name__:
27
                                del PROCESSOR_MAPPING[model_cls]
28
                                break
29
                    PROCESSOR_MAPPING[arch] = cls

1
def get_mm_processor_cls(hf_config, model_config=None):
2
    """The class :func:`get_mm_processor` would instantiate, or ``None`` when the
3
    architecture has no registered processor."""
4
    model_impl = str(get_model().model_impl).lower()
5
    uses_transformers_backend = model_impl == "transformers"
6
    if model_impl == "auto" and model_config is not None:
7
        from sglang.srt.model_loader.utils import get_resolved_model_impl
8

9
        uses_transformers_backend = (
10
            get_resolved_model_impl(model_config) == ModelImpl.TRANSFORMERS
11
        )
12

13
    for model_cls, processor_cls in PROCESSOR_MAPPING.items():
14
        if model_cls.__name__ not in hf_config.architectures:
15
            continue
16
        if not uses_transformers_backend or getattr(
17
            processor_cls, "supports_transformers_backend", False
18
        ):
19
            return processor_cls
20

21
    if uses_transformers_backend:
22
        from sglang.srt.multimodal.processors.transformers_auto import (
23
            TransformersAutoMultimodalProcessor,
24
        )
25

26
        return TransformersAutoMultimodalProcessor
27

28
    return None

1
def get_mm_processor(
2
    hf_config,
3
    server_args: ServerArgs,
4
    processor,
5
    transport_mode,
6
    model_config=None,
7
    **kwargs,
8
) -> BaseMultimodalProcessor:
9
    processor_cls = get_mm_processor_cls(hf_config, model_config)
10
    if processor_cls is None:
11
        raise ValueError(
12
            f"No processor registered for architecture: {hf_config.architectures}.\n"
13
            f"Registered architectures: "
14
            f"{[model_cls.__name__ for model_cls in PROCESSOR_MAPPING.keys()]}"
15
        )
16
    return processor_cls(hf_config, server_args, processor, transport_mode, **kwargs)

1
@staticmethod
2
def _try_apply_padded_mm_input_ids(recv_req, req, image_inputs) -> bool:
3
    """setup origin_input_ids with trying to reuse existing
4
    MultimodalInputs.padded_input_ids first, if absent, call pad_input_ids_func"""
5
    padded_input_ids = image_inputs.padded_input_ids
6
    if padded_input_ids is None or recv_req.input_ids is None:
7
        return False
8

9
    recv_input_len = len(recv_req.input_ids)
10
    if len(padded_input_ids) != recv_input_len:
11
        return False
12

13
    prefix_len = len(req.origin_input_ids) - recv_input_len
14
    if prefix_len < 0:
15
        return False
16

17
    padded_input_ids = array("q", padded_input_ids)
18
    if prefix_len == 0:
19
        req.origin_input_ids = padded_input_ids
20
    else:
21
        req.origin_input_ids = req.origin_input_ids[:prefix_len] + padded_input_ids
22
    return True

1
def prepare_encoder_info_extend(
2
    self, input_ids: List[array[int]], seq_lens: List[int]
3
):
4
    _pin = is_pin_memory_available(self.device)
5
    encoder_lens_cpu = []
6
    encoder_cached = []
7

8
    for req in self.reqs:
9
        im = req.multimodal_inputs
10
        if im is None or im.num_image_tokens is None:
11
            # No image input
12
            encoder_lens_cpu.append(0)
13
            encoder_cached.append(True)
14
        else:
15
            encoder_lens_cpu.append(im.num_image_tokens)
16
            encoder_cached.append(
17
                self.forward_mode.is_decode()
18
                or req.prefix_len >= im.num_image_tokens
19
            )
20
    self.encoder_lens_cpu = encoder_lens_cpu
21
    self.encoder_cached = encoder_cached
22

23
    self.encoder_lens = torch.tensor(
24
        self.encoder_lens_cpu, dtype=torch.int64, pin_memory=_pin
25
    ).to(self.device, non_blocking=True)

1
# Strip encoder infos
2
pt = 0
3
decoder_out_cache_loc = []
4
encoder_out_cache_loc = []
5
extend_lens = self.extend_lens[:]
6
prefix_lens = self.prefix_lens[:]
7
for i, req in enumerate(self.reqs):
8
    encoder_len = self.encoder_lens_cpu[i]
9
    seq_lens[i] -= encoder_len
10

11
    if req.prefix_len < encoder_len:
12
        # NOTE: the encoder part should be considered as a whole
13
        assert req.prefix_len == 0
14
        input_ids[i] = input_ids[i][encoder_len:]
15
        encoder_out_cache_loc.append(self.out_cache_loc[pt : pt + encoder_len])
16
        decoder_out_cache_loc.append(
17
            self.out_cache_loc[pt + encoder_len : pt + req.extend_len]
18
        )
19
        extend_lens[i] -= encoder_len
20
        self.extend_num_tokens = self.extend_num_tokens - encoder_len
21
    else:
22
        decoder_out_cache_loc.append(
23
            self.out_cache_loc[pt : pt + req.extend_len]
24
        )
25
        prefix_lens[i] -= encoder_len

1
def general_mm_embed_routine(
2
    input_ids: torch.Tensor,
3
    forward_batch: ForwardBatch,
4
    language_model: nn.Module,
5
    multimodal_model: Optional[nn.Module] = None,
6
    data_embedding_funcs: Dict[Modality, DataEmbeddingFunc] = None,
7
    placeholder_tokens: Optional[dict[Modality, List[int]]] = None,
8
    use_deepstack: Dict[Modality, bool] = {},
9
    feature_dtypes: Optional[Dict[Modality, torch.dtype]] = None,
10
    **kwargs,
11
) -> torch.Tensor:
12
    """
13
    Process multimodal inputs and forward through language model.
14
    """
15
    input_embeds = prepare_mm_inputs(
16
        input_ids,
17
        forward_batch,
18
        language_model,
19
        kwargs,
20
        multimodal_model,
21
        data_embedding_funcs,
22
        placeholder_tokens,
23
        use_deepstack,
24
        feature_dtypes,
25
    )
26

27
    with torch.profiler.record_function("sglang.vlm.language_model_prefill"):
28
        hidden_states = language_model(
29
            input_ids=None,
30
            forward_batch=forward_batch,
31
            input_embeds=input_embeds,
32
            **kwargs,
33
        )
34
    return hidden_states

1
def embed_mm_inputs(
2
    mm_inputs_list: List[MultimodalInputs],
3
    extend_prefix_lens: List[int],
4
    extend_seq_lens: List[int],
5
    input_ids: torch.Tensor,
6
    input_embedding: nn.Embedding,
7
    multimodal_model: nn.Module = None,
8
    data_embedding_func_mapping: Dict[Modality, DataEmbeddingFunc] = None,
9
    placeholder_tokens: dict[Modality, List[int]] = None,
10
    use_deepstack: Dict[Modality, bool] = {},
11
) -> Optional[torch.Tensor]:
12
    """
13
    Embed multimodal inputs and integrate them with text token embeddings.
14
    """
15
    other_info = {}
16
    if mm_inputs_list is None:
17
        return None
18

19
    # 1. Calculate the multimodal data which exists in input_ids,
20
    # with the help of pad_values
21
    # we assume that multimodal data are represented with its pad_values
22
    # in input_ids
23
    item_flatten_list = []
24
    for mm_inputs in mm_inputs_list:
25
        item_flatten_list += [item for item in mm_inputs.mm_items if item is not None]

1
    # deepstack_embeddings: per-modality
2
    modalities, embeddings, masks, deepstack_embeddings = [], [], [], []
3

4
    # 2. Get multimodal embedding separately
5
    for modality in Modality.all():
6
        items = [
7
            item for item in item_flatten_list if item.is_modality(modality=modality)
8
        ]
9
        embedder = (
10
            None
11
            if data_embedding_func_mapping is None
12
            else data_embedding_func_mapping.get(modality, None)
13
        )
14
        if embedder is None:
15
            # "image", "video", etc
16
            modality_id = modality.name.lower()
17
            embedder = getattr(multimodal_model, f"get_{modality_id}_feature", None)
18
        if len(items) != 0:
19
            assert embedder is not None, f"no embedding method found for {modality}"