sgl-project/sglang · error · ValueError

`initial_state` must be a 4D tensor (got ndim={initial_state

Error message

`initial_state` must be a 4D tensor (got ndim={initial_state.ndim}).

What it means

initial_state is the recurrent SSM state cache slice and must be 4D of shape (num_states, HV, V, K) (or batched (B, HV, V, K)); the last three dims are parsed as HV, V, K. The wrapper raises when the state pool tensor has any other rank, e.g. a flattened pool or a per-layer stacked 5D tensor.

Source

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

    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()})."
        )
    if out.shape != (B, 1, HV, V):
        raise ValueError(
            f"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)})."
        )

View on GitHub (pinned to 0132848349)

Solutions

1. Index/select down to 4D: initial_state = cache[layer_idx] giving (N, HV, V, K)
2. If states are stored per-request per-head 3D, re-add the leading dim via unsqueeze(0) and index with ssm_state_indices

Example fix

# before
initial_state = mamba_cache  # (L, N, HV, V, K)
# after
initial_state = mamba_cache[layer_idx]  # (N, HV, V, K)

assert initial_state.ndim == 4, initial_state.shape
initial_state = initial_state if initial_state.ndim == 4 else initial_state[0]

def state_4d(s: torch.Tensor) -> bool:
    return s.ndim == 4

When it happens

Trigger: Passing the entire mamba cache of shape (num_layers, N, HV, V, K); passing a flattened (N, HV*V*K) buffer; passing a single sequence's state as 3D.

Common situations: Slicing the wrong dim of a layered hybrid-attention state cache; writing a custom backend that stores states packed differently than the kernel contract.

Related errors

• `ssm_state_indices` must be 1D for packed decode (got ndim={
• `ssm_state_indices` must have shape [B] (got {tuple(ssm_stat
• `mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim}).
• `a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim=
• `A_log`/`dt_bias` must be 1D tensors.

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