Repository navigation
Description
Problem / Background
#2441 left three load paths that mishandle an NVFP4 tensor scale. Implement all three in one PR, run the verification block once at the end, and keep pushes to a minimum.
1. Dense ModelOpt NVFP4 in fused QKV
FusedQKVLinear::from_weights_separate (src/lib/mlxcel-core/src/layers.rs:3286-3297) hardcodes "affine", and the fused arm sets global_scale: None (layers.rs:3510-3519), dropping each plane's repacked weight_scale_2 ({p}.global_scale, src/models/sanitize.rs:805). nvidia/Qwen3-30B-A3B-NVFP4 stacks its experts, then aborts in this projection. group_size() also falls back to 64 for a ModelOpt-only config (src/models/qwen3_moe.rs:96, src/models/qwen3.rs:94).
Fix: infer the mode with infer_quantization_mode (layers.rs:2241) as UnifiedLinear::from_weights does. When any of q/k/v carries .global_scale, return QkvProjection::Split built through UnifiedLinear::from_weights_with_mode, which already applies the scalar after the matmul (layers.rs:2683, 2795-2828). That is the inv_global_scale convention, not MLX's native global_scale. Add the modelopt_nvfp4_group_size fallback as src/models/phimoe.rs:127-137 does. Add no field to QuantizedWeight, UnifiedLinear, QkvProjection or FusedQKVLinear (the debug Pipeline CI smoke test sits at the 2 MiB stack edge, #2434), and keep the #2438 claim_linear pinning on each plane. A fused per-row scale vector is rejected: the fused rope kernels take no scale.
-
layers.rstest: nvfp4 q/k/v planes, no biases, distinct.global_scalescalars, load asSplitwith modenvfp4, andforwardmatchesx @ dequantize(w).T * sper plane. TheSplit/mode assertion fails on current code. - Qwen3 and Qwen3-MoE args from a ModelOpt-only config return
group_size() == 16. -
~/models_archive/hf-src/Qwen3-30B-A3B-NVFP4loads on CUDA and generates coherent greedy text.
2. VLM loaders and the Gemma 4 unified expert split
apply_native_global_scale_flags (src/models/switch_layers.rs:945) runs only in load_text_weights (sanitize.rs:2030) and the Gemma 4 sanitizer (sanitize.rs:1181). Call it in finish_vlm_weights_common (src/loading/vlm.rs:407) and in load_gemma4_unified on the raw map before sanitize_gemma4_unified_weights (src/loading/vlm_gemma_unified.rs:250-253). Extend split_fused_moe_key (vlm_gemma_unified.rs:128) to .native_global_scale and .global_scale: copy the [num_experts] vector unsplit to both gate_proj and up_proj, and rename it under down_proj. An unflagged .global_scale then reaches a SwitchLinear prefix and fails there (switch_layers.rs:649).
-
vlm_gemma_unified_tests.rs: a synthetic fusedexperts.gate_up_projnvfp4 export with a flagged.global_scalebuilds gate and upSwitchLinearplanes that both carry the native scale; without the flag the load fails naming theswitch_glu.gate_projprefix. Current code orphans the key. - A
load_vlm_weights_commoncheckpoint shipping.native_global_scaleis refused.
3. DeepSeek v1 bias-free expert plane
SwitchLinear::from_stacked_parts (src/models/deepseek.rs:972-995) maps .scales without .biases to Regular, so a ModelOpt or mlx-lm nvfp4 expert plane aborts in gather_mm at the first forward. Return a load error naming the prefix when only one of the two is present.
-
deepseek_tests.rs: such a plane failsSwitchLinear::from_weightsnaming the prefix.
Out of scope: the O(planes x keys) sidecar scan cost.
Verification
On GB10, wrap GPU runs in gpu-lock run.
cargo fmt --check
cargo clippy --profile test-fast --features cuda --lib --tests -- -D warnings
cargo test -p mlxcel-core --profile test-fast --features cuda --lib layers:: -- --test-threads=1
cargo test --profile test-fast --features cuda --lib -- --test-threads=1 qwen3 vlm deepseek:: switch_layers sanitize
MLXCEL_SDPA_DETERMINISTIC=1 mlxcel generate -m ~/models_archive/hf-src/Qwen3-30B-A3B-NVFP4 -p "The capital of France is" -n 40 --temp 0
Activity
fix(nvfp4): guard fused block-float QKV and filtered VLM flag pass
fix(nvfp4): close global-scale load gaps in QKV, VLM and DeepSeek
fix(nvfp4): guard fused block-float QKV and filtered VLM flag pass
fix: close NVFP4 global-scale load gaps left after lablup#2441 (lablu…