{"record":{"id":"fe816f5515f9bd0c","repo":"sgl-project/sglang","slug":"f-recurrent-kda-needs-a-n-hv-v-k-state-pool","errorCode":null,"errorMessage":"f\"recurrent_kda needs a [N, HV, V, K] state pool; got shape {tuple(ssm_states.shape)}\"","messagePattern":"f\"recurrent_kda needs a \\[N, HV, V, K\\] state pool; got shape (.+?)\"","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py","lineNumber":105,"sourceCode":"    def _check_state_stride_contract(self, ssm_states: torch.Tensor) -> None:\n        \"\"\"One-time (per pool view) check that ``ssm_states`` matches the\n        layout ``recurrent_kda`` was compiled for.\n\n        The kernel's state argument is a CuTe fake tensor of shape\n        ``[N, HV, V, K]`` with stride ``(sym_int64(divisibility=16), V*K, K, 1)``\n        and ``assumed_align=32`` (flashinfer ``kda_kernels/recurrent_kda.py``):\n        the slot stride is free — which is what lets the envelope-strided pools\n        (unified memory / page-major layout, slot stride = per-slot envelope\n        pitch) be passed in and updated IN PLACE on the cu_seqlens path — but\n        the inner strides are compiled-in constants and the divisibility /\n        alignment are hard assumptions. A pool violating them would mis-address\n        state in-kernel without any error; fail loudly here instead.\n        \"\"\"\n        key = id(ssm_states)\n        if key in self._state_contract_ok:\n            return\n        if ssm_states.dim() != 4:\n            raise ValueError(\n                f\"recurrent_kda needs a [N, HV, V, K] state pool; got \"\n                f\"shape {tuple(ssm_states.shape)}\"\n            )\n        _, hv, v, k = ssm_states.shape\n        if ssm_states.stride()[1:] != (v * k, k, 1):\n            raise ValueError(\n                \"recurrent_kda state inner strides must be compact \"\n                f\"(V*K, K, 1)=({v * k}, {k}, 1); got {ssm_states.stride()[1:]} \"\n                \"(only the slot stride may be non-compact)\"\n            )\n        base_bytes = ssm_states.storage_offset() * ssm_states.element_size()\n        if ssm_states.stride(0) % 16 != 0 or base_bytes % 32 != 0:\n            raise ValueError(\n                \"recurrent_kda state pool breaks the compiled stride contract: \"\n                f\"slot stride {ssm_states.stride(0)} elements must be a multiple \"\n                f\"of 16 and the base byte offset {base_bytes} a multiple of 32 \"\n                \"(sym_int64(divisibility=16) / assumed_align=32)\"\n            )","sourceCodeStart":87,"sourceCodeEnd":123,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py#L87-L123","documentation":"recurrent_kda expects the SSM state pool shaped [N, HV, V, K]; this ValueError fires when the pool passed to decode is not 4-D. The check runs once per tensor identity before the first decode call.","triggerScenarios":"Calling FlashInferKDAKernel.decode (or a test harness) with an ssm_states tensor whose dim() != 4, e.g. a flat pool or a [B, T, H, V, K] 5-D buffer.","commonSituations":"Custom MambaCache / state-pool layouts that don't match SGLang's [N, HV, V, K] contract; refactors of the hybrid-attention cache that reshape the state pool.","solutions":["Reshape/allocate the state pool as [num_slots, num_v_heads, head_v_dim, head_k_dim]","If the pool came from a cache refactor, verify the mamba cache allocation code path","Add a unit assert on ssm_states.dim()==4 before dispatching to decode"],"exampleFix":"# before\nssm_states = pool.view(-1, hv * v * k)  # flat\n# after\nssm_states = pool.view(num_slots, hv, v, k)","handlingStrategy":"validation","validationCode":"assert ssm_states.dim() == 4, f'expected [N,HV,V,K], got {tuple(ssm_states.shape)}'\nN, HV, V, K = ssm_states.shape\nassert (HV, V, K) == (num_v_heads, head_v_dim, head_k_dim)","typeGuard":"def is_valid_kda_state_pool(t: torch.Tensor) -> bool:\n    return t.dim() == 4","tryCatchPattern":"try:\n    kernel.decode(...)\nexcept ValueError as e:\n    if 'state pool' in str(e):\n        ssm_states = ssm_states.reshape(num_slots, hv, v, k)\n    else:\n        raise","preventionTips":["Validate state pool shape at allocation time in custom caches","Add shape asserts in test harnesses before calling decode"],"tags":["sglang","kda","tensor-shape","validation","state-pool"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}