sgl-project/sglang · error · ValueError

The batch size is expected to be 1 rather than {q.shape[0]}

Error message

The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`.Please flatten variable-length inputs before processing.

What it means

fused_recurrent_kda in the KDA (Kimi Delta Attention) FLA kernel requires that when cu_seqlens (variable-length cumulative sequence lengths) is passed, the q tensor's batch dimension must be 1. Variable-length sequences must be flattened (packed) into a single [1, total_tokens, ...] tensor with cu_seqlens describing boundaries.

Source

Thrown at python/sglang/kernels/ops/attention/fla/kda.py:154

    return o, final_state


def fused_recurrent_kda(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor = None,
    scale: float = None,
    initial_state: torch.Tensor = None,
    inplace_final_state: bool = True,
    use_qk_l2norm_in_kernel: bool = True,
    cu_seqlens: torch.LongTensor | None = None,
    # ssm_state_indices: torch.LongTensor | None = None,
    **kwargs,
) -> tuple[torch.Tensor, torch.Tensor]:
    if cu_seqlens is not None and 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 scale is None:
        scale = k.shape[-1] ** -0.5

    o, final_state = fused_recurrent_kda_fwd(
        q=q.contiguous(),
        k=k.contiguous(),
        v=v.contiguous(),
        g=g.contiguous(),
        beta=beta.contiguous(),
        scale=scale,
        initial_state=initial_state,
        inplace_final_state=inplace_final_state,
        cu_seqlens=cu_seqlens,
        # ssm_state_indices=ssm_state_indices,
        use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,

View on GitHub (pinned to 0132848349)

Solutions

  1. Flatten q/k/v to shape [1, total_tokens, ...] and pass cu_seqlens describing per-sequence boundaries
  2. Ensure batch dims are collapsed: q.reshape(1, -1, q.shape[-1])
  3. If you need per-sequence states, use ssm_state_indices / loop over sequences instead of a batch dim

Example fix

// before
q = q  # [B, T, D] with cu_seqlens set
// after
q = q.reshape(1, -1, q.shape[-1])
k = k.reshape(1, -1, k.shape[-1])
v = v.reshape(1, -1, v.shape[-1])
out, _ = fused_recurrent_kda(q, k, v, cu_seqlens=cu_seqlens)
Defensive patterns

Strategy: validation

Validate before calling

assert cu_seqlens is None or q.shape[0] == 1, 'flatten varlen inputs to [1, total_T, D] before fused_recurrent_kda'

Prevention

When it happens

Trigger: Calling fused_recurrent_kda(q, k, ...) with cu_seqlens set and q.shape[0] != 1, i.e. passing a batched [B, T, ...] tensor instead of a flattened varlen tensor.

Common situations: Batching multiple sequences of different lengths into a padded [B, T, ...] tensor and passing cu_seqlens alongside; migrating from chunked KDA kernels that accept batched inputs.

Related errors


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