diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 8dc32839a..7266c0271 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -31,7 +31,7 @@ class KVArgs: state_data_ptrs: List[int] state_data_lens: List[int] state_item_lens: List[int] - state_type: str # "none", "mamba", "swa", "nsa", "dsv4" + state_type: str # "none", "mamba", "swa", "nsa" # for mamba state different tp slice transfer state_dim_per_tensor: List[int] # dimension to slice for each state tensor ib_device: str diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 5cd9558e2..0e7e7aac1 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -995,13 +995,8 @@ class MooncakeKVManager(CommonKVManager): prefill_state_indices, dst_state_data_ptrs, ) - elif state_type in ["swa", "nsa", "dsv4"]: - # SWA / NSA / DSv4 hybrid models do not support different TP sizes - # yet. (DSv4 carries a flat heterogeneous state pool of - # SWA + compress + indexer buffers; reusing this branch routes it - # through the same ``_send_kvcache_generic`` path that - # ``get_mla_kv_ptrs_with_pp`` already handles for compressed-MLA - # PP/MTP layouts.) + elif state_type in ["swa", "nsa"]: + # Non-MLA SWA / NSA hybrid models do not support different TP sizes yet. if ( target_rank_registration_info is not None and not self.is_mla_backend diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 97520499c..cb21fd5c9 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -1010,62 +1010,6 @@ class NixlKVManager(CommonKVManager): raise Exception("Failed to post Mamba state slice transfer") return xfer_handle - def _send_state_pages_flat( - self, - peer_name: str, - prefill_state_indices: List[int], - dst_state_data_ptrs: list[int], - dst_state_indices: List[int], - dst_state_item_lens: list[int], - dst_gpu_id: int, - notif: str, - ): - """Per-page WRITE transfer of a flat (heterogeneous) state pool. - - Used by V4 whose state pool is a flat list of buffers (SWA + compress - + indexer pools) that does not match the per-layer K/V layout assumed - by ``_send_kvcache_generic``. Both sides must have identical - ``state_item_lens`` (no TP-slicing path). - """ - src_state_ptrs = self.kv_args.state_data_ptrs - src_state_item_lens = self.kv_args.state_item_lens - assert len(src_state_ptrs) == len(dst_state_data_ptrs) - assert len(src_state_item_lens) == len(dst_state_item_lens) - assert len(prefill_state_indices) == len(dst_state_indices), ( - f"State index length mismatch: prefill={len(prefill_state_indices)}, " - f"dst={len(dst_state_indices)}" - ) - for i in range(len(src_state_item_lens)): - assert src_state_item_lens[i] == dst_state_item_lens[i], ( - f"V4 state item length mismatch at index {i}: " - f"{src_state_item_lens[i]} != {dst_state_item_lens[i]}" - ) - - src_addrs = [] - dst_addrs = [] - for i in range(len(src_state_ptrs)): - item_len = src_state_item_lens[i] - for src_idx, dst_idx in zip(prefill_state_indices, dst_state_indices): - src_addr = src_state_ptrs[i] + int(src_idx) * item_len - dst_addr = dst_state_data_ptrs[i] + int(dst_idx) * item_len - src_addrs.append((src_addr, item_len, self.kv_args.gpu_id)) - dst_addrs.append((dst_addr, item_len, dst_gpu_id)) - - if not src_addrs: - return None - - src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM") - dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM") - xfer_handle = self.agent.initialize_xfer( - "WRITE", src_descs, dst_descs, peer_name, notif.encode("ascii") - ) - if not xfer_handle: - raise Exception("KVSender failed to create state transfer") - state = self.agent.transfer(xfer_handle) - if state == "ERR": - raise Exception("KVSender failed to post state transfer") - return xfer_handle - def maybe_send_extra( self, peer_name: str, @@ -1124,16 +1068,6 @@ class NixlKVManager(CommonKVManager): dst_gpu_id=dst_gpu_id, notif=notif, ) - elif state_type == "dsv4": - return self._send_state_pages_flat( - peer_name, - prefill_state_indices, - dst_state_data_ptrs, - dst_state_indices, - dst_state_item_lens or [], - dst_gpu_id, - notif, - ) else: if state_type != "none": raise RuntimeError( diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 37cb47422..ea5e37999 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -543,7 +543,6 @@ def setup_state_kv_args( lives in one place. """ from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool - from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, NSATokenToKVPool if not hasattr(token_to_kv_pool, "get_state_buf_infos"): @@ -560,12 +559,9 @@ def setup_state_kv_args( kv_args.state_data_lens = state_data_lens kv_args.state_item_lens = state_item_lens - # V4 must be checked before BaseSWAKVPool: V4's state pool is a flat - # heterogeneous list (SWA + compress + indexer), so the per-layer K/V - # transfer path used for "swa"/"nsa" does not apply. - if isinstance(token_to_kv_pool, DeepSeekV4TokenToKVPool): - kv_args.state_type = "dsv4" - elif isinstance(token_to_kv_pool, BaseSWAKVPool): + # DeepSeekV4TokenToKVPool inherits BaseSWAKVPool; its heterogeneous + # state list is described per-entry via get_state_buf_infos. + if isinstance(token_to_kv_pool, BaseSWAKVPool): kv_args.state_type = "swa" elif isinstance(token_to_kv_pool, HybridLinearKVPool): kv_args.state_type = "mamba"