Skip to content

[CUDA] Add 64-column fast path for gather_qmm - #4853

Open
dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:cuda-perf-n64
Open

dhiltgen wants to merge 1 commit into
ml-explore:mainfrom
dhiltgen:cuda-perf-n64

Conversation

@dhiltgen

@dhiltgen dhiltgen commented Nov 8, 2026

Copy link
Copy Markdown
Contributor

This PR adds a new smaller 64-column tile fast path for sorted nvfp4 gather_qmm to improve performance for models with expert widths that are divisible by 64 but not by 128.

Tested with mlx-lm p2048/g128 using the mlx-community/gemma-4-26b-a4b-it-nvfp4
and mlx-community/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4

GPU Model Prompt TPS, base -> candidate speedup
RTX 6000 Ada Gemma26 1713 -> 6567 3.8x
DGX Spark Gemma26 906 -> 2075 2.3x
DGX Spark Nemotron3 Nano 30B-A3B NVFP4 966 -> 1727 1.8x
RTX 5090 Gemma26 3613 -> 10567 2.9x
  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: co-developed with coding agent

This branch has not been deployed

No deployments
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.

1 participant