[Apple Silicon] Add a custom Metal RMSNorm kernel - #30163
SasankYadati wants to merge 6 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a custom Metal RMSNorm kernel for Apple Silicon (MLX backend) to optimize performance, including the Metal shader, C++ bindings, Python wrappers, model patching integration, correctness tests, and benchmarks. The review feedback identifies a critical memory leak in the C++ nanobind wrapper due to improper placement new usage, the need for __getattr__ forwarding in the model wrapper, and missing row-contiguity checks on inputs and weights to prevent undefined behavior or crashes.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| auto* dst = nb::inst_ptr<array>(py_obj); | ||
| new (dst) array(std::move(outs[0])); | ||
| nb::inst_mark_ready(py_obj); |
There was a problem hiding this comment.
Using placement new on dst (which points to an already fully constructed array object created by py_array_type(0)) without calling its destructor will leak the resources of the old array (such as its internal shared pointers). Since this is called on every forward pass of every layer, it will cause a severe memory leak during inference. Using standard move assignment is much safer and avoids any memory leaks.
| auto* dst = nb::inst_ptr<array>(py_obj); | |
| new (dst) array(std::move(outs[0])); | |
| nb::inst_mark_ready(py_obj); | |
| auto* dst = nb::inst_ptr<array>(py_obj); | |
| *dst = std::move(outs[0]); |
| def __init__(self, inner: nn.Module, kernel: Callable): | ||
| super().__init__() | ||
| # Bypass nn.Module.__setattr__ so these are plain attributes, not | ||
| # re-registered submodules/params (mirrors MLXAttentionWrapper). | ||
| object.__setattr__(self, "_inner", inner) | ||
| object.__setattr__(self, "_kernel", kernel) |
There was a problem hiding this comment.
MLXRMSNormWrapper wraps nn.RMSNorm but does not forward attribute lookups (such as weight or eps) to the wrapped module. This can cause AttributeError if other parts of the model or SGLang attempt to access attributes of the norm module. Implementing __getattr__ to forward lookups to self._inner makes the wrapper transparent.
| def __init__(self, inner: nn.Module, kernel: Callable): | |
| super().__init__() | |
| # Bypass nn.Module.__setattr__ so these are plain attributes, not | |
| # re-registered submodules/params (mirrors MLXAttentionWrapper). | |
| object.__setattr__(self, "_inner", inner) | |
| object.__setattr__(self, "_kernel", kernel) | |
| def __init__(self, inner: nn.Module, kernel: Callable): | |
| super().__init__() | |
| # Bypass nn.Module.__setattr__ so these are plain attributes, not | |
| # re-registered submodules/params (mirrors MLXAttentionWrapper). | |
| object.__setattr__(self, "_inner", inner) | |
| object.__setattr__(self, "_kernel", kernel) | |
| def __getattr__(self, name: str) -> Any: | |
| return getattr(self._inner, name) |
| def _supported(self, x: mx.array) -> bool: | ||
| """Whether the custom kernel can handle this call; else we fall back.""" | ||
| w = self._inner.weight | ||
| return ( | ||
| x.dtype in (mx.float16, mx.bfloat16, mx.float32) | ||
| and w.dtype == x.dtype | ||
| and w.ndim == 1 | ||
| and x.shape[-1] == w.shape[0] | ||
| ) |
There was a problem hiding this comment.
The custom Metal kernel assumes that the input array x is row-contiguous and the weight array w is contiguous. If x or w is non-contiguous (e.g., due to slicing or transposing), the kernel will read incorrect memory, leading to silent correctness bugs or crashes. We should check x.flags.row_contiguous and w.flags.row_contiguous in _supported to safely fall back to mx.fast.rms_norm for non-contiguous inputs.
| def _supported(self, x: mx.array) -> bool: | |
| """Whether the custom kernel can handle this call; else we fall back.""" | |
| w = self._inner.weight | |
| return ( | |
| x.dtype in (mx.float16, mx.bfloat16, mx.float32) | |
| and w.dtype == x.dtype | |
| and w.ndim == 1 | |
| and x.shape[-1] == w.shape[0] | |
| ) | |
| def _supported(self, x: mx.array) -> bool: | |
| """Whether the custom kernel can handle this call; else we fall back.""" | |
| w = self._inner.weight | |
| return ( | |
| x.dtype in (mx.float16, mx.bfloat16, mx.float32) | |
| and w.dtype == x.dtype | |
| and w.ndim == 1 | |
| and x.shape[-1] == w.shape[0] | |
| and x.flags.row_contiguous | |
| and w.flags.row_contiguous | |
| ) |
| if x.dtype != w.dtype: | ||
| raise ValueError( | ||
| f"rms_norm x/w dtypes must match, got {x.dtype} vs {w.dtype}" | ||
| ) |
There was a problem hiding this comment.
If the Python API sgl_kernel.metal.rms_norm is called directly with non-contiguous arrays, it will bypass the wrapper's safety checks and call the C++ kernel directly, leading to undefined behavior or crashes. We should add checks for x.flags.row_contiguous and w.flags.row_contiguous here as well.
| if x.dtype != w.dtype: | |
| raise ValueError( | |
| f"rms_norm x/w dtypes must match, got {x.dtype} vs {w.dtype}" | |
| ) | |
| if x.dtype != w.dtype: | |
| raise ValueError( | |
| f"rms_norm x/w dtypes must match, got {x.dtype} vs {w.dtype}" | |
| ) | |
| if not x.flags.row_contiguous: | |
| raise ValueError("rms_norm expects x to be row-contiguous") | |
| if not w.flags.row_contiguous: | |
| raise ValueError("rms_norm expects w to be row-contiguous") |
07c593e to
1d937a7
Compare
jlee5814
left a comment
jlee5814
left a comment
There was a problem hiding this comment.
Please run pre-commit run --all-files for the lint failure, and add a test_norm_patching.py (mirroring test_attention_patching.py) covering patch idempotency, the Gemma rejection, and the unsupported-shape fallback. The kernel's tested but the wrapper isn't.
934f59f to
f4fc8ac
Compare
f4fc8ac to
33bedaf
Compare
AOT Metal rms_norm kernel (f16/bf16/f32) mirroring the merged rope_pool_fused path: kernel + Primitive in sgl-kernel/csrc/metal/ (rms_norm.metal, rms_norm.cpp, shared metal_common.h), registered in setup_metal.py, Python wrapper in sgl_kernel/metal.py. Vectorized (vec<T,4>) loads; at parity-or-better vs mx.fast.rms_norm across the hidden sizes Apple hardware runs, for f16 and bf16. SRT integration behind SGLANG_MLX_USE_CUSTOM_RMSNORM: a norm wrapper and patch_model_norms pass (srt/hardware_backend/mlx/norm_wrapper.py) with mx.fast.rms_norm fallback, hooked in model_runner after patch_model_attention. Tests in sgl-kernel/tests/test_metal_norm.py, benchmark in sgl-kernel/benchmark/bench_metal_rmsnorm.py.
Cache eps, weight metadata, and the contiguous weight at wrap time instead of recomputing per call; refresh the cache only when the module's weight array is rebound (identity check, so update_weights still takes effect on the next call); call the nanobind kernel entry directly, skipping duplicate validation in sgl_kernel.metal. 3-D serving-path dispatch overhead drops from ~2.6us to ~1.7us per call; e2e Qwen3-0.6B decode b1 regression halves (+2.9% -> +1.4%). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
clang-format rewrites rms_norm.cpp (include order, indentation); black collapses two statements in metal.py and test_metal_norm.py; trailing newline in rms_norm.metal; whitespace cleanups. Fixes the lint job on sgl-project#30163.
33bedaf to
98f87c4
Compare
- metal.py: unquote mx.array annotations (file already uses from __future__ import annotations + TYPE_CHECKING import). - bench_metal_rmsnorm.py: remove dead bench()/bench_chain() and wire --warmup/--repeats through to bench_pool (--iters was parsed but ignored). - Update run commands in docstrings for the python/sglang/kernels/aot/ tree location.
2 participants
What this adds
This adds a custom Metal RMSNorm kernel for the Apple Silicon MLX backend. It uses the same compiled ahead of time build path as the merged RoPE kernel (#22868 on the #23449 infra). The kernel is off by default. You turn it on with an environment variable, and it falls back to
mx.fast.rms_normwhen it is off or when a call is not supported.Code
Kernel and build:
sgl-kernel/csrc/metal/rms_norm.metal: the kernel. It computesy = x * rsqrt(mean(x^2, axis=-1) + eps) * w. It accumulates in fp32, upcasts the weight to fp32, and casts back to the input dtype on store. It loads four elements at a time when the hidden size is a multiple of 4, and uses a scalar path otherwise.sgl-kernel/csrc/metal/rms_norm.cpp: the MLX Primitive, the dispatch, and the Python entry point.sgl-kernel/csrc/metal/metal_common.h: a small shared header for the dtype suffix helper.sgl-kernel/setup_metal.py: registers the two new source files.sgl-kernel/python/sgl_kernel/metal.py: the Python wrapper.Backend integration:
python/sglang/srt/hardware_backend/mlx/norm_wrapper.py: a wrapper that replaces each plainnn.RMSNormmodule and calls the kernel, pluspatch_model_normswhich installs it. It is gated bySGLANG_MLX_USE_CUSTOM_RMSNORM.model_runner.pyright after the attention wrapper.Correctness
sgl-kernel/tests/test_metal_norm.py: 99 cases pass. 96 of them compare the kernel to an fp32 reference across the same sweep astest_norm.py: batch sizes {1, 19, 99, 989}, hidden sizes {111, 500, 1024, 3072, 3584, 4096, 8192, 16384}, dtypes {f16, bf16, f32}. The other 3 check input validation. Tolerances matchtest_norm.py: 1e-3 for f16, 2e-2 for bf16, 1e-4 for f32. Worst observed max difference: 3.9e-3 for f16, 3.1e-2 for bf16, 1.9e-6 for f32. The hidden size list covers both the vectorized path (multiples of 4) and the scalar path (111, 500). The test is skipped when the Metal extension is not built, so it is a no-op off Apple Silicon.output_idsare identical with the flag on and off. In a standalone model-level check that swaps only the norm implementation, the argmax token matches at every prompt position. The last-token logits differ by up to 0.34 in absolute value, which is accumulated rounding across 113 low precision norms and does not change any greedy token.Performance
Measured with
sgl-kernel/benchmark/bench_metal_rmsnorm.pyagainstmx.fast.rms_normon an M4 Pro (48 GB), mlx 0.31.2, macOS 26. The harness uses a pool of distinct inputs, takes the minimum over 7 repeats, and clears the MLX cache between configs. Numbers are microseconds per call, amortized inside one lazy evaluation. Speedup above 1.0 means the custom kernel is faster.Summary:
mx.fast.rms_normis notably strong for bf16 here.The fuller 84-config sweep is below. Given this profile, I would keep the kernel opt-in rather than on by default for now. The standalone op wins in a region, and the case for default-on is potentially the fused add + norm follow up, which removes a memory round trip. An alternative is to route by shape in the wrapper (e.g., small row counts go to
mx.fast.rms_norm), but that threshold would be tuned on one chip only.Open to hear feedback.
Full sweep (84 configs, f16 and bf16, batch 1 to 2048, hidden 896 to 8192)
The batch 1, hidden 896, f16 cell reads 0.35x when it is the first config a fresh process runs, which looks like GPU clock ramp up. The value above is from a warm re-run. All other cells are stable across repeated runs.
How to test
python sgl-kernel/setup_metal.py installon Apple Silicon with Python 3.11.python -m pytest sgl-kernel/tests/test_metal_norm.py.SGLANG_MLX_USE_CUSTOM_RMSNORM=1alongsideSGLANG_USE_MLX=1.CI States
Latest PR Test (Base): ❌ Run #30717815848
Latest PR Test (Extra): ❌ Run #30717815724