```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp = r"""
#include <torch/extension.h>
torch::Tensor small_k_mm(torch::Tensor A, torch::Tensor B);
"""

_cuda = 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 <algorithm>

namespace {

constexpr int BM = 128;
constexpr int BN = 128;
constexpr int KK = 64;
constexpr int AS = 72;
constexpr int BS = 136;
constexpr int CS = 136;
constexpr int NT = 256;
constexpr int SMEM = 2 * (BM * AS + KK * BS) * sizeof(__half);

static_assert(SMEM >= BM * CS * sizeof(float),
              "Insufficient shared memory for output staging");

__device__ __forceinline__ unsigned half2_bits(__half2 value)
{
    return static_cast<unsigned>(__half_as_ushort(__low2half(value))) |
           (static_cast<unsigned>(__half_as_ushort(__high2half(value))) << 16);
}

__global__ void pack_planar(
    const float* __restrict__ A,
    const float* __restrict__ B,
    __half* __restrict__ high,
    __half* __restrict__ low,
    int64_t na,
    int64_t total)
{
    for (int64_t i = 4LL * (int64_t(blockIdx.x) * blockDim.x + threadIdx.x);
         i < total;
         i += 4LL * blockDim.x * gridDim.x) {
        const float4 x = i < na
            ? *reinterpret_cast<const float4*>(A + i)
            : *reinterpret_cast<const float4*>(B + i - na);

        const __half2 h0 = __floats2half2_rn(x.x, x.y);
        const __half2 h1 = __floats2half2_rn(x.z, x.w);
        const float2 f0 = __half22float2(h0);
        const float2 f1 = __half22float2(h1);

        const __half2 l0 = __floats2half2_rn(x.x - f0.x, x.y - f0.y);
        const __half2 l1 = __floats2half2_rn(x.z - f1.x, x.w - f1.y);

        *reinterpret_cast<uint2*>(high + i) =
            make_uint2(half2_bits(h0), half2_bits(h1));
        *reinterpret_cast<uint2*>(low + i) =
            make_uint2(half2_bits(l0), half2_bits(l1));
    }
}

__device__ __forceinline__ void copy16(
    __half* destination, const __half* source)
{
#if __CUDA_ARCH__ >= 800
    const unsigned address =
        static_cast<unsigned>(__cvta_generic_to_shared(destination));
    asm volatile(
        "cp.async.cg.shared.global [%0], [%1], 16;\n"
        :: "r"(address), "l"(source) : "memory");
#endif
}

__device__ __forceinline__ void load_a(
    unsigned (&r)[4], const __half* p)
{
#if __CUDA_ARCH__ >= 800
    const unsigned address =
        static_cast<unsigned>(__cvta_generic_to_shared(p));
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x4.shared.b16 "
        "{%0, %1, %2, %3}, [%4];"
        : "=r"(r[0]), "=r"(r[1]), "=r"(r[2]), "=r"(r[3])
        : "r"(address));
#endif
}

__device__ __forceinline__ void load_b(
    unsigned (&r)[2], const __half* p)
{
#if __CUDA_ARCH__ >= 800
    const unsigned address =
        static_cast<unsigned>(__cvta_generic_to_shared(p));
    asm volatile(
        "ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 "
        "{%0, %1}, [%2];"
        : "=r"(r[0]), "=r"(r[1])
        : "r"(address));
#endif
}

__device__ __forceinline__ void mma(
    float (&d)[4], const unsigned (&a)[4], const unsigned (&b)[2])
{
#if __CUDA_ARCH__ >= 800
    asm volatile(
        "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
        "{%0, %1, %2, %3}, "
        "{%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
        : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3])
        : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
          "r"(b[0]), "r"(b[1]));
#endif
}

__global__ __launch_bounds__(NT, 2)
void multiply(
    const __half* __restrict__ AH,
    const __half* __restrict__ AL,
    const __half* __restrict__ BH,
    const __half* __restrict__ BL,
    float* __restrict__ C,
    int N)
{
#if __CUDA_ARCH__ >= 800
    extern __shared__ __align__(32) unsigned char storage[];
    __half* ah = reinterpret_cast<__half*>(storage);
    __half* al = ah + BM * AS;
    __half* bh = al + BM * AS;
    __half* bl = bh + KK * BS;

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

    #pragma unroll
    for (int v = tid; v < BM * KK / 8; v += NT) {
        const int row = v / (KK / 8);
        const int k = (v % (KK / 8)) * 8;
        const int64_t src = int64_t(row0 + row) * KK + k;
        copy16(ah + row * AS + k, AH + src);
        copy16(al + row * AS + k, AL + src);
    }

    #pragma unroll
    for (int v = tid; v < KK * BN / 8; v += NT) {
        const int k = v / (BN / 8);
        const int col = (v % (BN / 8)) * 8;
        const int64_t src = int64_t(k) * N + col0 + col;
        copy16(bh + k * BS + col, BH + src);
        copy16(bl + k * BS + col, BL + src);
    }

    asm volatile("cp.async.commit_group;\n" ::: "memory");

    float accum[2][8][4];

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            #pragma unroll
            for (int q = 0; q < 4; ++q)
                accum[i][j][q] = 0.0f;
        }
    }

    asm volatile("cp.async.wait_group 0;\n" ::: "memory");
    __syncthreads();

    #pragma unroll
    for (int kb = 0; kb < KK; kb += 16) {
        unsigned a_high[2][4];
        unsigned a_low[2][4];

        #pragma unroll
        for (int i = 0; i < 2; ++i) {
            const int offset =
                (wm + i * 16 + (lane & 15)) * AS +
                kb + (lane >> 4) * 8;
            load_a(a_high[i], ah + offset);
            load_a(a_low[i], al + offset);
        }

        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            unsigned b_high[2];
            unsigned b_low[2];
            const int offset =
                (kb + (lane & 15)) * BS + wn + j * 8;
            load_b(b_high, bh + offset);
            load_b(b_low, bl + offset);

            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                mma(accum[i][j], a_low[i], b_high);
                mma(accum[i][j], a_high[i], b_low);
                mma(accum[i][j], a_high[i], b_high);
            }
        }
    }

    __syncthreads();

    float* tile = reinterpret_cast<float*>(storage);
    const int lane_row = lane >> 2;
    const int lane_col = (lane & 3) * 2;

    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        const int row = wm + i * 16 + lane_row;

        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            const int col = wn + j * 8 + lane_col;
            *reinterpret_cast<float2*>(tile + row * CS + col) =
                make_float2(accum[i][j][0], accum[i][j][1]);
            *reinterpret_cast<float2*>(tile + (row + 8) * CS + col) =
                make_float2(accum[i][j][2], accum[i][j][3]);
        }
    }

    __syncthreads();

    #pragma unroll
    for (int index = tid * 4; index < BM * BN; index += NT * 4) {
        const int row = index / BN;
        const int col = index % BN;
        const float4 value =
            *reinterpret_cast<const float4*>(tile + row * CS + col);
        *reinterpret_cast<float4*>(
            C + int64_t(row0 + row) * N + col0 + col) = value;
    }
#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, "Expected float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Expected matrices");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Expected contiguous inputs");
    TORCH_CHECK(A.size(1) == KK && B.size(0) == KK, "Expected K=64");
    TORCH_CHECK(A.size(0) % BM == 0 && B.size(1) % BN == 0,
                "Dimensions must be multiples of 128");

    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));
    auto C = torch::empty({M, N}, A.options());

    if (M == 0 || N == 0)
        return C;

    C10_CUDA_CHECK(cudaFuncSetAttribute(
        multiply, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM));

    const int64_t na = A.numel();
    const int64_t total = na + B.numel();
    auto packed = torch::empty({total}, A.options());
    auto* high = reinterpret_cast<__half*>(packed.data_ptr<float>());
    auto* low = high + total;

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    const int blocks = static_cast<int>(
        std::min<int64_t>((total + 1023) / 1024, 4096));

    pack_planar<<<blocks, 256, 0, stream>>>(
        A.data_ptr<float>(), B.data_ptr<float>(), high, low, na, total);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    multiply<<<dim3(N / BN, M / BM), NT, SMEM, stream>>>(
        high, low, high + na, low + na, C.data_ptr<float>(), N);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return C;
}
"""

_extension = load_inline(
    name="small_k_vectorized_mma_v5",
    cpp_sources=_cpp,
    cuda_sources=_cuda,
    functions=["small_k_mm"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()
        self._device_support = {}

    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] == 64
            and B.shape[0] == 64
            and A.shape[0] % 128 == 0
            and B.shape[1] % 128 == 0
            and A.shape[0] <= 128 * 65535
            and B.shape[1] < 2**31
            and A.data_ptr() % 16 == 0
            and B.data_ptr() % 16 == 0
            and not (
                torch.is_grad_enabled()
                and (A.requires_grad or B.requires_grad)
            )
        ):
            device = A.device.index
            if device not in self._device_support:
                props = torch.cuda.get_device_properties(device)
                self._device_support[device] = (
                    props.major >= 8
                    and props.shared_memory_per_block_optin >= 71680
                )
            if self._device_support[device]:
                return _extension.small_k_mm(A, B)
        return torch.matmul(A, B)
```