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_gated_fwd uses a single Triton fused kernel whose block size BD is capped at MAX_FUSED_SIZE = 65536 / element_size, i.e. 64KB of registers/shared memory per feature row. If the normalized feature dimension D exceeds that cap (e.g. D > 16384 for fp32 or > 8192 for bf16... specifically 65536//itemsize), the function raises because no fused block config can cover the row.

Source

Thrown at python/sglang/kernels/ops/attention/fla/fused_norm_gate.py:222

    # allocate output
    y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype)
    if residual is not None or (
        residual_dtype is not None and residual_dtype != x.dtype
    ):
        residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)
    else:
        residual_out = None
    mean = (
        torch.empty((T,), dtype=torch.float, device=x.device)
        if not is_rms_norm
        else None
    )
    rstd = torch.empty((T,), dtype=torch.float, device=x.device)
    # Less than 64KB per feature: enqueue fused kernel
    MAX_FUSED_SIZE = 65536 // x.element_size()
    BD = min(MAX_FUSED_SIZE, next_power_of_2(D))
    if D > BD:
        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
    # heuristics for number of warps

    if D <= 512:
        BT = 32
        pdl_kwargs = (
            {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
        )
        layer_norm_gated_fwd_kernel[(cdiv(T, BT),)](
            x=x,
            g=g,
            y=y,
            w=weight,
            b=bias,
            residual=residual,
            residual_out=residual_out,
            mean=mean,
            rstd=rstd,
            eps=eps,

View on GitHub (pinned to 0132848349)

Solutions

1. Check the last dimension of x; if it accidentally concatenates heads, split to the true per-head hidden size
2. If D is genuinely huge, use a non-fused normalization path (e.g. torch.nn.functional.rms_norm applied manually with the gate) instead of this kernel
3. Reduce D by fixing the model config (hidden_size / head layout) if it's a misconfiguration

Example fix

# before
y = rms_norm_gated(x.view(T, -1), gate, weight, eps)  # D = H*dh too large
# after
y = rms_norm_gated(x.reshape(T, H, dh), gate, weight, eps)  # normalized per-head dim dh

D = x.shape[-1]
assert D * x.element_size() <= 65536, f"feature dim {D} too large for fused gated norm"

def norm_dim_supported(x: torch.Tensor) -> bool:
    return x.shape[-1] * x.element_size() <= 65536

try:
    y = rms_norm_gated(x, gate, weight, eps)
except RuntimeError:
    y = fallback_rms_norm_gated(x, gate, weight, eps)

When it happens

Trigger: Calling rms_norm_gated / layer_norm_gated_fwd on a hidden state whose last dimension D has next_power_of_2(D) * element_size > 65536 bytes; e.g. a hidden_size of 32768 with bfloat16 (65536 bytes) or 16384 with float32.

Common situations: Testing the gated RMSNorm used by Qwen3-Next/GatedDeltaNet MLPs with an abnormally large head_dim or hidden dim; accidentally passing the whole (T, H*D) concatenated projection instead of per-head slices; dtype changes (fp32 debug runs) halving the allowed D.

Related errors

• Unsupported activation: {self.activation}
• This layer norm doesn't support feature dim >= 64KB.
• Triton is not supported on current platform, roll back to CP
• {self._op_label()}: no triton backend
• sparse_attn_v4_paged_decode expects fp16/bf16 q, got {q.dtyp

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