sgl-project/sglang · error · RuntimeError

unsupported input for usp_merge_heads CUDA

Error message

unsupported input for usp_merge_heads CUDA

What it means

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

Source

Thrown at python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py:72

    return (
        isinstance(x, torch.Tensor)
        and torch.version.hip is None
        and x.is_cuda
        and x.dtype in _SUPPORTED_DTYPES
        and x.dim() == 5
        and x.numel() > 0
        and x.is_contiguous()
    )


def _usp_merge_heads_cuda(x: torch.Tensor) -> torch.Tensor:
    """[W, S, B, h_local, D] -> [B, S, W, h_local, D] contiguous.

    Bit-exact single-pass replacement for
    ``x.permute(2, 1, 0, 3, 4).contiguous()`` on the Ulysses output path.
    """
    if not can_use_usp_merge_heads(x):
        raise RuntimeError("unsupported input for usp_merge_heads CUDA")
    return _usp_merge_heads_custom_op(x)


def usp_merge_heads(x: torch.Tensor) -> torch.Tensor:
    """Merge Ulysses output heads with an exact eager fallback.

    The backend selection lives here so callers only express the layout
    transformation. Unsupported devices, layouts, and compiled regions retain
    the original PyTorch operation.
    """
    if not torch.compiler.is_compiling() and can_use_usp_merge_heads(x):
        return _usp_merge_heads_cuda(x)
    return x.permute(2, 1, 0, 3, 4).contiguous()

View on GitHub (pinned to 0132848349)

Solutions

1. Pre-check with can_use_usp_merge_heads(x) and fall back to x.permute(2,1,0,3,4).contiguous()
2. Prefer the public usp_merge_heads(x) API, which handles the fallback automatically
3. Ensure x is a 5D contiguous CUDA tensor in fp16/bf16/fp32
4. Fix upstream reshape/dtype logic if tensors arrive malformed

Example fix

# before
y = _usp_merge_heads_cuda(x)
# after
y = usp_merge_heads(x)  # has eager fallback
# or: y = _usp_merge_heads_cuda(x) if can_use_usp_merge_heads(x) else x.permute(2,1,0,3,4).contiguous()

from sglang.kernels.ops.diffusion.layout.usp_relayout_jit import can_use_usp_merge_heads
if not can_use_usp_merge_heads(x):
    y = x.permute(2, 1, 0, 3, 4).contiguous()

def can_merge(x) -> bool:
    return can_use_usp_merge_heads(x)

try:
    y = _usp_merge_heads_cuda(x)
except RuntimeError:
    y = x.permute(2, 1, 0, 3, 4).contiguous()

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• unsupported input for modulate_scale_shift CUDA
• unsupported input for causal Conv3D cat/pad CUDA
• q, k, and v must be CUDA tensors
• Unsupported usp_merge_heads dtype: {dtype}
• unsupported input for LTX2 QKNorm split-RoPE CUDA

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