From 3f1e5c58965bde4a62a63ec466dfa509514d666a Mon Sep 17 00:00:00 2001 From: ThomasX Date: Mon, 11 May 2026 02:51:22 +0800 Subject: [PATCH] fix: initialize cp_shared_kv_layout for draft model runners The CpSharedKVLayout initialization was gated behind `if not self.is_draft_worker`, leaving the draft model runner with `uses_cp_shared_kv=True` but `cp_shared_kv_layout=None`. This caused an AssertionError in NSAIndexer._filter_shared_index_write during DRAFT_EXTEND under CP shared KV + MTP. Move the layout init after the `is_draft_worker` guard so both target and draft model runners have it available. --- .../sglang/srt/model_executor/model_runner.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 4ae426b3c..02330c0e9 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -922,14 +922,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): server_args=self.server_args, model_config=self.model_config, ) - if self.uses_cp_shared_kv: - from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout - - self.cp_shared_kv_layout = CpSharedKVLayout( - page_size=self.page_size, - cp_size=self.server_args.attn_cp_size, - cp_rank=get_attention_cp_rank(), - ) if is_npu(): register_sgl_tp_rank(self.gpu_id) @@ -952,6 +944,15 @@ class ModelRunner(ModelRunnerKVCacheMixin): f"(tp_size={self.tp_size}, pp_size={self.pp_size}, ep_size={self.moe_ep_size})" ) + if self.uses_cp_shared_kv: + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + self.cp_shared_kv_layout = CpSharedKVLayout( + page_size=self.page_size, + cp_size=self.server_args.attn_cp_size, + cp_rank=get_attention_cp_rank(), + ) + pre_model_load_memory = get_available_gpu_memory( self.device, self.gpu_id,