sgl-project/sglang · error · ValueError

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

Error message

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

What it means

The KDA gate/input tensor `b` must be packed as [B, HV]: one scalar gate per value head per batch row. validate_packed_decode_inputs checks b.shape[1] == HV (HV inferred from initial_state) and raises when b carries per-key or per-token values instead.

Source

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

    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):
        raise ValueError(
            "Helion KDA decode requires power-of-two key and value head "
            f"dimensions (got K={K}, V={V})."
        )
    if a.shape[1] != HV * K:
        raise ValueError(
            f"`a` must have shape [B, HV*K] with HV={HV}, K={K} "
            f"(got a.shape={tuple(a.shape)})."
        )
    if b.shape[1] != HV:
        raise ValueError(
            f"`b` must have shape [B, HV] with HV={HV} (got b.shape={tuple(b.shape)})."
        )
    if A_log.numel() != HV:
        raise ValueError(f"`A_log` must have {HV} elements (got {A_log.numel()}).")
    if dt_bias.numel() != HV * K:
        raise ValueError(
            f"`dt_bias` must have {HV * K} elements (got {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}."

View on GitHub (pinned to 0132848349)

Solutions

Reshape b to [B, HV]: b = b_proj_out.reshape(B, -1) and confirm the projection width equals num_v_heads
Double-check that HV comes from initial_state.shape[-3] and matches the model's num_v_heads, not num_q_heads
Write a shape check in the model forward: assert b.shape == (B, num_v_heads)

Example fix

// before
b = gates  # [B, HV*K] accidentally
// after
b = gates.reshape(B, -1)
assert b.shape[1] == initial_state.shape[-3]

Defensive patterns

Strategy: validation

Validate before calling

HV = initial_state.shape[-3]
b = b.reshape(b.size(0), -1)
assert b.shape[1] == HV, (b.shape, HV)

Type guard

def valid_packed_b(b: torch.Tensor, hv: int) -> bool:
    return b.ndim == 2 and b.shape[1] == hv

Prevention

Never reuse the `a` packing width for `b`
Log num_v_heads once at startup and use it for all gate widths

When it happens

Trigger: Passing `b` with shape [B, HV*K] (per-key gates), [B, 1], or [B, H] where H is the query-head count rather than the value-head count.

Common situations: Reusing the packing used for `a` for `b` as well; confusing num_k_heads with num_v_heads in GQA-style KDA models (H != HV); slicing a projection output with the wrong width.

Related errors

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