minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Decouple ragged CUDA graph request and token capacities (#37300)

Browse Source
This commit is contained in:
Yichao Fu
2026-09-01 18:48:11 -07:00
committed by GitHub
parent dde0ecdb90
commit 66de38f30c
1 changed files with 6 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+6 -2
View File
@@ -338,6 +338,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
self.capture_bs, self.compile_bs = get_batch_sizes_to_capture( self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(
model_runner, self.captured_req_width model_runner, self.captured_req_width
) )
self.max_bs = max(self.capture_bs)
if KTRANSFORMERS_AVAILABLE: if KTRANSFORMERS_AVAILABLE:
KTMoEWrapper.set_capture_batch_sizes(self.capture_bs) KTMoEWrapper.set_capture_batch_sizes(self.capture_bs)
@@ -368,8 +369,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
) )
# Attention backend # Attention backend
self.max_bs = max(self.capture_bs) self.max_num_token = (
self.max_num_token = self.max_bs * self.captured_req_width max(self.capture_num_tokens)
if self.capture_num_tokens is not None
else self.max_bs * self.captured_req_width
)
self.attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token) self.attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token)
# Init PDMux if needed # Init PDMux if needed