# Copyright 2023-2026 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== """Layer-sharded DSA KV cache pool for context-parallel prefill. ``LayerSplitDSATokenToKVPool`` splits the DSA (DeepSeek Sparse Attention) GPU KV/indexer cache layers across context-parallel (CP) ranks so that each rank only materializes the layers it owns, reducing per-rank KV memory. When a rank needs to read a layer it does not own, the owning rank broadcasts that layer's buffer into a small per-rank remote scratch buffer. This subclass keeps the core ``KVCache`` / ``MLATokenToKVPool`` / ``DSATokenToKVPool`` pools untouched: all sharding, broadcast, and remote-scratch bookkeeping lives here. Layer split is only ever enabled for DSA MLA models on PD prefill workers under prefill-CP (see ``sglang.srt.layers.cp.utils.is_glm_dsa_cache_layer_split_enabled``). """ from __future__ import annotations import logging from contextlib import nullcontext from typing import TYPE_CHECKING, Optional import torch from sglang.srt.layers.attention.dsa import index_buf_accessor from sglang.srt.layers.cp.utils import get_layer_owner, get_layer_shard_range from sglang.srt.mem_cache.memory_pool import ( GPU_MEMORY_TYPE_KV_CACHE, DSATokenToKVPool, RadixAttention, get_tensor_size_bytes, maybe_detect_oob, unwrap_write_loc, ) from sglang.srt.runtime_context import get_parallel if TYPE_CHECKING: from sglang.srt.managers.cache_controller import LayerDoneCounter logger = logging.getLogger(__name__) class LayerSplitDSATokenToKVPool(DSATokenToKVPool): """DSA KV pool that shards layers across CP ranks with owner-broadcast reads.""" def __init__( self, *args, layer_shard_rank: int, layer_shard_size: int, **kwargs, ): assert ( layer_shard_rank is not None and layer_shard_size > 1 ), "LayerSplitDSATokenToKVPool requires layer_shard_size > 1" self.layer_shard_rank = layer_shard_rank self.layer_shard_size = layer_shard_size self.layer_shard_enabled = True self.layer_broadcast_comm = None super().__init__(*args, **kwargs) # First global layer index owned by this rank (used by PD transfer to # label the contiguous owned-buffer range). my_start, _ = self._owned_local_layer_range() self.layer_shard_start = self.start_layer + my_start # ---- layer ownership helpers ------------------------------------------ def _local_layer_idx(self, layer_id: int) -> int: return layer_id - self.start_layer def _owned_local_layer_range(self) -> tuple[int, int]: return get_layer_shard_range( self.layer_shard_rank, self.layer_shard_size, self.layer_num ) def _is_layer_owned(self, layer_id: int) -> bool: local_idx = self._local_layer_idx(layer_id) owned_start, owned_end = self._owned_local_layer_range() return owned_start <= local_idx < owned_end def _get_layer_owner_rank(self, layer_id: int) -> int: return get_layer_owner( self._local_layer_idx(layer_id), self.layer_shard_size, self.layer_num ) def _log_layer_shard_plan(self) -> None: partitions = [] for rank in range(self.layer_shard_size): st, ed = get_layer_shard_range(rank, self.layer_shard_size, self.layer_num) partitions.append(f"r{rank}:[{st},{ed})") my_start, my_end = self._owned_local_layer_range() logger.info( "Layer shard plan (continuous): " f"layer_num={self.layer_num}, shard_size={self.layer_shard_size}, " f"rank={self.layer_shard_rank}, local=[{my_start},{my_end}), " f"global=[{self.start_layer + my_start},{self.start_layer + my_end}), " f"partitions={'; '.join(partitions)}" ) # ---- broadcast plumbing ----------------------------------------------- def _init_layer_broadcast_comm(self) -> None: cp_group = get_parallel().attn_cp_group if cp_group.world_size <= 1 or cp_group.pynccl_comm is None: return from sglang.srt.distributed.device_communicators.pynccl import ( PyNcclCommunicator, ) self.layer_broadcast_comm = PyNcclCommunicator( group=cp_group.cpu_group, device=cp_group.device, ) logger.info( "Initialized dedicated layer-shard broadcast NCCL communicator: " f"rank={cp_group.rank_in_group}, world_size={cp_group.world_size}" ) def _broadcast_tensor_from_owner( self, tensor: torch.Tensor, layer_id: int, src_tensor: Optional[torch.Tensor] = None, use_layer_broadcast_comm: bool = False, ) -> torch.Tensor: owner_rank = self._get_layer_owner_rank(layer_id) if self.layer_shard_rank == owner_rank: assert src_tensor is not None if tensor.data_ptr() != src_tensor.data_ptr(): tensor.copy_(src_tensor) cp_group = get_parallel().attn_cp_group comm = ( self.layer_broadcast_comm if use_layer_broadcast_comm and self.layer_broadcast_comm is not None else cp_group.pynccl_comm ) if comm is not None: # PyNcclCommunicator defaults to disabled=True (it is only enabled # inside CUDA-graph capture via change_state). Without re-enabling it # here, comm.broadcast() is a silent no-op and non-owner CP ranks read # stale remote buffers, corrupting layer-split attention. Mirror the # standard usage in parallel_state.py. with comm.change_state(enable=True): comm.broadcast(tensor, src=owner_rank) else: torch.distributed.broadcast( tensor, src=owner_rank, group=cp_group.cpu_group ) return tensor # ---- buffer allocation (owned-only + remote scratch) ------------------ def _create_buffers(self): self._log_layer_shard_plan() with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): with ( torch.cuda.use_mem_pool(self.custom_mem_pool) if self.custom_mem_pool else nullcontext() ): # Owned layers get the full buffer; non-owned layers allocate a # 0-row placeholder so ``kv_buffer`` stays index-aligned by layer. self.kv_buffer = [ torch.zeros( ( ( (self.size + self.page_size) if self._is_layer_owned(self.start_layer + i) else 0 ), 1, self.kv_cache_dim, ), dtype=self.store_dtype, device=self.device, ) for i in range(self.layer_num) ] self.remote_kv_buffer = torch.empty( (self.size + self.page_size, 1, self.kv_cache_dim), dtype=self.store_dtype, device=self.device, ) self.remote_kv_layer_id: Optional[int] = None self.device_module = torch.get_device_module(self.device) self.kv_broadcast_stream = self.device_module.Stream() self.pending_remote_kv_layer_id: Optional[int] = None self.pending_remote_kv_broadcast = False self._init_layer_broadcast_comm() def _create_index_buffers(self): num_pages = (self.index_buf_size + self.page_size + 1) // self.page_size with ( torch.cuda.use_mem_pool(self.custom_mem_pool) if self.custom_mem_pool else nullcontext() ): self.index_k_with_scale_buffer = [ torch.zeros( self._index_buffer_shape( num_pages if self._is_layer_owned(self.start_layer + i) else 0 ), dtype=self.index_k_with_scale_buffer_dtype, device=self.device, ) for i in range(self.layer_num) ] self.remote_index_k_with_scale_buffer = torch.empty( self._index_buffer_shape(num_pages), dtype=self.index_k_with_scale_buffer_dtype, device=self.device, ) self.remote_index_layer_id: Optional[int] = None def _clear_buffers(self): del self.kv_buffer del self.remote_kv_buffer del self.remote_index_k_with_scale_buffer del self.index_k_with_scale_buffer # ---- MLA latent KV: owned-only writes, owner-broadcast reads ---------- def get_kv_size_bytes(self): kv_size_bytes = 0 for kv_cache in self.kv_buffer: kv_size_bytes += get_tensor_size_bytes(kv_cache) for index_k_cache in self.index_k_with_scale_buffer: kv_size_bytes += get_tensor_size_bytes(index_k_cache) return kv_size_bytes def get_contiguous_buf_infos(self): # Only report buffers owned by the current CP rank; non-owned layers # are empty and are pulled from their owner via PD transfer. owned_layer_ids = [ i for i in range(self.layer_num) if self._is_layer_owned(self.start_layer + i) ] kv_data_ptrs = [self.kv_buffer[i].data_ptr() for i in owned_layer_ids] kv_data_lens = [self.kv_buffer[i].nbytes for i in owned_layer_ids] kv_item_lens = [ self.kv_buffer[i][0].nbytes * self.page_size for i in owned_layer_ids ] return kv_data_ptrs, kv_data_lens, kv_item_lens def get_key_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) kv_buffer = self._get_broadcastable_kv_buffer(layer_id) if self.store_dtype != self.dtype: return kv_buffer.view(self.dtype) return kv_buffer def get_value_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) kv_buffer = self._get_broadcastable_kv_buffer(layer_id) if self.store_dtype != self.dtype: return kv_buffer[..., : self.kv_lora_rank].view(self.dtype) return kv_buffer[..., : self.kv_lora_rank] def set_kv_buffer( self, layer: RadixAttention, loc_info, cache_k: torch.Tensor, cache_v: torch.Tensor, ): loc, _, _ = unwrap_write_loc(loc_info) maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)") layer_id = layer.layer_id assert not self.dsa_kv_cache_store_fp8 # A write invalidates any cached remote copy for this layer. if self.pending_remote_kv_layer_id == layer_id: self._finalize_pending_kv_broadcast(set_remote_layer_id=False) if self.remote_kv_layer_id == layer_id: self.remote_kv_layer_id = None if not self._is_layer_owned(layer_id): return if cache_k.dtype != self.dtype: cache_k = cache_k.to(self.dtype) if self.store_dtype != self.dtype: self.kv_buffer[layer_id - self.start_layer][loc] = cache_k.view( self.store_dtype ) else: self.kv_buffer[layer_id - self.start_layer][loc] = cache_k def set_mla_kv_buffer( self, layer: RadixAttention, loc: torch.Tensor, cache_k_nope: torch.Tensor, cache_k_rope: torch.Tensor, ): maybe_detect_oob(loc, 0, self.size + self.page_size, "set_mla_kv_buffer (MLA)") layer_id = layer.layer_id if self.pending_remote_kv_layer_id == layer_id: self._finalize_pending_kv_broadcast(set_remote_layer_id=True) remote_kv_updatable = self.remote_kv_layer_id == layer_id if remote_kv_updatable: self._write_mla_kv_buffer( self.remote_kv_buffer, loc, cache_k_nope, cache_k_rope ) if not self._is_layer_owned(layer_id): return self._write_mla_kv_buffer( self.kv_buffer[layer_id - self.start_layer], loc, cache_k_nope, cache_k_rope, ) if not remote_kv_updatable and self.remote_kv_layer_id == layer_id: self.remote_kv_layer_id = None def _finalize_pending_kv_broadcast( self, *, set_remote_layer_id: bool = True ) -> None: if not self.pending_remote_kv_broadcast: return self.device_module.current_stream().wait_stream(self.kv_broadcast_stream) self.pending_remote_kv_broadcast = False if set_remote_layer_id and self.pending_remote_kv_layer_id is not None: self.remote_kv_layer_id = self.pending_remote_kv_layer_id self.pending_remote_kv_layer_id = None def prefetch_kv_buffer( self, layer_id: int, layer_transfer_counter: Optional[LayerDoneCounter] = None, layer_transfer_idx: Optional[int] = None, ) -> None: """Kick off an async owner-broadcast of ``layer_id``'s latent KV. Called ahead of the layer's attention so the remote scratch buffer is ready by the time a non-owner rank reads it (see the prefetch wiring in ``DeepseekV2DecoderLayer``). """ if self.remote_kv_layer_id == layer_id: return if self.pending_remote_kv_broadcast: if self.pending_remote_kv_layer_id == layer_id: return self._finalize_pending_kv_broadcast(set_remote_layer_id=False) local_idx = self._local_layer_idx(layer_id) src_tensor = ( self.kv_buffer[local_idx] if self._is_layer_owned(layer_id) else None ) if self.layer_broadcast_comm is None: self._broadcast_tensor_from_owner( self.remote_kv_buffer, layer_id, src_tensor=src_tensor, use_layer_broadcast_comm=True, ) self.remote_kv_layer_id = layer_id return self.kv_broadcast_stream.wait_stream(self.device_module.current_stream()) with self.device_module.stream(self.kv_broadcast_stream): if layer_transfer_counter is not None and layer_transfer_idx is not None: layer_transfer_counter.wait_until(layer_transfer_idx) self._broadcast_tensor_from_owner( self.remote_kv_buffer, layer_id, src_tensor=src_tensor, use_layer_broadcast_comm=True, ) self.pending_remote_kv_layer_id = layer_id self.pending_remote_kv_broadcast = True def _get_broadcastable_kv_buffer(self, layer_id: int) -> torch.Tensor: if self.pending_remote_kv_broadcast: self._finalize_pending_kv_broadcast( set_remote_layer_id=self.pending_remote_kv_layer_id == layer_id ) if self.remote_kv_layer_id != layer_id: local_idx = self._local_layer_idx(layer_id) src_tensor = ( self.kv_buffer[local_idx] if self._is_layer_owned(layer_id) else None ) self._broadcast_tensor_from_owner( self.remote_kv_buffer, layer_id, src_tensor=src_tensor, use_layer_broadcast_comm=True, ) self.remote_kv_layer_id = layer_id return self.remote_kv_buffer def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): size_limit = self.size + self.page_size maybe_detect_oob(tgt_loc, 0, size_limit, "move_kv_cache tgt_loc") maybe_detect_oob(src_loc, 0, size_limit, "move_kv_cache src_loc") if tgt_loc.numel() == 0: return tgt_loc_flat = tgt_loc.view(-1).long() src_loc_flat = src_loc.view(-1).long() for kv_cache in self.kv_buffer: if kv_cache.shape[0] == 0: continue kv_cache[tgt_loc_flat] = kv_cache[src_loc_flat] for index_k in self.index_k_with_scale_buffer: if index_k.shape[0] == 0: continue index_k[tgt_loc_flat] = index_k[src_loc_flat] # ---- DSA indexer buffer: owned-only writes, owner-broadcast reads ----- def get_broadcastable_index_k_with_scale_buffer( self, layer_id: int ) -> torch.Tensor: if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) return self._get_broadcastable_index_buffer(layer_id) def get_index_k_continuous(self, layer_id, seq_len, page_indices): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) buf = self._get_broadcastable_index_buffer(layer_id) return index_buf_accessor.GetK.execute( self, buf, seq_len=seq_len, page_indices=page_indices ) def get_index_k_scale_continuous(self, layer_id, seq_len, page_indices): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) buf = self._get_broadcastable_index_buffer(layer_id) return index_buf_accessor.GetS.execute( self, buf, seq_len=seq_len, page_indices=page_indices ) def get_index_k_scale_buffer( self, layer_id, seq_len_tensor, page_indices, seq_len_sum, max_seq_len ): if self.layer_transfer_counter is not None: self.layer_transfer_counter.wait_until(layer_id - self.start_layer) buf = self._get_broadcastable_index_buffer(layer_id) # Overlap the latent-KV owner-broadcast with the indexer read. self.prefetch_kv_buffer(layer_id) return index_buf_accessor.GetKAndS.execute( self, buf, page_indices=page_indices, seq_len_tensor=seq_len_tensor, seq_len_sum=seq_len_sum, max_seq_len=max_seq_len, ) def set_index_k_scale_buffer(self, layer_id, loc, index_k, index_k_scale) -> None: self.invalidate_index_buffer_for_layer(layer_id) if not self._is_layer_owned(layer_id): return buf = self.index_k_with_scale_buffer[layer_id - self.start_layer] index_buf_accessor.SetKAndS.execute( pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale ) def invalidate_index_buffer_for_layer(self, layer_id: int) -> None: if self.remote_index_layer_id == layer_id: self.remote_index_layer_id = None def _get_broadcastable_index_buffer(self, layer_id: int) -> torch.Tensor: if self.remote_index_layer_id != layer_id: local_idx = self._local_layer_idx(layer_id) src_tensor = ( self.index_k_with_scale_buffer[local_idx] if self._is_layer_owned(layer_id) else None ) self._broadcast_tensor_from_owner( self.remote_index_k_with_scale_buffer, layer_id, src_tensor=src_tensor, ) self.remote_index_layer_id = layer_id return self.remote_index_k_with_scale_buffer def get_state_buf_infos(self): owned_layer_ids = [ i for i in range(self.layer_num) if self._is_layer_owned(self.start_layer + i) ] data_ptrs = [ self.index_k_with_scale_buffer[i].data_ptr() for i in owned_layer_ids ] data_lens = [self.index_k_with_scale_buffer[i].nbytes for i in owned_layer_ids] item_lens = [ self.index_k_with_scale_buffer[i][0].nbytes for i in owned_layer_ids ] return data_ptrs, data_lens, item_lens # ---- HiCache CPU offload: skip empty (non-owned) layers --------------- def get_cpu_copy(self, indices, mamba_indices=None): from sglang.srt.utils import current_platform current_platform.synchronize() kv_cache_cpu = [] chunk_size = self.cpu_offloading_chunk_size for layer_id in range(self.layer_num): kv_cache_cpu.append([]) if self.kv_buffer[layer_id].shape[0] == 0: continue for i in range(0, len(indices), chunk_size): chunk_indices = indices[i : i + chunk_size] kv_cpu = self.kv_buffer[layer_id][chunk_indices].to( "cpu", non_blocking=True ) kv_cache_cpu[-1].append(kv_cpu) current_platform.synchronize() page_indices = indices[:: self.page_size] // self.page_size torch.cuda.synchronize() index_k_cpu = [] page_chunk_size = max(1, chunk_size // self.page_size) for layer_id in range(self.layer_num): index_k_cpu.append([]) if self.index_k_with_scale_buffer[layer_id].shape[0] == 0: continue for i in range(0, len(page_indices), page_chunk_size): chunk_page_indices = page_indices[i : i + page_chunk_size] idx_cpu = self.index_k_with_scale_buffer[layer_id][ chunk_page_indices ].to("cpu", non_blocking=True) index_k_cpu[-1].append(idx_cpu) torch.cuda.synchronize() return {"kv": kv_cache_cpu, "index_k": index_k_cpu} def load_cpu_copy(self, kv_cache_cpu_dict, indices, mamba_indices=None): from sglang.srt.utils import current_platform kv_cache_cpu = kv_cache_cpu_dict["kv"] current_platform.synchronize() chunk_size = self.cpu_offloading_chunk_size for layer_id in range(self.layer_num): if self.kv_buffer[layer_id].shape[0] == 0: continue for i in range(0, len(indices), chunk_size): chunk_indices = indices[i : i + chunk_size] kv_cpu = kv_cache_cpu[layer_id][i // chunk_size] assert kv_cpu.shape[0] == len(chunk_indices) kv_chunk = kv_cpu.to(self.kv_buffer[layer_id].device, non_blocking=True) self.kv_buffer[layer_id][chunk_indices] = kv_chunk current_platform.synchronize() page_indices = indices[:: self.page_size] // self.page_size index_k_cpu = kv_cache_cpu_dict["index_k"] torch.cuda.synchronize() page_chunk_size = max(1, chunk_size // self.page_size) for layer_id in range(self.layer_num): if self.index_k_with_scale_buffer[layer_id].shape[0] == 0: continue for i in range(0, len(page_indices), page_chunk_size): chunk_page_indices = page_indices[i : i + page_chunk_size] idx_cpu = index_k_cpu[layer_id][i // page_chunk_size] assert idx_cpu.shape[0] == len(chunk_page_indices) idx_chunk = idx_cpu.to( self.index_k_with_scale_buffer[layer_id].device, non_blocking=True ) self.index_k_with_scale_buffer[layer_id][chunk_page_indices] = idx_chunk torch.cuda.synchronize()