fix MLATokenToKVPoolHost get_size_per_token bug (#5161)

Co-authored-by: AniZpZ <zhuangsen.zp@antgroup.com>

fix MLATokenToKVPoolHost get_size_per_token bug (#5161)
Co-authored-by: AniZpZ <zhuangsen.zp@antgroup.com>
5fbafbb8 · huangtingwei · GitHub · a9499885 · 5fbafbb8
Unverified Commit 5fbafbb8 authored Apr 14, 2025 by huangtingwei Committed by GitHub Apr 13, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 6 additions and 1 deletion

python/sglang/srt/mem_cache/memory_pool.py python/sglang/srt/mem_cache/memory_pool.py +6 -1

No files found.
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -879,7 +879,12 @@ class MLATokenToKVPoolHost(HostKVCache):
        self.qk_rope_head_dim = self.device_pool.qk_rope_head_dim
        self.layer_num = self.device_pool.layer_num

-        return (self.kv_lora_rank + self.qk_rope_head_dim) * 1 * self.dtype.itemsize
+        return (
+            (self.kv_lora_rank + self.qk_rope_head_dim)
+            * 1
+            * self.dtype.itemsize
+            * self.layer_num
+        )

    def init_kv_buffer(self):
        return torch.empty(