sgl-project/sglang · error · ValueError

Mismatched batch sizes: mixed_qkv.shape[0]={B}, a.shape[0]={

Error message

Mismatched batch sizes: mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}.

What it means

In packed decode, B (num tokens) is defined by mixed_qkv.shape[0], and the per-token gate tensors a and b must both have exactly B rows. The wrapper raises with a diagnostic showing all three row counts when either gate tensor covers a different token set than the qkv tensor.

Source

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

    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
        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)})."

View on GitHub (pinned to 0132848349)

Solutions

1. Recompute or re-slice a and b from the same token layout as mixed_qkv: assert a.shape[0] == b.shape[0] == mixed_qkv.shape[0]
2. Build qkv and gates from the same forward of the same input batch (same projection buffer) so they cannot diverge

Example fix

# before
out, s = ...(mixed_qkv=qkv[pad_mask], a=a, b=b, ...)  # gates not masked
# after
out, s = ...(mixed_qkv=qkv[pad_mask].reshape(-1, D), a=a[pad_mask].reshape(-1, HV), b=b[pad_mask].reshape(-1, HV), ...)

B = mixed_qkv.shape[0]
assert a.shape[0] == B == b.shape[0], (B, a.shape[0], b.shape[0])

def same_batch(mixed_qkv, a, b) -> bool:
    B = mixed_qkv.shape[0]
    return a.shape[0] == B and b.shape[0] == B

When it happens

Trigger: a/b computed for a subset of tokens (padding stripped from qkv but not gates), or gates built from a previous batch after a schedule change; concat order differences between qkv projection and gate projection.

Common situations: Packing multiple ranks'/sequences' tokens where one path appends padding tokens and the other doesn't; ragged-batch reassembly bugs; caching gate projections across iterations.

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}).
• `mixed_qkv` must be contiguous in the last dim.
• `a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim=
• `a`/`b` must be contiguous in the last dim.

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