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

__device__ __forceinline__ float round_tf32(float x) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
    unsigned int bits;
    asm("cvt.rna.tf32.f32 %0, %1;" : "=r"(bits) : "f"(x));
    return __uint_as_float(bits);
#else
    return x;
#endif
}

__global__ void small_k_matmul_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 BM = 128;
    constexpr int BN = 128;
    constexpr int BK = 16;

    __shared__ __align__(32) float a_hi[BM * BK];
    __shared__ __align__(32) float a_lo[BM * BK];
    __shared__ __align__(32) float b_hi[BK * BN];
    __shared__ __align__(32) float b_lo[BK * BN];

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int warp_row = (warp >> 1) * 32;
    const int warp_col = (warp & 1) * 64;

    const int block_row = blockIdx.y * BM;
    const int block_col = blockIdx.x * BN;

    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 accum[2][4];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            wmma::fill_fragment(accum[i][j], 0.0f);
        }
    }

    for (int k_base = 0; k_base < K; k_base += BK) {
        #pragma unroll
        for (int offset = tid; offset < BM * BK; offset += 256) {
            const int row = block_row + offset / BK;
            const int k = k_base + offset % BK;
            float value = 0.0f;
            if (row < M && k < K) {
                value = A[static_cast<int64_t>(row) * K + k];
            }

            const float hi = round_tf32(value);
            a_hi[offset] = hi;
            a_lo[offset] = round_tf32(value - hi);
        }

        #pragma unroll
        for (int offset = tid; offset < BK * BN; offset += 256) {
            const int k = k_base + offset / BN;
            const int col = block_col + offset % BN;
            float value = 0.0f;
            if (k < K && col < N) {
                value = B[static_cast<int64_t>(k) * N + col];
            }

            const float hi = round_tf32(value);
            b_hi[offset] = hi;
            b_lo[offset] = round_tf32(value - hi);
        }

        __syncthreads();

        #pragma unroll
        for (int kk = 0; kk < BK; kk += 8) {
            AFragment ah[2], al[2];
            BFragment bh[4], bl[4];

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

            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                const int offset = kk * BN + warp_col + j * 16;
                wmma::load_matrix_sync(bh[j], b_hi + offset, BN);
                wmma::load_matrix_sync(bl[j], b_lo + offset, BN);
            }

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

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int row = block_row + warp_row + i * 16;
            const int col = block_col + warp_col + j * 16;

            if (row < M && col < N) {
                wmma::store_matrix_sync(
                    C + static_cast<int64_t>(row) * N + col,
                    accum[i][j],
                    N,
                    wmma::mem_row_major);
            }
        }
    }
#endif
}

torch::Tensor small_k_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(), "Device mismatch");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Expected matrices");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible matrix dimensions");

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

    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    const int64_t N = B.size(1);

    if (A.scalar_type() != at::kFloat ||
        B.scalar_type() != at::kFloat ||
        !A.is_contiguous() || !B.is_contiguous() ||
        M % 16 != 0 || N % 16 != 0 ||
        M > 65535LL * 128 || N > 2147483520LL ||
        K > 2147483632LL ||
        at::cuda::getCurrentDeviceProperties()->major < 8) {
        return at::matmul(A, B);
    }

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

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

    small_k_matmul_kernel<<<
        grid, 256, 0, c10::cuda::getCurrentCUDAStream()
    >>>(
        A.data_ptr<float>(),
        B.data_ptr<float>(),
        C.data_ptr<float>(),
        static_cast<int>(M),
        static_cast<int>(N),
        static_cast<int>(K));

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_small_k_matmul = load_inline(
    name="small_k_matmul_tf32_residual_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["small_k_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:
            return torch.matmul(A, B)
        return _small_k_matmul.small_k_matmul_cuda(A, B)
```