MLX backend: missing aten.native_group_norm and upsample_nearest2d fragment a diffusion UNet into 28 subgraphs

[msluszniak]

🐛 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