sgl-project/sglang · error · ValueError

`ssm_state_indices` must be 1D for packed decode

Error message

`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim}).

What it means

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.

Solutions

1. Flatten to 1D: ssm_state_indices = ssm_state_indices.reshape(-1)
2. Ensure its length equals the number of packed tokens B = mixed_qkv.shape[0] and uses -1 for tokens without state

Example fix

# before
out, s = ...(ssm_state_indices=locs)  # locs is (B, 1)
# after
out, s = ...(ssm_state_indices=locs.reshape(-1))  # (B,)

ssm_state_indices = ssm_state_indices.reshape(-1)
assert ssm_state_indices.ndim == 1

def indices_1d(idx: torch.Tensor) -> bool:
    return idx.ndim == 1

When it happens

Trigger: 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.

Common situations: 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.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

• `initial_state` must be a 4D tensor
• `ssm_state_indices` must have shape [B]
• `a` and `b` must be 2D tensors
• `A_log`/`dt_bias` must be 1D tensors.
• `initial_state` must be contiguous in the last dim.

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

Appendix: source

Thrown at python/sglang/kernels/ops/attention/fla/fused_recurrent.py:297

) -> tuple[torch.Tensor, torch.Tensor]:
    if mixed_qkv.ndim != 2:
        raise ValueError(
            f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})."
        )
    if mixed_qkv.stride(-1) != 1:
        raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
    if a.ndim != 2 or b.ndim != 2:
        raise ValueError(
            f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})."
        )
    if a.stride(-1) != 1 or b.stride(-1) != 1:
        raise ValueError("`a`/`b` must be contiguous in the last dim.")
    if A_log.ndim != 1 or dt_bias.ndim != 1:
        raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
    if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
        raise ValueError("`A_log`/`dt_bias` must be contiguous.")
    if ssm_state_indices.ndim != 1:
        raise ValueError(
            f"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim})."
        )
    if not out.is_contiguous():
        raise ValueError("`out` must be contiguous.")

    dev = mixed_qkv.device
    if any(
        t.device != dev
        for t in (a, b, A_log, dt_bias, initial_state, out, ssm_state_indices)
    ):
        raise ValueError("All inputs must be on the same device.")

    B = mixed_qkv.shape[0]
    if a.shape[0] != B or b.shape[0] != B:
        raise ValueError(
            "Mismatched batch sizes: "
            f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}."
        )

View on GitHub (pinned to 0132848349)