sgl-project/sglang · error · RuntimeError

This layer norm doesn't support feature dim >= 64KB.

Error message

This layer norm doesn't support feature dim >= 64KB.

What it means

_layer_norm_fwd (backing rms_norm_gated) fuses normalization over a group into one Triton block limited to 64KB per feature (65536 // element_size). If group_size exceeds MAX_FUSED_SIZE there is no fallback path, so it raises RuntimeError.

Source

Thrown at python/sglang/kernels/ops/attention/fla/layernorm_gated.py:250

        assert bias.stride(-1) == 1
        assert bias.shape == (N,)
    # allocate output
    if out is not None:
        assert out.shape == x.shape
    else:
        out = torch.empty_like(x)
    assert out.stride(-1) == 1
    mean = (
        torch.empty((ngroups * M,), dtype=torch.float32, device=x.device)
        if not is_rms_norm
        else None
    )
    rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device)
    # Less than 64KB per feature: enqueue fused kernel
    MAX_FUSED_SIZE = 65536 // x.element_size()
    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size))
    if group_size > BLOCK_N:
        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
    # heuristics for number of warps
    num_warps = min(max(BLOCK_N // 256, 1), 8)
    # Calculate rows per block based on SM count
    rows_per_block = calc_rows_per_block(M, x.device)
    # Update grid to use rows_per_block
    grid = (cdiv(M, rows_per_block), ngroups)
    pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
    # Workaround for PyTorch <= 2.12: torch.xpu.device is not Dynamo-compatible
    # in that release — it creates a DynamoConfigPatchProxy that
    # SourcelessBuilder cannot wrap, causing a hard error under
    # torch.compile(fullgraph=True).  The device context is a functional no-op
    # for Triton kernel launches (device is determined by the tensor, not the
    # surrounding context), so we simply skip it when Dynamo is tracing.
    # PyTorch main already has the proper fix (XPUDeviceVariable registered in
    # torch/_dynamo/variables/ctx_manager.py analogous to CUDADeviceVariable).
    # TODO: remove this branch once we upgrade from PyTorch 2.12.
    device_ctx = (
        nullcontext()

View on GitHub (pinned to 0132848349)

Solutions

1. Cast x to bf16/fp16 to double the allowed feature size
2. Reduce hidden dim or increase ngroups so group_size shrinks
3. Split along the feature dimension, normalize chunks, reassemble

Example fix

// before
y = rms_norm_gated(x, g, weight, bias)  # fp32, group_size 20000 -> raises
// after
y = rms_norm_gated(x.to(torch.bfloat16), g, weight, bias)

group_size = x.shape[-1] // ngroups\nassert group_size * x.element_size() <= 65536, 'rms_norm_gated group size exceeds 64KB'

try:\n    y = rms_norm_gated(x, g, w)\nexcept RuntimeError:\n    y = rms_norm_gated(x.to(torch.bfloat16), g, w)

When it happens

Trigger: Calling rms_norm_gated on x where group_size = x.shape[-1] // ngroups exceeds 65536 // element_size bytes, e.g. fp32 with group_size > 16384.

Common situations: Gated RMS norm layers (gated deltanet / KDA models) with very wide hidden dims or fp32 dtype during debugging/export.

Related errors

• This layer doesn't support feature dim >= 64KB.
• This layer norm doesn't support feature dim >= 64KB.
• `mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim}).
• {self._op_label()}: no triton backend
• Unexpected A_log shape: {A_log.shape}; expected numel={HV}

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