{"record":{"id":"3b988825d2376304","repo":"sgl-project/sglang","slug":"ssm-state-indices-must-be-1d-for-packed-decode","errorCode":null,"errorMessage":"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim}).","messagePattern":"`ssm_state_indices` must be 1D for packed decode \\(got ndim=(.+?)\\)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent.py","lineNumber":297,"sourceCode":") -> tuple[torch.Tensor, torch.Tensor]:\n    if mixed_qkv.ndim != 2:\n        raise ValueError(\n            f\"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim}).\"\n        )\n    if mixed_qkv.stride(-1) != 1:\n        raise ValueError(\"`mixed_qkv` must be contiguous in the last dim.\")\n    if a.ndim != 2 or b.ndim != 2:\n        raise ValueError(\n            f\"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim}).\"\n        )\n    if a.stride(-1) != 1 or b.stride(-1) != 1:\n        raise ValueError(\"`a`/`b` must be contiguous in the last dim.\")\n    if A_log.ndim != 1 or dt_bias.ndim != 1:\n        raise ValueError(\"`A_log`/`dt_bias` must be 1D tensors.\")\n    if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:\n        raise ValueError(\"`A_log`/`dt_bias` must be contiguous.\")\n    if ssm_state_indices.ndim != 1:\n        raise ValueError(\n            f\"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim}).\"\n        )\n    if not out.is_contiguous():\n        raise ValueError(\"`out` must be contiguous.\")\n\n    dev = mixed_qkv.device\n    if any(\n        t.device != dev\n        for t in (a, b, A_log, dt_bias, initial_state, out, ssm_state_indices)\n    ):\n        raise ValueError(\"All inputs must be on the same device.\")\n\n    B = mixed_qkv.shape[0]\n    if a.shape[0] != B or b.shape[0] != B:\n        raise ValueError(\n            \"Mismatched batch sizes: \"\n            f\"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}.\"\n        )","sourceCodeStart":279,"sourceCodeEnd":315,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py#L279-L315","documentation":"ssm_state_indices maps each packed decode token to its recurrent state slot (or -1 for stateless tokens) in the SSM state cache, and must be 1D of length num_tokens. The wrapper raises when it is passed with an extra dimension, e.g. (B, 1) or (1, B), which happens when the scheduler's location indices are not flattened.","triggerScenarios":"Passing loc indices shaped (B, 1) from a decode batch, or (1, B) from an unsqueezed tensor; passing the 2D cache index table instead of the per-token flat vector.","commonSituations":"sglang hybrid-attention decode paths where mamba cache location indices come from torch.arange(B).unsqueeze(-1); mismatch between the packed (all ranks' tokens concatenated) layout and a per-rank 2D index tensor.","solutions":["Flatten to 1D: ssm_state_indices = ssm_state_indices.reshape(-1)","Ensure its length equals the number of packed tokens B = mixed_qkv.shape[0] and uses -1 for tokens without state"],"exampleFix":"# before\nout, s = ...(ssm_state_indices=locs)  # locs is (B, 1)\n# after\nout, s = ...(ssm_state_indices=locs.reshape(-1))  # (B,)","handlingStrategy":"validation","validationCode":"ssm_state_indices = ssm_state_indices.reshape(-1)\nassert ssm_state_indices.ndim == 1","typeGuard":"def indices_1d(idx: torch.Tensor) -> bool:\n    return idx.ndim == 1","tryCatchPattern":null,"preventionTips":["Flatten scheduler-provided location indices at the model/kernel boundary","Treat every packed-decode input as (num_tokens, …) and reshape all of them together"],"tags":["fla","fused-recurrent","state-cache","shape-validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}