{"record":{"id":"3ece7525df7ffe13","repo":"sgl-project/sglang","slug":"ssm-state-indices-must-have-shape-b-got-tupl","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":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent.py","lineNumber":317,"sourceCode":"        )\n    if not out.is_contiguous():\n        raise ValueError(\"`out` must be contiguous.\")\n\n    dev = mixed_qkv.device\n    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        )","sourceCodeStart":299,"sourceCodeEnd":335,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py#L299-L335","documentation":"ssm_state_indices must be exactly shape [B] where B is the packed token count from mixed_qkv.shape[0]. This check (shape[0] != B) catches index tensors whose first dim disagrees even if they are 1D — e.g. sized for a different batch or containing one index per sequence rather than per token.","triggerScenarios":"Passing per-sequence indices (length num_seqs) instead of per-token indices (length num_tokens) for decode where each token is a sequence; stale indices from a previous batch; concatenating rank indices in a different order/count than the qkv tokens.","commonSituations":"Decode batches where num_tokens == num_seqs making the bug invisible in tests but failing under continuous batching with mixed lengths; DP gather order mismatches.","solutions":["Regenerate indices per packed batch: torch.full((B,), -1) filled with cache slots for tokens carrying state; assert ssm_state_indices.shape[0] == mixed_qkv.shape[0]","Keep index construction adjacent to qkv packing so they always share the same token list"],"exampleFix":"# before\nidx = cache_loc  # per-sequence\n# after\nidx = cache_loc.reshape(-1)\nassert idx.shape[0] == mixed_qkv.shape[0]","handlingStrategy":"validation","validationCode":"assert ssm_state_indices.shape[0] == mixed_qkv.shape[0], (\n    ssm_state_indices.shape, mixed_qkv.shape)","typeGuard":"def indices_match_batch(idx, mixed_qkv) -> bool:\n    return idx.ndim == 1 and idx.shape[0] == mixed_qkv.shape[0]","tryCatchPattern":null,"preventionTips":["Use one index per decode token, -1 for tokens without state","Rebuild indices whenever the packed token set changes (continuous batching)"],"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"}