From cda9cb99014949c7ad45456555006cac37a4874f Mon Sep 17 00:00:00 2001 From: "augusto.yjh" Date: Thu, 2 Jul 2026 10:56:32 +0800 Subject: [PATCH 1/4] fix bugs for dsv3.1 with pp in get kv_buffer_shape --- python/sglang/srt/model_executor/runner/eager_runner.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 3137a8eeae53..0c465f73efc1 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -275,7 +275,9 @@ def _execute_extend( forward_batch.req_pool_indices, get_req_to_token_pool().req_to_token, forward_batch.seq_lens_sum, - get_token_to_kv_pool().get_key_buffer(0).shape, + get_token_to_kv_pool() + .get_key_buffer(model_runner.start_layer) + .shape, model_runner.kv_cache_dtype, model_runner.device, create_chunked_prefix_cache_kv_indices, From c69b7ccd2f3426bbc23a2a7afc0c0673bc28aad1 Mon Sep 17 00:00:00 2001 From: "augusto.yjh" Date: Thu, 2 Jul 2026 12:07:48 +0800 Subject: [PATCH 2/4] add method get_key_buffer_shape for KVCache --- python/sglang/srt/mem_cache/memory_pool.py | 3 +++ python/sglang/srt/model_executor/runner/eager_runner.py | 4 +--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 82c9c98d244a..b03747fb7a6d 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1249,6 +1249,9 @@ def _finalize_allocation_log(self, num_tokens: int): ) self.mem_usage = kv_size_GB + def get_key_buffer_shape(self) -> torch.Size: + return self.get_key_buffer(self.start_layer).shape + @abc.abstractmethod def get_key_buffer(self, layer_id: int) -> torch.Tensor: raise NotImplementedError() diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 0c465f73efc1..2dfdfbf00d5c 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -275,9 +275,7 @@ def _execute_extend( forward_batch.req_pool_indices, get_req_to_token_pool().req_to_token, forward_batch.seq_lens_sum, - get_token_to_kv_pool() - .get_key_buffer(model_runner.start_layer) - .shape, + get_token_to_kv_pool().get_key_buffer_shape(), model_runner.kv_cache_dtype, model_runner.device, create_chunked_prefix_cache_kv_indices, From a2f738f9f8fc160fe8365a4f071741c995157bad Mon Sep 17 00:00:00 2001 From: "augusto.yjh" Date: Thu, 2 Jul 2026 12:13:21 +0800 Subject: [PATCH 3/4] add get_kv_buffer_shape instead of get_key_buffer_shape --- python/sglang/srt/mem_cache/memory_pool.py | 5 +++-- python/sglang/srt/model_executor/runner/eager_runner.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index b03747fb7a6d..9583c0ae243e 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1249,8 +1249,9 @@ def _finalize_allocation_log(self, num_tokens: int): ) self.mem_usage = kv_size_GB - def get_key_buffer_shape(self) -> torch.Size: - return self.get_key_buffer(self.start_layer).shape + def get_kv_buffer_shape(self) -> torch.Size: + k_buffer, v_buffer = self.get_kv_buffer(self.start_layer) + return k_buffer.shape, v_buffer.shape @abc.abstractmethod def get_key_buffer(self, layer_id: int) -> torch.Tensor: diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 2dfdfbf00d5c..ae1a4925bb6e 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -275,7 +275,7 @@ def _execute_extend( forward_batch.req_pool_indices, get_req_to_token_pool().req_to_token, forward_batch.seq_lens_sum, - get_token_to_kv_pool().get_key_buffer_shape(), + get_token_to_kv_pool().get_kv_buffer_shape()[0], model_runner.kv_cache_dtype, model_runner.device, create_chunked_prefix_cache_kv_indices, From cdd29946a14e1371c9ec2784457b7bd3167d4a03 Mon Sep 17 00:00:00 2001 From: "augusto.yjh" Date: Thu, 2 Jul 2026 12:15:27 +0800 Subject: [PATCH 4/4] update type of return value of get_kv_buffer_shape --- python/sglang/srt/mem_cache/memory_pool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 9583c0ae243e..2efdf598bdf9 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1249,7 +1249,7 @@ def _finalize_allocation_log(self, num_tokens: int): ) self.mem_usage = kv_size_GB - def get_kv_buffer_shape(self) -> torch.Size: + def get_kv_buffer_shape(self) -> Tuple[torch.Size, torch.Size]: k_buffer, v_buffer = self.get_kv_buffer(self.start_layer) return k_buffer.shape, v_buffer.shape