sgl-project/sglang · error · RuntimeError

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

Error message

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

What it means

l2norm_fwd uses a Triton fused kernel that loads an entire feature row into one block; the maximum fused size is 65536 bytes divided by element size. If the feature dimension D exceeds that (e.g. >16384 for fp32 or >32768 for bf16), no kernel path exists and it raises RuntimeError.

Source

Thrown at python/sglang/kernels/ops/attention/fla/l2norm.py:90

def l2norm_fwd(
    x: torch.Tensor, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None
):
    x_shape_og = x.shape
    x = x.view(-1, x.shape[-1])
    # allocate output
    if output_dtype is None:
        y = torch.empty_like(x)
    else:
        y = torch.empty_like(x, dtype=output_dtype)
    assert y.stride(-1) == 1
    T, D = x.shape[0], x.shape[-1]
    # rstd = torch.empty((T,), dtype=torch.float32, device=x.device)
    # Less than 64KB per feature: enqueue fused kernel
    MAX_FUSED_SIZE = 65536 // x.element_size()
    BD = min(MAX_FUSED_SIZE, triton.next_power_of_2(D))
    if D > BD:
        raise RuntimeError("This layer doesn't support feature dim >= 64KB.")

    if D <= 512:

        def grid(meta):
            return (triton.cdiv(T, meta["BT"]),)

        l2norm_fwd_kernel[grid](
            x,
            y,
            eps,
            T=T,
            D=D,
            BD=BD,
            BT=16,
            num_warps=8,
            num_stages=3,
        )
    else:

View on GitHub (pinned to 0132848349)

Solutions

  1. Reduce the feature dimension to below 64KB per row (e.g. <16384 for fp32)
  2. Split the tensor along the feature dim, l2norm each chunk, and concatenate
  3. Use a smaller dtype (bf16/fp16 halves the byte limit)

Example fix

// before
y = l2norm_fwd(x)  # x: [T, 65536] fp32 -> raises
// after
x = x.to(torch.bfloat16)  # or split feature dim
y = l2norm_fwd(x)
Defensive patterns

Strategy: validation

Validate before calling

assert x.shape[-1] * x.element_size() <= 65536, f'feature dim {x.shape[-1]} exceeds 64KB fused-kernel limit'

Try / catch

try:\n    y = l2norm_fwd(x)\nexcept RuntimeError:\n    x = x.to(torch.bfloat16); y = l2norm_fwd(x)

Prevention

When it happens

Trigger: Calling l2norm_fwd (via chunk_kda, forward, or cutedsl/flashinfer wrappers) with x.shape[-1] * element_size > 65536 bytes.

Common situations: Using an unusually large head_dim or hidden dim (e.g. D=32768 in fp32, D=65536+ in bf16) in KDA/attention normalization layers.

Related errors


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