```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 matrix-vector multiplication C = A @ B.
# This operation is heavily memory-bandwidth bound (A is streamed once),
# so the kernel is designed to saturate DRAM bandwidth:
#   * float4 vectorized loads for A and B (fully coalesced).
#   * Each thread block processes ROWS=4 rows so the B vector is loaded
#     from L2 once and reused for 4 rows (cuts L2 traffic for B by 4x).
#   * __ldcs (streaming, evict-first) loads for A so it does not evict B
#     from cache; __ldg (read-only) loads for the reused B vector.
#   * 4 independent accumulators + unrolled main loop for load ILP.
#   * Warp-shuffle + shared-memory block reduction at the end.
_matvec_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BLOCK 256
#define ROWS 4
#define UNROLL 4

__global__ void __launch_bounds__(BLOCK)
matvec_f4_kernel(const float4* __restrict__ A,
                 const float4* __restrict__ B,
                 float* __restrict__ C,
                 int64_t K4) {
    const int64_t row0 = (int64_t)blockIdx.x * ROWS;
    const float4* aRow = A + row0 * K4;

    float acc[ROWS];
    #pragma unroll
    for (int r = 0; r < ROWS; ++r) acc[r] = 0.0f;

    int64_t i = threadIdx.x;
    const int64_t step = (int64_t)BLOCK * UNROLL;

    for (; i + (UNROLL - 1) * BLOCK < K4; i += step) {
        #pragma unroll
        for (int u = 0; u < UNROLL; ++u) {
            const float4 b = __ldg(B + i + u * BLOCK);
            #pragma unroll
            for (int r = 0; r < ROWS; ++r) {
                const float4 a = __ldcs(aRow + (int64_t)r * K4 + i + u * BLOCK);
                acc[r] += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
            }
        }
    }
    for (; i < K4; i += BLOCK) {
        const float4 b = __ldg(B + i);
        #pragma unroll
        for (int r = 0; r < ROWS; ++r) {
            const float4 a = __ldcs(aRow + (int64_t)r * K4 + i);
            acc[r] += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
        }
    }

    __shared__ float sm[BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;

    #pragma unroll
    for (int r = 0; r < ROWS; ++r) {
        float sum = acc[r];
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, o);
        if (lane == 0) sm[wid] = sum;
        __syncthreads();
        if (wid == 0) {
            float v = (lane < (BLOCK / 32)) ? sm[lane] : 0.0f;
            #pragma unroll
            for (int o = 16; o > 0; o >>= 1)
                v += __shfl_down_sync(0xffffffffu, v, o);
            if (lane == 0) C[row0 + r] = v;
        }
        __syncthreads();
    }
}

// Generic fallback (any M, any K): one block per row, scalar loads.
__global__ void __launch_bounds__(BLOCK)
matvec_generic_kernel(const float* __restrict__ A,
                      const float* __restrict__ B,
                      float* __restrict__ C,
                      int64_t K) {
    const int64_t row = blockIdx.x;
    const float* aRow = A + row * K;

    float sum = 0.0f;
    for (int64_t i = threadIdx.x; i < K; i += BLOCK)
        sum += __ldg(aRow + i) * __ldg(B + i);

    #pragma unroll
    for (int o = 16; o > 0; o >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, o);

    __shared__ float sm[BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) sm[wid] = sum;
    __syncthreads();
    if (wid == 0) {
        float v = (lane < (BLOCK / 32)) ? sm[lane] : 0.0f;
        #pragma unroll
        for (int o = 16; o > 0; o >>= 1)
            v += __shfl_down_sync(0xffffffffu, v, o);
        if (lane == 0) C[row] = v;
    }
}

torch::Tensor matvec_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.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");
    TORCH_CHECK(A.dim() == 2 && B.dim() == 2 && B.size(1) == 1 &&
                A.size(1) == B.size(0), "expecting A (M,K) and B (K,1)");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();
    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);

    auto C = at::empty({M, 1}, Ac.options());
    if (M == 0) return C;

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

    if ((K & 3) == 0 && (M & 3) == 0) {
        const int64_t K4 = K >> 2;
        const int grid = (int)(M / ROWS);
        matvec_f4_kernel<<<grid, BLOCK, 0, stream>>>(
            reinterpret_cast<const float4*>(Ac.data_ptr<float>()),
            reinterpret_cast<const float4*>(Bc.data_ptr<float>()),
            C.data_ptr<float>(), K4);
    } else {
        matvec_generic_kernel<<<(unsigned int)M, BLOCK, 0, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(),
            C.data_ptr<float>(), K);
    }
    return C;
}
"""

_matvec_cpp_src = "torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B);"

_matvec_ext = load_inline(
    name="fast_matvec_ext",
    cpp_sources=_matvec_cpp_src,
    cuda_sources=_matvec_cuda_src,
    functions=["matvec_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Matrix-vector multiplication accelerated with a custom bandwidth-
    saturating CUDA kernel (falls back to torch.matmul when the inputs
    are not CUDA float32 tensors of shape (M,K) @ (K,1)).
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._matvec = _matvec_ext.matvec_cuda

    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
                and B.size(1) == 1 and A.size(1) == B.size(0)):
            return self._matvec(A, B)
        return torch.matmul(A, B)
```