From 669fd4b8a55b678d5419d0668d3906eb5854dcdc Mon Sep 17 00:00:00 2001 From: Augusto Yao Date: Tue, 7 Jul 2026 14:54:38 +0800 Subject: [PATCH] [PP] Fix start_layer_id with pp in get kv_buffer_shape (#29887) Co-authored-by: ybyang <10629930+whybeyoung@users.noreply.github.com> Co-authored-by: Shangming Cai --- python/sglang/srt/mem_cache/memory_pool.py | 4 ++++ python/sglang/srt/model_executor/runner/eager_runner.py | 2 +- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index f218f1ac8..87e2638c1 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1249,6 +1249,10 @@ class KVCache(abc.ABC): ) self.mem_usage = kv_size_GB + 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 + @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 a0d9fb9d0..4d6f430e3 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -278,7 +278,7 @@ class EagerRunner(BaseRunner): 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_kv_buffer_shape()[0], model_runner.kv_cache_dtype, model_runner.device, create_chunked_prefix_cache_kv_indices,