```python
import torch
import torch.nn as nn
from functools import lru_cache
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_runtime.h>
#include <mma.h>

namespace wmma = nvcuda::wmma;

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

__global__ void matmul_abt_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 = 128;
    constexpr int BN = 64;
    constexpr int BK = 32;

    __shared__ __align__(32) float ah[BM * BK];
    __shared__ __align__(32) float al[BM * BK];
    __shared__ __align__(32) float bh[BN * BK];
    __shared__ __align__(32) float bl[BN * BK];

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int block_m = blockIdx.y * BM;
    const int block_n = blockIdx.x * BN;
    const int warp_m = (warp / 2) * 32;
    const int warp_n = (warp % 2) * 32;

    using AFrag = wmma::fragment<
        wmma::matrix_a, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using BFrag = wmma::fragment<
        wmma::matrix_b, 16, 16, 8,
        wmma::precision::tf32, wmma::col_major>;
    using AccFrag = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;

    AccFrag sums[2][2];
    AccFrag corrections[2][2];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::fill_fragment(sums[i][j], 0.0f);
            wmma::fill_fragment(corrections[i][j], 0.0f);
        }
    }

    for (int base_k = 0; base_k < K; base_k += BK) {
        // Vectorized loads; each thread converts four adjacent values.
        for (int v = tid; v < BM * BK / 4; v += 256) {
            const int p = v * 4;
            const int row = p / BK;
            const int col = p % BK;
            const float4 x = *reinterpret_cast<const float4*>(
                A + static_cast<int64_t>(block_m + row) * K
                  + base_k + col);

            float4 h, l;
            h.x = round_tf32(x.x);
            h.y = round_tf32(x.y);
            h.z = round_tf32(x.z);
            h.w = round_tf32(x.w);
            l.x = round_tf32(x.x - h.x);
            l.y = round_tf32(x.y - h.y);
            l.z = round_tf32(x.z - h.z);
            l.w = round_tf32(x.w - h.w);

            *reinterpret_cast<float4*>(ah + p) = h;
            *reinterpret_cast<float4*>(al + p) = l;
        }

        for (int v = tid; v < BN * BK / 4; v += 256) {
            const int p = v * 4;
            const int row = p / BK;
            const int col = p % BK;
            const float4 x = *reinterpret_cast<const float4*>(
                B + static_cast<int64_t>(block_n + row) * K
                  + base_k + col);

            float4 h, l;
            h.x = round_tf32(x.x);
            h.y = round_tf32(x.y);
            h.z = round_tf32(x.z);
            h.w = round_tf32(x.w);
            l.x = round_tf32(x.x - h.x);
            l.y = round_tf32(x.y - h.y);
            l.z = round_tf32(x.z - h.z);
            l.w = round_tf32(x.w - h.w);

            *reinterpret_cast<float4*>(bh + p) = h;
            *reinterpret_cast<float4*>(bl + p) = l;
        }

        __syncthreads();

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

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int offset = (warp_m + 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 = (warp_n + j * 16) * BK + kk;
                wmma::load_matrix_sync(b_hi[j], bh + offset, BK);
                wmma::load_matrix_sync(b_lo[j], bl + offset, BK);
            }

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(
                        sums[i][j], a_hi[i], b_hi[j], sums[i][j]);
                    wmma::mma_sync(
                        corrections[i][j], a_hi[i], b_lo[j],
                        corrections[i][j]);
                    wmma::mma_sync(
                        corrections[i][j], a_lo[i], b_hi[j],
                        corrections[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 t = 0; t < AccFrag::num_elements; ++t) {
                sums[i][j].x[t] += corrections[i][j].x[t];
            }

            float* destination =
                C + static_cast<int64_t>(block_m + warp_m + i * 16) * N
                  + block_n + warp_n + j * 16;
            wmma::store_matrix_sync(
                destination, sums[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(), "CUDA tensors required");
    TORCH_CHECK(A.device() == B.device(), "Device mismatch");
    TORCH_CHECK(A.scalar_type() == at::kFloat &&
                B.scalar_type() == at::kFloat, "float32 tensors required");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Matrices required");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Contiguous tensors required");
    TORCH_CHECK(A.size(1) == B.size(1), "Inner dimension mismatch");

    const int M = static_cast<int>(A.size(0));
    const int N = static_cast<int>(B.size(0));
    const int K = static_cast<int>(A.size(1));

    TORCH_CHECK(M % 128 == 0 && N % 64 == 0 && K % 32 == 0,
                "Unsupported tile dimensions");

    c10::cuda::CUDAGuard guard(A.device());
    TORCH_CHECK(at::cuda::getCurrentDeviceProperties()->major >= 8,
                "Ampere or newer GPU required");

    auto C = torch::empty({M, N}, A.options());
    if (M == 0 || N == 0) {
        return C;
    }

    const dim3 grid(N / 64, M / 128);
    matmul_abt_kernel<<<
        grid, 256, 0, at::cuda::getCurrentCUDAStream()
    >>>(
        A.data_ptr<float>(),
        B.data_ptr<float>(),
        C.data_ptr<float>(),
        M, N, K
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""


@lru_cache(maxsize=1)
def _extension():
    return load_inline(
        name="matmul_abt_tf32_residual_v1",
        cpp_sources=(
            "torch::Tensor matmul_abt_cuda("
            "torch::Tensor A, torch::Tensor B);"
        ),
        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:
        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.is_contiguous()
            and B.is_contiguous()
            and A.shape[1] == B.shape[1]
            and A.shape[0] % 128 == 0
            and B.shape[0] % 64 == 0
            and A.shape[1] % 32 == 0
            and A.data_ptr() % 16 == 0
            and B.data_ptr() % 16 == 0
            and not (
                torch.is_grad_enabled()
                and (A.requires_grad or B.requires_grad)
            )
        )
        if supported and torch.cuda.get_device_capability(A.device)[0] >= 8:
            return _extension().matmul_abt_cuda(A, B)
        return torch.matmul(A, B.T)
```