Unverified Commit b29a026e authored by Simon (Jiyou) Li's avatar Simon (Jiyou) Li Committed by GitHub
Browse files

KV‑Cache (MHA, MLA): add missing start_layer / end_layer fields to...


KV‑Cache (MHA, MLA): add missing start_layer / end_layer fields to MHATokenToKVPoolHost and MLATokenToKVPoolHost (#6016)
Co-authored-by: default avatar继优 <jiyou.ljy@alibaba-inc.com>
Co-authored-by: default avatarchus-chus <chus-chus@users.noreply.github.com>
Co-authored-by: default avatarZhiqiang Xie <xiezhq@stanford.edu>
parent 678d8cc9
......@@ -762,6 +762,8 @@ class HostKVCache(abc.ABC):
self.size = int(device_pool.size * host_to_device_ratio)
# Align the host memory pool size to the page size
self.size = self.size - (self.size % self.page_size)
self.start_layer = device_pool.start_layer
self.end_layer = device_pool.end_layer
assert (
self.size > device_pool.size
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment