diff --git a/python/sglang/srt/configs/mamba_utils.py b/python/sglang/srt/configs/mamba_utils.py index 702fabd31..54d8cf95f 100644 --- a/python/sglang/srt/configs/mamba_utils.py +++ b/python/sglang/srt/configs/mamba_utils.py @@ -148,6 +148,12 @@ class Mamba2StateShape: # GDN kernels infer from `mixed_qkv`). Used by the GDN ReplaySSM ring # buffer (k_cache) to size/stride exactly like the kernel expects. num_k_heads_per_tp: int = 1 + # Full (unsharded) conv sub-block dims, e.g. GDN's [key_dim, key_dim, + # value_dim] for conv_state == cat([query, key, value]). Each sub-block is + # head-sharded INDEPENDENTLY across attn-TP, so PD transfer across different + # attn_tp_size must slice per sub-block. None when the single contiguous + # slice already matches the layout (e.g. standard Mamba2 conv order differs). + conv_shard_groups: Optional[List[int]] = None @staticmethod def create( @@ -159,6 +165,7 @@ class Mamba2StateShape: head_dim: int, state_size: int, conv_kernel: int, + conv_shard_groups: Optional[List[int]] = None, ) -> "Mamba2StateShape": # The q/k projections are sharded by `num_k_heads // tp` heads (the # ORIGINAL n_groups, before the conv head-shard extension below), so the @@ -196,6 +203,7 @@ class Mamba2StateShape: state_size=state_size, conv_kernel=conv_kernel, num_k_heads_per_tp=num_k_heads_per_tp, + conv_shard_groups=conv_shard_groups, ) diff --git a/python/sglang/srt/configs/qwen3_next.py b/python/sglang/srt/configs/qwen3_next.py index 4c8996d26..8ae26297c 100644 --- a/python/sglang/srt/configs/qwen3_next.py +++ b/python/sglang/srt/configs/qwen3_next.py @@ -291,6 +291,12 @@ class Qwen3NextConfig(PretrainedConfig): world_size = get_parallel().attn_tp_size adjust_tp_num_heads_if_necessary(self, world_size, False) + # GDN conv_state == cat([query, key, value]); each sub-block is + # head-sharded independently across attn-TP, so record the full + # (unsharded) sub-block dims for correct PD transfer between prefill and + # decode with different attn_tp_size (see _send_mamba_state_slice). + key_dim = self.linear_key_head_dim * self.linear_num_key_heads + value_dim = self.linear_value_head_dim * self.linear_num_value_heads shape = Mamba2StateShape.create( tp_world_size=get_parallel().attn_tp_size, intermediate_size=self.linear_value_head_dim * self.linear_num_value_heads, @@ -299,6 +305,7 @@ class Qwen3NextConfig(PretrainedConfig): head_dim=self.linear_value_head_dim, state_size=self.linear_key_head_dim, conv_kernel=self.linear_conv_kernel_dim, + conv_shard_groups=[key_dim, key_dim, value_dim], ) return Mamba2CacheParams( diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index b0566e05f..f50110f16 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -50,6 +50,10 @@ class KVArgs: # Per-tensor TP slice dim, used when prefill/decode attn_tp_size differ. state_dim_per_tensor: List[List[int]] is_hybrid_mla_backend: bool + # Per-tensor conv sub-block dims (GDN: [key_dim, key_dim, value_dim]) so the + # scatter transfer can slice each independently head-sharded sub-block; None + # per tensor when the single contiguous slice already matches the layout. + state_conv_shard_groups: List[List[Optional[List[int]]]] ib_device: str ib_traffic_class: str gpu_id: int diff --git a/python/sglang/srt/disaggregation/common/staging_buffer.py b/python/sglang/srt/disaggregation/common/staging_buffer.py index b6fb6ba1a..19eccbf4a 100644 --- a/python/sglang/srt/disaggregation/common/staging_buffer.py +++ b/python/sglang/srt/disaggregation/common/staging_buffer.py @@ -708,9 +708,11 @@ def compute_head_slice_params( unique_head_idx = local_tp_rank // src_replication dst_head_start = (unique_head_idx * src_heads_per_rank) % dst_heads_per_rank else: - src_head_start = ( - dst_tp_rank_in_group * dst_heads_per_rank - ) % src_heads_per_rank + # GQA replication: consecutive decode ranks share a KV head + # (tp_rank // num_kv_head_replicas), so map by integer division not modulo. + dst_replication = max(1, dst_attn_tp_size // total_kv_heads) + unique_dst_head_idx = dst_tp_rank_in_group // dst_replication + src_head_start = (unique_dst_head_idx * dst_heads_per_rank) % src_heads_per_rank num_heads_to_send = dst_heads_per_rank dst_head_start = 0 diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index f14d98aba..7d306a4d0 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -39,7 +39,10 @@ from sglang.srt.disaggregation.common.utils import ( from sglang.srt.disaggregation.mooncake.utils import ( check_mooncake_custom_mem_pool_enabled, ) -from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + compute_mamba_state_slice_blocks, +) from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine from sglang.srt.environ import envs from sglang.srt.observability.mooncake_trace import ( @@ -760,8 +763,13 @@ class MooncakeKVManager(CommonKVManager): ) % dst_heads_per_rank else: # Send KVCache from 1 prefill instance to multiple decode instances + # GQA replication (total_kv_heads < dst_attn_tp_size): consecutive decode + # ranks share one KV head (QKVParallelLinear: tp_rank // num_kv_head_replicas), + # so map by integer division NOT modulo or ranks 1..r-1 fetch the wrong head. + dst_replication = max(1, dst_attn_tp_size // total_kv_heads) + unique_dst_head_idx = dst_tp_rank_in_group // dst_replication src_head_start_offset = ( - dst_tp_rank_in_group * dst_heads_per_rank + unique_dst_head_idx * dst_heads_per_rank ) % src_heads_per_rank num_heads_to_send = dst_heads_per_rank dst_head_start_offset = 0 @@ -971,6 +979,10 @@ class MooncakeKVManager(CommonKVManager): if i < len(self.kv_args.state_dim_per_tensor) else [] ) + src_conv_shard_groups = getattr(self.kv_args, "state_conv_shard_groups", []) + src_conv_shard_groups = ( + src_conv_shard_groups[i] if i < len(src_conv_shard_groups) else [] + ) if target_rank_registration_info is not None: dst_data_ptrs = ( target_rank_registration_info.dst_state_data_ptrs[i] @@ -1012,6 +1024,7 @@ class MooncakeKVManager(CommonKVManager): dst_dim_per_tensor, target_rank_registration_info.dst_tp_rank, target_rank_registration_info.dst_attn_tp_size, + src_conv_shard_groups, ) or rc ) @@ -1150,6 +1163,7 @@ class MooncakeKVManager(CommonKVManager): dst_state_dim_per_tensor: list[int], dst_tp_rank: int, dst_attn_tp_size: int, + src_state_conv_shard_groups: list = None, ): """Transfer Mamba states with TP slice support. @@ -1158,7 +1172,10 @@ class MooncakeKVManager(CommonKVManager): - temporal_state: [num_layers, size+1, num_heads/tp, head_dim, state_size] The 3rd dimension is sliced by TP. When prefill and decode have different - attn_tp_size, we need to slice the state accordingly. + attn_tp_size, we slice the state accordingly. GDN conv_state is the + concatenation [query | key | value] with each sub-block head-sharded + independently, so on the scatter path it is sliced per sub-block via + ``src_state_conv_shard_groups`` (see compute_mamba_state_slice_blocks). """ logger.warning_once( "Using Mamba state slice transfer for different TP sizes between prefill and decode. " @@ -1192,33 +1209,42 @@ class MooncakeKVManager(CommonKVManager): src_bytes_per_dim = src_item_len // src_dim dst_bytes_per_dim = dst_item_len // dst_dim - if self.attn_tp_size > dst_attn_tp_size: - # Multiple prefill ranks send to 1 decode rank - src_dim_start = 0 - num_dims_to_send = src_dim - writers_per_decode = self.attn_tp_size // dst_attn_tp_size - local_writer_idx = local_tp_rank_in_group % writers_per_decode - dst_dim_start = local_writer_idx * src_dim - else: - # 1 prefill rank sends to multiple decode ranks - src_dim_start = (dst_tp_rank_in_group * dst_dim) % src_dim - num_dims_to_send = dst_dim - dst_dim_start = 0 - - src_dim_offset = src_dim_start * src_bytes_per_dim - dst_dim_offset = dst_dim_start * dst_bytes_per_dim - bytes_to_send = num_dims_to_send * src_bytes_per_dim - - src_addr = ( - src_state_data_ptrs[i] - + src_item_len * int(prefill_mamba_index[0]) - + src_dim_offset - ) - dst_addr = ( - dst_state_ptr + dst_item_len * int(dst_mamba_index[0]) + dst_dim_offset + conv_shard_groups = ( + src_state_conv_shard_groups[i] + if src_state_conv_shard_groups and i < len(src_state_conv_shard_groups) + else None ) + # One block for single-axis states; three (q/k/v) for GDN conv_state + # on the scatter path. + for ( + src_dim_start, + dst_dim_start, + num_dims_to_send, + ) in compute_mamba_state_slice_blocks( + src_dim=src_dim, + dst_dim=dst_dim, + src_attn_tp_size=self.attn_tp_size, + dst_attn_tp_size=dst_attn_tp_size, + dst_tp_rank_in_group=dst_tp_rank_in_group, + local_tp_rank_in_group=local_tp_rank_in_group, + conv_shard_groups=conv_shard_groups, + ): + src_dim_offset = src_dim_start * src_bytes_per_dim + dst_dim_offset = dst_dim_start * dst_bytes_per_dim + bytes_to_send = num_dims_to_send * src_bytes_per_dim - transfer_blocks.append((src_addr, dst_addr, bytes_to_send)) + src_addr = ( + src_state_data_ptrs[i] + + src_item_len * int(prefill_mamba_index[0]) + + src_dim_offset + ) + dst_addr = ( + dst_state_ptr + + dst_item_len * int(dst_mamba_index[0]) + + dst_dim_offset + ) + + transfer_blocks.append((src_addr, dst_addr, bytes_to_send)) return self._transfer_data(req.mooncake_session_id, transfer_blocks) diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 1f2597178..f60a8766a 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -32,7 +32,10 @@ from sglang.srt.disaggregation.common.utils import ( pack_int_lists, unpack_int_lists, ) -from sglang.srt.disaggregation.utils import DisaggregationMode +from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + compute_mamba_state_slice_blocks, +) from sglang.srt.environ import envs from sglang.srt.server_args import ServerArgs @@ -1868,12 +1871,16 @@ class NixlKVManager(CommonKVManager): notif: str, decode_tp_size: int, decode_tp_rank: int, + src_state_conv_shard_groups: list = None, ): """Transfer Mamba states with TP slice support via RDMA. When prefill and decode have different attn_tp_size, we slice the TP-sharded dimension (3rd dim) of conv_state and temporal_state - accordingly, mirroring Mooncake's _send_mamba_state_slice. + accordingly, mirroring Mooncake's _send_mamba_state_slice. GDN + conv_state is [query | key | value] with each sub-block head-sharded + independently, so on the scatter path it is sliced per sub-block via + ``src_state_conv_shard_groups`` (see compute_mamba_state_slice_blocks). """ logger.warning_once( "Using Mamba state slice transfer for different TP sizes. " @@ -1911,33 +1918,42 @@ class NixlKVManager(CommonKVManager): src_bytes_per_dim = src_item_len // src_dim dst_bytes_per_dim = dst_item_len // dst_dim - if self.attn_tp_size > decode_tp_size: - src_dim_start = 0 - num_dims_to_send = src_dim - writers_per_decode = self.attn_tp_size // decode_tp_size - local_writer_idx = local_tp_rank_in_group % writers_per_decode - dst_dim_start = local_writer_idx * src_dim - else: - src_dim_start = (dst_tp_rank_in_group * dst_dim) % src_dim - num_dims_to_send = dst_dim - dst_dim_start = 0 - - src_dim_offset = src_dim_start * src_bytes_per_dim - dst_dim_offset = dst_dim_start * dst_bytes_per_dim - bytes_to_send = num_dims_to_send * src_bytes_per_dim - - src_addr = ( - src_state_data_ptrs[i] - + src_item_len * int(prefill_state_indices[0]) - + src_dim_offset + conv_shard_groups = ( + src_state_conv_shard_groups[i] + if src_state_conv_shard_groups and i < len(src_state_conv_shard_groups) + else None ) - dst_addr = ( - dst_state_ptr - + dst_item_len * int(dst_state_indices[0]) - + dst_dim_offset - ) - src_addrs.append((src_addr, bytes_to_send, self.kv_args.gpu_id)) - dst_addrs.append((dst_addr, bytes_to_send, dst_gpu_id)) + # One block for single-axis states; three (q/k/v) for GDN conv_state + # on the scatter path. + for ( + src_dim_start, + dst_dim_start, + num_dims_to_send, + ) in compute_mamba_state_slice_blocks( + src_dim=src_dim, + dst_dim=dst_dim, + src_attn_tp_size=self.attn_tp_size, + dst_attn_tp_size=decode_tp_size, + dst_tp_rank_in_group=dst_tp_rank_in_group, + local_tp_rank_in_group=local_tp_rank_in_group, + conv_shard_groups=conv_shard_groups, + ): + src_dim_offset = src_dim_start * src_bytes_per_dim + dst_dim_offset = dst_dim_start * dst_bytes_per_dim + bytes_to_send = num_dims_to_send * src_bytes_per_dim + + src_addr = ( + src_state_data_ptrs[i] + + src_item_len * int(prefill_state_indices[0]) + + src_dim_offset + ) + dst_addr = ( + dst_state_ptr + + dst_item_len * int(dst_state_indices[0]) + + dst_dim_offset + ) + src_addrs.append((src_addr, bytes_to_send, self.kv_args.gpu_id)) + dst_addrs.append((dst_addr, bytes_to_send, dst_gpu_id)) src_descs = self.agent.get_xfer_descs(src_addrs, "VRAM") dst_descs = self.agent.get_xfer_descs(dst_addrs, "VRAM") @@ -1976,6 +1992,9 @@ class NixlKVManager(CommonKVManager): src_state_dim_per_tensor = ( getattr(self.kv_args, "state_dim_per_tensor", []) or [] ) + src_state_conv_shard_groups = ( + getattr(self.kv_args, "state_conv_shard_groups", []) or [] + ) dst_state_item_lens = dst_state_item_lens or [] dst_state_dim_per_tensor = dst_state_dim_per_tensor or [] @@ -1991,6 +2010,11 @@ class NixlKVManager(CommonKVManager): src_dims = ( src_state_dim_per_tensor[i] if i < len(src_state_dim_per_tensor) else [] ) + src_conv = ( + src_state_conv_shard_groups[i] + if i < len(src_state_conv_shard_groups) + else [] + ) dst_ptrs = dst_state_data_ptrs[i] if i < len(dst_state_data_ptrs) else [] dst_indices = dst_state_indices[i] if i < len(dst_state_indices) else [] dst_lens = dst_state_item_lens[i] if i < len(dst_state_item_lens) else [] @@ -2015,6 +2039,7 @@ class NixlKVManager(CommonKVManager): comp_notif, decode_tp_size, decode_tp_rank, + src_conv, ) else: h = self._send_mamba_state( diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 755d1e00c..8d741804d 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -737,6 +737,78 @@ def is_mla_backend(target_kv_pool) -> bool: return isinstance(target_kv_pool, (MLATokenToKVPool, DeepSeekV4TokenToKVPool)) +def compute_mamba_state_slice_blocks( + src_dim: int, + dst_dim: int, + src_attn_tp_size: int, + dst_attn_tp_size: int, + dst_tp_rank_in_group: int, + local_tp_rank_in_group: int, + conv_shard_groups: Optional[List[int]] = None, +) -> List[Tuple[int, int, int]]: + """Blocks to copy one mamba state item across differing attn-TP sizes. + + Returns ``(src_dim_start, dst_dim_start, num_dims)`` triples in units of the + sliceable (3rd) dimension. Single-axis states (temporal_state, or when + ``conv_shard_groups`` is None) return one contiguous block -- byte-identical to + the legacy behavior. + + GDN conv_state is ``cat([query | key | value])`` where each sub-block (full + dims == ``conv_shard_groups``, e.g. ``[key_dim, key_dim, value_dim]``) is + head-sharded INDEPENDENTLY across attn-TP. In the SCATTER direction + (1 prefill rank -> several decode ranks) a single contiguous slice straddles + the q/k/v boundaries and delivers wrong channels. The AGGREGATION direction + (several prefill ranks -> 1 decode rank) has the symmetric problem: a single + contiguous write interleaves the sub-blocks by writer. Both directions emit one + block per sub-block for conv_state; temporal_state and non-GDN states (when + ``conv_shard_groups`` is None) keep the single contiguous slice. + """ + use_subdims = ( + conv_shard_groups is not None + and sum(conv_shard_groups) == src_dim * src_attn_tp_size + ) + + if src_attn_tp_size > dst_attn_tp_size: + # Aggregation: several prefill ranks each write their shard into one decode slot. + writers_per_decode = src_attn_tp_size // dst_attn_tp_size + local_writer_idx = local_tp_rank_in_group % writers_per_decode + if not use_subdims: + return [(0, local_writer_idx * src_dim, src_dim)] + # conv_state: a plain contiguous write would interleave the sub-blocks by + # writer ([q0,k0,v0,q1,k1,v1,...]); place this writer's shard of each + # independently head-sharded sub-block at its grouped offset so the decode + # buffer is [q0,q1,...,k0,k1,...,v0,v1,...]. + blocks: List[Tuple[int, int, int]] = [] + src_off = 0 + dst_off = 0 + for full_sd in conv_shard_groups: + src_sub = full_sd // src_attn_tp_size + dst_sub = full_sd // dst_attn_tp_size + blocks.append((src_off, dst_off + local_writer_idx * src_sub, src_sub)) + src_off += src_sub + dst_off += dst_sub + return blocks + + # Scatter: 1 prefill rank feeds several decode ranks. + if not use_subdims: + src_dim_start = (dst_tp_rank_in_group * dst_dim) % src_dim + return [(src_dim_start, 0, dst_dim)] + + # conv_state: gather the decode rank's [q | k | v] shard from the three + # independently head-sharded sub-blocks of the src tensor. dst is contiguous. + blocks: List[Tuple[int, int, int]] = [] + src_off = 0 + dst_off = 0 + for full_sd in conv_shard_groups: + src_sub = full_sd // src_attn_tp_size # this prefill rank's shard of sub-block + dst_sub = full_sd // dst_attn_tp_size # this decode rank's shard of sub-block + src_start = src_off + (dst_tp_rank_in_group * dst_sub) % src_sub + blocks.append((src_start, dst_off, dst_sub)) + src_off += src_sub + dst_off += dst_sub + return blocks + + def append_state_component( kv_args: KVArgs, state_type: StateType, @@ -744,6 +816,7 @@ def append_state_component( data_lens: List[int], item_lens: List[int], dim_per_tensor: Optional[List[int]] = None, + conv_shard_groups: Optional[List[Optional[List[int]]]] = None, ) -> None: """Append one state component. Caller orders state_types consistently on prefill and decode sides.""" @@ -752,6 +825,7 @@ def append_state_component( kv_args.state_data_lens.append(data_lens) kv_args.state_item_lens.append(item_lens) kv_args.state_dim_per_tensor.append(dim_per_tensor or []) + kv_args.state_conv_shard_groups.append(conv_shard_groups or []) def setup_state_kv_args( @@ -781,6 +855,7 @@ def setup_state_kv_args( kv_args.state_item_lens = [] kv_args.state_dim_per_tensor = [] kv_args.is_hybrid_mla_backend = False + kv_args.state_conv_shard_groups = [] if isinstance(token_to_kv_pool, MiniMaxSparseKVPool): if token_to_kv_pool.index_kv_pool is not None: @@ -837,8 +912,19 @@ def setup_state_kv_args( kv_args.is_hybrid_mla_backend = is_mla_backend( token_to_kv_pool.full_kv_pool ) + conv_shard_groups = ( + token_to_kv_pool.get_state_conv_shard_groups() + if hasattr(token_to_kv_pool, "get_state_conv_shard_groups") + else None + ) append_state_component( - kv_args, StateType.MAMBA, data_ptrs, data_lens, item_lens, dim + kv_args, + StateType.MAMBA, + data_ptrs, + data_lens, + item_lens, + dim, + conv_shard_groups, ) elif isinstance(token_to_kv_pool, (DSATokenToKVPool, NPUMLATokenToKVPool)): if draft_token_to_kv_pool is not None and isinstance( @@ -950,8 +1036,19 @@ def setup_state_kv_args( if hasattr(req_to_token_pool, "get_state_dim_per_tensor") else None ) + conv_shard_groups = ( + req_to_token_pool.get_state_conv_shard_groups() + if hasattr(req_to_token_pool, "get_state_conv_shard_groups") + else None + ) append_state_component( - kv_args, StateType.MAMBA, data_ptrs, data_lens, item_lens, dim + kv_args, + StateType.MAMBA, + data_ptrs, + data_lens, + item_lens, + dim, + conv_shard_groups, ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 882cfa54f..844fd175d 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -677,6 +677,9 @@ class MambaPool: ) self.mem_usage = mem_usage_bytes / GB self.num_mamba_layers = num_mamba_layers + # Full (unsharded) conv sub-block dims for PD transfer across different + # attn_tp_size (GDN: [key_dim, key_dim, value_dim]); None otherwise. + self.conv_shard_groups = getattr(cache_params.shape, "conv_shard_groups", None) def get_speculative_mamba2_params_all_layers(self) -> SpeculativeState: assert isinstance(self.mamba_cache, self.SpeculativeState) @@ -832,6 +835,43 @@ class MambaPool: dim_per_tensor += [sliceable_dim] * self.num_mamba_layers return dim_per_tensor + def get_state_conv_shard_groups(self): + """Per-tensor conv sub-block dims, aligned element-wise with + get_state_dim_per_tensor(). + + For GDN, conv_state's sliceable axis is cat([query, key, value]) with + each sub-block head-sharded independently across attn-TP; the full + (unsharded) sub-block dims are returned so PD transfer across different + attn_tp_size can slice each sub-block. Returns None for temporal_state + (single head-sharded axis) and whenever no descriptor is available, so + those tensors keep the single contiguous slice. + """ + subdims_per_tensor = [] + for field in vars(self.mamba_cache): + # Mirror the exclusions in get_state_dim_per_tensor so the returned + # sub-dims line up element-wise with the RDMA buffer list. + if field in ( + "intermediate_ssm", + "intermediate_conv_window", + "replayssm_d", + "replayssm_k", + "replayssm_g", + ): + continue + value = getattr(self.mamba_cache, field) + if value is None: + continue + tensors = value if isinstance(value, list) else [value] + for _ in tensors: + # 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 + return subdims_per_tensor + class HybridReqToTokenPool(ReqToTokenPool): """A memory pool that maps a request to its token locations.""" @@ -1022,6 +1062,9 @@ class HybridReqToTokenPool(ReqToTokenPool): def get_state_dim_per_tensor(self): return self.mamba_pool.get_state_dim_per_tensor() + def get_state_conv_shard_groups(self): + return self.mamba_pool.get_state_conv_shard_groups() + def get_mamba_ping_pong_other_idx(self, mamba_next_track_idx: int) -> int: if self.mamba_ping_pong_track_buffer_size == 2: return 1 - mamba_next_track_idx @@ -2598,6 +2641,10 @@ class HybridLinearKVPool(KVCache): """Get the sliceable dimension size for each mamba state tensor.""" return self.mamba_pool.get_state_dim_per_tensor() + def get_state_conv_shard_groups(self): + """Per-tensor conv sub-block dims (GDN) aligned with the state list.""" + return self.mamba_pool.get_state_conv_shard_groups() + def maybe_get_custom_mem_pool(self): return self.full_kv_pool.maybe_get_custom_mem_pool() diff --git a/test/registered/disaggregation/test_disaggregation_different_tp.py b/test/registered/disaggregation/test_disaggregation_different_tp.py index e227cd73a..c3f47ea65 100644 --- a/test/registered/disaggregation/test_disaggregation_different_tp.py +++ b/test/registered/disaggregation/test_disaggregation_different_tp.py @@ -9,6 +9,7 @@ from sglang.test.server_fixtures.disaggregation_fixture import ( PDDisaggregationServerBase, ) from sglang.test.test_utils import ( + DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST, DEFAULT_MODEL_NAME_FOR_TEST, DEFAULT_MODEL_NAME_FOR_TEST_MLA, DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, @@ -504,5 +505,94 @@ class TestDisaggregationStagingDecodeLargerTP(PDDisaggregationServerBase): self.assertGreater(metrics["score"], 0.60) +class TestDisaggregationGDNHybridHeteroTP(PDDisaggregationServerBase): + """Prefill TP=1 -> Decode TP=4 on a GDN-hybrid (gated-delta-net) model. + + Exercises the heterogeneous attn-TP *scatter* path (prefill attn_tp < + decode attn_tp), where two independent bugs corrupt accuracy without the + fix in this change set: + 1. GDN conv_state is cat([query|key|value]) with each sub-block + independently head-sharded; a single contiguous slice straddles the + q/k/v boundaries and delivers wrong channels. + 2. GQA KV heads are replicated when num_key_value_heads < decode attn_tp; + the scatter head map must use integer division (tp_rank // + num_kv_head_replicas), not modulo. + Without the fix gsm8k collapses (~0.4); with it, it recovers to agg level. + """ + + @classmethod + def setUpClass(cls): + super().setUpClass() + envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False) + + cls.model = try_cached_model(DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST) + + cls.start_prefill() + cls.start_decode() + + cls.wait_server_ready(cls.prefill_url + "/health", process=cls.process_prefill) + cls.wait_server_ready(cls.decode_url + "/health", process=cls.process_decode) + + cls.launch_lb() + + @classmethod + def start_prefill(cls): + prefill_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "prefill", + "--disaggregation-bootstrap-port", + cls.bootstrap_port, + "--tp", + "1", + "--enable-metrics", + "--enable-request-time-stats-logging", + ] + prefill_args += cls.transfer_backend + cls.rdma_devices + cls.process_prefill = popen_launch_pd_server( + cls.model, + cls.prefill_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=prefill_args, + ) + + @classmethod + def start_decode(cls): + decode_args = [ + "--trust-remote-code", + "--disaggregation-mode", + "decode", + "--disaggregation-bootstrap-port", + cls.bootstrap_port, + "--tp", + "4", + "--base-gpu-id", + "4", + "--enable-metrics", + "--enable-request-time-stats-logging", + ] + decode_args += cls.transfer_backend + cls.rdma_devices + cls.process_decode = popen_launch_pd_server( + cls.model, + cls.decode_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=decode_args, + ) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + api="completion", + max_tokens=512, + num_examples=200, + num_threads=128, + ) + metrics = run_eval(args) + print(f"[GDNHybridHeteroTP] Evaluation metrics: {metrics}") + self.assertGreater(metrics["score"], 0.60) + + if __name__ == "__main__": unittest.main()