{"record":{"id":"41be1c95ab942595","repo":"sgl-project/sglang","slug":"this-layer-norm-doesn-t-support-feature-dim-64k","errorCode":null,"errorMessage":"This layer norm doesn't support feature dim >= 64KB.","messagePattern":"This layer norm doesn't support feature dim >= 64KB\\.","errorType":"validation","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_norm_gate.py","lineNumber":222,"sourceCode":"    # allocate output\n    y = x if out_dtype is None else torch.empty_like(x, dtype=out_dtype)\n    if residual is not None or (\n        residual_dtype is not None and residual_dtype != x.dtype\n    ):\n        residual_out = torch.empty(T, D, device=x.device, dtype=residual_dtype)\n    else:\n        residual_out = None\n    mean = (\n        torch.empty((T,), dtype=torch.float, device=x.device)\n        if not is_rms_norm\n        else None\n    )\n    rstd = torch.empty((T,), dtype=torch.float, device=x.device)\n    # Less than 64KB per feature: enqueue fused kernel\n    MAX_FUSED_SIZE = 65536 // x.element_size()\n    BD = min(MAX_FUSED_SIZE, next_power_of_2(D))\n    if D > BD:\n        raise RuntimeError(\"This layer norm doesn't support feature dim >= 64KB.\")\n    # heuristics for number of warps\n\n    if D <= 512:\n        BT = 32\n        pdl_kwargs = (\n            {\"USE_GDC\": True, \"launch_pdl\": True} if is_arch_support_pdl() else {}\n        )\n        layer_norm_gated_fwd_kernel[(cdiv(T, BT),)](\n            x=x,\n            g=g,\n            y=y,\n            w=weight,\n            b=bias,\n            residual=residual,\n            residual_out=residual_out,\n            mean=mean,\n            rstd=rstd,\n            eps=eps,","sourceCodeStart":204,"sourceCodeEnd":240,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py#L204-L240","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Check the last dimension of x; if it accidentally concatenates heads, split to the true per-head hidden size","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","Reduce D by fixing the model config (hidden_size / head layout) if it's a misconfiguration"],"exampleFix":"# before\ny = rms_norm_gated(x.view(T, -1), gate, weight, eps)  # D = H*dh too large\n# after\ny = rms_norm_gated(x.reshape(T, H, dh), gate, weight, eps)  # normalized per-head dim dh","handlingStrategy":"validation","validationCode":"D = x.shape[-1]\nassert D * x.element_size() <= 65536, f\"feature dim {D} too large for fused gated norm\"","typeGuard":"def norm_dim_supported(x: torch.Tensor) -> bool:\n    return x.shape[-1] * x.element_size() <= 65536","tryCatchPattern":"try:\n    y = rms_norm_gated(x, gate, weight, eps)\nexcept RuntimeError:\n    y = fallback_rms_norm_gated(x, gate, weight, eps)","preventionTips":["Compute 65536 // element_size once per dtype as the max supported D","Split concatenated head dims into per-head slices before normalization"],"tags":["fla","rms-norm","triton","feature-dim-limit"],"backgroundTag":"kernel-dimension-limit","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}