sgl-project/sglang · error · RuntimeError

PD state transfer failed: mamba requires single state index,

Error message

PD state transfer failed: mamba requires single state index, got src={src_state_indices.size}, dst={dst_state_indices.size}

What it means

_send_mamba_state requires exactly one state index on each side (src and dst) because Mamba recurrent state is per-sequence single-slot, but got a tensor whose size differs from 1. The transfer engine cannot batch multiple Mamba states in one call.

Source

Thrown at python/sglang/srt/disaggregation/mori/conn.py:1280

            else:
                raise RuntimeError(f"PD state transfer failed: unknown state_type={st}")

        return statuses

    def _send_mamba_state(
        self,
        peer_info: KVArgsRegisterInfo,
        src_state_indices: npt.NDArray[np.int32],
        dst_state_indices: npt.NDArray[np.int32],
        src_state_mem_descs: List[MemoryDesc],
        dst_state_mem_descs: List[MemoryDesc],
        src_state_item_lens: List[int],
        dst_state_item_lens: List[int],
        src_state_dim_per_tensor: List[int],
        dst_state_dim_per_tensor: List[int],
    ) -> List[TransferStatus]:
        if src_state_indices.size != 1 or dst_state_indices.size != 1:
            raise RuntimeError(
                f"PD state transfer failed: mamba requires single state index, "
                f"got src={src_state_indices.size}, dst={dst_state_indices.size}"
            )

        tp_mismatch = peer_info.decode_tp_size != self.attn_tp_size

        # If dim info missing, silently degrade to whole-item copy (Mooncake compat)
        if tp_mismatch and (
            not src_state_dim_per_tensor or not dst_state_dim_per_tensor
        ):
            tp_mismatch = False

        if tp_mismatch:
            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={peer_info.decode_tp_size}. "
                "Performance may be affected."
            )

View on GitHub (pinned to 0132848349)

Solutions

  1. Ensure Mamba state transfers are issued one sequence at a time (loop over per-request indices before calling send)
  2. Filter out empty/invalid state index tensors before invoking the transfer
  3. Check scheduler logic that populates state_indices for finished/cancelled requests

Example fix

# before
statuses = conn.send_state(peer_info, all_src_idx, all_dst_idx, ...)  # batched tensor

# after: per-sequence transfer for mamba
for src_i, dst_i in zip(src_idx_list, dst_idx_list):
    if src_i.numel() == 1 and dst_i.numel() == 1:
        statuses.extend(conn.send_state(peer_info, src_i, dst_i, ...))
Defensive patterns

Strategy: validation

Validate before calling

if state_type == "mamba":
    assert src_state_indices.size == 1 and dst_state_indices.size == 1, (
        "mamba state transfer requires exactly one index per side"
    )

Prevention

When it happens

Trigger: Calling the Mamba state send path with src_state_indices or dst_state_indices tensors containing 0 or >1 entries — e.g. a batch of requests collapsed into one transfer call, or an empty state index passed when a sequence has no allocated Mamba slot.

Common situations: Batched PD transfer for hybrid models where the scheduler aggregates multiple requests' state indices into one send, or a scheduling edge case (cancelled/finished request) yielding an empty index tensor.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/de13d3e9a46f4262. Report an issue: GitHub.