sgl-project/sglang · error · ValueError

`a` must have shape [B, HV*K] with HV={HV}, K={K} (got a.sha

Error message

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

What it means

The KDA decay-input tensor `a` must be packed as [B, HV*K]: one row per batch item with all key-dim decays for every value head concatenated. validate_packed_decode_inputs derives HV and K from initial_state and checks a.shape[1] == HV*K, raising when the packed layout doesn't match.

Source

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

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

View on GitHub (pinned to 0132848349)

Solutions

Flatten the head dims: a = a.reshape(B, -1) so shape[1] == HV*K
Verify HV and K in initial_state match the model's num_v_heads and head_k_dim used to produce `a`
Add a shape assert in your model forward before calling the fused kernel

Example fix

// before
a = a_proj(x)  # [B, HV, K]
helion_fused_recurrent_kda_packed_decode(qkv, a, b, ...)
// after
a = a_proj(x).reshape(a.size(0), -1)  # [B, HV*K]
helion_fused_recurrent_kda_packed_decode(qkv, a, b, ...)

Defensive patterns

Strategy: validation

Validate before calling

HV, V, K = initial_state.shape[-3:]
assert a.reshape(a.size(0), -1).shape[1] == HV * K
a = a.reshape(a.size(0), -1)

Type guard

def valid_packed_a(a: torch.Tensor, hv: int, k: int) -> bool:
    return a.ndim == 2 and a.shape[1] == hv * k

Prevention

Standardize on [B, HV*K] packing at the projection site
Derive HV/K from the same source (initial_state) everywhere

When it happens

Trigger: Calling helion_fused_recurrent_kda_packed_decode with `a` shaped [B, HV, K] (3D), [B, K] (single head), or with HV/K taken from a mismatched initial_state.

Common situations: Forgetting to .view(B, HV*K) / .reshape(B, -1) the conv/gate projection output before the fused decode; using a different num_v_heads in the projection than in the state pool; test fixtures built with the wrong packing.

Related errors

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