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

namespace wmma = nvcuda::wmma;

constexpr int TILE = 64;
constexpr int STEP = 32;
constexpr int STRIDE = 40;
constexpr int PLANE = TILE * STRIDE;

__global__ void matmul_nt_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int64_t M, int64_t N, int64_t K)
{
    const int64_t row_base = int64_t(blockIdx.y) * TILE;
    const int64_t col_base = int64_t(blockIdx.x) * TILE;

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
    __shared__ __align__(32) unsigned char storage[
        4 * PLANE * sizeof(half)];

    half* ah = reinterpret_cast<half*>(storage);
    half* al = ah + PLANE;
    half* bh = al + PLANE;
    half* bl = bh + PLANE;

    const int warp = threadIdx.x / 32;
    const int warp_row = (warp / 2) * 32;
    const int warp_col = (warp % 2) * 32;

    wmma::fragment<wmma::accumulator, 16, 16, 16, float>
        main_acc[2][2], 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 (int64_t kb = 0; kb < K; kb += STEP) {
        for (int t = threadIdx.x; t < TILE * STEP; t += blockDim.x) {
            const int r = t / STEP;
            const int k = t % STEP;
            const int64_t gk = kb + k;

            float av = 0.0f;
            float bv = 0.0f;
            if (row_base + r < M && gk < K)
                av = A[(row_base + r) * K + gk];
            if (col_base + r < N && gk < K)
                bv = B[(col_base + r) * K + gk];

            const half av_hi = __float2half_rn(av);
            const half bv_hi = __float2half_rn(bv);
            const int dst = r * STRIDE + k;

            ah[dst] = av_hi;
            al[dst] = __float2half_rn(av - __half2float(av_hi));
            bh[dst] = bv_hi;
            bl[dst] = __float2half_rn(bv - __half2float(bv_hi));
        }
        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < STEP; kk += 16) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 16, half, wmma::row_major>
                a_hi[2], a_lo[2];
            wmma::fragment<
                wmma::matrix_b, 16, 16, 16, half, wmma::col_major>
                b_hi[2], b_lo[2];

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int offset = (warp_row + 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_col + 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(
                        main_acc[i][j], a_hi[i], b_hi[j], main_acc[i][j]);
                    wmma::mma_sync(
                        correction[i][j], a_lo[i], b_hi[j], correction[i][j]);
                    wmma::mma_sync(
                        correction[i][j], a_hi[i], b_lo[j], correction[i][j]);
                }
            }
        }
        __syncthreads();
    }

    float* result = reinterpret_cast<float*>(storage);

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            #pragma unroll
            for (int x = 0; x < main_acc[i][j].num_elements; ++x)
                main_acc[i][j].x[x] += correction[i][j].x[x];

            wmma::store_matrix_sync(
                result + (warp_row + i * 16) * TILE + warp_col + j * 16,
                main_acc[i][j], TILE, wmma::mem_row_major);
        }
    }
    __syncthreads();

    for (int t = threadIdx.x; t < TILE * TILE; t += blockDim.x) {
        const int64_t row = row_base + t / TILE;
        const int64_t col = col_base + t % TILE;
        if (row < M && col < N)
            C[row * N + col] = result[t];
    }
#else
    for (int t = threadIdx.x; t < TILE * TILE; t += blockDim.x) {
        const int64_t row = row_base + t / TILE;
        const int64_t col = col_base + t % TILE;
        if (row < M && col < N) {
            float sum = 0.0f;
            for (int64_t k = 0; k < K; ++k)
                sum = fmaf(A[row * K + k], B[col * K + k], sum);
            C[row * N + col] = sum;
        }
    }
#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(), "Inputs must share a device");
    TORCH_CHECK(A.scalar_type() == at::kFloat &&
                B.scalar_type() == at::kFloat, "Expected float32 inputs");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Expected matrices");
    TORCH_CHECK(A.size(1) == B.size(1), "Incompatible dimensions");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Expected contiguous inputs");

    const c10::cuda::CUDAGuard guard(A.device());
    const int64_t M = A.size(0);
    const int64_t N = B.size(0);
    const int64_t K = A.size(1);

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

    const dim3 grid((N + TILE - 1) / TILE, (M + TILE - 1) / TILE);
    matmul_nt_kernel<<<grid, 128, 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;
}
"""

_matmul_ext = load_inline(
    name="matmul_nt_split_precision_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
        ):
            return torch.matmul(A, B.T)
        return _matmul_ext.matmul_nt_cuda(A.contiguous(), B.contiguous())
```