sgl-project/sglang · error · ValueError

`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim=

Error message

`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim}).

What it means

In the packed decode path, the gating log-signals a and b must each be 2D of shape (num_tokens, HV) — one scalar gate per token per value head. The wrapper raises when either has a different rank, such as the common (B, T, HV) layout or per-head (B, H, T) head-first layout.

Source

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

    mixed_qkv: torch.Tensor,
    a: torch.Tensor,
    b: torch.Tensor,
    A_log: torch.Tensor,
    dt_bias: torch.Tensor,
    scale: float,
    initial_state: torch.Tensor,
    out: torch.Tensor,
    ssm_state_indices: torch.Tensor,
    use_qk_l2norm_in_kernel: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    if mixed_qkv.ndim != 2:
        raise ValueError(
            f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})."
        )
    if mixed_qkv.stride(-1) != 1:
        raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
    if a.ndim != 2 or b.ndim != 2:
        raise ValueError(
            f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})."
        )
    if a.stride(-1) != 1 or b.stride(-1) != 1:
        raise ValueError("`a`/`b` must be contiguous in the last dim.")
    if A_log.ndim != 1 or dt_bias.ndim != 1:
        raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
    if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
        raise ValueError("`A_log`/`dt_bias` must be contiguous.")
    if ssm_state_indices.ndim != 1:
        raise ValueError(
            f"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim})."
        )
    if not out.is_contiguous():
        raise ValueError("`out` must be contiguous.")

    dev = mixed_qkv.device
    if any(
        t.device != dev

View on GitHub (pinned to 0132848349)

Solutions

  1. Reshape: a = a.view(-1, HV), b = b.view(-1, HV) (or .squeeze(1) for decode's T=1)
  2. Split the 2*HV gate projection with .chunk(2, dim=-1) and flatten tokens

Example fix

# before
a, b = a_log[..., 0], a_log[..., 1]  # (B, 1, HV) each
# after
a = a_log[..., 0].reshape(-1, a_log.shape[-1])
b = a_log[..., 1].reshape(-1, a_log.shape[-1])
Defensive patterns

Strategy: validation

Validate before calling

assert a.ndim == 2 and b.ndim == 2, (a.shape, b.shape)
a = a.reshape(-1, a.shape[-1]); b = b.reshape(-1, b.shape[-1])

Type guard

def gates_are_2d(*ts: torch.Tensor) -> bool:
    return all(t.ndim == 2 for t in ts)

Prevention

When it happens

Trigger: Passing a/b shaped (B, 1, HV) from a decode step without squeezing the time dim, or (B, HV, T) head-first tensors from fla-style code.

Common situations: Adapter code bridging fla's head_first conventions to sglang's packed decode; forgetting .squeeze(1)/.view(B, -1) after computing a and b from a Linear projection of shape (B, T, 2*HV).

Related errors

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