lablup / lablup/mlxcel

perf(cuda/quant): qmm_sm70 — Volta tensor-core MMA path for quantized GEMM

Open
#1,543 3 comments 0 reactions 0 assignees View on GitHub
area:core platform:linux priority:high status:investigation type:performance
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

Part of #1536. Phase 3. **Depends on #1542** (f16 activations) — Volta MMA has no bf16 variant, so there is no tensor-core path until activations are f16. Baseline from #1538; rebases onto #1541 if that lands first (same launch path).

This is the centerpiece of the program: it targets the single kernel holding **82.0% of GPU time** on the Volta baseline.

## Context

Every quantized GEMM on a pre-Ampere part is executed with scalar FP32 FMA on CUDA cores. `mlx/backend/cuda/device/gemm_sm70.cuh:42-63`:

```cpp
template
inline constexpr auto make_tiled_mma(CtaTiler cta_tiler) {
using Atom = cuda::std::conditional_t<
SM80,
cuda::std::conditional_t,
SM80_16x8x16_F32F16F16F32_TN,
cuda::std::conditional_t,
SM80_16x8x16_F32BF16BF16F32_TN,
UniversalFMA>>,
UniversalFMA>; // <- every cc < 8
if constexpr (!SM80 || is_same_v) {
return make_tiled_mma(Atom{}, Layout>{});
```

The file is named `gemm_sm70` but contains no Volta tensor-core path; the name marks the minimum arch it *compiles* for. Consumers: `device/qmm_naive.cuh:63,233` and `gemms/gather_gemm.cu:54`.

