sgl-project/sglang · error · ValueError

`ssm_state_indices` must have shape [B] (got {tuple(ssm_stat

Error message

`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},)).

What it means

ssm_state_indices must be exactly shape [B] where B is the packed token count from mixed_qkv.shape[0]. This check (shape[0] != B) catches index tensors whose first dim disagrees even if they are 1D — e.g. sized for a different batch or containing one index per sequence rather than per token.

Source

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

        )
    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]}."
        )
    if ssm_state_indices.shape[0] != B:
        raise ValueError(
            f"`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},))."
        )

    if initial_state.ndim != 4:
        raise ValueError(
            f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
        )
    if initial_state.stride(-1) != 1:
        raise ValueError("`initial_state` must be contiguous in the last dim.")
    HV, V, K = initial_state.shape[-3:]
    if a.shape[1] != HV or b.shape[1] != HV:
        raise ValueError(
            f"`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)})."
        )
    if A_log.numel() != HV or dt_bias.numel() != HV:
        raise ValueError(
            f"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()})."
        )

View on GitHub (pinned to 0132848349)

Solutions

Regenerate indices per packed batch: torch.full((B,), -1) filled with cache slots for tokens carrying state; assert ssm_state_indices.shape[0] == mixed_qkv.shape[0]
Keep index construction adjacent to qkv packing so they always share the same token list

Example fix

# before
idx = cache_loc  # per-sequence
# after
idx = cache_loc.reshape(-1)
assert idx.shape[0] == mixed_qkv.shape[0]

Defensive patterns

Strategy: validation

Validate before calling

assert ssm_state_indices.shape[0] == mixed_qkv.shape[0], (
    ssm_state_indices.shape, mixed_qkv.shape)

Type guard

def indices_match_batch(idx, mixed_qkv) -> bool:
    return idx.ndim == 1 and idx.shape[0] == mixed_qkv.shape[0]

Prevention

Use one index per decode token, -1 for tokens without state
Rebuild indices whenever the packed token set changes (continuous batching)

When it happens

Trigger: Passing per-sequence indices (length num_seqs) instead of per-token indices (length num_tokens) for decode where each token is a sequence; stale indices from a previous batch; concatenating rank indices in a different order/count than the qkv tokens.

Common situations: Decode batches where num_tokens == num_seqs making the bug invisible in tests but failing under continuous batching with mixed lengths; DP gather order mismatches.

Related errors

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