diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh index c0d5a817040a..399ff566253b 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hisparse.cuh @@ -350,6 +350,7 @@ template < bool TopKIsBlocks, bool RecordMissPlan, bool SkipIO, + int NUM_TAIL_SLOTS, typename SeqLensT, typename ReqPoolIndicesT> __global__ void load_cache_to_device_buffer_kernel( @@ -378,6 +379,8 @@ __global__ void load_cache_to_device_buffer_kernel( int32_t* __restrict__ miss_count_out, int64_t plan_stride) { static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); + static_assert(NUM_TAIL_SLOTS >= 0 && NUM_TAIL_SLOTS < NUM_TOP_K, "Invalid top-k tail slots."); + static_assert(!TopKIsBlocks || NUM_TAIL_SLOTS == 0, "Block top-k does not use tail columns."); static_assert(SPARSE_BLOCK_SIZE > 0, "SPARSE_BLOCK_SIZE must be positive."); // Cache residency and LRU replacement remain token-granular even when the // sparse-attention selection arrives as block ids. @@ -425,7 +428,7 @@ __global__ void load_cache_to_device_buffer_kernel( if (token_pos >= 0 && token_pos < seq_len) { device_loc = req_device_buffer_locs[token_pos]; } - } else if (i < count && token_pos >= 0) { + } else if ((i < count || i >= NUM_TOP_K_TOKENS - NUM_TAIL_SLOTS) && token_pos >= 0 && token_pos < seq_len) { device_loc = req_device_buffer_locs[token_pos]; } req_top_k_device_locs[i] = device_loc; @@ -483,12 +486,10 @@ __global__ void load_cache_to_device_buffer_kernel( // Insert top-k tokens into shared-memory hash table. for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) { const int32_t token_idx = resolve_selected_token(req_top_k, i); - if constexpr (TopKIsBlocks) { - if (token_idx < 0 || token_idx >= seq_len) { - s_top_k_tokens[i] = TOKEN_HIT; - req_top_k_device_locs[i] = -1; - continue; - } + if (token_idx < 0 || token_idx >= seq_len) { + s_top_k_tokens[i] = TOKEN_HIT; + req_top_k_device_locs[i] = -1; + continue; } if (token_idx == newest_token) { // If topk includes the latest token, bind its canonical occurrence to newest_slot (at HOT_BUFFER_SIZE) and mark @@ -558,23 +559,13 @@ __global__ void load_cache_to_device_buffer_kernel( __syncthreads(); if (warp_id == 0) { -#ifdef USE_ROCM - // ROCm wavefront64: WARP_SIZE (64) > NUM_WARPS (16 at block_size=1024), - // so the wide-count form below would let lanes beyond this iteration's - // NUM_WARPS-wide window write the accumulator into s_chunk_offset - // positions belonging to future iterations, corrupting their reads. - // Bound the scan window to NUM_WARPS lanes. + // Each iteration produces NUM_WARPS chunks, not a full warp of chunks. const int scan_offset = iter * NUM_WARPS + 1; const int scan_count = min(scan_offset + NUM_WARPS, NUM_BUFFER_CHUNKS + 1); total_hit_count = warp_inclusive_scan(s_chunk_offset, lane_id, scan_offset, scan_count, total_hit_count); total_evict_count = warp_inclusive_scan(s_evict_chunk_offset, lane_id, scan_offset, scan_count, total_evict_count); -#else - total_hit_count = - warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_hit_count); - total_evict_count = - warp_inclusive_scan(s_evict_chunk_offset, lane_id, chunk_idx + 1, NUM_BUFFER_CHUNKS + 1, total_evict_count); -#endif + if (tid == 0) { s_total_hits = total_hit_count; } @@ -633,13 +624,10 @@ __global__ void load_cache_to_device_buffer_kernel( __syncthreads(); if (warp_id == 0) { -#ifdef USE_ROCM const int scan_offset = iter * NUM_WARPS + 1; const int scan_count = min(scan_offset + NUM_WARPS, NUM_TOKEN_CHUNKS + 1); total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, scan_offset, scan_count, total_misses); -#else - total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses); -#endif + if (tid == 0) { s_total_misses = total_misses; } @@ -730,7 +718,8 @@ template < int SPARSE_BLOCK_SIZE, bool TopKIsBlocks, bool RecordMissPlan, - bool SkipIO> + bool SkipIO, + int NUM_TAIL_SLOTS = 0> void load_cache_to_device_buffer( tvm::ffi::TensorView top_k, tvm::ffi::TensorView device_buffer_tokens, @@ -830,6 +819,7 @@ void load_cache_to_device_buffer( TopKIsBlocks, RecordMissPlan, SkipIO, + NUM_TAIL_SLOTS, int64_t, int64_t>, static_cast(seq_lens.data_ptr()), @@ -846,6 +836,7 @@ void load_cache_to_device_buffer( TopKIsBlocks, RecordMissPlan, SkipIO, + NUM_TAIL_SLOTS, int64_t, int32_t>, static_cast(seq_lens.data_ptr()), @@ -862,6 +853,7 @@ void load_cache_to_device_buffer( TopKIsBlocks, RecordMissPlan, SkipIO, + NUM_TAIL_SLOTS, int32_t, int64_t>, static_cast(seq_lens.data_ptr()), @@ -878,6 +870,7 @@ void load_cache_to_device_buffer( TopKIsBlocks, RecordMissPlan, SkipIO, + NUM_TAIL_SLOTS, int32_t, int32_t>, static_cast(seq_lens.data_ptr()), diff --git a/python/sglang/kernels/ops/kvcache/hisparse.py b/python/sglang/kernels/ops/kvcache/hisparse.py index 9294de22dd30..41101bf675ef 100644 --- a/python/sglang/kernels/ops/kvcache/hisparse.py +++ b/python/sglang/kernels/ops/kvcache/hisparse.py @@ -178,9 +178,9 @@ def _jit_sparse_module( top_k_is_blocks: bool = False, record_miss_plan: bool = False, skip_io: bool = False, + num_tail_slots: int = 0, ) -> Module: - # record_miss_plan / skip_io are compile-time kernel flags; the - # (False, False) production instantiation stays byte-identical. + # Compile-time options preserve the ordinary top-k layout by default. template_args = make_cpp_args( block_size, num_top_k, @@ -191,6 +191,7 @@ def _jit_sparse_module( top_k_is_blocks, record_miss_plan, skip_io, + num_tail_slots, ) cache_args = make_cpp_args( item_size_bytes, @@ -203,6 +204,7 @@ def _jit_sparse_module( top_k_is_blocks, record_miss_plan, skip_io, + num_tail_slots, ) return load_jit( "sparse_cache", @@ -297,11 +299,26 @@ def _load_cache_to_device_buffer_mla( miss_dst: torch.Tensor | None, miss_count: torch.Tensor | None, skip_io: bool, + num_tail_slots: int = 0, ) -> None: assert hot_buffer_size >= num_top_k, ( f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})" ) + if ( + top_k_tokens.ndim != 2 + or top_k_tokens.shape[1] != num_top_k + or top_k_device_locs.shape != top_k_tokens.shape + or top_k_tokens.dtype != torch.int32 + or top_k_device_locs.dtype != torch.int32 + or top_k_tokens.stride(1) != 1 + or top_k_device_locs.stride(1) != 1 + ): + raise ValueError( + "HiSparse top-k must be int32 [batch, num_top_k] with contiguous columns." + ) + if not 0 <= num_tail_slots < num_top_k: + raise ValueError("HiSparse tail slots must be in [0, num_top_k).") record_miss_plan = miss_src is not None module = _jit_sparse_module( item_size_bytes, @@ -312,6 +329,7 @@ def _load_cache_to_device_buffer_mla( is_dsv4_layout=is_dsv4_layout, record_miss_plan=record_miss_plan, skip_io=skip_io, + num_tail_slots=num_tail_slots, ) empty = torch.empty(0, device=top_k_tokens.device) @@ -374,9 +392,13 @@ def load_cache_to_device_buffer_mla( miss_dst: torch.Tensor | None = None, miss_count: torch.Tensor | None = None, skip_io: bool = False, + num_tail_slots: int = 0, ) -> None: """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes). + num_tail_slots counts trailing KPool columns, valid independently of the + ordinary top-k prefix for short sequences. + Optional miss_src/miss_dst/miss_count record the miss plan for replay by copy_cache_planned_mla; skip_io elides only the KV bytes (timing probe). """ @@ -402,6 +424,7 @@ def load_cache_to_device_buffer_mla( miss_dst=miss_dst, miss_count=miss_count, skip_io=skip_io, + num_tail_slots=num_tail_slots, ) diff --git a/python/sglang/srt/arg_groups/hisparse_hook.py b/python/sglang/srt/arg_groups/hisparse_hook.py index 0c8e95184955..8eecc235764c 100644 --- a/python/sglang/srt/arg_groups/hisparse_hook.py +++ b/python/sglang/srt/arg_groups/hisparse_hook.py @@ -17,7 +17,7 @@ HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE = { "bfloat16": {"flashmla_sparse"}, - "fp8_e4m3": {"flashmla_kv", "flashinfer_sparse_mla"}, + "fp8_e4m3": {"flashmla_kv", "flashinfer_sparse_mla", "tilelang"}, } HISPARSE_ROCM_DSA_BACKENDS = {"tilelang", "triton", "aiter"} HISPARSE_XPU_DSA_BACKENDS_BY_DTYPE = {"bfloat16": {"intel_xpu"}} diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 91b02acfd48c..6c55b62964ff 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -684,15 +684,7 @@ def _dsa_split_backend_resolution(view: Any) -> dict: declared["dsa_decode_backend"] = backend prefill = declared.get("dsa_prefill_backend", view.dsa_prefill_backend) decode = declared.get("dsa_decode_backend", view.dsa_decode_backend) - if ( - not get_platform().is_hip - and kv_cache_dtype == "fp8_e4m3" - and "tilelang" in (prefill, decode) - ): - raise ValueError("CUDA TileLang FP8 KV does not support HiSparse.") - # The hisparse allow-list in hisparse_hook is platform- but not - # dtype-aware, so an explicitly requested backend still has to clear the - # shared backend/kv-cache-dtype rules before this arm returns early. + # HiSparse also requires compatible KV layouts on both DSA backends. _check_dsa_backend_constraints( kv_cache_dtype, prefill, diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index e6ea8381f59a..8888bcb0d2ca 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -12,6 +12,7 @@ if TYPE_CHECKING: from sglang.srt.disaggregation.mode import DisaggregationMode + from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost class StateType(str, enum.Enum): @@ -66,6 +67,8 @@ class KVArgs: host_kv_data_ptrs: Optional[List[int]] = None host_kv_data_lens: Optional[List[int]] = None host_kv_item_lens: Optional[List[int]] = None + # Decode-owned host KV pool exported for same-node NIXL transfers. + local_shm_pool: Optional[MLATokenToKVPoolHost] = None host_state_data_ptrs: Optional[List[List[int]]] = None host_state_data_lens: Optional[List[List[int]]] = None host_state_item_lens: Optional[List[List[int]]] = None @@ -163,6 +166,9 @@ def register_to_bootstrap(self): # (ABORT_ACK carrying the sender rank). Without it a hold only ends on timeout. supports_deferred_decode_kv_release: bool = False + def release_host_resources(self) -> None: + """Release backend-owned host mappings on graceful shutdown.""" + def teardown(self) -> None: raise NotImplementedError( f"{type(self).__name__} does not support PD role switch teardown" diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index a3cde4755c57..6d43fedb0c1c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -749,6 +749,11 @@ def _init_kv_manager(self) -> CommonKVManager: ) if self.transfer_backend == TransferBackend.NIXL: kv_args.kv_data_mem_kinds = kv_data_mem_kinds + if ( + envs.SGLANG_DISAGGREGATION_NIXL_LOCAL_SHM.get() + and self.scheduler.enable_hisparse + ): + kv_args.local_shm_pool = transfer_kv_pool kv_args.page_size = self.token_to_kv_pool.page_size kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( @@ -3131,6 +3136,13 @@ def _init_staging_handler(self, kv_manager): ) kv_manager._staging_handler = self.staging_handler + def _requires_host_drain(self, decode_req: DecodeRequest) -> bool: + # HiSparse receives KV directly into its own host pool, independently + # of the HiCache host-staging path. + return self.scheduler.enable_hisparse or ( + self.enable_host_receive and decode_req.host_staged + ) + def _release_request(self, decode_req: DecodeRequest) -> None: # HiSparse host pages are the transfer destination; free them with the KV. if self.scheduler.enable_hisparse: @@ -3262,7 +3274,7 @@ def pop_transferred_with_outcomes( [decode_req.req], decode_req.req.return_logprob, ) - requires_host_drain = self._is_host_destination(decode_req) + requires_host_drain = self._requires_host_drain(decode_req) deferred_release_enabled = ( self.enable_deferred_kv_release and decode_req.kv_receiver.kv_mgr.enable_deferred_decode_kv_release @@ -3282,6 +3294,10 @@ def pop_transferred_with_outcomes( # already hold destinations; this path always reaches # _do_release, which clears the tracker. decode_req.kv_receiver.ensure_abort_notified(force_arm=True) + if self.scheduler.enable_hisparse and deferred_release_enabled: + # No notified peer means metadata was never published; no + # remote write can still reference these host pages. + requires_host_drain = decode_req.kv_receiver.abort_notified failed.append( ( i, @@ -3458,9 +3474,15 @@ def resolve_deferred_releases(self) -> None: ready = [] releases = list(self._deferred_releases.items()) for decode_req, _, deadline, _, required_acks in releases: - if self._is_host_destination(decode_req): + if self._requires_host_drain(decode_req): # A timeout cannot prove that a remote write has stopped. - is_drained = decode_req.kv_receiver.is_abort_release_safe() + is_drained = ( + decode_req.kv_receiver.kv_mgr.is_abort_release_safe( + decode_req.req.bootstrap_room, required_acks + ) + if self.scheduler.enable_hisparse + else decode_req.kv_receiver.is_abort_release_safe() + ) is_ready = is_drained else: is_drained = decode_req.kv_receiver.kv_mgr.is_abort_release_safe( diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 85f00eb7c4d3..2a125a388ad7 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -43,6 +43,13 @@ unpack_int_lists, ) from sglang.srt.disaggregation.mode import DisaggregationMode +from sglang.srt.disaggregation.nixl.local_shm import ( + LocalShmExport, + LocalShmMapping, + LocalShmPeer, + LocalShmTransfer, + same_host, +) from sglang.srt.disaggregation.utils import ( build_dsa_tail_transfer_blocks, build_transfer_entry_pairs, @@ -241,6 +248,10 @@ def _homogeneous_kv_mem_kind(kinds: List[str], context: str) -> str: return next(iter(unique)) +_LOCAL_SHM_REAP_INTERVAL_S = 5 +_LOCAL_SHM_RETRY_BACKOFF_S = 2 + + @dataclasses.dataclass(frozen=True) class _KVXferMemSegment: start: int @@ -373,6 +384,7 @@ class KVArgsRegisterInfo: staging_total_size: int = 0 dst_state_data_lens: List[List[int]] = dataclasses.field(default_factory=list) dst_aux_item_lens: Optional[List[int]] = None + local_shm: Optional[dict] = None @classmethod def from_zmq(cls, msg: List[bytes]): @@ -421,6 +433,7 @@ def from_zmq(cls, msg: List[bytes]): agent_name=msg[3].decode("ascii"), agent_metadata=msg[4], dst_kv_ptrs=dst_kv_ptrs, + local_shm=json.loads(msg[25]) if len(msg) > 25 and msg[25] else None, dst_kv_mem_kinds=dst_kv_mem_kinds, dst_aux_ptrs=list(struct.unpack(f"{len(msg[6]) // 8}Q", msg[6])), dst_aux_item_lens=( @@ -633,6 +646,35 @@ def create_backend(): ) logger.info(f"NIXL KVManager initialized with backend: {backend}") + self.local_shm_enabled = envs.SGLANG_DISAGGREGATION_NIXL_LOCAL_SHM.get() + self.local_shm_descriptor = None + self._local_shm_peers: Dict[str, LocalShmPeer] = {} + self._local_shm_lock = threading.Lock() + self._local_shm_shutdown = threading.Event() + self._local_shm_reaper = None + if self.local_shm_enabled: + if not envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get(): + raise ValueError("Local SHM requires deferred decode KV release") + if ( + self.pp_size != 1 + or self.dcp_size != 1 + or envs.SGLANG_DISAGG_STAGING_BUFFER.get() + ): + raise ValueError("Local SHM requires PP1, DCP1 and no staging buffer") + pool = args.local_shm_pool + if disaggregation_mode == DisaggregationMode.DECODE: + from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost + + if not isinstance(pool, MLATokenToKVPoolHost) or any( + kind != "DRAM" for kind in args.kv_data_mem_kinds + ): + raise ValueError( + "Local SHM requires a HiSparse MLA host pool without draft KV" + ) + export = LocalShmExport(pool) + pool.register_host_resource_cleanup(export.close) + self.local_shm_descriptor = export.descriptor + self.register_buffer_to_engine() self.enable_staging = envs.SGLANG_DISAGG_STAGING_BUFFER.get() @@ -692,6 +734,11 @@ def create_backend(): daemon=True, ).start() self._start_bootstrap_thread() + if self.local_shm_enabled: + self._local_shm_reaper = threading.Thread( + target=self._reap_local_shm_loop, daemon=True + ) + self._local_shm_reaper.start() elif self.disaggregation_mode == DisaggregationMode.DECODE: if self.enable_staging: self._init_staging_decode_ctx() @@ -1094,6 +1141,67 @@ def _prefetch_staging_reqs(self, room: int): ) self._staging_ctx.prefetched_rooms.add(room) + def _get_local_shm_peer(self, peer_name: str) -> LocalShmPeer: + # This lock only protects the table, never socket or CUDA operations. + with self._local_shm_lock: + peer = self._local_shm_peers.get(peer_name) + if peer is None: + peer = LocalShmPeer() + self._local_shm_peers[peer_name] = peer + return peer + + def _local_shm_peer_snapshot(self) -> List[Tuple[str, LocalShmPeer]]: + with self._local_shm_lock: + return list(self._local_shm_peers.items()) + + def _reap_local_shm_loop(self) -> None: + while not self._local_shm_shutdown.wait(_LOCAL_SHM_REAP_INTERVAL_S): + self._reap_local_shm_mappings() + + def _reap_local_shm_mappings(self) -> None: + for peer_name, peer in self._local_shm_peer_snapshot(): + # Busy peers are revisited next time; one slow initialization must + # not prevent the reaper from servicing other peers. + if not peer.lock.acquire(blocking=False): + continue + try: + mapping = peer.mapping + if mapping is None or mapping.pending or not mapping.exporter_closed(): + continue + try: + mapping.close() + except Exception: + logger.exception( + "Retaining local SHM mapping for peer %s", peer_name + ) + continue + peer.mapping = None + peer.error = RuntimeError("Local SHM exporter has closed") + peer.retry_at = None + logger.info("Released local SHM mapping for closed peer %s", peer_name) + finally: + peer.lock.release() + + def release_host_resources(self) -> None: + self._local_shm_shutdown.set() + if self._local_shm_reaper is not None: + self._local_shm_reaper.join() + for peer_name, peer in self._local_shm_peer_snapshot(): + with peer.lock: + if peer.mapping is None: + continue + try: + peer.mapping.close() + except Exception: + # Never unmap unresolved writes or still-pinned memory. + logger.exception( + "Retaining local SHM mapping for peer %s", peer_name + ) + else: + peer.mapping = None + peer.error = RuntimeError("Local SHM manager has shut down") + peer.retry_at = None + def check_status(self, bootstrap_room: int): return self.request_status.get(bootstrap_room, KVPoll.WaitingForInput) @@ -1123,7 +1231,11 @@ def _await_handles( any_failed = failure_seen try: for handle in handles: - state = self.agent.check_xfer_state(handle) + state = ( + handle.poll() + if isinstance(handle, LocalShmTransfer) + else self.agent.check_xfer_state(handle) + ) if state == "ERR": any_failed = True elif state != "DONE": @@ -1712,6 +1824,7 @@ def transfer_worker(self, queue: FastQueue, staging_buffer=None, worker_index=0) dst_info.dst_aux_ptrs, dst_info.dst_aux_item_lens ) decode_tp_size = dst_info.decode_tp_size + self._ensure_local_shm_mapping(dst_info) # Skip KV RDMA transfer when there are no pages to send # (e.g., decode-side radix cache matched the entire prefix). @@ -2059,6 +2172,55 @@ def register_buffer_to_engine(self): if not self.state_descs: raise Exception("NIXL memory registration failed for state tensors") + def _ensure_local_shm_mapping( + self, decode_kv_args: KVArgsRegisterInfo + ) -> Optional[LocalShmMapping]: + if not self.local_shm_enabled or decode_kv_args.local_shm is None: + return None + agent_name = decode_kv_args.agent_name + peer = self._get_local_shm_peer(agent_name) + with peer.lock: + mapping = peer.mapping + if mapping is not None: + return mapping + error = peer.error + if error is not None: + retry_at = peer.retry_at + if retry_at is None or time.monotonic() < retry_at: + raise RuntimeError(str(error)) from error + try: + if not same_host(decode_kv_args.local_shm): + return None + if ( + not (self.is_mla_backend or self.is_hybrid_mla_backend) + or self.attn_tp_size != decode_kv_args.decode_tp_size + or decode_kv_args.dst_dcp_size != 1 + or decode_kv_args.staging_base_ptr + or decode_kv_args.dst_kv_mem_kinds + != ["DRAM"] * len(self.kv_args.kv_data_ptrs) + or decode_kv_args.dst_kv_item_lens != self.kv_args.kv_item_lens + or decode_kv_args.local_shm["items"] != self.kv_args.kv_item_lens + or decode_kv_args.dst_kv_layer_ids != self.kv_args.kv_layer_ids + ): + raise ValueError( + "Local SHM requires matching flat MLA KV layouts and TP sizes" + ) + mapping = LocalShmMapping(decode_kv_args.local_shm, self.kv_args.gpu_id) + except Exception as exc: + peer.error = exc + error = f"Local SHM peer initialization failed: {exc}" + if isinstance(exc, OSError): + peer.retry_at = time.monotonic() + _LOCAL_SHM_RETRY_BACKOFF_S + else: + peer.retry_at = None + logger.error("%s", error) + raise RuntimeError(error) from exc + peer.mapping = mapping + peer.error = None + peer.retry_at = None + logger.info("Local SHM D2H enabled for peer %s", agent_name) + return mapping + def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo): agent_name = decode_kv_args.agent_name if agent_name in self.decode_kv_args_table: @@ -2069,6 +2231,12 @@ def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo): ) if decode_kv_args.requires_dcp_relayout: self._init_dcp_pack_buffers_once(decode_kv_args.dst_dcp_size) + try: + self._ensure_local_shm_mapping(decode_kv_args) + except RuntimeError: + # Preserve registration so a later request can retry a transient + # connection failure. The transfer path still fails closed. + pass self.decode_kv_args_table[agent_name] = decode_kv_args self.agent.add_remote_agent(decode_kv_args.agent_metadata) if self.disaggregation_mode == DisaggregationMode.PREFILL: @@ -2339,6 +2507,22 @@ def send_kvcache( dst_gpu_id: int, dst_mem_kind: str = "VRAM", ): + if self.local_shm_enabled: + peer = self._get_local_shm_peer(peer_name) + with peer.lock: + mapping = self._ensure_local_shm_mapping( + self.decode_kv_args_table[peer_name] + ) + if mapping is not None: + return mapping.submit( + self.kv_args.kv_data_ptrs, + self.kv_args.kv_data_lens, + prefill_kv_indices, + dst_kv_indices, + # The worker publishes status through conclude_transfer only + # after both local DMA and NIXL state transfers have drained. + lambda: None, + ) assert self.src_mem_kind is not None return self._send_kvcache_generic( peer_name=peer_name, @@ -3328,6 +3512,10 @@ def add_transfer_request( info.dst_port for info in self.transfer_infos[bootstrap_room].values() ) shard_idx = session_port_sum % len(self.transfer_queues) + ready_event = wait_event + if ready_event is None and self.local_shm_enabled: + ready_event = torch.cuda.Event() + ready_event.record() self.transfer_queues[shard_idx].put( TransferKVChunk( room=bootstrap_room, @@ -3335,9 +3523,9 @@ def add_transfer_request( index_slice=index_slice, is_last_chunk=is_last_chunk, prefill_aux_index=aux_index, + wait_event=ready_event, state_indices=state_indices, num_kv_tokens=num_kv_tokens, - wait_event=wait_event, ) ) return None @@ -3535,6 +3723,11 @@ def build_receiver_registration_message(self) -> List[bytes]: str(self.dcp_size).encode("ascii"), str(self.dcp_rank).encode("ascii"), packed_state_data_lens, + ( + json.dumps(self.local_shm_descriptor).encode("ascii") + if self.local_shm_descriptor is not None + else b"" + ), ] def build_receiver_metadata_message( diff --git a/python/sglang/srt/disaggregation/nixl/local_shm.py b/python/sglang/srt/disaggregation/nixl/local_shm.py new file mode 100644 index 000000000000..5e5b36e9a198 --- /dev/null +++ b/python/sglang/srt/disaggregation/nixl/local_shm.py @@ -0,0 +1,331 @@ +"""Opt-in same-node HiSparse host KV writes. NIXL still carries notifications. + +An exported pool has a fresh endpoint and capability for its entire lifetime. +Imported mappings retain the backing file even if the exporting process exits. +""" + +from __future__ import annotations + +import logging +import mmap +import operator +import os +import secrets +import socket +import struct +import threading +from contextlib import ExitStack +from pathlib import Path +from typing import TYPE_CHECKING, Callable, Optional, Sequence + +import msgspec + +if TYPE_CHECKING: + import torch + + from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost + +# A failed unregister must not turn a still-pinned mapping into freed memory. +_PINNED_MAPPING_QUARANTINE = [] +logger = logging.getLogger(__name__) + + +# Internal state, never serialized; one lock protects one peer's lifecycle. +class LocalShmPeer(msgspec.Struct): + lock: threading.RLock = msgspec.field(default_factory=threading.RLock) + mapping: Optional[LocalShmMapping] = None + error: Optional[Exception] = None + retry_at: Optional[float] = None + + +def host_identity() -> list[str]: + # Abstract AF_UNIX sockets are scoped to the network namespace. + return [ + Path("/proc/sys/kernel/random/boot_id").read_text().strip(), + os.readlink("/proc/self/ns/net"), + ] + + +def same_host(descriptor: dict) -> bool: + return descriptor.get("host") == host_identity() + + +class LocalShmExport: + def __init__(self, pool: MLATokenToKVPoolHost) -> None: + buffer = pool.kv_buffer + if pool.fd is None or not buffer.is_contiguous(): + raise ValueError("Local SHM requires a contiguous shared host KV pool") + ptrs, lengths, items = pool.get_contiguous_buf_infos() + self.token = secrets.token_bytes(32) + self.stopped = threading.Event() + name = "sglang-pd-" + secrets.token_hex(16) + with ExitStack() as rollback: + self.fd = os.dup(pool.fd) + rollback.callback(os.close, self.fd) + self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) + rollback.callback(self.socket.close) + self.socket.bind("\0" + name) + self.socket.listen(16) + self.socket.settimeout(0.2) + self.descriptor = dict( + version=1, + host=host_identity(), + endpoint=name, + token=self.token.hex(), + size=buffer.numel() * buffer.element_size(), + offsets=[ptr - buffer.data_ptr() for ptr in ptrs], + lengths=lengths, + items=items, + ) + self.thread = threading.Thread(target=self._serve, daemon=True) + self.thread.start() + rollback.pop_all() + + def _serve(self) -> None: + while not self.stopped.is_set(): + try: + conn, _ = self.socket.accept() + except socket.timeout: + continue + except OSError: + return + with conn: + try: + conn.settimeout(1) + _, uid, _ = struct.unpack( + "3i", conn.getsockopt(socket.SOL_SOCKET, socket.SO_PEERCRED, 12) + ) + if uid != os.getuid(): + continue + token = bytearray() + while len(token) < len(self.token): + part = conn.recv(len(self.token) - len(token)) + if not part: + break + token.extend(part) + if secrets.compare_digest(token, self.token): + socket.send_fds(conn, [b"F"], [self.fd]) + except (OSError, ValueError): + continue + + def close(self) -> None: + self.stopped.set() + self.socket.close() + self.thread.join() + if self.fd is not None: + os.close(self.fd) + self.fd = None + + +def build_copy_plan( + src_lengths_bytes: Sequence[int], + dst_lengths_bytes: Sequence[int], + page_sizes_bytes: Sequence[int], + src_indices: Sequence[int], + dst_indices: Sequence[int], +) -> list[tuple[int, int, int, int]]: + """Validate a chunk and return (region, src offset, dst offset, byte count). + + Offsets and sizes are bytes; source/destination indices select KV pages. + No writes may be submitted until the entire plan has been validated. + """ + if len(src_lengths_bytes) != len(dst_lengths_bytes) or len(page_sizes_bytes) != len( + src_lengths_bytes + ): + raise ValueError("Local SHM KV region count mismatch") + if len(src_indices) != len(dst_indices): + raise ValueError("Local SHM page count mismatch") + pairs = [ + (operator.index(s), operator.index(d)) for s, d in zip(src_indices, dst_indices) + ] + plan = [] + for layer, (src_size, dst_size, item) in enumerate( + zip(src_lengths_bytes, dst_lengths_bytes, page_sizes_bytes) + ): + if item <= 0: + raise ValueError("Invalid local SHM page size") + for src, dst in pairs: + if ( + src < 0 + or dst < 0 + or (src + 1) * item > src_size + or (dst + 1) * item > dst_size + ): + raise ValueError("Local SHM page index outside KV region") + plan.append((layer, src * item, dst * item, item)) + return plan + + +class LocalShmTransfer: + def __init__( + self, + owner: LocalShmMapping, + event: Optional[torch.cuda.Event], + stream: torch.cuda.Stream, + notify: Callable[[], None], + error: Optional[Exception] = None, + ) -> None: + self.owner = owner # Keep the mapping and stream alive until DMA ends. + self.event = event + self.stream = stream + self.notify = notify + self.error = error + self.state = "PROC" + + def poll(self) -> str: + if self.state != "PROC": + return self.state + if self.event is None: + raise RuntimeError("Local SHM completion event could not be recorded") + if not self.event.query(): + return "PROC" + if self.error is None: + try: + self.notify() + except Exception as exc: + self.error = exc + self.state = "ERR" if self.error is not None else "DONE" + self.owner.pending.remove(self) + return self.state + + +class LocalShmMapping: + def __init__(self, descriptor: dict, gpu_id: int) -> None: + import torch + + from sglang.srt.mem_cache.pool_host.common import _cuda_host_register + + if descriptor.get("version") != 1 or not same_host(descriptor): + raise ValueError("Incompatible local SHM descriptor") + self.gpu_id = gpu_id + self.endpoint = descriptor["endpoint"] + self.pending = set() + self.buffer = None + self.mapping = None + self.fd = None + self.lengths = descriptor["lengths"] + self.items = descriptor["items"] + offsets = descriptor["offsets"] + size = descriptor["size"] + if ( + size <= 0 + or len(offsets) != len(self.lengths) + or len(offsets) != len(self.items) + ): + raise ValueError("Invalid local SHM regions") + for offset, length, item in zip(offsets, self.lengths, self.items): + if offset < 0 or length <= 0 or offset + length > size or item <= 0: + raise ValueError("Invalid local SHM region bounds") + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn: + conn.settimeout(5) + conn.connect("\0" + descriptor["endpoint"]) + conn.sendall(bytes.fromhex(descriptor["token"])) + message, fds, _, _ = socket.recv_fds(conn, 1, 1, socket.MSG_CMSG_CLOEXEC) + if message != b"F" or len(fds) != 1: + for fd in fds: + os.close(fd) + raise RuntimeError("Local SHM FD exchange failed") + self.fd = fds[0] + try: + if os.fstat(self.fd).st_size < size: + raise ValueError("Local SHM backing size mismatch") + self.mapping = mmap.mmap(self.fd, size, flags=mmap.MAP_SHARED) + self.buffer = torch.frombuffer(self.mapping, dtype=torch.uint8) + with torch.cuda.device(gpu_id): + _cuda_host_register( + self.buffer, registration_granularity_bytes=max(self.items) + ) + self.ptrs = [self.buffer.data_ptr() + offset for offset in offsets] + except Exception: + if getattr(self.buffer, "_sglang_cuda_host_registered_ranges", []): + _PINNED_MAPPING_QUARANTINE.append(self) + raise + self.buffer = None + if self.mapping is not None: + self.mapping.close() + os.close(self.fd) + raise + + def submit( + self, + src_ptrs: Sequence[int], + src_lengths: Sequence[int], + src_indices: Sequence[int], + dst_indices: Sequence[int], + notify: Callable[[], None], + ) -> LocalShmTransfer: + import torch + from cuda.bindings import runtime + + if len(src_ptrs) != len(self.ptrs): + raise ValueError("Local SHM source region count mismatch") + plan = build_copy_plan( + src_lengths, self.lengths, self.items, src_indices, dst_indices + ) + with torch.cuda.device(self.gpu_id): + stream = torch.cuda.Stream() + event = torch.cuda.Event() + handle = LocalShmTransfer(self, event, stream, notify) + self.pending.add(handle) + try: + for layer, src, dst, size in plan: + (rc,) = runtime.cudaMemcpyAsync( + self.ptrs[layer] + dst, + src_ptrs[layer] + src, + size, + runtime.cudaMemcpyKind.cudaMemcpyDeviceToHost, + stream.cuda_stream, + ) + if int(rc): + raise RuntimeError(f"Local SHM cudaMemcpyAsync failed: {rc}") + except Exception as exc: + logger.exception( + "Local SHM submission failed; waiting for queued writes" + ) + handle.error = exc + finally: + try: + event.record(stream) + except Exception as exc: + # Unknown completion: retain mapping and never acknowledge drain. + logger.exception( + "Local SHM completion is unknown; retaining mapping" + ) + handle.error = exc + handle.event = None + return handle + + def exporter_closed(self) -> bool: + """Only connection refusal proves this abstract endpoint is gone.""" + try: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn: + conn.settimeout(1) + conn.connect("\0" + self.endpoint) + except ConnectionRefusedError: + return True + except OSError: + # Timeout/resource pressure does not establish peer death. + return False + return False + + def close(self) -> None: + from sglang.srt.mem_cache.pool_host.common import _cuda_host_unregister + + if self.pending: + raise RuntimeError("Cannot close local SHM while writes may be in flight") + if self.buffer is not None: + import torch + + with torch.cuda.device(self.gpu_id): + _cuda_host_unregister(self.buffer) + if getattr(self.buffer, "_sglang_cuda_host_registered_ranges", []): + raise RuntimeError( + "Local SHM host unregister failed; retaining mapping" + ) + self.buffer = None + if self.mapping is not None: + self.mapping.close() + self.mapping = None + if self.fd is not None: + os.close(self.fd) + self.fd = None diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index bd4646cb582c..849fd9b9406b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -812,6 +812,9 @@ class Envs: SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300) # A wedged RDMA stack fails startup here instead of at the scheduler watchdog. SGLANG_DISAGGREGATION_ENGINE_INIT_TIMEOUT = EnvInt(60) + # Enable same-node shared-memory D2H for HiSparse CPU KV pools. + # Other transfers continue through NIXL. + SGLANG_DISAGGREGATION_NIXL_LOCAL_SHM = EnvBool(False) SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX") SGLANG_DISAGGREGATION_NIXL_BACKEND_PARAMS = EnvStr("{}") SGLANG_DISAGG_PREFILL_EARLY_SEND_CACHED_PREFIX = EnvBool(True) diff --git a/python/sglang/srt/managers/hisparse_coordinator.py b/python/sglang/srt/managers/hisparse_coordinator.py index 9d31623e2916..da8ee2eaaa0b 100644 --- a/python/sglang/srt/managers/hisparse_coordinator.py +++ b/python/sglang/srt/managers/hisparse_coordinator.py @@ -59,11 +59,15 @@ def load_blocks_to_device_buffer_mha(*args, **kwargs): from sglang.srt.mem_cache.hisparse_memory_pool import ( HiSparseDSATokenToKVPool, ) -from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool, ReqToTokenPool +from sglang.srt.mem_cache.memory_pool import ( + HybridLinearKVPool, + MiniMaxSparseKVPool, + ReqToTokenPool, +) from sglang.srt.mem_cache.memory_pool_host import DeepSeekV4PagedHostPool from sglang.srt.mem_cache.pool_host.mha import HiSparseMHATokenToKVPoolHost from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost -from sglang.srt.runtime_context import get_parallel +from sglang.srt.runtime_context import get_disagg, get_parallel if TYPE_CHECKING: from sglang.kernels.ops.kvcache.hisparse import HiSparseSpecState @@ -563,6 +567,7 @@ def __init__( self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.top_k = top_k + self.num_tail_slots = 0 self.device_buffer_size = device_buffer_size self.device = device self.swap_in_block_size = swap_in_block_size @@ -577,6 +582,7 @@ def __init__( ) self.compress_ratio = self.token_to_kv_pool_allocator.compress_ratio + self.full_attention_layer_id_mapping = None kvcache = self.token_to_kv_pool_allocator.get_kvcache() self.is_dsv4_hisparse = isinstance( self.token_to_kv_pool_allocator, DeepSeekV4HiSparseTokenToKVPoolAllocator @@ -627,7 +633,28 @@ def __init__( ) self.item_size_bytes = self.mem_pool_device.bytes_per_token_k else: - self.mem_pool_device: HiSparseDSATokenToKVPool = kvcache + self.mem_pool_device: HiSparseDSATokenToKVPool = ( + self.token_to_kv_pool_allocator.hisparse_kvcache + ) + if isinstance(kvcache, HybridLinearKVPool): + self.full_attention_layer_id_mapping = ( + kvcache.full_attention_layer_id_mapping + ) + if shared_index_layers is not None: + shared_index_layers = [ + shared_index_layers[layer_id] + for layer_id in self.full_attention_layer_id_mapping + ] + self.num_tail_slots = self.mem_pool_device.index_kpool - 1 + if self.num_tail_slots and _is_xpu: + raise ValueError( + "HiSparse KPool tail columns are not supported on XPU yet." + ) + self.top_k += self.num_tail_slots + if self.device_buffer_size < self.top_k: + raise ValueError( + "HiSparse device_buffer_size must cover top-k and tail columns." + ) self.mem_pool_host = MLATokenToKVPoolHost( buffer_info=self.mem_pool_device.get_device_pool_infos()[ 0 @@ -643,6 +670,13 @@ def __init__( page_size=self.mem_pool_device.page_size, layout="layer_first", override_kv_cache_dim=self.mem_pool_device.kv_cache_dim, + allocator_type=( + "shm" + if envs.SGLANG_DISAGGREGATION_NIXL_LOCAL_SHM.get() + and get_disagg().disaggregation_mode == "decode" + and get_disagg().disaggregation_transfer_backend == "nixl" + else "default" + ), ) self.item_size_bytes = self.mem_pool_host.token_stride_size self.page_size = self.mem_pool_device.page_size @@ -1311,6 +1345,8 @@ def naive_load_topk( assert not self.is_dsv4_hisparse, ( "naive_load_topk is not implemented for dsv4 hisparse" ) + if self.full_attention_layer_id_mapping is not None: + layer_id = self.full_attention_layer_id_mapping[layer_id] num_reqs = req_pool_indices.size(0) top_k_indices = torch.full( (num_reqs, self.top_k), -1, dtype=torch.int32, device=self.device @@ -1318,20 +1354,16 @@ def naive_load_topk( for i in range(num_reqs): seq_len = int(seq_lens[i].item()) - top_n = min(seq_len, self.top_k) + row = top_k_tokens[i].to(dtype=torch.int64) + valid = (row >= 0) & (row < seq_len) + if seq_len <= self.device_buffer_size: + history_end = self.top_k - self.num_tail_slots + valid[min(seq_len, history_end) : history_end] = False + selected_tokens = row[valid] + top_n = selected_tokens.numel() if top_n == 0: continue - req_idx = int(req_pool_indices[i].item()) - selected_tokens = top_k_tokens[i, :top_n].to(dtype=torch.int64) - - assert torch.all(selected_tokens >= 0), ( - f"Req {req_idx}: selected tokens contain negative positions" - ) - assert torch.all(selected_tokens < seq_len), ( - f"Req {req_idx}: selected tokens {selected_tokens.tolist()} " - f"out of range for seq_len={seq_len}" - ) if seq_len <= self.device_buffer_size: device_indices = self.req_to_device_buffer[req_idx, selected_tokens] @@ -1371,7 +1403,7 @@ def naive_load_topk( io_backend="kernel", ) - top_k_indices[i, :top_n] = device_indices.to(torch.int32) + top_k_indices[i, valid] = device_indices.to(torch.int32) return top_k_indices @@ -1518,6 +1550,8 @@ def _run_swap_in_kernel( if self.is_dsv4_hisparse else load_cache_to_device_buffer_mla ) + if not self.is_dsv4_hisparse and not _is_xpu: + plan["num_tail_slots"] = self.num_tail_slots skip_io_kwargs = {} if _is_xpu else dict(skip_io=self.skip_io) swap_in_fn( top_k_tokens=top_k_result, @@ -1609,6 +1643,8 @@ def swap_in_selected_pages( With prefetch enabled, anchors swap in synchronously (recording the miss plan) and prefetch their skip layers' copies; skip layers just wait. """ + if self.full_attention_layer_id_mapping is not None: + layer_id = self.full_attention_layer_id_mapping[layer_id] if not self.enable_prefetch: return self._run_swap_in_kernel( req_pool_indices, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 7fd47e760428..3e4f4ff62b3b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2036,6 +2036,8 @@ def release_host_resources(self) -> None: self.uocr_runtime.close() # Release pinned host buffers in userspace on graceful shutdown; see # HostKVCache.destroy. Called from run_scheduler_process's finally. + if self.disagg_prefill_bootstrap_queue is not None: + self.disagg_prefill_bootstrap_queue.kv_manager.release_host_resources() if self.hisparse_coordinator is not None: self.hisparse_coordinator.destroy() self.tree_cache.release_host_resources() diff --git a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py index a81cc25e8c57..6b1add81abb8 100644 --- a/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py +++ b/python/sglang/srt/managers/scheduler_components/pool_stats_observer.py @@ -102,6 +102,10 @@ def get_max_pool_usage(self) -> float: usage = max(usage, self.swa_token_usage) if self.is_hybrid_ssm: usage = max(usage, self.mamba_usage) + if self.is_hisparse: + usage = max( + usage, self.hisparse_device_token_usage, self.hisparse_host_token_usage + ) assert usage is not None and usage >= 0, f"{usage=} is not valid" return usage diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 6a9f8fad8b00..ba743a2fc31d 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -12,6 +12,7 @@ HiSparseC4DevicePool, ) from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool from sglang.srt.platforms import current_platform from sglang.srt.utils.common import get_num_new_pages @@ -26,11 +27,19 @@ def __init__( page_size: int, dtype: torch.dtype, device: torch.device, - kvcache: HiSparseDSATokenToKVPool | MiniMaxSparseKVPool, + kvcache: HiSparseDSATokenToKVPool | MiniMaxSparseKVPool | HybridLinearKVPool, need_sort: bool, host_to_device_ratio: int = 2, ): self._kvcache = kvcache + # Keep the wrapper for Mamba state access; map only DSA token locations. + self.hisparse_kvcache = ( + kvcache.full_kv_pool if isinstance(kvcache, HybridLinearKVPool) else kvcache + ) + if isinstance(kvcache, HybridLinearKVPool) and not isinstance( + self.hisparse_kvcache, HiSparseDSATokenToKVPool + ): + raise TypeError("Hybrid HiSparse requires a HiSparse DSA backing pool") self._size_full = size * host_to_device_ratio self._size_hisparse = size self.compress_ratio = 1 @@ -75,7 +84,7 @@ def __init__( self.release_pages = None self.free_group = None self.clear() - self._kvcache.register_mapping( + self.hisparse_kvcache.register_mapping( weakref.proxy(self.full_to_hisparse_device_index_mapping) ) @@ -203,7 +212,7 @@ def get_last_loc_compressed(self, last_locs: torch.Tensor): return last_locs def get_last_loc_hisparse_device(self, last_locs: torch.Tensor): - return self._kvcache._translate_loc_to_hisparse_device(last_locs) + return self.hisparse_kvcache._translate_loc_to_hisparse_device(last_locs) def alloc_extend( self, @@ -265,7 +274,9 @@ def alloc_decode( ) def free_hisparse(self, free_indices: torch.Tensor): - hisparse_indices = self._kvcache._translate_loc_to_hisparse_device(free_indices) + hisparse_indices = self.hisparse_kvcache._translate_loc_to_hisparse_device( + free_indices + ) hisparse_indices = hisparse_indices[hisparse_indices > 0] self.free_hisparse_indices(hisparse_indices) self.full_to_hisparse_device_index_mapping[free_indices] = 0 diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index c938f5486c49..a6190d16af10 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -2208,6 +2208,7 @@ def _build_dsa_kv_pool( max_total_num_tokens: int, max_running_requests: int, dsa_pool_class: type, + full_attention_layer_ids: Optional[list[int]] = None, ) -> KVCache: from sglang.srt.layers.cp.utils import get_glm_dsa_cp_layer_shard_info @@ -2250,11 +2251,20 @@ def _build_dsa_kv_pool( PoolCls = dsa_pool_class if self.loc_space_scale > 1: pool_kwargs["dcp_replicated"] = True + # Hybrid pools expose global layer IDs but store only DSA layers. + start_layer = self.layer_info.start_layer + end_layer = self.layer_info.end_layer + layer_num = self.layer_info.num_effective_layers + if full_attention_layer_ids is not None: + start_layer = 0 + end_layer = layer_num = len(full_attention_layer_ids) if _should_elide_dsa_index_k(is_draft_worker=self.is_draft_worker): pool_kwargs["skip_topk_layers"] = [ dsa_layer_skips_topk(self.model_config.hf_config, layer_id) - for layer_id in range( - self.layer_info.start_layer, self.layer_info.end_layer + for layer_id in ( + full_attention_layer_ids + if full_attention_layer_ids is not None + else range(start_layer, end_layer) ) ] token_to_kv_pool = PoolCls( @@ -2263,15 +2273,15 @@ def _build_dsa_kv_pool( dtype=self.kv_cache_dtype, kv_lora_rank=self.model_config.kv_lora_rank, qk_rope_head_dim=self.model_config.qk_rope_head_dim, - layer_num=self.layer_info.num_effective_layers, + layer_num=layer_num, device=self.device, kv_cache_dim=calculate_mla_kv_cache_dim( model_config=self.model_config, kv_cache_dtype=self.kv_cache_dtype, ), enable_memory_saver=get_exec().features.enable_memory_saver, - start_layer=self.layer_info.start_layer, - end_layer=self.layer_info.end_layer, + start_layer=start_layer, + end_layer=end_layer, index_head_dim=get_dsa_index_head_dim(self.model_config.hf_config), index_kpool=get_dsa_index_kpool(self.model_config.hf_config), index_kpool_compress=get_dsa_index_kpool_compress( @@ -2563,6 +2573,13 @@ def _build_hybrid_linear_kv_pool( tail_extra_slots=(max_speculative_num_draft_tokens() or 0), max_running_requests=(req_to_token_pool.req_to_token.shape[0]), ) + if get_memory().enable_hisparse and not self.is_draft_worker: + extra_args["full_kv_pool"] = self._build_dsa_kv_pool( + max_total_num_tokens=max_total_num_tokens, + max_running_requests=req_to_token_pool.req_to_token.shape[0], + dsa_pool_class=dsa_pool_class, + full_attention_layer_ids=full_attention_layer_ids, + ) quant_method = self._build_mha_quant_method( num_layers=len(full_attention_layer_ids) ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 954d78c0f370..eca43f89a29f 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -4131,7 +4131,7 @@ def __init__( start_layer: Optional[int] = None, full_kv_pool_class: Optional[type] = None, quant_method=None, - # When provided (shared-KV-pool path), use this pool for the + # When provided (shared-KV-pool or HiSparse path), use this pool for the # full-attention layers instead of constructing one internally. full_kv_pool: Optional[KVCache] = None, post_capture_active: bool = False, @@ -4153,8 +4153,7 @@ def __init__( self.use_mla = use_mla self.use_dsa = use_dsa if full_kv_pool is not None: - # Shared-KV-pool path: the caller built an MHA view - # aliasing the shared byte buffer. + # Reuse the supplied shared-memory view or HiSparse backing pool. self.full_kv_pool = full_kv_pool elif not use_mla: TokenToKVPoolClass = ( @@ -4304,6 +4303,10 @@ def get_kv_layer_ids(self): return layer_ids * (len(data_ptrs) // len(layer_ids)) return layer_ids if self.use_mla else layer_ids * 2 + def translate_loc_to_hisparse_device(self, indices: torch.Tensor): + # Mamba state slots are request-indexed and do not use this mapping. + return self.full_kv_pool.translate_loc_to_hisparse_device(indices) + def get_state_buf_infos(self): mamba_data_ptrs, mamba_data_lens, mamba_item_lens = ( self.mamba_pool.get_contiguous_buf_infos() diff --git a/python/sglang/srt/mem_cache/pool_host/base.py b/python/sglang/srt/mem_cache/pool_host/base.py index b6cbf2901c01..2fbbc7feb3a4 100644 --- a/python/sglang/srt/mem_cache/pool_host/base.py +++ b/python/sglang/srt/mem_cache/pool_host/base.py @@ -7,7 +7,7 @@ from contextlib import contextmanager from contextvars import ContextVar from functools import wraps -from typing import Optional, TypeGuard +from typing import Callable, Optional, TypeGuard import torch @@ -202,6 +202,7 @@ def wrapper(self, *args, **kwargs): class HostKVCache(abc.ABC): + _host_resource_cleanup: Optional[list[Callable[[], None]]] = None dcp_size = 1 dcp_rank = 0 shared_allocation_domain = None @@ -322,6 +323,12 @@ def _allocate_host_buffer(self) -> None: self.kv_buffer = self.init_kv_buffer() self.fd = getattr(self.allocator, "fd", None) + def register_host_resource_cleanup(self, cleanup: Callable[[], None]) -> None: + """Close users of the host allocation before unregistering its memory.""" + if self._host_resource_cleanup is None: + self._host_resource_cleanup = [] + self._host_resource_cleanup.append(cleanup) + def destroy(self): """Unregister pinned host buffers in userspace before process exit. @@ -332,7 +339,12 @@ def destroy(self): """ if getattr(self, "_destroyed", False): return - self._destroyed = True + # A failed callback keeps both itself and the allocation alive for retry. + # Completed callbacks are removed so retries do not close them twice. + while self._host_resource_cleanup: + self._host_resource_cleanup[0]() + self._host_resource_cleanup.pop(0) + self._host_resource_cleanup = None buffers = getattr(self, "kv_buffer", None) if buffers is not None and self.pin_memory and (_is_cuda or _is_hip): if not isinstance(buffers, (list, tuple)): @@ -341,6 +353,7 @@ def destroy(self): if buf is not None: _cuda_host_unregister(buf) self.kv_buffer = None + self._destroyed = True # Bytes between strided device KV rows; None (the default, and what a pool # built without `_init_device_row_stride` keeps) means packed rows. diff --git a/python/sglang/srt/mem_cache/static_kv_storage.py b/python/sglang/srt/mem_cache/static_kv_storage.py index 06f680da4728..61b84ed99e12 100644 --- a/python/sglang/srt/mem_cache/static_kv_storage.py +++ b/python/sglang/srt/mem_cache/static_kv_storage.py @@ -16,6 +16,7 @@ import msgspec import torch +from sglang.srt.mem_cache.allocator.hisparse import HiSparseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.page_interleave import page_interleave_shard_size from sglang.srt.mem_cache.kv.plan import ( TREE_STATE, @@ -488,18 +489,23 @@ def usage(self) -> Tuple[SubPoolUsage, ...]: ), ] else: + # Admission also checks the hot pool; occupancy counts logical KV. + is_hisparse = isinstance(allocator, HiSparseTokenToKVPoolAllocator) + token_allocator = ( + allocator.logical_attn_allocator if is_hisparse else allocator + ) # A state model's token capacity is what its requests can schedule. capacity = ( - self.requests.schedulable_token_capacity(allocator.size) + self.requests.schedulable_token_capacity(token_allocator.size) if shape.has_state - else shape.token_capacity + else (token_allocator.size if is_hisparse else shape.token_capacity) ) usage = [ SubPoolUsage( name="full", metrics="full", capacity=capacity, - available=allocator.available_size(), + available=token_allocator.available_size(), ) ] if shape.has_state: diff --git a/test/registered/kernels/ops/kvcache/test_hisparse.py b/test/registered/kernels/ops/kvcache/test_hisparse.py index a41618222618..383480a935c8 100644 --- a/test/registered/kernels/ops/kvcache/test_hisparse.py +++ b/test/registered/kernels/ops/kvcache/test_hisparse.py @@ -852,3 +852,72 @@ def test_load_cache_to_device_buffer_rocm_large_lru_writeback() -> None: if __name__ == "__main__": sys.exit(pytest.main([__file__, "-v", "-s"])) + + +@pytest.mark.skipif(is_xpu(), reason="KPool tail columns use the JIT kernel.") +@pytest.mark.parametrize("block_size", [256, 960, 1024]) +@pytest.mark.parametrize("seq_len", [3, 8192]) +def test_kpool_tail_columns_and_chunked_scan(block_size, seq_len): + hot_size, top_k, tail_slots = 4096, 2051, 3 + # Raw FP8 rows are copied byte-for-byte, without interpreting their values. + item_bytes = 512 + host = torch.empty((8192, 1, item_bytes), dtype=torch.uint8, pin_memory=True) + host.copy_( + torch.arange(host.numel() // 4, dtype=torch.int32) + .view(torch.uint8) + .view_as(host) + ) + device = torch.zeros( + (hot_size + 1, 1, item_bytes), dtype=torch.uint8, device=DEVICE + ) + device[:hot_size].copy_(host[:hot_size]) + device[hot_size].copy_(host[seq_len - 1]) + locs = torch.arange(hot_size + 1, dtype=torch.int32, device=DEVICE)[None, :] + tokens = locs.clone() + lru = torch.arange(hot_size, dtype=torch.int16, device=DEVICE)[None, :] + selected = torch.full((1, top_k), -1, dtype=torch.int32, device=DEVICE) + if seq_len == 3: + selected[0, :3] = torch.tensor([0, 1, 2], device=DEVICE) + selected[0, 3] = 0 # Stale padding must not become an attention entry. + selected[0, -3:] = torch.tensor([1, 2, -1], device=DEVICE) + else: + # Mix hits and misses across multiple scan iterations, plus duplicates. + selected[0, :2048] = torch.arange(3072, 5120, device=DEVICE) + selected[0, -3:] = torch.tensor([8191, 5000, -1], device=DEVICE) + out = torch.full_like(selected, 123) + kwargs = dict( + top_k_tokens=selected, + device_buffer_tokens=tokens, + host_cache_locs=torch.arange(8192, dtype=torch.int64, device=DEVICE)[None, :], + device_buffer_locs=locs, + host_cache=host, + device_buffer=device, + top_k_device_locs=out, + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=DEVICE), + lru_slots=lru, + item_size_bytes=item_bytes, + num_top_k=top_k, + hot_buffer_size=hot_size, + page_size=1, + block_size=block_size, + num_tail_slots=tail_slots, + ) + valid = selected[0] >= 0 + if seq_len == 3: + valid[3:2048] = False + for _ in range(2): + load_cache_to_device_buffer_mla(**kwargs) + get_device_module().synchronize() + assert torch.all(out[0, ~valid] == -1) + assert torch.all((out[0, valid] >= 0) & (out[0, valid] <= hot_size)) + torch.testing.assert_close( + device[out[0, valid].long()].cpu(), + host[selected[0, valid].long().cpu()], + rtol=0, + atol=0, + ) + torch.testing.assert_close( + lru[0].sort().values, + torch.arange(hot_size, dtype=torch.int16, device=DEVICE), + ) diff --git a/test/registered/unit/disaggregation/test_common_sender_lifecycle.py b/test/registered/unit/disaggregation/test_common_sender_lifecycle.py index 4224a7cc5362..7e9166061363 100644 --- a/test/registered/unit/disaggregation/test_common_sender_lifecycle.py +++ b/test/registered/unit/disaggregation/test_common_sender_lifecycle.py @@ -28,6 +28,8 @@ def make_sender(manager_type, room=42, *, sender_type=CommonKVSender): manager.attn_cp_size = 1 manager.attn_cp_rank = 0 manager.enable_staging = False + if manager_type is NixlKVManager: + manager.local_shm_enabled = False manager._staging_outstanding = {} manager._deferred_ack_targets = {} manager._deferred_ack_poisoned_rooms = set() diff --git a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py index 274de35a0c07..a3fee802c75f 100644 --- a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py +++ b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py @@ -923,6 +923,85 @@ def test_non_abort_failure_notifies_and_defers(self): self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, []) self.assertEqual(list(q.queue.items()), []) + def test_hisparse_holds_all_pools_until_every_rank_drains(self): + mgr = _make_manager() + mgr.enable_deferred_decode_kv_release = True + q = _make_pop_queue(mgr) + q.scheduler.enable_hisparse = True + q.scheduler.hisparse_coordinator = SimpleNamespace(request_finished=MagicMock()) + q.gloo_group = object() + q.deferred_kv_release_timeout = -1 + entry = _make_failed_req(mgr, room=7, idx=3, abort_notified=False) + q.queue.join([entry]) + generation = mgr.register_deferred_abort_room(7) + host_release = q.scheduler.hisparse_coordinator.request_finished + + with ( + patch.object(decode_mod, "prepare_abort"), + patch.object(decode_mod, "release_kv_cache") as device_release, + patch("torch.distributed.get_world_size", return_value=2), + patch("torch.distributed.all_reduce") as reduce, + ): + q.pop_transferred() + host_release.assert_not_called() + device_release.assert_not_called() + self.assertEqual(list(q.queue.items()), []) + self.assertEqual(len(q._deferred_releases), 1) + self.assertEqual(entry.kv_receiver.force_arm_calls, [True]) + + # Neither timeout nor a partial/stale ACK permits slot reuse. + mgr.note_abort_ack(7, 0, generation) + mgr.note_abort_ack(7, 1, generation - 1) + q.resolve_deferred_releases() + host_release.assert_not_called() + device_release.assert_not_called() + self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, []) + + mgr.note_abort_ack(7, 1, generation) + reduce.side_effect = lambda ready, **_: ready.zero_() + q.resolve_deferred_releases() + host_release.assert_not_called() + device_release.assert_not_called() + + # All decode ranks have drained: release host KV before the device + # KV/request row (which also owns the hybrid state), exactly once. + def release_device(*args, **kwargs): + host_release.assert_called_once_with(entry.req) + + device_release.side_effect = release_device + reduce.side_effect = None + q.resolve_deferred_releases() + device_release.assert_called_once_with( + entry.req, q.tree_cache, checkpoint=False + ) + self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [3]) + self.assertIsNone(entry.kv_receiver) + self.assertEqual(list(q._deferred_releases.items()), []) + q.resolve_deferred_releases() + host_release.assert_called_once_with(entry.req) + + # Reusing the room/metadata slot must not inherit the previous ACKs. + next_entry = _make_failed_req(mgr, room=7, idx=3, abort_notified=False) + q.queue.join([next_entry]) + next_generation = mgr.register_deferred_abort_room(7) + self.assertNotEqual(generation, next_generation) + q.pop_transferred() + mgr.note_abort_ack(7, 0, generation) + mgr.note_abort_ack(7, 1, generation) + q.resolve_deferred_releases() + host_release.assert_called_once_with(entry.req) + self.assertEqual(len(q._deferred_releases), 1) + device_release.side_effect = None + mgr.note_abort_ack(7, 0, next_generation) + mgr.note_abort_ack(7, 1, next_generation) + q.resolve_deferred_releases() + self.assertEqual( + host_release.call_args_list, [call(entry.req), call(next_entry.req)] + ) + self.assertEqual(device_release.call_count, 2) + self.assertEqual(q.req_to_metadata_buffer_idx_allocator.freed, [3, 3]) + self.assertEqual(list(q._deferred_releases.items()), []) + def test_unnotifiable_failure_releases_immediately(self): # Metadata never published: no prefill holds this destination, so no # write can be in flight and holding would waste a full timeout. diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 3a9640fed6ed..428762087a32 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -330,6 +330,7 @@ class TestNixlTransferWorker(CustomTestCase): def _make_manager(self, room): mgr = object.__new__(NixlKVManager) mgr.customized_info_enabled = False + mgr.local_shm_enabled = False mgr.request_status = {room: KVPoll.WaitingForInput} mgr.transfer_infos = { room: { @@ -832,6 +833,7 @@ def test_sender_queues_the_early_send_event_on_the_next_chunk_only(self): mgr = object.__new__(NixlKVManager) mgr.disaggregation_mode = DisaggregationMode.PREFILL mgr.enable_staging = False + mgr.local_shm_enabled = False mgr.transfer_infos = {room: {"agent": SimpleNamespace(dst_port=5555)}} queue = SimpleNamespace(put=MagicMock()) mgr.transfer_queues = [queue] @@ -1252,6 +1254,7 @@ class TestNixlInvalidatedPeerReload(CustomTestCase): def _make_manager(self, peers): mgr = object.__new__(NixlKVManager) mgr.customized_info_enabled = False + mgr.local_shm_enabled = False mgr.agent = ReloadFakeAgent() mgr.disaggregation_mode = DisaggregationMode.PREFILL mgr.is_mla_backend = False diff --git a/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py b/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py index 9e61c2174d01..eb19145d1755 100644 --- a/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py +++ b/test/registered/unit/disaggregation/test_nixl_deferred_kv_release.py @@ -38,6 +38,7 @@ def _manager( ) -> NixlKVManager: manager = NixlKVManager.__new__(NixlKVManager) manager.customized_info_enabled = False + manager.local_shm_enabled = False manager.enable_deferred_decode_kv_release = enabled manager._deferred_ack_targets = {} manager._deferred_ack_poisoned_rooms = set() diff --git a/test/registered/unit/disaggregation/test_nixl_local_shm.py b/test/registered/unit/disaggregation/test_nixl_local_shm.py new file mode 100644 index 000000000000..d5b54f644a0c --- /dev/null +++ b/test/registered/unit/disaggregation/test_nixl_local_shm.py @@ -0,0 +1,581 @@ +"""Completion and bounds contracts for same-node HiSparse writes.""" + +import os +import socket +import sys +import tempfile +import threading +import unittest +from types import SimpleNamespace +from unittest.mock import Mock, patch + +from sglang.srt.disaggregation.nixl.local_shm import ( + LocalShmExport, + LocalShmMapping, + LocalShmPeer, + LocalShmTransfer, + build_copy_plan, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestLocalShm(CustomTestCase): + def handle(self, error=None): + owner = SimpleNamespace(pending=set()) + event = Mock() + event.query.return_value = False + notify = Mock() + handle = LocalShmTransfer(owner, event, object(), notify, error) + owner.pending.add(handle) + return handle + + def test_completion_notifies_once(self): + handle = self.handle() + self.assertEqual(handle.poll(), "PROC") + handle.notify.assert_not_called() + self.assertIn(handle, handle.owner.pending) + handle.event.query.return_value = True + self.assertEqual(handle.poll(), "DONE") + self.assertEqual(handle.poll(), "DONE") + handle.notify.assert_called_once() + self.assertFalse(handle.owner.pending) + + def test_partial_submit_error_still_waits_for_dma(self): + handle = self.handle(RuntimeError("submission failed")) + self.assertEqual(handle.poll(), "PROC") + self.assertIn(handle, handle.owner.pending) + handle.event.query.return_value = True + self.assertEqual(handle.poll(), "ERR") + handle.notify.assert_not_called() + self.assertFalse(handle.owner.pending) + + def test_unknown_completion_retains_mapping(self): + handle = self.handle() + handle.event.query.side_effect = RuntimeError("device unavailable") + with self.assertRaises(RuntimeError): + handle.poll() + self.assertIn(handle, handle.owner.pending) + handle.notify.assert_not_called() + + def test_missing_event_retains_mapping(self): + handle = self.handle() + handle.event = None + with self.assertRaises(RuntimeError): + handle.poll() + self.assertIn(handle, handle.owner.pending) + + def test_notification_failure_is_settled_error(self): + handle = self.handle() + handle.event.query.return_value = True + handle.notify.side_effect = RuntimeError("peer gone") + self.assertEqual(handle.poll(), "ERR") + self.assertFalse(handle.owner.pending) + + def test_shuffled_pages_and_reused_slots(self): + self.assertEqual( + build_copy_plan([64], [128], [16], [3, 1], [0, 7]), + [(0, 48, 0, 16), (0, 16, 112, 16)], + ) + self.assertEqual( + build_copy_plan([64], [128], [16], [0], [7]), [(0, 0, 112, 16)] + ) + + def test_invalid_indices_and_geometry(self): + for src, dst in [([-1], [0]), ([4], [0]), ([0], [8]), ([0, 1], [0])]: + with self.subTest(src=src, dst=dst), self.assertRaises(ValueError): + build_copy_plan([64], [128], [16], src, dst) + with self.assertRaises(ValueError): + build_copy_plan([64, 64], [128], [16], [0], [0]) + with self.assertRaises(TypeError): + build_copy_plan([64], [128], [16], [0.5], [0]) + with self.assertRaises(ValueError): + build_copy_plan([64], [128], [0], [0], [0]) + + def test_barrier_waits_for_local_copy_and_nixl_state(self): + from sglang.srt.disaggregation.nixl.conn import NixlKVManager + + manager = NixlKVManager.__new__(NixlKVManager) + manager.agent = Mock() + manager.agent.check_xfer_state.side_effect = ["PROC", "DONE"] + handle = self.handle() + handle.event.query.side_effect = [False, True] + self.assertEqual( + manager._await_handles([handle, "state"], failure_seen=False), + (True, False), + ) + handle.notify.assert_called_once() + + def test_barrier_does_not_ack_unknown_local_completion(self): + from sglang.srt.disaggregation.nixl.conn import NixlKVManager + + manager = NixlKVManager.__new__(NixlKVManager) + handle = self.handle() + handle.event = None + self.assertEqual( + manager._await_handles([handle], failure_seen=True), (False, True) + ) + self.assertIn(handle, handle.owner.pending) + + def test_unenabled_send_keeps_existing_backend(self): + from sglang.srt.disaggregation.nixl.conn import NixlKVManager + + manager = NixlKVManager.__new__(NixlKVManager) + manager.local_shm_enabled = False + manager.src_mem_kind = "VRAM" + manager.kv_args = SimpleNamespace(kv_data_ptrs=[10], kv_item_lens=[16]) + manager._send_kvcache_generic = Mock(return_value="nixl") + self.assertEqual( + manager.send_kvcache("peer", [0], [20], [0], 0, "done"), "nixl" + ) + manager._send_kvcache_generic.assert_called_once() + + def test_failed_mapping_cannot_fall_back_to_nixl_write(self): + + manager, info = self.mapping_manager() + manager._get_local_shm_peer("peer").error = RuntimeError("mapping failed") + manager.decode_kv_args_table = {"peer": info} + manager._send_kvcache_generic = Mock() + with self.assertRaisesRegex(RuntimeError, "mapping failed"): + manager.send_kvcache("peer", [0], [20], [0], 0, "done") + manager._send_kvcache_generic.assert_not_called() + + def test_graceful_cleanup_releases_only_settled_mappings(self): + from sglang.srt.disaggregation.nixl.conn import NixlKVManager + + manager = NixlKVManager.__new__(NixlKVManager) + settled = Mock() + pending = Mock() + pending.close.side_effect = RuntimeError("DMA unresolved") + manager._local_shm_shutdown = threading.Event() + manager._local_shm_reaper = None + manager._local_shm_lock = threading.RLock() + manager._local_shm_peers = { + "settled": LocalShmPeer(mapping=settled), + "pending": LocalShmPeer(mapping=pending), + } + with self.assertLogs("sglang.srt.disaggregation.nixl.conn", level="ERROR"): + manager.release_host_resources() + settled.close.assert_called_once() + self.assertIsNone(manager._local_shm_peers["settled"].mapping) + self.assertIs(manager._local_shm_peers["pending"].mapping, pending) + + def test_sender_preserves_early_send_producer_event(self): + from sglang.srt.disaggregation.nixl.conn import NixlKVSender + + sender = NixlKVSender.__new__(NixlKVSender) + sender.conclude_state = None + sender.trace_ctx = Mock() + sender._prepare_send_indices = Mock( + return_value=([0], slice(0, 1), False, False) + ) + sender._transfer_start_time = None + sender.kv_mgr = SimpleNamespace( + local_shm_enabled=True, + add_transfer_request=Mock(), + prepare_send_state_indices=lambda indices: indices, + ) + sender.bootstrap_room = 1 + sender.chunk_id = 0 + sender.aux_index = 0 + sender._record_transfer_indices = Mock() + event = object() + sender._early_send_wait_event = event + sender.send([0]) + self.assertIs( + sender.kv_mgr.add_transfer_request.call_args.kwargs["wait_event"], event + ) + self.assertIsNone(sender._early_send_wait_event) + + def test_host_pool_closes_export_before_releasing_buffer(self): + from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost + + pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost) + pool.pin_memory = False + buffer = object() + pool.kv_buffer = buffer + observed = [] + pool.register_host_resource_cleanup(lambda: observed.append(pool.kv_buffer)) + other = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost) + self.assertIsNone(other._host_resource_cleanup) + pool.destroy() + pool.destroy() + self.assertEqual(observed, [buffer]) + self.assertIsNone(pool.kv_buffer) + self.assertFalse(pool._host_resource_cleanup) + + def mapping_manager(self): + from sglang.srt.disaggregation.nixl.conn import NixlKVManager + + manager = NixlKVManager.__new__(NixlKVManager) + manager.local_shm_enabled = True + manager._local_shm_peers = {} + manager._local_shm_lock = threading.Lock() + manager.is_mla_backend = True + manager.attn_tp_size = 4 + manager.kv_args = SimpleNamespace( + kv_data_ptrs=[10], kv_item_lens=[16], kv_layer_ids=[0], gpu_id=0 + ) + info = SimpleNamespace( + agent_name="peer", + local_shm={"items": [16]}, + decode_tp_size=4, + dst_dcp_size=1, + staging_base_ptr=0, + dst_kv_mem_kinds=["DRAM"], + dst_kv_item_lens=[16], + dst_kv_layer_ids=[0], + ) + return manager, info + + def test_transient_mapping_failure_recovers_after_backoff(self): + manager, info = self.mapping_manager() + module = "sglang.srt.disaggregation.nixl.conn" + mapping = object() + with ( + patch(module + ".same_host", return_value=True), + patch( + module + ".LocalShmMapping", + side_effect=[TimeoutError("FD timeout"), mapping], + ) as create, + patch(module + ".time.monotonic", return_value=10) as clock, + ): + with self.assertRaisesRegex(RuntimeError, "FD timeout"): + manager._ensure_local_shm_mapping(info) + with self.assertRaises(RuntimeError): + manager._ensure_local_shm_mapping(info) + self.assertEqual(create.call_count, 1) + clock.return_value = 12 + self.assertIs(manager._ensure_local_shm_mapping(info), mapping) + self.assertEqual(create.call_count, 2) + self.assertIsNone(manager._get_local_shm_peer("peer").error) + self.assertIsNone(manager._get_local_shm_peer("peer").retry_at) + self.assertIs(manager._ensure_local_shm_mapping(info), mapping) + self.assertEqual(create.call_count, 2) + + def test_layout_failure_is_not_retried(self): + manager, info = self.mapping_manager() + info.decode_tp_size = 8 + module = "sglang.srt.disaggregation.nixl.conn" + with ( + patch(module + ".same_host", return_value=True), + patch(module + ".LocalShmMapping") as create, + ): + for _ in range(2): + with self.assertRaisesRegex(RuntimeError, "matching flat MLA"): + manager._ensure_local_shm_mapping(info) + create.assert_not_called() + self.assertIsNone(manager._get_local_shm_peer("peer").retry_at) + + def test_dead_exporter_released_only_after_pending_copy_settles(self): + manager, info = self.mapping_manager() + mapping = Mock(pending={"copy"}) + mapping.exporter_closed.return_value = True + manager._get_local_shm_peer("peer").mapping = mapping + manager._reap_local_shm_mappings() + mapping.close.assert_not_called() + mapping.pending.clear() + manager._reap_local_shm_mappings() + mapping.close.assert_called_once() + self.assertIsNone(manager._get_local_shm_peer("peer").mapping) + # Late queued requests must not silently use NIXL or reopen the dead pool. + with self.assertRaisesRegex(RuntimeError, "exporter has closed"): + manager._ensure_local_shm_mapping(info) + + def test_live_exporter_and_failed_unregister_are_retained(self): + manager, _ = self.mapping_manager() + live = Mock(pending=set()) + live.exporter_closed.return_value = False + dead = Mock(pending=set()) + dead.exporter_closed.return_value = True + dead.close.side_effect = RuntimeError("still pinned") + manager._local_shm_peers = { + "live": LocalShmPeer(mapping=live), + "dead": LocalShmPeer(mapping=dead), + } + with self.assertLogs("sglang.srt.disaggregation.nixl.conn", level="ERROR"): + manager._reap_local_shm_mappings() + live.close.assert_not_called() + self.assertIs(manager._get_local_shm_peer("live").mapping, live) + self.assertIs(manager._get_local_shm_peer("dead").mapping, dead) + + def test_reaper_cannot_race_submission(self): + manager, _ = self.mapping_manager() + mapping = Mock(pending=set()) + mapping.exporter_closed.return_value = True + manager._get_local_shm_peer("peer").mapping = mapping + started = threading.Event() + + def reap(): + started.set() + manager._reap_local_shm_mappings() + + with manager._get_local_shm_peer("peer").lock: + worker = threading.Thread(target=reap) + worker.start() + self.assertTrue(started.wait(2)) + mapping.pending.add("submitted") + worker.join(2) + self.assertFalse(worker.is_alive()) + mapping.close.assert_not_called() + + def test_initializing_peer_does_not_block_healthy_peer(self): + manager, info = self.mapping_manager() + healthy = SimpleNamespace(**vars(info)) + healthy.agent_name = "healthy" + healthy_mapping = Mock() + manager._get_local_shm_peer("healthy").mapping = healthy_mapping + manager.decode_kv_args_table = {"healthy": healthy} + manager.kv_args.kv_data_lens = [64] + manager.agent = Mock() + entered, release, completed = ( + threading.Event(), + threading.Event(), + threading.Event(), + ) + errors = [] + + def create(*args): + entered.set() + if not release.wait(5): + raise TimeoutError("test initialization gate") + return Mock() + + def initialize(): + try: + manager._ensure_local_shm_mapping(info) + except Exception as exc: + errors.append(exc) + + def send(): + try: + manager.send_kvcache("healthy", [0], [20], [0], 0, "done") + except Exception as exc: + errors.append(exc) + finally: + completed.set() + + module = "sglang.srt.disaggregation.nixl.conn" + with ( + patch(module + ".same_host", return_value=True), + patch(module + ".LocalShmMapping", side_effect=create), + ): + initializing = threading.Thread(target=initialize) + sending = threading.Thread(target=send) + initializing.start() + try: + self.assertTrue(entered.wait(2)) + sending.start() + self.assertTrue( + completed.wait(2), + "healthy peer blocked by unrelated initialization", + ) + healthy_mapping.submit.assert_called_once() + finally: + release.set() + initializing.join(5) + if sending.ident is not None: + sending.join(5) + self.assertFalse(errors) + + def test_cleanup_callback_failure_can_be_retried(self): + from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost + + pool = MLATokenToKVPoolHost.__new__(MLATokenToKVPoolHost) + pool.pin_memory = False + buffer = object() + pool.kv_buffer = buffer + first, retry, last = ( + Mock(), + Mock(side_effect=[OSError("cleanup"), None]), + Mock(), + ) + for callback in (first, retry, last): + pool.register_host_resource_cleanup(callback) + with self.assertRaises(OSError): + pool.destroy() + self.assertIs(pool.kv_buffer, buffer) + self.assertFalse(getattr(pool, "_destroyed", False)) + last.assert_not_called() + pool.destroy() + pool.destroy() + first.assert_called_once() + self.assertEqual(retry.call_count, 2) + last.assert_called_once() + self.assertIsNone(pool.kv_buffer) + + def test_export_initialization_rolls_back_resources(self): + with tempfile.TemporaryFile() as file: + pool = SimpleNamespace( + fd=file.fileno(), + kv_buffer=Mock(), + get_contiguous_buf_infos=lambda: ([], [], []), + ) + pool.kv_buffer.is_contiguous.return_value = True + pool.kv_buffer.data_ptr.return_value = 1000 + pool.kv_buffer.numel.return_value = 4096 + pool.kv_buffer.element_size.return_value = 1 + module = "sglang.srt.disaggregation.nixl.local_shm" + duplicated = [] + real_dup = os.dup + + def duplicate_fd(fd): + result = real_dup(fd) + duplicated.append(result) + return result + + for stage in ("socket", "thread"): + with ( + self.subTest(stage=stage), + patch(module + ".os.dup", side_effect=duplicate_fd), + patch(module + ".socket.socket") as socket_factory, + patch(module + ".host_identity", return_value=["host", "net"]), + patch( + module + ".threading.Thread.start", + side_effect=RuntimeError("thread startup"), + ), + ): + if stage == "socket": + socket_factory.side_effect = OSError("socket allocation") + with self.assertRaises((OSError, RuntimeError)): + LocalShmExport(pool) + with self.assertRaises(OSError): + os.fstat(duplicated[-1]) + if stage == "thread": + socket_factory.return_value.close.assert_called_once() + + def test_export_close_retries_after_join_failure(self): + with tempfile.TemporaryFile() as file: + export = LocalShmExport.__new__(LocalShmExport) + export.fd = os.dup(file.fileno()) + fd = export.fd + export.stopped = threading.Event() + export.socket = Mock() + export.thread = Mock() + export.thread.join.side_effect = [RuntimeError("join"), None, None] + with self.assertRaises(RuntimeError): + export.close() + os.fstat(fd) + export.close() + export.close() + with self.assertRaises(OSError): + os.fstat(fd) + + def test_submit_failure_records_completion_for_already_queued_writes(self): + self.check_submit_failure(event_record_fails=False) + + def test_submit_event_failure_keeps_mapping_pending(self): + self.check_submit_failure(event_record_fails=True) + + def check_submit_failure(self, *, event_record_fails): + mapping = LocalShmMapping.__new__(LocalShmMapping) + mapping.ptrs, mapping.lengths, mapping.items = [100], [64], [16] + mapping.gpu_id, mapping.pending = 0, set() + event, stream, notify = Mock(), Mock(cuda_stream=1), Mock() + event.query.return_value = False + if event_record_fails: + event.record.side_effect = RuntimeError("record event") + runtime = SimpleNamespace( + cudaMemcpyAsync=Mock(side_effect=[(0,), (1,)]), + cudaMemcpyKind=SimpleNamespace(cudaMemcpyDeviceToHost=2), + ) + bindings = SimpleNamespace(runtime=runtime) + with ( + patch.dict( + sys.modules, + {"cuda": SimpleNamespace(bindings=bindings), "cuda.bindings": bindings}, + ), + patch("torch.cuda.device"), + patch("torch.cuda.Stream", return_value=stream), + patch("torch.cuda.Event", return_value=event), + ): + handle = mapping.submit([200], [64], [0, 1], [0, 1], notify) + self.assertEqual(runtime.cudaMemcpyAsync.call_count, 2) + event.record.assert_called_once_with(stream) + self.assertIn(handle, mapping.pending) + if event_record_fails: + with self.assertRaisesRegex(RuntimeError, "could not be recorded"): + handle.poll() + self.assertIn(handle, mapping.pending) + else: + self.assertEqual(handle.poll(), "PROC") + event.query.return_value = True + self.assertEqual(handle.poll(), "ERR") + self.assertFalse(mapping.pending) + notify.assert_not_called() + + def test_empty_chunk(self): + self.assertEqual(build_copy_plan([64], [128], [16], [], []), []) + + +@unittest.skipUnless(sys.platform == "linux", "FD exchange requires Linux") +class TestLocalShmFDExchange(CustomTestCase): + def setUp(self): + self.fd = os.memfd_create("sglang-shm-test", os.MFD_CLOEXEC) + os.ftruncate(self.fd, 4096) + os.pwrite(self.fd, b"test", 0) + buffer = Mock() + buffer.is_contiguous.return_value = True + buffer.data_ptr.return_value = 1000 + buffer.numel.return_value = 4096 + buffer.element_size.return_value = 1 + self.pool = SimpleNamespace( + fd=self.fd, + kv_buffer=buffer, + get_contiguous_buf_infos=lambda: ([1000], [4096], [256]), + ) + self.export = LocalShmExport(self.pool) + + def tearDown(self): + self.export.close() + os.close(self.fd) + + def receive(self, token): + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn: + conn.settimeout(2) + conn.connect("\0" + self.export.descriptor["endpoint"]) + conn.sendall(token) + return socket.recv_fds(conn, 1, 1, socket.MSG_CMSG_CLOEXEC) + + def test_wrong_generation_rejected_then_valid_fd_works(self): + message, fds, _, _ = self.receive(b"x" * 32) + self.assertEqual((message, fds), (b"", [])) + message, fds, _, _ = self.receive(self.export.token) + self.assertEqual(message, b"F") + self.assertEqual(len(fds), 1) + try: + self.assertEqual(os.pread(fds[0], 4, 0), b"test") + finally: + os.close(fds[0]) + + def test_exporter_liveness_tracks_endpoint_shutdown(self): + mapping = LocalShmMapping.__new__(LocalShmMapping) + mapping.endpoint = self.export.descriptor["endpoint"] + self.assertFalse(mapping.exporter_closed()) + self.export.close() + self.assertTrue(mapping.exporter_closed()) + + def test_probe_timeout_does_not_establish_peer_death(self): + mapping = LocalShmMapping.__new__(LocalShmMapping) + mapping.endpoint = self.export.descriptor["endpoint"] + with patch("socket.socket") as sock: + sock.return_value.__enter__.return_value.connect.side_effect = ( + TimeoutError() + ) + self.assertFalse(mapping.exporter_closed()) + + def test_pool_restart_has_new_identity(self): + old = self.export.descriptor + self.export.close() + self.export = LocalShmExport(self.pool) + self.assertNotEqual(old["endpoint"], self.export.descriptor["endpoint"]) + self.assertNotEqual(old["token"], self.export.descriptor["token"]) + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as conn: + with self.assertRaises(ConnectionRefusedError): + conn.connect("\0" + old["endpoint"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/disaggregation/test_receiver_connection_pool.py b/test/registered/unit/disaggregation/test_receiver_connection_pool.py index e576f98b2265..0df50248b68e 100644 --- a/test/registered/unit/disaggregation/test_receiver_connection_pool.py +++ b/test/registered/unit/disaggregation/test_receiver_connection_pool.py @@ -61,6 +61,7 @@ def _manager(manager_class=CommonKVManager): manager._deferred_ack_poisoned_rooms = set() manager.enable_staging = False manager.customized_info_enabled = False + manager.local_shm_descriptor = None return manager @@ -315,6 +316,8 @@ def test_registration_wire_is_preserved_for_all_peers(self): expected_frames.index(aux) + 1, struct.pack("Q", 64) if enabled else b"", ) + if backend == "nixl": + expected_frames.append(b"") self.assertTrue(receiver._register_kv_args()) self.assertEqual( sent, [(port, expected_frames) for port in (6000, 6001, 6002)] @@ -333,6 +336,18 @@ def test_registration_wire_is_preserved_for_all_peers(self): self.assertEqual(info.dst_state_item_lens, [[32]]) if backend == "nixl": self.assertEqual(info.dst_state_data_lens, [[512]]) + self.assertIsNone(info.local_shm) + self.assertIsNone(registration_type.from_zmq(frames[:-1]).local_shm) + descriptor = {"version": 1, "endpoint": "test-pool"} + receiver.kv_mgr.local_shm_descriptor = descriptor + registered = registration_type.from_zmq( + receiver.kv_mgr.build_receiver_registration_message()[1:] + ) + self.assertEqual(registered.local_shm, descriptor) + self.assertEqual( + registered.dst_aux_item_lens, info.dst_aux_item_lens + ) + self.assertEqual(registered.dst_state_data_lens, [[512]]) def test_metadata_wire_is_preserved_for_all_peers(self): for backend in self.backends: diff --git a/test/registered/unit/managers/scheduler_components/test_invariant_checker.py b/test/registered/unit/managers/scheduler_components/test_invariant_checker.py index c5cc22014622..ec4aad930e9f 100644 --- a/test/registered/unit/managers/scheduler_components/test_invariant_checker.py +++ b/test/registered/unit/managers/scheduler_components/test_invariant_checker.py @@ -39,6 +39,55 @@ register_cpu_ci(est_time=11, suite="base-a-test-cpu") +class TestHiSparseLogicalOccupancy(CustomTestCase): + def test_logical_free_count_is_not_capped_by_hot_pool(self): + from sglang.srt.mem_cache.allocator.hisparse import ( + HiSparseTokenToKVPoolAllocator, + ) + + allocator = MagicMock(spec=HiSparseTokenToKVPoolAllocator) + allocator.size = 2048 + allocator.available_size.return_value = 768 + allocator.logical_attn_allocator = SimpleNamespace( + size=2048, available_size=lambda: 1792 + ) + pool = SimpleNamespace( + schedulable_token_capacity=lambda size: size, + mamba_allocator=SimpleNamespace(available_size=lambda: 7), + mamba_pool=SimpleNamespace(size=8), + ) + tree = MagicMock() + tree.supports_mamba.return_value = False + tree.evictable_size.return_value = 0 + allocator.state.ckpt_pool = None + allocator.state.pool.size = 8 + allocator.state_slots.available_size.return_value = 7 + for has_state in (False, True): + storage = StaticKVStorage( + pool, + allocator, + shape=StaticPoolShape( + has_window=False, + has_state=has_state, + token_capacity=1024, + ), + ) + observer = SchedulerPoolStatsObserver( + tree_cache=tree, + storage=storage, + token_to_kv_pool_allocator=allocator, + session_controller=None, + hisparse_coordinator=None, + enable_hisparse=True, + ) + stats = observer.get_pool_stats() + self.assertEqual(stats.full_capacity, 2048) + self.assertEqual(stats.full_available_size, 1792) + self.assertEqual(stats.full_num_used, 256) + self.assertEqual(stats.full_token_usage, 0.125) + allocator.available_size.assert_not_called() + + class TestCheckTreeCacheGate(CustomTestCase): @contextmanager def _without_explicit_sanity_check_setting(self): diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index 0d7e65e7056b..dc9c4abc13f3 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -20,6 +20,74 @@ register_cpu_ci(est_time=10, suite="base-a-test-cpu") +class TestHybridHiSparseBackingPool(CustomTestCase): + def test_mapping_uses_dsa_pool_and_preserves_hybrid_wrapper(self): + from sglang.srt.mem_cache.allocator.hisparse import ( + HiSparseTokenToKVPoolAllocator, + ) + from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool + + backing = object.__new__(HiSparseDSATokenToKVPool) + backing.register_mapping = MagicMock() + wrapper = object.__new__(HybridLinearKVPool) + wrapper.full_kv_pool = backing + wrapper.mamba_pool = object() + allocator = HiSparseTokenToKVPoolAllocator( + size=16, + page_size=1, + dtype=torch.bfloat16, + device="cpu", + kvcache=wrapper, + need_sort=False, + host_to_device_ratio=2, + ) + self.assertIs(allocator.get_kvcache(), wrapper) + self.assertIs(allocator.hisparse_kvcache, backing) + backing.register_mapping.assert_called_once() + self.assertEqual( + backing.register_mapping.call_args.args[0].data_ptr(), + allocator.full_to_hisparse_device_index_mapping.data_ptr(), + ) + self.assertEqual(allocator.logical_attn_allocator.size, 32) + self.assertEqual(allocator.hisparse_attn_allocator.size, 16) + indices = torch.tensor([1, 2]) + backing.translate_loc_to_hisparse_device = MagicMock(return_value=indices + 4) + torch.testing.assert_close( + wrapper.translate_loc_to_hisparse_device(indices), indices + 4 + ) + + def test_rejects_dense_hybrid_backing_pool(self): + from sglang.srt.mem_cache.allocator.hisparse import ( + HiSparseTokenToKVPoolAllocator, + ) + from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool + + wrapper = object.__new__(HybridLinearKVPool) + wrapper.full_kv_pool = object() + with self.assertRaisesRegex(TypeError, "HiSparse DSA backing pool"): + HiSparseTokenToKVPoolAllocator( + size=16, + page_size=1, + dtype=torch.bfloat16, + device="cpu", + kvcache=wrapper, + need_sort=False, + ) + + def test_swap_maps_global_layer_to_backing_layer(self): + from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator + + coordinator = object.__new__(HiSparseCoordinator) + coordinator.full_attention_layer_id_mapping = {3: 0, 7: 1} + coordinator.enable_prefetch = False + coordinator._run_swap_in_kernel = MagicMock(return_value="locations") + self.assertEqual( + coordinator.swap_in_selected_pages(None, None, None, 7), "locations" + ) + coordinator._run_swap_in_kernel.assert_called_once_with(None, None, None, 1) + + class TestHiSparseDecodeRemap(CustomTestCase): def test_spec_scratch_is_charged_only_before_request_slot_allocation(self): allocator = object.__new__(HiSparseTokenToKVPoolAllocator) diff --git a/test/registered/unit/model_executor/test_hisparse_pool_configurator.py b/test/registered/unit/model_executor/test_hisparse_pool_configurator.py index a870eeb2392c..287094f2b607 100644 --- a/test/registered/unit/model_executor/test_hisparse_pool_configurator.py +++ b/test/registered/unit/model_executor/test_hisparse_pool_configurator.py @@ -1,6 +1,6 @@ import unittest from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import torch @@ -13,6 +13,71 @@ class TestHiSparsePoolConfigurator(CustomTestCase): + def test_hybrid_dsa_pool_uses_dense_layer_ids(self): + from sglang.srt.mem_cache import kv_cache_configurator as module + + kvc = SimpleNamespace() + kvc.layer_info = SimpleNamespace( + start_layer=4, end_layer=12, num_effective_layers=8 + ) + kvc.model_config = SimpleNamespace( + kv_lora_rank=512, qk_rope_head_dim=0, hf_config=SimpleNamespace() + ) + kvc.loc_space_scale = 1 + kvc.kv_shard_rank = None + kvc.pool_page_size = 64 + kvc.kv_cache_dtype = torch.bfloat16 + kvc.device = "cpu" + kvc.is_draft_worker = False + with ( + patch.object( + module, "get_memory", return_value=SimpleNamespace(enable_hisparse=True) + ), + patch.object( + module, + "get_exec", + return_value=SimpleNamespace( + features=SimpleNamespace(enable_memory_saver=False) + ), + ), + patch.object(module, "reject_out_of_tree_path"), + patch( + "sglang.srt.layers.cp.utils.get_glm_dsa_cp_layer_shard_info", + return_value=(None, None), + ), + patch( + "sglang.srt.mem_cache.sparsity.parse_hisparse_config", + return_value=SimpleNamespace(host_to_device_ratio=2), + ), + patch.object(module, "_should_elide_dsa_index_k", return_value=True), + patch.object( + module, + "dsa_layer_skips_topk", + side_effect=lambda config, layer: layer == 11, + ) as skips, + patch.object(module, "calculate_mla_kv_cache_dim", return_value=512), + patch.object(module, "get_dsa_index_head_dim", return_value=128), + patch.object(module, "get_dsa_index_kpool", return_value=4), + patch.object(module, "get_dsa_index_kpool_compress", return_value=False), + patch.object(module, "max_speculative_num_draft_tokens", return_value=None), + patch.object(module, "HiSparseDSATokenToKVPool") as pool, + ): + result = module.KVCacheConfigurator._build_dsa_kv_pool( + kvc, + max_total_num_tokens=1024, + max_running_requests=8, + dsa_pool_class=object, + full_attention_layer_ids=[7, 11], + ) + self.assertIs(result, pool.return_value) + kwargs = pool.call_args.kwargs + self.assertEqual( + (kwargs["start_layer"], kwargs["end_layer"], kwargs["layer_num"]), (0, 2, 2) + ) + self.assertEqual(kwargs["skip_topk_layers"], [False, True]) + self.assertEqual([call.args[1] for call in skips.call_args_list], [7, 11]) + self.assertEqual(kwargs["host_to_device_ratio"], 2) + def _compute_cell_size( self, kv_cache_dtype: torch.dtype, diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 4b1cf29a5b49..2c2d83a9802c 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -2754,6 +2754,29 @@ def _view(arch="DeepseekV32ForCausalLM", learnable_sink=False, **kw): "dsa_decode_backend": "flashmla_kv", }, ) + self.assertEqual( + _dsa_split_backend_resolution( + _view( + enable_hisparse=True, + dsa_prefill_backend="tilelang", + dsa_decode_backend="tilelang", + ) + ), + {}, + ) + with self.assertRaisesRegex(ValueError, "both DSA backends"): + _dsa_split_backend_resolution( + _view(enable_hisparse=True, dsa_decode_backend="tilelang") + ) + with self.assertRaisesRegex(ValueError, "dcp-size"): + _dsa_split_backend_resolution( + _view( + enable_hisparse=True, + dsa_prefill_backend="tilelang", + dsa_decode_backend="tilelang", + dcp_size=2, + ) + ) for arch in ("HYV4ForCausalLM", "HYV4ForCausalLMNextN"): with self.subTest(arch=arch, backends="default"): self.assertEqual(