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 have exactly shape [B] where B is mixed_qkv.shape[0]; the validator reports both got and expected shapes. A length mismatch means state-slot indices don't cover the batch (or over-cover it).

Source

Thrown at python/sglang/kernels/ops/attention/helion/kda_decode.py:275

            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]}, "
            f"b.shape[0]={b.shape[0]}."
        )
    if ssm_state_indices.shape[0] != B:
        raise ValueError(
            f"`ssm_state_indices` must have shape [B] "
            f"(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 not _is_power_of_two(K) or not _is_power_of_two(V):
        raise ValueError(
            "Helion KDA decode requires power-of-two key and value head "
            f"dimensions (got K={K}, V={V})."
        )
    if a.shape[1] != HV * K:
        raise ValueError(

View on GitHub (pinned to 0132848349)

Solutions

Recompute/filter ssm_state_indices to match the current batch length: idx[keep] or idx[:B]
Ensure the scheduler updates state indices whenever the batch changes
Reshape any [B,1] tensor to [B]

Example fix

# before
out = decode(qkv, a, b, ..., ssm_state_indices=all_idx, ...)  # len(all_idx) > B
# after
out = decode(qkv, a, b, ..., ssm_state_indices=all_idx[keep], ...)  # len == B

Defensive patterns

Strategy: validation

Validate before calling

B = mixed_qkv.shape[0]
assert ssm_state_indices.shape == (B,), f'indices must be [{B}]'

Type guard

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

Prevention

Recompute state indices whenever the running batch changes
Include a shape assert right before the kernel call

When it happens

Trigger: Passing ssm_state_indices with more entries than batch rows (e.g. indices from a larger scheduler batch) or wrong shape like [B,1].

Common situations: Indices cached from a previous, larger batch; not filtering indices when finished requests are removed; reshaping errors elsewhere.

Related errors

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