Add NAX Split-K GEMM for large-K matmuls to improve performance (#3018)

Co-authored-by: Huan <[email protected]>
2026-01-26 11:23:20 -08:00
parent a828e769be
commit 7ed2b6b935
14 changed files with 605 additions and 20 deletions
@@ -0,0 +1,119 @@
+# Copyright © 2026 Apple Inc.
+
+import math
+import time
+
+import mlx.core as mx
+import numpy as np
+import torch
+
+N_WARMUP = 5
+N_BENCH = 20
+
+
+def bench_mlx(a, b):
+    for _ in range(N_WARMUP):
+        mx.eval(a @ b)
+
+    times = []
+    for _ in range(N_BENCH):
+        start = time.perf_counter_ns()
+        mx.eval(a @ b)
+        end = time.perf_counter_ns()
+        times.append((end - start) * 1e-9)
+
+    return np.mean(times), np.std(times)
+
+
+@torch.no_grad()
+def bench_torch(a, b):
+    for _ in range(N_WARMUP):
+        _ = a @ b
+        torch.mps.synchronize()
+
+    times = []
+    for _ in range(N_BENCH):
+        start = time.perf_counter_ns()
+        _ = a @ b
+        torch.mps.synchronize()
+        end = time.perf_counter_ns()
+        times.append((end - start) * 1e-9)
+
+    return np.mean(times), np.std(times)
+
+
+def check_correctness(out_mx, out_pt, rtol, M, N, K):
+    if not np.allclose(out_pt, out_mx, rtol=rtol, atol=0):
+        abs_diff = np.abs(out_pt - out_mx)
+        rel_diff = abs_diff / np.maximum(np.abs(out_pt), 1e-10)
+
+        print(
+            f"  WARNING: Correctness failed at {M}x{N}x{K}: "
+            f"max_abs={np.max(abs_diff):.6e}, max_rel={np.max(rel_diff):.6e}"
+        )
+
+
+def bench_gemm(M, N, K, dtype, rtol):
+    scale = 0.5 / math.sqrt(K)
+    a_np = np.random.uniform(0, scale, (M, K)).astype(np.float32)
+    b_np = np.random.uniform(0, scale, (K, N)).astype(np.float32)
+
+    a_mx = mx.array(a_np).astype(getattr(mx, dtype))
+    b_mx = mx.array(b_np).astype(getattr(mx, dtype))
+
+    a_pt = torch.from_numpy(a_np).to(dtype=getattr(torch, dtype), device="mps")
+    b_pt = torch.from_numpy(b_np).to(dtype=getattr(torch, dtype), device="mps")
+    torch.mps.synchronize()
+
+    torch_mean, torch_std = bench_torch(a_pt, b_pt)
+    mlx_mean, mlx_std = bench_mlx(a_mx, b_mx)
+
+    out_mx = (a_mx @ b_mx).astype(mx.float32)
+    out_pt = (a_pt @ b_pt).to(torch.float32).to("cpu").numpy(force=True)
+    check_correctness(out_mx, out_pt, rtol, M, N, K)
+
+    return mlx_mean, mlx_std, torch_mean, torch_std
+
+
+if __name__ == "__main__":
+    dtypes = ("bfloat16", "float16", "float32")
+
+    rtols = {
+        "float32": 1e-3,
+        "float16": 5e-3,
+        "bfloat16": 1e-2,
+    }
+
+    shapes = (
+        (2048, 2048, 10240),
+        (2048, 3072, 10240),
+        (3072, 3072, 10240),
+        (3072, 3072, 12288),
+        (3072, 4096, 12288),
+        (4096, 4096, 12288),
+        (4096, 4096, 18432),
+        (4096, 4096, 21504),
+        (4096, 6144, 21504),
+        (6144, 6144, 21504),
+    )
+
+    for dtype in dtypes:
+        print(f"\nPerformance ({dtype}):")
+        print(
+            f"{'M':>5s} {'N':>5s} {'K':>6s}  "
+            f"{'MLX (ms)':>15s}  {'Torch (ms)':>15s}  {'Speedup':>10s}"
+        )
+        print("-" * 80)
+
+        for M, N, K in shapes:
+            mlx_mean, mlx_std, torch_mean, torch_std = bench_gemm(
+                M, N, K, dtype, rtols[dtype]
+            )
+            speedup = torch_mean / mlx_mean
+
+            print(
+                f"{M:5d} {N:5d} {K:6d}  "
+                f"{mlx_mean*1000:7.2f}±{mlx_std*1000:5.2f}  "
+                f"{torch_mean*1000:7.2f}±{torch_std*1000:5.2f}  "
+                f"{speedup:8.2f}x"
+            )
@@ -87,6 +87,7 @@ if(MLX_METAL_JIT)
    kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h)
  make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax)
  make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax)
