[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>
2026-07-07 14:54:38 +08:00
parent fefc1743a9
commit 669fd4b8a5
2 changed files with 5 additions and 1 deletions
@@ -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,