sgl-project/sglang · error · RuntimeError

Unsupported dtype {dtype}. Supported: float16, bfloat16, flo

Error message

Unsupported dtype {dtype}. Supported: float16, bfloat16, float32

What it means

The causal-conv1d JIT module factory only compiles kernels for float16, bfloat16, and float32. Any other dtype (float64, int, etc.) raises before TVM-FFI module build. Both causal_conv1d_fwd and causal_conv1d_update route through this cached factory, so the error surfaces on first use of a given dtype.

Source

Thrown at python/sglang/kernels/ops/mamba/causal_conv1d.py:28

import torch

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

if TYPE_CHECKING:
    from tvm_ffi.module import Module


@cache_once
def _jit_causal_conv1d_module(dtype: torch.dtype) -> Module:
    if dtype not in (torch.float16, torch.bfloat16, torch.float32):
        raise RuntimeError(
            f"Unsupported dtype {dtype}. Supported: float16, bfloat16, float32"
        )
    # The AOT wheel ships an SM90 build compiled with `-use_fast_math` and a
    # precise-math build for every other arch; match that split so the SiLU
    # epilogue keeps producing the same bits as the op being replaced.
    arch = get_jit_cuda_arch()
    use_fast_math = (arch.major, arch.minor) == (9, 0)
    math_mode = "fast_math" if use_fast_math else "precise_math"
    args = make_cpp_args(dtype)
    return load_jit(
        "causal_conv1d",
        math_mode,
        *args,
        cuda_files=["mamba/causal_conv1d.cuh"],
        cuda_wrappers=[
            ("causal_conv1d_fwd", f"causal_conv1d_fwd<{args}>"),
            ("causal_conv1d_update", f"causal_conv1d_update<{args}>"),
        ],

View on GitHub (pinned to 0132848349)

Solutions

  1. Cast the conv inputs and states to one of float16/bfloat16/float32 (usually the model dtype, e.g. .to(torch.bfloat16))
  2. Fix the tensor creation site (from_numpy(...).float() or dtype= in torch.randn) rather than casting at call time
  3. Keep hidden states and conv states in the same dtype as the model to avoid downstream mismatches

Example fix

// before
y = causal_conv1d_fwd(torch.from_numpy(x_np), weight, ...)

// after
y = causal_conv1d_fwd(torch.from_numpy(x_np).to(torch.bfloat16), weight, ...)
Defensive patterns

Strategy: type-guard

Validate before calling

assert x.dtype in (torch.float16, torch.bfloat16, torch.float32), f'bad dtype {x.dtype}'

Type guard

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

Prevention

When it happens

Trigger: Calling causal_conv1d_fwd or causal_conv1d_update with x/weight in float64 (numpy default), or an integer/other dtype tensor.

Common situations: Mamba/conv states created via torch.from_numpy(...) which defaults to float64; test fixtures with wrong dtype; mixed input pipelines after .double() debugging.

Related errors


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