perf(cuda/quant): Turing (sm_75) needs its own tensor-core arm
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Follow-on to epic #1536 (NVIDIA Volta sm_70 inference acceleration) and to #1543 (the Volta tensor-core MMA path for quantized GEMM). This issue records work that is understood but deliberately not started, because the project has no Turing hardware to build on or measure against. It is written so that whoever picks it up on an sm_75 part has the whole decision already made.
## Problem / Background
Two different architecture boundaries are in play in the pre-Ampere work, and they are not the same boundary. Conflating them is the failure mode this issue exists to prevent.
The first boundary is the **activation dtype policy**, and it falls at compute capability 8.0. Neither Volta nor Turing has any bf16 hardware, so a single `cc < 8` predicate covers both and no split is needed. Verified against the CUTLASS 4.4.2 headers vendored into this tree: `cute/arch/mma_sm70.hpp` and `cute/arch/mma_sm75.hpp` contain zero BF16 MMA atoms, while `cute/arch/mma_sm80.hpp` contains two (`SM80_16x8x8_F32BF16BF16F32_TN` at line 191 and `SM80_16x8x16_F32BF16BF16F32_TN` at line 224). The same split shows at the CUTLASS level: `cutlass/arch/mma_sm80.h` names `bfloat16_t` 12 times, `cutlass/arch/mma_sm70.h` and `cutlass/arch/mma_sm75.h` name it zero times. CUDA agrees: `/usr/local/cuda/include/cuda_bf16.hpp:2557` and following guard bf16 arithmetic on `__CUDA_ARCH__ >= 800`. So the f16 activation policy tracked in #1542 correctly uses one pre-Ampere predicate for both generations.
The second boundary is the **tensor-core MMA atom**, and it is a three-way split, because the atom shapes differ per architecture. Volta at sm_70 has `SM70_8x8x4_F32F16F16F32_TN` (the `884` HMMA shape, plus its NN/NT/TT and F16-accumulate siblings, 8 atoms in total). Turing at sm_75 has `SM75_16x8x8_F32F16F16F32_TN` and additionally an int8 atom, `SM75_8x8x16_S32S8S8S32_TN`, that Volta does not have at all. Ampere and later have `SM80_16x8x16`. A single pre-Ampere arm cannot serve both Volta and Turing once tensor cores are actually involved.
## Current Behavior
`src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu:389-402` tags every part below compute capability 8.0 with `cutlass::arch::Sm70`, through `mlxcel::grouped_gemm_arch_for` in `src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_arch.h`. Nothing is broken today, and #1544 measured why rather than assuming it: the configuration that arm selects is the primary `GemmConfiguration` template at `grouped_gemm_unaligned.cu:205-217`, which is `cutlass::arch::OpClassSimt` with `InstructionShape<1, 1, 1>`, so there is no MMA atom for the tag to choose, CUTLASS erases the tag, and one arm can serve 7.0 through 7.5. Retagging that arm from `Sm75` to `Sm70` left the same 51 device symbols with byte-identical bodies across 58,211,476 bytes of per-symbol SASS text (`docs/benchmark_results/grouped-gemm-arch-v100-2026-08-31.md`).
The moment that arm is given a tensor-core operator, the tag becomes load-bearing and a Turing part tagged `Sm70` silently loses `m16n8k8`. Two `static_assert`s added by #1544 at `grouped_gemm_unaligned.cu:262-275` fail the build on that day rather than letting it pass silently, and their message names this exact follow-up ("Turing needs its own arm again"). This issue is the work those assertions point at.
`GroupedGemmArch` in `grouped_gemm_arch.h` deliberately has no `Sm75` enumerator today, and its comment says why: an enumerator there is a template argument the whole GEMM gets instantiated over, so a tag the function never returns would emit a dead copy of the arm (measured at 26 host-side functions and 194,704 bytes of object) for device code that is byte-identical to the `Sm70` arm's.
## Proposed Solution
Split the pre-Ampere arm into two, so compute capability 7.0 and 7.5 select different configurations, and give Turing a tensor-core `GemmConfiguration` specialization built on its own atom rather than borrowing Volta's.
Concretely, in `grouped_gemm_arch.h`: add `Sm75 = 75` to `GroupedGemmArch` and replace the comment that explains its deliberate absence with the reason it now exists. `grouped_gemm_arch_for` currently takes only the compute capability **major** version, which is exactly what makes 7.0 and 7.5 indistinguishable to it, so its signature has to grow the minor version (`grouped_gemm_arch_for(int major, int minor)`). The device side already has the value: `cu::Device::compute_capability_minor()` is used at `src/lib/mlx-cpp/patches/mlx/backend/cuda/jit_module.cpp:287`. Update the caller at `grouped_gemm_unaligned.cu:390` to pass both, and add the `case mlxcel::GroupedGemmArch::Sm75:` arm to the switch at `grouped_gemm_unaligned.cu:391-401`.
The C shim `mlxcel_grouped_gemm_arch_for` in `src/lib/mlxcel-core/cpp/grouped_gemm_arch_probe.cpp` takes the same signature change, keeping its `int` in, `int` out ABI so the host-only enumeration test keeps working with no CUDA toolkit. `src/lib/mlxcel-core/src/grouped_gemm_arch_tests.rs` then needs its `extern "C"` declaration and `arch_for` helper updated, `turing_shares_the_pre_ampere_arm` (line 151) inverted into a test that asserts 7.5 selects `SM75` while 7.0 still selects `SM70`, and `only_the_pre_ampere_arm_changed` and `arch_mapping_table` extended so the enumeration still covers every architecture and still proves that nothing at compute capability 8 or above moved.
Then add the Turing tensor-core `GemmConfiguration` specialization next to the two existing sm_80 ones at `grouped_gemm_unaligned.cu:219-248`, constrained on the Turing tag rather than on `Arch::kMinComputeCapability >= 80`, with `OpClass = cutlass::arch::OpClassTensorOp` and an `InstructionShape` matching `m16n8k8` rather than sm_80's `m16n8k16`. Keep `kStages = 2`: `cp.async` arrived with Ampere, and `grouped_gemm_unaligned.cu:277-286` asserts that no pre-Ampere configuration asks for the 3-stage pipeline.
Rejected alternative, so it is not relitigated: keeping one pre-Ampere arm and selecting the atom inside the configuration by `Arch::kMinComputeCapability`. That is what the current arm effectively does, and it works only while the answer is "no atom at all". Once the arm has an operator, CUTLASS resolves the atom from the tag it was handed, so the tag has to be the real one.
## Scope
**In scope:** `src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_arch.h` (the `Sm75` enumerator, the minor-version parameter, the rewritten rationale comment), `src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu` (the third switch arm, the Turing `GemmConfiguration` specialization, and the `static_assert` pair), `src/lib/mlxcel-core/cpp/grouped_gemm_arch_probe.cpp` (the shim signature), `src/lib/mlxcel-core/src/grouped_gemm_arch_tests.rs` (the enumeration tests), and a benchmark record under `docs/benchmark_results/` carrying the measured sm_75 numbers.
**In scope, same problem:** any tensor-core work landing in the dense quantized path under #1543. That path has the identical per-architecture atom problem: `src/lib/mlx-cpp/patches/mlx/backend/cuda/quantized/quantized.cpp:190` and `:262-285` gate `qmm_sm80` on compute capability and fall back to `qmm_naive` below it, so a `qmm_sm70` built on the `884` atom is not a Turing kernel and a Turing part must not be routed into it.
**Out of scope:** the f16 activation policy itself, which is #1542 and needs no per-generation split. The Volta MMA path itself, which is #1543. Turing int8 quantized work through `SM75_8x8x16_S32S8S8S32_TN`, which has no issue yet and should get its own if anyone wants it.
## Implementation Notes
- **Prerequisite**: #1542, the f16 activation policy below Ampere. Turing tensor cores are f16-only exactly as Volta's are, so a bf16 activation reaching this path defeats the whole change on sm_75 for the same reason it does on sm_70.
- **Blocked on hardware**: this cannot start until the project has a Turing device or runner. See "Why this is not started" below.
- **Reuse**: `grouped_gemm_arch_for` stays the single definition shared by the CUDA overlay and the host test through `grouped_gemm_arch_probe.cpp`, so the tested function stays the shipped one. Do not add a Rust-side restatement of the mapping; the existing `arch_before_1544` helper in the test file is the only permitted restatement and exists solely to be compared against.
- **Constraints**: adding an enumerator to `GroupedGemmArch` instantiates the entire GEMM over one more tag. #1544 measured that cost at 26 extra host-side template instantiations and 194,704 bytes of object for a dead arm. Here the arm is not dead, so the cost is justified, but it must be recorded in the same terms rather than waved through.
- **Constraints**: `kStages` must remain 2 on both pre-Ampere arms. A 3-stage pipeline without `cp.async` is either a build failure or a silent serialization, and neither announces itself.
- **Edge cases**: compute capability 7.2 (Xavier) and anything else in the 7.x range that is neither 7.0 nor 7.5 must land on a defined arm rather than falling through the switch. Pick `Sm70` as the floor for `minor < 5`, matching the existing "the tag names the floor of the range the arm covers" rule, and assert it in the enumeration test. Anything below compute capability 7.0 keeps landing on `Sm70` as it does today.
- **Edge cases**: the `switch` in `dispatch_cutlass_arch` has no `default`, so a new enumerator that is not handled is a compiler warning at best and a silently unset function pointer at worst. `get_grouped_mm_funcion` initializes `fun` to `nullptr` precisely so an unhandled path cannot be misread as an architecture decision, so verify the null is impossible rather than merely unlikely.
- **Error handling**: no new runtime failure path. A wrong tag does not throw, it produces wrong numbers or loses the atom, which is why the guard is a `static_assert` at build time and a numeric test at run time, not a check.
## Acceptance Criteria
- [ ] A Turing device or CI runner is actually available to the project, and the benchmark record names it (model, compute capability, driver, toolkit) the way `docs/benchmark_results/volta-sm70-baseline-2026-08-31.md` names the V100.
- [ ] `grouped_gemm_arch_for` takes the compute capability minor version and returns a three-way split that separates 7.0 from 7.5, with each arm selecting an MMA atom that exists on its target: `SM70_8x8x4` for 7.0, `SM75_16x8x8` for 7.5, `SM80_16x8x16` for 8.0 and later.
- [ ] The split reaches the live dispatch, not a standalone configuration: `dispatch_cutlass_arch` in `grouped_gemm_unaligned.cu` handles the new arm, and a Turing part running a quantized MoE checkpoint is observed selecting the Turing kernel (nsys, or `MLXCEL_TRACE_ARCH=1` plus the kernel name in the profile), not the Volta one and not the SIMT fallback.
- [ ] The two `static_assert`s at `grouped_gemm_unaligned.cu:262-275` are updated to guard the new invariant, with their comments rewritten to describe what is now true. Deleting them silently is not acceptable; if the invariant they encode no longer exists, the replacement assertion and the reason for the change must be in the diff.
- [ ] `grouped_gemm_arch_tests.rs` enumerates the new mapping over every architecture, `turing_shares_the_pre_ampere_arm` is replaced by tests that assert 7.5 and 7.0 diverge, and the tests still prove that compute capability 8 and above are untouched. These run with no GPU and no `cuda` feature.
- [ ] Measured before and after numbers on sm_75 hardware, written to `docs/benchmark_results/`, following the six methodology rules in `volta-sm70-baseline-2026-08-31.md` (five repetitions, warm cache, explicit architecture list, contention checked). Numbers inferred from Volta, or from an Ampere part with the Turing arm forced, do not satisfy this.
- [ ] Numeric correctness on the device, not just plausible text: `grouped_gemm_numeric_tests.rs` passes on sm_75, and a greedy continuation through the grouped path is byte-identical to the same continuation with `MLXCEL_GATHER_QMM_GROUPED=0`.
- [ ] No regression on sm_80 and later, verified by the per-symbol SASS comparison technique from #1544 rather than by inspection: compile the translation unit at each target with the flags taken from the build's `compile_commands.json`, and diff `cuobjdump --dump-sass` split by `Function :`. Whole-file diffs are misleading here, because the cubin emits the same functions in a different order.
## Verification
```bash
# Host-only enumeration of the new mapping. No GPU, no CUDA toolkit, no cuda feature.
cargo test --lib grouped_gemm_arch_tests::
# Build for Turing, and confirm the archive holds sm_75 and nothing else.
MLX_CUDA_ARCHITECTURES=75 cargo build --release --features cuda
cuobjdump --list-elf target/release/build/mlxcel-core-*/out/build/lib/libmlx.a \
| grep -oE 'sm_[0-9]+a?' | sort | uniq -c
# Numeric correctness on the device.
cargo test --release --features cuda --lib grouped_gemm_arch_tests::gather_mm
cargo test --release --features cuda --lib grouped_gemm_numeric_tests::
# The arm is actually reached. The gate is prompt_tokens * top_k >= 8 * num_experts,
# which is 128 prompt tokens for this checkpoint. --cuda-graph-trace=node is mandatory.
M=./models/mlx-community/gemma-4-26b-a4b-it-4bit
LONG=$(python3 -c "print(('Virtual memory lets an operating system give each process its own address space. '*40).strip())")
nsys profile -t cuda,nvtx --cuda-graph-trace=node -o turing \
./target/release/mlxcel generate -m $M -p "$LONG" -n 4
nsys stats --report cuda_gpu_kern_sum --format csv turing.nsys-rep | grep GemmGrouped
# Greedy parity against the legacy path. Expect byte-identical output.
./target/release/mlxcel generate -m $M -p "$LONG" -n 64 -t 0.0 --profile > grouped.txt
MLXCEL_GATHER_QMM_GROUPED=0 ./target/release/mlxcel generate -m $M -p "$LONG" -n 64 -t 0.0 --profile > legacy.txt
diff grouped.txt legacy.txt
# No sm_80-and-later regression, per symbol. Take the nvcc line from
# compile_commands.json so the flags are the shipped ones, swap --generate-code
# for the target under test, compile before and after, and split the dump by "Function :".
cuobjdump --dump-sass grouped_gemm_unaligned.cu.o
```
A pass is: the enumeration tests green on any host, the numeric tests green on sm_75, `GemmGrouped` kernels present in the sm_75 profile with the Turing configuration in their demangled name, `diff` silent, per-symbol SASS identical at `compute_80` and `compute_121`, and a measured sm_75 delta recorded with its repetition spread.
## Technical Considerations
**Why this is not started.** There is no Turing device available to the project. The development machine is a single Tesla V100 at compute capability 7.0, and the CUDA runners carry CUDA 13, which removed Volta support entirely: `nvcc` there fails with `nvcc fatal : Unsupported gpu architecture 'compute_70'` before compiling anything. `.github/workflows/ci.yml:690-717` already probes for this and skips the `cuda-sm70-compile` job loudly rather than red-lighting PRs for a toolkit limitation. So the fleet cannot build for sm_70, let alone measure sm_75. Nothing in this issue can be validated without hardware, and it deliberately makes **no performance claim for Turing**.
**Supporting measurement from Volta, for context only.** On a Tesla V100 with cuBLAS at the GEMM shapes this project's prefill actually runs (m=2048), bf16 reaches 9.6 TFLOPs, f32 reaches 14.0, and f16 reaches 95 to 100. bf16 is the slowest of the three because it is emulated rather than executed, which is what motivates #1542. The corresponding measurement for Turing has not been taken and must not be assumed to match: Turing's `m16n8k8` is a different atom on a different SM, and its f32 and f16 peaks differ per SKU.
**Related:** epic #1536 (the Volta program and its measured baseline), #1542 (f16 activation policy, prerequisite), #1543 (Volta tensor-core MMA path, the sibling with the same atom problem in the dense quantized path), #1544 (the `Sm70` retag that established the current invariant and left the assertions this issue answers). Records: `docs/benchmark_results/volta-sm70-baseline-2026-08-31.md` and `docs/benchmark_results/grouped-gemm-arch-v100-2026-08-31.md`.
Contributor guide
Research direction
Start with grouped_gemm_arch.h and grouped_gemm_arch_tests.rs to understand the architecture mapping, then trace the dispatch and configuration in grouped_gemm_unaligned.cu and the C shim in grouped_gemm_arch_probe.cpp. Run the host mapping tests first; done requires an available sm_75 device, passing grouped_gemm_numeric_tests.rs, live Turing dispatch, and measured results recorded under docs/benchmark_results/.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, rust
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 30/100