{"record":{"id":"82ac1d0cd50d709f","repo":"sgl-project/sglang","slug":"q-k-and-v-must-have-the-same-3d-shape","errorCode":null,"errorMessage":"q, k, and v must have the same 3D shape","messagePattern":"q, k, and v must have the same 3D shape","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/diffusion/layout/ulysses_qkv_triton.py","lineNumber":64,"sourceCode":"        v_ptr + row * stride_v_row + global_head * stride_v_head + dim,\n        mask=mask,\n    )\n    output_base = head_slot * (3 * head_size) + dim\n    tl.store(output_ptr + output_base, q, mask=mask)\n    tl.store(output_ptr + output_base + head_size, k, mask=mask)\n    tl.store(output_ptr + output_base + 2 * head_size, v, mask=mask)\n\n\ndef pack_qkv_destination_major(\n    q: torch.Tensor,\n    k: torch.Tensor,\n    v: torch.Tensor,\n    world_size: int,\n    out: torch.Tensor | None = None,\n) -> torch.Tensor:\n    \"\"\"Pack matching ``[rows, global_heads, head_size]`` Q/K/V tensors.\"\"\"\n    if q.dim() != 3 or q.shape != k.shape or q.shape != v.shape:\n        raise ValueError(\"q, k, and v must have the same 3D shape\")\n    if not (q.is_cuda and k.is_cuda and v.is_cuda):\n        raise ValueError(\"q, k, and v must be CUDA tensors\")\n    if not (q.device == k.device == v.device and q.dtype == k.dtype == v.dtype):\n        raise ValueError(\"q, k, and v must have the same device and dtype\")\n    if q.stride(-1) != 1 or k.stride(-1) != 1 or v.stride(-1) != 1:\n        raise ValueError(\"q, k, and v must be contiguous in head_size\")\n    if world_size < 1 or q.shape[1] % world_size != 0:\n        raise ValueError(\"world_size must be positive and divide global_heads\")\n\n    rows, global_heads, head_size = q.shape\n    local_heads = global_heads // world_size\n    expected_shape = (world_size, rows, local_heads, 3 * head_size)\n    if out is not None:\n        if not (\n            out.shape == expected_shape\n            and out.is_contiguous()\n            and out.dtype == q.dtype\n            and out.device == q.device","sourceCodeStart":46,"sourceCodeEnd":82,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/diffusion/layout/ulysses_qkv_triton.py#L46-L82","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Flatten to 3D: q = q.view(-1, num_heads, head_dim) for q, k, v","If using GQA/MQA, repeat_interleave K/V heads so all three have global_heads before packing","Assert q.shape == k.shape == v.shape and dim()==3 before the call","Route through the higher-level _usp_input_all_to_all_qkv helpers which pre-normalize shapes"],"exampleFix":"# before\npacked = pack_qkv_destination_major(q_4d, k_4d, v_4d, world_size)\n# after\nq = q_4d.reshape(-1, H, D); k = k_4d.reshape(-1, H, D); v = v_4d.reshape(-1, H, D)\npacked = pack_qkv_destination_major(q, k, v, world_size)","handlingStrategy":"validation","validationCode":"assert q.dim() == 3 and q.shape == k.shape == v.shape, (q.shape, k.shape, v.shape)","typeGuard":"def qkv_same_3d(q, k, v) -> bool:\n    return q.dim() == 3 and q.shape == k.shape == v.shape","tryCatchPattern":null,"preventionTips":["Flatten [B,S,H,D] to [B*S,H,D] before packing","Expand GQA K/V heads to match Q before Ulysses packing"],"tags":["attention","ulysses","shape-mismatch","sequence-parallelism"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}