fix(pd): fix kv pools without end_layer (#25476)

This commit is contained in:
ybyang
2026-05-16 19:49:26 +08:00
committed by GitHub
parent 90d3d42ac1
commit 0f50ed86c9
2 changed files with 2 additions and 2 deletions
@@ -56,7 +56,7 @@ class KVArgs:
# reconstruct PP sub-ranges when kv_data_ptrs does not use a flat
# layer-indexed layout (e.g. DeepSeek V4's buffer-type-organized flat
# list).
prefill_end_layer: int
prefill_end_layer: Optional[int]
# For DeepSeek V4 (and other compressed-MLA) memory pools only.
# Full-model compression ratio per layer (entries are 0/4/128). Used by
# the connection layer to slice the buffer-type-organized flat list in a
+1 -1
View File
@@ -147,7 +147,7 @@ class PrefillBootstrapQueue:
kv_args.pp_rank = self.pp_rank
kv_args.system_dp_rank = self.scheduler.ps.dp_rank
kv_args.prefill_start_layer = self.token_to_kv_pool.start_layer
kv_args.prefill_end_layer = self.token_to_kv_pool.end_layer
kv_args.prefill_end_layer = getattr(self.token_to_kv_pool, "end_layer", None)
kv_args.mla_compression_ratios = None
kv_data_ptrs, kv_data_lens, kv_item_lens = (
self.token_to_kv_pool.get_contiguous_buf_infos()