diff --git a/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py b/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py index ae8e2af72..8cac79f4f 100644 --- a/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py +++ b/python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py @@ -28,8 +28,16 @@ def _qwen4_exp_overrides(server_args: Any, hf_config: Any) -> dict: which MambaRadixCache allows only with mamba extra-buffer or --disable-radix-cache. """ cfg = resolving_view(server_args) - if cfg.disaggregation_mode != "null": - raise ValueError("Qwen4-Exp does not support PD disaggregation yet") + if ( + cfg.disaggregation_mode != "null" + and cfg.disaggregation_transfer_backend == "mori" + and cfg.pp_size > 1 + ): + raise ValueError( + "Qwen4-Exp PD with MORI requires --pp-size 1; MORI does not yet " + "exchange the global QSA layer metadata needed to pair compact " + "state descriptors across pipeline stages." + ) if cfg.enable_unified_memory: raise ValueError("Qwen4-Exp does not support --enable-unified-memory yet") overrides: Dict[str, Any] = {} diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 2e3f5dfe2..239504e88 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -16,6 +16,8 @@ if TYPE_CHECKING: class StateType(str, enum.Enum): MAMBA = "mamba" + QSA_PENDING = "qsa_pending" + QSA_COMPRESSED = "qsa_compressed" SWA = "swa" DSA = "dsa" # DSA kpool-compress tail: one per-request ring row. The indices encode diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index ebb553272..82ce1ff5c 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -57,6 +57,7 @@ from sglang.srt.disaggregation.utils import ( get_dsa_tail_state_indices, get_dsv4_c128_state_indices, get_kv_class, + get_qsa_pending_state_indices, is_dsv4_c128_online_enabled, is_mla_backend, poll_and_all_reduce, @@ -251,6 +252,10 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): linear_replayssm_cache_len: int = 16, mamba_envelope_layout: bool = False, enable_linear_replayssm_spec: bool = False, + short_conv_layer_ids: Optional[List[int]] = None, + short_conv_state_shape: Optional[Tuple[int, int]] = None, + ngram_context_len: int = 0, + ngram_eos_token_id: int = 0, ): DecodeReqToTokenPool.__init__( self, @@ -298,6 +303,10 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): linear_replayssm_cache_len=linear_replayssm_cache_len, mamba_envelope_layout=mamba_envelope_layout, enable_linear_replayssm_spec=enable_linear_replayssm_spec, + short_conv_layer_ids=short_conv_layer_ids, + short_conv_state_shape=short_conv_state_shape, + ngram_context_len=ngram_context_len, + ngram_eos_token_id=ngram_eos_token_id, ) def clear(self): @@ -1428,6 +1437,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): seq_len, ) + def _qsa_pending_payload(): + # Match the prefill request-pool row positionally; the two + # req_pool_idx values need not be equal. + return get_qsa_pending_state_indices(decode_req.req) + def _swa_ring_payload(): # Mirror of prefill _swa_ring_payload using this side's req_pool_idx. # Same window positions and order -> positional match with prefill. @@ -1458,6 +1472,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): clear_c128_state(int(decode_req.req.kv.req_pool_idx)) payloads = { StateType.MAMBA: _mamba_payload, + StateType.QSA_PENDING: _qsa_pending_payload, + StateType.QSA_COMPRESSED: _full_kv_pages_payload, StateType.SWA: _swa_payload, StateType.DSA: _full_kv_pages_payload, StateType.DSA_TAIL: _dsa_tail_payload, diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 98a1c289a..5c75146f7 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -51,6 +51,7 @@ from sglang.srt.disaggregation.utils import ( build_transfer_entry_pairs, compute_mamba_state_slice_byte_blocks, resolve_dcp_dst_entry_indices, + should_send_replicated_state, slice_dsa_tail_dst_ptrs_for_pp, ) from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine @@ -1361,6 +1362,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): return st in ( StateType.SWA, StateType.DSA, + StateType.QSA_PENDING, + StateType.QSA_COMPRESSED, StateType.SWA_RING, StateType.DSV4_REQUEST_STATE, StateType.BLOCK_SCALE, @@ -1369,7 +1372,12 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): def _requires_exact_state_index_match(self, st: StateType) -> bool: """State types whose page lists are positional and must not be truncated.""" - return st in (StateType.SWA_RING, StateType.DSV4_REQUEST_STATE) + return st in ( + StateType.QSA_PENDING, + StateType.QSA_COMPRESSED, + StateType.SWA_RING, + StateType.DSV4_REQUEST_STATE, + ) def maybe_send_extra( self, @@ -1501,16 +1509,60 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): or rc ) elif self._is_generic_kvcache_state_type(st): - if ( + is_qwen4_qsa_state = st in ( + StateType.QSA_PENDING, + StateType.QSA_COMPRESSED, + ) + has_heterogeneous_attn_tp = ( target_rank_registration_info is not None - and not self.is_mla_backend - and not self.is_hybrid_mla_backend and self.attn_tp_size != target_rank_registration_info.dst_attn_tp_size + ) + if ( + has_heterogeneous_attn_tp + and not self.is_mla_backend + and not self.is_hybrid_mla_backend + and not is_qwen4_qsa_state ): raise RuntimeError( f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {st.upper()} hybrid models yet." ) + if has_heterogeneous_attn_tp and is_qwen4_qsa_state: + if len(dst_item_lens) != len(dst_data_ptrs): + raise RuntimeError( + f"Replicated {st.upper()} destination pointer/item-length " + "metadata is inconsistent: " + f"dst ptrs={len(dst_data_ptrs)} lens={len(dst_item_lens)}" + ) + qsa_entry_pairs = build_transfer_entry_pairs( + src_state_layer_ids, + dst_state_layer_ids, + len(src_data_ptrs), + len(dst_data_ptrs), + allow_positional_fallback=self.pp_size == 1, + ) + layout_mismatches = [ + (i, j, src_item_lens[i], dst_item_lens[j]) + for i, j in qsa_entry_pairs + if src_item_lens[i] != dst_item_lens[j] + ] + if layout_mismatches: + raise RuntimeError( + f"Replicated {st.upper()} layout differs between mapped " + "prefill and decode entries: " + f"{layout_mismatches}" + ) + local_tp_rank_in_group = ( + self.kv_args.engine_rank % self.attn_tp_size + ) + if not should_send_replicated_state( + src_attn_tp_size=self.attn_tp_size, + dst_attn_tp_size=( + target_rank_registration_info.dst_attn_tp_size + ), + local_tp_rank_in_group=local_tp_rank_in_group, + ): + continue src_indices = list(indices) dst_indices_local = list(dst_indices) if ( @@ -1546,6 +1598,10 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): dst_data_indices=np.array(dst_indices_local, dtype=np.int32), executor=executor, state_type=st, + force_flat=st + in (StateType.QSA_PENDING, StateType.QSA_COMPRESSED), + src_layer_ids=src_state_layer_ids, + dst_layer_ids=dst_state_layer_ids, ) or rc ) @@ -1685,8 +1741,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): compute_mamba_state_slice_byte_blocks). """ logger.warning_once( - "Using Mamba state slice transfer for different TP sizes between prefill and decode. " - f"Prefill attn_tp_size={self.attn_tp_size}, Decode attn_tp_size={dst_attn_tp_size}. " + "Using Mamba state slice transfer for different runtime attention TP " + f"sizes: prefill={self.attn_tp_size}, decode={dst_attn_tp_size}. " "Performance may be affected." ) assert len(prefill_mamba_index) == 1, "Mamba should have single state index" diff --git a/python/sglang/srt/disaggregation/mori/conn.py b/python/sglang/srt/disaggregation/mori/conn.py index 42c3409eb..2b62d97ab 100644 --- a/python/sglang/srt/disaggregation/mori/conn.py +++ b/python/sglang/srt/disaggregation/mori/conn.py @@ -1171,6 +1171,11 @@ class MoriKVManager(CommonKVManager): ) if st == "mamba": + if peer_info.decode_tp_size != self.attn_tp_size and 0 in src_dims: + raise RuntimeError( + "Replicated Mamba PD state transfer currently requires " + "matching prefill/decode attention TP sizes" + ) statuses.extend( self._send_mamba_state( peer_info, @@ -1184,7 +1189,15 @@ class MoriKVManager(CommonKVManager): dst_dims, ) ) - elif st in ("swa", "dsa", "swa_ring", "c128_state", "minimax_index_k"): + elif st in ( + "swa", + "dsa", + "qsa_pending", + "qsa_compressed", + "swa_ring", + "c128_state", + "minimax_index_k", + ): statuses.extend( self._send_swa_dsa_state( peer_info, @@ -1311,14 +1324,19 @@ class MoriKVManager(CommonKVManager): f"PD state transfer does not support TP-mismatched non-MLA SWA models " f"(prefill_tp_size={self.attn_tp_size}, decode_tp_size={peer_info.decode_tp_size})" ) - if state_type == "minimax_index_k": + if state_type in ("qsa_pending", "qsa_compressed", "minimax_index_k"): if self.pp_size is not None and self.pp_size > 1: + # MORI registration does not exchange state_layer_ids. Compact + # sparse-state lists therefore cannot be paired safely across + # pipeline stages until that metadata is added to its protocol. raise RuntimeError( - "PD disagg: PP>1 not supported for MiniMax sparse index yet." + f"MORI PD disaggregation requires PP=1 for {state_type}; " + "PP>1 needs peer state_layer_ids for global-layer descriptor " + "pairing." ) if peer_info.decode_tp_size != self.attn_tp_size: raise RuntimeError( - "PD disagg: heterogeneous TP not supported for MiniMax sparse index yet." + f"PD disagg: heterogeneous TP not supported for {state_type} yet." ) common_len = min(src_state_indices.size, dst_state_indices.size) @@ -1337,7 +1355,12 @@ class MoriKVManager(CommonKVManager): # These components are position- or request-indexed: truncating # silently misaligns rows and corrupts KV. Paged swa/dsa tolerate # a 1-page drift -> keep truncation. - if state_type in ("swa_ring", "c128_state"): + if state_type in ( + "qsa_pending", + "qsa_compressed", + "swa_ring", + "c128_state", + ): raise RuntimeError( f"{state_type.upper()} state index length mismatch: " f"src={src_state_indices.size}, dst={dst_state_indices.size}" diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 2b88fe520..3705c1661 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -1574,6 +1574,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): dst_mem_kind: str = "VRAM", force_flat: bool = False, bypass_prepped: bool = False, + src_layer_ids: Optional[List[int]] = None, + dst_layer_ids: Optional[List[int]] = None, + dst_item_lens: Optional[List[int]] = None, ): """Generic KV cache transfer supporting both MHA and MLA architectures. Used by both send_kvcache and maybe_send_extra. @@ -1635,17 +1638,41 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): logger.debug(f"sending kvcache to {peer_name} with notif {notif}") # Make descs 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, state_type) - ) - layers_params = [ - ( - src_kv_ptrs[layer_id], - dst_kv_ptrs[layer_id], - item_lens[layer_id], + if src_layer_ids or dst_layer_ids: + pairs = build_transfer_entry_pairs( + src_layer_ids or [], + dst_layer_ids or [], + len(src_data_ptrs), + len(dst_data_ptrs), + allow_positional_fallback=self.pp_size == 1, ) - for layer_id in range(layers_current_pp_stage) - ] + # The source item length is used as the destination stride, so + # the paired entries must have identical layouts. + if dst_item_lens is not None: + for i, j in pairs: + if item_lens[i] != dst_item_lens[j]: + raise RuntimeError( + f"{state_type} item length mismatch for paired " + f"entries src[{i}]={item_lens[i]} " + f"dst[{j}]={dst_item_lens[j]}" + ) + layers_params = [ + (src_data_ptrs[i], dst_data_ptrs[j], item_lens[i]) for i, j in pairs + ] + else: + src_kv_ptrs, dst_kv_ptrs, layers_current_pp_stage = ( + self.get_mla_kv_ptrs_with_pp( + src_data_ptrs, dst_data_ptrs, state_type + ) + ) + layers_params = [ + ( + src_kv_ptrs[layer_id], + dst_kv_ptrs[layer_id], + item_lens[layer_id], + ) + for layer_id in range(layers_current_pp_stage) + ] else: src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = ( self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs) @@ -1667,6 +1694,9 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): for layer_id in range(layers_current_pp_stage) ] + if not layers_params: + return None + src_addrs = [] src_lens = [] dst_addrs = [] @@ -2507,6 +2537,11 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): if st == StateType.MAMBA: if self.attn_tp_size != decode_tp_size: + if 0 in src_dims: + raise RuntimeError( + "Replicated Mamba PD state transfer currently requires " + "matching prefill/decode attention TP sizes" + ) h = self._send_mamba_state_slice( peer_name, src_indices, @@ -2571,6 +2606,8 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): ) elif st in ( StateType.SWA, + StateType.QSA_PENDING, + StateType.QSA_COMPRESSED, StateType.SWA_RING, StateType.DSV4_REQUEST_STATE, ): @@ -2599,6 +2636,10 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): dst_gpu_id=dst_gpu_id, notif=comp_notif, state_type=st, + force_flat=st in (StateType.QSA_PENDING, StateType.QSA_COMPRESSED), + src_layer_ids=src_lids, + dst_layer_ids=dst_lids, + dst_item_lens=dst_lens, ) elif st == StateType.MINIMAX_INDEX_K: # Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 05c08ff7d..255ae893d 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -48,6 +48,7 @@ from sglang.srt.disaggregation.utils import ( get_dsa_tail_state_indices, get_dsv4_c128_state_indices, get_kv_class, + get_qsa_pending_state_indices, is_aborted, is_dsv4_c128_online_enabled, is_mla_backend, @@ -1346,6 +1347,11 @@ class SchedulerDisaggregationPrefillMixin: seq_len, ) + def _qsa_pending_payload(): + # Raw index-K/RoPE state is one full compression-group ring per + # request, addressed by the request-pool slot rather than KV pages. + return get_qsa_pending_state_indices(req) + def _swa_ring_payload(): # Unified_kv SWA ring rows (req_pool_idx*ring_stride + pos%ring_stride) # for the last `window` positions, in ascending position order so @@ -1380,6 +1386,8 @@ class SchedulerDisaggregationPrefillMixin: ) payloads = { StateType.MAMBA: _mamba_payload, + StateType.QSA_PENDING: _qsa_pending_payload, + StateType.QSA_COMPRESSED: _full_kv_pages_payload, StateType.SWA: _swa_payload, StateType.DSA: _full_kv_pages_payload, StateType.DSA_TAIL: _dsa_tail_payload, diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 4ed40446d..c6bd1080c 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -122,6 +122,14 @@ def get_dsv4_c128_state_indices( return np.array([page], dtype=np.int32) +def get_qsa_pending_state_indices(req: Req) -> np.ndarray: + """Return the request-pool row that owns a QSA pending-state ring.""" + req_pool_idx = req.kv.req_pool_idx + if req_pool_idx is None: + raise ValueError("QSA pending-state transfer requires an allocated request row") + return np.array([int(req_pool_idx)], dtype=np.int32) + + class DisaggregationMode(Enum): NULL = "null" PREFILL = "prefill" @@ -762,6 +770,41 @@ def is_mla_backend(target_kv_pool) -> bool: return isinstance(target_kv_pool, (MLATokenToKVPool, DeepSeekV4TokenToKVPool)) +def should_send_replicated_state( + *, + src_attn_tp_size: int, + dst_attn_tp_size: int, + local_tp_rank_in_group: int, +) -> bool: + """Elect writers for state replicated within an attention-TP group. + + Scatter (one source rank to several destination ranks) is a broadcast, so + the source sends to every destination registration. Aggregation has several + equivalent source copies targeting one destination; only the first source + in each aggregation group writes it. + """ + if src_attn_tp_size <= 0 or dst_attn_tp_size <= 0: + raise ValueError( + "Attention TP sizes must be positive for replicated-state transfer" + ) + larger_tp_size = max(src_attn_tp_size, dst_attn_tp_size) + smaller_tp_size = min(src_attn_tp_size, dst_attn_tp_size) + if larger_tp_size % smaller_tp_size != 0: + raise ValueError( + "One attention TP size must divide the other for replicated-state " + f"transfer: src={src_attn_tp_size}, dst={dst_attn_tp_size}" + ) + if not 0 <= local_tp_rank_in_group < src_attn_tp_size: + raise ValueError( + "Source attention TP rank is out of range for replicated-state " + f"transfer: rank={local_tp_rank_in_group}, size={src_attn_tp_size}" + ) + if src_attn_tp_size <= dst_attn_tp_size: + return True + writers_per_decode = src_attn_tp_size // dst_attn_tp_size + return local_tp_rank_in_group % writers_per_decode == 0 + + def compute_mamba_state_slice_blocks( src_dim: int, dst_dim: int, @@ -851,8 +894,28 @@ def compute_mamba_state_slice_byte_blocks( ``outer_count`` is one for the usual ``[slice_dim, ...]`` layout. Kimi conv state is ``[K - 1, slice_dim]``, so each logical channel slice expands - into one byte block per convolution row. + into one byte block per convolution row. A zero src/dst dim marks an item + replicated across attention TP and copies the whole item from an elected + source rank. """ + if (src_dim == 0) != (dst_dim == 0): + raise ValueError( + "Mamba state replication metadata differs between prefill and decode" + ) + if src_dim == 0: + if src_item_len != dst_item_len: + raise ValueError( + "Replicated Mamba state item lengths differ between prefill and " + f"decode: {src_item_len} != {dst_item_len}" + ) + if not should_send_replicated_state( + src_attn_tp_size=src_attn_tp_size, + dst_attn_tp_size=dst_attn_tp_size, + local_tp_rank_in_group=local_tp_rank_in_group, + ): + return [] + return [(0, 0, src_item_len)] + src_bytes_per_dim = src_item_len // (src_dim * outer_count) dst_bytes_per_dim = dst_item_len // (dst_dim * outer_count) logical_blocks = compute_mamba_state_slice_blocks( @@ -1281,6 +1344,7 @@ def setup_state_kv_args( MHATokenToKVPoolMXFP8, MiniMaxSparseKVPool, ) + from sglang.srt.mem_cache.qsa_kv_pool import QSATokenToKVPool from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool kv_args.state_types = [] @@ -1422,6 +1486,29 @@ def setup_state_kv_args( dsa_item_lens, ) append_dsa_tail(dsa_pool) + if isinstance(token_to_kv_pool, QSATokenToKVPool): + qsa_ptrs, qsa_lens, qsa_item_lens = ( + token_to_kv_pool.get_qsa_pending_state_buf_infos() + ) + append_state_component( + kv_args, + StateType.QSA_PENDING, + qsa_ptrs, + qsa_lens, + qsa_item_lens, + layer_ids=token_to_kv_pool.get_qsa_pending_state_layer_ids(), + ) + compressed_ptrs, compressed_lens, compressed_item_lens = ( + token_to_kv_pool.get_qsa_compressed_state_buf_infos() + ) + append_state_component( + kv_args, + StateType.QSA_COMPRESSED, + compressed_ptrs, + compressed_lens, + compressed_item_lens, + layer_ids=token_to_kv_pool.get_qsa_compressed_state_layer_ids(), + ) elif isinstance(token_to_kv_pool, (DSATokenToKVPool, NPUMLATokenToKVPool)): tail_ptrs, tail_lens, tail_item_lens = [], [], [] if isinstance(token_to_kv_pool, DSATokenToKVPool): diff --git a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py index 54b28ceff..d42d6695b 100644 --- a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py +++ b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py @@ -256,7 +256,16 @@ class QwenSparseAttnBackend(AttentionBackend): ) return max(1, int(sequence_lengths.max())) spec_info = forward_batch.spec_info - draft_window = int(spec_info.draft_token_num) if spec_info is not None else 0 + # Target verify exposes ``draft_token_num`` while draft-extend exposes + # ``num_tokens_per_req``. Both modes use this gather-width bound. + draft_window = int( + getattr( + spec_info, + "draft_token_num", + getattr(spec_info, "num_tokens_per_req", 0), + ) + or 0 + ) return max(1, int(seq_lens_cpu.max()) + draft_window) @staticmethod diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 94d6113e7..82784f11e 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -1024,6 +1024,22 @@ class KVCacheConfigurator: mamba_layer_ids.append(layer_id) return mamba_layer_ids + def _get_ple_req_pool_kwargs(self) -> dict[str, Any]: + from sglang.srt.configs.qwen4_exp import Qwen4ExpTextConfig + + if not isinstance(self.mambaish_config, Qwen4ExpTextConfig): + return {} + return { + "short_conv_layer_ids": [ + i + for i in self.mambaish_config.short_conv_layer_ids + if self.layer_info.start_layer <= i < self.layer_info.end_layer + ], + "short_conv_state_shape": self.mambaish_config.short_conv_state_shape, + "ngram_context_len": self.mambaish_config.ngram_context_len, + "ngram_eos_token_id": int(self.mambaish_config.eos_token_id), + } + def _build_hybrid_mamba_decode_req_pool( self, *, @@ -1049,6 +1065,7 @@ class KVCacheConfigurator: enable_overlap_schedule=not get_schedule().disable_overlap_schedule, mamba_size=get_schedule().max_mamba_cache_size, start_layer=self.layer_info.start_layer, + **self._get_ple_req_pool_kwargs(), linear_replayssm_cache_len=get_exec().mamba.linear_replayssm_cache_len, mamba_envelope_layout=get_memory().enable_page_major_kv_layout, # ReplaySSM spec-verify is for linear-attn models (GDN fold or KDA @@ -1106,20 +1123,6 @@ class KVCacheConfigurator: "--enable-linear-replayssm-spec with DSPARK/DFLASH requires a KDA " "(kimi_linear) model; got a non-KDA model." ) - from sglang.srt.configs.qwen4_exp import Qwen4ExpTextConfig - - ple_kwargs = {} - if isinstance(self.mambaish_config, Qwen4ExpTextConfig): - ple_kwargs = dict( - short_conv_layer_ids=[ - i - for i in self.mambaish_config.short_conv_layer_ids - if self.layer_info.start_layer <= i < self.layer_info.end_layer - ], - short_conv_state_shape=self.mambaish_config.short_conv_state_shape, - ngram_context_len=self.mambaish_config.ngram_context_len, - ngram_eos_token_id=int(self.mambaish_config.eos_token_id), - ) req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=get_schedule().max_mamba_cache_size, @@ -1131,7 +1134,7 @@ class KVCacheConfigurator: mamba_layer_ids=self._get_mamba_layer_ids_for_req_pool(), enable_mamba_extra_buffer=get_exec().mamba.enable_mamba_extra_buffer, enable_mamba_extra_buffer_lazy=get_exec().mamba.enable_mamba_extra_buffer_lazy, - **ple_kwargs, + **self._get_ple_req_pool_kwargs(), # A PD prefill server never runs TARGET_VERIFY, so skip the # verify-only per-draft-token state snapshots (see the draft-head # case above: None => the pool skips SpeculativeState). diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 0b23b9156..db9d8654f 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1113,8 +1113,8 @@ class MambaPool: } ) - def _iter_transfer_state_tensors(self): - """Yield transferable state tensors with their per-slot slice axis.""" + def _iter_transfer_state_entries(self): + """Yield ``[slot, ...]`` state entries and their transfer metadata.""" for field, value in vars(self.mamba_cache).items(): if field in self._NON_TRANSFER_STATE_FIELDS or value is None: continue @@ -1125,20 +1125,20 @@ class MambaPool: # empty. Advertising it fails the whole batch registration. if state_tensor.numel() == 0: continue - yield field, state_tensor, slice_axis + for layer_index, layer_id in enumerate(self.mamba_layer_ids): + yield field, state_tensor[layer_index], slice_axis, layer_id + + for sibling in self._slot_siblings: + yield from sibling.iter_transfer_state_entries() def get_contiguous_buf_infos(self): """Get transferable state buffer information for RDMA registration.""" data_ptrs, data_lens, item_lens = [], [], [] - for _, state_tensor, _ in self._iter_transfer_state_tensors(): - data_ptrs += [ - state_tensor[i].data_ptr() for i in range(self.num_mamba_layers) - ] - data_lens += [state_tensor[i].nbytes for i in range(self.num_mamba_layers)] - item_lens += [ - state_tensor[i][0].nbytes for i in range(self.num_mamba_layers) - ] + for _, state_tensor, _, _ in self._iter_transfer_state_entries(): + data_ptrs.append(state_tensor.data_ptr()) + data_lens.append(state_tensor.nbytes) + item_lens.append(state_tensor[0].nbytes) return data_ptrs, data_lens, item_lens def get_state_dim_per_tensor(self): @@ -1148,13 +1148,17 @@ class MambaPool: while Kimi conv state uses the second per-slot axis. """ dim_per_tensor = [] - for _, state_tensor, slice_axis in self._iter_transfer_state_tensors(): - # state_tensor shape: [num_layers, size+1, sliceable_dim, ...] - # Kimi conv state transposes the two per-slot axes to [K-1, dim]. - axis = 2 + slice_axis - sliceable_dim = state_tensor.shape[axis] - # Repeat for each layer since we have per-layer data_ptrs - dim_per_tensor += [sliceable_dim] * self.num_mamba_layers + for _, state_tensor, slice_axis, _ in self._iter_transfer_state_entries(): + # Zero is a protocol marker for request state replicated across the + # attention-TP group. Heterogeneous PD copies the whole item from one + # elected source rank instead of slicing it as a TP-sharded tensor. + if slice_axis is None: + dim_per_tensor.append(0) + continue + # state_tensor shape: [size+1, sliceable_dim, ...]. Kimi conv state + # transposes the two per-slot axes to [K-1, dim]. + axis = 1 + slice_axis + dim_per_tensor.append(state_tensor.shape[axis]) return dim_per_tensor def get_state_layer_ids(self): @@ -1164,15 +1168,18 @@ class MambaPool: the state list tensor-major x layer. Lets PD transfer match entries by layer id when prefill (PP stage) holds a subset of the mamba layers. """ - state_tensor_count = sum(1 for _ in self._iter_transfer_state_tensors()) - return list(self.mamba_layer_ids) * state_tensor_count + return [layer_id for _, _, _, layer_id in self._iter_transfer_state_entries()] def get_state_slice_outer_counts(self): """Get the number of rows preceding each tensor's TP slice axis.""" outer_counts = [] - for _, state_tensor, slice_axis in self._iter_transfer_state_tensors(): - outer_count = math.prod(state_tensor.shape[2 : 2 + slice_axis]) - outer_counts += [outer_count] * self.num_mamba_layers + for _, state_tensor, slice_axis, _ in self._iter_transfer_state_entries(): + outer_count = ( + 1 + if slice_axis is None + else math.prod(state_tensor.shape[1 : 1 + slice_axis]) + ) + outer_counts.append(outer_count) return outer_counts def get_state_conv_shard_groups(self): @@ -1187,14 +1194,14 @@ class MambaPool: those tensors keep the single contiguous slice. """ subdims_per_tensor = [] - for field, _, _ in self._iter_transfer_state_tensors(): + for field, _, _, _ in self._iter_transfer_state_entries(): # Only conv_state carries a q/k/v decomposition. subdims = ( list(self.conv_shard_groups) if field == "conv" and self.conv_shard_groups is not None else None ) - subdims_per_tensor += [subdims] * self.num_mamba_layers + subdims_per_tensor.append(subdims) return subdims_per_tensor def get_kv_size_bytes(self): diff --git a/python/sglang/srt/mem_cache/ple_state_pool.py b/python/sglang/srt/mem_cache/ple_state_pool.py index a4242636d..a3ad34538 100644 --- a/python/sglang/srt/mem_cache/ple_state_pool.py +++ b/python/sglang/srt/mem_cache/ple_state_pool.py @@ -13,6 +13,11 @@ from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.mem_cache.utils import maybe_init_custom_mem_pool from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter +# State layer IDs are serialized as uint32 by the disaggregation protocols. +# Reserve the largest value for PLE's request-wide N-gram state, which is not +# owned by any model layer. +PLE_NGRAM_STATE_LAYER_ID = (1 << 32) - 1 + class SlotIndexedState(Protocol): """Per-request state addressed by MambaPool slot index, @@ -27,6 +32,8 @@ class SlotIndexedState(Protocol): def load_cpu_slots(self, data: Any, indices: torch.Tensor) -> None: ... + def iter_transfer_state_entries(self): ... + class ShortConvPool: def __init__( @@ -119,6 +126,18 @@ class ShortConvPool: return self.conv_state[:, indices] = data.to(self.conv_state.device, non_blocking=True) + def iter_transfer_state_entries(self): + """Yield replicated per-layer state for PD transfer.""" + if self.conv_state is None: + return + for layer_id, layer_index in self.layer_map.items(): + yield ( + "ple_short_conv", + self.conv_state[layer_index], + None, + layer_id, + ) + class NGramPool: def __init__( @@ -216,3 +235,8 @@ class NGramPool: self.context[indices.to(dtype=torch.long)] = data.to( self.context.device, non_blocking=True ) + + def iter_transfer_state_entries(self): + """Yield replicated request-wide N-gram history for PD transfer.""" + if self.context is not None: + yield "ple_ngram", self.context, None, PLE_NGRAM_STATE_LAYER_ID diff --git a/python/sglang/srt/mem_cache/qsa_kv_pool.py b/python/sglang/srt/mem_cache/qsa_kv_pool.py index f8c1c0e8e..49d69bdb7 100644 --- a/python/sglang/srt/mem_cache/qsa_kv_pool.py +++ b/python/sglang/srt/mem_cache/qsa_kv_pool.py @@ -1,13 +1,26 @@ -"""KV pools carrying the QSA sparse-attention indexer caches.""" +"""KV pools carrying the QSA sparse-attention indexer caches. + +``QSATokenToKVPool`` (compressed, Qwen4-Exp) adds the per-request pending +index-key/RoPE ring and the paged compressed-K cache on top of the hybrid +full/linear KV pool. ``QwenDSATokenToKVPool`` (tokenwise, +Qwen3Next-DSA) adds only the flat per-token index-K cache. +""" from __future__ import annotations +from contextlib import nullcontext from typing import List, Optional import torch +from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.mem_cache.memory_pool import GB, HybridLinearKVPool, MambaPool +# State layer IDs are serialized as uint32 by the disaggregation protocols. +# Reserve the value below PLE's request-wide sentinel for QSA's request-wide +# RoPE ring, which is shared by all full-attention layers. +QSA_ROPE_STATE_LAYER_ID = (1 << 32) - 2 + def _index_k_bytes(*, kv_heads: int, head_dim: int, dtype: torch.dtype) -> int: return kv_heads * head_dim * dtype.itemsize @@ -119,31 +132,49 @@ class QSATokenToKVPool(HybridLinearKVPool): ) self.qsa_num_request_slots = int(num_request_slots) ring_slots = self.qsa_num_request_slots * self.qsa_compress_ratio - self.qsa_key_state_buffer_pool = [ - torch.zeros( - (ring_slots, self.qsa_index_kv_heads, self.qsa_index_head_dim), + # These buffers participate in Mooncake PD transfer just like the base + # KV and Mamba pools. Keep their allocation in the same memory-saver + # and Mooncake custom-pool regions; otherwise MNNVL cannot resolve the + # ordinary CUDA allocation when the first QSA state page is sent. + allocation_pool = self.full_kv_pool + with ( + allocation_pool.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE), + ( + torch.cuda.use_mem_pool(allocation_pool.custom_mem_pool) + if allocation_pool.enable_custom_mem_pool + else nullcontext() + ), + ): + self.qsa_key_state_buffer_pool = [ + torch.zeros( + ( + ring_slots, + self.qsa_index_kv_heads, + self.qsa_index_head_dim, + ), + dtype=self.index_state_dtype, + device=device, + ) + for _ in full_attention_layer_ids + ] + # RoPE coordinates are layer-independent. Keep the exact Qwen4-Exp + # MRoPE position of every incomplete key so compression can rotate + # the pooled key with the group's real starting coordinate. + self.qsa_rope_position_buffer = torch.zeros( + (ring_slots, 3), dtype=torch.int64, device=device + ) + # One contiguous allocation behind per-layer views: every layer's + # compressed pages are addressable from a single base pointer. + self.qsa_compressed_flat = torch.zeros( + ( + len(full_attention_layer_ids), + self.qsa_compressed_capacity + * self.qsa_index_kv_heads + * self.qsa_index_head_dim, + ), dtype=self.index_state_dtype, device=device, ) - for _ in full_attention_layer_ids - ] - # Layer-independent MRoPE coordinate of every pending key; - # the compress kernel rotates the pooled key at the group's real start position. - self.qsa_rope_position_buffer = torch.zeros( - (ring_slots, 3), dtype=torch.int64, device=device - ) - # One contiguous allocation behind per-layer views: every layer's - # compressed pages are addressable from a single base pointer. - self.qsa_compressed_flat = torch.zeros( - ( - len(full_attention_layer_ids), - self.qsa_compressed_capacity - * self.qsa_index_kv_heads - * self.qsa_index_head_dim, - ), - dtype=self.index_state_dtype, - device=device, - ) self.qsa_compressed_k_buffer_pool = [ self.qsa_compressed_flat[layer_offset].view( self.qsa_compressed_capacity, @@ -192,6 +223,52 @@ class QSATokenToKVPool(HybridLinearKVPool): buffer = self.get_qsa_compressed_k_buffer(layer_id) buffer[loc.long()] = compressed_k.to(buffer.dtype) + @staticmethod + def _get_paged_state_buf_infos(tensors, page_size: int): + return ( + [tensor.data_ptr() for tensor in tensors], + [tensor.nbytes for tensor in tensors], + [tensor[0].nbytes * page_size for tensor in tensors], + ) + + def get_qsa_pending_state_buf_infos(self): + """Per-request pending key-state and RoPE ring transfer buffers.""" + # A PP stage without a local QSA layer never writes the shared RoPE + # ring. Do not register it as a transfer source: otherwise that stage + # can race with a QSA-owning stage and overwrite valid positions with + # its zero-initialized or stale contents. + if not self.full_attention_layer_id_mapping: + return [], [], [] + tensors = [*self.qsa_key_state_buffer_pool, self.qsa_rope_position_buffer] + return self._get_paged_state_buf_infos( + tensors, + self.qsa_compress_ratio, + ) + + def get_qsa_pending_state_layer_ids(self): + """Global layer metadata for the compact QSA pending-state list.""" + if not self.full_attention_layer_id_mapping: + return [] + return [ + *self.full_attention_layer_id_mapping.keys(), + QSA_ROPE_STATE_LAYER_ID, + ] + + def get_qsa_compressed_state_layer_ids(self): + """Global layer metadata for the compact compressed-K list.""" + return list(self.full_attention_layer_id_mapping.keys()) + + def get_qsa_compressed_state_buf_infos(self): + """Per-full-page compressed-K transfer buffers. + + One full KV page maps to one compressed page because the full page size + is an integer multiple of the compression ratio. + """ + return self._get_paged_state_buf_infos( + self.qsa_compressed_k_buffer_pool, + self.qsa_compressed_page_size, + ) + def get_kv_size_bytes(self): k_size, v_size = super().get_kv_size_bytes() qsa_k_size = ( diff --git a/test/registered/unit/disaggregation/test_disaggregation_wire.py b/test/registered/unit/disaggregation/test_disaggregation_wire.py index c5cb2bdba..2c8709483 100644 --- a/test/registered/unit/disaggregation/test_disaggregation_wire.py +++ b/test/registered/unit/disaggregation/test_disaggregation_wire.py @@ -33,15 +33,23 @@ from sglang.srt.disaggregation.mooncake.conn import ( ) from sglang.srt.disaggregation.utils import ( MetadataBuffers, + build_transfer_entry_pairs, + compute_mamba_state_slice_byte_blocks, get_dsv4_c4_state_indices, get_dsv4_c128_state_indices, + get_qsa_pending_state_indices, setup_state_kv_args, + should_send_replicated_state, ) from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.utils import should_use_dsa_fused_topk from sglang.srt.managers.overlap_utils import FutureMap, RelayPayload from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool +from sglang.srt.mem_cache.qsa_kv_pool import ( + QSA_ROPE_STATE_LAYER_ID, + QSATokenToKVPool, +) from sglang.srt.runtime_context import get_context from sglang.srt.speculative.eagle_disaggregation import ( build_eagle_disagg_draft_input, @@ -189,6 +197,132 @@ class TestCPReplicatedStateTransfer(unittest.TestCase): ) +class TestQwen4StateWire(unittest.TestCase): + def test_qsa_pending_payload_uses_nested_request_pool_row(self): + req = SimpleNamespace(kv=ReqKvInfo(req_pool_idx=7)) + + np.testing.assert_array_equal( + get_qsa_pending_state_indices(req), + np.array([7], dtype=np.int32), + ) + + def test_qsa_registers_request_ring_and_page_state_separately(self): + pool = object.__new__(QSATokenToKVPool) + pool.full_kv_pool = object() + pool.get_state_buf_infos = lambda: ([10], [100], [20]) + pool.get_state_dim_per_tensor = lambda: [4] + pool.get_state_conv_shard_groups = lambda: [None] + pool.get_state_slice_outer_counts = lambda: [1] + pool.get_state_layer_ids = lambda: [2] + pool.page_size = 4 + pool.qsa_compress_ratio = 2 + pool.qsa_compressed_page_size = 2 + pool.full_attention_layer_id_mapping = {24: 0} + pool.qsa_key_state_buffer_pool = [torch.zeros((6, 1, 8), dtype=torch.bfloat16)] + pool.qsa_rope_position_buffer = torch.zeros((6, 3), dtype=torch.int64) + pool.qsa_compressed_k_buffer_pool = [ + torch.zeros((6, 1, 8), dtype=torch.bfloat16) + ] + + kv_args = SimpleNamespace() + setup_state_kv_args(kv_args, pool) + + self.assertEqual( + kv_args.state_types, + [StateType.MAMBA, StateType.QSA_PENDING, StateType.QSA_COMPRESSED], + ) + # Pending entries are whole two-row request rings; compressed-K remains + # a two-row compressed page corresponding to one four-token KV page. + self.assertEqual(kv_args.state_item_lens[1:], [[32, 48], [32]]) + self.assertEqual( + kv_args.state_layer_ids[1:], + [[24, QSA_ROPE_STATE_LAYER_ID], [24]], + ) + + def test_qsa_stage_without_qsa_layers_does_not_register_rope_ring(self): + pool = object.__new__(QSATokenToKVPool) + pool.full_kv_pool = object() + pool.get_state_buf_infos = lambda: ([10], [100], [20]) + pool.get_state_dim_per_tensor = lambda: [4] + pool.get_state_conv_shard_groups = lambda: [None] + pool.get_state_slice_outer_counts = lambda: [1] + pool.get_state_layer_ids = lambda: [2] + pool.page_size = 4 + pool.qsa_compress_ratio = 2 + pool.qsa_compressed_page_size = 2 + pool.full_attention_layer_id_mapping = {} + pool.qsa_key_state_buffer_pool = [] + pool.qsa_rope_position_buffer = torch.zeros((6, 3), dtype=torch.int64) + pool.qsa_compressed_k_buffer_pool = [] + + kv_args = SimpleNamespace() + setup_state_kv_args(kv_args, pool) + + # Keep the component slots aligned across PP stages, but expose no QSA + # buffers or layer ids from a stage that cannot produce their contents. + self.assertEqual( + kv_args.state_types, + [StateType.MAMBA, StateType.QSA_PENDING, StateType.QSA_COMPRESSED], + ) + self.assertEqual(kv_args.state_data_ptrs[1:], [[], []]) + self.assertEqual(kv_args.state_data_lens[1:], [[], []]) + self.assertEqual(kv_args.state_item_lens[1:], [[], []]) + self.assertEqual(kv_args.state_layer_ids[1:], [[], []]) + + def test_compact_qsa_entries_map_by_global_layer_id(self): + self.assertEqual( + build_transfer_entry_pairs( + [24, QSA_ROPE_STATE_LAYER_ID], + [0, 12, 24, QSA_ROPE_STATE_LAYER_ID], + 2, + 4, + ), + [(0, 2), (1, 3)], + ) + + def test_replicated_state_tp_policy(self): + for src_tp, dst_tp, rank, expected in ( + (4, 1, 0, True), + (4, 1, 1, False), + (1, 4, 0, True), + (4, 4, 3, True), + ): + with self.subTest(src_tp=src_tp, dst_tp=dst_tp, rank=rank): + self.assertEqual( + should_send_replicated_state( + src_attn_tp_size=src_tp, + dst_attn_tp_size=dst_tp, + local_tp_rank_in_group=rank, + ), + expected, + ) + + common = dict( + src_item_len=96, + dst_item_len=96, + src_dim=0, + dst_dim=0, + outer_count=1, + src_attn_tp_size=4, + dst_attn_tp_size=1, + dst_tp_rank_in_group=0, + ) + self.assertEqual( + compute_mamba_state_slice_byte_blocks(**common, local_tp_rank_in_group=0), + [(0, 0, 96)], + ) + self.assertEqual( + compute_mamba_state_slice_byte_blocks(**common, local_tp_rank_in_group=1), + [], + ) + with self.assertRaisesRegex(ValueError, "must divide"): + should_send_replicated_state( + src_attn_tp_size=3, + dst_attn_tp_size=2, + local_tp_rank_in_group=0, + ) + + class TestMooncakeTransferInfoIsDummy(unittest.TestCase): """Truth table for mooncake's payload-inferred is_dummy, with frames built as KVSender sends them: kv and aux are empty iff dummy, state indices are diff --git a/test/registered/unit/disaggregation/test_nixl_backend_basic.py b/test/registered/unit/disaggregation/test_nixl_backend_basic.py index 503a33dc0..d8ebcde29 100644 --- a/test/registered/unit/disaggregation/test_nixl_backend_basic.py +++ b/test/registered/unit/disaggregation/test_nixl_backend_basic.py @@ -11,7 +11,7 @@ from unittest.mock import MagicMock, patch import numpy as np -from sglang.srt.disaggregation.base.conn import KVPoll +from sglang.srt.disaggregation.base.conn import KVPoll, StateType from sglang.srt.disaggregation.common.conn import CommonKVManager from sglang.srt.disaggregation.common.staging_handler import PrefillStagingContext from sglang.srt.disaggregation.common.utils import pack_int_lists @@ -439,6 +439,59 @@ class TestNixlKVSenderChunkPolicy(CustomTestCase): self.assertTrue(sender.should_send_kv_chunk(3, last_chunk=False)) +class TestNixlEmptyStateTransfer(CustomTestCase): + def test_empty_pp_state_component_is_a_noop(self): + mgr = object.__new__(NixlKVManager) + mgr.agent = StagingFakeAgent() + mgr.is_mla_backend = False + mgr.pp_size = 2 + mgr.kv_args = SimpleNamespace(prefill_start_layer=0, kv_data_ptrs=[1]) + + handle = mgr._send_kvcache_generic( + peer_name="decode", + src_data_ptrs=[], + dst_data_ptrs=[], + item_lens=[], + prefill_data_indices=np.array([3], dtype=np.int32), + dst_data_indices=np.array([5], dtype=np.int32), + dst_gpu_id=0, + notif="qsa-empty", + state_type=StateType.QSA_PENDING, + force_flat=True, + src_layer_ids=[], + dst_layer_ids=[], + ) + + self.assertIsNone(handle) + self.assertEqual(mgr.agent.get_xfer_descs_calls, []) + self.assertEqual(mgr.agent.initialize_xfer_calls, []) + + def test_paired_state_entries_reject_item_length_mismatch(self): + mgr = object.__new__(NixlKVManager) + mgr.agent = StagingFakeAgent() + mgr.is_mla_backend = False + mgr.pp_size = 1 + mgr.kv_args = SimpleNamespace(prefill_start_layer=0, kv_data_ptrs=[1]) + + with self.assertRaisesRegex(RuntimeError, "item length mismatch"): + mgr._send_kvcache_generic( + peer_name="decode", + src_data_ptrs=[10], + dst_data_ptrs=[20], + item_lens=[32], + prefill_data_indices=np.array([3], dtype=np.int32), + dst_data_indices=np.array([5], dtype=np.int32), + dst_gpu_id=0, + notif="qsa-mismatch", + state_type=StateType.QSA_PENDING, + force_flat=True, + src_layer_ids=[24], + dst_layer_ids=[24], + dst_item_lens=[48], + ) + self.assertEqual(mgr.agent.initialize_xfer_calls, []) + + class TestNixlAbortHandling(CustomTestCase): def _make_manager(self, request_status=None): mgr = object.__new__(NixlKVManager) diff --git a/test/registered/unit/mem_cache/test_mamba_state_transfer_buffers.py b/test/registered/unit/mem_cache/test_mamba_state_transfer_buffers.py index 18e862f5b..4b889a735 100644 --- a/test/registered/unit/mem_cache/test_mamba_state_transfer_buffers.py +++ b/test/registered/unit/mem_cache/test_mamba_state_transfer_buffers.py @@ -15,6 +15,8 @@ def _pool(temporal: torch.Tensor, num_conv: int = 2) -> MambaPool: """A MambaPool stub carrying only what the transfer accessors read.""" pool = object.__new__(MambaPool) pool.num_mamba_layers = NUM_LAYERS + pool.mamba_layer_ids = list(range(NUM_LAYERS)) + pool._slot_siblings = [] pool.conv_slice_axis = 0 pool.mamba_cache = MambaPool.State( conv=[torch.zeros(NUM_LAYERS, NUM_SLOTS, 4, 5) for _ in range(num_conv)], @@ -54,6 +56,18 @@ class TestMambaStateTransferBuffers(unittest.TestCase): self.assertEqual(len(pool.get_state_dim_per_tensor()), len(lens)) + def test_sibling_declares_replicated_transfer_without_field_name_coupling(self): + pool = _pool(torch.zeros(NUM_LAYERS, NUM_SLOTS, 6, 7, 8)) + + class ReplicatedSibling: + def iter_transfer_state_entries(self): + yield "future_sibling", torch.zeros(NUM_SLOTS, 9), None, 123 + + pool._slot_siblings = [ReplicatedSibling()] + + self.assertEqual(pool.get_state_dim_per_tensor()[-1], 0) + self.assertEqual(pool.get_state_slice_outer_counts()[-1], 1) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/mem_cache/test_qsa_kv_pool.py b/test/registered/unit/mem_cache/test_qsa_kv_pool.py new file mode 100644 index 000000000..94507a686 --- /dev/null +++ b/test/registered/unit/mem_cache/test_qsa_kv_pool.py @@ -0,0 +1,68 @@ +import sys +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest +import torch + +from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool +from sglang.srt.mem_cache.qsa_kv_pool import QSATokenToKVPool +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def test_qsa_allocations_follow_parent_mooncake_scope(monkeypatch): + active_scopes = set() + allocations = 0 + original_zeros = torch.zeros + + @contextmanager + def scope(name): + active_scopes.add(name) + try: + yield + finally: + active_scopes.remove(name) + + def init_parent(pool, **_): + pool.full_kv_pool = SimpleNamespace( + memory_saver_adapter=SimpleNamespace( + region=lambda _: scope("memory_saver") + ), + enable_custom_mem_pool=True, + custom_mem_pool=object(), + ) + + def allocate(*args, **kwargs): + nonlocal allocations + assert active_scopes == {"memory_saver", "custom_pool"} + allocations += 1 + return original_zeros(*args, **kwargs) + + monkeypatch.setattr(HybridLinearKVPool, "__init__", init_parent) + monkeypatch.setattr(QSATokenToKVPool, "get_kv_size_bytes", lambda _: (0, 0)) + monkeypatch.setattr(torch.cuda, "use_mem_pool", lambda _: scope("custom_pool")) + monkeypatch.setattr("sglang.srt.mem_cache.qsa_kv_pool.torch.zeros", allocate) + + QSATokenToKVPool( + size=8, + dtype=torch.bfloat16, + page_size=4, + head_num=1, + head_dim=8, + full_attention_layer_ids=[1, 3], + device="cpu", + mamba_pool=object(), + qsa_index_kv_heads=1, + qsa_index_head_dim=8, + qsa_compress_ratio=2, + qsa_token_topk=4, + num_request_slots=3, + ) + + assert allocations == 4 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index 7f146151f..8e0933311 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -615,16 +615,22 @@ class TestGoldenModelOverrides(_IsolatedPublish): # value: readers only ever read flags. self.assertEqual((self._publish(sa), self._leaf("dtype"))[1], "auto") - def test_qwen4_rejects_pd_and_unified_memory(self): + def test_qwen4_pd_support_and_remaining_limits(self): qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp") - for kwargs, message in ( - ({"disaggregation_mode": "prefill"}, "PD disaggregation"), - ({"disaggregation_mode": "decode"}, "PD disaggregation"), - ({"enable_unified_memory": True}, "enable-unified-memory"), - ): - with self.subTest(**kwargs): - with self.assertRaisesRegex(ValueError, message): - self._construct(*qwen4, **kwargs) + with override_platform(is_cuda=True): + for mode in ("prefill", "decode"): + with self.subTest(mode=mode): + self._construct(*qwen4, disaggregation_mode=mode) + + with self.assertRaisesRegex(ValueError, "enable-unified-memory"): + self._construct(*qwen4, enable_unified_memory=True) + with self.assertRaisesRegex(ValueError, "MORI requires --pp-size 1"): + self._construct( + *qwen4, + disaggregation_mode="prefill", + disaggregation_transfer_backend="mori", + pp_size=2, + ) def test_qwen4_ple_offload_default(self): qwen4 = ("Qwen4ExpForConditionalGeneration", "qwen4_exp")