sgl-project/sglang · error · ValueError

`A_log` and `dt_bias` must have {HV} elements (got A_log.num

Error message

`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()}).

What it means

A_log and dt_bias must each contain exactly HV elements (one per value head), where HV comes from initial_state.shape[-3]. The wrapper checks numel and raises when the parameter vectors were sized for a different head configuration — e.g. sharded or full-head versions mismatched with the state cache.

Source

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

        )
    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
    if q_dim % K != 0:
        raise ValueError(f"Invalid packed Q size {q_dim}: must be divisible by K={K}.")
    H = q_dim // K
    if H <= 0 or HV % H != 0:

View on GitHub (pinned to 0132848349)

Solutions

  1. Shard consistently: A_log = A_log_full[hv_start:hv_end] per TP rank so numel == HV
  2. If fused, split: A_log, dt_bias = param.chunk(2) then verify each has HV elements

Example fix

# before
A_log_rank = A_log_full  # numel = num_heads, but HV = num_heads//tp
# after
A_log_rank = A_log_full.chunk(tp, dim=-1)[rank].contiguous()  # numel == HV
Defensive patterns

Strategy: validation

Validate before calling

HV = initial_state.shape[-3]
assert A_log.numel() == HV == dt_bias.numel(), (A_log.numel(), dt_bias.numel(), HV)

Type guard

def param_heads_match(A_log, dt_bias, HV: int) -> bool:
    return A_log.numel() == HV and dt_bias.numel() == HV

Prevention

When it happens

Trigger: Passing full-model A_log (num_heads) with a TP-sharded state cache (num_heads/tp); passing dt_bias of shape (HV*2) from a fused gate projection without splitting; loading params from a checkpoint with a different head count.

Common situations: Tensor-parallel rank 1..N passing unsharded parameters; fusing A_log and dt_bias into one vector and forgetting to chunk; Heterogeneous head configs (Qwen3-Next KDA linear head ratios).

Related errors

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