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

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

Browse Source
This commit is contained in:
YAMY
2026-09-11 22:31:45 -07:00
committed by GitHub
parent dbd4302bbb
commit 55e5e21c88
18 changed files with 734 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/model_overrides/qwen4_exp.py
+10 -2
View File
@@ -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
+2
View File
@@ -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
+16
View File
@@ -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
+62 -6
View File
@@ -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
+28 -5
View File
@@ -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
+42 -1
View File
@@ -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,8 +1638,32 @@ 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:
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,
)
# 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)
self.get_mla_kv_ptrs_with_pp(
src_data_ptrs, dst_data_ptrs, state_type
)
)
layers_params = [
(
@@ -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
+8
View File
@@ -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
+88 -1
View File
@@ -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
+10 -1
View File
@@ -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
+18 -15
View File
@@ -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
+32 -25
View File
@@ -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
+24
View File
@@ -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
+81 -4
View File
@@ -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,16 +132,34 @@ class QSATokenToKVPool(HybridLinearKVPool):
)
self.qsa_num_request_slots = int(num_request_slots)
ring_slots = self.qsa_num_request_slots * self.qsa_compress_ratio
# 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),
(
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
]
# Layer-independent MRoPE coordinate of every pending key;
# the compress kernel rotates the pooled key at the group's real start position.
# 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
)
@@ -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 = (
test/registered/unit/disaggregation/test_disaggregation_wire.py
+134
View File
@@ -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
test/registered/unit/disaggregation/test_nixl_backend_basic.py
+54 -1
View File
@@ -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)
test/registered/unit/mem_cache/test_mamba_state_transfer_buffers.py
+14
View File
@@ -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()
test/registered/unit/mem_cache/test_qsa_kv_pool.py
+68
View File
@@ -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"]))
test/registered/unit/test_model_overrides.py
+15 -9
View File
@@ -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")