{"record":{"id":"c93f9de1821561fd","repo":"sgl-project/sglang","slug":"a-and-b-must-be-2d-tensors-got-a-ndim-a-ndim","errorCode":null,"errorMessage":"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim}).","messagePattern":"`a` and `b` must be 2D tensors \\(got a\\.ndim=(.+?), b\\.ndim=(.+?)\\)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent.py","lineNumber":287,"sourceCode":"    mixed_qkv: torch.Tensor,\n    a: torch.Tensor,\n    b: torch.Tensor,\n    A_log: torch.Tensor,\n    dt_bias: torch.Tensor,\n    scale: float,\n    initial_state: torch.Tensor,\n    out: torch.Tensor,\n    ssm_state_indices: torch.Tensor,\n    use_qk_l2norm_in_kernel: bool = False,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    if mixed_qkv.ndim != 2:\n        raise ValueError(\n            f\"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim}).\"\n        )\n    if mixed_qkv.stride(-1) != 1:\n        raise ValueError(\"`mixed_qkv` must be contiguous in the last dim.\")\n    if a.ndim != 2 or b.ndim != 2:\n        raise ValueError(\n            f\"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim}).\"\n        )\n    if a.stride(-1) != 1 or b.stride(-1) != 1:\n        raise ValueError(\"`a`/`b` must be contiguous in the last dim.\")\n    if A_log.ndim != 1 or dt_bias.ndim != 1:\n        raise ValueError(\"`A_log`/`dt_bias` must be 1D tensors.\")\n    if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:\n        raise ValueError(\"`A_log`/`dt_bias` must be contiguous.\")\n    if ssm_state_indices.ndim != 1:\n        raise ValueError(\n            f\"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim}).\"\n        )\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","sourceCodeStart":269,"sourceCodeEnd":305,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py#L269-L305","documentation":"In the packed decode path, the gating log-signals a and b must each be 2D of shape (num_tokens, HV) — one scalar gate per token per value head. The wrapper raises when either has a different rank, such as the common (B, T, HV) layout or per-head (B, H, T) head-first layout.","triggerScenarios":"Passing a/b shaped (B, 1, HV) from a decode step without squeezing the time dim, or (B, HV, T) head-first tensors from fla-style code.","commonSituations":"Adapter code bridging fla's head_first conventions to sglang's packed decode; forgetting .squeeze(1)/.view(B, -1) after computing a and b from a Linear projection of shape (B, T, 2*HV).","solutions":["Reshape: a = a.view(-1, HV), b = b.view(-1, HV) (or .squeeze(1) for decode's T=1)","Split the 2*HV gate projection with .chunk(2, dim=-1) and flatten tokens"],"exampleFix":"# before\na, b = a_log[..., 0], a_log[..., 1]  # (B, 1, HV) each\n# after\na = a_log[..., 0].reshape(-1, a_log.shape[-1])\nb = a_log[..., 1].reshape(-1, a_log.shape[-1])","handlingStrategy":"validation","validationCode":"assert a.ndim == 2 and b.ndim == 2, (a.shape, b.shape)\na = a.reshape(-1, a.shape[-1]); b = b.reshape(-1, b.shape[-1])","typeGuard":"def gates_are_2d(*ts: torch.Tensor) -> bool:\n    return all(t.ndim == 2 for t in ts)","tryCatchPattern":null,"preventionTips":["Squeeze the T=1 decode dim eagerly when producing gate tensors","Centralize gate reshaping next to the gate projection"],"tags":["fla","fused-recurrent","shape-validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}