perf(nvfp4): direct-transcode ModelOpt triplets to MLX native #693

Context

PR #692 repairs the CUDA ModelOpt NVFP4 path for gemma-4-31b-it-nvfp4 by accepting the narrow NVIDIA ModelOpt triplet format and repacking it at load time. The final CUDA path repacks weight/weight_scale/weight_scale_2 into MLX native NVFP4 tensors under --features cuda; non-CUDA builds intentionally keep the affine fallback until Apple Silicon is measured separately.

The current repair path is intentionally conservative, but cold-load expensive. It reconstructs dense f16 values first, then calls MLX quantize(..., mode="nvfp4"): ModelOpt NVFP4 -> dense f16 -> MLX native NVFP4. This is safe because MLX owns scale rounding and packing, but it pays a large one-time load cost and temporarily materializes dense tensors.

Goal

Implement a direct ModelOpt triplet -> MLX native NVFP4 transcode path that avoids materializing a dense f16 weight matrix for each linear weight group while preserving the source two-level scale semantics.

Technical notes

MLX native NVFP4 stores weights as uint32 with eight FP4 E2M1 nibbles per word and stores per-block scales as uint8 FP8 E4M3 with group_size=16. The ModelOpt FP4 nibble encoding used by the local Gemma 4 checkpoint matches the MLX E2M1 LUT order (0, 0.5, 1, 1.5, 2, 3, 4, 6 with sign bit), so the weight payload can likely be repacked directly by combining four source U8 bytes into one little-endian U32 per eight input columns.

The scale side should prefer preserving weight_scale_2 as an explicit global-scale sidecar rather than folding it into every E4M3 block scale. Folding weight_scale * weight_scale_2 into each block scale is smaller, but it can introduce avoidable overflow or rounding drift when the folded value exceeds what one E4M3 byte can represent. A sidecar keeps the source ModelOpt two-level scale model intact and lines up with MLX/CUDA NVFP4 paths that already understand optional global scales.

Implementation plan

1. Extend the weight representation and UnifiedLinear loading path so a quantized weight can carry an optional per-linear global scale sidecar for native NVFP4.
2. Thread that sidecar through dequantize, quantized_matmul, and any qmv/qmm bridge surfaces that currently pass mode="nvfp4" without a global scale.
3. Add a direct transcode helper for ModelOpt NVFP4 triplets: pack four source U8 bytes into one little-endian U32, preserve per-block weight_scale as E4M3 bytes, and preserve weight_scale_2 as the global-scale sidecar.
4. Keep the dense f16 -> MLX quantize(..., mode="nvfp4") path behind a debug/test fallback until direct transcode parity is validated.
5. Measure cold-load wall time, MLX peak memory, short-prompt throughput, and 2048-prompt throughput on GB10.

Acceptance criteria

• Direct transcode avoids materializing a full dense f16 matrix for each ModelOpt triplet.
• weight_scale_2 is preserved as a global-scale sidecar and used by native NVFP4 dequantize/matmul rather than folded into each E4M3 block scale.
• A focused fixture compares dequantized direct-transcode values against the dense f16 -> MLX native NVFP4 path and documents the expected FP8 scale rounding tolerance.
• Greedy continuation parity is spot-checked against the dense f16 -> MLX native NVFP4 path for at least 40 generated tokens, allowing only expected FP8 rounding drift.
• Cold-load wall time or peak load-time memory improves by at least 20% versus the dense f16 -> MLX quantize path on gemma-4-31b-it-nvfp4.
• Results are recorded under benchmarks/ and summarized under docs/benchmark_results/.

Relationship to Apple Silicon

This issue should land before enabling native NVFP4 repack broadly on Apple Silicon. Once direct transcode and sidecar plumbing are validated, a separate Apple Silicon issue should benchmark whether Metal/non-CUDA should switch from the affine fallback to the native NVFP4 path.