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

namespace {

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

__global__ void small_k_kernel(
    const float* __restrict__ a,
    const float* __restrict__ b,
    float* __restrict__ c,
    int64_t n)
{
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    using namespace nvcuda;

    // Input tiles and the output staging tile share storage.
    __shared__ __align__(32) float storage[BM * BK + BK * BN];
    float* as = storage;
    float* bs = storage + BM * BK;

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int warp_m = (warp >> 2) * 32;
    const int warp_n = (warp & 3) * 32;

    const int64_t row0 = static_cast<int64_t>(blockIdx.y) * BM;
    const int64_t col0 = static_cast<int64_t>(blockIdx.x) * BN;

    for (int v = tid; v < BM * BK / 4; v += THREADS) {
        reinterpret_cast<float4*>(as)[v] =
            reinterpret_cast<const float4*>(a + row0 * BK)[v];
    }

    for (int v = tid; v < BK * BN / 4; v += THREADS) {
        const int k = v / (BN / 4);
        const int j = (v % (BN / 4)) * 4;
        reinterpret_cast<float4*>(bs)[v] =
            *reinterpret_cast<const float4*>(b + k * n + col0 + j);
    }

    __syncthreads();

    using AFrag = wmma::fragment<
        wmma::matrix_a, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using BFrag = wmma::fragment<
        wmma::matrix_b, 16, 16, 8,
        wmma::precision::tf32, wmma::row_major>;
    using CFrag = wmma::fragment<
        wmma::accumulator, 16, 16, 8, float>;

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

    #pragma unroll
    for (int k = 0; k < BK; k += 8) {
        AFrag ah[2], al[2];
        BFrag bh[2], bl[2];

        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            wmma::load_matrix_sync(
                ah[i], as + (warp_m + i * 16) * BK + k, BK);

            #pragma unroll
            for (int t = 0; t < AFrag::num_elements; ++t) {
                const float value = ah[i].x[t];
                const float high = __float_to_tf32(value);
                ah[i].x[t] = high;
                al[i].x[t] = __float_to_tf32(value - high);
            }
        }

        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::load_matrix_sync(
                bh[j], bs + k * BN + warp_n + j * 16, BN);

            #pragma unroll
            for (int t = 0; t < BFrag::num_elements; ++t) {
                const float value = bh[j].x[t];
                const float high = __float_to_tf32(value);
                bh[j].x[t] = high;
                bl[j].x[t] = __float_to_tf32(value - high);
            }
        }

        // Three TF32 products retain substantially more precision than
        // a single TF32 multiplication.
        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            #pragma unroll
            for (int j = 0; j < 2; ++j) {
                wmma::mma_sync(acc[i][j], al[i], bh[j], acc[i][j]);
                wmma::mma_sync(acc[i][j], ah[i], bl[j], acc[i][j]);
                wmma::mma_sync(acc[i][j], ah[i], bh[j], acc[i][j]);
            }
        }
    }

    // All warps must finish reading the inputs before reusing storage.
    __syncthreads();

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

    __syncthreads();

    for (int v = tid; v < BM * BN / 4; v += THREADS) {
        const int row = v / (BN / 4);
        const int col = (v % (BN / 4)) * 4;
        *reinterpret_cast<float4*>(c + (row0 + row) * n + col0 + col) =
            reinterpret_cast<const float4*>(storage)[v];
    }
#endif
}

}  // namespace

torch::Tensor small_k_matmul(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");

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

    const bool supported =
        a.dim() == 2 && b.dim() == 2 &&
        a.scalar_type() == at::kFloat &&
        b.scalar_type() == at::kFloat &&
        a.is_contiguous() && b.is_contiguous() &&
        a.size(1) == BK && b.size(0) == BK &&
        a.size(0) > 0 && b.size(1) > 0 &&
        a.size(0) % BM == 0 && b.size(1) % BN == 0 &&
        a.size(0) / BM <= 65535 &&
        b.size(1) / BN <= 2147483647LL &&
        reinterpret_cast<uintptr_t>(a.data_ptr()) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(b.data_ptr()) % 16 == 0;

    if (!supported ||
        at::cuda::getDeviceProperties(a.get_device())->major < 8) {
        return at::matmul(a, b);
    }

    auto c = torch::empty({a.size(0), b.size(1)}, a.options());

    const dim3 grid(
        static_cast<unsigned int>(b.size(1) / BN),
        static_cast<unsigned int>(a.size(0) / BM));

    small_k_kernel<<<
        grid, THREADS, 0, at::cuda::getCurrentCUDAStream()
    >>>(
        a.data_ptr<float>(),
        b.data_ptr<float>(),
        c.data_ptr<float>(),
        b.size(1));

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return c;
}
"""

_extension = load_inline(
    name="small_k64_tf32_compensated_matmul",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["small_k_matmul"],
    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:
            return torch.matmul(A, B)
        if torch.is_grad_enabled() and (A.requires_grad or B.requires_grad):
            return torch.matmul(A, B)
        return _extension.small_k_matmul(A, B)
```