sgl-project/sglang · error · ValueError

f"SGLang KDA state must be [B,H,V,K] with (V,K)={expected},

Error message

f"SGLang KDA state must be [B,H,V,K] with (V,K)={expected}, got {tuple(state.shape)}"

What it means

_to_nvidia_kda_state_layout validates that SGLang-format state is 4-D [B,H,V,K] with the last two dims exactly (head_v_dim, head_k_dim) before transposing to the vendor [B,H,K,V] layout. Any other shape raises ValueError.

Source

Thrown at python/sglang/srt/layers/attention/linear/kernels/kda_nvidia.py:54

from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
    LinearAttnKernelBase,
)

logger = logging.getLogger(__name__)

_BUCKETS = (2048, 4096, 8192, 16384)
_MAX_NVIDIA_KDA_BATCH = 8


def _to_nvidia_kda_state_layout(
    state: torch.Tensor, *, head_k_dim: int, head_v_dim: int
) -> torch.Tensor:
    """Materialize SGLang [B,H,V,K] state as vendor [B,H,K,V]."""
    expected = (head_v_dim, head_k_dim)
    if state.ndim != 4 or tuple(state.shape[-2:]) != expected:
        raise ValueError(
            "SGLang KDA state must be [B,H,V,K] with "
            f"(V,K)={expected}, got {tuple(state.shape)}"
        )
    return state.transpose(-1, -2).float().contiguous()


def _from_nvidia_kda_state_layout(
    state: torch.Tensor,
    *,
    head_k_dim: int,
    head_v_dim: int,
    dtype: torch.dtype,
) -> torch.Tensor:
    """Materialize vendor [B,H,K,V] state as SGLang [B,H,V,K]."""
    expected = (head_k_dim, head_v_dim)
    if state.ndim != 4 or tuple(state.shape[-2:]) != expected:
        raise ValueError(
            "NVIDIA KDA state must be [B,H,K,V] with "

View on GitHub (pinned to 0132848349)

Solutions

1. Ensure the input state is [B,H,V,K] with V=head_v_dim, K=head_k_dim
2. If you hold a [B,H,K,V] tensor, convert with the inverse helper first
3. Assert the trailing shape before calling extend

Example fix

# before
state = torch.empty(B, H, K, V)  # vendor layout passed directly
# after
state = torch.empty(B, H, V, K)  # SGLang layout

assert state.ndim == 4 and tuple(state.shape[-2:]) == (head_v_dim, head_k_dim), (
    f'expected [B,H,V,K] ending ({head_v_dim},{head_k_dim}), got {tuple(state.shape)}')

def is_sglang_kda_state(t: torch.Tensor, head_v_dim: int, head_k_dim: int) -> bool:
    return t.ndim == 4 and tuple(t.shape[-2:]) == (head_v_dim, head_k_dim)

When it happens

Trigger: Calling extend on NvidiaKDAKernel (or the layout unit tests) with a state tensor whose trailing dims are swapped (K,V) or whose ndim != 4.

Common situations: Passing a vendor-layout state straight back in, or a cache allocated with head dims in the wrong order; refactors of the state pool layout.

Related errors

• NVIDIA KDA state must be [B,H,K,V] with (K,V)={expected}, go
• f"recurrent_kda needs a [N, HV, V, K] state pool; got shape
• NvidiaKDAKernel is prefill-only
• NvidiaKDAKernel does not support target_verify
• `dt_bias` must have {HV * K} elements (got {dt_bias.numel()}

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