Skip to content

fix: Close NVFP4 global-scale load gaps left after #2441 #2444

Closed

Description

@inureyes

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.rs test: nvfp4 q/k/v planes, no biases, distinct .global_scale scalars, load as Split with mode nvfp4, and forward matches x @ dequantize(w).T * s per plane. The Split/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-NVFP4 loads 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 fused experts.gate_up_proj nvfp4 export with a flagged .global_scale builds gate and up SwitchLinear planes that both carry the native scale; without the flag the load fails naming the switch_glu.gate_proj prefix. Current code orphans the key.
  • A load_vlm_weights_common checkpoint shipping .native_global_scale is 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 fails SwitchLinear::from_weights naming 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

  1. added
    type:bugBug fixes, error corrections, or issue resolutions
    area:modelsModel architectures, weights, loading, metadata
    and removed on Oct 18, 2026
  2. added 2 commits that reference this issue on Oct 19, 2026
    397205f
    485f0a4
  3. added a commit that references this issue on Oct 21, 2026
    e772cf2
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:modelsModel architectures, weights, loading, metadatapriority:mediumMedium prioritystatus:doneCompletedtype:bugBug fixes, error corrections, or issue resolutions

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions