Improve CUDA graph and speculative execution output handling (#37329)

Co-authored-by: jiayisuse <jiayisuse@fb.com>
Co-authored-by: Yinghai Lu <yinghai@meta.com>
Co-authored-by: Hao Zhang <zhisbug@users.noreply.github.com>
Co-authored-by: Yichao Fu <yichaofu@meta.com>

python/sglang/srt/model_executor/model_runner.py

@@ -124,8 +124,10 @@ from sglang.srt.model_executor.model_runner_components.kv_pool_runtime import (
     is_post_capture_kv_active,
 )
 from sglang.srt.model_executor.model_runner_components.layer_setup import (
+    AttentionAndMoeLayers,
     ModelLayerInfo,
     adjust_hybrid_swa_layer_ids,
+    compute_attention_and_moe_layers,
     resolve_layer_indices,
 )
 from sglang.srt.model_executor.model_runner_components.load_model_utils import (
@@ -1435,6 +1437,10 @@ class ModelRunner:
 
         return DecodeCudaGraphRunner
 
+    def get_cuda_graph_layers(self, layer_model) -> AttentionAndMoeLayers:
+        """Return the model layers used by prefill CUDA graph execution."""
+        return compute_attention_and_moe_layers(layer_model)
+
     def init_decode_cuda_graph(self):
         self.decode_cuda_graph_runner = None
         capture = capture_decode_graph(model_runner=self)

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

@@ -31,9 +31,6 @@ from sglang.srt.model_executor.graph_memory_usage import (
 )
 from sglang.srt.model_executor.graph_shared_output import GraphSharedOutput
 from sglang.srt.model_executor.hook_manager import register_forward_hooks
-from sglang.srt.model_executor.model_runner_components.layer_setup import (
-    compute_attention_and_moe_layers,
-)
 from sglang.srt.model_executor.runner import (
     EagerRunner,
     PrefillCudaGraphRunner,
@@ -444,7 +441,7 @@ def capture_prefill_graph(
         model_runner.moe_fusions,
         model_runner.dsa_indexers,
         model_runner.mha_companion_layers,
-    ) = compute_attention_and_moe_layers(layer_model)
+    ) = model_runner.get_cuda_graph_layers(layer_model)
     (
         model_runner.attention_layers,
         model_runner.mha_companion_layers,