sgl-project/sglang · error · ValueError

image QKV shapes must match

Error message

image QKV shapes must match

What it means

img_q, img_k, img_v must share the exact same (B, img_tokens, num_heads, head_dim) shape, taken from img_q, since the kernel packs them into one interleaved buffer.

Source

Thrown at python/sglang/kernels/ops/diffusion/rope/hunyuan_qkv_pack_triton.py:173

    txt_q: torch.Tensor,
    txt_k: torch.Tensor,
    txt_v: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    tensors = (img_q, img_k, img_v, txt_q, txt_k, txt_v)
    if any(x.ndim != 4 for x in tensors):
        raise ValueError("QKV tensors must have shape [B, S, H, D]")
    if any(not x.is_cuda or x.dtype != torch.bfloat16 for x in tensors):
        raise ValueError("QKV tensors must be CUDA bfloat16 tensors")
    if any(x.device != img_q.device for x in tensors):
        raise ValueError("QKV tensors must be on the same CUDA device")
    batch, img_tokens, num_heads, head_dim = img_q.shape
    txt_tokens = txt_q.shape[1]
    expected_img = (batch, img_tokens, num_heads, head_dim)
    expected_txt = (batch, txt_tokens, num_heads, head_dim)
    if any(tuple(x.shape) != expected_img for x in (img_q, img_k, img_v)):
        raise ValueError("image QKV shapes must match")
    if any(tuple(x.shape) != expected_txt for x in (txt_q, txt_k, txt_v)):
        raise ValueError("text QKV shapes must match")
    if any(x.stride(-1) != 1 for x in tensors):
        raise ValueError("QKV last dimensions must be contiguous")
    if head_dim <= 0 or head_dim > 128 or head_dim % 2:
        raise ValueError("head_dim must be positive, even, and <= 128")
    if cos.ndim != 2 or sin.ndim != 2 or cos.shape != sin.shape:
        raise ValueError("cos and sin must have matching [S, D/2] shapes")
    if cos.shape[0] < img_tokens or cos.shape[1] != head_dim // 2:
        raise ValueError("cos/sin shape does not cover image tokens and head_dim")
    if not cos.is_cuda or not sin.is_cuda or cos.stride(-1) != 1 or sin.stride(-1) != 1:
        raise ValueError("cos and sin must be CUDA and last-dim contiguous")
    if cos.device != img_q.device or sin.device != img_q.device:
        raise ValueError("QKV and cos/sin tensors must be on the same CUDA device")

    total_tokens = img_tokens + txt_tokens
    storage = torch.empty(
        (3, batch, total_tokens, num_heads, head_dim),

View on GitHub (pinned to 0132848349)

Solutions

1. Repeat/expand K/V heads to match Q heads (e.g. torch.repeat_interleave on the head dim)
2. Ensure img_q/k/v come from chunking the same projection with consistent shapes

Example fix

# before
img_k = img_k  # (B, S, H_kv, D) with H_kv < H
# after
img_k = img_k.repeat_interleave(H // H_kv, dim=2)  # (B, S, H, D)

expected = tuple(img_q.shape)
assert all(tuple(t.shape) == expected for t in (img_k, img_v))

def img_qkv_match(q, k, v) -> bool:
    return q.shape == k.shape == v.shape

When it happens

Trigger: Any of img_k or img_v having a different token count, head count, or head dim than img_q (e.g. GQA where kv heads weren't expanded to match, or mismatched sequence lengths after a split).

Common situations: Using grouped-query/MLA attention without repeating K/V heads to match Q heads before the fused pack; img tensor split at wrong boundaries.

Related errors

• QKV tensors must have shape [B, S, H, D]
• text QKV shapes must match
• q shape must be [num_tokens, num_qo_heads, head_dim], got {q
• k shape must be [num_tokens, num_kv_heads, head_dim], got {k
• QKV tensors must be CUDA bfloat16 tensors

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