sgl-project/sglang · error · ValueError

kv must have shape (s_kv, h_kv, d_qk), got {tuple(kv.shape)}

Error message

kv must have shape (s_kv, h_kv, d_qk), got {tuple(kv.shape)}

What it means

Same pre-launch rank validation as for q: kv must be a 3-D (s_kv, h_kv, d_qk) tensor. Any other rank triggers ValueError including the observed shape.

Source

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

) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Run Q8KV8 (FP8) sparse prefill attention on SM90.

    The kernel writes into three output tensors. By default fresh tensors
    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")

View on GitHub (pinned to 0132848349)

Solutions

1. Gather/reshape kv to (s_kv, h_kv, d_qk) contiguous before the call
2. Verify the kv layout produced by your quantization/qprep step matches the kernel contract

Example fix

# before
kv3 = kv_cache[layer]  # [num_blocks, ...] paged view
# after
kv3 = gather_kv_dense(kv_cache, indices_meta)  # [s_kv, h_kv, d_qk]
out = sparse_mla_q8kv8_prefill_fwd(q, kv3, indices, ...)

assert kv.ndim == 3, f'kv must be (s_kv, h_kv, d_qk), got {kv.shape}'

def is_kv3d(kv: torch.Tensor) -> bool:
    return kv.ndim == 3

When it happens

Trigger: Passing a 2-D packed KV cache or 4-D batched kv to sparse_mla_q8kv8_prefill_fwd.

Common situations: Feeding the paged KV cache view directly without gathering to the dense [s_kv, h_kv, d] layout; porting from an API that expects (B, H, S, D).

Related errors

• q must have shape (s_q, h_q, d_qk), got {tuple(q.shape)}
• rope_pool_fused expects pool tensors to be 3-D
• k_pool has incompatible shape {k_pool.shape}
• v_pool shape must match k_pool shape, got {v_pool.shape} vs
• D={D_check} must be divisible by GROUP_SIZE={_FP8_GROUP_SIZE

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