Skip to content

Add global scale support to quantized layers#426

Open
aleroot wants to merge 1 commit into
ml-explore:mainfrom
aleroot:quant_scale
Open

Add global scale support to quantized layers#426
aleroot wants to merge 1 commit into
ml-explore:mainfrom
aleroot:quant_scale

Conversation

@aleroot

@aleroot aleroot commented Jun 19, 2026

Copy link
Copy Markdown
Contributor

Store optional globalScale on QuantizedLinear and QuantizedEmbedding so nvfp4 weights can preserve the scale needed by lower-level MLX quantize/dequantize operations.

Forward the scale when creating or dequantizing weights, add a direct pre-quantized QuantizedEmbedding initializer to match QuantizedLinear, and guard global-scale execution paths on Metal because MLX does not support globalScale dequantization there.

Add focused tests for the new layer state and parameter exposure without broadening the generic quantization API surface.

Proposed changes

Please include a description of the problem or feature this PR is addressing. If there is a corresponding issue, include the issue #.

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

Comment thread Source/MLXNN/Quantized.swift Outdated
Comment thread Source/MLXNN/Quantized.swift Outdated
public let mode: QuantizationMode
public let scales: MLXArray
public let biases: MLXArray?
public let globalScale: MLXArray?

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One thing we need to be careful of here is that there is currently no python support for this. That means:

  • maybe we want to set the key as global_scale or global_scales to match the likely python naming
  • if this is non-nil after quantizing the model, it will be required when loading weights
    • will this cause a problem or is nvfp4 just not used?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point. I changed the serialised parameter key to global_scale

@davidkoski davidkoski left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like it, but see my questions -- I think these need some thought before merging.

Store optional globalScale on QuantizedLinear and QuantizedEmbedding so nvfp4 weights can preserve the scale needed by lower-level MLX quantize/dequantize operations.

Forward the scale when creating or dequantizing weights, add a direct pre-quantized QuantizedEmbedding initializer to match QuantizedLinear, and guard global-scale execution paths on Metal because MLX does not support globalScale dequantization there.

Add focused tests for the new layer state and parameter exposure without broadening the generic quantization API surface.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants