Skip to content

Use matrix kernels for global-scale gather_qqmm - #4481

Open
dhiltgen wants to merge 2 commits into
ml-explore:mainfrom
dhiltgen:g-qqmm-dispatch
Open

Use matrix kernels for global-scale gather_qqmm#4481
dhiltgen wants to merge 2 commits into
ml-explore:mainfrom
dhiltgen:g-qqmm-dispatch

Conversation

@dhiltgen

@dhiltgen dhiltgen commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

This wires up the new gather_qmm kernels to optimize gather_qqmm with global scale.

Performance

Using the NVIDIA Model Optimizer on Qwen/Qwen3.6-35B-A3B with nvfp4_mlp_only p2048/g128

main prompt tps this PR prompt tps main generation tps this PR generation tps
M5 Max 1,257.0 3,105.3 81.9 81.7
M3 Ultra 1,627.3 2,673.2 68.5 68.6
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: co-developed with coding agent

@zcbenz

zcbenz commented Sep 11, 2026

Copy link
Copy Markdown
Member

The tests are failing for CUDA backend, not sure of the reason, we can probably disable the test for CUDA for now and create an issue for it so we can look into it later.

@zcbenz

zcbenz commented Sep 11, 2026

Copy link
Copy Markdown
Member

It seems that the CI failure is just a numeric precision issue, I think we can just allow for a bigger error in CUDA.

@nastya236

nastya236 commented Sep 11, 2026

Copy link
Copy Markdown
Collaborator

I think it is happening, because we don't have proper global scaling for cuda.. In metal gather_qqmm always uses gather_qmm (because for metal we don't have quantized x quantized instructions). Because we plan to switch to grouped_mm instead of gather_mm (and grouped_qqmm instead of gather_qqmm) specifically on cuda, I think for now we can just skip this test on cuda and fix this with proper scaling in the future grouped_qqmm. What do you think?

@dhiltgen
dhiltgen force-pushed the g-qqmm-dispatch branch 2 times, most recently from 50b106a to 5fc7dd0 Compare September 11, 2026 19:30
@zcbenz

zcbenz commented Sep 12, 2026

Copy link
Copy Markdown
Member

Sounds good to me!

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.

3 participants