fix(cuda/moe): tag the pre-Ampere grouped GEMM arm Sm70, not Sm75

CHANGELOG.md

@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 
 ### Fixed
 
+- **CUDA grouped GEMM no longer tells CUTLASS a Volta part is a Turing part** (#1544). `dispatch_cutlass_arch` mapped every device below compute capability 8.0 to `cutlass::arch::Sm75`, which names the `m16n8k8` MMA that Turing introduced and an sm_70 part does not have, and `get_grouped_mm_funcion` opened with a matching `Sm75` placeholder. The pre-Ampere arm now selects `cutlass::arch::Sm70` and the placeholder is gone. **The branch is live, and on the checkpoint the epic benchmarks.** #629's sorted-MoE prefill fast path routes a quantized `GatherQMM` into `cutlass_grouped_gemm_unaligned` once the batch clears `B >= 8 * num_experts`, which for `gemma-4-26b-a4b-it-4bit` means a prompt of 128 tokens or more; an nsys profile at 573 prompt tokens shows 180 `cutlass::Kernel<GemmGrouped>` launches taking 3.8% of GPU time, and the same profile at the 46-token prompt the Volta baseline used shows none, which is why nobody had seen it. **Nothing was computing wrong.** The pre-Ampere arm resolves to `GemmConfiguration`'s primary template, which is `OpClassSimt` with `InstructionShape<1, 1, 1>`, so no arch tag can reach an MMA atom there and CUTLASS erases the tag: the running kernel is an `MmaSimt` / `OpMultiplyAdd` / `MmaPipelined` instantiation whose mangled name carries no architecture token at all. Correctness was measured rather than assumed, because a mismatched CUTLASS path can return plausible wrong numbers that greedy decoding hides: the grouped path and the legacy `qmm_naive` path produce a byte-identical 64-token greedy continuation, and new unit tests compare `gather_mm` against an `f64` dense per-expert reference on the model's real expert dims across both entry points, both output-alignment arms, both operand layouts and f32, bf16 and f16. The retag moves no device code on any architecture, which was checked rather than argued: compiling the translation unit before and after at `compute_70`, `compute_80` and `compute_121` yields the same 51 device symbols with byte-identical SASS bodies at all three, over 144 MB of dump compared per symbol. The architecture decision now lives in `gemms/grouped_gemm_arch.h` as a pure function of the compute capability major version, enumerated over every architecture by a host-side test through a C shim with no GPU involved, which is what closes this issue's "zero change on sm_80+" criterion locally instead of deferring it to GB10. Two `static_assert`s pin the preconditions: that the pre-Ampere configuration is still SIMT, without which one arm could not cover Volta and Turing together, and that its stage count is 2, since the 3-stage `SM80_CP_ASYNC` pipeline is bound to an explicit `cutlass::arch::Sm80` specialization and `cp.async` does not exist before Ampere. MoE decode and prefill are unmoved on a V100, as byte-identical device code requires. Full record in `docs/benchmark_results/grouped-gemm-arch-v100-2026-08-31.md`.
 - **`mlxcel serve` and `mlxcel-server` match the low-risk llama-server b10621 deployment surface** (#1430). Canonical `LLAMA_ARG_*` variables and `--temperature` now work on both entry points, server sampling defaults match the nightly, OpenAI-shaped requests accept llama aliases, scalar stops and seed `-1`, safe health/reranking route aliases reach the existing handlers, `f16` maps exactly to FP16 split-cache storage, and unsupported cache-reuse chunk sizes or GGML cache quantizers fail instead of changing unrelated behavior.
 
 ## [v0.6.0] - 2026-08-22

docs/benchmark_results/volta-sm70-baseline-2026-08-31.md

@@ -227,6 +227,8 @@ Net over the decode kernels is +1.77 ms/token favoring 8-bit against a measured
 
 This also **corrects a claim made while this issue was open**, that MoE decode routes through `GatherQMM` to `qmm_naive` because the expert count makes `M * B >= 8`. On this build it does not. The expert path in decode is mlxcel's own fused kernels, `custom_kernel_moe_gateup_kernel_cu_*` and `custom_kernel_moe_down_kernel_cu_*`, firing 3,570 times each, which is 30 layers times 119 steps. `qmm_naive` appears with a prefill-sized 326 instances in both arms. The `qmv` inversion is present in the MoE profile at the same 1.76x ratio the earlier audit found (2.6224 s against 1.4921 s over identical 28,084 instances), but the fused expert kernels move the other way and nearly cancel it, which is why the MoE advantage is 1.13x rather than the dense pair's 1.89x.
 
+**What this profile does not show, and why** (added by #1544). No CUTLASS grouped GEMM appears anywhere in it, which reads as though MoE on this build never touches that path. It does, in prefill, above a prompt-length gate this profile sits below. #629's sorted-MoE prefill fast path (`patches/mlx/backend/cuda/quantized/quantized.cpp:359`) routes a quantized `GatherQMM` into `cutlass_grouped_gemm_unaligned` once `B >= min_rows * E`, with `min_rows` defaulting to 8; for this checkpoint `E = 128` and `top_k = 8`, so the gate is 128 prompt tokens and the 46-token prompt used here gives `B = 368`. Re-profiled at 573 prompt tokens, the same binary and checkpoint launch `cutlass::Kernel<GemmGrouped>` 180 times for 3.8% of GPU time, and prefill runs 2.22x faster than with the path disabled. The 326 `qmm_naive` instances above reproduce exactly at the short prompt, which is the cross-check that the two profiles are the same measurement at different prompt lengths. See `grouped-gemm-arch-v100-2026-08-31.md`.
+
 ### Practical guidance
 
 Prefer an 8-bit checkpoint over a 4-bit one on Volta wherever it fits in 32 GB. This is a workaround for #1539 rather than a substitute for it: once `qmv` uses float accumulators below Ampere, the 4-bit arm should regain its bandwidth advantage and this guidance should be re-measured and probably reversed.
@@ -335,7 +337,7 @@ To be filled as epic #1536's remaining items land. Each row is a re-run of the m
 | #1541 | `qmm_naive` tile sized from the device shared-memory budget | dense prefill `qmm_naive` 10.0733 s at 329 inst | dense prefill `qmm_naive` 10.867 s at 329 inst | **No change, by measurement.** The 128-wide tile the issue set out to unlock fits Volta's budget four times over, and loses anyway: 255 registers against 224 and a 128-byte spill, no occupancy gain (2 blocks/SM either way), and a halved CTA count on a grid already smaller than 80 SMs. `qmm_naive` measures 1.50x slower at a 106-token prompt, 1.06x at 516 and 1.02x at 4,106, so the width stays where upstream had it and the shared-memory decision becomes a real device query. The 7.9% gap between the two columns is between sessions, not between tile widths: both widths agree to 0.04% in every profile and `qmv` in the same runs reproduces #1539 to 0.04%. Full record: `qmm-naive-tile-v100-2026-08-31.md` |
 | #1542 | f16 activation policy below Ampere | dense 4-bit 124.41 ms/tok; qwen 220.33 ms/tok | | |
 | #1543 | `qmm_sm70`, Volta tensor-core MMA for quantized GEMM | qwen prefill 8.00 tok/s marginal, 2.83% of FP32 peak | | |
-| #1544 | grouped GEMM arch tag on an sm_70 part | MoE 4-bit 38.61 ms/tok, 8-bit 34.29 ms/tok | | |
+| #1544 | grouped GEMM arch tag on an sm_70 part | re-measured on this branch: MoE 4-bit 29.66 ms/tok, 8-bit 31.68 ms/tok; 4-bit prefill 12,720 ms at a 285-token prompt | MoE 4-bit 29.73 ms/tok, 8-bit 32.71 ms/tok; 4-bit prefill 12,644 ms | **No change, and byte-identical device code is why.** The pre-Ampere arm was tagged `cutlass::arch::Sm75` on a part that has no `m16n8k8` MMA. The branch is live here, not dead: #629's sorted-MoE prefill path routes a quantized `GatherQMM` into `cutlass_grouped_gemm_unaligned` once `B >= 8 * num_experts`, which for this checkpoint is a 128-token prompt, and an nsys profile at 573 prompt tokens shows 180 grouped-GEMM launches at 3.8% of GPU time. The kernel it selects is `MmaSimt` / `OpMultiplyAdd`, so the tag never reached an MMA atom and the output was never wrong; the grouped path and the legacy `qmm_naive` path give a byte-identical 64-token greedy continuation. Retagging to `Sm70` leaves the same 51 device symbols with byte-identical SASS at `compute_70`, `compute_80` and `compute_121`, so both columns above are the same machine code and every delta sits inside the before arm's own repetition spread. **The baseline column is re-measured rather than quoted**: this document's own 38.61 and 34.29 predate #1539. Full record: `grouped-gemm-arch-v100-2026-08-31.md` |
 | #1545 | CUDA graph instantiation and JIT module load in TTFT | qwen fixed prefill cost 18.41 s; graph APIs 3.3 s of a 24-token run | | |
 
 ## Reproduce

src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_arch.h

@@ -0,0 +1,105 @@
+// Added by mlxcel (lablup/mlxcel#1544). Not an upstream MLX file.
+//
+// The CUTLASS architecture tag selection for the grouped GEMM in
+// `grouped_gemm_unaligned.cu`, factored out of that file as a pure integer
+// function so it can be enumerated on the host with no CUDA device, no CUDA
+// toolkit and no NVIDIA hardware at all. `grouped_gemm_unaligned.cu` calls it
+// with `cu::Device::compute_capability_major()`, and
+// `src/lib/mlxcel-core/cpp/grouped_gemm_arch_probe.cpp` exposes the same
+// function to `grouped_gemm_arch_tests.rs` through a C shim. Both callers
+// share this one definition, so the tested function is the shipped one.
+//
+// # What was wrong
+//
+// Upstream wrote the pre-Ampere arm as:
+//
+//   if (device.compute_capability_major() < 8) {
+//     f(type_identity<cutlass::arch::Sm75>{});
+//
+// `cutlass::arch::Sm75` names Turing. A Tesla V100 is compute capability 7.0,
+// one generation below Turing, and has no `m16n8k8` MMA at all; its tensor
+// cores expose only the `8x8x4` HMMA shape. The tag therefore described
+// hardware the device does not have on every Volta part that reached this
+// dispatch.
+//
+// # Why nothing broke
+//
+// Measured rather than assumed, because "the model still produces correct
+// text" is not evidence about a GEMM. On a compute-70 build of the shipped
+// translation unit the tag never reaches device code:
+//
+//   - The configuration this arm selects is `GemmConfiguration`'s primary
+//     template, which is `cutlass::arch::OpClassSimt` with
+//     `InstructionShape<1, 1, 1>`. That is plain FFMA. The two tensor-core
+//     specializations are constrained on `Arch::kMinComputeCapability >= 80`,
+//     so no pre-Ampere tag can select an MMA atom, `m16n8k8` included.
+//   - Consequently CUTLASS erases the tag. No `cutlass::arch::Sm*` token
+//     survives into any of the 51 device symbols the translation unit emits,
+//     and retagging the arm from `Sm75` to `Sm70` leaves the same 51 device
+//     symbols with byte-identical bodies, 58,211,476 bytes of
+//     `cuobjdump --dump-sass` text compared per symbol. Only the order the
+//     cubin emits them in moves.
+//   - Compiling the file with a separate `Sm70` arm added *alongside* the
+//     `Sm75` one is the control: same 51 device symbols again, whole dump
+//     byte-identical at 493,538 lines, against 26 extra host-side template
+//     instantiations and 194,704 more bytes of object.
+//
+// So the pre-Ampere arm is one arm for a reason. Splitting Turing back out
+// buys nothing but a second host-side copy of instantiations that emit the
+// same device code. The tag names the floor of the range the arm covers rather
+// than a target, and `Sm70` is the correct floor. Full record:
+// `docs/benchmark_results/grouped-gemm-arch-v100-2026-08-31.md`.
+//
+// # When this stops being safe
+//
+// The erasure is a property of the configuration, not of CUTLASS. The moment
+// the pre-Ampere `GemmConfiguration` selects `OpClassTensorOp` (the obvious
+// candidate is Volta `884` MMA work of the kind #1543 is doing for
+// `qmm_naive`), the tag becomes load-bearing and a Turing part tagged `Sm70`
+// would lose `m16n8k8`. `grouped_gemm_unaligned.cu` carries a `static_assert`
+// on that configuration's `OpClass` so the build fails on the day that changes
+// rather than silently downgrading Turing, and the assert names this file.
+
+#pragma once
+
+namespace mlxcel {
+
+// The CUTLASS architecture tags the grouped GEMM dispatches over. The values
+// are the compute capability each tag names, so a reader does not have to know
+// CUTLASS to read a test failure.
+//
+// `Sm75` is deliberately absent. An enumerator here is a template argument
+// `grouped_gemm_unaligned.cu` instantiates the whole GEMM over, so listing a
+// tag this function never returns would emit a dead copy of every
+// instantiation in the arm: 26 host-side functions and 195 KB of object, for
+// device code that is byte identical to the `Sm70` arm's. Turing is covered by
+// `Sm70` for the reason given above.
+enum class GroupedGemmArch {
+  Sm70 = 70,
+  Sm80 = 80,
+  Sm90 = 90,
+};
+
+// Maps a device's compute capability major version to the CUTLASS arch tag the
+// grouped GEMM instantiates for it.
+//
+// Only the major version is consulted, which is the whole reason this is a
+// pure function worth testing: the pre-Ampere arm deliberately covers 7.0
+// through 7.5 with one tag (see the header comment), so a minor version would
+// be an argument the function does not use.
+//
+// Anything below compute capability 7.0 also lands on `Sm70`. MLX's CUDA
+// backend does not support those parts, so this is a floor rather than a
+// claim; it is still strictly closer to the truth than the `Sm75` they used to
+// receive.
+constexpr GroupedGemmArch grouped_gemm_arch_for(int compute_capability_major) {
+  if (compute_capability_major > 8) {
+    return GroupedGemmArch::Sm90;
+  }
+  if (compute_capability_major == 8) {
+    return GroupedGemmArch::Sm80;
+  }
+  return GroupedGemmArch::Sm70;
+}
+
+} // namespace mlxcel

src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu

@@ -7,10 +7,12 @@
 #include "mlx/backend/cuda/cutlass_utils.cuh"
 #include "mlx/backend/cuda/device.h"
 #include "mlx/backend/cuda/gemms/grouped_gemm.h"
+#include "mlx/backend/cuda/gemms/grouped_gemm_arch.h"
 #include "mlx/backend/cuda/kernel_utils.cuh"
 #include "mlx/dtype_utils.h"
 
 #include <cooperative_groups.h>
+#include <stdexcept>
 #include <cutlass/gemm/device/default_gemm_configuration.h>
 #include <cutlass/gemm/device/gemm_grouped.h>
 #include <cutlass/gemm/kernel/default_gemm_grouped.h>
@@ -245,6 +247,44 @@ struct GemmConfiguration<float, cutlass::arch::Sm80, kAlignmentC, true>
   static const int kStages = 3; // use SM80_CP_ASYNC
 };
 
+// [mlxcel #1544] The pre-Ampere arm of `dispatch_cutlass_arch` tags every part
+// below compute capability 8.0 with `cutlass::arch::Sm70`, Turing included.
+// That holds only while the configuration the arm selects stays SIMT: with
+// `OpClassSimt` and `InstructionShape<1, 1, 1>` there is no MMA atom for the
+// tag to choose, CUTLASS erases the tag, and one arm can serve 7.0 through
+// 7.5. Give the pre-Ampere arm a tensor-core operator and the tag becomes
+// load-bearing, at which point a Turing part tagged `Sm70` silently loses
+// `m16n8k8` and Turing needs its own arm again. These two assertions fail the
+// build on that day instead. The second one carries `kEnableTF32 = true`,
+// which is the arm `MLX_ENABLE_TF32` reaches, and is what rules out a
+// pre-Ampere tag ever selecting one of the two tensor-core specializations
+// above: both are constrained on `Arch::kMinComputeCapability >= 80`.
+static_assert(
+    std::is_same_v<
+        GemmConfiguration<float, cutlass::arch::Sm70, 1, false>::OpClass,
+        cutlass::arch::OpClassSimt>,
+    "pre-Ampere grouped GEMM is no longer SIMT: the architecture tag in "
+    "dispatch_cutlass_arch became load-bearing, so Turing needs its own arm "
+    "again (see gemms/grouped_gemm_arch.h)");
+static_assert(
+    std::is_same_v<
+        GemmConfiguration<float, cutlass::arch::Sm70, 8, true>::OpClass,
+        cutlass::arch::OpClassSimt>,
+    "a pre-Ampere tag selected a tensor-core configuration under "
+    "MLX_ENABLE_TF32 (see gemms/grouped_gemm_arch.h)");
+
+// [mlxcel #1544] `cp.async` arrived with Ampere. The 3-stage pipeline above is
+// commented "use SM80_CP_ASYNC" and is bound to `cutlass::arch::Sm80` by an
+// explicit full specialization, so no pre-Ampere tag can reach it; a
+// pre-Ampere configuration stays at the 2 stages `MmaPipelined` implements
+// with ordinary global-to-shared copies. Asserted rather than reasoned about,
+// because a 3-stage pipeline without `cp.async` is either a build failure or a
+// silent serialization and neither announces itself.
+static_assert(
+    GemmConfiguration<float, cutlass::arch::Sm70, 8, true>::kStages == 2,
+    "a pre-Ampere grouped GEMM configuration asked for a 3-stage pipeline, "
+    "which needs the cp.async that Ampere introduced");
+
 // Get direct access to kernel.
 template <typename GemmKernel>
 class GemmGroupedEncoder
@@ -335,19 +375,57 @@ void grouped_gemm_v2(
   });
 }
 
