[Qwen3.8-Next] Add PD state transfer for Flash Next (#36651)

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] = {}

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

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,

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"

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}"

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

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,

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):

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

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).

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):

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

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 = (