sgl-project/sglang · error · ValueError

text QKV shapes must match

Error message

text QKV shapes must match

What it means

txt_q, txt_k, txt_v must share the exact (B, txt_tokens, num_heads, head_dim) shape derived from txt_q (note: txt_tokens may differ from img_tokens, but head count/dim and batch must match the img side too).

Solutions

1. Derive all three txt tensors from the same projection output chunked on the last dim, guaranteeing identical shapes
2. Assert txt_q.shape == txt_k.shape == txt_v_v.shape before the call

Example fix

# before
q, k, v = txt_proj(x).split([H*D, H_kv*D, H_kv*D], dim=-1)
# after
q, k, v = txt_proj(x).chunk(3, dim=-1)  # equal sizes -> matching shapes
q = q.view(B, S_txt, H, D); k = k.view(B, S_txt, H, D); v = v.view(B, S_txt, H, D)

expected = tuple(txt_q.shape)
assert all(tuple(t.shape) == expected for t in (txt_k, txt_v))

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

When it happens

Trigger: txt_k/txt_v shaped differently from txt_q — different text token counts (mismatched text masks), or head counts mismatched after a projection split.

Common situations: Text conditioning stream truncated/padded inconsistently with the img stream; chunk boundaries off when splitting a fused QKV projection.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

• image QKV shapes must match
• QKV tensors must have shape [B, S, H, D]
• head_dim must be positive, even, and <= 128
• QKV last dimensions must be contiguous
• QKV tensors must be CUDA bfloat16 tensors

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

Appendix: source

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

    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),
        device=img_q.device,
        dtype=img_q.dtype,

View on GitHub (pinned to 0132848349)