+// [mlxcel #1544] The architecture decision itself lives in
+// `gemms/grouped_gemm_arch.h` as a pure function of the compute capability
+// major version, so it can be enumerated over every architecture on a host
+// with no GPU (`grouped_gemm_arch_tests.rs`, through
+// `cpp/grouped_gemm_arch_probe.cpp`). Upstream mapped every pre-Ampere part to
+// `cutlass::arch::Sm75`, which names Turing and so described hardware that a
+// compute capability 7.0 part does not have. That header records why the
+// corrected tag is `Sm70` for the whole pre-Ampere arm, what it measured to
+// establish that the retag moves no device code, and what has to change before
+// Turing needs an arm of its own.
 template <typename F>
 void dispatch_cutlass_arch(cu::Device& device, F&& f) {
-  if (device.compute_capability_major() < 8) {
-    f(type_identity<cutlass::arch::Sm75>{});
-  } else if (device.compute_capability_major() == 8) {
-    f(type_identity<cutlass::arch::Sm80>{});
-  } else {
-    f(type_identity<cutlass::arch::Sm90>{});
+  switch (mlxcel::grouped_gemm_arch_for(device.compute_capability_major())) {
+    case mlxcel::GroupedGemmArch::Sm70:
+      f(type_identity<cutlass::arch::Sm70>{});
+      return;
+    case mlxcel::GroupedGemmArch::Sm80:
+      f(type_identity<cutlass::arch::Sm80>{});
+      return;
+    case mlxcel::GroupedGemmArch::Sm90:
+      f(type_identity<cutlass::arch::Sm90>{});
+      return;
   }
 }
 
-auto* get_grouped_mm_funcion(Dtype dtype, int N, cu::Device& device) {
-  auto* fun = grouped_gemm_v2<GemmConfiguration<float, cutlass::arch::Sm75>>;
+// The signature every `grouped_gemm_v2` instantiation shares, named so the
+// selection below can start from no kernel at all rather than from a
+// placeholder instantiation.
+using GroupedGemmFn = void (*)(
+    bool,
+    bool,
+    int,
+    ProblemSize*,
+    int64_t*,
+    int64_t*,
+    int64_t*,
+    void*,
+    void*,
+    void*,
+    cu::CommandEncoder&);
+
+GroupedGemmFn get_grouped_mm_funcion(Dtype dtype, int N, cu::Device& device) {
+  // [mlxcel #1544] This was
+  // `grouped_gemm_v2<GemmConfiguration<float, cutlass::arch::Sm75>>`, a
+  // placeholder that reads as a pre-Ampere default and names Turing on parts
+  // that are not Turing. It was never the value returned, because
+  // `dispatch_float_types` throws on a non-float dtype and every float dtype
+  // assigns `fun`, but it did force one more template instantiation into the
+  // binary purely to serve as an initializer. `nullptr` says what the code
+  // means and cannot be misread as an architecture decision.
+  GroupedGemmFn fun = nullptr;
   dispatch_float_types(dtype, "grouped_gemm_v2", [&](auto type_tag) {
     using DataType = cutlass_type_t<MLX_GET_TYPE(type_tag)>;
     dispatch_cutlass_arch(device, [&](auto arch_tag) {
@@ -361,6 +439,14 @@ auto* get_grouped_mm_funcion(Dtype dtype, int N, cu::Device& device) {
       });
     });
   });
+  // Unreachable while `dispatch_float_types` throws on every dtype it does not
+  // dispatch, which is the contract today. Kept so that a future widening of
+  // that dispatch surfaces as a named error instead of a null call.
+  if (fun == nullptr) {
+    throw std::runtime_error(
+        "[grouped_gemm_v2] no grouped GEMM kernel was selected for the "
+        "requested dtype");
+  }
   return fun;
 }