sgl-project/sglang · error · RuntimeError

Unsupported modulate_scale_shift dtype: {dtype}

Error message

Unsupported modulate_scale_shift dtype: {dtype}

What it means

The JIT-compiled modulate_scale_shift kernel is only built for the dtypes listed in its _SUPPORTED_DTYPES (fp16/bf16/fp32 family). _jit_modulate_scale_shift_module raises when asked to compile for anything else.

Source

Thrown at python/sglang/kernels/ops/diffusion/modulate/modulate_scale_shift_jit.py:32

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

logger = logging.getLogger(__name__)


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


def _fake_impl(
    x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor
) -> torch.Tensor:
    return torch.empty_like(x)

View on GitHub (pinned to 0132848349)

Solutions

1. Cast x, scale, shift to a supported dtype (bf16/fp16/fp32) before calling
2. Trace where the unsupported dtype was introduced (often a .double() or float promotion upstream)
3. Use the eager expression x * (1 + scale[:, None]) + shift[:, None] as fallback
4. Extend _SUPPORTED_DTYPES plus the .cuh kernel if a new dtype is required

Example fix

# before
y = modulate_scale_shift_cuda(x_fp64, s, b)
# after
x = x.to(torch.bfloat16); s = s.to(torch.bfloat16); b = b.to(torch.bfloat16)
y = modulate_scale_shift_cuda(x, s, b)

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

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

When it happens

Trigger: Calling modulate_scale_shift_cuda with x/scale/shift in an unsupported dtype (fp64, int, fp8) — the dtype propagates to the JIT module factory and fails the check.

Common situations: Diffusion model runs where activations or AdaLN scale/shift tensors end up in fp64 (e.g. after operations promoting precision) or unquantized fp8 paths; dtype drift between x and scale/shift.

Related errors

• Unsupported usp_merge_heads 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/53da0697e92c5e33. Report an issue: GitHub.