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

13
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -1301,7 +1301,6 @@ def triton_mrope(
return q, k
@torch._dynamo.disable()
def triton_mrope_wrapper(
query,
key,
@@ -1428,15 +1427,18 @@ class MRotaryEmbedding(RotaryEmbedding):
dim=-1,
)
seq_len_q = query.shape[0]
query_shape = query.shape
query = query.view(num_tokens, -1, self.head_size)
query = query.view(seq_len_q, -1, self.head_size)
query_rot = query[..., : self.rotary_dim]
query_pass = query[..., self.rotary_dim :]
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
seq_len_k = key.shape[0]
key_shape = key.shape
key = key.view(num_tokens, -1, self.head_size)
key = key.view(seq_len_k, -1, self.head_size)
key_rot = key[..., : self.rotary_dim]
key_pass = key[..., self.rotary_dim :]
key_rot = _apply_rotary_emb(key_rot, cos, sin, self.is_neox_style)
@@ -1467,7 +1469,6 @@ class MRotaryEmbedding(RotaryEmbedding):
else:
return self._forward_native(positions, query, key)
@torch.compile(dynamic=True, backend=get_compiler_backend())
def _forward_triton(
self,
positions: torch.Tensor,
@@ -1502,7 +1503,9 @@ class MRotaryEmbedding(RotaryEmbedding):
return q.reshape(query_shape), k.reshape(key_shape)
query = query.view(num_tokens, -1, self.head_size)
seq_len_q = query.shape[0]
query = query.view(seq_len_q, -1, self.head_size)
query_rot = query[..., : self.rotary_dim]
query_pass = query[..., self.rotary_dim :]
query_rot = _apply_rotary_emb(query_rot, cos, sin, self.is_neox_style)

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

4
python/sglang/srt/model_executor/cuda_graph_runner.py
View File

@@ -318,6 +318,10 @@ class CudaGraphRunner:
# Graph inputs
with torch.device(self.device):
self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
self.input_embeds = torch.zeros(
(self.max_num_token, self.model_runner.model_config.hidden_size),
dtype=self.model_runner.model_config.dtype,
)
self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int32)
self.seq_lens = torch.full(
(self.max_bs,), self.seq_len_fill_value, dtype=torch.int32

11
python/sglang/srt/model_executor/model_runner.py
View File

@@ -92,6 +92,10 @@ from sglang.srt.layers.sampler import Sampler
from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model
from sglang.srt.lora.lora_manager import LoRAManager
from sglang.srt.lora.lora_registry import LoRARef
from sglang.srt.managers.mm_utils import (
external_mm_preprocess_routine,
should_use_external_mm_preprocess,
)
from sglang.srt.mem_cache.allocator import (
BaseTokenToKVPoolAllocator,
PagedTokenToKVPoolAllocator,
@@ -2139,6 +2143,13 @@ class ModelRunner:
skip_attn_backend_init: bool = False,
pp_proxy_tensors=None,
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
if self.is_multimodal and should_use_external_mm_preprocess(self.model):
forward_batch = external_mm_preprocess_routine(
forward_batch=forward_batch,
multimodal_model=self.model,
)
kwargs = {}
if self.support_pp:
kwargs["pp_proxy_tensors"] = pp_proxy_tensors

72
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
View File

@@ -166,9 +166,15 @@ class PiecewiseCudaGraphRunner:
self.max_num_tokens = max(self.capture_num_tokens)
self.use_input_embeds = model_runner.is_multimodal
# Graph inputs
with torch.device(self.device):
self.input_ids = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
self.input_embeds = torch.zeros(
(self.max_num_tokens, self.model_runner.model_config.hidden_size),
dtype=self.model_runner.dtype,
)
self.out_cache_loc = torch.zeros(
(self.max_num_tokens,), dtype=self._cache_loc_dtype()
)
@@ -176,6 +182,9 @@ class PiecewiseCudaGraphRunner:
(self.max_num_tokens,), dtype=self._cache_loc_dtype()
)
self.positions = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
self.mrope_positions = torch.zeros(
(3, self.max_num_tokens), dtype=torch.int64
)
self.tbo_plugin = TboCudaGraphRunnerPlugin()
self.attention_layers = self.model_runner.attention_layers
@@ -216,7 +225,21 @@ class PiecewiseCudaGraphRunner:
forward_batch = ForwardBatch(
forward_mode=ForwardMode.EXTEND,
batch_size=1,
input_ids=torch.randint(0, 100, (num_tokens,), device=self.device),
input_ids=(
torch.randint(0, 100, (num_tokens,), device=self.device)
if not self.use_input_embeds
else None
),
input_embeds=(
torch.randn(
num_tokens,
self.model_runner.model_config.hidden_size,
dtype=self.model_runner.dtype,
device=self.device,
)
if self.use_input_embeds
else None
),
req_pool_indices=torch.arange(1, device=self.device),
seq_lens=torch.tensor([num_tokens], device=self.device),
next_token_logits_buffer=None,
@@ -246,7 +269,7 @@ class PiecewiseCudaGraphRunner:
global_num_tokens_for_logprob_gpu=None,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=None,
mrope_positions=None,
mrope_positions=self.mrope_positions[:, :num_tokens],
spec_algorithm=None,
spec_info=None,
capture_hidden_mode=CaptureHiddenMode.NULL,
@@ -326,10 +349,17 @@ class PiecewiseCudaGraphRunner:
bs = 1
# Graph inputs
input_ids = self.input_ids[:num_tokens]
if self.use_input_embeds:
input_ids = None
input_embeds = self.input_embeds[:num_tokens]
else:
input_ids = self.input_ids[:num_tokens]
input_embeds = None
out_cache_loc = self.out_cache_loc[:num_tokens]
out_cache_loc_swa = self.out_cache_loc_swa[:num_tokens]
positions = self.positions[:num_tokens]
mrope_positions = self.mrope_positions[:, :num_tokens]
# pipeline parallelism
if self.pp_size > 1:
@@ -351,6 +381,7 @@ class PiecewiseCudaGraphRunner:
forward_mode=ForwardMode.EXTEND,
batch_size=bs,
input_ids=input_ids,
input_embeds=input_embeds,
req_pool_indices=torch.arange(bs, device=self.device),
seq_lens=torch.tensor([num_tokens], device=self.device),
next_token_logits_buffer=None,
@@ -376,7 +407,7 @@ class PiecewiseCudaGraphRunner:
global_num_tokens_for_logprob_gpu=None,
dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
global_dp_buffer_len=None,
mrope_positions=None,
mrope_positions=mrope_positions,
spec_algorithm=None,
spec_info=None,
capture_hidden_mode=CaptureHiddenMode.NULL,
@@ -428,7 +459,11 @@ class PiecewiseCudaGraphRunner:
forward_batch: ForwardBatch,
**kwargs,
):
num_tokens = len(forward_batch.input_ids)
if self.use_input_embeds:
num_tokens = forward_batch.input_embeds.shape[0]
else:
num_tokens = len(forward_batch.input_ids)
index = bisect.bisect_left(self.capture_num_tokens, num_tokens)
static_num_tokens = self.capture_num_tokens[index]
self.raw_num_tokens = num_tokens
@@ -437,7 +472,11 @@ class PiecewiseCudaGraphRunner:
self.out_cache_loc_swa.zero_()
bs = forward_batch.batch_size
self.input_ids[:num_tokens].copy_(forward_batch.input_ids)
if self.use_input_embeds:
self.input_embeds[:num_tokens].copy_(forward_batch.input_embeds)
else:
self.input_ids[:num_tokens].copy_(forward_batch.input_ids)
self.positions[:num_tokens].copy_(forward_batch.positions)
self.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc)
if forward_batch.out_cache_loc_swa is not None:
@@ -452,13 +491,32 @@ class PiecewiseCudaGraphRunner:
else None
)
if forward_batch.mrope_positions is not None:
self.mrope_positions[:, :num_tokens].copy_(forward_batch.mrope_positions)
if self.use_input_embeds:
input_ids = None
input_embeds = self.input_embeds[:static_num_tokens]
else:
input_ids = self.input_ids[:static_num_tokens]
input_embeds = None
positions = self.positions[:static_num_tokens]
out_cache_loc = self.out_cache_loc[:static_num_tokens]
mrope_positions = (
self.mrope_positions[:, :static_num_tokens]
if forward_batch.mrope_positions is not None
else None
)
next_token_logits_buffer = None
mrope_positions = None
static_forward_batch = ForwardBatch(
forward_mode=forward_batch.forward_mode,
batch_size=bs,
input_ids=input_ids,
input_embeds=input_embeds,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
next_token_logits_buffer=next_token_logits_buffer,

45
python/sglang/srt/models/qwen2_5_vl.py
View File

@@ -57,11 +57,12 @@ from sglang.srt.layers.pooler import Pooler, PoolingType
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer, get_layer_id
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.managers.mm_utils import (
MultiModalityDataPaddingPatternMultimodalTokens,
general_mm_embed_routine,
from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
from sglang.srt.managers.schedule_batch import (
Modality,
MultimodalDataItem,
MultimodalInputs,
)
from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.qwen2 import Qwen2Model
@@ -566,6 +567,25 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module):
video_embeds = self.visual(pixel_values, grid_thw=video_grid_thw)
return video_embeds
def post_process(
self,
inputs_embeds,
modalities: List[Modality],
embeddings: List[torch.Tensor],
indices: List[torch.Tensor],
forward_batch: ForwardBatch,
) -> torch.Tensor:
# Placeholder for post_process
new_embeddings = []
for i, (modality, embedding, index) in enumerate(
zip(modalities, embeddings, indices)
):
if embedding is None or index is None:
continue
new_embeddings.append(embedding)
return new_embeddings, forward_batch
def get_input_embeddings(self):
return self.model.embed_tokens
@@ -575,6 +595,7 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module):
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds=None,
get_embedding: bool = False,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
):
@@ -603,11 +624,21 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module):
f"(3, seq_len) positions, but got {positions.size()}"
)
hidden_states = general_mm_embed_routine(
input_embeds = forward_batch.input_embeds
# It may seem strange to assign input_embeds again even after passing it as an argument.
# This is for compatibility considerations.
# In the 'extend' scenario, this forward function is called from two places:
# 1. model_runner calls forward directly,
# 2. piece_wise_cuda_graph_runner calls forward and replay.
# Currently,
# In 'extend', input_embeds is passed in.
# In 'decode', input_ids is passed in.
hidden_states = self.model(
input_ids=input_ids,
forward_batch=forward_batch,
language_model=self.model,
multimodal_model=self,
input_embeds=input_embeds,
positions=positions,
pp_proxy_tensors=pp_proxy_tensors,
)

23
python/sglang/srt/models/qwen2_vl.py
View File

@@ -39,10 +39,7 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.pooler import Pooler, PoolingType
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.managers.mm_utils import (
MultiModalityDataPaddingPatternMultimodalTokens,
general_mm_embed_routine,
)
from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternMultimodalTokens
from sglang.srt.managers.schedule_batch import MultimodalDataItem, MultimodalInputs
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.model_loader.weight_utils import default_weight_loader
@@ -509,6 +506,7 @@ class Qwen2VLForConditionalGeneration(nn.Module):
input_ids: torch.Tensor,
positions: torch.Tensor,
forward_batch: ForwardBatch,
input_embeds=None,
get_embedding: bool = False,
):
"""Run forward pass for Qwen2-VL.
@@ -535,11 +533,22 @@ class Qwen2VLForConditionalGeneration(nn.Module):
"multimodal section rotary embedding requires "
f"(3, seq_len) positions, but got {positions.size()}"
)
hidden_states = general_mm_embed_routine(
input_embeds = forward_batch.input_embeds
# It may seem strange to assign input_embeds again even after passing it as an argument.
# This is for compatibility considerations.
# In the 'extend' scenario, this forward function is called from two places:
# 1. model_runner calls forward directly,
# 2. piece_wise_cuda_graph_runner calls forward and replay.
# Currently,
# In 'extend', input_embeds is passed in.
# In 'decode', input_ids is passed in.
hidden_states = self.model(
input_ids=input_ids,
forward_batch=forward_batch,
language_model=self.model,
multimodal_model=self,
input_embeds=input_embeds,
positions=positions,
)