```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 GEMV kernel (C = A @ B, A: (M, K), B: (K, 1)).
# The operation is purely memory-bandwidth bound (A is streamed once,
# B fits in L2 and is reused across rows), so the kernel is designed to:
#   * use one thread block per row with fully coalesced 128-bit (float4) loads
#   * keep many loads in flight via 4x unrolling and independent accumulators
#   * use streaming loads (__ldcs) for A so it does not thrash L2 (B is cached)
#   * use a fast warp-shuffle + shared-memory block reduction
gemv_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 256

__device__ __forceinline__ float block_reduce_sum(float val) {
    __shared__ float smem[BLOCK_SIZE / 32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

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

    if (lane == 0) smem[warp] = val;
    __syncthreads();

    float sum = 0.f;
    if (warp == 0) {
        sum = (lane < BLOCK_SIZE / 32) ? smem[lane] : 0.f;
        #pragma unroll
        for (int offset = (BLOCK_SIZE / 32) / 2; offset > 0; offset >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, offset);
    }
    return sum;
}

// Vectorized GEMV, one block per row. Requires K % 4 == 0.
__global__ void __launch_bounds__(BLOCK_SIZE) gemv_f32_vec4_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    const int M,
    const int K)
{
    const int row = blockIdx.x;
    const int K4 = K >> 2;
    const float4* __restrict__ A4 =
        reinterpret_cast<const float4*>(A) + (size_t)row * (size_t)K4;
    const float4* __restrict__ B4 = reinterpret_cast<const float4*>(B);

    float acc0 = 0.f, acc1 = 0.f, acc2 = 0.f, acc3 = 0.f;

    int i = threadIdx.x;
    // 4x unrolled main loop: keeps 8 outstanding 128-bit loads per thread.
    for (; i + 3 * BLOCK_SIZE < K4; i += 4 * BLOCK_SIZE) {
        const float4 a0 = __ldcs(&A4[i]);
        const float4 a1 = __ldcs(&A4[i + BLOCK_SIZE]);
        const float4 a2 = __ldcs(&A4[i + 2 * BLOCK_SIZE]);
        const float4 a3 = __ldcs(&A4[i + 3 * BLOCK_SIZE]);
        const float4 b0 = __ldg(&B4[i]);
        const float4 b1 = __ldg(&B4[i + BLOCK_SIZE]);
        const float4 b2 = __ldg(&B4[i + 2 * BLOCK_SIZE]);
        const float4 b3 = __ldg(&B4[i + 3 * BLOCK_SIZE]);

        acc0 += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
        acc1 += a1.x * b1.x + a1.y * b1.y + a1.z * b1.z + a1.w * b1.w;
        acc2 += a2.x * b2.x + a2.y * b2.y + a2.z * b2.z + a2.w * b2.w;
        acc3 += a3.x * b3.x + a3.y * b3.y + a3.z * b3.z + a3.w * b3.w;
    }
    // Tail (not taken when K4 is a multiple of 4 * BLOCK_SIZE)
    for (; i < K4; i += BLOCK_SIZE) {
        const float4 a = __ldcs(&A4[i]);
        const float4 b = __ldg(&B4[i]);
        acc0 += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
    }

    float sum = (acc0 + acc1) + (acc2 + acc3);
    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) C[row] = sum;
}

// Scalar fallback for the (unused) case where K % 4 != 0.
__global__ void __launch_bounds__(BLOCK_SIZE) gemv_f32_scalar_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    const int M,
    const int K)
{
    const int row = blockIdx.x;
    const float* __restrict__ Arow = A + (size_t)row * (size_t)K;

    float sum = 0.f;
    for (int i = threadIdx.x; i < K; i += BLOCK_SIZE)
        sum += __ldcs(&Arow[i]) * __ldg(&B[i]);

    sum = block_reduce_sum(sum);
    if (threadIdx.x == 0) C[row] = sum;
}

torch::Tensor gemv_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, "A must be 2D");

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

    auto A_c = A.contiguous();
    auto B_c = B.contiguous();

    auto C = torch::empty({M, 1}, A.options());
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((K & 3) == 0) {
        gemv_f32_vec4_kernel<<<(int)M, BLOCK_SIZE, 0, stream>>>(
            A_c.data_ptr<float>(), B_c.data_ptr<float>(), C.data_ptr<float>(),
            (int)M, (int)K);
    } else {
        gemv_f32_scalar_kernel<<<(int)M, BLOCK_SIZE, 0, stream>>>(
            A_c.data_ptr<float>(), B_c.data_ptr<float>(), C.data_ptr<float>(),
            (int)M, (int)K);
    }
    return C;
}
"""

gemv_cpp_source = "torch::Tensor gemv_cuda(torch::Tensor A, torch::Tensor B);"

gemv_module = load_inline(
    name="gemv_cuda_ext",
    cpp_sources=gemv_cpp_source,
    cuda_sources=gemv_cuda_source,
    functions=["gemv_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs matrix-vector multiplication (C = A * B)
    using a custom fused CUDA GEMV kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.gemv = gemv_module

    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):
            return self.gemv.gemv_cuda(A, B)
        return torch.matmul(A, B)
```