{"record":{"id":"4662b4d3793e3bb0","repo":"sgl-project/sglang","slug":"unsupported-usp-merge-heads-dtype-dtype","errorCode":null,"errorMessage":"Unsupported usp_merge_heads dtype: {dtype}","messagePattern":"Unsupported usp_merge_heads dtype: (.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py","lineNumber":20,"sourceCode":"\nfrom typing import TYPE_CHECKING\n\nimport torch\n\nfrom sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args\nfrom sglang.srt.utils.custom_op import register_custom_op\n\nif TYPE_CHECKING:\n    from tvm_ffi.module import Module\n\n\n_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16, torch.float32)\n\n\n@cache_once\ndef _jit_usp_relayout_module(dtype: torch.dtype) -> Module:\n    if dtype not in _SUPPORTED_DTYPES:\n        raise RuntimeError(f\"Unsupported usp_merge_heads dtype: {dtype}\")\n    args = make_cpp_args(dtype)\n    return load_jit(\n        \"diffusion_usp_relayout\",\n        *args,\n        cuda_files=[\"diffusion/usp_relayout.cuh\"],\n        cuda_wrappers=[\n            (\n                \"usp_merge_heads\",\n                \"usp_relayout::\" f\"UspMergeHeadsKernel<{args}>::run\",\n            ),\n        ],\n    )\n\n\ndef _fake_merge_heads(x: torch.Tensor) -> torch.Tensor:\n    world, seq, batch, h_local, head_dim = x.shape\n    return x.new_empty((batch, seq, world, h_local, head_dim))\n","sourceCodeStart":2,"sourceCodeEnd":38,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/diffusion/layout/usp_relayout_jit.py#L2-L38","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Cast x to bf16/fp16/fp32 before calling usp_merge_heads","Check the upstream attention output dtype and fix the conversion that produced an unsupported one","Use the eager fallback path (x.permute(2,1,0,3,4).contiguous()) for unsupported dtypes","Extend _SUPPORTED_DTYPES and the .cuh kernel if a new dtype is genuinely required"],"exampleFix":"# before\ny = _usp_merge_heads_cuda(x_fp8_dequantized_as_fp64)\n# after\nx = x.to(torch.bfloat16)\ny = _usp_merge_heads_cuda(x)","handlingStrategy":"type-guard","validationCode":"if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):\n    x = x.to(torch.bfloat16)","typeGuard":"def relayout_dtype_ok(x: torch.Tensor) -> bool:\n    return x.dtype in (torch.float16, torch.bfloat16, torch.float32)","tryCatchPattern":null,"preventionTips":["Keep Ulysses path tensors in bf16","Use the public usp_merge_heads wrapper with eager fallback"],"tags":["dtype","jit","ulysses","cuda"],"backgroundTag":"unsupported-dtype","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}