sgl-project/sglang · error · ValueError

q, k, and v must have the same 3D shape

Error message

q, k, and v must have the same 3D shape

What it means

pack_qkv_destination_major packs Q, K, V tensors of identical [rows, global_heads, head_size] shape into a destination-major layout for Ulysses sequence parallelism. It requires all three to be 3D and exactly equal in shape; any mismatch raises this ValueError.

Source

Thrown at python/sglang/kernels/ops/diffusion/layout/ulysses_qkv_triton.py:64

        v_ptr + row * stride_v_row + global_head * stride_v_head + dim,
        mask=mask,
    )
    output_base = head_slot * (3 * head_size) + dim
    tl.store(output_ptr + output_base, q, mask=mask)
    tl.store(output_ptr + output_base + head_size, k, mask=mask)
    tl.store(output_ptr + output_base + 2 * head_size, v, mask=mask)


def pack_qkv_destination_major(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    world_size: int,
    out: torch.Tensor | None = None,
) -> torch.Tensor:
    """Pack matching ``[rows, global_heads, head_size]`` Q/K/V tensors."""
    if q.dim() != 3 or q.shape != k.shape or q.shape != v.shape:
        raise ValueError("q, k, and v must have the same 3D shape")
    if not (q.is_cuda and k.is_cuda and v.is_cuda):
        raise ValueError("q, k, and v must be CUDA tensors")
    if not (q.device == k.device == v.device and q.dtype == k.dtype == v.dtype):
        raise ValueError("q, k, and v must have the same device and dtype")
    if q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1:
        raise ValueError("q, k, and v must be contiguous in head_size")
    if world_size < 1 or q.shape[1] % world_size != 0:
        raise ValueError("world_size must be positive and divide global_heads")

    rows, global_heads, head_size = q.shape
    local_heads = global_heads // world_size
    expected_shape = (world_size, rows, local_heads, 3 * head_size)
    if out is not None:
        if not (
            out.shape == expected_shape
            and out.is_contiguous()
            and out.dtype == q.dtype
            and out.device == q.device

View on GitHub (pinned to 0132848349)

Solutions

1. Flatten to 3D: q = q.view(-1, num_heads, head_dim) for q, k, v
2. If using GQA/MQA, repeat_interleave K/V heads so all three have global_heads before packing
3. Assert q.shape == k.shape == v.shape and dim()==3 before the call
4. Route through the higher-level _usp_input_all_to_all_qkv helpers which pre-normalize shapes

Example fix

# before
packed = pack_qkv_destination_major(q_4d, k_4d, v_4d, world_size)
# after
q = q_4d.reshape(-1, H, D); k = k_4d.reshape(-1, H, D); v = v_4d.reshape(-1, H, D)
packed = pack_qkv_destination_major(q, k, v, world_size)

assert q.dim() == 3 and q.shape == k.shape == v.shape, (q.shape, k.shape, v.shape)

def qkv_same_3d(q, k, v) -> bool:
    return q.dim() == 3 and q.shape == k.shape == v.shape

When it happens

Trigger: Passing q/k/v with different ranks (e.g. 4D [B,S,H,D] tensors instead of flattened 3D), or shapes that differ across q, k, v (e.g. GQA where K/V have fewer heads without prior expansion/flattening).

Common situations: Feeding per-batch attention tensors directly instead of the flattened [rows=B*S, heads, head_size] form; using grouped-query attention heads where K/V head counts differ from Q without repeating them first.

Related errors

• {tensor_name}{context_clause} with shape {tensor.shape} cann
• q, k, and v must be CUDA tensors
• out must be a contiguous tensor with the expected shape, dev
• Unexpected q_proj output shape {q_proj_output.shape} for {ty
• num_heads must be divisible by num_epi_subtiles

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