sgl-project/sglang · error · ValueError

`dt_bias` must have {HV * K} elements (got {dt_bias.numel()}

Error message

`dt_bias` must have {HV * K} elements (got {dt_bias.numel()}).

What it means

The KDA packed decode kernel requires the dt_bias tensor to contain exactly HV*K elements, one delta-time bias per (head, key-dim) pair. This is a shape contract check inside fused_recurrent_kda_packed_decode before launching the Triton kernel. A mismatch means the model's dt bias projection does not line up with the configured head/value geometry.

Source

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

        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 * 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}."
        )
    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}. "
            "KDA packed decode requires num_q_heads == num_k_heads and "

View on GitHub (pinned to 0132848349)

Solutions

1. Verify dt_bias.numel() == HV * K (e.g. num_v_heads * head_v_dim * head_k_dim) and reshape/reproject the checkpoint's dt bias to that size
2. Double-check the HV and K values you derive from mixed_qkv/config; they must match the checkpoint geometry
3. If the model genuinely has a smaller dt bias, expand/repeat it per key dim before calling the kernel

Example fix

// before
dt_bias = ckpt['dt_bias']  # [HV]
out = fused_recurrent_kda_packed_decode(..., dt_bias=dt_bias, ...)
// after
assert dt_bias.numel() == HV * K, (dt_bias.shape, HV, K)
dt_bias = dt_bias.repeat_interleave(K)  # [HV*K] if checkpoint stores per-head
out = fused_recurrent_kda_packed_decode(..., dt_bias=dt_bias, ...)

HV = num_v_heads * head_v_dim
assert dt_bias.numel() == HV * K, f'dt_bias {dt_bias.numel()} != HV*K {HV*K}'

def valid_kda_dt_bias(dt_bias: torch.Tensor, HV: int, K: int) -> bool:
    return dt_bias.numel() == HV * K

When it happens

Trigger: Calling fused_recurrent_kda_packed_decode with a dt_bias tensor whose numel() differs from HV*K, e.g. passing a per-head bias (HV elements) instead of per-(head, key) bias, or using HV/V/K values inconsistent with the checkpoint.

Common situations: Porting a KDA (Kimi Delta Attention) checkpoint whose head_dim or num_key_heads differs from the config used to build dt_bias; slicing dt_bias incorrectly when unpacking a fused projection; mismatched num_v_heads vs num_k_heads settings.

Related errors

• Invalid packed Q size {q_dim}: must be divisible by K={K}. K
• The batch size is expected to be 1 rather than {q.shape[0]}
• `mixed_qkv` must be 2D (got ndim={mixed_qkv.ndim}).
• pairs must be a torch.Tensor of shape [N, 2]
• f"recurrent_kda needs a [N, HV, V, K] state pool; got shape

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