diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 9811e663c..73e75293f 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -1905,11 +1905,17 @@ multimodal_piecewise_cuda_graph_supported_model_archs = [ # Multimodal archs whose LM prefill is validated under breakable CUDA graph; # embed-carrying batches are rejected at replay (can_run_graph) and run eager. +# The Kimi archs are structurally multimodal -- their configs always carry a +# vision_config, so is_multimodal is True even for text-only serving -- and the +# generic multimodal rule disabled prefill CG for them despite the LM prefill +# capturing cleanly. multimodal_breakable_cuda_graph_supported_model_archs = [ "InternS2MobiusForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", "MuseGlimmerForConditionalGeneration", + "KimiK3ForConditionalGeneration", + "KimiK25ForConditionalGeneration", ] if external_mm_model_arch := envs.SGLANG_EXTERNAL_MM_MODEL_ARCH.get(): diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 18d16fa28..ed46b25d2 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -4520,13 +4520,14 @@ class ServerArgs: if (Phase.PREFILL, "backend") in self._cuda_graph_config_locked: return - # Breakable is the general CUDA default, but it is not compatible with - # multimodal prefill. Models on this allowlist have had their decoder - # prefill validated under tc_piecewise; the vision encoder remains - # eager outside the captured LM forward. + # Breakable is the CUDA default but not multimodal-compatible; + # piecewise-allowlisted archs run their validated decoder prefill + # there instead. Archs also on the breakable allowlist keep it -- + # this runs first, so piecewise would otherwise silently win. if ( self.cuda_graph_config.prefill.backend == Backend.BREAKABLE and self.get_model_config().is_multimodal_piecewise_cuda_graph_supported + and not self.get_model_config().is_multimodal_breakable_cuda_graph_supported # Keep trtllm_mla on the preferred breakable path, which now serves # MLA by falling back to the flashinfer MLA impl for extend. and self._resolved_attention_backends()[0] != "trtllm_mla" diff --git a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py index 6591fe1c1..b81df8b7b 100644 --- a/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py +++ b/test/registered/unit/configs/test_multimodal_piecewise_cuda_graph.py @@ -72,7 +72,8 @@ class TestMultimodalPiecewiseCudaGraph(CustomTestCase): def test_supported_multimodal_model_upgrades_default_to_tc_piecewise(self): args = ServerArgs(model_path="dummy") args.model_config = SimpleNamespace( - is_multimodal_piecewise_cuda_graph_supported=True + is_multimodal_piecewise_cuda_graph_supported=True, + is_multimodal_breakable_cuda_graph_supported=False, ) args.cuda_graph_config = CudaGraphConfig( prefill=PhaseConfig(backend=Backend.BREAKABLE)