sgl-project/sglang · error · RuntimeError

Unsupported usp_merge_heads dtype: {dtype}

Error message

Unsupported usp_merge_heads dtype: {dtype}

What it means

The JIT-compiled usp_relayout kernel (used by usp_merge_heads on the Ulysses output path) is only built for float16, bfloat16, and float32 (see _SUPPORTED_DTYPES). Requesting any other dtype raises this RuntimeError at module-build time.

Source

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

from typing import TYPE_CHECKING

import torch

from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
from sglang.srt.utils.custom_op import register_custom_op

if TYPE_CHECKING:
    from tvm_ffi.module import Module


_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)


@cache_once
def _jit_usp_relayout_module(dtype: torch.dtype) -> Module:
    if dtype not in _SUPPORTED_DTYPES:
        raise RuntimeError(f"Unsupported usp_merge_heads dtype: {dtype}")
    args = make_cpp_args(dtype)
    return load_jit(
        "diffusion_usp_relayout",
        *args,
        cuda_files=["diffusion/usp_relayout.cuh"],
        cuda_wrappers=[
            (
                "usp_merge_heads",
                "usp_relayout::" f"UspMergeHeadsKernel<{args}>::run",
            ),
        ],
    )


def _fake_merge_heads(x: torch.Tensor) -> torch.Tensor:
    world, seq, batch, h_local, head_dim = x.shape
    return x.new_empty((batch, seq, world, h_local, head_dim))

View on GitHub (pinned to 0132848349)

Solutions

1. Cast x to bf16/fp16/fp32 before calling usp_merge_heads
2. Check the upstream attention output dtype and fix the conversion that produced an unsupported one
3. Use the eager fallback path (x.permute(2,1,0,3,4).contiguous()) for unsupported dtypes
4. Extend _SUPPORTED_DTYPES and the .cuh kernel if a new dtype is genuinely required

Example fix

# before
y = _usp_merge_heads_cuda(x_fp8_dequantized_as_fp64)
# after
x = x.to(torch.bfloat16)
y = _usp_merge_heads_cuda(x)

if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):
    x = x.to(torch.bfloat16)

def relayout_dtype_ok(x: torch.Tensor) -> bool:
    return x.dtype in (torch.float16, torch.bfloat16, torch.float32)

When it happens

Trigger: Calling _usp_merge_heads_cuda / usp_merge_heads with a tensor in fp64, fp8, or an integer dtype, which propagates the dtype into _jit_usp_relayout_module and fails the _SUPPORTED_DTYPES check.

Common situations: Running Ulysses sequence-parallel attention with activations cast to fp8 or fp64; model checkpoints or LoRA adapters that leave hidden states in an unexpected precision; new quantization paths not yet supported by the relayout kernel.

Related errors

• Unsupported modulate_scale_shift dtype: {dtype}
• Unsupported residual_gate_add dtype: {dtype}
• SplitKV partial output (mO) must be Float32
• All tensors must have the same data type
• Only Float16 or BFloat16 is supported

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