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).

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)

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.

Related errors

• QKV tensors must have shape [B, S, H, D]
• image QKV shapes must match
• QKV tensors must be CUDA bfloat16 tensors
• QKV tensors must be on the same CUDA device
• QKV last dimensions must be contiguous

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