fix(cuda/moe): grouped GEMM selects cutlass::arch::Sm75 on an sm_70 part #1544

Part of #1536. Phase 4. Independent of #1539-#1543; can land any time after #1538's baseline. Coordinate with #1543 if that also touches gather_gemm.

Verification-first: this may be a latent defect, a dead branch, or benign. The point of the issue is to find out, because "MoE on Volta silently takes an unsupported-arch path" is not a state anyone should have to guess about.

Context

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

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>{});
  }
}

and the default at :349: auto* fun = grouped_gemm_v2<GemmConfiguration<float, cutlass::arch::Sm75>>;

cutlass::arch::Sm75 selects Turing MMA (m16n8k8), which does not exist on sm_70. The < 8 branch is written as though pre-Ampere means Turing; Volta is a generation below that and has only the 8x8x4 HMMA shape. CUTLASS guards these behind CUTLASS_ARCH_MMA_SM75_SUPPORTED, so compiling an Sm75 configuration for an sm_70 target does not necessarily fail the build — it can compile to a path that traps or degenerates at runtime.

Observed during the #1536 audit: gemma-4-26b-a4b-it-4bit (MoE) generates correct output on the V100. So this is not an outright crash. Either the branch is unreachable for that checkpoint (mlxcel's fused-MoE path may bypass grouped_gemm_v2 entirely), or CUTLASS is falling back to something functional. Neither has been confirmed.

Scope

Determine what actually executes on sm_70, and make the arch tag correct for whatever that is.

Out of scope: MoE performance tuning beyond fixing the arch tag. If the path turns out to be dead on Volta, saying so precisely and adding a guard is the whole deliverable.

Implementation plan

1. Determine reachability first. Instrument dispatch_cutlass_arch (or nsys the MoE run with --cuda-graph-trace=node) to record whether grouped_gemm_v2 is entered at all on sm_70 for gemma-4-26b-a4b-it-4bit, and which kernel names appear. mlxcel has a fused decode-MoE path and a sorted-MoE prefill fast path (quantized.cpp, the [mlxcel #629] block); establish which of these the checkpoint takes before touching the dispatch.
2. If reachable, verify correctness rather than assuming it: compare grouped-GEMM output against a reference (dense per-expert matmul, or the same run on CPU) on the shapes the model actually uses. A silently-degenerate CUTLASS path can produce plausible-looking but wrong numbers, which greedy text output will not reveal.
3. Fix the tag. Replace the < 8 arm with cutlass::arch::Sm70 for compute_capability_major() == 7, keeping Sm75 for == 7.5 if Turing is ever a target (it is not today — do not add speculative arms). Also fix the Sm75 default initializer at :349, which is the value used when dispatch_float_types finds no match.
4. Check the sibling. GemmConfiguration at :219-245 carries static const int kStages = 3; // use SM80_CP_ASYNC. Confirm the stage count is not applied to a pre-Ampere arch tag, where cp.async does not exist; a 3-stage pipeline without async copy is either a compile error or a silent serialization.
5. If unreachable on Volta, add an explicit guard that fails loudly (or routes to a documented fallback) rather than leaving an arch tag that describes hardware the device does not have, and document the actual MoE path on Volta in the test(benchmark): Volta (sm_70) baseline record and build coverage #1538 baseline doc.

Acceptance criteria

• Written determination of whether grouped_gemm_v2 is reached on sm_70, with the evidence (instrumentation output or nsys kernel names) for at least gemma-4-26b-a4b-it-4bit. Reached. 180 cutlass::Kernel<GemmGrouped> launches at 3.8% of GPU time on a 573-token prompt, absent under MLXCEL_GATHER_QMM_GROUPED=0 and absent at a 46-token prompt, which is below perf(cuda/moe): fix MoE prefill collapse on CUDA (gather_qmm / grouped GEMM path) #629's B >= 8 * num_experts gate and is why the test(benchmark): Volta (sm_70) baseline record and build coverage #1538 profile showed none.
• If reached: numerical correctness of grouped-GEMM output verified against a reference on the model's real shapes. Byte-identical 64-token greedy continuation against the legacy qmm_naive path, plus new unit tests against an f64 dense per-expert reference at k = 2816, n = 704 and its transpose, over both entry points, both alignment arms, both operand layouts and f32/bf16/f16.
• Arch tag matches the hardware for every branch, including the default initializer at :349. The pre-Ampere arm is cutlass::arch::Sm70; the :349 placeholder is replaced by a named function-pointer type initialised to nullptr with a loud guard, so no architecture is named there at all.
• kStages / cp.async consistency confirmed for the pre-Ampere configuration. kStages = 3 belongs to an explicit full specialization on cutlass::arch::Sm80 and is unreachable from any pre-Ampere tag; the pre-Ampere arm gets kStages = 2 and runs MmaPipelined. Pinned by static_assert.
• MoE decode and prefill tok/s on Volta reported before and after against test(benchmark): Volta (sm_70) baseline record and build coverage #1538. 4-bit decode 29.66 to 29.73 ms/token, 8-bit 31.68 to 32.71, 4-bit prefill 12,720 to 12,644 ms; every delta inside the before arm's own repetition spread, which byte-identical device code requires. The test(benchmark): Volta (sm_70) baseline record and build coverage #1538 row is re-measured rather than quoted, because perf(cuda/quant): use float accumulators in qmv below Ampere at bits < 8 #1539 moved MoE decode after that record was written.
• Zero change on sm_80+: the == 8 and > 8 branches untouched; GB10 MoE output byte-identical and throughput unmoved. Left unticked because its GB10 half needs an sm_80+ part and this host has none. The rest is closed here rather than deferred: the == 8 and > 8 arms are textually untouched, only_the_pre_ampere_arm_changed enumerates the tag mapping over every compute capability and shows nothing at or above 8 moves, and compiling the translation unit before and after at compute_80 and compute_121 yields the same 51 device symbols with byte-identical SASS bodies at both.
• If the path is dead on Volta, that is stated in the test(benchmark): Volta (sm_70) baseline record and build coverage #1538 baseline doc along with what runs instead. The path is not dead, so the test(benchmark): Volta (sm_70) baseline record and build coverage #1538 doc is annotated with the opposite finding: what actually runs, the prompt-length gate that decides it, and why that document's own MoE profile showed no grouped GEMM.
• cargo test --features cuda green on sm_70 and sm_121. sm_121 needs a GB10 host. On sm_70 the new grouped_gemm_arch_tests module is green (7 of 7) and the mlxcel-core library suite was run in full; see the PR test plan for what was executed.

Validation

MLX_CUDA_ARCHITECTURES=70 make release-cuda
nsys profile -t cuda,nvtx --cuda-graph-trace=node -o moe_volta \
  ./target/release/mlxcel generate -m ./models/mlx-community/gemma-4-26b-a4b-it-4bit -p "Hi." -n 32
nsys stats --report cuda_gpu_kern_sum --format csv moe_volta.nsys-rep

References

• Dispatch and default: src/lib/mlx-cpp/patches/mlx/backend/cuda/gemms/grouped_gemm_unaligned.cu:338-355.
• Stage count comment: same file, :219-245.
• Upstream's separate sm80 flag in the gather path, which perf(cuda/quant): qmm_sm70 — Volta tensor-core MMA path for quantized GEMM #1543 also touches: mlx/backend/cuda/gemms/gather_gemm.cu:54.
• Prior MoE prefill work on CUDA: perf(cuda/moe): fix MoE prefill collapse on CUDA (gather_qmm / grouped GEMM path) #629.