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

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

__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,
    int64_t K)
{
    __shared__ float As[BK][BM + 1];
    __shared__ float Bs[BK][BN];

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

    const int64_t row_base = int64_t(blockIdx.y) * BM;
    const int64_t col_base = int64_t(blockIdx.x) * BN;

    float acc[8][4];

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

    for (int64_t kb = 0; kb < K; kb += BK) {
        #pragma unroll
        for (int offset = 0; offset < BM * BK; offset += 256) {
            const int index = tid + offset;
            const int r = index / BK;
            const int k = index % BK;
            const int64_t global_r = row_base + r;
            const int64_t global_k = kb + k;

            As[k][r] = (global_r < M && global_k < K)
                ? A[global_r * K + global_k]
                : 0.0f;
        }

        #pragma unroll
        for (int offset = 0; offset < BK * BN; offset += 256) {
            const int index = tid + offset;
            const int k = index / BN;
            const int c = index % BN;
            const int64_t global_k = kb + k;
            const int64_t global_c = col_base + c;

            Bs[k][c] = (global_k < K && global_c < N)
                ? B[global_k * N + global_c]
                : 0.0f;
        }

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            float b[4];

            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                b[j] = Bs[k][lane + j * 32];
            }

            #pragma unroll
            for (int i = 0; i < 8; ++i) {
                const float a = As[k][warp + i * 8];

                #pragma unroll
                for (int j = 0; j < 4; ++j) {
                    acc[i][j] = fmaf(a, b[j], acc[i][j]);
                }
            }
        }

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        const int64_t row = row_base + warp + i * 8;
        if (row < M) {
            #pragma unroll
            for (int j = 0; j < 4; ++j) {
                const int64_t col = col_base + lane + j * 32;
                if (col < N) {
                    C[row * N + col] = acc[i][j];
                }
            }
        }
    }
}

torch::Tensor small_k_matmul_cuda(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Inputs must share a device");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be matrices");
    TORCH_CHECK(A.scalar_type() == at::kFloat &&
                B.scalar_type() == at::kFloat, "Inputs must be float32");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Inputs must be contiguous");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible matrix dimensions");

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

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

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

    const dim3 grid(
        static_cast<unsigned int>((N + BN - 1) / BN),
        static_cast<unsigned int>((M + BM - 1) / BM)
    );

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

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_small_k_matmul = load_inline(
    name="small_k_matmul_tiled_fp32_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.ndim == 2
            and B.ndim == 2
            and A.is_contiguous()
            and B.is_contiguous()
            and A.shape[0] <= 64 * 65535
            and not (
                torch.is_grad_enabled()
                and (A.requires_grad or B.requires_grad)
            )
        ):
            return _small_k_matmul.small_k_matmul_cuda(A, B)
        return torch.matmul(A, B)
```