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

[piecewise] Refactor VLM to support input embed buffer and remove external embedder hack (#14155)

Browse Source
This commit is contained in:
Byron Hsu
2025-11-30 21:43:09 -08:00
committed by GitHub
parent 0b9dbea593
commit 0825d7f4c6
7 changed files with 145 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -97,11 +97,6 @@ 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,
resolve_external_mm_data_embedding_funcs,
should_use_external_mm_preprocess,
)
from sglang.srt.mem_cache.allocator import (
BaseTokenToKVPoolAllocator,
PagedTokenToKVPoolAllocator,
@@ -2548,15 +2543,6 @@ class ModelRunner:
skip_attn_backend_init: bool = False,
pp_proxy_tensors=None,
) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
if self.is_multimodal and should_use_external_mm_preprocess(self.model):
data_embedding_funcs = resolve_external_mm_data_embedding_funcs(self.model)
forward_batch = external_mm_preprocess_routine(
forward_batch=forward_batch,
multimodal_model=self.model,
data_embedding_funcs=data_embedding_funcs,
)
kwargs = {}
if self.support_pp:
kwargs["pp_proxy_tensors"] = pp_proxy_tensors

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

@@ -70,15 +70,38 @@ def disable_ca_comm(tp_group):
TODO(yuwei): Fix this
"""
old_disabled = None
if tp_group.ca_comm is None:
yield
return
original_disabled = tp_group.ca_comm.disabled
tp_group.ca_comm.original_disabled = original_disabled
try:
if tp_group.ca_comm is not None:
old_disabled = tp_group.ca_comm.disabled
tp_group.ca_comm.disabled = True
tp_group.ca_comm.disabled = True
yield
finally:
if tp_group.ca_comm is not None and old_disabled is not None:
tp_group.ca_comm.disabled = old_disabled
tp_group.ca_comm.disabled = original_disabled
@contextmanager
def use_original_ca_comm(tp_group):
"""
For the module not in piecewise cuda graph capture, use the original custom allreduce communication.
This is a no-op if not using piecewise cuda graph because .disabled == .original_disabled
TODO(Byron): remove this once custom allreduce is enabled in piecewise cuda graph
"""
if tp_group.ca_comm is None:
yield
return
current_disabled = tp_group.ca_comm.disabled
original_disabled = tp_group.ca_comm.original_disabled
try:
tp_group.ca_comm.disabled = original_disabled
yield
finally:
tp_group.ca_comm.disabled = current_disabled
@contextmanager
@@ -189,15 +212,11 @@ class PiecewiseCudaGraphRunner:
self.max_num_tokens = max(self.capture_num_tokens)
self.use_input_embeds = model_runner.is_multimodal
self.is_multimodal = 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()
)
@@ -207,11 +226,23 @@ class PiecewiseCudaGraphRunner:
else None
)
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()
if (
self.is_multimodal
): # Only create input_embeds and mrope_positions for multimodal model to save memory
# 1. In multimodal, we only compile and capture the language model part.
# 2. The embedder is outside of the graph, but cuda graph requires the input embeds to have a fixed memory address.
# 3. Input embeds is a pre-allocated buffer. In model.forward, we copy the embed output to this buffer.
self.input_embeds = torch.zeros(
(self.max_num_tokens, self.model_runner.model_config.hidden_size),
dtype=self.model_runner.dtype,
)
self.mrope_positions = torch.zeros(
(3, self.max_num_tokens), dtype=torch.int64
)
self.attention_layers = self.model_runner.attention_layers
if get_global_graph_memory_pool() is None:
@@ -258,11 +289,7 @@ class PiecewiseCudaGraphRunner:
forward_batch = ForwardBatch(
forward_mode=ForwardMode.EXTEND,
batch_size=1,
input_ids=(
torch.randint(0, 100, (num_tokens,), device=self.device)
if not self.use_input_embeds
else None
),
input_ids=(torch.randint(0, 100, (num_tokens,), device=self.device)),
input_embeds=(
torch.randn(
num_tokens,
@@ -270,7 +297,7 @@ class PiecewiseCudaGraphRunner:
dtype=self.model_runner.dtype,
device=self.device,
)
if self.use_input_embeds
if self.is_multimodal
else None
),
req_pool_indices=torch.arange(1, device=self.device),
@@ -298,7 +325,9 @@ 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=self.mrope_positions[:, :num_tokens],
mrope_positions=(
self.mrope_positions[:, :num_tokens] if self.is_multimodal else None
),
spec_algorithm=None,
spec_info=None,
capture_hidden_mode=CaptureHiddenMode.NULL,
@@ -374,12 +403,8 @@ class PiecewiseCudaGraphRunner:
bs = 1
# Graph inputs
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
input_ids = self.input_ids[:num_tokens]
input_embeds = self.input_embeds[:num_tokens] if self.is_multimodal else None
out_cache_loc = self.out_cache_loc[:num_tokens]
out_cache_loc_swa = (
@@ -388,13 +413,9 @@ class PiecewiseCudaGraphRunner:
else None
)
positions = self.positions[:num_tokens]
mrope_positions = self.mrope_positions[:, :num_tokens]
# pipeline parallelism
if self.pp_size > 1:
pp_proxy_tensors = PPProxyTensors(
{k: v[:num_tokens] for k, v in self.pp_proxy_tensors.items()}
)
mrope_positions = (
self.mrope_positions[:, :num_tokens] if self.is_multimodal else None
)
global_dp_buffer_len = None
@@ -488,11 +509,7 @@ class PiecewiseCudaGraphRunner:
forward_batch: ForwardBatch,
**kwargs,
):
if self.use_input_embeds:
num_tokens = forward_batch.input_embeds.shape[0]
else:
num_tokens = len(forward_batch.input_ids)
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
@@ -502,11 +519,7 @@ class PiecewiseCudaGraphRunner:
self.out_cache_loc_swa.zero_()
bs = forward_batch.batch_size
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.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 self.out_cache_loc_swa is not None:
@@ -528,12 +541,10 @@ class PiecewiseCudaGraphRunner:
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
input_ids = self.input_ids[:static_num_tokens]
input_embeds = (
self.input_embeds[:static_num_tokens] if self.is_multimodal else None
)
mrope_positions = (
self.mrope_positions[:, :static_num_tokens]