```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_cuda(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 <climits>

namespace {

constexpr int TILE = 64;
constexpr int BK = 32;
constexpr int SPLIT_K = 4096;
constexpr int THREADS = 256;

__global__ void splitk_gemm(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partial,
    int M, int N, int K,
    int padded_m, int padded_n)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    using namespace nvcuda;

    __shared__ __align__(32) float As[TILE * BK];
    __shared__ __align__(32) float Bs[BK * TILE];

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int warp_row = (warp >> 1) * 16;
    const int warp_col = (warp & 1) * 32;

    const int row_base = blockIdx.y * TILE;
    const int col_base = blockIdx.x * TILE;
    const int k_begin = blockIdx.z * SPLIT_K;
    const int k_end = min(k_begin + SPLIT_K, K);

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

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

    for (int kb = k_begin; kb < k_end; kb += BK) {
        #pragma unroll
        for (int i = tid; i < TILE * BK; i += THREADS) {
            const int r = i / BK;
            const int k = i % BK;
            As[i] = (row_base + r < M && kb + k < k_end)
                ? A[static_cast<int64_t>(row_base + r) * K + kb + k]
                : 0.0f;
        }

        #pragma unroll
        for (int i = tid; i < BK * TILE; i += THREADS) {
            const int k = i / TILE;
            const int c = i % TILE;
            Bs[i] = (kb + k < k_end && col_base + c < N)
                ? B[static_cast<int64_t>(kb + k) * N + col_base + c]
                : 0.0f;
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major> ah, al;

            wmma::load_matrix_sync(ah, As + warp_row * BK + kk, BK);

            #pragma unroll
            for (int e = 0; e < ah.num_elements; ++e) {
                const float value = ah.x[e];
                const float high = wmma::__float_to_tf32(value);
                ah.x[e] = high;
                al.x[e] = wmma::__float_to_tf32(value - high);
            }

            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::fragment<
                    wmma::matrix_b, 16, 16, 8,
                    wmma::precision::tf32, wmma::row_major> bh, bl;

                wmma::load_matrix_sync(
                    bh, Bs + kk * TILE + warp_col + j * 16, TILE);

                #pragma unroll
                for (int e = 0; e < bh.num_elements; ++e) {
                    const float value = bh.x[e];
                    const float high = wmma::__float_to_tf32(value);
                    bh.x[e] = high;
                    bl.x[e] = wmma::__float_to_tf32(value - high);
                }

                wmma::mma_sync(acc[j], ah, bh, acc[j]);
                wmma::mma_sync(correction[j], ah, bl, correction[j]);
                wmma::mma_sync(correction[j], al, bh, correction[j]);
            }
        }
        __syncthreads();
    }

    const int64_t offset =
        static_cast<int64_t>(blockIdx.z) * padded_m * padded_n +
        static_cast<int64_t>(row_base + warp_row) * padded_n +
        col_base + warp_col;

    #pragma unroll
    for (int j = 0; j < 2; ++j) {
        #pragma unroll
        for (int e = 0; e < acc[j].num_elements; ++e) {
            acc[j].x[e] += correction[j].x[e];
        }
        wmma::store_matrix_sync(
            partial + offset + j * 16,
            acc[j], padded_n, wmma::mem_row_major);
    }
#endif
}

__global__ void reduce_splits(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int M, int N, int padded_m, int padded_n, int splits)
{
    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 stride = static_cast<int64_t>(padded_m) * padded_n;
    const int64_t offset = static_cast<int64_t>(row) * padded_n + col;

    float sums[8] = {0.0f, 0.0f, 0.0f, 0.0f,
                     0.0f, 0.0f, 0.0f, 0.0f};

    for (int s = 0; s < splits; s += 8) {
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            if (s + j < splits) {
                sums[j] += partial[static_cast<int64_t>(s + j) * stride + offset];
            }
        }
    }

    output[index] =
        ((sums[0] + sums[1]) + (sums[2] + sums[3])) +
        ((sums[4] + sums[5]) + (sums[6] + sums[7]));
}

} // namespace

torch::Tensor splitk_matmul_cuda(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Inputs must share a device");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be matrices");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible matrix dimensions");

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

    if (A.scalar_type() != at::kFloat ||
        B.scalar_type() != at::kFloat ||
        at::cuda::getCurrentDeviceProperties()->major < 8 ||
        A.size(0) == 0 || A.size(1) == 0 || B.size(1) == 0 ||
        A.size(0) > INT_MAX - 64 ||
        B.size(1) > INT_MAX - 64 ||
        A.size(1) > static_cast<int64_t>(SPLIT_K) * 65535) {
        return at::matmul(A, B);
    }

    A = A.contiguous();
    B = B.contiguous();

    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));
    const int padded_m = ((M + TILE - 1) / TILE) * TILE;
    const int padded_n = ((N + TILE - 1) / TILE) * TILE;
    const int splits = (K + SPLIT_K - 1) / SPLIT_K;

    if (padded_m / TILE > 65535) {
        return at::matmul(A, B);
    }

    auto output = torch::empty({M, N}, A.options());
    auto partial = torch::empty({splits, padded_m, padded_n}, A.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const dim3 grid(padded_n / TILE, padded_m / TILE, splits);
    splitk_gemm<<<grid, THREADS, 0, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), partial.data_ptr<float>(),
        M, N, K, padded_m, padded_n);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const int64_t elements = static_cast<int64_t>(M) * N;
    reduce_splits<<<(elements + THREADS - 1) / THREADS, THREADS, 0, stream>>>(
        partial.data_ptr<float>(), output.data_ptr<float>(),
        M, N, padded_m, padded_n, splits);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_extension = load_inline(
    name="large_k_splitk_tf32_corrected",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["splitk_matmul_cuda"],
    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:
            return torch.matmul(A, B)
        return _extension.splitk_matmul_cuda(A, B)
```