[BUG] NAX split-K GEMM (steel_gemm_splitk_axpby_nax) returns garbage for bf16 on M5 Max at the M·N ≥ 2048² dispatch boundary

[xocialize]

Describe the bug

On an M5 Max, half-precision matmuls that hit the NAX split-K GEMM path introduced via #3017 / #3018 (steel_gemm_splitk_axpby_nax) return garbage (cosine ≈ 0 vs an fp32 reference, magnitudes up to 1e33, sometimes NaN). The failure switches on bit-exactly at the dispatch boundary in mlx/backend/metal/matmul.cpp:

constexpr int min_mn_threshold = 2048 * 2048;
constexpr int min_k_threshold = 10240;
if (batch_size_out == 1 && metal::is_nax_available() && ... &&
    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>(...);

For [M, 12288] @ [12288, 4096]ᵀ in bf16 (N=4096): M=896 → 896·4096 = 3.67M < 2048² → regular kernel, correct; M=1024 → 1024·4096 = 2048² exactly → NAX split-K, garbage.

To Reproduce

Self-contained SPM package (also reproduces via mlx-swift 0.31.3 and 0.31.4; the kernel isn't in any PyPI Python release yet — 0.31.2 on the same machine is correct on identical inputs):

// Package.swift
// swift-tools-version: 5.9
import PackageDescription
let package = Package(
    name: "nax-splitk-repro",
    platforms: [.macOS(.v14)],
    dependencies: [.package(url: "https://github.com/ml-explore/mlx-swift.git", exact: "0.31.6")],
    targets: [.executableTarget(name: "repro", dependencies: [
        .product(name: "MLX", package: "mlx-swift"),
        .product(name: "MLXRandom", package: "mlx-swift"),
    ])]
)

// Sources/repro/main.swift
import MLX
import MLXRandom

MLXRandom.seed(0)
let (K, N) = (12288, 4096)
let b = MLXRandom.normal([N, K]).asType(.bfloat16)
for m in [512, 896, 1024, 2048] {
    let a = MLXRandom.normal([m, K]).asType(.bfloat16)
    let y = matmul(a, b.T)                                       // bf16 path under test
    let yRef = matmul(a.asType(.float32), b.asType(.float32).T)  // fp32 reference
    eval(y, yRef)
    let yf = y.asType(.float32)
    let cos = (yf * yRef).sum() / (sqrt((yf * yf).sum()) * sqrt((yRef * yRef).sum()))
    print("M=\(m) K=\(K) N=\(N) bf16-vs-fp32: cos \(cos.item(Float.self)) max_abs \(abs(yf - yRef).max().item(Float.self))")
}

Output (swift run -c release):

M=512  K=12288 N=4096 bf16-vs-fp32: cos 0.99999857     max_abs 1.96        ← normal bf16 accumulation
M=896  K=12288 N=4096 bf16-vs-fp32: cos 0.99999875     max_abs 1.62
M=1024 K=12288 N=4096 bf16-vs-fp32: cos -0.0           max_abs 4.2e+33     ← M·N crosses 2048² exactly
M=2048 K=12288 N=4096 bf16-vs-fp32: cos -0.00064606126 max_abs 850.06

With real model tensors (Qwen3-VL-8B down_proj weights + activations) a finer sweep shows the same clean boundary: rows ≤ 896 correct, rows ≥ 1024 garbage/NaN; batched [1, M, K] and 2-D [M, K] both affected; fp32 (excluded from the dispatch) correct at all sizes. Deterministic across runs and reboots.

Expected behavior

bf16 result within normal accumulation error of the fp32 reference (as at M ≤ 896), or fallback to the regular kernel.

Desktop

• OS Version: macOS 27.0 (beta, build 26A5368g) — beta-OS caveat: not yet verified on release macOS; the Metal compiler/driver interplay may matter. Happy to run any diagnostics or test patches.
• Hardware: Apple M5 Max
• Version: mlx-swift 0.31.3 / 0.31.4 / 0.31.6 (mlx core as vendored in those tags); Python mlx 0.31.2 (pre-kernel) is clean on the same machine

Additional context

Real-world impact: this silently corrupts prefill for ordinary LLM inference — any model whose FFN down-projection has K ≥ 10240 (e.g. Qwen3 8B K=12288, Gemma-3-12B K=15360) produces garbage hidden states once the prompt reaches ~2048²/N tokens (≈1024 for N=4096). We root-caused it from a Qwen3-VL-8B vision-conditioning pipeline that degraded exactly at 1024-token image grids while shorter prompts were fine. Workaround we're shipping: chunk the GEMM along M to stay under min_mn_threshold (row-chunks of ≤896), which is exact and avoids the kernel.