sgl-project/sglang · error · ValueError

indices must have shape (s_q, h_kv, topk), got {tuple(indice

Error message

indices must have shape (s_q, h_kv, topk), got {tuple(indices.shape)}

What it means

indices must be 3-D (s_q, h_kv, topk) — the per-query/per-head selected KV positions. Other ranks raise ValueError before unpacking.

Source

Thrown at python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py:305

    are allocated and returned; callers that want to reuse buffers may pass
    pre-allocated ``out`` / ``max_logits`` / ``lse`` tensors of the expected
    shape/dtype/device. The three output tensors must not alias each other.

    Returns:
        out:        [s_q, h_q, d_v], bfloat16
        max_logits: [s_q, h_q], float32
        lse:        [s_q, h_q], float32
    """
    # Validate ranks before unpacking shapes so malformed callers fail with a
    # clear error instead of a Python unpacking/indexing exception.
    if q.ndim != 3:
        raise ValueError(f"q must have shape (s_q, h_q, d_qk), got {tuple(q.shape)}")
    if kv.ndim != 3:
        raise ValueError(
            f"kv must have shape (s_kv, h_kv, d_qk), got {tuple(kv.shape)}"
        )
    if indices.ndim != 3:
        raise ValueError(
            "indices must have shape (s_q, h_kv, topk), " f"got {tuple(indices.shape)}"
        )

    s_q, h_q, d_qk = q.shape
    s_kv, h_kv, kv_d_qk = kv.shape
    topk = indices.shape[2]
    device = q.device

    # entry.cuh interprets q/kv as contiguous FP8 buffers and launches all
    # accesses on q's CUDA device. Reject contract violations before launch.
    if not q.is_cuda:
        raise ValueError("q must be a CUDA tensor")
    if not kv.is_cuda:
        raise ValueError("kv must be a CUDA tensor")
    if not indices.is_cuda:
        raise ValueError("indices must be a CUDA tensor")

    if kv.device != device:

View on GitHub (pinned to 0132848349)

Solutions

1. Expand indices to (s_q, h_kv, topk), e.g. repeat along the head dim if shared: indices.expand(s_q, h_kv, topk)
2. Check indices.shape == (q.shape[0], kv.shape[1], topk) before the call

Example fix

# before
idx = topk_indices  # [s_q, topk]
# after
idx = topk_indices[:, None, :].expand(s_q, h_kv, topk)
out = sparse_mla_q8kv8_prefill_fwd(q, kv, idx, ...)

assert idx.ndim == 3 and idx.shape[:2] == (s_q, h_kv), idx.shape

def is_valid_indices(idx: torch.Tensor, s_q: int, h_kv: int, topk: int) -> bool:
    return idx.ndim == 3 and tuple(idx.shape) == (s_q, h_kv, topk)

When it happens

Trigger: Passing 2-D indices (e.g. shared across heads) or 4-D indices (batched) to sparse_mla_q8kv8_prefill_fwd.

Common situations: Using NSA top-k output without broadcasting to h_kv heads; assuming one index set per query is enough.

Related errors

• q must have shape (s_q, h_q, d_qk), got {tuple(q.shape)}
• kv must have shape (s_kv, h_kv, d_qk), got {tuple(kv.shape)}
• indices must be a CUDA tensor
• {name} must have shape {tuple(shape)}, got {tuple(t.shape)}
• {name} must have dtype {dtype}, got {t.dtype}

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/4c81d26cc0c7b9d7. Report an issue: GitHub.