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

Share one logits output buffer across prefill/decode/draft cuda-graph runners (#29779)

Browse Source
This commit is contained in:
cctry
2026-07-01 23:30:56 -07:00
committed by GitHub
parent 926140d789
commit 697b400d70
9 changed files with 130 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py
+6
View File
@@ -17,6 +17,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.graph_shared_output import GraphSharedOutput
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from sglang.srt.server_args import set_global_server_args_for_scheduler
@@ -403,6 +404,11 @@ class MockModelRunner(ModelRunner):
self.is_hybrid_swa = case.sliding_window_size is not None
self.sliding_window_size = case.sliding_window_size
self.use_mla_backend = False
# Runner-mode helpers mutate speculative graph sizes after construction.
self.graph_shared_output = GraphSharedOutput(
device=self.device,
max_rows=pool_batch_size * max_context_len,
)
@property
def hybrid_gdn_config(self):
python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py
+6
View File
@@ -21,6 +21,7 @@ from sglang.srt.model_executor.forward_context import (
forward_context,
get_token_to_kv_pool,
)
from sglang.srt.model_executor.graph_shared_output import GraphSharedOutput
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_parallel
from sglang.srt.server_args import set_global_server_args_for_scheduler
@@ -308,6 +309,11 @@ class MockMLAModelRunner(ModelRunner):
self.use_mla_backend = True
self.is_draft_worker = False
self._kernel_warmed_up = True
# Runner-mode helpers mutate speculative graph sizes after construction.
self.graph_shared_output = GraphSharedOutput(
device=self.device,
max_rows=pool_batch_size * max_context_len,
)
@property
def hybrid_gdn_config(self):