From 2d288ba8c92f249d5d0d8c98b70751396ebf6f57 Mon Sep 17 00:00:00 2001 From: XDaoHong Date: Mon, 23 Mar 2026 15:56:47 +0800 Subject: [PATCH] =?UTF-8?q?[Bugfix]=20fix=20npu=20get=20kv=5Fitem=5Flens?= =?UTF-8?q?=20in=20PD=20separation=20when=20use=20ASCEND=5FUS=E2=80=A6=20(?= =?UTF-8?q?#15852)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: ZhengdQin --- .../hardware_backend/npu/memory_pool_npu.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index 82e7ce693..ea81f4e58 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -100,13 +100,22 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): self.get_value_buffer(i).nbytes for i in range(self.start_layer, self.start_layer + self.layer_num) ] - kv_item_lens = [ - self.get_key_buffer(i)[0].nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self.get_value_buffer(i)[0].nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + if self.use_fia: + kv_item_lens = [ + self.get_key_buffer(i)[0].nbytes * self.page_size + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + [ + self.get_value_buffer(i)[0].nbytes * self.page_size + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + else: + kv_item_lens = [ + self.get_key_buffer(i)[0].nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + [ + self.get_value_buffer(i)[0].nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] return kv_data_ptrs, kv_data_lens, kv_item_lens def set_kv_buffer(