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

template<bool FULL_TILES>
__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 row_group = tid >> 5;
    const int64_t block_m = int64_t(blockIdx.y) * BM;
    const int64_t block_n = int64_t(blockIdx.x) * BN;

    float accum[8][4] = {};

    for (int64_t k0 = 0; k0 < K; k0 += BK) {
        #pragma unroll
        for (int i = tid; i < BM * BK; i += 256) {
            const int row = i / BK;
            const int k = i % BK;
            if (FULL_TILES || (block_m + row < M && k0 + k < K)) {
                As[k][row] = A[(block_m + row) * K + k0 + k];
            } else {
                As[k][row] = 0.0f;
            }
        }

        #pragma unroll
        for (int i = tid; i < BK * BN; i += 256) {
            const int k = i / BN;
            const int col = i % BN;
            if (FULL_TILES || (k0 + k < K && block_n + col < N)) {
                Bs[k][col] = B[(k0 + k) * N + block_n + col];
            } else {
                Bs[k][col] = 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][row_group * 8 + i];

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

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        const int64_t row = block_m + row_group * 8 + i;

        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            const int64_t col = block_n + lane + j * 32;
            if (FULL_TILES || (row < M && col < N)) {
                C[row * N + col] = accum[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.scalar_type() == at::kFloat &&
                B.scalar_type() == at::kFloat, "Inputs must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "Inputs must be matrices");
    TORCH_CHECK(A.size(1) == B.size(0), "Incompatible matrix dimensions");
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Inputs must be contiguous");

    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;
    }

    TORCH_CHECK((M + BM - 1) / BM <= 65535,
                "Matrix has too many rows for this kernel");

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

    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (M % BM == 0 && N % BN == 0 && K % BK == 0) {
        small_k_matmul_kernel<true><<<grid, 256, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(), C.data_ptr<float>(),
            M, N, K);
    } else {
        small_k_matmul_kernel<false><<<grid, 256, 0, stream>>>(
            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_fp32_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.ndim == 2
            and B.ndim == 2
            and A.shape[1] == B.shape[0]
            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)
```