sgl-project/sglang · error · ValueError

NVIDIA KDA state must be [B,H,K,V] with (K,V)={expected}, go

Error message

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

What it means

_from_nvidia_kda_state_layout is the inverse check: it validates the vendor state is [B,H,K,V] with trailing dims (head_k_dim, head_v_dim) before converting back to SGLang [B,H,V,K]. Mismatched shape raises ValueError.

Source

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

    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 "
            f"(K,V)={expected}, got {tuple(state.shape)}"
        )
    return state.transpose(-1, -2).to(dtype=dtype).contiguous()


class NvidiaKDAKernel(LinearAttnKernelBase):
    def __init__(self):
        # This kernel uses tcgen05 + TMEM, which are available on datacenter
        # Blackwell (SM100/SM103, reported as capability major 10), but not on
        # Blackwell desktop SM120 even though its capability number is larger.
        self.supports_prefill = torch.cuda.is_available() and (
            torch.cuda.get_device_capability()[0] == 10
        )
        self._fwd = None
        self._l2norm = None
        self._triton = TritonKDAKernel()
        # Stable detached fp32 views of the (frozen) gate params: nn.Parameters

View on GitHub (pinned to 0132848349)

Solutions

1. Pass the raw [B,H,K,V] vendor-layout tensor, not an already-converted one
2. Check tuple(state.shape[-2:]) == (head_k_dim, head_v_dim) before the call
3. Use the paired helper functions rather than manual transposes

Example fix

# before
sglang_state = _from_nvidia_kda_state_layout(already_converted, ...)  # wrong
# after
sglang_state = _from_nvidia_kda_state_layout(vendor_state, ...)  # [B,H,K,V]

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

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

When it happens

Trigger: Calling extend with a NVIDIA-kernel output state (or test fixture) whose trailing dims are (V,K) instead of (K,V), or ndim != 4.

Common situations: Round-tripping states between NVIDIA kernels and SGLang caches with one side already converted; hand-built test tensors with swapped dims.

Related errors

• f"SGLang KDA state must be [B,H,V,K] with (V,K)={expected},
• 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/85e956a0a0c5fde5. Report an issue: GitHub.