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

The validator derives batch size B from mixed_qkv.shape[0] and requires a.shape[0] and b.shape[0] to equal it, since each batch row has corresponding gating tensors. A mismatch means inconsistent tensor batches.

Source

Thrown at python/sglang/kernels/ops/attention/helion/kda_decode.py:269

    device = mixed_qkv.device
    if any(
        tensor.device != device
        for tensor 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]}, "
            f"b.shape[0]={b.shape[0]}."
        )
    if ssm_state_indices.shape[0] != B:
        raise ValueError(
            f"`ssm_state_indices` must have shape [B] "
            f"(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 not _is_power_of_two(K) or not _is_power_of_two(V):

View on GitHub (pinned to 0132848349)

Solutions

Re-align a and b to the current batch: recompute them from the same request set as mixed_qkv
Check batch filtering logic (finished-request removal) updates all decode tensors consistently

Example fix

# before
out = decode(qkv[keep], a, b, ...)  # a,b still full batch
# after
a, b = a[keep], b[keep]
out = decode(qkv[keep], a, b, ...)

Defensive patterns

Strategy: validation

Validate before calling

B = mixed_qkv.shape[0]
assert a.shape[0] == B and b.shape[0] == B, 'batch mismatch'

Type guard

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

Prevention

Filter all decode tensors with the same keep-mask atomically
Derive batched tensors from one canonical batch structure

When it happens

Trigger: Calling packed decode where a or b has fewer/more rows than mixed_qkv (e.g. B=8 qkv with B=4 a).

Common situations: Dropped/finished requests removed from one tensor but not others; slicing a/b by a different batch dimension after scheduling.

Related errors

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