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


_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>

namespace {

constexpr int TILE = 64;
constexpr int BK = 32;
constexpr int SPLIT_K = 2048;

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

    __shared__ __align__(32) float storage[8192];
    float* ah = storage;
    float* al = storage + 2048;
    float* bh = storage + 4096;
    float* bl = storage + 6144;

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int warp_row = (warp / 2) * 32;
    const int warp_col = (warp % 2) * 32;
    const int row_base = blockIdx.y * TILE;
    const int col_base = blockIdx.x * TILE;
    const int begin = blockIdx.z * SPLIT_K;
    const int end = min(begin + SPLIT_K, K);

    using Acc = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;
    using FragA = wmma::fragment<
        wmma::matrix_a, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using FragB = wmma::fragment<
        wmma::matrix_b, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;

    Acc acc[2][2];
    Acc correction[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);
            wmma::fill_fragment(correction[i][j], 0.0f);
        }
    }

    for (int kb = begin; kb < end; kb += BK) {
        #pragma unroll
        for (int p = tid; p < TILE * BK; p += 128) {
            const int ar = row_base + p / BK;
            const int ak = kb + p % BK;
            const float av = (ar < M && ak < end)
                ? A[static_cast<int64_t>(ar) * K + ak] : 0.0f;
            const float avh = wmma::__float_to_tf32(av);
            ah[p] = avh;
            al[p] = wmma::__float_to_tf32(av - avh);

            const int bk = kb + p / TILE;
            const int bc = col_base + p % TILE;
            const float bv = (bk < end && bc < N)
                ? B[static_cast<int64_t>(bk) * N + bc] : 0.0f;
            const float bvh = wmma::__float_to_tf32(bv);
            bh[p] = bvh;
            bl[p] = wmma::__float_to_tf32(bv - bvh);
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            FragA a_hi[2], a_lo[2];
            FragB b_hi[2], b_lo[2];

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

            #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);
            }

            #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_hi[i], b_hi[j], acc[i][j]);
                    wmma::mma_sync(
                        correction[i][j], a_hi[i], b_lo[j],
                        correction[i][j]);
                    wmma::mma_sync(
                        correction[i][j], a_lo[i], b_hi[j],
                        correction[i][j]);
                }
            }
        }
        __syncthreads();
    }

    // Reuse the input staging memory for the completed output tile.
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            #pragma unroll
            for (int t = 0; t < Acc::num_elements; ++t) {
                acc[i][j].x[t] += correction[i][j].x[t];
            }
            wmma::store_matrix_sync(
                storage + (warp_row + i * 16) * TILE
                        + warp_col + j * 16,
                acc[i][j], TILE, wmma::mem_row_major);
        }
    }
    __syncthreads();

    const int64_t split_offset =
        static_cast<int64_t>(blockIdx.z) * M * N;

    for (int p = tid; p < TILE * TILE; p += 128) {
        const int row = row_base + p / TILE;
        const int col = col_base + p % TILE;
        if (row < M && col < N) {
            partial[split_offset + static_cast<int64_t>(row) * N + col]
                = storage[p];
        }
    }
#endif
}

__global__ void reduce_splits(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int64_t elements,
    int splits)
{
    const int64_t idx =
        static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
    if (idx >= elements) return;

    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[static_cast<int64_t>(s) * elements + idx];
        s1 += partial[static_cast<int64_t>(s + 1) * elements + idx];
        s2 += partial[static_cast<int64_t>(s + 2) * elements + idx];
        s3 += partial[static_cast<int64_t>(s + 3) * elements + idx];
    }
    for (; s < splits; ++s) {
        s0 += partial[static_cast<int64_t>(s) * elements + idx];
    }
    output[idx] = (s0 + s1) + (s2 + s3);
}

} // namespace

torch::Tensor large_k_matmul(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Device mismatch");
    TORCH_CHECK(A.scalar_type() == at::kFloat &&
                B.scalar_type() == at::kFloat, "Expected float32 tensors");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Expected matrices");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible matrix shapes");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Expected contiguous tensors");

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

    TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8,
                "This kernel requires compute capability 8.0 or newer");

    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)
        return torch::empty({M, N}, A.options());
    if (K == 0)
        return torch::zeros({M, N}, A.options());

    const int splits = (K + SPLIT_K - 1) / SPLIT_K;
    TORCH_CHECK(splits <= 65535, "K exceeds supported launch size");

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

    splitk_matmul<<<dim3((N + TILE - 1) / TILE,
                        (M + TILE - 1) / TILE, splits),
                      128, 0, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(),
        partial.data_ptr<float>(), M, N, K);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

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

    return output;
}
"""

_extension = load_inline(
    name="large_k_splitk_tf32_corrected_v1",
    cpp_sources=(
        "torch::Tensor large_k_matmul(torch::Tensor A, torch::Tensor B);"
    ),
    cuda_sources=_cuda_source,
    functions=["large_k_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 A.requires_grad
            or B.requires_grad
            or torch.cuda.get_device_capability(A.device)[0] < 8
        ):
            return torch.matmul(A, B)

        return _extension.large_k_matmul(A.contiguous(), B.contiguous())
```