+  make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax)

  make_jit_source(quantized_nax kernels/quantized_utils.h)
  make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h
@@ -48,6 +48,7 @@ const char* steel_attention();
 const char* gemm_nax();
 const char* steel_gemm_fused_nax();
 const char* steel_gemm_gather_nax();
+const char* steel_gemm_splitk_nax();

 const char* quantized_nax();
 const char* fp_quantized_nax();
@@ -949,6 +949,40 @@ MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
 }

+MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn) {
+  const auto& lib_name = kernel_name;
+  auto lib = d.get_library(lib_name, [&]() {
+    std::ostringstream kernel_source;
+    kernel_source << metal::utils() << metal::gemm_nax()
+                  << metal::steel_gemm_splitk_nax()
+                  << get_template_definition(
+                         lib_name,
+                         "gemm_splitk_nax",
+                         get_type_string(out.dtype()),
+                         bm,
+                         bn,
+                         bk,
+                         wm,
+                         wn,
+                         transpose_a,
+                         transpose_b);
+    return kernel_source.str();
+  });
+  return d.get_kernel(kernel_name, lib, hash_name, func_consts);
+}
+
 MTL::ComputePipelineState* get_qmm_nax_kernel(
    metal::Device& d,
    const std::string& kernel_name,
@@ -286,6 +286,20 @@ MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
    int wn,
    bool rhs);

+MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array& out,
+    bool transpose_a,
+    bool transpose_b,
+    int bm,
+    int bn,
+    int bk,
+    int wm,
+    int wn);
+
 MTL::ComputePipelineState* get_qmm_nax_kernel(
    metal::Device& d,
    const std::string& kernel_name,
@@ -107,7 +107,8 @@ set(STEEL_NAX_HEADERS
    steel/utils/type_traits.h
    steel/utils/integral_constant.h
    steel/gemm/kernels/steel_gemm_fused_nax.h
-    steel/gemm/kernels/steel_gemm_gather_nax.h)
+    steel/gemm/kernels/steel_gemm_gather_nax.h
+    steel/gemm/kernels/steel_gemm_splitk_nax.h)

 set(STEEL_NAX_ATTN_HEADERS
    steel/defines.h
@@ -156,6 +157,7 @@ if(NOT MLX_METAL_JIT)

    build_kernel(steel/gemm/kernels/steel_gemm_fused_nax ${STEEL_NAX_HEADERS})
    build_kernel(steel/gemm/kernels/steel_gemm_gather_nax ${STEEL_NAX_HEADERS})
+    build_kernel(steel/gemm/kernels/steel_gemm_splitk_nax ${STEEL_NAX_HEADERS})

    build_kernel(quantized_nax quantized_nax.h ${STEEL_NAX_HEADERS})
    build_kernel(fp_quantized_nax fp4.h fp8.h fp_quantized_nax.h
@@ -29,7 +29,10 @@ template <
 auto gemm_loop(
    const device T* A,
    const device T* B,
-    const constant GEMMParams* params [[buffer(4)]],
+    int lda,
+    int ldb,
+    int K,
+    int gemm_k_iterations_aligned,
    const short sgp_sm,
    const short sgp_sn) {
  constexpr short TM = SM / UM;
@@ -51,7 +54,7 @@ auto gemm_loop(
  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
  Dtile.clear();

-  int gemm_k_iterations_ = params->gemm_k_iterations_aligned;
+  int gemm_k_iterations_ = gemm_k_iterations_aligned;

  STEEL_PRAGMA_NO_UNROLL
  for (int kk0 = 0; kk0 < gemm_k_iterations_; kk0++) {
@@ -65,23 +68,23 @@ auto gemm_loop(

      volatile int compiler_barrier;

-      const int A_offset = transpose_a ? k * params->lda : k;
-      const int B_offset = transpose_b ? k : k * params->ldb;
+      const int A_offset = transpose_a ? k * lda : k;
+      const int B_offset = transpose_b ? k : k * ldb;

      if constexpr (kAlignedM) {
-        Atile.load(A + A_offset, params->lda);
+        Atile.load(A + A_offset, lda);
      } else {
        const short rmax = transpose_a ? SK : sgp_sm;
        const short cmax = transpose_a ? sgp_sm : SK;
-        Atile.load_safe(A + A_offset, params->lda, short2(cmax, rmax));
+        Atile.load_safe(A + A_offset, lda, short2(cmax, rmax));
      }

      if constexpr (kAlignedN) {
-        Btile.load(B + B_offset, params->ldb);
+        Btile.load(B + B_offset, ldb);
      } else {
        const short rmax = transpose_b ? sgp_sn : SK;
        const short cmax = transpose_b ? SK : sgp_sn;
-        Btile.load_safe(B + B_offset, params->ldb, short2(cmax, rmax));
+        Btile.load_safe(B + B_offset, ldb, short2(cmax, rmax));
      }

      tile_matmad_nax(
@@ -94,14 +97,14 @@ auto gemm_loop(
      (void)compiler_barrier;
    }

-    A += transpose_a ? (BK * params->lda) : BK;
-    B += transpose_b ? BK : (BK * params->ldb);
+    A += transpose_a ? (BK * lda) : BK;
+    B += transpose_b ? BK : (BK * ldb);
  }

  if constexpr (!kAlignedK) {
    simdgroup_barrier(mem_flags::mem_none);

-    const short rem_bk = params->K - gemm_k_iterations_ * BK;
+    const short rem_bk = K - gemm_k_iterations_ * BK;

    STEEL_PRAGMA_NO_UNROLL
    for (int kk1 = 0; kk1 < rem_bk; kk1 += SK) {
@@ -119,23 +122,21 @@ auto gemm_loop(
            const int k = kk1 + kk * UK;
            const short psk = max(0, rem_bk - k);

-            const int A_offset =
-                transpose_a ? (m + k * params->lda) : (m * params->lda + k);
-            const int B_offset =
-                transpose_b ? (k + n * params->ldb) : (k * params->ldb + n);
+            const int A_offset = transpose_a ? (m + k * lda) : (m * lda + k);
+            const int B_offset = transpose_b ? (k + n * ldb) : (k * ldb + n);

            {
              const short psm = kAlignedM ? SM : max(0, sgp_sm - m);
              const short rmax = transpose_a ? psk : psm;
              const short cmax = transpose_a ? psm : psk;
-              Atile.load_safe(A + A_offset, params->lda, short2(cmax, rmax));
+              Atile.load_safe(A + A_offset, lda, short2(cmax, rmax));
            }

            {
              const short psn = kAlignedN ? SN : max(0, sgp_sn - n);
              const short rmax = transpose_b ? psn : psk;
              const short cmax = transpose_b ? psk : psn;
-              Btile.load_safe(B + B_offset, params->ldb, short2(cmax, rmax));
+              Btile.load_safe(B + B_offset, ldb, short2(cmax, rmax));
            }

            subtile_matmad_nax(
@@ -191,7 +191,15 @@ template <
            UM,
            UN,
            UK,
-            AccumType>(A, B, params, sgp_sm, sgp_sn);
+            AccumType>(
+            A,
+            B,
+            params->lda,
+            params->ldb,
+            params->K,
+            params->gemm_k_iterations_aligned,
+            sgp_sm,
+            sgp_sn);
        if (use_out_source) {
          gemm_epilogue<kAlignedM.value, kAlignedN.value>(
              Dtile, C, params, addmm_params, sgp_sm, sgp_sn);
@@ -106,7 +106,14 @@ gather_mm_rhs_nax(
              UK,
              AccumType>;
          Ctile = do_gemm(
-              A, B + index * params->batch_stride_b, params, sgp_sm, sgp_sn);
+              A,
+              B + index * params->batch_stride_b,
+              params->lda,
+              params->ldb,
+              params->K,
+              params->gemm_k_iterations_aligned,
+              sgp_sm,
+              sgp_sn);

          if constexpr (kAlignedN.value) {
            if (offset_next - offset == SM) {
@@ -0,0 +1,148 @@
+// Copyright © 2026 Apple Inc.
+
+using namespace mlx::steel;
+
+constant bool align_M [[function_constant(200)]];
+constant bool align_N [[function_constant(201)]];
+
+///////////////////////////////////////////////////////////////////////////////
+// NAX Split-K GEMM kernel
+///////////////////////////////////////////////////////////////////////////////
+
+// clang-format off
+template <
+    typename T,
+    int BM,
+    int BN,
+    int BK,
+    int WM,
+    int WN,
+    bool transpose_a,
+    bool transpose_b,
+    typename AccumType = float>
+[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gemm_splitk_nax(
+    const device T* A [[buffer(0)]],
+    const device T* B [[buffer(1)]],
+    device AccumType* C [[buffer(2)]],
+    const constant GEMMSpiltKParams* params [[buffer(3)]],
+    uint simd_group_id [[simdgroup_index_in_threadgroup]],
+    uint3 tid [[threadgroup_position_in_grid]]) { // clang-format on
+
+  const int linear_tid = tid.x;
+
+  // Compute swizzled tile dimensions
+  const int tn_swizzled = params->tiles_n << params->swizzle_log;
+  const int tm_swizzled =
+      (params->tiles_m + (1 << params->swizzle_log) - 1) >> params->swizzle_log;
+  const int tiles_per_partition = tn_swizzled * tm_swizzled;
+
+  const int tid_z = linear_tid / tiles_per_partition;
+  const int xy_flat = linear_tid % tiles_per_partition;
+
+  // Decode 2D grid coordinates in swizzled space
+  const int grid_x = xy_flat % tn_swizzled;
+  const int grid_y = xy_flat / tn_swizzled;
+
+  // Apply X-Y swizzle
+  const int tid_y = (grid_y << params->swizzle_log) +
+      (grid_x & ((1 << params->swizzle_log) - 1));
+  const int tid_x = grid_x >> params->swizzle_log;
+
+  // Exit early
+  if (params->tiles_n <= tid_x || params->tiles_m <= tid_y) {
+    return;
+  }
+
+  // Calculate partition bounds
+  const int c_row = tid_y * BM;
+  const int c_col = tid_x * BN;
+  const int k_start = params->split_k_partition_size * tid_z;
+  const int k_end = min(k_start + params->split_k_partition_size, params->K);
+
+  const size_t c_row_long = size_t(c_row);
+  const size_t c_col_long = size_t(c_col);
+  const size_t k_start_long = size_t(k_start);
+
+  // Adjust pointers for split-K partition
+  A += transpose_a ? (c_row_long + k_start_long * params->lda)
+                   : (k_start_long + c_row_long * params->lda);
+  B += transpose_b ? (k_start_long + c_col_long * params->ldb)
+                   : (c_col_long + k_start_long * params->ldb);
+  C += (size_t(params->split_k_partition_stride) * tid_z) +
+      (c_row_long * params->ldc + c_col_long);
+
+  // NAX tile configuration
+  constexpr short UM = 16;
+  constexpr short UN = 32;
+  constexpr short UK = 16;
+  constexpr short SM = BM / WM;
+  constexpr short SN = BN / WN;
+  constexpr short SK = 32;
+
+  constexpr short TM = SM / UM;
+  constexpr short TN = SN / UN;
+
+  // Calculate simdgroup offsets and alignment
+  const short tm = SM * (simd_group_id / WN);
+  const short tn = SN * (simd_group_id % WN);
+
+  const short sgp_sm = align_M ? SM : min(SM, short(params->M - (c_row + tm)));
+  const bool is_unaligned_sm = align_M ? false : (sgp_sm != SM);
+
+  const short sgp_sn = align_N ? SN : min(SN, short(params->N - (c_col + tn)));
+  const bool is_unaligned_sn = align_N ? false : (sgp_sn != SN);
+
+  A += transpose_a ? tm : (tm * params->lda);
+  B += transpose_b ? (tn * params->ldb) : tn;
+  C += tm * params->ldc + tn;
+
+  using DSubTile = NAXSubTile<AccumType, UM, UN>;
+  NAXTile<AccumType, TM, TN, DSubTile> Dtile;
+
+  // gemm_loop through the partition
+  // Check K-alignment at runtime (partition-specific)
+  const int partition_k_size = k_end - k_start;
+  const int partition_k_iters = partition_k_size / BK;
+  const bool partition_k_aligned = (partition_k_size % BK) == 0;
+
+  dispatch_bool(partition_k_aligned, [&](auto kAlignedK) {
+    dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) {
+      dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) {
+        Dtile = gemm_loop<
+            T,
+            SM,
+            SN,
+            SK,
+            BK,
+            transpose_a,
+            transpose_b,
+            kAlignedM.value,
+            kAlignedN.value,
+            kAlignedK.value,
+            UM,
+            UN,
+            UK,
+            AccumType>(
+            A,
+            B,
+            params->lda,
+            params->ldb,
+            partition_k_size,
+            partition_k_iters,
+            sgp_sm,
+            sgp_sn);
+      });
+    });
+  });
+
+  // Store result
+  dispatch_bool(align_M || !is_unaligned_sm, [&](auto kAlignedM) {
+    dispatch_bool(align_N || !is_unaligned_sn, [&](auto kAlignedN) {
+      if constexpr (kAlignedM && kAlignedN) {
+        Dtile.store(C, int(params->ldc));
+      } else {
+        Dtile.store_safe(C, int(params->ldc), short2(sgp_sn, sgp_sm));
+      }
+    });
+  });
+}
@@ -0,0 +1,30 @@
+// Copyright © 2026 Apple Inc.
+
+#include <metal_stdlib>
+
+#include "mlx/backend/metal/kernels/utils.h"
+
+#include "mlx/backend/metal/kernels/steel/gemm/gemm_nax.h"
+#include "mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk_nax.h"
+
+// clang-format off
+#define instantiate_gemm_splitk(tname, trans_a, trans_b, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+  instantiate_kernel(                                                                             \
+      "steel_gemm_splitk_nax_" #tname "_"  #iname "_" #oname                                      \
+      "_bm" #bm "_bn" #bn "_bk" #bk "_wm" #wm "_wn" #wn,                                          \
+  gemm_splitk_nax, itype, bm, bn, bk, wm, wn, trans_a, trans_b, float)
+
+#define instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_splitk(nn, false, false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_splitk(nt, false, true , iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_splitk(tn, true , false, iname, itype, oname, otype, bm, bn, bk, wm, wn) \
+    instantiate_gemm_splitk(tt, true , true , iname, itype, oname, otype, bm, bn, bk, wm, wn)
+
+#define instantiate_gemm_splitk_shapes_helper(iname, itype, oname, otype) \
+    instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype,  64,  64, 256, 2, 2) \
+    instantiate_gemm_splitk_transpose_helper(iname, itype, oname, otype, 128, 128, 512, 4, 4)
+
+instantiate_gemm_splitk_shapes_helper(float16, half, float32, float);
+instantiate_gemm_splitk_shapes_helper(bfloat16, bfloat, float32, float);
+instantiate_gemm_splitk_shapes_helper(float32, float, float32, float);
+// clang-format on
@@ -47,6 +47,7 @@ struct GEMMSpiltKParams {
  const int split_k_partition_stride;
  const int split_k_partition_size;

+  const int swizzle_log;
  const int gemm_k_iterations_aligned;
 };

@@ -604,6 +604,7 @@ void steel_gemm_splitk_axpby(
      /* const int split_k_partitions = */ split_k_partitions,
      /* const int split_k_partition_stride = */ split_k_partition_stride,
      /* const int split_k_partition_size = */ split_k_partition_size,
+      /* const int swizzle_log = */ 0, // no swizzle
      /* const int gemm_k_iterations_aligned = */ gemm_k_iterations};

  MTL::Size group_dims = MTL::Size(32, wn, wm);
@@ -662,6 +663,178 @@ void steel_gemm_splitk_axpby(
  d.add_temporaries(std::move(copies), s.index);
 }

+///////////////////////////////////////////////////////////////////////////////
+// NAX Split k steel matmul
+///////////////////////////////////////////////////////////////////////////////
+
+template <bool CHECK_AB = true>
+void steel_gemm_splitk_axpby_nax(
+    const Stream& s,
+    metal::Device& d,
+    const array& a,
+    const array& b,
+    const array& c,
+    array& out,
+    int M,
+    int N,
+    int K,
+    int batch_size_out,
+    int lda,
+    int ldb,
+    bool transpose_a,
+    bool transpose_b,
+    std::vector<array>& copies,
+    float alpha = 1.0f,
+    float beta = 0.0f) {
+  using namespace mlx::steel;
+
+  constexpr int bm = 128, bn = 128, bk = 512;
+  constexpr int wm = 4, wn = 4;
+
+  // Determine how many partitions to split K into
+  constexpr int split_k_partition_size = 3072;
+  int split_k_partitions =
+      (K + split_k_partition_size - 1) / split_k_partition_size;
+
+  const int bk_iters_per_partition = split_k_partition_size / bk;
+  const int split_k_partition_stride = M * N;
+
+  array C_split({split_k_partitions, M, N}, float32, nullptr, {});
+  C_split.set_data(allocator::malloc(C_split.nbytes()));
+  copies.push_back(C_split);
+
+  const bool align_M = (M % bm) == 0;
+  const bool align_N = (N % bn) == 0;
+  const bool align_K = (K % bk) == 0;
+
+  // Per-tile align_K is checked at runtime; only the last tile can be unaligned
+  metal::MTLFCList func_consts = {
+      {&align_M, MTL::DataType::DataTypeBool, 200},
+      {&align_N, MTL::DataType::DataTypeBool, 201}};
+
+  std::ostringstream kname;
+
+  // clang-format off
+  kname << "steel_gemm_splitk_nax_"
+        << (transpose_a ? 't' : 'n')
+        << (transpose_b ? 't' : 'n')
+        << "_" << type_to_name(a)
+        << "_" << type_to_name(C_split)
+        << "_bm" << bm << "_bn" << bn << "_bk" << bk
+        << "_wm" << wm << "_wn" << wn; // clang-format on
+
+  std::string base_name = kname.str();
+
+  // clang-format off
+  kname << "_align_M_" << (align_M ? 't' : 'n')
+        << "_align_N_" << (align_N ? 't' : 'n')
+        << "_align_K_" << (align_K ? 't' : 'n'); // clang-format on
+
+  std::string hash_name = kname.str();
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  auto kernel = get_steel_gemm_splitk_nax_kernel(
+      /* metal::Device& d = */ d,
+      /* const std::string& kernel_name = */ base_name,
+      /* const std::string& hash_name = */ hash_name,
+      /* const metal::MTLFCList& func_consts = */ func_consts,
+      /* const array& out = */ C_split,
+      /* bool transpose_a = */ transpose_a,
+      /* bool transpose_b = */ transpose_b,
+      /* int bm = */ bm,
+      /* int bn = */ bn,
+      /* int bk = */ bk,
+      /* int wm = */ wm,
+      /* int wn = */ wn);
+
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  int tn = (N + bn - 1) / bn;
+  int tm = (M + bm - 1) / bm;
+
+  int swizzle_log = tm <= 3 ? 0 : 1;
+
+  // Compute swizzled tile counts
+  int tile = 1 << swizzle_log;
+  int tm_swizzled = (tm + tile - 1) / tile;
+  int tn_swizzled = tn * tile;
+
+  GEMMSpiltKParams params{
+      /* const int M = */ M,
+      /* const int N = */ N,
+      /* const int K = */ K,
+      /* const int lda = */ lda,
+      /* const int ldb = */ ldb,
+      /* const int ldc = */ N,
+      /* const int tiles_n = */ tn,
+      /* const int tiles_m = */ tm,
+      /* const int split_k_partitions = */ split_k_partitions,
+      /* const int split_k_partition_stride = */ split_k_partition_stride,
+      /* const int split_k_partition_size = */ split_k_partition_size,
+      /* const int swizzle_log = */ swizzle_log,
+      /* const int gemm_k_iterations_aligned = */ bk_iters_per_partition};
+
+  MTL::Size group_dims = MTL::Size(32, wn, wm);
+  // Use 1D grid with K-partition-major layout: [Partition0: M×N
+  // tiles][Partition1: M×N tiles]... Grid size is 1D to prevent driver/HW from
+  // using its own heuristic to exploit 2D locality by launching threadgroups in
+  // a non-linear order
+  MTL::Size grid_dims =
+      MTL::Size(tn_swizzled * tm_swizzled * split_k_partitions, 1, 1);
+
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_output_array(C_split, 2);
+
+  compute_encoder.set_bytes(params, 3);
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+
+  // Do accum kernel
+  {
+    const bool do_axpby = CHECK_AB && (alpha != 1.0f || beta != 0.0f);
+
+    auto kernel_name = "steel_gemm_splitk_accum_" + type_to_name(out) + "_" +
+        type_to_name(C_split);
+
+    if (do_axpby) {
+      kernel_name = kernel_name + "_axbpy";
+    }
+
+    auto kernel = get_steel_gemm_splitk_accum_kernel(
+        /* metal::Device& d = */ d,
+        /* const std::string& kernel_name = */ kernel_name,
+        /* const array& in = */ C_split,
+        /* const array& out = */ out,
+        /* bool axbpy = */ do_axpby);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    // Set the arguments for the kernel
+    compute_encoder.set_input_array(C_split, 0);
+    compute_encoder.set_output_array(out, 1);
+    compute_encoder.set_bytes(split_k_partitions, 2);
+    compute_encoder.set_bytes(split_k_partition_stride, 3);
+    compute_encoder.set_bytes(N, 4);
+
+    if (do_axpby) {
+      int ldc = c.strides()[c.ndim() - 2];
+      int fdc = c.strides()[c.ndim() - 1];
+
+      compute_encoder.set_input_array(c, 5);
+      compute_encoder.set_bytes(ldc, 6);
+      compute_encoder.set_bytes(fdc, 7);
+      compute_encoder.set_bytes(alpha, 8);
+      compute_encoder.set_bytes(beta, 9);
+    }
+
+    // Launch enough thread groups for each output
+    MTL::Size grid_dims = MTL::Size(N, M, 1);
+    auto group_dims = get_block_dims(N, M, 1);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // Split matmul routing
 ///////////////////////////////////////////////////////////////////////////////
@@ -740,6 +913,7 @@ void steel_matmul_axpby(
  int _tn = N / 16;
  int _tk = K / 16;

+  // Case 1: Small M×N with large K, use SIMD split-K
  if (batch_size_out == 1 && (_tm * _tn) <= 32 && _tk >= 8) {
    return steel_gemm_splitk_axpby<CHECK_AB>(
        /* const Stream& s = */ s,
@@ -761,6 +935,35 @@ void steel_matmul_axpby(
        /* float beta = */ beta);
  }

+  // Case 2: Large K with sufficient M, N, and NAX is available, use NAX split-K
+  // TODO: Add device-specific tuning for more NAX GPUs in the future
+  constexpr int min_mn_threshold = 2048 * 2048;
+  constexpr int min_k_threshold = 10240;
+  if (batch_size_out == 1 && metal::is_nax_available() &&
+      !issubdtype(a.dtype(), complexfloating) &&
+      (env::enable_tf32() || a.dtype() != float32) &&
+      int64_t(M) * N >= min_mn_threshold && K >= min_k_threshold &&
+      K >= (3 * std::max(M, N))) {
+    return steel_gemm_splitk_axpby_nax<CHECK_AB>(
+        /* const Stream& s = */ s,
+        /* metal::Device& d = */ d,
+        /* const array& a = */ a,
+        /* const array& b = */ b,
+        /* const array& c = */ c,
+        /* array& out = */ out,
+        /* int M = */ M,
+        /* int N = */ N,
+        /* int K = */ K,
+        /* int batch_size_out = */ batch_size_out,
+        /* int lda = */ lda,
+        /* int ldb = */ ldb,
+        /* bool transpose_a = */ transpose_a,
+        /* bool transpose_b = */ transpose_b,
+        /* std::vector<array>& copies = */ copies,
+        /* float alpha = */ alpha,
+        /* float beta = */ beta);
+  }
+
  /////////////////////////////////////////////////////////////////////////////
  // Regular kernel dispatch
  auto batch_strides = A_batch_stride;
@@ -339,6 +339,22 @@ MTL::ComputePipelineState* get_steel_gemm_gather_nax_kernel(
  return d.get_kernel(kernel_name, hash_name, func_consts);
 }

+MTL::ComputePipelineState* get_steel_gemm_splitk_nax_kernel(
+    metal::Device& d,
+    const std::string& kernel_name,
+    const std::string& hash_name,
+    const metal::MTLFCList& func_consts,
+    const array&,
+    bool,
+    bool,
+    int,
+    int,
+    int,
+    int,
+    int) {
+  return d.get_kernel(kernel_name, hash_name, func_consts);
+}
+
 MTL::ComputePipelineState* get_qmm_nax_kernel(
    metal::Device& d,
    const std::string& kernel_name,