From 35f2d4f76161c6bd8c5d8d18a54fd758d3bb99a4 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Mon, 20 Jul 2026 10:28:35 +0800 Subject: [PATCH] Fix no-padding CUDA graph admission (#31273) --- .../srt/model_executor/runner/decode_cuda_graph_runner.py | 8 +++++--- .../srt/speculative/eagle_draft_cuda_graph_runner.py | 2 +- .../speculative/eagle_draft_extend_cuda_graph_runner.py | 2 +- .../multi_layer_eagle_draft_extend_cuda_graph_runner.py | 2 +- 4 files changed, 8 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 68c5ee820..ec1275999 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -532,9 +532,11 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): else: cuda_graph_bs = forward_batch.batch_size - graph_key = cuda_graph_bs - if self.enable_pdmux: - graph_key = f"{get_current_stream_idx()}_{cuda_graph_bs}" + graph_key = self._make_graph_key( + cuda_graph_bs, + stream_idx=get_current_stream_idx() if self.enable_pdmux else None, + variant_label=self._resolve_lora_variant(forward_batch), + ) is_bs_supported = ( self.backend.can_run(forward_batch, graph_key) diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 3d09b2ef6..5f1f46894 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -310,7 +310,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): cuda_graph_bs = forward_batch.batch_size is_bs_supported = ( - self.backend.can_run(forward_batch, cuda_graph_bs) + self.backend.can_run(forward_batch, self._make_graph_key(cuda_graph_bs)) if self.disable_padding else cuda_graph_bs <= self.max_bs ) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index f95ae9b1e..183997310 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -307,7 +307,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): cuda_graph_bs = forward_batch.seq_lens.numel() is_bs_supported = ( - self.backend.can_run(forward_batch, cuda_graph_bs) + self.backend.can_run(forward_batch, self._make_graph_key(cuda_graph_bs)) if self.disable_padding else cuda_graph_bs <= self.max_bs ) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 07aa50dec..ea2b92135 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -218,7 +218,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): cuda_graph_bs = forward_batch.seq_lens.numel() is_bs_supported = ( - self.backend.can_run(forward_batch, cuda_graph_bs) + self.backend.can_run(forward_batch, self._make_graph_key(cuda_graph_bs)) if self.disable_padding else cuda_graph_bs <= self.max_bs )