{"record":{"id":"de13d3e9a46f4262","repo":"sgl-project/sglang","slug":"pd-state-transfer-failed-mamba-requires-single-st","errorCode":null,"errorMessage":"PD state transfer failed: mamba requires single state index, got src={src_state_indices.size}, dst={dst_state_indices.size}","messagePattern":"PD state transfer failed: mamba requires single state index, got src=(.+?), dst=(.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/disaggregation/mori/conn.py","lineNumber":1280,"sourceCode":"            else:\n                raise RuntimeError(f\"PD state transfer failed: unknown state_type={st}\")\n\n        return statuses\n\n    def _send_mamba_state(\n        self,\n        peer_info: KVArgsRegisterInfo,\n        src_state_indices: npt.NDArray[np.int32],\n        dst_state_indices: npt.NDArray[np.int32],\n        src_state_mem_descs: List[MemoryDesc],\n        dst_state_mem_descs: List[MemoryDesc],\n        src_state_item_lens: List[int],\n        dst_state_item_lens: List[int],\n        src_state_dim_per_tensor: List[int],\n        dst_state_dim_per_tensor: List[int],\n    ) -> List[TransferStatus]:\n        if src_state_indices.size != 1 or dst_state_indices.size != 1:\n            raise RuntimeError(\n                f\"PD state transfer failed: mamba requires single state index, \"\n                f\"got src={src_state_indices.size}, dst={dst_state_indices.size}\"\n            )\n\n        tp_mismatch = peer_info.decode_tp_size != self.attn_tp_size\n\n        # If dim info missing, silently degrade to whole-item copy (Mooncake compat)\n        if tp_mismatch and (\n            not src_state_dim_per_tensor or not dst_state_dim_per_tensor\n        ):\n            tp_mismatch = False\n\n        if tp_mismatch:\n            logger.warning_once(\n                \"Using Mamba state slice transfer for different TP sizes between prefill and decode. \"\n                f\"Prefill attn_tp_size={self.attn_tp_size}, Decode attn_tp_size={peer_info.decode_tp_size}. \"\n                \"Performance may be affected.\"\n            )","sourceCodeStart":1262,"sourceCodeEnd":1298,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/disaggregation/mori/conn.py#L1262-L1298","documentation":"_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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Ensure Mamba state transfers are issued one sequence at a time (loop over per-request indices before calling send)","Filter out empty/invalid state index tensors before invoking the transfer","Check scheduler logic that populates state_indices for finished/cancelled requests"],"exampleFix":"# before\nstatuses = conn.send_state(peer_info, all_src_idx, all_dst_idx, ...)  # batched tensor\n\n# after: per-sequence transfer for mamba\nfor src_i, dst_i in zip(src_idx_list, dst_idx_list):\n    if src_i.numel() == 1 and dst_i.numel() == 1:\n        statuses.extend(conn.send_state(peer_info, src_i, dst_i, ...))","handlingStrategy":"validation","validationCode":"if state_type == \"mamba\":\n    assert src_state_indices.size == 1 and dst_state_indices.size == 1, (\n        \"mamba state transfer requires exactly one index per side\"\n    )\n","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Issue Mamba state sends per-sequence, never batched","Guard scheduler code against empty state index tensors"],"tags":["disaggregation","mamba","state-transfer","batching"],"backgroundTag":"invalid-argument-shape","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}