```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 tf32x3_matmul(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 {

constexpr int BM = 128;
constexpr int BN = 64;
constexpr int BK = 32;
constexpr int THREADS = 256;

__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
}

union SharedStorage {
    struct {
        float ah[BM * BK];
        float al[BM * BK];
        float bh[BK * BN];
        float bl[BK * BN];
    } operands;
    float output[BM * BN];
};

__global__ void matmul_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
    using namespace nvcuda;

    __shared__ __align__(32) SharedStorage shared;

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

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

    for (int kb = 0; kb < k; kb += BK) {
        for (int idx = tid; idx < BM * BK; idx += THREADS) {
            const int row = row_base + idx / BK;
            const int col = kb + idx % BK;
            const float value =
                (row < m && col < k) ? a[(int64_t)row * k + col] : 0.0f;
            const float high = round_tf32(value);
            shared.operands.ah[idx] = high;
            shared.operands.al[idx] = round_tf32(value - high);
        }

        for (int idx = tid; idx < BK * BN; idx += THREADS) {
            const int row = kb + idx / BN;
            const int col = col_base + idx % BN;
            const float value =
                (row < k && col < n) ? b[(int64_t)row * n + col] : 0.0f;
            const float high = round_tf32(value);
            shared.operands.bh[idx] = high;
            shared.operands.bl[idx] = round_tf32(value - high);
        }

        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            FragA ah[2], al[2];
            FragB bh[2], bl[2];

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

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

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                #pragma unroll
                for (int j = 0; j < 2; ++j) {
                    wmma::mma_sync(
                        main_acc[i][j], ah[i], bh[j], main_acc[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 t = 0; t < Acc::num_elements; ++t) {
                main_acc[i][j].x[t] += correction[i][j].x[t];
            }
        }
    }

    const bool direct_store =
        row_base + BM <= m && col_base + BN <= n && n % 8 == 0;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            const int row = warp_m + i * 16;
            const int col = warp_n + j * 16;
            if (direct_store) {
                wmma::store_matrix_sync(
                    c + (int64_t)(row_base + row) * n + col_base + col,
                    main_acc[i][j], n, wmma::mem_row_major);
            } else {
                wmma::store_matrix_sync(
                    shared.output + row * BN + col,
                    main_acc[i][j], BN, wmma::mem_row_major);
            }
        }
    }

    if (!direct_store) {
        __syncthreads();
        for (int idx = tid; idx < BM * BN; idx += THREADS) {
            const int row = row_base + idx / BN;
            const int col = col_base + idx % BN;
            if (row < m && col < n) {
                c[(int64_t)row * n + col] = shared.output[idx];
            }
        }
    }
#endif
}

} // namespace

torch::Tensor tf32x3_matmul(torch::Tensor a, torch::Tensor b) {
    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(a.device() == b.device(), "Devices must match");
    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 dimensions");
    TORCH_CHECK(a.is_contiguous() && b.is_contiguous(),
                "Expected contiguous matrices");

    c10::cuda::CUDAGuard guard(a.device());
    TORCH_CHECK(
        at::cuda::getCurrentDeviceProperties()->major >= 8,
        "TF32 kernels require 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));
    auto c = torch::empty({m, n}, a.options());

    if (m == 0 || n == 0) {
        return c;
    }

    const dim3 grid((n + BN - 1) / BN, (m + BM - 1) / BM);
    matmul_kernel<<<grid, THREADS, 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;
}
"""

_extension = None


def _get_extension():
    global _extension
    if _extension is None:
        _extension = load_inline(
            name="tf32x3_wmma_matmul_128x64",
            cpp_sources=_cpp_source,
            cuda_sources=_cuda_source,
            functions=["tf32x3_matmul"],
            extra_cflags=["-O3"],
            extra_cuda_cflags=["-O3"],
            verbose=False,
        )
    return _extension


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.device != B.device
            or A.dtype != torch.float32
            or B.dtype != torch.float32
            or A.ndim != 2
            or B.ndim != 2
            or not A.is_contiguous()
            or not B.is_contiguous()
            or (torch.is_grad_enabled() and (A.requires_grad or B.requires_grad))
            or torch.cuda.get_device_capability(A.device)[0] < 8
        ):
            return torch.matmul(A, B)
        return _get_extension().tf32x3_matmul(A, B)
```