🐛 Describe the bug
The MLX backend has no handler for aten.native_group_norm or for
aten.upsample_nearest2d. Neither appears in backends/mlx/ops.py on main,
while layer_norm, native_layer_norm, rms_norm and
_native_batch_norm_legit_no_training are all registered.
GroupNorm sits in every ResBlock of a Stable-Diffusion-style UNet, and nearest
upsampling sits in every decoder stage, so these two gaps are enough to shatter
a diffusion model rather than merely slow it down.
Exporting SDXS-512-DreamShaper (an SD-1.5-style distilled one-step pipeline:
CLIP text encoder, UNet on 4x64x64 latents, TAESD decoder) with
MLXPartitioner gives:
| method | subgraphs | cause |
|---|
denoise (UNet) | 28 | 25x native_group_norm, 2x upsample_nearest2d |
decode (TAESD) | 4 | native_group_norm |
encode (CLIP) | 1 | fully delegated, LayerNorm only |
encode is the control: it is pure transformer, uses only LayerNorm, and
partitions into a single delegate as expected.
Why this is op coverage and not quantization
The artifact we measured is 4-bit weight-only on linears. To rule out the
dequant pattern as the cause, we also exported an unquantized fp32 MLX build of
the same pipeline. It fragments identically, same subgraph counts. No dtype or
quantization change affects the split.
Impact
Measured on an iPhone 17 Pro, the MLX build ran about 3x slower than the Core ML
fp16 build of the same pipeline (it was about 4x faster than XNNPACK fp32). The
per-subgraph boundary crossings dominate: 28 delegate handoffs per denoise
call, each one leaving and re-entering the MLX runtime.
The size story does not rescue it either. At 4-bit weight-only the MLX artifact
is 1034.8 MB against 880.7 MB for Core ML fp16, so it is both larger and slower
than the alternative on the same device.
We have removed the MLX artifact from our published model repo as a result. The
exporter is kept so the decision can be re-tested against a future release.
The primitive already exists in MLX
mlx.nn.layers.normalization already ships a GroupNorm with an explicit
_pytorch_compatible_group_norm path, so this looks like a missing binding in
the ExecuTorch backend rather than missing functionality in MLX itself.
Versions
- ExecuTorch 1.4.1, verified against
backends/mlx/ops.py on main at time of
filing (5244 lines, neither op present) - Host export: macOS, Apple Silicon
- Runtime: iOS, iPhone 17 Pro
Suggested fix
Register handlers for aten.native_group_norm and aten.upsample_nearest2d.vec
in backends/mlx/ops.py. native_group_norm returns (output, mean, rstd);
the existing _native_layer_norm_handler already demonstrates the pattern of
computing only the normalized output and asserting mean/rstd go unused, which
covers the inference case.
cc @metascroy
🐛 Describe the bug
The MLX backend has no handler for
aten.native_group_normor foraten.upsample_nearest2d. Neither appears inbackends/mlx/ops.pyonmain,while
layer_norm,native_layer_norm,rms_normand_native_batch_norm_legit_no_trainingare all registered.GroupNorm sits in every ResBlock of a Stable-Diffusion-style UNet, and nearest
upsampling sits in every decoder stage, so these two gaps are enough to shatter
a diffusion model rather than merely slow it down.
Exporting SDXS-512-DreamShaper (an SD-1.5-style distilled one-step pipeline:
CLIP text encoder, UNet on 4x64x64 latents, TAESD decoder) with
MLXPartitionergives:denoise(UNet)native_group_norm, 2xupsample_nearest2ddecode(TAESD)native_group_normencode(CLIP)encodeis the control: it is pure transformer, uses only LayerNorm, andpartitions into a single delegate as expected.
Why this is op coverage and not quantization
The artifact we measured is 4-bit weight-only on linears. To rule out the
dequant pattern as the cause, we also exported an unquantized fp32 MLX build of
the same pipeline. It fragments identically, same subgraph counts. No dtype or
quantization change affects the split.
Impact
Measured on an iPhone 17 Pro, the MLX build ran about 3x slower than the Core ML
fp16 build of the same pipeline (it was about 4x faster than XNNPACK fp32). The
per-subgraph boundary crossings dominate: 28 delegate handoffs per
denoisecall, each one leaving and re-entering the MLX runtime.
The size story does not rescue it either. At 4-bit weight-only the MLX artifact
is 1034.8 MB against 880.7 MB for Core ML fp16, so it is both larger and slower
than the alternative on the same device.
We have removed the MLX artifact from our published model repo as a result. The
exporter is kept so the decision can be re-tested against a future release.
The primitive already exists in MLX
mlx.nn.layers.normalizationalready ships a GroupNorm with an explicit_pytorch_compatible_group_normpath, so this looks like a missing binding inthe ExecuTorch backend rather than missing functionality in MLX itself.
Versions
backends/mlx/ops.pyonmainat time offiling (5244 lines, neither op present)
Suggested fix
Register handlers for
aten.native_group_normandaten.upsample_nearest2d.vecin
backends/mlx/ops.py.native_group_normreturns(output, mean, rstd);the existing
_native_layer_norm_handleralready demonstrates the pattern ofcomputing only the normalized output and asserting mean/rstd go unused, which
covers the inference case.
cc @metascroy