Measured cost on `qwen3.8-27B-4bit` (#1538): `qmm_naive` at bf16 with a 64x64x64 tile is 82.0% of GPU time (994 launches, 29.3 ms average), and prefill achieves roughly **0.41 TFLOPS against a 14 TFLOPS FP32 peak (~3%) and a 112 TFLOPS FP16 tensor-core peak (~0.4%)**.

The headroom is 8x on peak alone (14 -> 112 TFLOPS), and the measured 3% attainment says the realizable gap is larger than that.

## Why this is tractable

Three things make an SM70 MMA path a contained change rather than a from-scratch kernel:

1. **The atoms already exist in the vendored CUTLASS.** `cute/arch/mma_sm70.hpp` defines `SM70_8x8x4_F32F16F16F32_TN` (and the `F16F16F16F16` and NT/NN/TT variants), with `MMA_Traits` specializations in `cute/atom/mma_traits_sm70.hpp`. Verified present in this tree.
2. **The copy path is already pre-Ampere.** `qmm_naive` stages through registers with `UniversalCopy` (`device/qmm_naive.cuh:68-85`, `tArA`/`tBrB` fragments), not `cp.async`. Unlike `qmm_sm80.cuh` — which is built on `SM80_CP_ASYNC_CACHEALWAYS` (`patches/.../device/qmm_sm80.cuh:139-140`) and would need its whole pipeline rewritten — `qmm_naive` needs no async-copy surgery.
3. **The dispatch seam exists.** `quantized.cpp` already branches on `can_use_qmm_sm90` / `can_use_qmm_sm80` / `can_use_qmm_naive` (`:229-255`) with per-path `supports_*` predicates in `qmm/qmm.cu`. An `sm70` arm follows the established shape.

So the minimum viable change is: teach `make_tiled_mma` a third branch selecting `SM70_8x8x4_F32F16F16F32_TN` for `half_t` on pre-Ampere, with the accompanying thread/tile layout.

## Scope

A Volta tensor-core path for quantized GEMM, reached by `qmm_naive` and `gather_gemm`.

Out of scope: attention (tracked separately), `qmm_sm80`/`qmm_sm90` (unchanged), and any sm_75 (Turing) specialization — Turing is not a target of #1536 and its `m16n8k8` atom is a different shape.

## Implementation plan

1. **Overlay `gemm_sm70.cuh`** into `src/lib/mlx-cpp/patches/mlx/backend/cuda/device/` (it is not currently overlaid — see the inventory in #1536). Add the SM70 branch. Note the `Layout>` in the non-SM80 return is sized for `UniversalFMA`; the `8x8x4` atom needs its own thread layout and `Tile<>` — derive it rather than reusing the FMA one.
2. **Dequantization placement.** `qmm_naive` dequantizes B into `tBrB_dq` per K-iteration (`device/qmm_naive.cuh:80-85`). With an MMA atom the operand must land in the fragment layout the atom expects, which is not the FMA layout. Decide whether to dequantize into a shared-memory staging tile and then load fragments, or dequantize directly into fragment registers. This is the main design question of the issue; measure both if the answer is not obvious.
3. **Tile and pipeline.** Volta HMMA wants deeper K-blocking than the current `tile_k = max(64, group_size)`. Sweep tile shapes with the shared-memory budget from #1541. Double-buffering through registers (no `cp.async` on Volta) is the available pipelining primitive.
4. **Dispatch.** Add `supports_qmm_sm70` beside the existing predicates in `qmm/qmm.cu` and a `can_use_qmm_sm70` arm in `quantized.cpp:229-255`, ordered before the `qmm_naive` fallback. Predicate must require `compute_capability_major() == 7`, `half_t` activations, and whatever alignment the chosen tile needs; anything failing it falls back to `qmm_naive` unchanged.
5. **`gather_gemm`.** Same atom selection reaches `gemms/gather_gemm.cu:54`. Confirm the MoE path benefits and that it interacts correctly with #1544's arch-tag fix.
6. **Fallback must survive.** If any of the above proves unworkable, `qmm_naive` stays the pre-Ampere path and this issue closes as a findings doc — see below.

## Allowed outcomes

Following the precedent of #637, this is research-then-implement and a negative result is a valid close:

- **Implemented**: an `sm70` quantized GEMM meeting the criteria below.
- **Not implemented**: a findings doc in `docs/benchmark_results/` with ncu tensor-core-utilization evidence, the specific blocker (fragment layout, register pressure, dequant placement), and an upstream MLX issue filed if the gap is upstream's to close. The doc is required either way.

## Acceptance criteria

- [ ] Findings doc in `docs/benchmark_results/` with achieved-vs-ceiling TFLOPS, the kernel dispatch map on sm_70, and the go/no-go rationale. **Required regardless of outcome.**
- [ ] If implementing: prefill on `qwen3.8-27B-4bit` reaches >= 10% of the 112 TFLOPS tensor-core peak (from ~0.4%), i.e. a >= 3x prefill improvement over the #1538 baseline of ~7.7 tok/s.
- [ ] If implementing: `ncu` confirms non-zero tensor-core pipe utilization (`sm__inst_executed_pipe_tensor` or the roofline set) on the new kernel. A speedup without tensor-core traffic means something else was measured.
- [ ] No decode regression: `qmv` still handles `M * B < 8` and its rate is unchanged.
- [ ] Parity: `qmm_sm70` output against `qmm_naive` across bits 4/8, group sizes 32/64/128, and both transpose settings, within the tolerance the existing quantized-matmul tests use. Model-level greedy parity on a dense and a MoE checkpoint.
- [ ] Predicate falls back cleanly: forcing `supports_qmm_sm70` false reproduces today's behavior exactly.
- [ ] **Zero change on sm_80+**: GB10 greedy output byte-identical and baseline throughput unmoved. The new branch must be unreachable above cc 7.
- [ ] `cargo test --features cuda` green on sm_70 and sm_121; both `70` and the full x86_64 release arch matrix build.

## Validation

```
MLX_CUDA_ARCHITECTURES=70 make release-cuda
P=$(python3 -c "print(' '.join(['The quick brown fox jumps over the lazy dog.']*60))")
./target/release/mlxcel generate -m ./models/qwen3.8-27B-4bit -p "$P Summarize." -n 8
nsys profile -t cuda,nvtx --cuda-graph-trace=node -o sm70_after \
./target/release/mlxcel generate -m ./models/qwen3.8-27B-4bit -p "$P Summarize." -n 8
ncu --set roofline --kernel-name regex:qmm_sm70 --launch-count 4 ./target/release/mlxcel generate \
-m ./models/qwen3.8-27B-4bit -p "$P Summarize." -n 1
```

Note: `ncu` on this workload conflicts with cuDNN handle creation (`cudnnCreate ... CUDNN_STATUS_NOT_INITIALIZED`) during the audit; run with `MLX_CUDA_USE_CUDNN_SDPA=0` if it reproduces.

## References

- The FMA-only selection this issue replaces: `mlx/backend/cuda/device/gemm_sm70.cuh:42-63`.
- Consumers: `mlx/backend/cuda/device/qmm_naive.cuh:63,233`; `mlx/backend/cuda/gemms/gather_gemm.cu:54`.
- Volta atoms and traits: `cute/arch/mma_sm70.hpp`, `cute/atom/mma_traits_sm70.hpp` (vendored CUTLASS).
- Contrast case showing why `qmm_sm80` is not the template to copy: `patches/mlx/backend/cuda/device/qmm_sm80.cuh:139-140` (`SM80_CP_ASYNC_CACHEALWAYS`).
- Dispatch seam: `patches/mlx/backend/cuda/quantized/quantized.cpp:229-255`; predicates `mlx/backend/cuda/quantized/qmm/qmm.cu:20-190`.
- Precedent for an arch-specialized GEMM issue with an allowed negative result: #637.

Contributor guide

Open the contributing guide

Research direction

Start with gemm_sm70.cuh, qmm_naive.cuh, qmm.cu, and quantized.cpp to trace the existing FMA path, dispatch predicates, and dequantization layout; inspect the SM70 atoms in the vendored CUTLASS files. Build with MLX_CUDA_ARCHITECTURES=70 and run the listed mlxcel, nsys, and ncu commands, using MLX_CUDA_USE_CUDNN_SDPA=0 if needed. Done means either a validated SM70 path meeting the parity, performance, tensor-core, fallback, and build criteria, or a findings document in docs/benchmark_results/.

Written by the indexing model from the issue text.

Assessment

Tech stack
cpp, rust
Domain
backend, machine-learning, performance
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.