```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_abt_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;

constexpr int TILE = 64;
constexpr int KTILE = 32;
constexpr int STRIDE = 40;

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

union Workspace {
    struct {
        float ah[TILE * STRIDE];
        float al[TILE * STRIDE];
        float bh[TILE * STRIDE];
        float bl[TILE * STRIDE];
    } inputs;
    float output[TILE * TILE];
};

__global__ void matmul_abt_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ c,
    int64_t m,
    int64_t n,
    int64_t k
) {
#if __CUDA_ARCH__ >= 800
    __shared__ __align__(32) Workspace shared;

    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 int64_t row_start = int64_t(blockIdx.y) * TILE;
    const int64_t col_start = int64_t(blockIdx.x) * TILE;

    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::col_major>;
    using CFragment = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;

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

    for (int64_t base = 0; base < k; base += KTILE) {
        for (int index = tid; index < TILE * KTILE; index += 128) {
            const int row = index / KTILE;
            const int col = index % KTILE;
            const int64_t kk = base + col;

            float av = 0.0f;
            float bv = 0.0f;
            if (row_start + row < m && kk < k)
                av = a[(row_start + row) * k + kk];
            if (col_start + row < n && kk < k)
                bv = b[(col_start + row) * k + kk];

            const float ah = round_tf32(av);
            const float bh = round_tf32(bv);
            const int offset = row * STRIDE + col;

            shared.inputs.ah[offset] = ah;
            shared.inputs.al[offset] = round_tf32(av - ah);
            shared.inputs.bh[offset] = bh;
            shared.inputs.bl[offset] = round_tf32(bv - bh);
        }
        __syncthreads();

        #pragma unroll
        for (int inner = 0; inner < KTILE; inner += 8) {
            AFragment ah[2], al[2];
            BFragment bh[2], bl[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int offset = (warp_row + i * 16) * STRIDE + inner;
                wmma::load_matrix_sync(ah[i], shared.inputs.ah + offset, STRIDE);
                wmma::load_matrix_sync(al[i], shared.inputs.al + offset, STRIDE);
            }

            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const int offset = (warp_col + j * 16) * STRIDE + inner;
                wmma::load_matrix_sync(bh[j], shared.inputs.bh + offset, STRIDE);
                wmma::load_matrix_sync(bl[j], shared.inputs.bl + offset, STRIDE);
            }

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

    #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)
                accum[i][j].x[element] += correction[i][j].x[element];

            wmma::store_matrix_sync(
                shared.output
                    + (warp_row + i * 16) * TILE
                    + warp_col + j * 16,
                accum[i][j],
                TILE,
                wmma::mem_row_major);
        }
    }
    __syncthreads();

    for (int index = tid; index < TILE * TILE; index += 128) {
        const int64_t row = row_start + index / TILE;
        const int64_t col = col_start + index % TILE;
        if (row < m && col < n)
            c[row * n + col] = shared.output[index];
    }
#endif
}

torch::Tensor matmul_abt_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(1), "Inner dimensions must match");
    TORCH_CHECK(
        a.scalar_type() == at::kFloat && b.scalar_type() == at::kFloat,
        "Inputs must have float32 dtype");

    c10::cuda::CUDAGuard guard(a.device());

    if (at::cuda::getDeviceProperties(a.get_device())->major < 8)
        return at::mm(a, b.transpose(0, 1));

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

    const int64_t m = a.size(0);
    const int64_t n = b.size(0);
    const int64_t k = a.size(1);
    auto result = torch::empty({m, n}, a.options());

    if (m == 0 || n == 0)
        return result;

    const dim3 grid((n + TILE - 1) / TILE, (m + TILE - 1) / TILE);
    const auto stream = at::cuda::getCurrentCUDAStream(a.get_device());

    matmul_abt_kernel<<<grid, 128, 0, stream.stream()>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        result.data_ptr<float>(),
        m, n, k);
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return result;
}
"""

_extension = load_inline(
    name="custom_abt_compensated_tf32_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["matmul_abt_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 (
            A.is_cuda
            and B.is_cuda
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and not (
                torch.is_grad_enabled()
                and (A.requires_grad or B.requires_grad)
            )
        ):
            return _extension.matmul_abt_cuda(A, B)
        return torch.matmul(A, B.T)
```