minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[minimax-m3] Split 3/4: disagg K-only index-K transfer (#28714)

Browse Source
This commit is contained in:
Xinyuan Tong
2026-06-28 13:54:42 +08:00
committed by GitHub
parent 6eedc8f376
commit ddc389cf09
7 changed files with 173 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/base/conn.py
+1
View File
@@ -18,6 +18,7 @@ class StateType(str, enum.Enum):
MAMBA = "mamba"
SWA = "swa"
DSA = "dsa"
MINIMAX_INDEX_K = "minimax_index_k"
# DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot
# (req_pool_idx * ring_stride + pos % ring_stride), needs its own component.
SWA_RING = "swa_ring"
python/sglang/srt/disaggregation/decode.py
+4
View File
@@ -1058,6 +1058,10 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
state_indices.append(_swa_payload())
elif st == StateType.DSA:
state_indices.append(_dsa_payload())
elif st == StateType.MINIMAX_INDEX_K:
# Index rows live at the same loc as main KV on the same
# page_size, so reuse the full-seq page-ids.
state_indices.append(_dsa_payload())
elif st == StateType.SWA_RING:
state_indices.append(_swa_ring_payload())
else:
python/sglang/srt/disaggregation/mooncake/conn.py
+41 -1
View File
@@ -586,10 +586,15 @@ class MooncakeKVManager(CommonKVManager):
prefill_data_indices: npt.NDArray[np.int32],
dst_data_indices: npt.NDArray[np.int32],
executor: concurrent.futures.ThreadPoolExecutor,
force_flat: bool = False,
) -> int:
"""
Generic KV cache transfer supporting both MHA and MLA architectures.
This method is used by both send_kvcache (full pool) and maybe_send_extra.
``force_flat`` uses the MLA-style flat (single-buffer-per-layer) layout
even on a non-MLA backend, for K-only state buffers (e.g. MiniMax sparse
index) whose per-layer list must not be half-split into K/V.
"""
# Group by indices for optimization
prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
@@ -599,7 +604,7 @@ class MooncakeKVManager(CommonKVManager):
layers_params = None
# Decode pp size should be equal to prefill pp size or 1
if self.is_mla_backend:
if self.is_mla_backend or force_flat:
src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
)
@@ -1036,6 +1041,41 @@ class MooncakeKVManager(CommonKVManager):
)
or rc
)
elif st == StateType.MINIMAX_INDEX_K:
# Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
# lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
if self.pp_size is not None and self.pp_size > 1:
raise RuntimeError(
"PD disagg: PP>1 not supported for MiniMax sparse index yet."
)
if (
target_rank_registration_info is not None
and self.attn_tp_size
!= target_rank_registration_info.dst_attn_tp_size
):
raise RuntimeError(
"PD disagg: heterogeneous TP not supported for MiniMax "
"sparse index yet."
)
src_indices = list(indices)
dst_indices_local = list(dst_indices)
if len(src_indices) > len(dst_indices_local):
src_indices = src_indices[: len(dst_indices_local)]
elif len(src_indices) < len(dst_indices_local):
dst_indices_local = dst_indices_local[: len(src_indices)]
rc = (
self._send_kvcache_generic(
mooncake_session_id=req.mooncake_session_id,
src_data_ptrs=src_data_ptrs,
dst_data_ptrs=dst_data_ptrs,
item_lens=src_item_lens,
prefill_data_indices=np.array(src_indices, dtype=np.int32),
dst_data_indices=np.array(dst_indices_local, dtype=np.int32),
executor=executor,
force_flat=True,
)
or rc
)
return rc
def _send_mamba_state(
python/sglang/srt/disaggregation/nixl/conn.py
+34 -2
View File
@@ -1279,10 +1279,14 @@ class NixlKVManager(CommonKVManager):
notif: str,
src_mem_kind: str = "VRAM",
dst_mem_kind: str = "VRAM",
force_flat: bool = False,
):
"""Generic KV cache transfer supporting both MHA and MLA architectures.
Used by both send_kvcache and maybe_send_extra."""
Used by both send_kvcache and maybe_send_extra.
``force_flat`` uses the MLA-style flat (single-buffer-per-layer) layout
even on a non-MLA backend, for K-only state buffers (e.g. MiniMax sparse
index) whose per-layer list must not be half-split into K/V."""
# Prepped path (KV only; state transfers use the non-prepped path below).
if (
src_data_ptrs is self.kv_args.kv_data_ptrs
@@ -1335,7 +1339,7 @@ class NixlKVManager(CommonKVManager):
logger.debug(f"sending kvcache to {peer_name} with notif {notif}")
# Make descs
if self.is_mla_backend:
if self.is_mla_backend or force_flat:
src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = (
self.get_mla_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs)
)
@@ -2010,6 +2014,34 @@ class NixlKVManager(CommonKVManager):
dst_gpu_id=dst_gpu_id,
notif=comp_notif,
)
elif st == StateType.MINIMAX_INDEX_K:
# Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
# lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
if self.pp_size is not None and self.pp_size > 1:
raise RuntimeError(
"PD disagg: PP>1 not supported for MiniMax sparse index yet."
)
if self.attn_tp_size != decode_tp_size:
raise RuntimeError(
"PD disagg: heterogeneous TP not supported for MiniMax "
"sparse index yet."
)
if len(src_indices) != len(dst_indices):
raise RuntimeError(
f"State index length mismatch at component {i}: "
f"prefill={len(src_indices)}, dst={len(dst_indices)}"
)
h = self._send_kvcache_generic(
peer_name=peer_name,
src_data_ptrs=src_ptrs,
dst_data_ptrs=dst_ptrs,
item_lens=src_lens,
prefill_data_indices=np.array(src_indices, dtype=np.int32),
dst_data_indices=np.array(dst_indices, dtype=np.int32),
dst_gpu_id=dst_gpu_id,
notif=comp_notif,
force_flat=True,
)
else:
raise RuntimeError(
f"PD Disaggregation via NIXL does NOT support {st} hybrid models yet."
python/sglang/srt/disaggregation/prefill.py
+4
View File
@@ -1070,6 +1070,10 @@ class SchedulerDisaggregationPrefillMixin:
state_indices.append(_swa_payload())
elif st == StateType.DSA:
state_indices.append(_dsa_payload())
elif st == StateType.MINIMAX_INDEX_K:
# Index rows live at the same loc as main KV on the same
# page_size, so reuse the full-seq page-ids.
state_indices.append(_dsa_payload())
elif st == StateType.SWA_RING:
state_indices.append(_swa_ring_payload())
else:
python/sglang/srt/disaggregation/utils.py
+15 -2
View File
@@ -643,7 +643,11 @@ def setup_state_kv_args(
from sglang.srt.disaggregation.base.conn import StateType
from sglang.srt.hardware_backend.npu.memory_pool_npu import NPUMLATokenToKVPool
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, HybridLinearKVPool
from sglang.srt.mem_cache.memory_pool import (
DSATokenToKVPool,
HybridLinearKVPool,
MiniMaxSparseKVPool,
)
kv_args.state_types = []
kv_args.state_data_ptrs = []
@@ -651,7 +655,16 @@ def setup_state_kv_args(
kv_args.state_item_lens = []
kv_args.state_dim_per_tensor = []
if hasattr(token_to_kv_pool, "get_state_buf_infos"):
if isinstance(token_to_kv_pool, MiniMaxSparseKVPool):
if token_to_kv_pool.index_kv_pool is not None:
raise NotImplementedError(
"PD disaggregation for MiniMax sparse layers with index value "
"(index_kv_pool) is not yet supported; only K-only sparse layers are."
)
if token_to_kv_pool.index_k_pool is not None:
dp, dl, il = token_to_kv_pool.get_index_k_state_buf_infos()
append_state_component(kv_args, StateType.MINIMAX_INDEX_K, dp, dl, il)
elif hasattr(token_to_kv_pool, "get_state_buf_infos"):
data_ptrs, data_lens, item_lens = token_to_kv_pool.get_state_buf_infos()
# DeepSeekV4TokenToKVPool inherits BaseSWAKVPool; its heterogeneous