chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[VLM] Support Piecewise CUDA Graph for Qwen2.5-VL (#13055)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: Yuhao Yang <yhyang201@gmail.com>
This commit is contained in:
Yuan Luo
2025-11-20 10:23:44 +08:00
committed by GitHub
parent 67fca6b297
commit af6bcadcf7
10 changed files with 710 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

261
python/sglang/srt/managers/mm_utils.py
View File

@@ -493,7 +493,7 @@ def get_embedding_and_mask(
return embedding, special_multimodal_mask
def embed_mm_inputs(
def general_embed_mm_inputs(
mm_inputs_list: List[MultimodalInputs],
extend_prefix_lens: List[int],
extend_seq_lens: List[int],
@@ -679,7 +679,7 @@ def general_mm_embed_routine(
for i, seq_len in enumerate(forward_batch.extend_seq_lens_cpu)
if forward_batch.mm_inputs[i] is not None
]
inputs_embeds, other_info = embed_mm_inputs(
inputs_embeds, other_info = general_embed_mm_inputs(
mm_inputs_list=mm_inputs_list,
extend_prefix_lens=extend_prefix_lens,
extend_seq_lens=extend_seq_lens,
@@ -816,3 +816,260 @@ def hash_feature(f):
reconstruct_t = f.reconstruct_on_target_device(torch.cuda.current_device())
return tensor_hash([reconstruct_t])
return data_hash(f)
def resolve_language_model(multimodal_model: nn.Module) -> Optional[nn.Module]:
# Qwen2-VL / Qwen3-VL Style
if hasattr(multimodal_model, "model"):
lm = getattr(multimodal_model, "model")
if hasattr(lm, "get_input_embeddings"):
return lm
# Llava / OneVision Style
if hasattr(multimodal_model, "language_model"):
lm = getattr(multimodal_model, "language_model")
if hasattr(lm, "get_input_embeddings"):
return lm
if hasattr(multimodal_model, "get_input_embeddings"):
return multimodal_model
return None
def external_embed_mm_inputs(
forward_batch: ForwardBatch,
mm_inputs_list: List[MultimodalInputs],
extend_prefix_lens: List[int],
extend_seq_lens: List[int],
input_ids: torch.Tensor,
input_embedding: nn.Embedding,
multimodal_model: nn.Module = None,
data_embedding_func_mapping: Dict[
Modality, Callable[[List[MultimodalDataItem]], torch.Tensor]
] = None,
) -> Optional[torch.Tensor]:
"""
Embed multimodal inputs and integrate them with text token embeddings.
Args:
mm_inputs_list: List of multimodal inputs to process
extend_prefix_lens: Prefix lengths for each request
extend_seq_lens: Sequence lengths for each request
input_ids: Input token IDs tensor
input_embedding: Embedding layer for text tokens
Returns:
Combined embedding tensor with multimodal content integrated
"""
if mm_inputs_list is None:
return None
# 1. Calculate the multimodal data which exists in input_ids, with the help of pad_values
# we assume that multimodal data are represented with its pad_values in input_ids
item_flatten_list = []
for mm_inputs in mm_inputs_list:
item_flatten_list += [item for item in mm_inputs.mm_items if item is not None]
modalities, embeddings, masks = [], [], []
# 2. Get multimodal embedding separately
# Try get mm embedding if any
for modality in Modality.all():
items = [
item for item in item_flatten_list if item.is_modality(modality=modality)
]
embedder = (
None
if data_embedding_func_mapping is None
else data_embedding_func_mapping.get(modality, None)
)
if embedder is None:
# "image", "video", etc
modality_id = modality.name.lower()
embedder = getattr(multimodal_model, f"get_{modality_id}_feature", None)
if len(items) != 0:
assert embedder is not None, f"no embedding method found for {modality}"
placeholder_tensor = torch.as_tensor(
[item.pad_value for item in items],
device=input_ids.device,
)
# calculate per request items length offset
items_size = torch.zeros(len(mm_inputs_list) + 1, dtype=int)
items_offsets = []
for i, mm_inputs in enumerate(mm_inputs_list):
mm_items = [
item
for item in mm_inputs.mm_items
if item.is_modality(modality=modality)
]
items_size[i + 1] = len(mm_items)
items_offsets.append(
flatten_nested_list([item.offsets for item in mm_items])
)
items_size = torch.cumsum(items_size, dim=0).tolist()
embedding, mask = get_embedding_and_mask(
data_embedding_func=embedder,
embedding_items=items,
placeholder_tensor=placeholder_tensor,
input_ids=input_ids,
items_size=items_size,
prefix_length=extend_prefix_lens,
extend_length=extend_seq_lens,
items_offset_list=items_offsets,
)
modalities += [modality]
embeddings += [embedding]
masks += [mask]
# 3. Get input embeddings
vocab_size = input_embedding.num_embeddings
# Important: clamp after getting original multimodal regions
# Clamp input ids. This is because the input_ids for the multimodal tokens are
# filled with the hash values of the multimodal for the prefix matching in the radix attention.
# There values are useless because their embeddings will be replaced by vision embeddings anyway.
input_ids.clamp_(min=0, max=vocab_size - 1)
inputs_embeds = input_embedding(input_ids)
indices = []
for mask in masks:
if mask is not None:
indices.append(torch.where(mask.squeeze(dim=-1))[0])
else:
indices.append(None)
# only for qwen3vl right now, replace the original use_deepstack with this method.
if hasattr(multimodal_model, "post_process"):
embeddings, forward_batch = multimodal_model.post_process(
inputs_embeds, modalities, embeddings, indices, forward_batch
)
# 4. scatter embeddings into input embedding
for i, modality, embedding, index in zip(
range(len(embeddings)), modalities, embeddings, indices
):
if embedding is None or index is None:
continue
# in-place update
inputs_embeds[index] = embedding.to(inputs_embeds.device, inputs_embeds.dtype)
return inputs_embeds, forward_batch
def should_use_external_mm_preprocess(multimodal_model: nn.Module) -> bool:
"""Decide whether we should use our generic "external_mm_preprocess_routine".
We are adapting VLM for piecewise CUDA graph. Since the encoder's forward
pass cannot be executed within the model's forward pass, we need to
precompute image embeddings using the encoder within the model runner.
For models that have already been adjusted, there is a member called
should_use_external_mm_preprocess, which is set to True. In practice,
the external_mm_preprocess_routine function will be called in the
model_runner.forward_extend to handle multimodal inputs.
For models that have not yet been adapted, the general_mm_embed_routine
will still be called in the model class's forward function for processing.
Current strategy:
- Llava family (models with vision_tower + multi_modal_projector):
Their forward already calls general_mm_embed_routine and includes
built-in multimodal processing. If we run it again in ModelRunner,
it will conflict with the internal logic, so we skip it here.
- Others (such as Qwen2-VL / Qwen2.5-VL): use the multimodal
preprocessing.
"""
cls_name = multimodal_model.__class__.__name__
qwen_vl_classes = {
"Qwen2VLForConditionalGeneration",
"Qwen2_5_VLForConditionalGeneration",
}
return cls_name in qwen_vl_classes
def external_mm_preprocess_routine(
forward_batch: ForwardBatch,
multimodal_model: Optional[nn.Module] = None,
data_embedding_funcs: Dict[
Modality, Callable[[List[MultimodalDataItem]], torch.Tensor]
] = None,
) -> torch.Tensor:
"""
Process multimodal inputs and forward through language model.
Args:
input_ids: Input token IDs tensor
forward_batch: Batch information for model forward pass
data_embedding_funcs: A dictionary mapping from modality type to the corresponding embedding function.
**kwargs: Additional arguments passed to language model
Returns:
Hidden states from language model forward pass
"""
language_model = resolve_language_model(multimodal_model)
if language_model is None:
raise ValueError(
f"Cannot resolve language model from {type(multimodal_model).__name__}. "
f"Please ensure the model has 'model' or 'language_model' attribute."
)
assert hasattr(language_model, "get_input_embeddings")
embed_tokens = language_model.get_input_embeddings()
if not hasattr(language_model, "pp_group") or language_model.pp_group.is_first_rank:
input_ids = forward_batch.input_ids
if (
not forward_batch.forward_mode.is_decode()
and not forward_batch.forward_mode.is_target_verify()
and forward_batch.contains_mm_inputs()
):
mm_inputs_list = [
mm_input for mm_input in forward_batch.mm_inputs if mm_input is not None
]
extend_prefix_lens = [
prefix_len
for i, prefix_len in enumerate(forward_batch.extend_prefix_lens_cpu)
if forward_batch.mm_inputs[i] is not None
]
extend_seq_lens = [
seq_len
for i, seq_len in enumerate(forward_batch.extend_seq_lens_cpu)
if forward_batch.mm_inputs[i] is not None
]
input_embeds, forward_batch = external_embed_mm_inputs(
forward_batch=forward_batch,
mm_inputs_list=mm_inputs_list,
extend_prefix_lens=extend_prefix_lens,
extend_seq_lens=extend_seq_lens,
input_ids=forward_batch.input_ids,
multimodal_model=multimodal_model,
input_embedding=embed_tokens,
data_embedding_func_mapping=data_embedding_funcs,
)
# once used, mm_inputs is useless, considering chunked-prefill is disabled for multimodal models
# just being defensive here
forward_batch.mm_inputs = None
else:
# NOTE: This may reduce the performance for only-text inputs.
# Using a fixed-address buffer might be better, though it could be a bit dirty.
input_embeds = embed_tokens(input_ids)
# only for qwen3vl
if getattr(multimodal_model, "use_deepstack", False):
forward_batch.input_deepstack_embeds = torch.zeros(
(
len(input_ids),
multimodal_model.config.hidden_size
* len(multimodal_model.deepstack_visual_indexes),
),
device=input_embeds.device,
dtype=input_embeds.dtype,
)
forward_batch.input_embeds = input_embeds
else:
forward_batch.input_embeds = None
return forward_batch