sgl-project/sglang · error · RuntimeError

rmsnorm_hf: unsupported hidden_size={hidden_size} (must be a

Error message

rmsnorm_hf: unsupported hidden_size={hidden_size} (must be a multiple of {_WARP_SIZE} in [{_WARP_SIZE}, {_CTA_BLOCK_SIZE}) or a multiple of {_CTA_BLOCK_SIZE})

What it means

rmsnorm_hf rejects hidden sizes that are not a multiple of the warp size (32) within [32, CTA_BLOCK_SIZE) or a multiple of the CTA block size. The Triton/CUDA kernel vectorizes loads over the hidden dimension, so arbitrary sizes (e.g. 4095, 1000) cannot be handled. This is a shape-contract error raised before any kernel launch.

Source

Thrown at python/sglang/kernels/ops/layernorm/rmsnorm_hf.py:68

    input: torch.Tensor,
    weight: torch.Tensor,
    eps: float = 1e-6,
    out: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """RMSNorm: ``out = weight * cast_dtype(rsqrt(mean(x^2) + eps) * x)``.

    ``input`` must be 2D ``(num_tokens, hidden_size)``; callers with
    higher-rank tensors should reshape first. ``hidden_size`` must satisfy
    :func:`is_supported_rmsnorm_hf_hidden_size`. Empty inputs return an empty
    output without launching the kernel.
    """
    if input.dtype not in (torch.float16, torch.bfloat16):
        raise RuntimeError(f"rmsnorm_hf: input must be fp16 or bf16, got {input.dtype}")
    if input.dim() != 2:
        raise RuntimeError(f"rmsnorm_hf: input must be 2D, got {input.dim()}D")
    hidden_size = input.size(-1)
    if not is_supported_rmsnorm_hf_hidden_size(hidden_size):
        raise RuntimeError(
            f"rmsnorm_hf: unsupported hidden_size={hidden_size} "
            f"(must be a multiple of {_WARP_SIZE} in [{_WARP_SIZE}, {_CTA_BLOCK_SIZE}) "
            f"or a multiple of {_CTA_BLOCK_SIZE})"
        )
    if out is None:
        out = torch.empty_like(input)
    if input.numel() == 0:
        return out
    module = _jit_rmsnorm_hf_module(hidden_size, input.dtype)
    module.rmsnorm_hf(input, weight, out, eps)
    return out

View on GitHub (pinned to 0132848349)

Solutions

  1. Check input.size(-1): round the model hidden size to a multiple of 32 (or the kernel's CTA block size); most real models (4096, 5120, 7168, 8192) already qualify
  2. Pad the last dimension to the next multiple of 32 with torch.nn.functional.pad and slice the output back
  3. If the size is genuinely unsupported, use the generic torch RMSNorm / HF implementation instead of this fused op

Example fix

// before
out = rmsnorm_hf(x, weight, eps)  # x: (N, 4095)

// after
pad = (-x.shape[-1]) % 32
out = rmsnorm_hf(torch.nn.functional.pad(x, (0, pad)), weight, eps)[:, : x.shape[-1]]
Defensive patterns

Strategy: validation

Validate before calling

from sglang.kernels.ops.layernorm.rmsnorm_hf import is_supported_rmsnorm_hf_hidden_size
assert is_supported_rmsnorm_hf_hidden_size(x.size(-1)), f"hidden_size {x.size(-1)} unsupported"

Type guard

def is_rmsnorm_hf_size_ok(x: torch.Tensor) -> bool:
    return is_supported_rmsnorm_hf_hidden_size(x.size(-1))

Prevention

When it happens

Trigger: Calling rmsnorm_hf(input) with input.size(-1) not a multiple of 32 (when < _CTA_BLOCK_SIZE) or not a multiple of _CTA_BLOCK_SIZE (when larger). Direct calls or via tests with odd hidden sizes.

Common situations: Porting a model whose hidden_size is not a power-of-two multiple of 32; slicing the last dim (e.g. x[:, :5120] on a 5122-wide tensor); unit tests iterating arbitrary sizes.

Related errors


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