{"record":{"id":"ff7c289a9ff1e4c3","repo":"sgl-project/sglang","slug":"unsupported-input-for-usp-merge-heads-cuda","errorCode":null,"errorMessage":"unsupported input for usp_merge_heads CUDA","messagePattern":"unsupported input for usp_merge_heads CUDA","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py","lineNumber":72,"sourceCode":"    return (\n        isinstance(x, torch.Tensor)\n        and torch.version.hip is None\n        and x.is_cuda\n        and x.dtype in _SUPPORTED_DTYPES\n        and x.dim() == 5\n        and x.numel() > 0\n        and x.is_contiguous()\n    )\n\n\ndef _usp_merge_heads_cuda(x: torch.Tensor) -> torch.Tensor:\n    \"\"\"[W, S, B, h_local, D] -> [B, S, W, h_local, D] contiguous.\n\n    Bit-exact single-pass replacement for\n    ``x.permute(2, 1, 0, 3, 4).contiguous()`` on the Ulysses output path.\n    \"\"\"\n    if not can_use_usp_merge_heads(x):\n        raise RuntimeError(\"unsupported input for usp_merge_heads CUDA\")\n    return _usp_merge_heads_custom_op(x)\n\n\ndef usp_merge_heads(x: torch.Tensor) -> torch.Tensor:\n    \"\"\"Merge Ulysses output heads with an exact eager fallback.\n\n    The backend selection lives here so callers only express the layout\n    transformation. Unsupported devices, layouts, and compiled regions retain\n    the original PyTorch operation.\n    \"\"\"\n    if not torch.compiler.is_compiling() and can_use_usp_merge_heads(x):\n        return _usp_merge_heads_cuda(x)\n    return x.permute(2, 1, 0, 3, 4).contiguous()\n","sourceCodeStart":54,"sourceCodeEnd":86,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py#L54-L86","documentation":"_usp_merge_heads_cuda is the strict CUDA fast path for merging Ulysses output heads ([W,S,B,h_local,D] -> [B,S,W,h_local,D]). It only runs when can_use_usp_merge_heads(x) accepts the tensor (5D, supported dtype, CUDA, contiguity constraints); otherwise it raises instead of falling back.","triggerScenarios":"Calling _usp_merge_heads_cuda directly with a tensor failing can_use_usp_merge_heads: wrong rank, non-CUDA device, unsupported dtype, or a layout the fused kernel cannot handle.","commonSituations":"Bypassing the public usp_merge_heads dispatcher (which has an exact eager fallback) and calling the CUDA entry directly; tensors reshaped incorrectly after the all-to-all; unsupported dtype activations reaching the merge step.","solutions":["Pre-check with can_use_usp_merge_heads(x) and fall back to x.permute(2,1,0,3,4).contiguous()","Prefer the public usp_merge_heads(x) API, which handles the fallback automatically","Ensure x is a 5D contiguous CUDA tensor in fp16/bf16/fp32","Fix upstream reshape/dtype logic if tensors arrive malformed"],"exampleFix":"# before\ny = _usp_merge_heads_cuda(x)\n# after\ny = usp_merge_heads(x)  # has eager fallback\n# or: y = _usp_merge_heads_cuda(x) if can_use_usp_merge_heads(x) else x.permute(2,1,0,3,4).contiguous()","handlingStrategy":"fallback","validationCode":"from sglang.kernels.ops.diffusion.layout.usp_relayout_jit import can_use_usp_merge_heads\nif not can_use_usp_merge_heads(x):\n    y = x.permute(2, 1, 0, 3, 4).contiguous()","typeGuard":"def can_merge(x) -> bool:\n    return can_use_usp_merge_heads(x)","tryCatchPattern":"try:\n    y = _usp_merge_heads_cuda(x)\nexcept RuntimeError:\n    y = x.permute(2, 1, 0, 3, 4).contiguous()","preventionTips":["Call usp_merge_heads (public) instead of the CUDA entry","Validate 5D/contiguous/CUDA before the fast path"],"tags":["cuda","fallback","ulysses","input-validation"],"backgroundTag":"tensor-input-validation-failed","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}