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

_cpp_source = r"""
#include <torch/extension.h>
torch::Tensor splitk_matmul(torch::Tensor a, torch::Tensor b);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <mma.h>
#include <algorithm>
#include <climits>

namespace {

constexpr int TILE = 64;
constexpr int KTILE = 32;

__global__ void splitk_tf32_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ partial,
    int m, int n, int k, int chunk,
    int tiles_m, int tiles_n)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    namespace wmma = nvcuda::wmma;

    __shared__ __align__(32) float ah[TILE * KTILE];
    __shared__ __align__(32) float al[TILE * KTILE];
    __shared__ __align__(32) float bh[KTILE * TILE];
    __shared__ __align__(32) float bl[KTILE * TILE];

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int warp_row = (warp >> 1) * 32;
    const int warp_col = (warp & 1) * 32;
    const int row_base = blockIdx.y * TILE;
    const int col_base = blockIdx.x * TILE;
    const int begin = blockIdx.z * chunk;
    const int end = min(begin + chunk, k);

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> acc[2][2];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::fill_fragment(acc[i][j], 0.0f);
        }
    }

    for (int kb = begin; kb < end; kb += KTILE) {
        #pragma unroll
        for (int p = tid; p < TILE * KTILE; p += 128) {
            const int r = p / KTILE;
            const int c = p % KTILE;
            float x = 0.0f;
            if (row_base + r < m && kb + c < end) {
                x = a[static_cast<int64_t>(row_base + r) * k + kb + c];
            }
            const float h = wmma::__float_to_tf32(x);
            ah[p] = h;
            al[p] = wmma::__float_to_tf32(x - h);
        }

        #pragma unroll
        for (int p = tid; p < KTILE * TILE; p += 128) {
            const int r = p / TILE;
            const int c = p % TILE;
            float x = 0.0f;
            if (kb + r < end && col_base + c < n) {
                x = b[static_cast<int64_t>(kb + r) * n + col_base + c];
            }
            const float h = wmma::__float_to_tf32(x);
            bh[p] = h;
            bl[p] = wmma::__float_to_tf32(x - h);
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < KTILE; kk += 8) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major> a_hi[2], a_lo[2];
            wmma::fragment<
                wmma::matrix_b, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major> b_hi[2], b_lo[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int offset = (warp_row + i * 16) * KTILE + kk;
                wmma::load_matrix_sync(a_hi[i], ah + offset, KTILE);
                wmma::load_matrix_sync(a_lo[i], al + offset, KTILE);
            }

            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const int offset = kk * TILE + warp_col + j * 16;
                wmma::load_matrix_sync(b_hi[j], bh + offset, TILE);
                wmma::load_matrix_sync(b_lo[j], bl + offset, TILE);
            }

            // Retain the two first-order TF32 rounding corrections.
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(acc[i][j], a_lo[i], b_hi[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], a_hi[i], b_lo[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], a_hi[i], b_hi[j], acc[i][j]);
                }
            }
        }
        __syncthreads();
    }

    const int64_t tile_index =
        (static_cast<int64_t>(blockIdx.z) * tiles_m + blockIdx.y)
        * tiles_n + blockIdx.x;
    float* dst = partial + tile_index * TILE * TILE;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::store_matrix_sync(
                dst + (warp_row + i * 16) * TILE + warp_col + j * 16,
                acc[i][j], TILE, wmma::mem_row_major);
        }
    }
#endif
}

__global__ void reduce_splitk_kernel(
    const float* __restrict__ partial,
    float* __restrict__ out,
    int m, int n, int splits, int tiles_m, int tiles_n)
{
    const int64_t index =
        static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    if (index >= static_cast<int64_t>(m) * n) {
        return;
    }

    const int row = index / n;
    const int col = index % n;
    const int64_t tile = static_cast<int64_t>(row / TILE) * tiles_n + col / TILE;
    const int64_t offset =
        tile * TILE * TILE + (row % TILE) * TILE + col % TILE;
    const int64_t stride =
        static_cast<int64_t>(tiles_m) * tiles_n * TILE * TILE;

    float s0 = 0.0f;
    float s1 = 0.0f;
    float s2 = 0.0f;
    float s3 = 0.0f;
    int s = 0;

    for (; s + 3 < splits; s += 4) {
        s0 += partial[offset + static_cast<int64_t>(s) * stride];
        s1 += partial[offset + static_cast<int64_t>(s + 1) * stride];
        s2 += partial[offset + static_cast<int64_t>(s + 2) * stride];
        s3 += partial[offset + static_cast<int64_t>(s + 3) * stride];
    }
    float sum = (s0 + s1) + (s2 + s3);
    for (; s < splits; ++s) {
        sum += partial[offset + static_cast<int64_t>(s) * stride];
    }
    out[index] = sum;
}

}  // namespace

torch::Tensor splitk_matmul(torch::Tensor a, torch::Tensor b) {
    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(a.device() == b.device(), "Inputs must share a device");
    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "Expected matrices");
    TORCH_CHECK(a.size(1) == b.size(0), "Incompatible matrix dimensions");
    TORCH_CHECK(
        a.scalar_type() == torch::kFloat32 &&
        b.scalar_type() == torch::kFloat32,
        "Expected float32 tensors");
    TORCH_CHECK(a.is_contiguous() && b.is_contiguous(), "Expected contiguous inputs");

    c10::cuda::CUDAGuard device_guard(a.device());

    if (at::cuda::getCurrentDeviceProperties()->major < 8) {
        return at::matmul(a, b);
    }

    TORCH_CHECK(
        a.size(0) <= INT_MAX && a.size(1) <= INT_MAX && b.size(1) <= INT_MAX,
        "Matrix dimensions exceed kernel limits");

    const int m = static_cast<int>(a.size(0));
    const int k = static_cast<int>(a.size(1));
    const int n = static_cast<int>(b.size(1));

    if (m == 0 || n == 0 || k == 0) {
        return torch::zeros({m, n}, a.options());
    }

    const int tiles_m = (m + TILE - 1) / TILE;
    const int tiles_n = (n + TILE - 1) / TILE;
    const int requested_splits = static_cast<int>(
        std::min<int64_t>(128, (static_cast<int64_t>(k) + 4095) / 4096));
    const int chunk = static_cast<int>(
        ((static_cast<int64_t>(k) + requested_splits - 1) /
         requested_splits + KTILE - 1) / KTILE * KTILE);
    const int splits = static_cast<int>(
        (static_cast<int64_t>(k) + chunk - 1) / chunk);

    auto partial = torch::empty(
        {splits, tiles_m, tiles_n, TILE, TILE}, a.options());
    auto out = torch::empty({m, n}, a.options());

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    splitk_tf32_kernel<<<dim3(tiles_n, tiles_m, splits), 128, 0, stream>>>(
        a.data_ptr<float>(), b.data_ptr<float>(), partial.data_ptr<float>(),
        m, n, k, chunk, tiles_m, tiles_n);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const int64_t elements = static_cast<int64_t>(m) * n;
    reduce_splitk_kernel<<<(elements + 255) / 256, 256, 0, stream>>>(
        partial.data_ptr<float>(), out.data_ptr<float>(),
        m, n, splits, tiles_m, tiles_n);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return out;
}
"""

_extension = load_inline(
    name="large_k_splitk_tf32_corrected",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["splitk_matmul"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (
            not A.is_cuda
            or not B.is_cuda
            or A.dtype != torch.float32
            or B.dtype != torch.float32
            or A.ndim != 2
            or B.ndim != 2
            or (torch.is_grad_enabled() and (A.requires_grad or B.requires_grad))
        ):
            return torch.matmul(A, B)
        return _extension.splitk_matmul(A.contiguous(), B.contiguous())
```