{"record":{"id":"184000a037f5bee8","repo":"sgl-project/sglang","slug":"unsupported-input-shape-g-shape-which-should-be","errorCode":null,"errorMessage":"Unsupported input shape {g.shape}, which should be (B, T, H, D) if `head_first=False` or (B, H, T, D) otherwise","messagePattern":"Unsupported input shape (.+?), which should be \\(B, T, H, D\\) if `head_first=False` or \\(B, H, T, D\\) otherwise","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/cumsum.py","lineNumber":290,"sourceCode":"            scale=scale,\n            cu_seqlens=cu_seqlens,\n            head_first=head_first,\n            output_dtype=output_dtype,\n            chunk_indices=chunk_indices,\n        )\n    elif len(g.shape) == 4:\n        return chunk_local_cumsum_vector(\n            g=g,\n            chunk_size=chunk_size,\n            reverse=reverse,\n            scale=scale,\n            cu_seqlens=cu_seqlens,\n            head_first=head_first,\n            output_dtype=output_dtype,\n            chunk_indices=chunk_indices,\n        )\n    else:\n        raise ValueError(\n            f\"Unsupported input shape {g.shape}, \"\n            f\"which should be (B, T, H, D) if `head_first=False` \"\n            f\"or (B, H, T, D) otherwise\"\n        )\n","sourceCodeStart":272,"sourceCodeEnd":295,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/cumsum.py#L272-L295","documentation":"chunk_local_cumsum computes log-decay cumulative sums and accepts g of shape (B, T, H) when head_first=False (optionally with an extra dim for two-component decays) or (B, H, T) when head_first=True. The trailing else branch raises when g.ndim does not match either supported layout for the chosen head_first flag.","triggerScenarios":"Calling chunk_local_cumsum with head_first=False but a 3D g that is actually (B, H, T), or with an ndim outside the accepted set (e.g. 2D or 5D); passing head_first=True with a (B, T, H) tensor.","commonSituations":"Porting code from the fla library where head_first defaulted to True into sglang where layouts are B,T,H-first; forgetting to transpose gate logits produced by an nn.Linear of shape (B, T, num_heads); passing per-head decay plus extra dims the kernel doesn't support.","solutions":["Reshape/transpose g to (B, T, H) (or (B, T, H, dim)) and call with head_first=False","If your tensor is (B, H, T), either pass head_first=True or permute with g.transpose(1, 2).contiguous()","Print g.shape right before the call and compare with the branch conditions in cumsum.py to see which layout the function expects"],"exampleFix":"# before\ng_cumsum = chunk_local_cumsum(g, chunk_size, head_first=False)  # g is (B, H, T)\n# after\ng_cumsum = chunk_local_cumsum(g.transpose(1, 2).contiguous(), chunk_size, head_first=False)","handlingStrategy":"validation","validationCode":"assert g.ndim == 3 and (head_first and g.shape[2] == 1 or not head_first and g.shape[2] == H), g.shape\n# head_first=False expects (B, T, H); transpose if needed\ng = g.transpose(1, 2).contiguous() if head_first and g.shape[1] != T else g","typeGuard":"def cumsum_layout_ok(g: torch.Tensor, head_first: bool) -> bool:\n    return g.ndim == 3 and ((not head_first) or True) and g.ndim in (3, 4)","tryCatchPattern":null,"preventionTips":["Standardize on head_first=False (B, T, H) throughout your model code","Keep the gate projection output contiguous and untransposed before cumsum"],"tags":["fla","cumsum","shape-validation","head-first"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}