{"record":{"id":"58cbe2add59d2653","repo":"sgl-project/sglang","slug":"this-layer-norm-doesn-t-support-feature-dim-64k-58cbe2","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/layernorm_gated.py","lineNumber":250,"sourceCode":"        assert bias.stride(-1) == 1\n        assert bias.shape == (N,)\n    # allocate output\n    if out is not None:\n        assert out.shape == x.shape\n    else:\n        out = torch.empty_like(x)\n    assert out.stride(-1) == 1\n    mean = (\n        torch.empty((ngroups * M,), dtype=torch.float32, device=x.device)\n        if not is_rms_norm\n        else None\n    )\n    rstd = torch.empty((ngroups * M,), dtype=torch.float32, device=x.device)\n    # Less than 64KB per feature: enqueue fused kernel\n    MAX_FUSED_SIZE = 65536 // x.element_size()\n    BLOCK_N = min(MAX_FUSED_SIZE, triton.next_power_of_2(group_size))\n    if group_size > BLOCK_N:\n        raise RuntimeError(\"This layer norm doesn't support feature dim >= 64KB.\")\n    # heuristics for number of warps\n    num_warps = min(max(BLOCK_N // 256, 1), 8)\n    # Calculate rows per block based on SM count\n    rows_per_block = calc_rows_per_block(M, x.device)\n    # Update grid to use rows_per_block\n    grid = (cdiv(M, rows_per_block), ngroups)\n    pdl_kwargs = {\"USE_GDC\": True, \"launch_pdl\": True} if is_arch_support_pdl() else {}\n    # Workaround for PyTorch <= 2.12: torch.xpu.device is not Dynamo-compatible\n    # in that release — it creates a DynamoConfigPatchProxy that\n    # SourcelessBuilder cannot wrap, causing a hard error under\n    # torch.compile(fullgraph=True).  The device context is a functional no-op\n    # for Triton kernel launches (device is determined by the tensor, not the\n    # surrounding context), so we simply skip it when Dynamo is tracing.\n    # PyTorch main already has the proper fix (XPUDeviceVariable registered in\n    # torch/_dynamo/variables/ctx_manager.py analogous to CUDADeviceVariable).\n    # TODO: remove this branch once we upgrade from PyTorch 2.12.\n    device_ctx = (\n        nullcontext()","sourceCodeStart":232,"sourceCodeEnd":268,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/layernorm_gated.py#L232-L268","documentation":"_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.","triggerScenarios":"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.","commonSituations":"Gated RMS norm layers (gated deltanet / KDA models) with very wide hidden dims or fp32 dtype during debugging/export.","solutions":["Cast x to bf16/fp16 to double the allowed feature size","Reduce hidden dim or increase ngroups so group_size shrinks","Split along the feature dimension, normalize chunks, reassemble"],"exampleFix":"// before\ny = rms_norm_gated(x, g, weight, bias)  # fp32, group_size 20000 -> raises\n// after\ny = rms_norm_gated(x.to(torch.bfloat16), g, weight, bias)","handlingStrategy":"validation","validationCode":"group_size = x.shape[-1] // ngroups\\nassert group_size * x.element_size() <= 65536, 'rms_norm_gated group size exceeds 64KB'","typeGuard":null,"tryCatchPattern":"try:\\n    y = rms_norm_gated(x, g, w)\\nexcept RuntimeError:\\n    y = rms_norm_gated(x.to(torch.bfloat16), g, w)","preventionTips":["Compute group_size in bytes before calling","Cast to half precision for wide hidden dims"],"tags":["rms-norm","triton","feature-dim","gated","kda"],"backgroundTag":"kernel-feature-dim-limit-exceeded","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}