```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(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 TILE = 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,
    int M, int N, int K)
{
    __shared__ float As[BK][TILE];
    __shared__ __align__(16) float Bs[BK][TILE];

    const int tid = threadIdx.x;
    const int tx = tid & 15;
    const int ty = tid >> 4;
    const int row_base = blockIdx.y * TILE;
    const int col_base = blockIdx.x * TILE;

    float acc[8][8] = {};

    for (int kb = 0; kb < K; kb += BK) {
        #pragma unroll
        for (int q = 0; q < 2; ++q) {
            const int v = tid + q * 256;
            const int ar = v >> 2;
            const int ak = (v & 3) * 4;
            float4 av = make_float4(0.f, 0.f, 0.f, 0.f);

            if (row_base + ar < M) {
                const int64_t offset =
                    int64_t(row_base + ar) * K + kb + ak;
                if ((K & 3) == 0 && kb + ak + 3 < K) {
                    av = *reinterpret_cast<const float4*>(A + offset);
                } else {
                    if (kb + ak     < K) av.x = A[offset];
                    if (kb + ak + 1 < K) av.y = A[offset + 1];
                    if (kb + ak + 2 < K) av.z = A[offset + 2];
                    if (kb + ak + 3 < K) av.w = A[offset + 3];
                }
            }

            As[ak    ][ar] = av.x;
            As[ak + 1][ar] = av.y;
            As[ak + 2][ar] = av.z;
            As[ak + 3][ar] = av.w;

            const int bk = v >> 5;
            const int bc = (v & 31) * 4;
            float4 bv = make_float4(0.f, 0.f, 0.f, 0.f);

            if (kb + bk < K) {
                const int64_t offset =
                    int64_t(kb + bk) * N + col_base + bc;
                if ((N & 3) == 0 && col_base + bc + 3 < N) {
                    bv = *reinterpret_cast<const float4*>(B + offset);
                } else {
                    if (col_base + bc     < N) bv.x = B[offset];
                    if (col_base + bc + 1 < N) bv.y = B[offset + 1];
                    if (col_base + bc + 2 < N) bv.z = B[offset + 2];
                    if (col_base + bc + 3 < N) bv.w = B[offset + 3];
                }
            }

            *reinterpret_cast<float4*>(&Bs[bk][bc]) = bv;
        }

        __syncthreads();

        #pragma unroll
        for (int k = 0; k < BK; ++k) {
            const float4 b0 =
                *reinterpret_cast<const float4*>(&Bs[k][tx * 4]);
            const float4 b1 =
                *reinterpret_cast<const float4*>(&Bs[k][tx * 4 + 64]);

            #pragma unroll
            for (int i = 0; i < 8; ++i) {
                const int r = ty * 4 + (i & 3) + (i >> 2) * 64;
                const float a = As[k][r];
                acc[i][0] = fmaf(a, b0.x, acc[i][0]);
                acc[i][1] = fmaf(a, b0.y, acc[i][1]);
                acc[i][2] = fmaf(a, b0.z, acc[i][2]);
                acc[i][3] = fmaf(a, b0.w, acc[i][3]);
                acc[i][4] = fmaf(a, b1.x, acc[i][4]);
                acc[i][5] = fmaf(a, b1.y, acc[i][5]);
                acc[i][6] = fmaf(a, b1.z, acc[i][6]);
                acc[i][7] = fmaf(a, b1.w, acc[i][7]);
            }
        }

        __syncthreads();
    }

    #pragma unroll
    for (int i = 0; i < 8; ++i) {
        const int row =
            row_base + ty * 4 + (i & 3) + (i >> 2) * 64;

        if (row < M) {
            #pragma unroll
            for (int group = 0; group < 2; ++group) {
                const int col = col_base + tx * 4 + group * 64;
                const int j = group * 4;
                const int64_t offset = int64_t(row) * N + col;

                if ((N & 3) == 0 && col + 3 < N) {
                    *reinterpret_cast<float4*>(C + offset) =
                        make_float4(
                            acc[i][j], acc[i][j + 1],
                            acc[i][j + 2], acc[i][j + 3]);
                } else {
                    #pragma unroll
                    for (int t = 0; t < 4; ++t) {
                        if (col + t < N)
                            C[offset + t] = acc[i][j + t];
                    }
                }
            }
        }
    }
}

torch::Tensor small_k_matmul(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.size(1) == B.size(0), "Incompatible matrix dimensions");
    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");

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

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

    const dim3 grid((N + TILE - 1) / TILE, (M + TILE - 1) / TILE);
    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;
}
"""

_extension = load_inline(
    name="small_k_tiled_matmul_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["small_k_matmul"],
    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.dtype == torch.float32
            and B.dtype == torch.float32
            and A.ndim == 2
            and B.ndim == 2
        ):
            return _extension.small_k_matmul(A.contiguous(), B.contiguous())
        return torch.matmul(A, B)
```