```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), C: (M,1)).
# This operation is purely memory-bandwidth bound (A is ~8GB), so the kernel
# is optimized for maximal coalesced throughput:
#   - one thread block per output row
#   - float4 vectorized loads (16B per thread per load)
#   - 4x ILP unrolling for memory-level parallelism
#   - warp-shuffle based block reduction
gemv_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void gemv_vec4_kernel(const float4* __restrict__ A,
                                 const float4* __restrict__ B,
                                 float* __restrict__ C,
                                 int M, int K4) {
    const int row = blockIdx.x;
    const float4* Arow = A + (size_t)row * (size_t)K4;

    float sum = 0.0f;
    int i = threadIdx.x;
    const int stride = blockDim.x;

    // 4x unrolled main loop for memory-level parallelism
    for (; i + 3 * stride < K4; i += 4 * stride) {
        float4 a0 = Arow[i];
        float4 a1 = Arow[i + stride];
        float4 a2 = Arow[i + 2 * stride];
        float4 a3 = Arow[i + 3 * stride];
        float4 b0 = B[i];
        float4 b1 = B[i + stride];
        float4 b2 = B[i + 2 * stride];
        float4 b3 = B[i + 3 * stride];
        sum += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
        sum += a1.x * b1.x + a1.y * b1.y + a1.z * b1.z + a1.w * b1.w;
        sum += a2.x * b2.x + a2.y * b2.y + a2.z * b2.z + a2.w * b2.w;
        sum += a3.x * b3.x + a3.y * b3.y + a3.z * b3.z + a3.w * b3.w;
    }
    // remainder
    for (; i < K4; i += stride) {
        float4 a = Arow[i];
        float4 b = B[i];
        sum += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
    }

    // warp-level reduction
    for (int offset = warpSize / 2; offset > 0; offset >>= 1)
        sum += __shfl_down_sync(0xffffffff, sum, offset);

    __shared__ float warp_sums[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) warp_sums[wid] = sum;
    __syncthreads();

    if (wid == 0) {
        const int nwarps = (blockDim.x + 31) >> 5;
        sum = (lane < nwarps) ? warp_sums[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            sum += __shfl_down_sync(0xffffffff, sum, offset);
        if (lane == 0) C[row] = sum;
    }
}

torch::Tensor gemv_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 &&
                B.scalar_type() == torch::kFloat32, "Inputs must be float32");

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

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

    auto C = torch::empty({M, 1}, Ac.options());

    if (K % 4 == 0 && K > 0 && M > 0) {
        const int K4 = (int)(K / 4);
        const int block = 256;
        const dim3 grid((unsigned int)M);
        auto stream = at::cuda::getCurrentCUDAStream();
        gemv_vec4_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>(),
            (int)M, K4);
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    } else {
        C = torch::matmul(Ac, Bc);
    }
    return C;
}
"""

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

gemv_ext = load_inline(
    name="gemv_cuda_ext",
    cpp_sources=gemv_cpp_source,
    cuda_sources=gemv_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 bandwidth-optimized CUDA GEMV kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.gemv = gemv_ext

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if not A.is_cuda:
            A = A.cuda()
            B = B.cuda()
        return self.gemv.gemv_cuda(A, B)
```