```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# Custom split-K TF32 tensor-core GEMM, specialized for the "small M/N, huge K" case.
# C(M,N) = A(M,K) @ B(K,N). The K dimension is decomposed across blocks (split-K),
# each block computes a full 128x128 tile partial product with WMMA TF32 MMAs, and
# partial results are accumulated with global atomic adds.

matmul_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <mma.h>

using namespace nvcuda;

#define BM 128
#define BN 128
#define BK 32
#define LDA (BK + 8)
#define LDB (BN + 8)
#define NTHREADS 256
#define SPLIT_K 128

#define A_STAGE_FLOATS (BM * LDA)
#define B_STAGE_FLOATS (BK * LDB)
#define SMEM_FLOATS (2 * (A_STAGE_FLOATS + B_STAGE_FLOATS))
#define SMEM_BYTES (SMEM_FLOATS * sizeof(float))

__global__ void __launch_bounds__(NTHREADS) splitk_tf32_gemm_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int M, int N, int K, int kChunk)
{
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
    extern __shared__ float smem[];
    float* As = smem;                      // two stages of BM x LDA
    float* Bs = smem + 2 * A_STAGE_FLOATS; // two stages of BK x LDB

    const int ntiles_n = N / BN;
    const int mt = blockIdx.x / ntiles_n;
    const int nt = blockIdx.x % ntiles_n;
    const int m0 = mt * BM;
    const int n0 = nt * BN;
    const int kBegin = blockIdx.y * kChunk;
    const int iters = kChunk / BK;

    const int tid = threadIdx.x;
    const float* Ag = A + (long)m0 * K;
    const float* Bg = B + n0;

    // register buffers used for the software pipeline (load t+1 while computing t)
    float4 ra[4];
    float4 rb[4];

    auto load_global = [&](int k0) {
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            int i = tid + r * NTHREADS;   // 1024 float4 = 128x32 A tile
            int m = i >> 3;
            int kq = i & 7;
            ra[r] = *reinterpret_cast<const float4*>(Ag + (long)m * K + (k0 + kq * 4));
        }
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            int i = tid + r * NTHREADS;   // 1024 float4 = 32x128 B tile
            int k = i >> 5;
            int nq = i & 31;
            rb[r] = *reinterpret_cast<const float4*>(Bg + (long)(k0 + k) * N + nq * 4);
        }
    };

    auto store_smem = [&](int stage) {
        float* AsS = As + stage * A_STAGE_FLOATS;
        float* BsS = Bs + stage * B_STAGE_FLOATS;
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            int i = tid + r * NTHREADS;
            *reinterpret_cast<float4*>(AsS + (i >> 3) * LDA + ((i & 7) << 2)) = ra[r];
        }
        #pragma unroll
        for (int r = 0; r < 4; ++r) {
            int i = tid + r * NTHREADS;
            *reinterpret_cast<float4*>(BsS + (i >> 5) * LDB + ((i & 31) << 2)) = rb[r];
        }
    };

    // 8 warps arranged 2(m) x 4(n), each warp computes a 64x32 tile via 16x16x8 TF32 MMAs
    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[4][2];
    #pragma unroll
    for (int i = 0; i < 4; ++i)
        #pragma unroll
        for (int j = 0; j < 2; ++j)
            wmma::fill_fragment(acc[i][j], 0.0f);

    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int wm = (warp >> 2) << 6;   // (warp/4)*64
    const int wn = (warp & 3) << 5;    // (warp%4)*32

    load_global(kBegin);
    store_smem(0);
    __syncthreads();

    for (int t = 0; t < iters; ++t) {
        const int stage = t & 1;
        if (t + 1 < iters) load_global(kBegin + (t + 1) * BK);

        const float* AsS = As + stage * A_STAGE_FLOATS;
        const float* BsS = Bs + stage * B_STAGE_FLOATS;

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<wmma::matrix_a, 16, 16, 8, wmma::precision::tf32, wmma::row_major> af[4];
            wmma::fragment<wmma::matrix_b, 16, 16, 8, wmma::precision::tf32, wmma::row_major> bf[2];
            #pragma unroll
            for (int i = 0; i < 4; ++i) {
                wmma::load_matrix_sync(af[i], AsS + (wm + i * 16) * LDA + kk, LDA);
                #pragma unroll
                for (int e = 0; e < af[i].num_elements; ++e)
                    af[i].x[e] = wmma::__float_to_tf32(af[i].x[e]);
            }
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::load_matrix_sync(bf[j], BsS + kk * LDB + wn + j * 16, LDB);
                #pragma unroll
                for (int e = 0; e < bf[j].num_elements; ++e)
                    bf[j].x[e] = wmma::__float_to_tf32(bf[j].x[e]);
            }
            #pragma unroll
            for (int i = 0; i < 4; ++i)
                #pragma unroll
                for (int j = 0; j < 2; ++j)
                    wmma::mma_sync(acc[i][j], af[i], bf[j], acc[i][j]);
        }

        if (t + 1 < iters) store_smem(stage ^ 1);
        __syncthreads();
    }

    // epilogue: accumulate this K-slice partial tile into C with atomics
    const int gr = lane >> 2;        // 0..7
    const int gc = (lane & 3) * 2;   // 0,2,4,6
    #pragma unroll
    for (int i = 0; i < 4; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            const int r = m0 + wm + i * 16 + gr;
            const int c = n0 + wn + j * 16 + gc;
            float* base = C + (long)r * N + c;
            atomicAdd(base,             acc[i][j].x[0]);
            atomicAdd(base + 1,         acc[i][j].x[1]);
            atomicAdd(base + 8 * N,     acc[i][j].x[2]);
            atomicAdd(base + 8 * N + 1, acc[i][j].x[3]);
            atomicAdd(base + 8,         acc[i][j].x[4]);
            atomicAdd(base + 9,         acc[i][j].x[5]);
            atomicAdd(base + 8 * N + 8, acc[i][j].x[6]);
            atomicAdd(base + 8 * N + 9, acc[i][j].x[7]);
        }
    }
#endif
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be 2D");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    const int M = (int)Ac.size(0);
    const int K = (int)Ac.size(1);
    const int N = (int)Bc.size(1);
    TORCH_CHECK(Bc.size(0) == K, "Inner dimensions must match");

    auto prop = at::cuda::getCurrentDeviceProperties();
    const bool tf32_ok = (prop->major >= 8);

    // Fallback for unsupported shapes / devices / dtypes
    if (!tf32_ok || Ac.scalar_type() != at::kFloat || Bc.scalar_type() != at::kFloat ||
        (M % BM) != 0 || (N % BN) != 0 || (K % (SPLIT_K * BK)) != 0) {
        return at::matmul(Ac, Bc);
    }

    auto C = torch::zeros({M, N}, Ac.options());
    const int kChunk = K / SPLIT_K;

    static bool attr_done = false;
    if (!attr_done) {
        cudaFuncSetAttribute(splitk_tf32_gemm_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             SMEM_BYTES);
        attr_done = true;
    }

    dim3 grid((M / BM) * (N / BN), SPLIT_K, 1);
    dim3 block(NTHREADS, 1, 1);
    auto stream = at::cuda::getCurrentCUDAStream();
    splitk_tf32_gemm_kernel<<<grid, block, SMEM_BYTES, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        M, N, K, kChunk);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

matmul_cpp_src = "torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B);"

_ext = load_inline(
    name="splitk_tf32_gemm",
    cpp_sources=[matmul_cpp_src],
    cuda_sources=[matmul_cuda_src],
    functions=["matmul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model: replaces torch.matmul with a custom split-K TF32 tensor-core
    CUDA kernel tuned for small M/N and very large K.
    """
    def __init__(self):
        super(ModelNew, self).__init__()

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if A.is_cuda and B.is_cuda and A.dtype == torch.float32 and B.dtype == torch.float32:
            return _ext.matmul_cuda(A, B)
        return torch.matmul(A, B)
```