[PD] Add different TP sizes support for no-MLA models (#6793)

Co-authored-by: shangmingc <csmthu@gmail.com> Co-authored-by: Shangming Cai <caishangming@linux.alibaba.com>
2025-06-25 02:00:22 -07:00
parent f3cbd24541
commit e21aa1df67
3 changed files with 265 additions and 27 deletions
@@ -122,6 +122,9 @@ class PrefillBootstrapQueue:
        kv_args.kv_data_ptrs = kv_data_ptrs
        kv_args.kv_data_lens = kv_data_lens
        kv_args.kv_item_lens = kv_item_lens
+        if not self.is_mla_backend:
+            kv_args.kv_head_num = self.token_to_kv_pool.head_num
+        kv_args.page_size = self.token_to_kv_pool.page_size

        kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (
            self.metadata_buffers.get_buf_infos()