perf(metal): close native NVFP4 prefill gap

[inureyes]

Context

Issue #694 and PR #702 measured the Metal/non-CUDA ModelOpt NVFP4 native path against the affine fallback on M1 Ultra using the local models/gemma-4-31b-it-nvfp4 ModelOpt NVFP4 checkpoint. Native is clearly better for load time, decode throughput, and peak memory, so it is the right default direction once the prefill gap is either closed or justified by the net load/decode/memory trade-off.

A longer sweep on main commit 80b4729ad used mlxcel-bench-decode with prompt lengths 512, 1024, 2048, 4096, 8192, and 16384. The trend was stable: native decode was about 2.6x faster and peak memory stayed 3.3-12.9 GB lower, but native prefill remained 3-9% slower.

Prompt tokens	Affine prefill tok/s	Native prefill tok/s	Native prefill gap	Affine decode tok/s	Native decode tok/s	Native decode speedup	Affine peak	Native peak	Affine load	Native load
512	121.77	118.15	-3.0%	4.92	13.14	2.67x	40.44 GB	37.17 GB	229.8s	82.6s
1024	128.72	118.66	-7.8%	5.02	13.18	2.63x	41.61 GB	37.17 GB	221.4s	82.0s
2048	126.47	115.29	-8.8%	4.95	13.20	2.67x	44.83 GB	37.17 GB	226.2s	84.1s
4096	121.69	112.20	-7.8%	4.87	12.77	2.62x	47.76 GB	37.43 GB	224.5s	82.8s
8192	118.62	109.79	-7.4%	4.73	12.20	2.58x	49.11 GB	37.91 GB	222.6s	82.7s
16384	114.19	107.45	-5.9%	4.44	11.54	2.60x	52.60 GB	39.70 GB	222.0s	83.4s

PR #704 has now updated the MLX pin to 0.32.1, so this work uses that merged baseline when re-testing the prefill path. The earlier wording that named Llama 4 Scout has been corrected here: the reproducible local checkpoint for #694/#702/#705 is models/gemma-4-31b-it-nvfp4.

Goal

Reduce the native ModelOpt NVFP4 prefill gap enough that mlxcel can prefer the native representation by default on Metal/non-CUDA, while preserving the load-time, decode-throughput, and memory advantages that motivated the native path.

Hypothesis

The likely gap is not the raw native representation itself. Decode already benefits from the native sidecar/global-scale path, but multi-token prefill appears to bypass the same fused/compiled sidecar helpers or falls back to less favorable C++ eager folding in Gemma/Llama4 MLP paths. The implementation should focus on the native NVFP4 prefill path rather than reverting to affine.

Activation-path follow-up

The remaining native prefill gap also prompted a focused re-check of the issue #701 path choice on the current Metal/MLX 0.32.1 baseline. At 2048 prompt tokens, the native default shape-specific scaled MLP measured 116.10 tok/s, the op-at-a-time sidecar path with compiled GeGLU activation (MLXCEL_DISABLE_FUSED_GLOBAL_SCALE=1) measured 116.11 tok/s, and the C++ eager fold (MLXCEL_COMPILED_QGELU_MLP=0) measured 113.35 tok/s. At 8192 prompt tokens, the same paths measured 110.66, 109.99, and 108.10 tok/s respectively. This does not support switching default native prefill back to the compiled-activation fallback; it keeps the #701 eager-fold warning intact.

Acceptance criteria

• Identify the current multi-token native NVFP4 prefill dispatch path and document the bottleneck in code comments, the PR body, or a small benchmark note.
• Add a prefill-friendly native NVFP4 fused/compiled path, or otherwise remove the measured dispatch overhead, without changing ModelOpt NVFP4 dequantization semantics.
• Keep affine fallback and explicit environment overrides available for rollback and comparison.
• Re-run the prompt-length sweep at 512, 1024, 2048, 4096, 8192, and 16384 tokens on the same local models/gemma-4-31b-it-nvfp4 ModelOpt NVFP4 checkpoint or document any reduced validation set used for iteration.
• Show native prefill within 3% of affine on the representative long-prefill cases, or provide measured evidence that the gap is materially reduced while native decode remains at least 2x faster.
• Show native peak memory remains lower than affine and native load time remains substantially faster.
• If the measurements pass, make native ModelOpt NVFP4 the default on Metal/non-CUDA while preserving the opt-out override for affine.

Out of scope

• CUDA behavior, which already uses the native NVFP4 path directly.
• Broad model-family rewrites unrelated to ModelOpt NVFP4 prefill.