{"record":{"id":"482f89b0c7ce81d7","repo":"sgl-project/sglang","slug":"a-b-must-have-shape-b-hv-with-hv-hv-got","errorCode":null,"errorMessage":"`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)}).","messagePattern":"`a`/`b` must have shape \\[B, HV\\] with HV=(.+?) \\(got a\\.shape=(.+?), b\\.shape=(.+?)\\)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent.py","lineNumber":329,"sourceCode":"    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        )\n    if out.shape != (B, 1, HV, V):\n        raise ValueError(\n            f\"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)}).\"\n        )\n\n    qkv_dim = mixed_qkv.shape[1]\n    qk_dim = qkv_dim - HV * V\n    if qk_dim <= 0 or qk_dim % 2 != 0:\n        raise ValueError(\n            f\"Invalid packed `mixed_qkv` last dim={qkv_dim} for HV={HV}, V={V}.\"\n        )\n    q_dim = qk_dim // 2","sourceCodeStart":311,"sourceCodeEnd":347,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py#L311-L347","documentation":"The head-count is derived from the state tensor: HV = initial_state.shape[-3]. The gate tensors a and b must then have exactly shape (B, HV) so each token has one gate scalar per value head. The wrapper raises when their head dim disagrees with the state cache's head count (a.shape[1] != HV or b.shape[1] != HV).","triggerScenarios":"Model config changed num_v_heads (e.g. 32) but the state cache was allocated with a different count (e.g. 8 for KDA/Qwen3-Next ratios); gates computed with num_k_heads instead of num_v_heads; TP sharding of gates not matching state sharding.","commonSituations":"Tensor-parallel setups where value heads are sharded but gate projections are replicated (or vice versa); loading a checkpoint with different head counts than the runtime config; unit tests with hand-made shapes that don't match the model.","solutions":["Align head counts: assert a.shape[1] == b.shape[1] == initial_state.shape[-3] and fix whichever producer (gate projection out_features or cache allocation) is wrong","Under TP, shard both the state cache and the gate projection by the same HV factor (typically hv // tp_size)"],"exampleFix":"# before\n# a, b: (B, K_heads); state: (N, HV, V, K)\n# after\n# project gates with out_features=2*HV so a, b: (B, HV); state: (N, HV, V, K)","handlingStrategy":"validation","validationCode":"HV = initial_state.shape[-3]\nassert a.shape == (a.shape[0], HV) and b.shape == (b.shape[0], HV), (a.shape, b.shape, HV)","typeGuard":"def gate_heads_match(a, b, initial_state) -> bool:\n    HV = initial_state.shape[-3]\n    return a.ndim == 2 and b.ndim == 2 and a.shape[1] == HV == b.shape[1]","tryCatchPattern":null,"preventionTips":["Derive gate projection out_features from the same num_v_heads used to allocate the state cache","Under TP, shard gates and state cache with an identical head split"],"tags":["fla","fused-recurrent","head-mismatch","tensor-parallel"],"backgroundTag":"head-count-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}