{"record":{"id":"4c81d26cc0c7b9d7","repo":"sgl-project/sglang","slug":"indices-must-have-shape-s-q-h-kv-topk-got-tu","errorCode":null,"errorMessage":"indices must have shape (s_q, h_kv, topk), got {tuple(indices.shape)}","messagePattern":"indices must have shape \\(s_q, h_kv, topk\\), got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py","lineNumber":305,"sourceCode":"    are allocated and returned; callers that want to reuse buffers may pass\n    pre-allocated ``out`` / ``max_logits`` / ``lse`` tensors of the expected\n    shape/dtype/device. The three output tensors must not alias each other.\n\n    Returns:\n        out:        [s_q, h_q, d_v], bfloat16\n        max_logits: [s_q, h_q], float32\n        lse:        [s_q, h_q], float32\n    \"\"\"\n    # Validate ranks before unpacking shapes so malformed callers fail with a\n    # clear error instead of a Python unpacking/indexing exception.\n    if q.ndim != 3:\n        raise ValueError(f\"q must have shape (s_q, h_q, d_qk), got {tuple(q.shape)}\")\n    if kv.ndim != 3:\n        raise ValueError(\n            f\"kv must have shape (s_kv, h_kv, d_qk), got {tuple(kv.shape)}\"\n        )\n    if indices.ndim != 3:\n        raise ValueError(\n            \"indices must have shape (s_q, h_kv, topk), \" f\"got {tuple(indices.shape)}\"\n        )\n\n    s_q, h_q, d_qk = q.shape\n    s_kv, h_kv, kv_d_qk = kv.shape\n    topk = indices.shape[2]\n    device = q.device\n\n    # entry.cuh interprets q/kv as contiguous FP8 buffers and launches all\n    # accesses on q's CUDA device. Reject contract violations before launch.\n    if not q.is_cuda:\n        raise ValueError(\"q must be a CUDA tensor\")\n    if not kv.is_cuda:\n        raise ValueError(\"kv must be a CUDA tensor\")\n    if not indices.is_cuda:\n        raise ValueError(\"indices must be a CUDA tensor\")\n\n    if kv.device != device:","sourceCodeStart":287,"sourceCodeEnd":323,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py#L287-L323","documentation":"indices must be 3-D (s_q, h_kv, topk) — the per-query/per-head selected KV positions. Other ranks raise ValueError before unpacking.","triggerScenarios":"Passing 2-D indices (e.g. shared across heads) or 4-D indices (batched) to sparse_mla_q8kv8_prefill_fwd.","commonSituations":"Using NSA top-k output without broadcasting to h_kv heads; assuming one index set per query is enough.","solutions":["Expand indices to (s_q, h_kv, topk), e.g. repeat along the head dim if shared: indices.expand(s_q, h_kv, topk)","Check indices.shape == (q.shape[0], kv.shape[1], topk) before the call"],"exampleFix":"# before\nidx = topk_indices  # [s_q, topk]\n# after\nidx = topk_indices[:, None, :].expand(s_q, h_kv, topk)\nout = sparse_mla_q8kv8_prefill_fwd(q, kv, idx, ...)","handlingStrategy":"type-guard","validationCode":"assert idx.ndim == 3 and idx.shape[:2] == (s_q, h_kv), idx.shape","typeGuard":"def is_valid_indices(idx: torch.Tensor, s_q: int, h_kv: int, topk: int) -> bool:\n    return idx.ndim == 3 and tuple(idx.shape) == (s_q, h_kv, topk)","tryCatchPattern":null,"preventionTips":["Expand shared top-k indices across the head dimension","Verify (s_q, h_kv, topk) shape after every selection step"],"tags":["rank-validation","indices","sparse-mla","topk"],"backgroundTag":"invalid-tensor-rank","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}