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

namespace wmma = nvcuda::wmma;

constexpr int BM = 64;
constexpr int BN = 128;
constexpr int BK = 16;

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

__global__ void small_k_mm_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
    __shared__ SharedStorage shared;

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int warp_row = (warp >> 2) * 32;
    const int warp_col = (warp & 3) * 32;
    const int row_base = blockIdx.y * BM;
    const int col_base = blockIdx.x * BN;

    wmma::fragment<wmma::accumulator, 16, 16, 8, float> 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 k_base = 0; k_base < K; k_base += BK) {
        #pragma unroll
        for (int index = tid; index < BM * BK; index += 256) {
            const int row = row_base + index / BK;
            const int k = k_base + index % BK;
            const float value =
                (row < M && k < K)
                    ? A[static_cast<int64_t>(row) * K + k]
                    : 0.0f;
            const float high = wmma::__float_to_tf32(value);
            shared.input.ah[index] = high;
            shared.input.al[index] = wmma::__float_to_tf32(value - high);
        }

        #pragma unroll
        for (int index = tid; index < BK * BN; index += 256) {
            const int k = k_base + index / BN;
            const int col = col_base + index % BN;
            const float value =
                (k < K && col < N)
                    ? B[static_cast<int64_t>(k) * N + col]
                    : 0.0f;
            const float high = wmma::__float_to_tf32(value);
            shared.input.bh[index] = high;
            shared.input.bl[index] = wmma::__float_to_tf32(value - high);
        }

        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            wmma::fragment<
                wmma::matrix_a, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major
            > ah[2], al[2];

            wmma::fragment<
                wmma::matrix_b, 16, 16, 8,
                wmma::precision::tf32, wmma::row_major
            > bh[2], bl[2];

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

            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                const int offset = kk * BN + warp_col + j * 16;
                wmma::load_matrix_sync(bh[j], shared.input.bh + offset, BN);
                wmma::load_matrix_sync(bl[j], shared.input.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], ah[i], bl[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], al[i], bh[j], acc[i][j]);
                    wmma::mma_sync(acc[i][j], ah[i], bh[j], acc[i][j]);
                }
            }
        }

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            float* destination =
                shared.output
                + (warp_row + i * 16) * BN
                + warp_col + j * 16;

            wmma::store_matrix_sync(
                destination, acc[i][j], BN, wmma::mem_row_major
            );
        }
    }

    __syncthreads();

    if ((N & 3) == 0) {
        #pragma unroll
        for (int index = tid; index < BM * BN / 4; index += 256) {
            const int local_row = index / (BN / 4);
            const int local_col = (index % (BN / 4)) * 4;
            const int row = row_base + local_row;
            const int col = col_base + local_col;

            if (row < M && col < N) {
                const float4 value =
                    reinterpret_cast<const float4*>(shared.output)[index];
                *reinterpret_cast<float4*>(
                    C + static_cast<int64_t>(row) * N + col
                ) = value;
            }
        }
    } else {
        #pragma unroll
        for (int index = tid; index < BM * BN; index += 256) {
            const int row = row_base + index / BN;
            const int col = col_base + index % BN;
            if (row < M && col < N) {
                C[static_cast<int64_t>(row) * N + col] = shared.output[index];
            }
        }
    }
#endif
}

torch::Tensor small_k_mm_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Expected CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Expected matching devices");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Expected matrices");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible matrix dimensions");
    TORCH_CHECK(
        A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
        "Expected float32 tensors"
    );
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Expected contiguous tensors");

    const c10::cuda::CUDAGuard guard(A.device());

    if (at::cuda::getCurrentDeviceProperties()->major < 8) {
        return at::matmul(A, B);
    }

    const int64_t M64 = A.size(0);
    const int64_t K64 = A.size(1);
    const int64_t N64 = B.size(1);

    if (M64 > 2147483520LL || N64 > 2147483520LL ||
        K64 > 2147483520LL || (M64 + BM - 1) / BM > 65535) {
        return at::matmul(A, B);
    }

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

    const dim3 grid(
        static_cast<unsigned>((N64 + BN - 1) / BN),
        static_cast<unsigned>((M64 + BM - 1) / BM)
    );

    small_k_mm_kernel<<<
        grid, 256, 0, c10::cuda::getCurrentCUDAStream()
    >>>(
        A.data_ptr<float>(),
        B.data_ptr<float>(),
        C.data_ptr<float>(),
        static_cast<int>(M64),
        static_cast<int>(N64),
        static_cast<int>(K64)
    );

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_small_k_mm = load_inline(
    name="small_k_tf32_compensated_mm",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["small_k_mm_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 (
            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 not (
                torch.is_grad_enabled()
                and (A.requires_grad or B.requires_grad)
            )
        ):
            return _small_k_mm.small_k_mm_cuda(A, B)
        return torch.matmul(A, B)
```