Repository navigation
[CUDA] GatherQMM matrix-matrix sm80/naive path - #3417
Conversation
The pre-gather approach (gather_slices + qmm_sm90) is a temporary workaround that benchmarks 2-4x slower than the fused sm80 kernel.
There was a problem hiding this comment.
Pull request overview
Adds gather-index support to the CUDA quantized matrix-matrix (QMM) implementations so GatherQMM can use the faster SM80 CUTLASS-based kernel (or the naive kernel) by selecting per-output batch slices via lhs_indices / rhs_indices.
Changes:
- Plumbs optional
lhs_indices/rhs_indicespointers throughqmm_sm80andqmm_naivedown to the CUDA kernels, and uses them to select the A/B (and scale/bias) batch slices. - Updates
GatherQMM::eval_gputo preferqmm_sm80/qmm_naivepaths (with a small-problem fallback togather_qmv). - Extends the QMM public headers/dispatch to accept the new optional gather-index parameters.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
mlx/backend/cuda/quantized/quantized.cpp |
Routes GatherQMM to SM80/naive QMM kernels and passes gather index buffers into the QMM calls. |
mlx/backend/cuda/quantized/qmm/qmm_impl_sm80.cuh |
Adds optional gather index parameters to SM80 kernel and uses them for batch slice selection of A/B/scales/biases. |
mlx/backend/cuda/quantized/qmm/qmm_impl_naive.cuh |
Adds optional gather index parameters to naive kernel and uses them for batch slice selection of A/B/scales/biases. |
mlx/backend/cuda/quantized/qmm/qmm.h |
Extends qmm_sm80 / qmm_naive APIs with optional gather index pointer parameters. |
mlx/backend/cuda/quantized/qmm/qmm.cu |
Wires new optional gather index pointer parameters through the QMM dispatch into qmm_impl_sm80 / qmm_impl_naive. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
an untested sm90 kernel using CUTLASS kArray mode (fused gather via per-batch pointer arrays) 5787da7 |
|
On sm90: we don't need to implement the |
Ideally we want to have full coverage. We shouldn't go out of our way to not implement it. But it is indeed lower priority.
|
…dLocalStream, linalg.det) Highlights: - ThreadLocalStream in C++ (ml-explore#3405) — joins StreamOrDevice variant - Segmented mm NAX kernel (ml-explore#3419) — new steel/gemm/kernels/steel_gemm_segmented{,_nax} - mlx.core.linalg additions: determinant, sign-log-determinant (ml-explore#3416) - NAX split-K tuning + addmm fix (ml-explore#3422) - gather_mm CUDA + qmm CUDA fixes (ml-explore#3414, ml-explore#3417, ml-explore#3445) - Scheduler::enqueue thread safety (ml-explore#3423) - Random kernel bytes_per_key truncation fix (ml-explore#3432) - Many smaller bug + CI fixes; nanobind 2.12.0 pin Conflict resolution (5 files; all auto-merged cleanly, no manual edits): - mlx/utils.h — `ThreadLocalStream` joins StreamOrDevice; our metal_retain_bound_buffers() helper preserved. - mlx/backend/metal/device.cpp — upstream's "no Metal device available" error message coexists with our signpost/MLX_METAL_PROFILE perf diagnostic + hazard-mode retain logic. - mlx/backend/metal/CMakeLists.txt — upstream's steel_gemm_segmented{,_nax} joins our kernel registrations. - mlx/backend/metal/kernels/CMakeLists.txt — same: upstream segmented + our fused_gate_activation kernel. - tests/scheduler_tests.cpp — upstream's "test thread local stream" joins our scheduler-test alignment for lazy command-encoder init. No disputed conflicts; the auto-merge produced a coherent result with both sides' work intact. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Followup to #3321
Proposed Changes
Add optional gather index parameters to
qmm_sm80_kernelandqmm_naive_kernel. The kernel readslhs_indices/rhs_indicesat the batch slice lookup.Performance
RTX 4070 SUPER, fp16, 4-bit affine quantization, K=4096. Baseline is
main(qmv-only path).Small M stays on qmv with no regression.