sgl-project/sglang · error · RuntimeError

Unsupported residual_gate_add dtype: {dtype}

Error message

Unsupported residual_gate_add dtype: {dtype}

What it means

The JIT-compiled residual_gate_add kernel supports only the dtypes in its _SUPPORTED_DTYPES (fp16/bf16/fp32). _jit_residual_gate_add_module raises when the requested dtype is outside that set, preventing compilation of an unsupported variant.

Source

Thrown at python/sglang/kernels/ops/diffusion/modulate/residual_gate_add_jit.py:25

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)
_BIT_EXACT_DTYPES = (torch.float16, torch.bfloat16)
_FAILED_RUNTIME_KEYS: set[tuple[int | None, torch.dtype]] = set()

logger = logging.getLogger(__name__)


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


def _fake_impl(
    residual: torch.Tensor, update: torch.Tensor, gate: torch.Tensor
) -> torch.Tensor:
    return torch.empty_like(residual)

View on GitHub (pinned to 0132848349)

Solutions

1. Cast the residual and gated tensors to bf16/fp16/fp32 before the fused call
2. Use the eager expression (residual + gate * x style) as fallback
3. Audit the preceding norm/modulate ops for unwanted dtype promotion
4. Extend the JIT kernel and _SUPPORTED_DTYPES if a new dtype is truly needed

Example fix

# before
y = _residual_gate_add_custom_op(x_fp64, res_fp64)
# after
x = x.to(torch.bfloat16); res = res.to(torch.bfloat16)
y = _residual_gate_add_custom_op(x, res)

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

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

When it happens

Trigger: Calling _residual_gate_add_custom_op (via the residual gate add CUDA path) with tensors in fp64, fp8, or integer dtype; the dtype is used as the JIT specialization key and rejected.

Common situations: Diffusion transformer residual paths where activations were promoted to float64 (e.g. by a norm in fp32 not cast back), or partially-quantized fp8 pipelines routing through the fused residual add.

Related errors

• Unsupported usp_merge_heads dtype: {dtype}
• Unsupported modulate_scale_shift 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/393d1e86409ab4b0. Report an issue: GitHub.