{"record":{"id":"53da0697e92c5e33","repo":"sgl-project/sglang","slug":"unsupported-modulate-scale-shift-dtype-dtype","errorCode":null,"errorMessage":"Unsupported modulate_scale_shift dtype: {dtype}","messagePattern":"Unsupported modulate_scale_shift dtype: (.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/diffusion/modulate/modulate_scale_shift_jit.py","lineNumber":32,"sourceCode":"\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)\n_ALIGN_BYTES = 16\n_FAILED_RUNTIME_KEYS: set[tuple[int | None, torch.dtype]] = set()\n\nlogger = logging.getLogger(__name__)\n\n\n@cache_once\ndef _jit_modulate_scale_shift_module(dtype: torch.dtype) -> Module:\n    if dtype not in _SUPPORTED_DTYPES:\n        raise RuntimeError(f\"Unsupported modulate_scale_shift dtype: {dtype}\")\n    args = make_cpp_args(dtype)\n    return load_jit(\n        \"diffusion_modulate_scale_shift\",\n        *args,\n        cuda_files=[\"diffusion/modulate_scale_shift.cuh\"],\n        cuda_wrappers=[\n            (\n                \"modulate_scale_shift\",\n                f\"modulate_scale_shift::ModulateScaleShiftKernel<{args}>::run\",\n            ),\n        ],\n    )\n\n\ndef _fake_impl(\n    x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor\n) -> torch.Tensor:\n    return torch.empty_like(x)","sourceCodeStart":14,"sourceCodeEnd":50,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/diffusion/modulate/modulate_scale_shift_jit.py#L14-L50","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Cast x, scale, shift to a supported dtype (bf16/fp16/fp32) before calling","Trace where the unsupported dtype was introduced (often a .double() or float promotion upstream)","Use the eager expression x * (1 + scale[:, None]) + shift[:, None] as fallback","Extend _SUPPORTED_DTYPES plus the .cuh kernel if a new dtype is required"],"exampleFix":"# before\ny = modulate_scale_shift_cuda(x_fp64, s, b)\n# after\nx = x.to(torch.bfloat16); s = s.to(torch.bfloat16); b = b.to(torch.bfloat16)\ny = modulate_scale_shift_cuda(x, s, b)","handlingStrategy":"type-guard","validationCode":"if x.dtype not in (torch.float16, torch.bfloat16, torch.float32):\n    x, scale, shift = x.to(torch.bfloat16), scale.to(torch.bfloat16), shift.to(torch.bfloat16)","typeGuard":"def mod_dtype_ok(x: torch.Tensor) -> bool:\n    return x.dtype in (torch.float16, torch.bfloat16, torch.float32)","tryCatchPattern":null,"preventionTips":["Avoid fp64 promotion upstream (check norm outputs)","Route unsupported dtypes through the eager modulate formula"],"tags":["dtype","jit","modulate","cuda"],"backgroundTag":"unsupported-dtype","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}