feat: Expose MLX's hadamard_transform scale on the mlxcel-core bridge

[inureyes]

Problem / Background

MLX's hadamard_transform takes an optional scale. mlxcel's bridge does not carry it, so every caller silently receives the orthonormal 1/sqrt(N) default. That is the correct default for everything mlxcel rotates today, and the wrong constraint for rotation combined with a block-scaled low-precision format, where the scale has to be chosen so it survives the target format's scale encoding.

Nothing is broken today, which is why this is priority:medium and not higher. Every current caller wants the orthonormal default and keeps it. What would make this urgent: any work that rotates a tensor before quantizing it to mxfp4, mxfp8, nvfp4 or fp8. Until then this is a missing argument, not a defect.

Current Behavior

• mlx/ops.h:166 declares hadamard_transform(const array& a, std::optional<float> scale = std::nullopt, StreamOrDevice s = {}), and mlx/ops.cpp:566-576 resolves float scale = scale_.has_value() ? *scale_ : 1.0f / std::sqrt(n) under the comment "Default to an orthonormal Hadamard matrix scaled by 1/sqrt(N)". The vendored copy (MLX 0.32.3) is under Cargo's per-profile OUT_DIR, target/<profile>/build/mlxcel-core-<hash>/out/build/_deps/mlx-src, not a single top-level out/.
• The bridge declares and defines only the one-argument form: src/lib/mlxcel-core/cpp/mlx_cxx_bridge.h:1593, src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp:6634-6636, src/lib/mlxcel-core/src/lib.rs:2453. src/lib/mlxcel-core/src/ops.rs:129-137 wraps it as wht() with a power-of-two assertion on the last axis, re-exported at src/lib/mlxcel-core/src/lib.rs:2999-3000.
• Callers are the TurboQuant KV cache: src/lib/mlxcel-core/src/cache/turbo/quant.rs:337,556,741,763,787, quant3.rs:214,370, sparse_v.rs:843,1037,1447, plus src/lib/mlxcel-core/examples/wht_microbench.rs and the tests. Their rotated output feeds a 4-bit or 3-bit Lloyd-Max codebook with nibble packing (quant.rs) or 24-bit-group packing (quant3.rs), not an MLX block-scaled format, so orthonormal is right there and nothing regresses.
• The convention is already pinned by tests: src/lib/mlxcel-core/src/ffi_tests.rs:3857-3876 (test_wht_matches_h4_reference) against a hand-computed H_4, and tests/wht_op.rs:121-144 (wht_preserves_l2_norm) by energy. ffi_tests.rs:3844-3851 records that MLX normalizes the radix-mixed sizes (m * 2^k, m in 12/20/28) differently, so those do not round-trip and stay out of scope.

Proposed Solution

Add a second entry point rather than making the existing one optional. cxx cannot carry std::optional<float> or Option<f32> (the bridge in lib.rs contains zero Option<...> types), and this crate's established idiom for an optional argument is a parallel function, as in quantized_linear_forward versus quantized_linear_forward_global_scale (src/lib/mlxcel-core/src/lib.rs:590,603).

• C++: std::unique_ptr<MlxArray> hadamard_transform_scaled(const MlxArray& a, float scale) declared beside mlx_cxx_bridge.h:1593 and defined beside mlx_cxx_bridge.cpp:6634, forwarding std::make_optional(scale) as the second argument.
• Rust extern: fn hadamard_transform_scaled(a: &MlxArray, scale: f32) -> UniquePtr<MlxArray>; beside lib.rs:2453.
• Rust wrapper: pub fn wht_scaled(x: &ffi::MlxArray, scale: f32) -> UniquePtr<ffi::MlxArray> in ops.rs, exported alongside wht at lib.rs:2999-3000.
• wht() keeps its signature and keeps calling the one-argument bridge fn, so no current caller changes behaviour.

Rejected: changing wht(x) to wht(x, Option<f32>). It churns ten call sites for no behaviour change, and cxx would still need a second extern underneath it.

Scope

In scope: the three bridge layers (mlx_cxx_bridge.h, mlx_cxx_bridge.cpp, lib.rs), the ops.rs wrapper and its re-export, the tests below, and the written per-format scale policy.

Out of scope: choosing or applying a non-default scale in any TurboQuant path (they stay orthonormal); radix-mixed head dimensions; building an mxfp/nvfp4/fp8 rotation pipeline. This issue supplies the argument such a pipeline will need, nothing more.

Implementation Notes

