sgl-project/sglang · error · ValueError

Unexpected A_log shape: {A_log.shape}; expected numel={HV}

Error message

Unexpected A_log shape: {A_log.shape}; expected numel={HV}

What it means

The KDA fused update helper _normalize_A_log requires the A_log gating parameter to contain exactly HV elements (one scalar per head) so it can be reshaped to (HV,). A mismatched element count raises ValueError with the expected numel.

Source

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

        HV=HV,
        use_initial_state=True,
        use_qk_l2norm=True,
        stream=stream,
    )

    _compiled_kernels[key] = compiled_kernel
    logger.info(
        "CuTe DSL KDA kernel compiled: "
        f"N={N}, H={H}, HV={HV}, K={K}, V={V}, pool_size={pool_size}, "
        f"pool_strides={tuple(h0_source.stride())}, "
        f"small_batch={use_small_batch}, varlen={is_varlen_decode}"
    )
    return compiled_kernel


def _normalize_A_log(A_log: torch.Tensor, HV: int) -> torch.Tensor:
    if A_log.numel() != HV:
        raise ValueError(f"Unexpected A_log shape: {A_log.shape}; expected numel={HV}")
    return A_log.reshape(HV).contiguous()


def _normalize_dt_bias(dt_bias: torch.Tensor, HV: int, K: int) -> torch.Tensor:
    if dt_bias.numel() != HV * K:
        raise ValueError(
            f"Unexpected dt_bias shape: {dt_bias.shape}; expected numel={HV * K}"
        )
    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:

View on GitHub (pinned to 0132848349)

Solutions

1. Set HV to match A_log.numel() if the weights are the source of truth (one gate per head)
2. If A_log has shape (HV, K) due to checkpoint layout, reduce over K or fix the loader so A_log is per-head
3. Verify with assert A_log.numel() == HV before calling the fused update

Example fix

# before: A_log shape (HV, K) -> HV*K elements -> ValueError
update = cutedsl_fused_sigmoid_gating_kda_update(A_log, ...)
# after
assert A_log.numel() == HV
update = cutedsl_fused_sigmoid_gating_kda_update(A_log.reshape(HV), ...)

assert A_log.numel() == HV, f'A_log numel {A_log.numel()} != HV {HV}'

def is_valid_a_log(A_log, HV: int) -> bool:
    return A_log.numel() == HV

When it happens

Trigger: Calling cutedsl_fused_sigmoid_gating_kda_update with A_log whose numel != HV — e.g. per-head-per-key layout of shape (HV, K) from a checkpoint, or a scalar/expanded tensor when the model uses fused heads.

Common situations: Loading a KDA (Kimi Delta Attention) checkpoint whose A_log was stored with a different head layout (e.g. (num_kv_heads*K,) vs fused (HV,)); configuring the kernel with the wrong HV value (unfused vs fused head count) for the same weights.

Related errors

• Unexpected dt_bias shape: {dt_bias.shape}; expected numel={H
• Unexpected initial_state_source shape: {initial_state_source
• Unexpected a shape for varlen: {a.shape}
• Unexpected a shape for dense: {a.shape}
• KDA cutedsl: safe_gate (lower_bound) not yet supported

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