Fix circular import when sglang.srt.model_executor.runner_backend is imported first (#28002)

python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py

@@ -20,10 +20,9 @@ from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional
 
 import torch
 
-from sglang.srt.model_executor.runner.shape_key import ShapeKey
-
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
 
 
 class BaseCudaGraphBackend(ABC):

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -27,7 +27,6 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     set_graph_pool_id,
 )
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
-from sglang.srt.model_executor.runner.shape_key import ShapeKey
 from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
     BaseCudaGraphBackend,
 )
@@ -45,6 +44,7 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
         BaseCudaGraphRunner,
     )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
 
 
 class BreakableCudaGraphBackend(BaseCudaGraphBackend):

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

@@ -27,7 +27,6 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     set_graph_pool_id,
 )
-from sglang.srt.model_executor.runner.shape_key import ShapeKey
 from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
     BaseCudaGraphBackend,
 )
@@ -39,6 +38,7 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
         BaseCudaGraphRunner,
     )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
 
 
 class FullCudaGraphBackend(BaseCudaGraphBackend):

python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py

@@ -40,7 +40,6 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
 )
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 from sglang.srt.layers.utils import MultiPlatformOp
-from sglang.srt.model_executor.runner.shape_key import ShapeKey
 from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
     BaseCudaGraphBackend,
 )
@@ -54,6 +53,7 @@ if TYPE_CHECKING:
     from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
         BaseCudaGraphRunner,
     )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
     from sglang.srt.server_args import ServerArgs