• Reuse: ops.rs:131-136 asserts the last axis is a non-zero power of two because the cxx extern is noexcept, so MLX's std::invalid_argument would reach std::terminate instead of unwinding. Factor that guard into a private helper and have both wht and wht_scaled call it; do not copy it.
• Bit-identical default: MLX evaluates 1.0f / std::sqrt(n) with n an int, so std::sqrt resolves to the double overload and the quotient is narrowed to float once. A caller reproducing the default must compute (1.0f64 / (n as f64).sqrt()) as f32; the naive 1.0f32 / (n as f32).sqrt() can differ by one ulp when log2(N) is odd (head_dim 128 included) and would make an exact-equality test flaky.
• Edge cases: scale = 0.0 is legal in MLX and returns zeros, so do not special-case it. n == 1 short-circuits at mlx/ops.cpp:580-586 to a plain multiply and passes the power-of-two guard, which is correct. MLX does not validate a non-finite scale, so wht_scaled must assert scale.is_finite() next to the shape guard.
• Error handling: identical contract to wht(), a Rust panic carrying the shape or the offending scale, never an exception crossing the noexcept extern.

Acceptance Criteria

• wht() behaviour is unchanged: test_wht_matches_h4_reference, test_wht_round_trip_power_of_two, test_wht_fp16_preserves_dtype and wht_preserves_l2_norm all still pass with their assertions untouched.
• For head_dim in {64, 128, 256}, wht_scaled(x, (1.0f64 / (n as f64).sqrt()) as f32) matches wht(x) elementwise within allclose(1e-6, 1e-6).
• For head_dim in {64, 128, 256}, ||wht(x)||_2 / ||x||_2 reads 1.0 and ||wht_scaled(x, 1.0)||_2 / ||x||_2 reads sqrt(N), both within 1e-4 relative.
• wht_scaled panics with a clear message (it does not abort the process) on a non-power-of-two last axis and on a NaN or infinite scale.
• The per-format scale policy is written where the next implementer will hit it: rustdoc on wht_scaled in ops.rs, plus a paragraph under docs/turbo-kv-cache.md:197 ("WHT head-dimension constraint"). It states orthonormal 1/sqrt(N) for the fp16 KV path as today; a power-of-two scale for mxfp4 and mxfp8 targets so it folds into the E8M0 block exponent exactly, noting that 1/sqrt(N) is a power of two only when log2(N) is even (64 gives 2^-3 and 256 gives 2^-4, but 128 is 2^-3.5 and is not E8M0-representable, costing an extra multiply and an extra rounding of the data); absorption into nvfp4's FP32 per-tensor global scale at a cost in range; and an explicit saturation check against E4M3's 448 for fp8, since scale 1 grows magnitudes by up to sqrt(N), which is 11.3x at head_dim 128.
• Integrated rather than stranded: wht_scaled is re-exported from mlxcel_core and exercised through the public path by a case in tests/wht_op.rs, mirroring wht_is_publicly_exported_and_runs.

Verification

cargo test --workspace --profile test-fast --features metal,accelerate -- wht
cargo test --workspace --profile test-fast --features metal,accelerate
cargo clippy --workspace --all-targets -- -D warnings
cargo fmt --check

A pass is every WHT test green (including the four pre-existing ones, unmodified), a clean clippy and fmt, and the C++ bridge rebuilding to pick up the new extern.

Technical Considerations

Cross-backend: MLX's Hadamard primitive stores the resolved scale as a plain float member (mlx/primitives.h:1229-1232) and each backend applies it (mlx/backend/{cpu,metal,cuda}/hadamard.*). ROCm still stubs the primitive out (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/primitives.cpp:33, NO_GPU(Hadamard)), and #1825 is implementing it now; that port must read scale_ rather than hardcoding 1/sqrt(N), or a scaled call will silently produce orthonormal output on AMD only.

Format encoding, for the policy note: MLX selects the block-scale encoding by group size at mlx/ops.cpp:5108-5133, group 16 to E4M3 with the optional global scale (nvfp4) and everything else to an E8M0 exponent (the mx family). mlxcel has already been bitten by E8M0 rounding: #1769 failed on every Metal host, and the pin bump in #1772 fixed it by rounding the exponent up, after about half the blocks had their maxima clipped by up to 29% (src/models/fp8_block.rs:37-43). CUDA carries native nvfp4 and mxfp8 paths and Metal gained the matching mxfp8 behaviour with that pin, so the scale choice is per backend as well as per format.