```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_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 BM = 128;
constexpr int BN = 64;
constexpr int BK = 32;
constexpr int THREADS = 256;
constexpr int SHARED_BYTES =
    2 * (BM * BK + BK * BN) * sizeof(float);

__global__ void matmul_tf32_corrected(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ c,
    int M, int N, int K)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    extern __shared__ __align__(32) float storage[];
    float* ah = storage;
    float* al = ah + BM * BK;
    float* bh = al + BM * BK;
    float* bl = bh + BK * BN;

    const int tid = threadIdx.x;
    const int warp = tid / 32;
    const int tile_m = blockIdx.y * BM;
    const int tile_n = blockIdx.x * BN;
    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::row_major>;
    using CFragment = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;

    CFragment 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 k0 = 0; k0 < K; k0 += BK) {
        #pragma unroll
        for (int idx = tid; idx < BM * BK; idx += THREADS) {
            const int row = idx / BK;
            const int col = idx % BK;
            float value = 0.0f;
            if (tile_m + row < M && k0 + col < K)
                value = a[(size_t)(tile_m + row) * K + k0 + col];

            const float high = wmma::__float_to_tf32(value);
            ah[idx] = high;
            al[idx] = wmma::__float_to_tf32(value - high);
        }

        #pragma unroll
        for (int idx = tid; idx < BK * BN; idx += THREADS) {
            const int row = idx / BN;
            const int col = idx % BN;
            float value = 0.0f;
            if (k0 + row < K && tile_n + col < N)
                value = b[(size_t)(k0 + row) * N + tile_n + col];

            const float high = wmma::__float_to_tf32(value);
            bh[idx] = high;
            bl[idx] = wmma::__float_to_tf32(value - high);
        }

        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            AFragment a_high[2], a_low[2];
            BFragment b_high[2], b_low[2];

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

            #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_low[i], b_high[j], acc[i][j]);
                    wmma::mma_sync(
                        acc[i][j], a_high[i], b_low[j], acc[i][j]);
                    wmma::mma_sync(
                        acc[i][j], a_high[i], b_high[j], acc[i][j]);
                }
            }
        }

        __syncthreads();
    }

    // Reuse the input staging memory for coalesced output stores.
    float* output_tile = storage;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::store_matrix_sync(
                output_tile + (warp_m + i * 16) * BN + warp_n + j * 16,
                acc[i][j], BN, wmma::mem_row_major);
        }
    }

    __syncthreads();

    #pragma unroll
    for (int idx = tid; idx < BM * BN; idx += THREADS) {
        const int row = tile_m + idx / BN;
        const int col = tile_n + idx % BN;
        if (row < M && col < N)
            c[(size_t)row * N + col] = output_tile[idx];
    }
#endif
}

torch::Tensor matmul_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 matrices");
    TORCH_CHECK(a.size(1) == b.size(0), "Incompatible matrix shapes");
    TORCH_CHECK(
        a.scalar_type() == at::kFloat && b.scalar_type() == at::kFloat,
        "Expected float32 inputs");
    TORCH_CHECK(a.is_contiguous() && b.is_contiguous(),
                "Expected contiguous inputs");

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

    const auto* properties = at::cuda::getCurrentDeviceProperties();
    if (properties->major < 8)
        return at::matmul(a, b);

    const int64_t m = a.size(0);
    const int64_t k = a.size(1);
    const int64_t n = b.size(1);

    TORCH_CHECK(m <= INT_MAX && n <= INT_MAX && k <= INT_MAX,
                "Matrix dimensions are too large");

    auto out = torch::empty({m, n}, a.options());
    if (m == 0 || n == 0)
        return out;
    if (k == 0)
        return out.zero_();

    const dim3 grid((n + BN - 1) / BN, (m + BM - 1) / BM);
    const auto stream = at::cuda::getCurrentCUDAStream();

    matmul_tf32_corrected<<<grid, THREADS, SHARED_BYTES, stream.stream()>>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        out.data_ptr<float>(),
        static_cast<int>(m),
        static_cast<int>(n),
        static_cast<int>(k));

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

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

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