Skip to content

[CUDA] Add fp_gather_qmv to optimize gather_qmm - #4508

Open
dhiltgen wants to merge 4 commits into
ml-explore:mainfrom
dhiltgen:cuda_fp_gather_qmv
Open

dhiltgen wants to merge 4 commits into
ml-explore:mainfrom
dhiltgen:cuda_fp_gather_qmv

Conversation

@dhiltgen

Copy link
Copy Markdown
Contributor

This speeds up gather_qmm on CUDA for MoE models.

mlx-lm benchmark examples with gemma4 26b and Nemotron 3 Nano 30B (p2048/g128)

GPU Model (mlx-community, NVFP4) main prompt tps branch tps
RTX 5090 NVIDIA-Nemotron-3-Nano-30B-A3B 1671.6 3177.8
RTX 5090 gemma-4-26b-a4b-it 1575.2 2553.0
RTX 6000 Ada NVIDIA-Nemotron-3-Nano-30B-A3B 906.7 1307.3

Generation speed unchanged.

Note: this carries a few lines from #4507 but is otherwise independent of the global scale work so these can merge in either order.

  • ☑️ I understand it is strictly prohibited to use AI to write PR description
  • AI usage disclosure: co-developed with coding agent

Comment thread mlx/backend/cuda/quantized/qmm/fp_qmv.cu
Comment thread mlx/backend/cuda/quantized/quantized.cpp Outdated
Comment thread mlx/backend/cuda/quantized/qmm/fp_qmv.cu Outdated
Comment thread mlx/backend/cuda/quantized/qmm/fp_qmv.cu Outdated
Comment thread mlx/backend/cuda/quantized/quantized.cpp Outdated
@nastya236

Copy link
Copy Markdown
Collaborator

Thank you for the pull request! Great improvements! I left some comments and questions.

@nastya236 nastya236 added the await response This pull request is waiting for response from the author. label Sep 15, 2026
bool can_use_qmv = supports(supports_qmv) || can_use_fp_gather_qmv;
// fp_gather_qmv wins at any batch size, gather_qmv only under 8 rows.
bool prefer_qmv = can_use_fp_gather_qmv || (can_use_qmv && M * B < 8);
auto call_qmm_sm80 = [&]() {

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.

If at any batch size, why M * B < 8 not M < 8?

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.

Sorry, that comment was misleading. I've removed it. The existing logic (M * B < 8) still holds based on my benchmarking tests - the old gather_qmv kernel loses to the tile kernel as soon as there are more than a few rows.

Comment thread mlx/backend/cuda/quantized/quantized.cpp Outdated

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

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