```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.
# The operation C = A @ B with A:(M,K), B:(K,1) is purely memory-bandwidth
# bound (the whole matrix A must be streamed exactly once). This kernel is
# optimized accordingly:
#   * one thread block per output row
#   * 128-bit vectorized (float4) loads with 4x unrolling for high MLP
#   * multiple accumulators to hide FMA latency
#   * warp-shuffle based block reduction (no shared-memory atomics)
#   * B is tiny and stays hot in L2 across blocks
matvec_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 256
#define NUM_WARPS (BLOCK_SIZE / 32)

__global__ void matvec_vec4_kernel(const float* __restrict__ A,
                                   const float* __restrict__ B,
                                   float* __restrict__ C,
                                   long long K) {
    const long long row = blockIdx.x;
    const float4* __restrict__ a4 = reinterpret_cast<const float4*>(A + row * K);
    const float4* __restrict__ b4 = reinterpret_cast<const float4*>(B);
    const long long K4 = K >> 2;
    const long long stride = blockDim.x;

    float sum0 = 0.0f, sum1 = 0.0f, sum2 = 0.0f, sum3 = 0.0f;

    long long i = threadIdx.x;
    // 4x unrolled main loop -> many outstanding 128-bit loads per thread.
    for (; i + 3 * stride < K4; i += 4 * stride) {
        float4 a0 = a4[i];
        float4 a1 = a4[i + stride];
        float4 a2 = a4[i + 2 * stride];
        float4 a3 = a4[i + 3 * stride];
        float4 b0 = b4[i];
        float4 b1 = b4[i + stride];
        float4 b2 = b4[i + 2 * stride];
        float4 b3 = b4[i + 3 * stride];
        sum0 += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
        sum1 += a1.x * b1.x + a1.y * b1.y + a1.z * b1.z + a1.w * b1.w;
        sum2 += a2.x * b2.x + a2.y * b2.y + a2.z * b2.z + a2.w * b2.w;
        sum3 += a3.x * b3.x + a3.y * b3.y + a3.z * b3.z + a3.w * b3.w;
    }
    for (; i < K4; i += stride) {
        float4 a = a4[i];
        float4 b = b4[i];
        sum0 += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
    }

    float sum = (sum0 + sum1) + (sum2 + sum3);

    // Intra-warp reduction.
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, offset);

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

    // Final reduction by the first warp.
    if (warp == 0) {
        sum = (lane < NUM_WARPS) ? warp_sums[lane] : 0.0f;
        #pragma unroll
        for (int offset = NUM_WARPS / 2; offset > 0; offset >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, offset);
        if (lane == 0) C[row] = sum;
    }
}

// Scalar fallback for the (unlikely) case K is not a multiple of 4.
__global__ void matvec_scalar_kernel(const float* __restrict__ A,
                                     const float* __restrict__ B,
                                     float* __restrict__ C,
                                     long long K) {
    const long long row = blockIdx.x;
    const float* __restrict__ a = A + row * K;

    float sum = 0.0f;
    for (long long i = threadIdx.x; i < K; i += blockDim.x)
        sum += a[i] * B[i];

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

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

    if (warp == 0) {
        sum = (lane < NUM_WARPS) ? warp_sums[lane] : 0.0f;
        #pragma unroll
        for (int offset = NUM_WARPS / 2; offset > 0; offset >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, offset);
        if (lane == 0) C[row] = sum;
    }
}

torch::Tensor matvec_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() == at::kFloat && B.scalar_type() == at::kFloat,
                "Inputs must be float32");
    TORCH_CHECK(A.dim() == 2, "A must be 2-D");

    auto Ac = A.contiguous();
    auto Bc = B.contiguous().view({-1});

    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);
    TORCH_CHECK(Bc.numel() == K, "B must have K elements");

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

    const dim3 grid(static_cast<unsigned int>(M));
    const dim3 block(BLOCK_SIZE);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if ((K & 3) == 0) {
        matvec_vec4_kernel<<<grid, block, 0, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
            static_cast<long long>(K));
    } else {
        matvec_scalar_kernel<<<grid, block, 0, stream>>>(
            Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(),
            static_cast<long long>(K));
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return C;
}
"""

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

_matvec_ext = load_inline(
    name="matvec_ext",
    cpp_sources=matvec_cpp_source,
    cuda_sources=matvec_cuda_source,
    functions=["matvec_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    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._ext = _matvec_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
            and B.size(1) == 1
        ):
            return self._ext.matvec_cuda(A, B)
        # Safe fallback for any configuration the custom kernel does not cover.
        return torch.matmul(A, B)
```