sgl-project/sglang · error · ValueError

sparse_mla_q8kv8_prefill_fwd requires h_kv=1, got {h_kv}

Error message

sparse_mla_q8kv8_prefill_fwd requires h_kv=1, got {h_kv}

What it means

This sparse MLA kernel only supports multi-head attention with a single shared KV head (MLA-style, h_kv=1) where all query heads attend to one latent KV stream. A kv tensor with h_kv != 1 is rejected.

Source

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

        raise ValueError("kv must be contiguous")
    if not indices.is_contiguous():
        raise ValueError("indices must be contiguous")

    if kv_d_qk != d_qk:
        raise ValueError(f"kv d_qk must match q d_qk={d_qk}, got {kv_d_qk}")

    # The CUDA implementation uses B_H=64 and launches h_q / B_H CTAs.
    # 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, "

View on GitHub (pinned to 0132848349)

Solutions

  1. Ensure the kv cache is squeezed to the single shared latent head (h_kv=1) as MLA produces
  2. Use the appropriate GQA attention backend for models with multiple KV heads
  3. Check that kv.shape is (s_kv, 1, d_qk...) before calling

Example fix

// before
kv = kv_cache  # shape (s_kv, 8, d)
out = sparse_mla_q8kv8_prefill_fwd(q, kv, indices)
// after
assert kv_cache.shape[1] == 1
out = sparse_mla_q8kv8_prefill_fwd(q, kv_cache, indices)
Defensive patterns

Strategy: validation

Validate before calling

assert kv.shape[1] == 1, f"h_kv must be 1 for sparse MLA, got {kv.shape[1]}"

Type guard

def is_mla_kv(kv: torch.Tensor) -> bool:
    return kv.shape[1] == 1

Prevention

When it happens

Trigger: Passing a kv cache with a leading head dim greater than 1, e.g. shape (s_kv, h_kv=8, d_qk) from a grouped-query attention layout.

Common situations: Reusing this MLA-specific kernel with a GQA/MQA model; cache shape assumptions from a different backend where the head dim is folded differently.

Related errors


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