sgl-project/sglang · error · ValueError

Unexpected a shape for dense: {a.shape}

Error message

Unexpected a shape for dense: {a.shape}

What it means

For dense (non-varlen) decode, _normalize_kda_a accepts `a` only as (N, HV*K) 2D, (N, HV, K) 3D (unsqueezed to (N,1,HV,K)), or exactly (N, 1, HV, K) 4D. Other shapes raise ValueError('Unexpected a shape for dense: ...').

Source

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

    """
    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,
    b: torch.Tensor,
    initial_state_source: torch.Tensor,
    initial_state_indices: torch.Tensor,
    cu_seqlens: Optional[torch.Tensor] = None,
    scale: Optional[float] = None,
    use_qk_l2norm_in_kernel: bool = True,
    softplus_beta: float = 1.0,
    softplus_threshold: float = 20.0,
) -> torch.Tensor:

View on GitHub (pinned to 0132848349)

Solutions

1. If a is (1, N, HV, K), squeeze dim 0 to get (N, HV, K)
2. Verify a.shape == (N, HV, K) or (N, HV*K) with the same N, HV, K passed to the kernel
3. Centralize the reshape in one helper so dense and varlen paths can't diverge

Example fix

# before (dense): a shape (1, N, HV, K) -> ValueError
# after
a = a.squeeze(0)  # -> (N, HV, K), then kernel unsqueezes to (N,1,HV,K)
update = cutedsl_fused_sigmoid_gating_kda_update(..., a=a, ...)

assert a.shape in [(N, HV*K), (N, HV, K), (N, 1, HV, K)], f'a shape {a.shape}'

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

When it happens

Trigger: Calling the KDA fused update in dense mode with a varlen-style tensor like (1, N, HV, K) (extra leading batch dim), or an N dimension that doesn't match the declared token/batch count.

Common situations: Reusing a varlen-prepared `a` tensor in the dense path; batch schedulers that prepend a size-1 dimension during scheduling; head-count mismatch between `a` layout (HV*K fused vs HV,K) and the passed HV/K values.

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 varlen: {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/abf71a600bca75c4. Report an issue: GitHub.