Skip to content

[CUDA] Add fp_gather_qmv to optimize gather_qmm - #4508

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

dhiltgen wants to merge 3 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

uint32_t blocks_y = (N + rows_per_block - 1) / rows_per_block;
const uint32_t* mat_ptr = gpu_ptr<uint32_t>(mat);
const T* vec_ptr = gpu_ptr<T>(vec);
int n = 1;

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.

this block is just a refactoring if I am reading it correctly?

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.

Correct, no logic change, since my change adds a second path that needs the same computation I refactored to DRY it out.

Comment thread mlx/backend/cuda/quantized/quantized.cpp
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
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