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

__global__ __launch_bounds__(256)
void small_k_matmul_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    int64_t M,
    int64_t N)
{
    constexpr int BM = 64;
    constexpr int BN = 128;
    constexpr int K = 64;

    __shared__ float4 shared_a[BM][K / 4];
    __shared__ float4 shared_b[K][BN / 4];

    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;

    const int64_t row_base = static_cast<int64_t>(blockIdx.y) * BM;
    const int64_t col_base = static_cast<int64_t>(blockIdx.x) * BN;
    const int64_t n4 = N / 4;

    const float4* A4 = reinterpret_cast<const float4*>(A);
    const float4* B4 = reinterpret_cast<const float4*>(B);
    float4* C4 = reinterpret_cast<float4*>(C);

    #pragma unroll
    for (int index = tid; index < BM * (K / 4); index += 256) {
        const int row = index / (K / 4);
        const int col4 = index % (K / 4);
        shared_a[row][col4] =
            row_base + row < M
                ? A4[(row_base + row) * (K / 4) + col4]
                : make_float4(0.f, 0.f, 0.f, 0.f);
    }

    #pragma unroll
    for (int index = tid; index < K * (BN / 4); index += 256) {
        const int k = index / (BN / 4);
        const int col4 = index % (BN / 4);
        const int64_t global_col4 = col_base / 4 + col4;
        shared_b[k][col4] =
            global_col4 < n4
                ? B4[static_cast<int64_t>(k) * n4 + global_col4]
                : make_float4(0.f, 0.f, 0.f, 0.f);
    }

    __syncthreads();

    float4 accum[8];

    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        accum[i] = make_float4(0.f, 0.f, 0.f, 0.f);
    }

    const float* tile_a = reinterpret_cast<const float*>(shared_a);

    #pragma unroll 4
    for (int k = 0; k < K; ++k) {
        const float4 b = shared_b[k][lane];

        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            const float a = tile_a[(warp * 8 + i) * K + k];
            accum[i].x = fmaf(a, b.x, accum[i].x);
            accum[i].y = fmaf(a, b.y, accum[i].y);
            accum[i].z = fmaf(a, b.z, accum[i].z);
            accum[i].w = fmaf(a, b.w, accum[i].w);
        }
    }

    const int64_t output_col4 = col_base / 4 + lane;

    if (output_col4 < n4) {
        #pragma unroll
        for (int i = 0; i < 8; ++i) {
            const int64_t row = row_base + warp * 8 + i;
            if (row < M) {
                C4[row * n4 + output_col4] = accum[i];
            }
        }
    }
}

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(), "Expected matching devices");
    TORCH_CHECK(
        A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
        "Expected float32 tensors");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Expected matrices");
    TORCH_CHECK(A.size(1) == 64 && B.size(0) == 64, "Expected K = 64");
    TORCH_CHECK(B.size(1) % 4 == 0, "Expected N divisible by four");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Expected contiguous tensors");

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

    const int64_t M = A.size(0);
    const int64_t N = B.size(1);
    auto C = torch::empty({M, N}, A.options());

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

    TORCH_CHECK((M + 63) / 64 <= 65535, "M exceeds grid capacity");
    TORCH_CHECK((N + 127) / 128 <= 2147483647LL,
                "N exceeds grid capacity");

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

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

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_small_k_extension = load_inline(
    name="small_k64_matmul_tiled_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 (
            A.is_cuda
            and B.is_cuda
            and A.device == B.device
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.dim() == 2
            and B.dim() == 2
            and A.shape[1] == 64
            and B.shape[0] == 64
            and B.shape[1] % 4 == 0
            and A.shape[0] <= 64 * 65535
        ):
            return _small_k_extension.small_k_matmul_cuda(
                A.contiguous(), B.contiguous()
            )
        return torch.matmul(A, B)
```