```python
import functools
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/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <mma.h>
#include <cstdint>

namespace {

constexpr int TILE_M = 64;
constexpr int TILE_N = 64;
constexpr int TILE_K = 32;
constexpr int SPLIT_K = 4096;
constexpr int THREADS = 128;

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

    __shared__ __align__(32) half a_hi[TILE_M * TILE_K];
    __shared__ __align__(32) half a_lo[TILE_M * TILE_K];
    __shared__ __align__(32) half b_hi[TILE_K * TILE_N];
    __shared__ __align__(32) half b_lo[TILE_K * TILE_N];

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int warp_m = (warp / 2) * 32;
    const int warp_n = (warp % 2) * 32;
    const int row_base = blockIdx.y * TILE_M;
    const int col_base = blockIdx.x * TILE_N;
    const int k_begin = blockIdx.z * SPLIT_K;
    const int k_end = min(k_begin + SPLIT_K, K);

    wmma::fragment<wmma::accumulator, 16, 16, 16, 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 = k_begin; kb < k_end; kb += TILE_K) {
        #pragma unroll
        for (int index = tid; index < TILE_M * TILE_K; index += THREADS) {
            const int row = index / TILE_K;
            const int kk = kb + index % TILE_K;
            const float value = kk < k_end
                ? A[static_cast<int64_t>(row_base + row) * K + kk]
                : 0.0f;
            const half high = __float2half_rn(value);
            a_hi[index] = high;
            a_lo[index] = __float2half_rn(value - __half2float(high));
        }

        #pragma unroll
        for (int index = tid; index < TILE_K * TILE_N; index += THREADS) {
            const int kk = kb + index / TILE_N;
            const int col = index % TILE_N;
            const float value = kk < k_end
                ? B[static_cast<int64_t>(kk) * N + col_base + col]
                : 0.0f;
            const half high = __float2half_rn(value);
            b_hi[index] = high;
            b_lo[index] = __float2half_rn(value - __half2float(high));
        }

        __syncthreads();

        #pragma unroll
        for (int offset = 0; offset < TILE_K; offset += 16) {
            wmma::fragment<wmma::matrix_a, 16, 16, 16, half,
                           wmma::row_major> ah[2], al[2];
            wmma::fragment<wmma::matrix_b, 16, 16, 16, half,
                           wmma::row_major> bh[2], bl[2];

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

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

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(acc[i][j], ah[i], bl[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], al[i], bh[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], ah[i], bh[j], acc[i][j]);
                }
            }
        }

        __syncthreads();
    }

    float* destination =
        partial + static_cast<int64_t>(blockIdx.z) * M * N;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            const int row = row_base + warp_m + i * 16;
            const int col = col_base + warp_n + j * 16;
            wmma::store_matrix_sync(
                destination + static_cast<int64_t>(row) * N + col,
                acc[i][j], N, wmma::mem_row_major);
        }
    }
#endif
}

__global__ void reduce_kernel(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int elements,
    int splits)
{
    const int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= elements) {
        return;
    }

    float sum = 0.0f;
    float correction = 0.0f;

    for (int split = 0; split < splits; ++split) {
        const float value =
            partial[static_cast<int64_t>(split) * elements + index];
        const float adjusted = __fsub_rn(value, correction);
        const float next = __fadd_rn(sum, adjusted);
        correction = __fsub_rn(__fsub_rn(next, sum), adjusted);
        sum = next;
    }

    output[index] = sum;
}

} // namespace

torch::Tensor splitk_matmul(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 be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be matrices");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible matrix dimensions");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Inputs must be contiguous");
    TORCH_CHECK(A.size(0) % TILE_M == 0 && B.size(1) % TILE_N == 0,
                "Output dimensions must be multiples of 64");

    const c10::cuda::CUDAGuard guard(A.device());

    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 splits = (K + SPLIT_K - 1) / SPLIT_K;

    auto output = torch::empty({M, N}, A.options());
    auto partial = torch::empty({splits, M, N}, A.options());

    cudaStream_t stream =
        c10::cuda::getCurrentCUDAStream(A.get_device()).stream();

    const dim3 grid(N / TILE_N, M / TILE_M, splits);
    splitk_kernel<<<grid, THREADS, 0, stream>>>(
        A.data_ptr<float>(),
        B.data_ptr<float>(),
        partial.data_ptr<float>(),
        M, N, K);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const int elements = M * N;
    reduce_kernel<<<(elements + 255) / 256, 256, 0, stream>>>(
        partial.data_ptr<float>(),
        output.data_ptr<float>(),
        elements,
        splits);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""


@functools.lru_cache(maxsize=1)
def _extension():
    return load_inline(
        name="large_k_splitk_wmma_v1",
        cpp_sources=_CPP_SOURCE,
        cuda_sources=_CUDA_SOURCE,
        functions=["splitk_matmul"],
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3"],
        verbose=False,
    )


@functools.lru_cache(maxsize=None)
def _supports_wmma(device_index):
    return torch.cuda.get_device_capability(device_index)[0] >= 7


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

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        supported = (
            A.is_cuda
            and B.is_cuda
            and A.device == B.device
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.ndim == 2
            and B.ndim == 2
            and A.shape[1] == B.shape[0]
            and A.shape[0] > 0
            and B.shape[1] > 0
            and A.shape[0] % 64 == 0
            and B.shape[1] % 64 == 0
            and 4096 <= A.shape[1] <= 4096 * 65535
            and A.is_contiguous()
            and B.is_contiguous()
            and not (
                torch.is_grad_enabled() and (A.requires_grad or B.requires_grad)
            )
        )
        if supported and _supports_wmma(A.device.index):
            return _extension().splitk_matmul(A, B)
        return torch.matmul(A, B)
```