{"record":{"id":"fcd6e935bdc0c736","repo":"sgl-project/sglang","slug":"a-log-and-dt-bias-must-have-hv-elements-got","errorCode":null,"errorMessage":"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()}).","messagePattern":"`A_log` and `dt_bias` must have (.+?) elements \\(got A_log\\.numel\\(\\)=(.+?), dt_bias\\.numel\\(\\)=(.+?)\\)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent.py","lineNumber":333,"sourceCode":"        )\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\n    if q_dim % K != 0:\n        raise ValueError(f\"Invalid packed Q size {q_dim}: must be divisible by K={K}.\")\n    H = q_dim // K\n    if H <= 0 or HV % H != 0:","sourceCodeStart":315,"sourceCodeEnd":351,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py#L315-L351","documentation":"A_log and dt_bias must each contain exactly HV elements (one per value head), where HV comes from initial_state.shape[-3]. The wrapper checks numel and raises when the parameter vectors were sized for a different head configuration — e.g. sharded or full-head versions mismatched with the state cache.","triggerScenarios":"Passing full-model A_log (num_heads) with a TP-sharded state cache (num_heads/tp); passing dt_bias of shape (HV*2) from a fused gate projection without splitting; loading params from a checkpoint with a different head count.","commonSituations":"Tensor-parallel rank 1..N passing unsharded parameters; fusing A_log and dt_bias into one vector and forgetting to chunk; Heterogeneous head configs (Qwen3-Next KDA linear head ratios).","solutions":["Shard consistently: A_log = A_log_full[hv_start:hv_end] per TP rank so numel == HV","If fused, split: A_log, dt_bias = param.chunk(2) then verify each has HV elements"],"exampleFix":"# before\nA_log_rank = A_log_full  # numel = num_heads, but HV = num_heads//tp\n# after\nA_log_rank = A_log_full.chunk(tp, dim=-1)[rank].contiguous()  # numel == HV","handlingStrategy":"validation","validationCode":"HV = initial_state.shape[-3]\nassert A_log.numel() == HV == dt_bias.numel(), (A_log.numel(), dt_bias.numel(), HV)","typeGuard":"def param_heads_match(A_log, dt_bias, HV: int) -> bool:\n    return A_log.numel() == HV and dt_bias.numel() == HV","tryCatchPattern":null,"preventionTips":["Shard A_log/dt_bias per TP rank alongside the state cache","Keep a single source-of-truth head-count constant used by cache allocation and parameter creation"],"tags":["fla","fused-recurrent","mamba-params","tensor-parallel"],"backgroundTag":"head-count-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}