Skip to content

[CUDA] Add global scale support to gather_qmm - #4507

Merged
nastya236 merged 2 commits into
ml-explore:mainfrom
dhiltgen:cuda_g_qmm_global_scle
Sep 17, 2026
Merged

nastya236 merged 2 commits into
ml-explore:mainfrom
dhiltgen:cuda_g_qmm_global_scle

Conversation

@dhiltgen

@dhiltgen dhiltgen commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

This adds support for global scales in gather_qmm on cuda.

mlx-lm perf data with qwen 3.6 35B tested on a 5090. (p2048/g128)

model / system prompt tps gen tps
mlx-community/Qwen3.6-35B-A3B-nvfp4 (no global scale) 394.3 75.9
Model Optimizer export (with global scale) 404.8 70.5
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: co-developed with a coding agent

Attached wrapper script loads the Model Optimizer exports and wires up gather_qmm global scale, and adjusts dtypes for a fair comparison with the standard community model.
modelopt_nvfp4_mlxlm.py

@nastya236

nastya236 commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator

Thanks! Why "optimizes performance"? Did you notice a speed up?

General overall note: after giving this some more thought, I realised that dequantization during inference would be correct if and only if it matches the training. I believe there are two possible ways to produce nvfp4 weights:

  1. quantize(W) -> use nvfp4 x nvfp4 instructions for matmul and fold global scale into epilogue — in this case we also must fold scale into epilogue during inference as well (like in this PR)
  2. If the global scale was applied when dequantizing the weights, before casting to BF16 and performing the matrix multiplication in bf16 (quantization aware training on sm < 100), we should follow that same sequence during inference — this is done in [Metal] global scale for qmm #4458

Given that in the future all quantized operations during training will be native, I think it is okay to agree on folding it into epilogue. However, probably we should do the same for metal..

@nastya236
nastya236 self-requested a review September 16, 2026 09:56

@nastya236 nastya236 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.

Some tests are failing can you have a look?

@nastya236 nastya236 added the await response This pull request is waiting for response from the author. label Sep 16, 2026
@dhiltgen
dhiltgen force-pushed the cuda_g_qmm_global_scle branch from fcd135f to d3012f2 Compare September 16, 2026 16:49
@dhiltgen

Copy link
Copy Markdown
Contributor Author

Why "optimizes performance"? Did you notice a speed up?

On further testing, I couldn't reproduce any speedup in the second commit so I've dropped it and the claim.

@nastya236
nastya236 merged commit 59d600b into ml-explore:main Sep 17, 2026
29 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

await response This pull request is waiting for response from the author.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants