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

namespace wmma = nvcuda::wmma;

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

// Four warps each compute a 32 x 32 output tile.
// Splitting each input into two TF32 components preserves FP32 accuracy
// much better than a single TF32 multiplication.
__global__ void matmul_abt_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ c,
    int K,
    int N
) {
#if __CUDA_ARCH__ >= 800
    __shared__ __align__(32) float ah[TILE * STRIDE];
    __shared__ __align__(32) float al[TILE * STRIDE];
    __shared__ __align__(32) float bh[TILE * STRIDE];
    __shared__ __align__(32) float bl[TILE * STRIDE];

    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 base_m = blockIdx.y * TILE;
    const int base_n = 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 AccFragment = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;

    AccFragment 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 = 0; kb < K; kb += TILE_K) {
        #pragma unroll
        for (int p = tid; p < TILE * TILE_K; p += 128) {
            const int row = p / TILE_K;
            const int col = p % TILE_K;
            const int s = row * STRIDE + col;

            const float av = a[(base_m + row) * K + kb + col];
            const float bv = b[(base_n + row) * K + kb + col];

            const float avh = wmma::__float_to_tf32(av);
            const float bvh = wmma::__float_to_tf32(bv);

            ah[s] = avh;
            al[s] = wmma::__float_to_tf32(av - avh);
            bh[s] = bvh;
            bl[s] = wmma::__float_to_tf32(bv - bvh);
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < TILE_K; kk += 8) {
            AFragment a_hi[2], a_lo[2];
            BFragment b_hi[2], b_lo[2];

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

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

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

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            float* dst =
                c + (base_m + warp_m + i * 16) * N
                  + base_n + warp_n + j * 16;
            wmma::store_matrix_sync(
                dst, acc[i][j], N, wmma::mem_row_major);
        }
    }
#endif
}

torch::Tensor matmul_abt_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.dim() == 2 && b.dim() == 2, "Expected two matrices");
    TORCH_CHECK(a.size(1) == b.size(1), "Inner dimensions must match");

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

    const int64_t M = a.size(0);
    const int64_t K = a.size(1);
    const int64_t N = b.size(0);

    const auto* properties = at::cuda::getCurrentDeviceProperties();

    if (properties->major < 8 ||
        a.scalar_type() != at::kFloat ||
        b.scalar_type() != at::kFloat ||
        !a.is_contiguous() || !b.is_contiguous() ||
        M == 0 || N == 0 || K == 0 ||
        M % TILE != 0 || N % TILE != 0 || K % TILE_K != 0 ||
        M > INT_MAX || N > INT_MAX || K > INT_MAX ||
        M * K > INT_MAX || N * K > INT_MAX || M * N > INT_MAX) {
        return at::matmul(a, b.transpose(0, 1));
    }

    auto c = torch::empty({M, N}, a.options());
    const dim3 grid(
        static_cast<unsigned int>(N / TILE),
        static_cast<unsigned int>(M / TILE));

    matmul_abt_kernel<<<
        grid, 128, 0, at::cuda::getCurrentCUDAStream()
    >>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        c.data_ptr<float>(),
        static_cast<int>(K),
        static_cast<int>(N)
    );

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return c;
}
"""

_extension = load_inline(
    name="compensated_tf32_matmul_abt",
    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 (
            not A.is_cuda
            or not B.is_cuda
            or (torch.is_grad_enabled() and (A.requires_grad or B.requires_grad))
            or torch.is_autocast_enabled()
        ):
            return torch.matmul(A, B.T)
        return _extension.matmul_abt_cuda(A, B)
```