{"record":{"id":"e2042778800c0663","repo":"sgl-project/sglang","slug":"expected-beta-shape-1-t-h","errorCode":null,"errorMessage":"expected beta shape {(1, T, H)}","messagePattern":"expected beta shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py","lineNumber":991,"sourceCode":"    commit-time exact fold, see kda_replayssm_spec_decode.py) and the per-step\n    intermediate_ssm state snapshots are skipped, so intermediate_ssm may be\n    None.\n\n    Passing all three onorm_* arguments fuses gated RMSNorm into the recurrence\n    kernel.\n    \"\"\"\n    import torch\n\n    H = x_q.shape[2]\n    N = cu_seqlens.numel() - 1\n    T = x_q.shape[1]\n    expected_shape = (1, T, H, TILE_K)\n    if tuple(x_q.shape) != expected_shape or tuple(x_k.shape) != expected_shape:\n        raise ValueError(f\"expected q/k shape {expected_shape}\")\n    if tuple(x_v.shape) != expected_shape or tuple(g.shape) != expected_shape:\n        raise ValueError(f\"expected v/g shape {expected_shape}\")\n    if tuple(beta.shape) != (1, T, H):\n        raise ValueError(f\"expected beta shape {(1, T, H)}\")\n    # T // N == 1 is num_spec == 0: one token per request, i.e. a plain decode\n    # step. The backend never dispatches here for it (that is the dedicated\n    # decode kernel's job), but the layout is legal and benchmarks compare the\n    # two at this point, so the wrapper accepts it.\n    if N <= 0 or T % N != 0 or T // N < 1:\n        raise ValueError(\n            f\"DSpARK KDA MTP requires a fixed 1 + num_spec dense tokens per \"\n            f\"request; got T={T}, N={N}\"\n        )\n    num_spec = T // N - 1\n    if recurrent_state.shape[1:] != (H, TILE_K, TILE_K):\n        raise ValueError(\"expected recurrent state layout [pool, H, V=128, K=128]\")\n    if (\n        recurrent_state.dtype != torch.float32\n        or tuple(recurrent_state.stride()[-3:]) != (TILE_K * TILE_K, TILE_K, 1)\n        or recurrent_state.stride(0) % 4 != 0\n        or recurrent_state.storage_offset() % 4 != 0\n    ):","sourceCodeStart":973,"sourceCodeEnd":1009,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py#L973-L1009","documentation":"beta is the per-token per-head decay scalar and must be exactly [1, T, H] — one scalar per (token, head), batch flattened. Any other rank or shape (e.g. [T, H], [1, T, H, 1], or per-head broadcast) is rejected.","triggerScenarios":"Passing beta with an extra unit dim, missing batch dim, or a per-token scalar [1, T, 1] to fused_kda_decode_mtp_dspark.","commonSituations":"Projections that return beta as [T, H, 1] from a [*, H, D] matmul and forget to squeeze; broadcasting assumptions from another KDA kernel that accepted [T, H].","solutions":["Squeeze/reshape beta to [1, T, H] before the call","Audit the beta projection: compute it as x @ beta_proj.reshape(H, D).sum(-1) style giving [T, H], then beta.unsqueeze(0)","Add a debug assert on beta.ndim == 3 in test wrappers"],"exampleFix":"# before\nfused_kda_decode_mtp_dspark(q, k, v, g, beta)  # beta is [1, T, H, 1]\n# after\nbeta = beta.reshape(1, T, H)\nfused_kda_decode_mtp_dspark(q, k, v, g, beta)","handlingStrategy":"validation","validationCode":"assert tuple(beta.shape) == (1, T, H), beta.shape\nbeta = beta.reshape(1, T, H)","typeGuard":"def is_dspark_beta(b: torch.Tensor, T: int, H: int) -> bool:\n    return tuple(b.shape) == (1, T, H)","tryCatchPattern":null,"preventionTips":["Squeeze the beta projection output explicitly after matmul","Keep a shape contract table for all DSpARK MTP inputs in one place","Add torch.testing.assert_close-style wrappers that validate shapes first"],"tags":["kda","mtp","shape-validation","beta-decay"],"backgroundTag":"invalid-shape-argument","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}