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

[piecewise] move piecewise_cuda_graph_runner init to model_runner initialize (#14034)

Browse Source 
Co-authored-by: Stefan He <hebiaobuaa@gmail.com>
Co-authored-by: Binyao Jiang <byjiang1996@gmail.com>
This commit is contained in:
Minglei Zhu
2025-11-30 22:16:04 -08:00
committed by GitHub
parent 6df1e8d628
commit e8542db558
1 changed files with 55 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -367,40 +367,6 @@ class ModelRunner:
self._model_update_group = {}
self._weights_send_group = {}
if (
self.server_args.enable_piecewise_cuda_graph
and self.can_run_piecewise_cuda_graph()
):
self.attention_layers = []
for layer in self.model.model.layers:
if hasattr(layer, "self_attn"):
if hasattr(layer.self_attn, "attn"):
self.attention_layers.append(layer.self_attn.attn)
elif hasattr(layer.self_attn, "attn_mqa"):
# For DeepSeek model
self.attention_layers.append(layer.self_attn.attn_mqa)
# For hybrid model
elif hasattr(layer, "attn"):
self.attention_layers.append(layer.attn)
elif hasattr(layer, "linear_attn"):
self.attention_layers.append(layer.linear_attn)
# For InternVL model
elif hasattr(layer, "attention"):
if hasattr(layer.attention, "attn"):
self.attention_layers.append(layer.attention.attn)
if len(self.attention_layers) < self.model_config.num_hidden_layers:
# TODO(yuwei): support Non-Standard GQA
log_info_on_rank0(
logger,
"Disable piecewise CUDA graph because some layers do not apply Standard GQA",
)
self.piecewise_cuda_graph_runner = None
else:
self.piecewise_cuda_graph_runner = PiecewiseCudaGraphRunner(self)
else:
self.piecewise_cuda_graph_runner = None
def init_mindspore_runner(self):
# Init the mindspore runner
# for now, there is only some communication initialization work
@@ -578,6 +544,9 @@ class ModelRunner:
self.model.set_eagle3_layers_to_capture(eagle_aux_hidden_state_layer_ids)
# Initialize piecewise CUDA graph
self.init_piecewise_cuda_graphs()
def model_specific_adjustment(self):
server_args = self.server_args
@@ -2470,6 +2439,58 @@ class ModelRunner:
f"mem usage={self.graph_mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
)
def init_piecewise_cuda_graphs(self):
"""Initialize piecewise CUDA graph runner."""
self.piecewise_cuda_graph_runner = None
if (
not self.server_args.enable_piecewise_cuda_graph
or not self.can_run_piecewise_cuda_graph()
):
return
# Collect attention layers from the model
self.attention_layers = []
for layer in self.model.model.layers:
if hasattr(layer, "self_attn"):
if hasattr(layer.self_attn, "attn"):
self.attention_layers.append(layer.self_attn.attn)
elif hasattr(layer.self_attn, "attn_mqa"):
# For DeepSeek model
self.attention_layers.append(layer.self_attn.attn_mqa)
# For hybrid model
elif hasattr(layer, "attn"):
self.attention_layers.append(layer.attn)
elif hasattr(layer, "linear_attn"):
self.attention_layers.append(layer.linear_attn)
# For InternVL model
elif hasattr(layer, "attention"):
if hasattr(layer.attention, "attn"):
self.attention_layers.append(layer.attention.attn)
if len(self.attention_layers) < self.model_config.num_hidden_layers:
# TODO(yuwei): support Non-Standard GQA
log_info_on_rank0(
logger,
"Disable piecewise CUDA graph because some layers do not apply Standard GQA",
)
return
tic = time.perf_counter()
before_mem = get_available_gpu_memory(self.device, self.gpu_id)
logger.info(
f"Capture piecewise CUDA graph begin. avail mem={before_mem:.2f} GB"
)
self.piecewise_cuda_graph_runner = PiecewiseCudaGraphRunner(self)
after_mem = get_available_gpu_memory(self.device, self.gpu_id)
mem_usage = before_mem - after_mem
logger.info(
f"Capture piecewise CUDA graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. "
f"mem usage={mem_usage:.2f} GB. avail mem={after_mem:.2f} GB."
)
def init_threads_binding(self):
omp_cpuids = os.environ.get("SGLANG_CPU_OMP_THREADS_BIND", "all")
cpu_ids_by_node = get_cpu_ids_by_node()