```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Custom CUDA kernel for GEMM with a small K dimension (K = 64 here).
#
# Because K is tiny, this problem is memory-bandwidth bound (the output
# C is M x N = 4 GB while the inputs are only ~8 MB each). This kernel is
# designed to saturate DRAM write bandwidth on C:
#   * Each thread block handles a 64 (rows) x 256 (cols) tile of C.
#   * The 64 x K tile of A is staged once in shared memory.
#   * B is streamed directly from global/L2 with vectorized float4 loads
#     (the whole B matrix is only ~8 MB, so it stays hot in L2).
#   * Each thread accumulates 8 rows x 4 columns in registers.
#   * C is written with streaming (write-through) float4 stores so the
#     output does not thrash the L2 cache that B relies on.
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BM 64   // rows of C per block
#define TM 8    // rows of C per thread
#define BTX 64  // threads in x (each owns one float4 column -> 256 cols)
#define BTY 8   // threads in y

__global__ void __launch_bounds__(BTX * BTY)
matmul_small_k_kernel(const float* __restrict__ A,
                      const float* __restrict__ B,
                      float* __restrict__ C,
                      int M, int N, int K)
{
    extern __shared__ float4 sA[];  // BM x K tile of A

    const int tx = threadIdx.x;
    const int ty = threadIdx.y;

    const int block_row = blockIdx.y * BM;
    const int n4 = N >> 2;
    const int k4 = K >> 2;
    const int col4 = blockIdx.x * BTX + tx;

    // Stage the (BM x K) tile of A into shared memory with float4 loads.
    const int tid = ty * BTX + tx;
    const int total = BM * k4;
    const float4* A4 = reinterpret_cast<const float4*>(A);
    for (int i = tid; i < total; i += BTX * BTY) {
        const int r = i / k4;
        const int c = i - r * k4;
        const int gr = block_row + r;
        sA[i] = (gr < M) ? A4[gr * k4 + c] : make_float4(0.f, 0.f, 0.f, 0.f);
    }
    __syncthreads();

    if (col4 >= n4) return;

    const float4* B4 = reinterpret_cast<const float4*>(B);
    const float*  sAf = reinterpret_cast<const float*>(sA);

    float4 acc[TM];
    #pragma unroll
    for (int i = 0; i < TM; ++i) acc[i] = make_float4(0.f, 0.f, 0.f, 0.f);

    #pragma unroll 4
    for (int k = 0; k < K; ++k) {
        const float4 b = __ldg(B4 + (long long)k * n4 + col4);
        #pragma unroll
        for (int i = 0; i < TM; ++i) {
            const float a = sAf[(ty + i * BTY) * K + k];
            acc[i].x = fmaf(a, b.x, acc[i].x);
            acc[i].y = fmaf(a, b.y, acc[i].y);
            acc[i].z = fmaf(a, b.z, acc[i].z);
            acc[i].w = fmaf(a, b.w, acc[i].w);
        }
    }

    float4* C4 = reinterpret_cast<float4*>(C);
    #pragma unroll
    for (int i = 0; i < TM; ++i) {
        const int r = block_row + ty + i * BTY;
        if (r < M) {
            // Streaming store: C is never re-read, keep L2 free for B.
            __stwt(C4 + (long long)r * n4 + col4, acc[i]);
        }
    }
}

torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2, "A and B must be 2-D");

    if (A.scalar_type() != at::kFloat || B.scalar_type() != at::kFloat) {
        return torch::matmul(A, B);
    }

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();

    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);
    const int64_t N = Bc.size(1);

    // Fast path requires 16-byte vectorizable rows and a small K that
    // fits the shared-memory tile; otherwise defer to cuBLAS.
    if ((K % 4) != 0 || (N % 4) != 0 || K > 192) {
        return torch::matmul(Ac, Bc);
    }

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

    dim3 block(BTX, BTY);
    dim3 grid((unsigned)((N / 4 + BTX - 1) / BTX),
              (unsigned)((M + BM - 1) / BM));
    const size_t shmem = (size_t)BM * K * sizeof(float);

    auto stream = at::cuda::getCurrentCUDAStream();
    matmul_small_k_kernel<<<grid, block, shmem, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
        (int)M, (int)N, (int)K);

    return C;
}
"""

cpp_source = "torch::Tensor matmul_cuda(torch::Tensor A, torch::Tensor B);"

matmul_ext = load_inline(
    name="matmul_small_k_ext",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = A * B for small-K GEMMs using a custom
    CUDA kernel tuned to saturate output write bandwidth.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._ext = matmul_ext

    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.dim() == 2 and B.dim() == 2):
            return self._ext.matmul_cuda(A, B)
        return torch.matmul(A, B)
```