```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(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 {

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

__global__ void pack_inputs(
    const float* __restrict__ A,
    const float* __restrict__ B,
    __half2* __restrict__ packed,
    int64_t na,
    int64_t total)
{
    for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         i < total;
         i += int64_t(blockDim.x) * gridDim.x) {
        float x = i < na ? A[i] : B[i - na];
        __half hi = __float2half_rn(x);
        __half lo = __float2half_rn(x - __half2float(hi));
        packed[i] = __halves2half2(hi, lo);
    }
}

union __align__(32) TileStorage {
    struct {
        __half ah[BM * BK];
        __half al[BM * BK];
        __half bh[BK * BN];
        __half bl[BK * BN];
    } input;
    float output[BM * BN];
};

template<bool FULL>
__global__ __launch_bounds__(THREADS)
void tensor_mm(
    const __half2* __restrict__ A,
    const __half2* __restrict__ B,
    float* __restrict__ C,
    int M,
    int N,
    int K)
{
#if __CUDA_ARCH__ >= 700
    using namespace nvcuda;

    __shared__ TileStorage smem;

    const int tid = threadIdx.x;
    const int warp = tid >> 5;
    const int wm = (warp >> 2) * 32;
    const int wn = (warp & 3) * 32;
    const int row0 = blockIdx.y * BM;
    const int col0 = blockIdx.x * BN;

    wmma::fragment<wmma::accumulator, 16, 16, 16, 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 kb = 0; kb < K; kb += BK) {
        #pragma unroll
        for (int index = tid; index < BM * BK; index += THREADS) {
            const int r = index / BK;
            const int k = index % BK;
            __half2 x = __float2half2_rn(0.0f);
            if (FULL || (row0 + r < M && kb + k < K))
                x = A[int64_t(row0 + r) * K + kb + k];
            smem.input.ah[index] = __low2half(x);
            smem.input.al[index] = __high2half(x);
        }

        #pragma unroll
        for (int index = tid; index < BK * BN; index += THREADS) {
            const int k = index / BN;
            const int c = index % BN;
            __half2 x = __float2half2_rn(0.0f);
            if (FULL || (kb + k < K && col0 + c < N))
                x = B[int64_t(kb + k) * N + col0 + c];
            smem.input.bh[index] = __low2half(x);
            smem.input.bl[index] = __high2half(x);
        }

        __syncthreads();

        wmma::fragment<wmma::matrix_a, 16, 16, 16,
                       __half, wmma::row_major> ah[2], al[2];
        wmma::fragment<wmma::matrix_b, 16, 16, 16,
                       __half, wmma::row_major> bh[2], bl[2];

        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            wmma::load_matrix_sync(
                ah[i], smem.input.ah + (wm + i * 16) * BK, BK);
            wmma::load_matrix_sync(
                al[i], smem.input.al + (wm + i * 16) * BK, BK);
            wmma::load_matrix_sync(
                bh[i], smem.input.bh + wn + i * 16, BN);
            wmma::load_matrix_sync(
                bl[i], smem.input.bl + wn + i * 16, 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], 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]);
            }
        }

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 2; ++j) {
            wmma::store_matrix_sync(
                smem.output + (wm + i * 16) * BN + wn + j * 16,
                acc[i][j], BN, wmma::mem_row_major);
        }
    }

    __syncthreads();

    if (FULL || ((N & 3) == 0 && col0 + BN <= N)) {
        #pragma unroll
        for (int index = tid; index < BM * BN / 4; index += THREADS) {
            const int r = index / (BN / 4);
            const int c = (index % (BN / 4)) * 4;
            if (FULL || row0 + r < M) {
                *reinterpret_cast<float4*>(
                    C + int64_t(row0 + r) * N + col0 + c) =
                    reinterpret_cast<const float4*>(smem.output)[index];
            }
        }
    } else {
        #pragma unroll
        for (int index = tid; index < BM * BN; index += THREADS) {
            const int r = index / BN;
            const int c = index % BN;
            if (row0 + r < M && col0 + c < N)
                C[int64_t(row0 + r) * N + col0 + c] = smem.output[index];
        }
    }
#endif
}

} // namespace

torch::Tensor small_k_mm(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "CUDA inputs required");
    TORCH_CHECK(A.device() == B.device(), "Device mismatch");
    TORCH_CHECK(A.scalar_type() == at::kFloat &&
                B.scalar_type() == at::kFloat, "float32 inputs required");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Matrices required");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Contiguous inputs required");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible dimensions");

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

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

    const int64_t na = A.numel();
    const int64_t total = na + B.numel();
    auto packed = torch::empty({total}, A.options());
    auto* p = reinterpret_cast<__half2*>(packed.data_ptr<float>());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (total != 0) {
        const int blocks = static_cast<int>(
            std::min<int64_t>((total + 255) / 256, 8192));
        pack_inputs<<<blocks, 256, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), p, na, total);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
    if (M % BM == 0 && N % BN == 0 && K % BK == 0) {
        tensor_mm<true><<<grid, THREADS, 0, stream>>>(
            p, p + na, C.data_ptr<float>(), M, N, K);
    } else {
        tensor_mm<false><<<grid, THREADS, 0, stream>>>(
            p, p + na, C.data_ptr<float>(), M, N, K);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_extension = load_inline(
    name="small_k_dual_half_tensorcore_v2",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["small_k_mm"],
    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 A.shape[1] == B.shape[0]
            and A.shape[0] <= 64 * 65535
            and B.shape[1] < 2**31 - 128
            and A.shape[1] < 2**31 - 16
            and torch.cuda.get_device_capability(A.device)[0] >= 7
            and not (
                torch.is_grad_enabled()
                and (A.requires_grad or B.requires_grad)
            )
        ):
            return _extension.small_k_mm(A, B)
        return torch.matmul(A, B)
```