diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index f5be51e83dde..8001c191f3eb 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -369,6 +369,18 @@ def _init_kv_manager(self) -> CommonKVManager: ) elif isinstance(self.token_to_kv_pool, NSATokenToKVPool): kv_args.state_type = "nsa" + if self.draft_token_to_kv_pool is not None and isinstance( + self.draft_token_to_kv_pool, NSATokenToKVPool + ): + ( + draft_state_data_ptrs, + draft_state_data_lens, + draft_state_item_lens, + ) = self.draft_token_to_kv_pool.get_state_buf_infos() + kv_args.state_data_ptrs += draft_state_data_ptrs + kv_args.state_data_lens += draft_state_data_lens + kv_args.state_item_lens += draft_state_item_lens + else: kv_args.state_type = "none" else: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index c389e44e58ee..f946a5c28067 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -192,6 +192,18 @@ def _init_kv_manager(self) -> CommonKVManager: ) elif isinstance(self.token_to_kv_pool, NSATokenToKVPool): kv_args.state_type = "nsa" + if self.draft_token_to_kv_pool is not None and isinstance( + self.draft_token_to_kv_pool, NSATokenToKVPool + ): + ( + draft_state_data_ptrs, + draft_state_data_lens, + draft_state_item_lens, + ) = self.draft_token_to_kv_pool.get_state_buf_infos() + kv_args.state_data_ptrs += draft_state_data_ptrs + kv_args.state_data_lens += draft_state_data_lens + kv_args.state_item_lens += draft_state_item_lens + else: kv_args.state_type = "none" else: