{"record":{"id":"2cb119fc9f7aae09","repo":"sgl-project/sglang","slug":"initial-state-must-be-a-4d-tensor-got-ndim-ini","errorCode":null,"errorMessage":"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim}).","messagePattern":"`initial_state` must be a 4D tensor \\(got ndim=(.+?)\\)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent.py","lineNumber":322,"sourceCode":"    if any(\n        t.device != dev\n        for t in (a, b, A_log, dt_bias, initial_state, out, ssm_state_indices)\n    ):\n        raise ValueError(\"All inputs must be on the same device.\")\n\n    B = mixed_qkv.shape[0]\n    if a.shape[0] != B or b.shape[0] != B:\n        raise ValueError(\n            \"Mismatched batch sizes: \"\n            f\"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}.\"\n        )\n    if ssm_state_indices.shape[0] != B:\n        raise ValueError(\n            f\"`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},)).\"\n        )\n\n    if initial_state.ndim != 4:\n        raise ValueError(\n            f\"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim}).\"\n        )\n    if initial_state.stride(-1) != 1:\n        raise ValueError(\"`initial_state` must be contiguous in the last dim.\")\n    HV, V, K = initial_state.shape[-3:]\n    if a.shape[1] != HV or b.shape[1] != HV:\n        raise ValueError(\n            f\"`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)}).\"\n        )\n    if A_log.numel() != HV or dt_bias.numel() != HV:\n        raise ValueError(\n            f\"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()}).\"\n        )\n    if out.shape != (B, 1, HV, V):\n        raise ValueError(\n            f\"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)}).\"\n        )\n","sourceCodeStart":304,"sourceCodeEnd":340,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py#L304-L340","documentation":"initial_state is the recurrent SSM state cache slice and must be 4D of shape (num_states, HV, V, K) (or batched (B, HV, V, K)); the last three dims are parsed as HV, V, K. The wrapper raises when the state pool tensor has any other rank, e.g. a flattened pool or a per-layer stacked 5D tensor.","triggerScenarios":"Passing the entire mamba cache of shape (num_layers, N, HV, V, K); passing a flattened (N, HV*V*K) buffer; passing a single sequence's state as 3D.","commonSituations":"Slicing the wrong dim of a layered hybrid-attention state cache; writing a custom backend that stores states packed differently than the kernel contract.","solutions":["Index/select down to 4D: initial_state = cache[layer_idx] giving (N, HV, V, K)","If states are stored per-request per-head 3D, re-add the leading dim via unsqueeze(0) and index with ssm_state_indices"],"exampleFix":"# before\ninitial_state = mamba_cache  # (L, N, HV, V, K)\n# after\ninitial_state = mamba_cache[layer_idx]  # (N, HV, V, K)","handlingStrategy":"validation","validationCode":"assert initial_state.ndim == 4, initial_state.shape\ninitial_state = initial_state if initial_state.ndim == 4 else initial_state[0]","typeGuard":"def state_4d(s: torch.Tensor) -> bool:\n    return s.ndim == 4","tryCatchPattern":null,"preventionTips":["Index layered caches per layer before passing to the kernel","Document the (states, HV, V, K) cache layout at allocation sites"],"tags":["fla","fused-recurrent","state-cache","shape-validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}