{"record":{"id":"85e956a0a0c5fde5","repo":"sgl-project/sglang","slug":"nvidia-kda-state-must-be-b-h-k-v-with-k-v-exp","errorCode":null,"errorMessage":"NVIDIA KDA state must be [B,H,K,V] with (K,V)={expected}, got {tuple(state.shape)}","messagePattern":"NVIDIA KDA state must be \\[B,H,K,V\\] with \\(K,V\\)=(.+?), got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/layers/attention/linear/kernels/kda_nvidia.py","lineNumber":71,"sourceCode":"    if state.ndim != 4 or tuple(state.shape[-2:]) != expected:\n        raise ValueError(\n            \"SGLang KDA state must be [B,H,V,K] with \"\n            f\"(V,K)={expected}, got {tuple(state.shape)}\"\n        )\n    return state.transpose(-1, -2).float().contiguous()\n\n\ndef _from_nvidia_kda_state_layout(\n    state: torch.Tensor,\n    *,\n    head_k_dim: int,\n    head_v_dim: int,\n    dtype: torch.dtype,\n) -> torch.Tensor:\n    \"\"\"Materialize vendor [B,H,K,V] state as SGLang [B,H,V,K].\"\"\"\n    expected = (head_k_dim, head_v_dim)\n    if state.ndim != 4 or tuple(state.shape[-2:]) != expected:\n        raise ValueError(\n            \"NVIDIA KDA state must be [B,H,K,V] with \"\n            f\"(K,V)={expected}, got {tuple(state.shape)}\"\n        )\n    return state.transpose(-1, -2).to(dtype=dtype).contiguous()\n\n\nclass NvidiaKDAKernel(LinearAttnKernelBase):\n    def __init__(self):\n        # This kernel uses tcgen05 + TMEM, which are available on datacenter\n        # Blackwell (SM100/SM103, reported as capability major 10), but not on\n        # Blackwell desktop SM120 even though its capability number is larger.\n        self.supports_prefill = torch.cuda.is_available() and (\n            torch.cuda.get_device_capability()[0] == 10\n        )\n        self._fwd = None\n        self._l2norm = None\n        self._triton = TritonKDAKernel()\n        # Stable detached fp32 views of the (frozen) gate params: nn.Parameters","sourceCodeStart":53,"sourceCodeEnd":89,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/layers/attention/linear/kernels/kda_nvidia.py#L53-L89","documentation":"_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.","triggerScenarios":"Calling extend with a NVIDIA-kernel output state (or test fixture) whose trailing dims are (V,K) instead of (K,V), or ndim != 4.","commonSituations":"Round-tripping states between NVIDIA kernels and SGLang caches with one side already converted; hand-built test tensors with swapped dims.","solutions":["Pass the raw [B,H,K,V] vendor-layout tensor, not an already-converted one","Check tuple(state.shape[-2:]) == (head_k_dim, head_v_dim) before the call","Use the paired helper functions rather than manual transposes"],"exampleFix":"# before\nsglang_state = _from_nvidia_kda_state_layout(already_converted, ...)  # wrong\n# after\nsglang_state = _from_nvidia_kda_state_layout(vendor_state, ...)  # [B,H,K,V]","handlingStrategy":"validation","validationCode":"assert state.ndim == 4 and tuple(state.shape[-2:]) == (head_k_dim, head_v_dim), (\n    f'expected vendor [B,H,K,V] ending ({head_k_dim},{head_v_dim}), got {tuple(state.shape)}')","typeGuard":"def is_nvidia_kda_state(t: torch.Tensor, head_k_dim: int, head_v_dim: int) -> bool:\n    return t.ndim == 4 and tuple(t.shape[-2:]) == (head_k_dim, head_v_dim)","tryCatchPattern":null,"preventionTips":["Never double-convert layouts; track which side produced the tensor","Round-trip test layout conversion in unit tests"],"tags":["sglang","nvidia","kda","tensor-shape","layout"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}