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

_cpp_source = """
#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 <cuda_fp16.h>
#include <mma.h>
#include <algorithm>

namespace {

constexpr int TM = 32;
constexpr int TN = 64;
constexpr int TK = 32;

__global__ void splitk_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ partial,
    int m, int n, int k, int chunk)
{
#if __CUDA_ARCH__ >= 700
    using namespace nvcuda;

    __shared__ __align__(32) half ah[TM * TK];
    __shared__ __align__(32) half al[TM * TK];
    __shared__ __align__(32) half bh[TK * TN];
    __shared__ __align__(32) half bl[TK * TN];
    __shared__ __align__(32) float result[TM * TN];

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int wr = (warp / 4) * 16;
    const int wc = (warp % 4) * 16;
    const int row0 = blockIdx.y * TM;
    const int col0 = blockIdx.x * TN;
    const int begin = blockIdx.z * chunk;
    const int end = min(begin + chunk, k);

    wmma::fragment<wmma::accumulator, 16, 16, 16, float> main_sum;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> correction;
    wmma::fill_fragment(main_sum, 0.0f);
    wmma::fill_fragment(correction, 0.0f);

    for (int base = begin; base < end; base += TK) {
        for (int i = tid; i < TM * TK; i += blockDim.x) {
            const int row = row0 + i / TK;
            const int kk = base + i % TK;
            const float v = (row < m && kk < end)
                ? a[static_cast<int64_t>(row) * k + kk] : 0.0f;
            const half hi = __float2half_rn(v);
            ah[i] = hi;
            al[i] = __float2half_rn(v - __half2float(hi));
        }

        for (int i = tid; i < TK * TN; i += blockDim.x) {
            const int kk = base + i / TN;
            const int col = col0 + i % TN;
            const float v = (kk < end && col < n)
                ? b[static_cast<int64_t>(kk) * n + col] : 0.0f;
            const half hi = __float2half_rn(v);
            bh[i] = hi;
            bl[i] = __float2half_rn(v - __half2float(hi));
        }
        __syncthreads();

        #pragma unroll
        for (int offset = 0; offset < TK; offset += 16) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 16, half, wmma::row_major
            > a_hi, a_lo;
            wmma::fragment<
                wmma::matrix_b, 16, 16, 16, half, wmma::row_major
            > b_hi, b_lo;

            wmma::load_matrix_sync(a_hi, ah + wr * TK + offset, TK);
            wmma::load_matrix_sync(a_lo, al + wr * TK + offset, TK);
            wmma::load_matrix_sync(b_hi, bh + offset * TN + wc, TN);
            wmma::load_matrix_sync(b_lo, bl + offset * TN + wc, TN);

            wmma::mma_sync(main_sum, a_hi, b_hi, main_sum);
            wmma::mma_sync(correction, a_hi, b_lo, correction);
            wmma::mma_sync(correction, a_lo, b_hi, correction);
        }
        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < main_sum.num_elements; ++i) {
        main_sum.x[i] += correction.x[i];
    }

    wmma::store_matrix_sync(
        result + wr * TN + wc, main_sum, TN, wmma::mem_row_major);
    __syncthreads();

    const int64_t partition_offset =
        static_cast<int64_t>(blockIdx.z) * m * n;

    for (int i = tid; i < TM * TN; i += blockDim.x) {
        const int row = row0 + i / TN;
        const int col = col0 + i % TN;
        if (row < m && col < n) {
            partial[partition_offset + static_cast<int64_t>(row) * n + col]
                = result[i];
        }
    }
#endif
}

__global__ void reduce_kernel(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int64_t count,
    int partitions)
{
    const int64_t i =
        static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    if (i >= count) return;

    float sum = 0.0f;
    for (int p = 0; p < partitions; ++p) {
        sum += partial[static_cast<int64_t>(p) * count + i];
    }
    output[i] = sum;
}

} // 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.scalar_type() == at::kFloat && b.scalar_type() == at::kFloat,
        "Inputs must have float32 dtype");
    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 guard(a.device());
    TORCH_CHECK(
        at::cuda::getCurrentDeviceProperties()->major >= 7,
        "This kernel requires compute capability 7.0 or newer");

    a = a.contiguous();
    b = b.contiguous();

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

    auto output = torch::empty({m, n}, a.options());
    if (m == 0 || n == 0) return output;
    if (k == 0) {
        output.zero_();
        return output;
    }

    const int requested_partitions = std::min(128, (k + 4095) / 4096);
    const int chunk =
        ((static_cast<int64_t>(k) + requested_partitions * TK - 1)
         / (requested_partitions * TK)) * TK;
    const int partitions = (k + chunk - 1) / chunk;

    auto partial = torch::empty({partitions, m, n}, a.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    const dim3 grid(
        (n + TN - 1) / TN,
        (m + TM - 1) / TM,
        partitions);

    splitk_kernel<<<grid, 256, 0, stream>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        partial.data_ptr<float>(),
        m, n, k, chunk);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const int64_t count = static_cast<int64_t>(m) * n;
    reduce_kernel<<<(count + 255) / 256, 256, 0, stream>>>(
        partial.data_ptr<float>(),
        output.data_ptr<float>(),
        count,
        partitions);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_splitk_extension = load_inline(
    name="large_k_splitk_compensated_matmul",
    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:
        return _splitk_extension.splitk_matmul_cuda(A, B)
```