{"record":{"id":"8cb5f5242b815ff9","repo":"sgl-project/sglang","slug":"expected-q-k-shape-expected-shape","errorCode":null,"errorMessage":"expected q/k shape {expected_shape}","messagePattern":"expected q/k shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py","lineNumber":987,"sourceCode":"    --speculative-dspark-block-size), inferred here from T // N - 1.\n\n    ReplaySSM: passing the four replayssm_* rings switches the kernel to\n    CACHE_RING mode — per-step raw inputs go to the rings (consumed by the\n    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","sourceCodeStart":969,"sourceCodeEnd":1005,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py#L969-L1005","documentation":"fused_kda_decode_mtp_dspark validates that q and k are dense [1, T, H, TILE_K] (TILE_K=128) tensors — batch 1, all tokens flattened into dim 1, heads H, head dim 128. A ragged/varlen layout or wrong head dim fails this check before the DSlab kernel launches.","triggerScenarios":"Passing x_q or x_k with a batch dim != 1, a 3D layout, a head dim != TILE_K (128), or varlen cu_seqlens-packed layout instead of the dense flattened one the wrapper requires.","commonSituations":"Reusing a varlen [1, total_tokens, H, D] path with per-request 4D tensors; models with head_dim != 128 (e.g. 64) routed to this kernel; accidental transpose.","solutions":["Reshape q/k to [1, T, H, 128] where T is the flattened token count","Verify head dim is exactly TILE_K=128; for other head dims use the generic KDA decode backend","Ensure batch size is 1 (all requests concatenated along dim 1), matching cu_seqlens"],"exampleFix":"# before\nq = q.reshape(num_tokens, H, 128)  # 3D -> fails\nfused_kda_decode_mtp_dspark(q, k, ...)\n# after\nq = q.reshape(1, num_tokens, H, 128)\nk = k.reshape(1, num_tokens, H, 128)\nfused_kda_decode_mtp_dspark(q, k, ...)","handlingStrategy":"validation","validationCode":"T = cu_seqlens[-1].item()\nassert x_q.shape == (1, T, H, 128) and x_k.shape == (1, T, H, 128)","typeGuard":"def is_dspark_qk(x: torch.Tensor, T: int, H: int) -> bool:\n    return tuple(x.shape) == (1, T, H, 128)","tryCatchPattern":null,"preventionTips":["Centralize the [1, T, H, 128] flatten in one helper used by q/k/v/g","Assert head_dim == 128 in the backend-selection code","Keep cu_seqlens consistent with the flattened token dim"],"tags":["kda","mtp","shape-validation","dspark"],"backgroundTag":"invalid-shape-argument","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}