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

__global__ void matmul_nt_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 TILE = 64;
    constexpr int BK = 32;

    __shared__ __align__(32) float storage[4 * TILE * BK];
    float* ah = storage;
    float* al = ah + TILE * BK;
    float* bh = al + TILE * BK;
    float* bl = bh + TILE * BK;

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int tile_m = blockIdx.y * TILE;
    const int tile_n = blockIdx.x * TILE;
    const int warp_m = (warp / 2) * 32;
    const int warp_n = (warp % 2) * 32;

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

    Accumulator main_acc[2][2];
    Accumulator correction_acc[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_acc[i][j], 0.0f);
        }
    }

    for (int base_k = 0; base_k < K; base_k += BK) {
        for (int index = tid; index < TILE * BK; index += 128) {
            const int row = index / BK;
            const int k = base_k + index % BK;

            float av = 0.0f;
            float bv = 0.0f;
            if (tile_m + row < M && k < K)
                av = a[static_cast<int64_t>(tile_m + row) * K + k];
            if (tile_n + row < N && k < K)
                bv = b[static_cast<int64_t>(tile_n + row) * K + k];

            const float av_hi = __float_to_tf32(av);
            const float bv_hi = __float_to_tf32(bv);
            ah[index] = av_hi;
            al[index] = __float_to_tf32(av - av_hi);
            bh[index] = bv_hi;
            bl[index] = __float_to_tf32(bv - bv_hi);
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; 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) * 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(
                        main_acc[i][j], a_hi[i], b_hi[j],
                        main_acc[i][j]);
                    wmma::mma_sync(
                        correction_acc[i][j], a_hi[i], b_lo[j],
                        correction_acc[i][j]);
                    wmma::mma_sync(
                        correction_acc[i][j], a_lo[i], b_hi[j],
                        correction_acc[i][j]);
                }
            }
        }
        __syncthreads();
    }

    // Reuse the input staging memory for the completed output tile.
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            #pragma unroll
            for (int e = 0; e < Accumulator::num_elements; ++e)
                main_acc[i][j].x[e] += correction_acc[i][j].x[e];

            wmma::store_matrix_sync(
                storage + (warp_m + i * 16) * TILE + warp_n + j * 16,
                main_acc[i][j], TILE, wmma::mem_row_major);
        }
    }
    __syncthreads();

    for (int index = tid; index < TILE * TILE; index += 128) {
        const int row = tile_m + index / TILE;
        const int col = tile_n + index % TILE;
        if (row < M && col < N)
            c[static_cast<int64_t>(row) * N + col] = storage[index];
    }
#endif
}

torch::Tensor matmul_nt_cuda(torch::Tensor a, torch::Tensor b) {
    TORCH_CHECK(a.is_cuda() && b.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(a.device() == b.device(), "Tensor 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(1), "Inner dimensions must match");
    TORCH_CHECK(a.is_contiguous() && b.is_contiguous(),
                "Expected contiguous tensors");
    TORCH_CHECK(a.size(0) <= INT_MAX && b.size(0) <= INT_MAX &&
                a.size(1) <= INT_MAX, "Matrix dimensions are too large");

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

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

    if (at::cuda::getDeviceProperties(a.get_device())->major < 8 ||
        K == 0 || M > 65535 * 64) {
        return at::matmul(a, b.transpose(0, 1));
    }

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

    const dim3 grid(
        (static_cast<unsigned int>(N) + 63) / 64,
        (static_cast<unsigned int>(M) + 63) / 64);

    matmul_nt_kernel<<<
        grid, 128, 0, at::cuda::getCurrentCUDAStream(a.get_device())
    >>>(
        a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(),
        M, N, K);

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_matmul_extension = load_inline(
    name="custom_matmul_nt_tf32_split_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["matmul_nt_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 (torch.is_grad_enabled() and (A.requires_grad or B.requires_grad))
        ):
            return torch.matmul(A, B.T)

        return _matmul_extension.matmul_nt_cuda(
            A.contiguous(), B.contiguous()
        )
```