sgl-project/sglang · error · ValueError

`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={t

Error message

`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)}).

What it means

The head-count is derived from the state tensor: HV = initial_state.shape[-3]. The gate tensors a and b must then have exactly shape (B, HV) so each token has one gate scalar per value head. The wrapper raises when their head dim disagrees with the state cache's head count (a.shape[1] != HV or b.shape[1] != HV).

Source

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

    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)})."
        )
    if A_log.numel() != HV or dt_bias.numel() != HV:
        raise ValueError(
            f"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()})."
        )
    if out.shape != (B, 1, HV, V):
        raise ValueError(
            f"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)})."
        )

    qkv_dim = mixed_qkv.shape[1]
    qk_dim = qkv_dim - HV * V
    if qk_dim <= 0 or qk_dim % 2 != 0:
        raise ValueError(
            f"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}."
        )
    q_dim = qk_dim // 2

View on GitHub (pinned to 0132848349)

Solutions

Align head counts: assert a.shape[1] == b.shape[1] == initial_state.shape[-3] and fix whichever producer (gate projection out_features or cache allocation) is wrong
Under TP, shard both the state cache and the gate projection by the same HV factor (typically hv // tp_size)

Example fix

# before
# a, b: (B, K_heads); state: (N, HV, V, K)
# after
# project gates with out_features=2*HV so a, b: (B, HV); state: (N, HV, V, K)

Defensive patterns

Strategy: validation

Validate before calling

HV = initial_state.shape[-3]
assert a.shape == (a.shape[0], HV) and b.shape == (b.shape[0], HV), (a.shape, b.shape, HV)

Type guard

def gate_heads_match(a, b, initial_state) -> bool:
    HV = initial_state.shape[-3]
    return a.ndim == 2 and b.ndim == 2 and a.shape[1] == HV == b.shape[1]

Prevention

Derive gate projection out_features from the same num_v_heads used to allocate the state cache
Under TP, shard gates and state cache with an identical head split

When it happens

Trigger: Model config changed num_v_heads (e.g. 32) but the state cache was allocated with a different count (e.g. 8 for KDA/Qwen3-Next ratios); gates computed with num_k_heads instead of num_v_heads; TP sharding of gates not matching state sharding.

Common situations: Tensor-parallel setups where value heads are sharded but gate projections are replicated (or vice versa); loading a checkpoint with different head counts than the runtime config; unit tests with hand-made shapes that don't match the model.

Related errors

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