fix(moe): import capture mode from cuda graph runner

Tested: PYTHONPYCACHEPREFIX=/private/tmp/sglang_pycache python3 -m py_compile python/sglang/srt/layers/moe/mega_moe.py test/registered/unit/moe/test_glm_megamoe.py test/registered/unit/server_args/test_server_args.py
This commit is contained in:
LuminolT
2026-07-06 14:13:55 +08:00
parent c00630088c
commit 0d3bf6b5c2

View File

@@ -26,7 +26,7 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers.dp_attention import get_dp_global_num_tokens from sglang.srt.layers.dp_attention import get_dp_global_num_tokens
from sglang.srt.layers.moe.topk import select_experts from sglang.srt.layers.moe.topk import select_experts
from sglang.srt.layers.moe.utils import get_moe_a2a_backend from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
if TYPE_CHECKING: if TYPE_CHECKING:
from deep_gemm import SymmBuffer from deep_gemm import SymmBuffer