# SPDX-License-Identifier: MIT AND Apache-2.0 # SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation # SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project # # Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. import abc import threading from functools import wraps from pathlib import Path import psutil import torch from tokenspeed_kernel.platform import current_platform from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool from tokenspeed.runtime.layers.attention.kv_cache.dsa import DSATokenToKVPool from tokenspeed.runtime.layers.attention.kv_cache.mha import MHATokenToKVPool from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool from tokenspeed.runtime.utils import get_colorful_logger logger = get_colorful_logger(__name__) _platform = current_platform() if _platform.is_nvidia: from tokenspeed_kernel.ops.kvcache.cuda import ( transfer_kv_all_layer_lf_pf, transfer_kv_all_layer_lf_ph, transfer_kv_all_layer_mla, transfer_kv_all_layer_mla_lf_pf, transfer_kv_direct, transfer_kv_per_layer_mla, transfer_kv_per_layer_mla_pf_lf, transfer_kv_per_layer_pf_lf, transfer_kv_per_layer_ph_lf, ) from tokenspeed_kernel.ops.kvcache.triton import ( transfer_kv_all_layer, transfer_kv_per_layer, ) if _platform.is_amd: from tokenspeed_kernel.ops.kvcache.triton import ( transfer_kv_all_layer_mla, transfer_kv_per_layer_mla, ) MLA_KVSTORE_LOADBACK_BLOCK_QUOTA = 16 MLA_KVSTORE_WRITEBACK_BLOCK_QUOTA = 16 def _read_cgroup_memory_value(path: Path) -> int | None: try: raw = path.read_text().strip() except OSError: return None if not raw or raw == "max": return None try: value = int(raw) except ValueError: return None if value <= 0 or value >= (1 << 60): return None return value def get_cgroup_memory_limit_and_current() -> tuple[int, int] | None: """Return the active cgroup memory limit/current bytes, if constrained.""" limit = _read_cgroup_memory_value(Path("/sys/fs/cgroup/memory.max")) current = _read_cgroup_memory_value(Path("/sys/fs/cgroup/memory.current")) if limit is None or current is None: limit = _read_cgroup_memory_value( Path("/sys/fs/cgroup/memory/memory.limit_in_bytes") ) current = _read_cgroup_memory_value( Path("/sys/fs/cgroup/memory/memory.usage_in_bytes") ) if limit is None or current is None: return None host_total = psutil.virtual_memory().total if limit >= host_total: return None return limit, current def get_available_host_memory_bytes( reserve_bytes: int, ) -> tuple[int, int, int | None]: host_available = max(psutil.virtual_memory().available - reserve_bytes, 0) cgroup_available = None cgroup_info = get_cgroup_memory_limit_and_current() if cgroup_info is not None: limit, current = cgroup_info cgroup_available = max(limit - current - reserve_bytes, 0) return min(host_available, cgroup_available), host_available, cgroup_available return host_available, host_available, None def synchronized(func): @wraps(func) def wrapper(self, *args, **kwargs): with self.lock: return func(self, *args, **kwargs) return wrapper class HostKVCache(abc.ABC): def __init__( self, device_pool: BaseTokenToKVPool, host_to_device_ratio: float, host_size: int, page_size: int, layout: str, device: str, host_size_tokens: int = 0, ): self.device_pool = device_pool self.page_size = page_size self.layout = layout self.device = device self.dtype = device_pool.store_dtype self.size_per_token = self.get_size_per_token() if host_size_tokens > 0: # Explicitly specified token count takes the highest priority. # Used when this pool must share the same page address space as # another host pool (e.g. draft model sharing base model pages). self.size = host_size_tokens elif host_size > 0: self.size = int(host_size * 1e9 // self.size_per_token) else: self.size = int(device_pool.size * host_to_device_ratio) # Align up the host memory pool size to the page size self.page_num = self.size // self.page_size + 1 self.size = self.page_num * self.page_size if self.size > device_pool.size: logger.warning( "The host memory is less than the device memory with the current protocol" ) # Verify there is enough available host memory. requested_bytes = self.size * self.size_per_token # preserve at least 10GB for other usage ten_gb = 10 * (1024**3) available_bytes, host_available, cgroup_available = ( get_available_host_memory_bytes(ten_gb) ) if requested_bytes > available_bytes: raise ValueError( f"Not enough host memory available. Requesting " f"{requested_bytes / 1e9:.2f} GB but only have " f"{available_bytes / 1e9:.2f} GB free. Please reduce the " f"size of the KVStore." ) else: logger.info( "Allocating %.2f GB host memory for KVStore. host_size=%r self.size_per_token=%r host_to_device_ratio=%r device_pool.size=%r host_mem.available=%r", requested_bytes / 1e9, host_size, self.size_per_token, host_to_device_ratio, device_pool.size, host_available, ) if cgroup_available is not None: logger.info( "KVStore cgroup-aware available host memory: %.2f GB", cgroup_available / 1e9, ) self.kv_buffer = self.init_kv_buffer() # A lock for synchronized operations on memory allocation and state transitions. self.lock = threading.RLock() self.clear() @abc.abstractmethod def get_size_per_token(self): raise NotImplementedError() @abc.abstractmethod def init_kv_buffer(self): raise NotImplementedError() @abc.abstractmethod def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ) -> None: """ Load KV data from the host memory pool to the device memory pool for a specific layer. """ raise NotImplementedError() @abc.abstractmethod def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend, block_quota: int | None = None, ) -> None: """ Backup KV data from the device memory pool to the host memory pool for all layers. """ raise NotImplementedError() @abc.abstractmethod def get_data_page(self, index, flat: bool = True) -> torch.Tensor: """ Get a flat data page from the host memory pool. """ raise NotImplementedError() @abc.abstractmethod def get_dummy_flat_data_page(self) -> torch.Tensor: """ Get a dummy flat data page from the host memory pool. This is used for prefetching or initializing empty pages. """ raise NotImplementedError() @abc.abstractmethod def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: """ Set a flat data page to the host memory pool. """ raise NotImplementedError() @synchronized def clear(self): # Initialize memory states and tracking structures. self.mem_state = torch.zeros( (self.size,), dtype=torch.uint8, device=self.device ) self.free_slots = torch.arange(self.size, dtype=torch.int64) def available_size(self): return len(self.free_slots) @synchronized def alloc(self, need_size: int) -> torch.Tensor | None: if need_size % self.page_size != 0: raise ValueError("The requested size should be a multiple of page_size.") if need_size > self.available_size(): return None select_index = self.free_slots[:need_size] self.free_slots = self.free_slots[need_size:] return select_index @synchronized def free(self, indices: torch.Tensor) -> int: self.free_slots = torch.cat([self.free_slots, indices]) return len(indices) class MHATokenToKVPoolHost(HostKVCache): device_pool: MHATokenToKVPool def __init__( self, device_pool: MHATokenToKVPool, host_to_device_ratio: float, host_size: int, page_size: int, layout: str, device: str = "cpu", host_size_tokens: int = 0, ): super().__init__( device_pool, host_to_device_ratio, host_size, page_size, layout, device, host_size_tokens=host_size_tokens, ) self.element_dim = self.device_pool.head_num * self.device_pool.head_dim self.k_data_refs = [self.k_buffer[i] for i in range(self.layer_num)] self.v_data_refs = [self.v_buffer[i] for i in range(self.layer_num)] platform = current_platform() self.k_data_ptrs = torch.tensor( [platform.device_visible_data_ptr(x) for x in self.k_data_refs], dtype=torch.uint64, device=self.device_pool.device, ) self.v_data_ptrs = torch.tensor( [platform.device_visible_data_ptr(x) for x in self.v_data_refs], dtype=torch.uint64, device=self.device_pool.device, ) def get_size_per_token(self): self.head_num = self.device_pool.head_num self.head_dim = self.device_pool.head_dim self.layer_num = self.device_pool.layer_num return self.head_dim * self.head_num * self.layer_num * self.dtype.itemsize * 2 def get_ksize_per_token(self): return self.get_size_per_token() // 2 def init_kv_buffer(self): if self.layout == "layer_first": dims = (2, self.layer_num, self.size, self.head_num, self.head_dim) elif self.layout == "page_first": dims = (2, self.size, self.layer_num, self.head_num, self.head_dim) elif self.layout == "page_head": dims = ( 2, self.page_num, self.head_num, self.page_size, self.layer_num, self.head_dim, ) else: raise ValueError(f"Unsupported layout: {self.layout}") self.token_stride_size = self.head_num * self.head_dim * self.dtype.itemsize self.layout_dim = self.token_stride_size * self.layer_num buffer = torch.empty( dims, dtype=self.dtype, device=self.device, ) current_platform().register_host_tensor_for_gpu_access(buffer) return buffer @property def k_buffer(self): return self.kv_buffer[0] @property def v_buffer(self): return self.kv_buffer[1] def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend, ): if io_backend == "kernel": if self.layout == "layer_first": transfer_kv_per_layer( src_k=self.k_buffer[layer_id], dst_k=device_pool.k_buffer[layer_id], src_v=self.v_buffer[layer_id], dst_v=device_pool.v_buffer[layer_id], src_indices=host_indices, dst_indices=device_indices, item_size=self.token_stride_size, ) elif self.layout == "page_first": transfer_kv_per_layer_pf_lf( src_k=self.k_buffer, dst_k=device_pool.k_buffer[layer_id], src_v=self.v_buffer, dst_v=device_pool.v_buffer[layer_id], src_indices=host_indices, dst_indices=device_indices, layer_id=layer_id, item_size=self.token_stride_size, src_layout_dim=self.layout_dim, ) elif self.layout == "page_head": transfer_kv_per_layer_ph_lf( src_k=self.k_buffer, dst_k=device_pool.k_buffer[layer_id], src_v=self.v_buffer, dst_v=device_pool.v_buffer[layer_id], src_indices=host_indices, dst_indices=device_indices, layer_id=layer_id, item_size=self.token_stride_size, src_layout_dim=self.layout_dim, page_size=self.page_size, head_num=self.head_num, ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( src_layers=[self.k_buffer[layer_id], self.v_buffer[layer_id]], dst_layers=[ device_pool.k_buffer[layer_id], device_pool.v_buffer[layer_id], ], src_indices=host_indices, dst_indices=device_indices, page_size=self.page_size, ) else: raise ValueError(f"Unsupported layout: {self.layout}") else: raise ValueError(f"Unsupported IO backend: {io_backend}") def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend, block_quota: int | None = None, ): if io_backend == "kernel": if self.layout == "layer_first": transfer_kv_all_layer( src_k_layers=device_pool.k_data_ptrs, dst_k_layers=self.k_data_ptrs, src_v_layers=device_pool.v_data_ptrs, dst_v_layers=self.v_data_ptrs, src_indices=device_indices, dst_indices=host_indices, item_size=self.token_stride_size, num_layers=self.layer_num, ) elif self.layout == "page_first": transfer_kv_all_layer_lf_pf( src_k_layers=device_pool.k_data_ptrs, dst_k=self.k_buffer, src_v_layers=device_pool.v_data_ptrs, dst_v=self.v_buffer, src_indices=device_indices, dst_indices=host_indices, item_size=self.token_stride_size, dst_layout_dim=self.layout_dim, num_layers=self.layer_num, ) elif self.layout == "page_head": transfer_kv_all_layer_lf_ph( src_k_layers=device_pool.k_data_ptrs, dst_k=self.k_buffer, src_v_layers=device_pool.v_data_ptrs, dst_v=self.v_buffer, src_indices=device_indices, dst_indices=host_indices, item_size=self.token_stride_size, dst_layout_dim=self.layout_dim, num_layers=self.layer_num, page_size=self.page_size, head_num=self.head_num, ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( src_layers=device_pool.k_buffer + device_pool.v_buffer, dst_layers=self.k_data_refs + self.v_data_refs, src_indices=device_indices, dst_indices=host_indices, page_size=self.page_size, ) else: raise ValueError(f"Unsupported layout: {self.layout}") else: raise ValueError(f"Unsupported IO backend: {io_backend}") def get_data_page(self, index, flat: bool = True) -> torch.Tensor: if self.layout == "layer_first": data_page = self.kv_buffer[:, :, index : index + self.page_size, :, :] elif self.layout == "page_first": data_page = self.kv_buffer[:, index : index + self.page_size, :, :, :] elif self.layout == "page_head": real_index = index // self.page_size data_page = self.kv_buffer[:, real_index : real_index + 1, :, :, :, :] else: raise ValueError(f"Unsupported layout: {self.layout}") if flat: data_page = data_page.flatten() return data_page def get_dummy_flat_data_page(self) -> torch.Tensor: return torch.zeros( (2, self.layer_num, self.page_size, self.head_num, self.head_dim), dtype=self.dtype, device=self.device, pin_memory=True, ).flatten() def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: if self.layout == "layer_first": self.kv_buffer[:, :, index : index + self.page_size, :, :] = ( data_page.reshape( 2, self.layer_num, self.page_size, self.head_num, self.head_dim, ) ) elif self.layout == "page_first": self.kv_buffer[:, index : index + self.page_size, :, :, :] = ( data_page.reshape( 2, self.page_size, self.layer_num, self.head_num, self.head_dim ) ) elif self.layout == "page_head": real_index = index // self.page_size self.kv_buffer[:, real_index : real_index + 1, :, :, :, :] = ( data_page.reshape( 2, 1, self.head_num, self.page_size, self.layer_num, self.head_dim ) ) else: raise ValueError(f"Unsupported layout: {self.layout}") def get_page_buffer_meta(self, indices): if len(indices) % self.page_size != 0: raise ValueError("indices length must be a multiple of page_size") ptr_list = [] kv_buffer_data_ptr = self.kv_buffer.data_ptr() indices = indices.tolist() v_offset = ( self.layer_num * self.size * self.head_num * self.head_dim * self.dtype.itemsize ) if self.layout == "layer_first": for index in range(0, len(indices), self.page_size): for layer_id in range(self.layer_num): k_ptr = ( kv_buffer_data_ptr + indices[index] * self.head_num * self.head_dim * self.dtype.itemsize + layer_id * self.size * self.head_num * self.head_dim * self.dtype.itemsize ) v_ptr = k_ptr + v_offset ptr_list.append(k_ptr) ptr_list.append(v_ptr) element_size = ( self.dtype.itemsize * self.page_size * self.head_num * self.head_dim ) element_size_list = [element_size] * len(ptr_list) elif self.layout in ["page_first", "page_head"]: for index in range(0, len(indices), self.page_size): k_ptr = ( kv_buffer_data_ptr + indices[index] * self.layer_num * self.head_num * self.head_dim * self.dtype.itemsize ) v_ptr = k_ptr + v_offset ptr_list.append(k_ptr) ptr_list.append(v_ptr) element_size = ( self.layer_num * self.dtype.itemsize * self.page_size * self.head_num * self.head_dim ) element_size_list = [element_size] * len(ptr_list) else: raise ValueError(f"Unsupported layout: {self.layout}") return ptr_list, element_size_list class MLATokenToKVPoolHost(HostKVCache): device_pool: MLATokenToKVPool def __init__( self, device_pool: MLATokenToKVPool, host_to_device_ratio: float, host_size: int, page_size: int, layout: str, device: str = "cpu", host_size_tokens: int = 0, ): super().__init__( device_pool, host_to_device_ratio, host_size, page_size, layout, device, host_size_tokens=host_size_tokens, ) self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] platform = current_platform() self.data_ptrs = torch.tensor( [platform.device_visible_data_ptr(x) for x in self.data_refs], dtype=torch.uint64, device=self.device_pool.device, ) def get_size_per_token(self): self.kv_lora_rank = self.device_pool.kv_lora_rank 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 * self.layer_num ) def get_ksize_per_token(self): return self.get_size_per_token() def init_kv_buffer(self): if self.layout == "layer_first": dims = ( self.layer_num, self.size, 1, self.kv_lora_rank + self.qk_rope_head_dim, ) elif self.layout == "page_first": dims = ( self.size, self.layer_num, 1, self.kv_lora_rank + self.qk_rope_head_dim, ) else: raise ValueError(f"Unsupported layout: {self.layout}") self.token_stride_size = ( self.kv_lora_rank + self.qk_rope_head_dim ) * self.dtype.itemsize self.layout_dim = self.token_stride_size * self.layer_num buffer = torch.empty( dims, dtype=self.dtype, device=self.device, ) current_platform().register_host_tensor_for_gpu_access(buffer) return buffer def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): if io_backend == "kernel": if self.layout == "layer_first": transfer_kv_per_layer_mla( src=self.kv_buffer[layer_id], dst=device_pool.kv_buffer[layer_id], src_indices=host_indices, dst_indices=device_indices, item_size=self.token_stride_size, block_quota=MLA_KVSTORE_LOADBACK_BLOCK_QUOTA, ) elif self.layout == "page_first": transfer_kv_per_layer_mla_pf_lf( src=self.kv_buffer, dst=device_pool.kv_buffer[layer_id], src_indices=host_indices, dst_indices=device_indices, layer_id=layer_id, item_size=self.token_stride_size, src_layout_dim=self.layout_dim, block_quota=MLA_KVSTORE_LOADBACK_BLOCK_QUOTA, ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( src_layers=[self.kv_buffer[layer_id]], dst_layers=[device_pool.kv_buffer[layer_id]], src_indices=host_indices, dst_indices=device_indices, page_size=self.page_size, ) else: raise ValueError(f"Unsupported layout: {self.layout}") else: raise ValueError(f"Unsupported IO backend: {io_backend}") def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend, block_quota: int | None = None, ): if block_quota is None: block_quota = MLA_KVSTORE_WRITEBACK_BLOCK_QUOTA if io_backend == "kernel": if self.layout == "layer_first": transfer_kv_all_layer_mla( src_layers=device_pool.data_ptrs, dst_layers=self.data_ptrs, src_indices=device_indices, dst_indices=host_indices, item_size=self.token_stride_size, num_layers=self.layer_num, block_quota=block_quota, ) elif self.layout == "page_first": transfer_kv_all_layer_mla_lf_pf( src_layers=device_pool.data_ptrs, dst=self.kv_buffer, src_indices=device_indices, dst_indices=host_indices, item_size=self.token_stride_size, dst_layout_dim=self.layout_dim, num_layers=self.layer_num, block_quota=block_quota, ) else: raise ValueError(f"Unsupported layout: {self.layout}") elif io_backend == "direct": if self.layout == "layer_first": transfer_kv_direct( src_layers=device_pool.kv_buffer, dst_layers=self.data_refs, src_indices=device_indices, dst_indices=host_indices, page_size=self.page_size, ) else: raise ValueError(f"Unsupported layout: {self.layout}") else: raise ValueError(f"Unsupported IO backend: {io_backend}") def get_data_page(self, index, flat: bool = True) -> torch.Tensor: if self.layout == "layer_first": data_page = self.kv_buffer[:, index : index + self.page_size, :, :] elif self.layout == "page_first": data_page = self.kv_buffer[index : index + self.page_size, :, :, :] else: raise ValueError(f"Unsupported layout: {self.layout}") if flat: data_page = data_page.flatten() return data_page def get_dummy_flat_data_page(self) -> torch.Tensor: return torch.zeros( ( self.layer_num, self.page_size, 1, self.kv_lora_rank + self.qk_rope_head_dim, ), dtype=self.dtype, device=self.device, pin_memory=True, ).flatten() def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: if self.layout == "layer_first": self.kv_buffer[:, index : index + self.page_size, :, :] = data_page.reshape( self.layer_num, self.page_size, 1, self.kv_lora_rank + self.qk_rope_head_dim, ) elif self.layout == "page_first": self.kv_buffer[index : index + self.page_size, :, :, :] = data_page.reshape( self.page_size, self.layer_num, 1, self.kv_lora_rank + self.qk_rope_head_dim, ) else: raise ValueError(f"Unsupported layout: {self.layout}") def get_page_buffer_meta(self, indices): if len(indices) % self.page_size != 0: raise ValueError("indices length must be a multiple of page_size") ptr_list = [] kv_buffer_data_ptr = self.kv_buffer.data_ptr() indices = indices.tolist() if self.layout == "layer_first": for index in range(0, len(indices), self.page_size): for layer_id in range(self.layer_num): k_ptr = ( kv_buffer_data_ptr + indices[index] * (self.kv_lora_rank + self.qk_rope_head_dim) * self.dtype.itemsize + layer_id * self.size * (self.kv_lora_rank + self.qk_rope_head_dim) * self.dtype.itemsize ) ptr_list.append(k_ptr) element_size = ( self.dtype.itemsize * self.page_size * (self.kv_lora_rank + self.qk_rope_head_dim) ) element_size_list = [element_size] * len(ptr_list) elif self.layout == "page_first": for index in range(0, len(indices), self.page_size): k_ptr = ( kv_buffer_data_ptr + indices[index] * self.layer_num * (self.kv_lora_rank + self.qk_rope_head_dim) * self.dtype.itemsize ) ptr_list.append(k_ptr) element_size = ( self.layer_num * self.dtype.itemsize * self.page_size * (self.kv_lora_rank + self.qk_rope_head_dim) ) element_size_list = [element_size] * len(ptr_list) else: raise ValueError(f"Unsupported layout: {self.layout}") return ptr_list, element_size_list class DSATokenToKVPoolHost(MLATokenToKVPoolHost): """Host (L2) mirror of the GLM DSA KV pool. Extends the MLA latent host pool with the DSA FP8 index-K buffer. Both buffers mirror the device row layout and transfer per token. The index-K buffers are in a block-split layout: each page is laid out as ``[page_size * head_dim FP8 values]`` followed by ``[page_size * num_groups FP32 scales]``. Hence, it requires the token indices are built as whole page-expanded blocks (see host_executor.page_ids_to_token_indices); otherwise the D<->H transfer would be corrupted. """ device_pool: DSATokenToKVPool def __init__( self, device_pool: DSATokenToKVPool, host_to_device_ratio: float, host_size: int, page_size: int, layout: str, device: str = "cpu", host_size_tokens: int = 0, ): if device_pool.quant_method == "per_token_head": raise NotImplementedError( "DSA KVStore does not support the per_token_head latent layout." ) if layout != "layer_first": raise NotImplementedError( f"DSA KVStore supports only the layer_first host layout, got {layout}." ) self.index_k_row_bytes = device_pool.index_k_row_bytes super().__init__( device_pool, host_to_device_ratio, host_size, page_size, layout, device, host_size_tokens=host_size_tokens, ) self.index_k_refs = [self.index_k_buffer[i] for i in range(self.layer_num)] platform = current_platform() self.index_k_data_ptrs = torch.tensor( [platform.device_visible_data_ptr(x) for x in self.index_k_refs], dtype=torch.uint64, device=self.device_pool.device, ) def get_size_per_token(self): return super().get_size_per_token() + self.index_k_row_bytes * self.layer_num def init_kv_buffer(self): kv_buffer = super().init_kv_buffer() # Mirror the device index-K layout: page p of layer L occupies rows # [p * page_size : (p + 1) * page_size], so a whole page is contiguous # and the block-split FP8/scale bytes within it survive a raw page copy. self.index_k_buffer = torch.zeros( (self.layer_num, self.size, self.index_k_row_bytes), dtype=torch.uint8, device=self.device, ) current_platform().register_host_tensor_for_gpu_access(self.index_k_buffer) return kv_buffer def load_to_device_per_layer( self, device_pool, host_indices, device_indices, layer_id, io_backend ): super().load_to_device_per_layer( device_pool, host_indices, device_indices, layer_id, io_backend ) if io_backend == "kernel": transfer_kv_per_layer_mla( src=self.index_k_buffer[layer_id], dst=device_pool.index_k_buffer[layer_id], src_indices=host_indices, dst_indices=device_indices, item_size=self.index_k_row_bytes, block_quota=MLA_KVSTORE_LOADBACK_BLOCK_QUOTA, ) elif io_backend == "direct": transfer_kv_direct( src_layers=[self.index_k_buffer[layer_id]], dst_layers=[device_pool.index_k_buffer[layer_id]], src_indices=host_indices, dst_indices=device_indices, page_size=self.page_size, ) else: raise ValueError(f"Unsupported IO backend: {io_backend}") def backup_from_device_all_layer( self, device_pool, host_indices, device_indices, io_backend, block_quota: int | None = None, ): super().backup_from_device_all_layer( device_pool, host_indices, device_indices, io_backend, block_quota ) if io_backend == "kernel": if block_quota is None: block_quota = MLA_KVSTORE_WRITEBACK_BLOCK_QUOTA transfer_kv_all_layer_mla( src_layers=device_pool.index_k_data_ptrs, dst_layers=self.index_k_data_ptrs, src_indices=device_indices, dst_indices=host_indices, item_size=self.index_k_row_bytes, num_layers=self.layer_num, block_quota=block_quota, ) elif io_backend == "direct": transfer_kv_direct( src_layers=list(device_pool.index_k_buffer), dst_layers=self.index_k_refs, src_indices=device_indices, dst_indices=host_indices, page_size=self.page_size, ) else: raise ValueError(f"Unsupported IO backend: {io_backend}")