From 0132848349585cfe6aae51c4941cbae872505f8a Mon Sep 17 00:00:00 2001 From: sky Date: Fri, 28 Aug 2026 05:19:27 +0800 Subject: [PATCH] Revert "[Fix] Disable --enable-symm-mem under CUDA graphs on Kimi hybrid models" (#34842) Signed-off-by: wangfakang Co-authored-by: Claude --- python/sglang/srt/arg_groups/kimi_k3_hook.py | 48 -------------------- python/sglang/srt/server_args.py | 4 -- 2 files changed, 52 deletions(-) diff --git a/python/sglang/srt/arg_groups/kimi_k3_hook.py b/python/sglang/srt/arg_groups/kimi_k3_hook.py index 464df72d0..1b86f4427 100644 --- a/python/sglang/srt/arg_groups/kimi_k3_hook.py +++ b/python/sglang/srt/arg_groups/kimi_k3_hook.py @@ -55,54 +55,6 @@ def apply_kimi_k3_spec_backend_defaults(server_args: ServerArgs) -> None: ) -def disable_kimi_k3_symm_mem(server_args: ServerArgs) -> None: - """Turn `--enable-symm-mem` back off unless every phase runs eager. - - Symm-mem allocations are per-forward, so an address captured into a graph is - neither reserved for its lifetime nor at the same offset on every rank. Under - capture that corrupts spec decode: accept collapses to 1.000, or the server - silently emits garbage with accept pinned at the ceiling. Prefill counts too -- - the same allocation sits in any captured RowParallelLinear. - - Gates on the arch itself: this runs from cuda-graph resolution, which is earlier - than the model-specific hook block. - """ - cfg = resolving_view(server_args) - from sglang.srt.connector import ConnectorType - from sglang.srt.model_executor.cuda_graph_config import Backend - from sglang.srt.utils import parse_connector_type - - if not cfg.enable_symm_mem: - return - if parse_connector_type(cfg.model_path) == ConnectorType.INSTANCE: - return - if server_args.get_model_config().hf_config.architectures[0] not in ( - "KimiLinearForCausalLM", - "KimiK3ForConditionalGeneration", - ): - return - graph = cfg.cuda_graph_config - if ( - graph.decode.backend == Backend.DISABLED - and graph.prefill.backend == Backend.DISABLED - ): - return - declare_resolution( - server_args, - "disable_kimi_k3_symm_mem", - enable_symm_mem=False, - ) - logger.warning( - "Kimi hybrid model: ignoring --enable-symm-mem because CUDA graphs are on. " - "The symmetric-memory pool's per-forward allocations are not valid for the " - "lifetime of a captured graph, which corrupts speculative decoding and can " - "silently produce wrong output. The auto-probed K3 fused all-reduce is faster " - "anyway. Disable capture on every phase " - "(--cuda-graph-backend-decode=disabled --cuda-graph-backend-prefill=disabled) " - "if you genuinely need symmetric memory." - ) - - def apply_kimi_k3_linear_attn_defaults(server_args: ServerArgs) -> None: """KDA decode-fallback default for Kimi hybrid models (spec-independent).""" cfg = resolving_view(server_args) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 57a49f665..d9c4f4a64 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4840,12 +4840,8 @@ class ServerArgs: def _handle_cuda_graph_config(self): cfg = resolving_view(self) - from sglang.srt.arg_groups.kimi_k3_hook import disable_kimi_k3_symm_mem self._parse_cuda_graph_config() - # Reads the resolved per-phase backends; must precede the compat rules - # below and _handle_gpu_memory_settings, which key off enable_symm_mem. - disable_kimi_k3_symm_mem(self) self._apply_cuda_graph_compatibility() self._apply_deepep_adjustments() self._apply_cuda_graph_disaggregation_roles()