{"record":{"id":"b173383fba510e0c","repo":"sgl-project/sglang","slug":"the-number-of-initial-states-is-expected-to-be-equ","errorCode":null,"errorMessage":"The number of initial states is expected to be equal to the number of input sequences, i.e., {len(cu_seqlens) - 1} rather than {initial_state_indices.shape[0]}.","messagePattern":"The number of initial states is expected to be equal to the number of input sequences, i\\.e\\., (.+?) rather than (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/chunk.py","lineNumber":242,"sourceCode":"        )\n    # if not head_first and q.shape[1] < q.shape[2]:\n    #     warnings.warn(\n    #         f\"Input tensor shape suggests potential format mismatch: seq_len ({q.shape[1]}) < num_heads ({q.shape[2]}). \"\n    #         \"This may indicate the inputs were passed in head-first format [B, H, T, ...] \"\n    #         \"when head_first=False was specified. \"\n    #         \"Please verify your input tensor format matches the expected shape [B, T, H, ...].\"\n    #     )\n    if cu_seqlens is not None:\n        if q.shape[0] != 1:\n            raise ValueError(\n                f\"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`.\"\n                f\"Please flatten variable-length inputs before processing.\"\n            )\n        if (\n            initial_state_indices is not None\n            and initial_state_indices.shape[0] != len(cu_seqlens) - 1\n        ):\n            raise ValueError(\n                f\"The number of initial states is expected to be equal to the number of input sequences, \"\n                f\"i.e., {len(cu_seqlens) - 1} rather than {initial_state_indices.shape[0]}.\"\n            )\n    if scale is None:\n        scale = k.shape[-1] ** -0.5\n    o, h = ChunkGatedDeltaRuleFunction.apply(\n        q,\n        k,\n        v,\n        g,\n        beta,\n        scale,\n        initial_state,\n        initial_state_indices,\n        cu_seqlens,\n        use_qk_l2norm_in_kernel,\n    )\n    if head_first:","sourceCodeStart":224,"sourceCodeEnd":260,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/chunk.py#L224-L260","documentation":"In chunk_gated_delta_rule (varlen path with cu_seqlens), the optional initial_state_indices tensor must contain one index per input sequence. The check compares initial_state_indices.shape[0] against len(cu_seqlens) - 1, which is the number of variable-length sequences described by the cumulative sequence lengths tensor. It throws when the caller passes a state pool index tensor sized for a different batch than the one described by cu_seqlens.","triggerScenarios":"Calling chunk_gated_delta_rule with both cu_seqlens and initial_state_indices where the number of packed sequences (len(cu_seqlens)-1) differs from initial_state_indices.shape[0]; e.g. cu_seqlens=[0, 5, 10] (2 sequences) but initial_state_indices has 1 or 3 entries.","commonSituations":"Hybrid linear-attention models (GatedDeltaNet/Qwen3-Next style) prefilling a batch whose request set changed between when the state indices were built and when the kernel was invoked; passing full-batch state indices while slicing the varlen input; off-by-one when constructing cu_seqlens or reusing indices from a previous scheduler batch.","solutions":["Verify len(cu_seqlens) - 1 equals initial_state_indices.shape[0] before the call and rebuild initial_state_indices from the current batch's request indices","Check how cu_seqlens is constructed (torch.cumsum of seq_lens with a leading 0) and confirm it matches the sequences the indices refer to","If you don't need prior states, pass initial_state_indices=None","Audit scheduler/batching code that slices inputs but not the state-pool index tensor"],"exampleFix":"# before\no, h = chunk_gated_delta_rule(q, k, v, g, cu_seqlens=cu_seqlens, initial_state_indices=idx)\n# after\nassert initial_state_indices.shape[0] == len(cu_seqlens) - 1, (\n    f\"{initial_state_indices.shape[0]} states vs {len(cu_seqlens)-1} sequences\")\no, h = chunk_gated_delta_rule(q, k, v, g, cu_seqlens=cu_seqlens, initial_state_indices=idx)","handlingStrategy":"validation","validationCode":"n_seqs = len(cu_seqlens) - 1\nassert initial_state_indices is None or initial_state_indices.shape[0] == n_seqs, (\n    initial_state_indices.shape[0], n_seqs)","typeGuard":"def valid_state_indices(idx: torch.Tensor, cu_seqlens: torch.Tensor) -> bool:\n    return idx is None or (idx.ndim == 1 and idx.shape[0] == len(cu_seqlens) - 1)","tryCatchPattern":null,"preventionTips":["Build initial_state_indices in the same function that builds cu_seqlens so they always describe the same batch","When slicing varlen inputs to a subset of sequences, slice the indices tensor with the same mask"],"tags":["fla","gated-delta-rule","varlen","batch-mismatch","shape-validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}