sgl-project/sglang · error · ValueError

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

Error message

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

What it means

indices must have leading shape (s_q, h_kv=1), i.e. one index row per query token for the single latent KV head, with the last dim being topk. This check compares indices.shape[:2] against q's s_q and the required h_kv.

Source

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

    # Reject unpadded TP-local head counts instead of launching zero CTAs and
    # returning uninitialized outputs, which can appear to callers as a hang or
    # a later collective failure.
    if h_q == 0 or h_q % 64 != 0:
        raise ValueError(
            "sparse_mla_q8kv8_prefill_fwd requires h_q padded to a positive "
            f"multiple of 64, got {h_q}"
        )

    if h_kv != 1:
        raise ValueError(f"sparse_mla_q8kv8_prefill_fwd requires h_kv=1, got {h_kv}")

    if d_qk not in (512, 576):
        raise ValueError(
            f"sparse_mla_q8kv8_prefill_fwd supports d_qk=512/576, got {d_qk}"
        )

    if indices.shape[:2] != (s_q, h_kv):
        raise ValueError(
            "indices must have shape "
            f"({s_q}, {h_kv}, topk), got {tuple(indices.shape)}"
        )

    if indices.dtype != torch.int32:
        raise ValueError(f"indices must be int32, got {indices.dtype}")

    if topk == 0 or topk % 128 != 0:
        raise ValueError(
            "Q8KV8 sparse-prefill topk width must be a positive multiple of 128, "
            f"got {topk}"
        )

    if topk_length is not None:
        if topk_length.shape != (s_q,) or topk_length.dtype != torch.int32:
            raise ValueError(
                f"topk_length must be int32 with shape ({s_q},), got "
                f"{tuple(topk_length.shape)}/{topk_length.dtype}"

View on GitHub (pinned to 0132848349)

Solutions

Regenerate indices at shape (s_q, 1, topk) matching the flattened prefill token count of q
Verify the top-k index producer runs over the same flattened token dimension as q
Print q.shape[0] and indices.shape to confirm they agree

Example fix

# before
indices.shape == (bs, 1, topk); s_q = bs * seq_len
# after
indices = indices.reshape(s_q, 1, topk).contiguous()

Defensive patterns

Strategy: validation

Validate before calling

s_q = q.shape[0]
assert indices.shape[:2] == (s_q, 1), f"indices {tuple(indices.shape)} vs s_q={s_q}"

Type guard

def indices_shape_ok(q: torch.Tensor, indices: torch.Tensor) -> bool:
    return tuple(indices.shape[:2]) == (q.shape[0], 1)

Prevention

Produce indices on the same flattened token axis as q
Flatten batch dims consistently between q and indices

When it happens

Trigger: Passing indices shaped for a decode batch (e.g. (bs, 1, topk) when s_q is the full prefill token count), or with head dim > 1.

Common situations: Reusing decode-path topk indices for prefill; producer that emits indices per-batch instead of per-token; mismatch between s_q used in q (flattened tokens) and indices (batched).

Related errors

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