```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 optimized for large-K matrix-vector multiplication.
# The operation is heavily memory-bound (must stream the entire M x K matrix),
# so the kernel is designed to saturate DRAM bandwidth:
#   - one thread block per output row (full coalescing across the block)
#   - float4 vectorized loads (16B per thread per load)
#   - 4x unrolling for memory-level parallelism
#   - streaming loads (__ldcs) for A (no reuse) and read-only loads (__ldg)
#     for B (reused across all rows, stays resident in cache)
#   - warp-shuffle + shared-memory block reduction
cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#define FULL_MASK 0xffffffffu

__inline__ __device__ float warpReduceSum(float val) {
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_down_sync(FULL_MASK, val, offset);
    return val;
}

template <int BLOCK>
__global__ void __launch_bounds__(BLOCK)
gemv_vec4_kernel(const float4* __restrict__ A,
                 const float4* __restrict__ B,
                 float* __restrict__ C,
                 int K4) {
    const int row = blockIdx.x;
    const float4* __restrict__ a = A + (size_t)row * (size_t)K4;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;

    int i = threadIdx.x;
    // 4x unrolled main loop: 4 independent load/dot chains for high MLP.
    for (; i + 3 * BLOCK < K4; i += 4 * BLOCK) {
        float4 a0 = __ldcs(a + i);
        float4 a1 = __ldcs(a + i + BLOCK);
        float4 a2 = __ldcs(a + i + 2 * BLOCK);
        float4 a3 = __ldcs(a + i + 3 * BLOCK);
        float4 b0 = __ldg(B + i);
        float4 b1 = __ldg(B + i + BLOCK);
        float4 b2 = __ldg(B + i + 2 * BLOCK);
        float4 b3 = __ldg(B + i + 3 * BLOCK);
        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);
    }
    // Tail loop.
    for (; i < K4; i += BLOCK) {
        float4 av = __ldcs(a + i);
        float4 bv = __ldg(B + i);
        s0 += (av.x * bv.x + av.y * bv.y) + (av.z * bv.z + av.w * bv.w);
    }

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

    // Block reduction.
    __shared__ float shared[BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    sum = warpReduceSum(sum);
    if (lane == 0) shared[wid] = sum;
    __syncthreads();
    if (wid == 0) {
        sum = (lane < BLOCK / 32) ? shared[lane] : 0.f;
        sum = warpReduceSum(sum);
        if (lane == 0) C[row] = sum;
    }
}

// Scalar fallback for shapes/alignments that don't support float4.
template <int BLOCK>
__global__ void __launch_bounds__(BLOCK)
gemv_scalar_kernel(const float* __restrict__ A,
                   const float* __restrict__ B,
                   float* __restrict__ C,
                   int64_t K) {
    const int64_t row = blockIdx.x;
    const float* __restrict__ a = A + row * K;

    float sum = 0.f;
    for (int64_t i = threadIdx.x; i < K; i += BLOCK) {
        sum += __ldcs(a + i) * __ldg(B + i);
    }

    __shared__ float shared[BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    sum = warpReduceSum(sum);
    if (lane == 0) shared[wid] = sum;
    __syncthreads();
    if (wid == 0) {
        sum = (lane < BLOCK / 32) ? shared[lane] : 0.f;
        sum = warpReduceSum(sum);
        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() == 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();

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

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

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

    constexpr int BLOCK = 256;
    const float* a_ptr = Ac.data_ptr<float>();
    const float* b_ptr = Bc.data_ptr<float>();
    float* c_ptr = C.data_ptr<float>();

    const bool vec4_ok =
        ((K % 4) == 0) &&
        ((reinterpret_cast<uintptr_t>(a_ptr) & 0xF) == 0) &&
        ((reinterpret_cast<uintptr_t>(b_ptr) & 0xF) == 0);

    if (vec4_ok) {
        const int K4 = static_cast<int>(K / 4);
        gemv_vec4_kernel<BLOCK><<<(unsigned int)M, BLOCK, 0, stream>>>(
            reinterpret_cast<const float4*>(a_ptr),
            reinterpret_cast<const float4*>(b_ptr),
            c_ptr, K4);
    } else {
        gemv_scalar_kernel<BLOCK><<<(unsigned int)M, BLOCK, 0, stream>>>(
            a_ptr, b_ptr, c_ptr, K);
    }
    return C;
}
"""

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

gemv_module = load_inline(
    name="custom_gemv_matmul",
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=["gemv_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-saturating 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
                and A.dim() == 2 and B.dim() == 2
                and B.size(0) == A.size(1) and B.size(1) == 1):
            return self.gemv.gemv_cuda(A, B)
        return torch.matmul(A, B)
```