{"record":{"id":"a2e35f21d3e11976","repo":"sgl-project/sglang","slug":"ssm-state-indices-must-have-shape-b-got-tupl-a2e35f","errorCode":null,"errorMessage":"`ssm_state_indices` must have shape [B] (got {tuple(ssm_state_indices.shape)}; expected ({B},)).","messagePattern":"`ssm_state_indices` must have shape \\[B\\] \\(got (.+?); expected \\((.+?),\\)\\)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_decode.py","lineNumber":275,"sourceCode":"            b,\n            A_log,\n            dt_bias,\n            initial_state,\n            out,\n            ssm_state_indices,\n        )\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]}, \"\n            f\"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] \"\n            f\"(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 not _is_power_of_two(K) or not _is_power_of_two(V):\n        raise ValueError(\n            \"Helion KDA decode requires power-of-two key and value head \"\n            f\"dimensions (got K={K}, V={V}).\"\n        )\n    if a.shape[1] != HV * K:\n        raise ValueError(","sourceCodeStart":257,"sourceCodeEnd":293,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_decode.py#L257-L293","documentation":"ssm_state_indices must have exactly shape [B] where B is mixed_qkv.shape[0]; the validator reports both got and expected shapes. A length mismatch means state-slot indices don't cover the batch (or over-cover it).","triggerScenarios":"Passing ssm_state_indices with more entries than batch rows (e.g. indices from a larger scheduler batch) or wrong shape like [B,1].","commonSituations":"Indices cached from a previous, larger batch; not filtering indices when finished requests are removed; reshaping errors elsewhere.","solutions":["Recompute/filter ssm_state_indices to match the current batch length: idx[keep] or idx[:B]","Ensure the scheduler updates state indices whenever the batch changes","Reshape any [B,1] tensor to [B]"],"exampleFix":"# before\nout = decode(qkv, a, b, ..., ssm_state_indices=all_idx, ...)  # len(all_idx) > B\n# after\nout = decode(qkv, a, b, ..., ssm_state_indices=all_idx[keep], ...)  # len == B","handlingStrategy":"validation","validationCode":"B = mixed_qkv.shape[0]\nassert ssm_state_indices.shape == (B,), f'indices must be [{B}]'","typeGuard":"def indices_match_batch(idx: torch.Tensor, B: int) -> bool:\n    return idx.ndim == 1 and idx.shape[0] == B","tryCatchPattern":null,"preventionTips":["Recompute state indices whenever the running batch changes","Include a shape assert right before the kernel call"],"tags":["kda","mamba","shape-mismatch","state-indices","helion"],"backgroundTag":"batch-size-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}