sgl-project/sglang · error · ValueError

The number of initial states is expected to be equal to the

Error message

The number of initial states is expected to be equal to the number of input sequences, i.e., {len(cu_seqlens) - 1} rather than {initial_state_indices.shape[0]}.

What it means

In chunk_gated_delta_rule (varlen path with cu_seqlens), the optional initial_state_indices tensor must contain one index per input sequence. The check compares initial_state_indices.shape[0] against len(cu_seqlens) - 1, which is the number of variable-length sequences described by the cumulative sequence lengths tensor. It throws when the caller passes a state pool index tensor sized for a different batch than the one described by cu_seqlens.

Source

Thrown at python/sglang/kernels/ops/attention/fla/chunk.py:242

        )
    # if not head_first and q.shape[1] < q.shape[2]:
    #     warnings.warn(
    #         f"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). "
    #         "This may indicate the inputs were passed in head-first format [B, H, T, ...] "
    #         "when head_first=False was specified. "
    #         "Please verify your input tensor format matches the expected shape [B, T, H, ...]."
    #     )
    if cu_seqlens is not None:
        if q.shape[0] != 1:
            raise ValueError(
                f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
                f"Please flatten variable-length inputs before processing."
            )
        if (
            initial_state_indices is not None
            and initial_state_indices.shape[0] != len(cu_seqlens) - 1
        ):
            raise ValueError(
                f"The number of initial states is expected to be equal to the number of input sequences, "
                f"i.e., {len(cu_seqlens) - 1} rather than {initial_state_indices.shape[0]}."
            )
    if scale is None:
        scale = k.shape[-1] ** -0.5
    o, h = ChunkGatedDeltaRuleFunction.apply(
        q,
        k,
        v,
        g,
        beta,
        scale,
        initial_state,
        initial_state_indices,
        cu_seqlens,
        use_qk_l2norm_in_kernel,
    )
    if head_first:

View on GitHub (pinned to 0132848349)

Solutions

  1. Verify len(cu_seqlens) - 1 equals initial_state_indices.shape[0] before the call and rebuild initial_state_indices from the current batch's request indices
  2. Check how cu_seqlens is constructed (torch.cumsum of seq_lens with a leading 0) and confirm it matches the sequences the indices refer to
  3. If you don't need prior states, pass initial_state_indices=None
  4. Audit scheduler/batching code that slices inputs but not the state-pool index tensor

Example fix

# before
o, h = chunk_gated_delta_rule(q, k, v, g, cu_seqlens=cu_seqlens, initial_state_indices=idx)
# after
assert initial_state_indices.shape[0] == len(cu_seqlens) - 1, (
    f"{initial_state_indices.shape[0]} states vs {len(cu_seqlens)-1} sequences")
o, h = chunk_gated_delta_rule(q, k, v, g, cu_seqlens=cu_seqlens, initial_state_indices=idx)
Defensive patterns

Strategy: validation

Validate before calling

n_seqs = len(cu_seqlens) - 1
assert initial_state_indices is None or initial_state_indices.shape[0] == n_seqs, (
    initial_state_indices.shape[0], n_seqs)

Type guard

def valid_state_indices(idx: torch.Tensor, cu_seqlens: torch.Tensor) -> bool:
    return idx is None or (idx.ndim == 1 and idx.shape[0] == len(cu_seqlens) - 1)

Prevention

When it happens

Trigger: Calling chunk_gated_delta_rule with both cu_seqlens and initial_state_indices where the number of packed sequences (len(cu_seqlens)-1) differs from initial_state_indices.shape[0]; e.g. cu_seqlens=[0, 5, 10] (2 sequences) but initial_state_indices has 1 or 3 entries.

Common situations: Hybrid linear-attention models (GatedDeltaNet/Qwen3-Next style) prefilling a batch whose request set changed between when the state indices were built and when the kernel was invoked; passing full-batch state indices while slicing the varlen input; off-by-one when constructing cu_seqlens or reusing indices from a previous scheduler batch.

Related errors


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