sgl-project/sglang · error · ValueError

Unexpected a shape for varlen: {a.shape}

Error message

Unexpected a shape for varlen: {a.shape}

What it means

For variable-length (varlen) decode, _normalize_kda_a accepts the gating tensor `a` only as (N, HV*K) 2D, (N, HV, K) 3D, or (1, N, HV, K) 4D. Any other shape raises ValueError('Unexpected a shape for varlen: ...').

Source

Thrown at python/sglang/kernels/ops/attention/cutedsl_kda.py:1373

        )
    return dt_bias.reshape(HV, K).contiguous()


def _normalize_kda_a(a, *, is_varlen_decode, N, HV, K):
    """Normalize `a` to match the compile-time shape expected by the kernel.

    varlen kernel compiled shape: (N, HV, K)  -- 3D
    dense kernel compiled shape:  (N, 1, HV, K) -- 4D
    """
    if is_varlen_decode:
        # Target: (N, HV, K) -- 3D
        if a.dim() == 2 and a.shape == (N, HV * K):
            return a.view(N, HV, K)
        if a.dim() == 3 and a.shape == (N, HV, K):
            return a  # already correct
        if a.dim() == 4 and a.shape == (1, N, HV, K):
            return a.squeeze(0)  # remove leading dim
        raise ValueError(f"Unexpected a shape for varlen: {a.shape}")
    else:
        # Target: (N, 1, HV, K) -- 4D
        if a.dim() == 2 and a.shape == (N, HV * K):
            return a.view(N, 1, HV, K)
        if a.dim() == 3 and a.shape == (N, HV, K):
            return a.unsqueeze(1)
        if a.dim() == 4 and a.shape == (N, 1, HV, K):
            return a
        raise ValueError(f"Unexpected a shape for dense: {a.shape}")


def cutedsl_fused_sigmoid_gating_kda_update(
    A_log: torch.Tensor,
    dt_bias: torch.Tensor,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    a: torch.Tensor,

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape to (N, HV, K): if dense layout (N,1,HV,K), call a.squeeze(1)
2. Ensure N equals a.shape[0] (token count) exactly; recompute N from the same source as the input tensors
3. Add a pre-call assert a.dim() in (2,3) and a.shape[0] == N

Example fix

# before (varlen): a shape (N, 1, HV, K) -> ValueError
# after
a = a.squeeze(1) if a.dim() == 4 else a  # -> (N, HV, K)
update = cutedsl_fused_sigmoid_gating_kda_update(..., a=a, ...)

assert a.dim() in (2, 3) or (a.dim() == 4 and a.shape[0] == 1), f'a shape {a.shape}'
assert a.shape[0] == N, f'a token dim {a.shape[0]} != N {N}'

def is_valid_varlen_a(a, N: int, HV: int, K: int) -> bool:
    return a.shape in [(N, HV*K), (N, HV, K), (1, N, HV, K)]

When it happens

Trigger: Calling the KDA fused update with is_varlen_decode=True and an `a` tensor that is (N, 1, HV, K) (the dense layout, missing squeeze of dim 1), or a per-request 3D tensor without the token dim N, or wrong N due to a cu_seqlen mismatch.

Common situations: Switching the same caller code between dense and varlen decode paths without reshaping `a`; ragged-batch schedulers producing `a` with token counts that disagree with the N argument computed elsewhere.

Related errors

• Unexpected A_log shape: {A_log.shape}; expected numel={HV}
• Unexpected dt_bias shape: {dt_bias.shape}; expected numel={H
• Unexpected a shape for dense: {a.shape}
• num_heads must be divisible by num_epi_subtiles
• num_heads // num_epi_subtiles must be divisible by 4 (FMA un

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