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

[PieceWise CUDA Graph] Support awq/gptq model in piecewise cudagraph (#12518)

Browse Source
This commit is contained in:
Xiaoyu Zhang
2025-11-11 11:56:15 +08:00
committed by GitHub
parent 08c805a85f
commit 9caca6a45c
7 changed files with 339 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -163,6 +163,7 @@ class PiecewiseCudaGraphRunner:
self.model_runner.server_args.piecewise_cuda_graph_tokens,
self.model_runner.server_args.piecewise_cuda_graph_compiler,
)
self.quant_config = getattr(self.model_runner.model, "quant_config", None)
# Batch sizes to capture
self.capture_num_tokens = self.compile_config.get_capture_sizes()
@@ -263,7 +264,9 @@ class PiecewiseCudaGraphRunner:
# Attention backend
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
with set_forward_context(forward_batch, self.attention_layers):
with set_forward_context(
forward_batch, self.attention_layers, self.quant_config
):
_ = self.model_runner.model.forward(
forward_batch.input_ids,
forward_batch.positions,
@@ -405,7 +408,9 @@ class PiecewiseCudaGraphRunner:
set_is_extend_in_batch(False)
kwargs = {}
with set_forward_context(forward_batch, self.attention_layers):
with set_forward_context(
forward_batch, self.attention_layers, self.quant_config
):
self.model_runner.model.forward(
forward_batch.input_ids,
forward_batch.positions,
@@ -503,7 +508,9 @@ class PiecewiseCudaGraphRunner:
self.model_runner.attn_backend.init_forward_metadata(forward_batch)
static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
# Replay
with set_forward_context(static_forward_batch, self.attention_layers):
with set_forward_context(
static_forward_batch, self.attention_layers, self.quant_config
):
with set_compiled(True):
output = self.model_runner.model.forward(
static_forward_batch.input_ids,