```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 fused GEMV kernel (C = A @ B) optimized for tall-K matrix-vector products.
# Strategy:
#   * One thread block per output row (M blocks) -> massive memory parallelism.
#   * 128-bit vectorized loads (float4) with 4 independent streams per thread to
#     saturate DRAM bandwidth (the op is purely memory-bound).
#   * Streaming (evict-first) loads for A so it does not thrash L2, cached loads
#     for the small B vector which is reused by every block.
#   * Warp-shuffle + shared-memory block reduction, single write per row.
gemv_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 256

__global__ void gemv_kernel(const float* __restrict__ A,
                            const float* __restrict__ B,
                            float* __restrict__ C,
                            int M, int K)
{
    const int row = blockIdx.x;
    const float* __restrict__ Arow = A + (long)row * (long)K;

    const int K4 = K >> 2;
    const float4* __restrict__ A4 = reinterpret_cast<const float4*>(Arow);
    const float4* __restrict__ B4 = reinterpret_cast<const float4*>(B);

    float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f;

    const int stride = BLOCK_SIZE;
    int i = threadIdx.x;

    // Main loop: 4 independent 128-bit load streams per thread.
    for (; i + 3 * stride < K4; i += 4 * stride) {
        const float4 a0 = __ldcs(A4 + i);
        const float4 a1 = __ldcs(A4 + i + stride);
        const float4 a2 = __ldcs(A4 + i + 2 * stride);
        const float4 a3 = __ldcs(A4 + i + 3 * stride);
        const float4 b0 = __ldg(B4 + i);
        const float4 b1 = __ldg(B4 + i + stride);
        const float4 b2 = __ldg(B4 + i + 2 * stride);
        const float4 b3 = __ldg(B4 + i + 3 * stride);
        s0 += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
        s1 += a1.x * b1.x + a1.y * b1.y + a1.z * b1.z + a1.w * b1.w;
        s2 += a2.x * b2.x + a2.y * b2.y + a2.z * b2.z + a2.w * b2.w;
        s3 += a3.x * b3.x + a3.y * b3.y + a3.z * b3.z + a3.w * b3.w;
    }
    // Remaining full float4 chunks.
    for (; i < K4; i += stride) {
        const float4 a0 = __ldcs(A4 + i);
        const float4 b0 = __ldg(B4 + i);
        s0 += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
    }
    // Scalar tail for K not divisible by 4.
    for (int k = (K4 << 2) + threadIdx.x; k < K; k += stride) {
        s0 += Arow[k] * B[k];
    }

    float sum = (s0 + s1) + (s2 + s3);

    // Warp-level reduction.
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, offset);
    }

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

    if (wid == 0) {
        sum = (lane < BLOCK_SIZE / 32) ? sdata[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, 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(), "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");

    const c10::cuda::CUDAGuard device_guard(A.device());

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

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

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

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    gemv_kernel<<<M, BLOCK_SIZE, 0, stream>>>(
        Ac.data_ptr<float>(), Bc.data_ptr<float>(), C.data_ptr<float>(), M, K);

    return C;
}
"""

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

gemv_ext = load_inline(
    name="fast_gemv_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 bandwidth-saturating 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 (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.gemv.gemv_cuda(A, B)
        return torch.matmul(A, B)
```