sgl-project/sglang · error · ValueError

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

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.

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)

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

Example fix

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

Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

Flatten scheduler-provided location indices at the model/kernel boundary
Treat every packed-decode input as (num_tokens, …) and reshape all of them together

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.

Related errors

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