sgl-project/sglang · error · ValueError

QKV tensors must have shape [B, S, H, D]

Error message

QKV tensors must have shape [B, S, H, D]

What it means

hunyuan_qkv_rope_pack fuses RoPE with QKV packing for Hunyuan video diffusion and requires all six img/txt Q/K/V tensors to be 4D (B, S, H, D). Anything else can't be packed by the Triton kernel.

Source

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

    tl.store(output_ptr + plane_stride + output_row + even[None, :], ok0, mask=mask)
    tl.store(output_ptr + plane_stride + output_row + odd[None, :], ok1, mask=mask)
    tl.store(output_ptr + 2 * plane_stride + output_row + even[None, :], v0, mask=mask)
    tl.store(output_ptr + 2 * plane_stride + output_row + odd[None, :], v1, mask=mask)


def hunyuan_qkv_rope_pack(
    img_q: torch.Tensor,
    img_k: torch.Tensor,
    img_v: torch.Tensor,
    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")

View on GitHub (pinned to 0132848349)

Solutions

  1. Reshape each projection: q.view(B, S, num_heads, head_dim) before calling
  2. Check upstream linear layer outputs are reshaped consistently for all six tensors

Example fix

# before
img_q, img_k, img_v = proj(x).chunk(3, dim=-1)  # (B,S,H*D)
out = hunyuan_qkv_rope_pack(img_q, ...)
# after
img_q = img_q.view(B, S, H, D)
img_k = img_k.view(B, S, H, D)
img_v = img_v.view(B, S, H, D)
out = hunyuan_qkv_rope_pack(img_q, ...)
Defensive patterns

Strategy: validation

Validate before calling

assert all(t.ndim == 4 for t in (img_q, img_k, img_v, txt_q, txt_k, txt_v))

Type guard

def is_4d_bshd(t: torch.Tensor) -> bool:
    return t.ndim == 4

Prevention

When it happens

Trigger: Calling hunyuan_qkv_rope_pack with any of img_q/img_k/img_v/txt_q/txt_k/txt_v having ndim != 4 (e.g. 3D (B, S, H*D) unreshaped projections).

Common situations: Forgetting to reshape attention projections from (B, S, H*D) to (B, S, H, D) before the fused RoPE/pack call.

Related errors


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