{"record":{"id":"1af99f68ff6dc98f","repo":"sgl-project/sglang","slug":"mismatched-batch-sizes-mixed-qkv-shape-0-b-a","errorCode":null,"errorMessage":"Mismatched batch sizes: mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, b.shape[0]={b.shape[0]}.","messagePattern":"Mismatched batch sizes: mixed_qkv\\.shape\\[0\\]=(.+?), a\\.shape\\[0\\]=(.+?), b\\.shape\\[0\\]=(.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent.py","lineNumber":312,"sourceCode":"    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\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)}).\"","sourceCodeStart":294,"sourceCodeEnd":330,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py#L294-L330","documentation":"In packed decode, B (num tokens) is defined by mixed_qkv.shape[0], and the per-token gate tensors a and b must both have exactly B rows. The wrapper raises with a diagnostic showing all three row counts when either gate tensor covers a different token set than the qkv tensor.","triggerScenarios":"a/b computed for a subset of tokens (padding stripped from qkv but not gates), or gates built from a previous batch after a schedule change; concat order differences between qkv projection and gate projection.","commonSituations":"Packing multiple ranks'/sequences' tokens where one path appends padding tokens and the other doesn't; ragged-batch reassembly bugs; caching gate projections across iterations.","solutions":["Recompute or re-slice a and b from the same token layout as mixed_qkv: assert a.shape[0] == b.shape[0] == mixed_qkv.shape[0]","Build qkv and gates from the same forward of the same input batch (same projection buffer) so they cannot diverge"],"exampleFix":"# before\nout, s = ...(mixed_qkv=qkv[pad_mask], a=a, b=b, ...)  # gates not masked\n# after\nout, s = ...(mixed_qkv=qkv[pad_mask].reshape(-1, D), a=a[pad_mask].reshape(-1, HV), b=b[pad_mask].reshape(-1, HV), ...)","handlingStrategy":"validation","validationCode":"B = mixed_qkv.shape[0]\nassert a.shape[0] == B == b.shape[0], (B, a.shape[0], b.shape[0])","typeGuard":"def same_batch(mixed_qkv, a, b) -> bool:\n    B = mixed_qkv.shape[0]\n    return a.shape[0] == B and b.shape[0] == B","tryCatchPattern":null,"preventionTips":["Produce qkv and gates in the same forward pass over the same token list","Apply the same padding mask/filter to every packed tensor before the kernel"],"tags":["fla","fused-recurrent","batch-mismatch"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}