```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 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_runtime.h>
#include <mma.h>

namespace wmma = nvcuda::wmma;

__device__ __forceinline__ float round_tf32(float x) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    float result;
    asm("cvt.rn.tf32.f32 %0, %1;" : "=f"(result) : "f"(x));
    return result;
#else
    return x;
#endif
}

__global__ void matmul_tf32_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ c,
    int m, int n, int k)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    constexpr int BM = 64;
    constexpr int BN = 64;
    constexpr int BK = 32;
    constexpr int AS = 40;
    constexpr int BS = 72;

    __shared__ __align__(32) float storage[2 * BM * AS + 2 * BK * BS];

    float* ah = storage;
    float* al = ah + BM * AS;
    float* bh = al + BM * AS;
    float* bl = bh + BK * BS;

    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 * BM;
    const int col_base = blockIdx.x * BN;

    using AFragment = wmma::fragment<
        wmma::matrix_a, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using BFragment = wmma::fragment<
        wmma::matrix_b, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using CFragment = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;

    CFragment acc[2][2];
    CFragment 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 = 0; kb < k; kb += BK) {
        for (int index = tid; index < BM * BK; index += 128) {
            const int row = index / BK;
            const int col = index % BK;
            float value = 0.0f;
            if (row_base + row < m && kb + col < k) {
                value = a[
                    static_cast<int64_t>(row_base + row) * k + kb + col];
            }
            const float high = round_tf32(value);
            ah[row * AS + col] = high;
            al[row * AS + col] = round_tf32(value - high);
        }

        for (int index = tid; index < BK * BN; index += 128) {
            const int row = index / BN;
            const int col = index % BN;
            float value = 0.0f;
            if (kb + row < k && col_base + col < n) {
                value = b[
                    static_cast<int64_t>(kb + row) * n + col_base + col];
            }
            const float high = round_tf32(value);
            bh[row * BS + col] = high;
            bl[row * BS + col] = round_tf32(value - high);
        }

        __syncthreads();

        #pragma unroll
        for (int inner = 0; inner < BK; inner += 8) {
            AFragment a_high[2], a_low[2];
            BFragment b_high[2], b_low[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                wmma::load_matrix_sync(
                    a_high[i], ah + (warp_row + i * 16) * AS + inner, AS);
                wmma::load_matrix_sync(
                    a_low[i], al + (warp_row + i * 16) * AS + inner, AS);
                wmma::load_matrix_sync(
                    b_high[i], bh + inner * BS + warp_col + i * 16, BS);
                wmma::load_matrix_sync(
                    b_low[i], bl + inner * BS + warp_col + i * 16, BS);
            }

            #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_high[i], b_high[j], acc[i][j]);
                    wmma::mma_sync(
                        correction[i][j],
                        a_high[i], b_low[j], correction[i][j]);
                    wmma::mma_sync(
                        correction[i][j],
                        a_low[i], b_high[j], correction[i][j]);
                }
            }
        }

        __syncthreads();
    }

    // Reuse the input staging memory for coalesced output stores.
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            #pragma unroll
            for (int element = 0; element < CFragment::num_elements; ++element) {
                acc[i][j].x[element] += correction[i][j].x[element];
            }
            wmma::store_matrix_sync(
                storage + (warp_row + i * 16) * BN + warp_col + j * 16,
                acc[i][j], BN, wmma::mem_row_major);
        }
    }

    __syncthreads();

    for (int index = tid; index < BM * BN; index += 128) {
        const int row = row_base + index / BN;
        const int col = col_base + index % BN;
        if (row < m && col < n) {
            c[static_cast<int64_t>(row) * n + col] = storage[index];
        }
    }
#endif
}

torch::Tensor matmul_cuda(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.scalar_type() == torch::kFloat32 &&
                b.scalar_type() == torch::kFloat32, "Expected float32 tensors");
    TORCH_CHECK(a.dim() == 2 && b.dim() == 2, "Expected matrices");
    TORCH_CHECK(a.size(1) == b.size(0), "Incompatible matrix dimensions");

    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 are too large");

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

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

    dim3 grid((n + 63LL) / 64, (m + 63LL) / 64);
    matmul_tf32_kernel<<<grid, 128, 0, at::cuda::getCurrentCUDAStream()>>>(
        a.data_ptr<float>(), b.data_ptr<float>(), output.data_ptr<float>(),
        m, n, k);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_extension = load_inline(
    name="custom_compensated_tf32_matmul",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["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
            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.matmul_cuda(A, B)
```