[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 <csmthu@gmail.com>

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()

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,