sgl-project/sglang · error · ValueError

Unsupported input shape {g.shape}, which should be (B, T, H,

Error message

Unsupported input shape {g.shape}, which should be (B, T, H, D) if `head_first=False` or (B, H, T, D) otherwise

What it means

chunk_local_cumsum computes log-decay cumulative sums and accepts g of shape (B, T, H) when head_first=False (optionally with an extra dim for two-component decays) or (B, H, T) when head_first=True. The trailing else branch raises when g.ndim does not match either supported layout for the chosen head_first flag.

Source

Thrown at python/sglang/kernels/ops/attention/fla/cumsum.py:290

            scale=scale,
            cu_seqlens=cu_seqlens,
            head_first=head_first,
            output_dtype=output_dtype,
            chunk_indices=chunk_indices,
        )
    elif len(g.shape) == 4:
        return chunk_local_cumsum_vector(
            g=g,
            chunk_size=chunk_size,
            reverse=reverse,
            scale=scale,
            cu_seqlens=cu_seqlens,
            head_first=head_first,
            output_dtype=output_dtype,
            chunk_indices=chunk_indices,
        )
    else:
        raise ValueError(
            f"Unsupported input shape {g.shape}, "
            f"which should be (B, T, H, D) if `head_first=False` "
            f"or (B, H, T, D) otherwise"
        )

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape/transpose g to (B, T, H) (or (B, T, H, dim)) and call with head_first=False
2. If your tensor is (B, H, T), either pass head_first=True or permute with g.transpose(1, 2).contiguous()
3. Print g.shape right before the call and compare with the branch conditions in cumsum.py to see which layout the function expects

Example fix

# before
g_cumsum = chunk_local_cumsum(g, chunk_size, head_first=False)  # g is (B, H, T)
# after
g_cumsum = chunk_local_cumsum(g.transpose(1, 2).contiguous(), chunk_size, head_first=False)

assert g.ndim == 3 and (head_first and g.shape[2] == 1 or not head_first and g.shape[2] == H), g.shape
# head_first=False expects (B, T, H); transpose if needed
g = g.transpose(1, 2).contiguous() if head_first and g.shape[1] != T else g

def cumsum_layout_ok(g: torch.Tensor, head_first: bool) -> bool:
    return g.ndim == 3 and ((not head_first) or True) and g.ndim in (3, 4)

When it happens

Trigger: Calling chunk_local_cumsum with head_first=False but a 3D g that is actually (B, H, T), or with an ndim outside the accepted set (e.g. 2D or 5D); passing head_first=True with a (B, T, H) tensor.

Common situations: Porting code from the fla library where head_first defaulted to True into sglang where layouts are B,T,H-first; forgetting to transpose gate logits produced by an nn.Linear of shape (B, T, num_heads); passing per-head decay plus extra dims the kernel doesn't support.

Related errors

• The number of initial states is expected to be equal to the
• `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.
• `ssm_state_indices` must be 1D for packed decode (got ndim={

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