```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 kernel for matrix-vector multiplication (GEMV): C = A @ B
# A: (M, K), B: (K, 1), C: (M, 1)
# The operation is memory-bandwidth bound, so the kernel is designed to
# saturate DRAM bandwidth:
#   - One thread block per output row with fully coalesced 128-bit (float4)
#     vectorized loads.
#   - Streaming loads (__ldcs) for A (read once, avoids L2 thrashing) and
#     cached read-only loads (__ldg) for B (reused by every block).
#   - Unrolled main loop for high memory-level parallelism.
#   - Fused multiply-add accumulation + warp-shuffle / shared-memory
#     block reduction, entirely fused in a single kernel (no global
#     intermediate reduction buffer needed).
gemv_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define BLOCK_THREADS 256
#define WARP_SIZE 32
#define NUM_WARPS (BLOCK_THREADS / WARP_SIZE)

__device__ __forceinline__ void block_reduce_store(float sum, float* __restrict__ out, int row) {
    __shared__ float warp_sums[NUM_WARPS];
    const int lane = threadIdx.x & (WARP_SIZE - 1);
    const int warp = threadIdx.x >> 5;

    #pragma unroll
    for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, offset);
    }
    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) out[row] = sum;
    }
}

// Vectorized kernel (requires K % 4 == 0). One block per row.
__global__ void __launch_bounds__(BLOCK_THREADS)
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* __restrict__ a_row = A + (size_t)row * (size_t)K4;

    float sum = 0.0f;
    int i = threadIdx.x;

    // Unrolled main loop: 4 float4 loads per thread per iteration.
    for (; i + 3 * BLOCK_THREADS < K4; i += 4 * BLOCK_THREADS) {
        float4 a0 = __ldcs(a_row + i);
        float4 a1 = __ldcs(a_row + i + BLOCK_THREADS);
        float4 a2 = __ldcs(a_row + i + 2 * BLOCK_THREADS);
        float4 a3 = __ldcs(a_row + i + 3 * BLOCK_THREADS);
        float4 b0 = __ldg(B + i);
        float4 b1 = __ldg(B + i + BLOCK_THREADS);
        float4 b2 = __ldg(B + i + 2 * BLOCK_THREADS);
        float4 b3 = __ldg(B + i + 3 * BLOCK_THREADS);
        sum = fmaf(a0.x, b0.x, sum); sum = fmaf(a0.y, b0.y, sum);
        sum = fmaf(a0.z, b0.z, sum); sum = fmaf(a0.w, b0.w, sum);
        sum = fmaf(a1.x, b1.x, sum); sum = fmaf(a1.y, b1.y, sum);
        sum = fmaf(a1.z, b1.z, sum); sum = fmaf(a1.w, b1.w, sum);
        sum = fmaf(a2.x, b2.x, sum); sum = fmaf(a2.y, b2.y, sum);
        sum = fmaf(a2.z, b2.z, sum); sum = fmaf(a2.w, b2.w, sum);
        sum = fmaf(a3.x, b3.x, sum); sum = fmaf(a3.y, b3.y, sum);
        sum = fmaf(a3.z, b3.z, sum); sum = fmaf(a3.w, b3.w, sum);
    }
    for (; i < K4; i += BLOCK_THREADS) {
        float4 a = __ldcs(a_row + i);
        float4 b = __ldg(B + i);
        sum = fmaf(a.x, b.x, sum);
        sum = fmaf(a.y, b.y, sum);
        sum = fmaf(a.z, b.z, sum);
        sum = fmaf(a.w, b.w, sum);
    }

    block_reduce_store(sum, C, row);
}

// Scalar fallback kernel (works for any K). One block per row.
__global__ void __launch_bounds__(BLOCK_THREADS)
gemv_scalar_kernel(const float* __restrict__ A,
                   const float* __restrict__ B,
                   float* __restrict__ C,
                   int M, int K)
{
    const int row = blockIdx.x;
    const float* __restrict__ a_row = A + (size_t)row * (size_t)K;

    float sum = 0.0f;
    for (int i = threadIdx.x; i < K; i += BLOCK_THREADS) {
        sum = fmaf(__ldcs(a_row + i), __ldg(B + i), sum);
    }

    block_reduce_store(sum, C, row);
}

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 && B.dim() == 2 && B.size(1) == 1,
                "A must be (M, K) and B must be (K, 1)");
    TORCH_CHECK(A.size(1) == B.size(0), "Inner dimensions must match");

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

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

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

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

    if ((K & 3) == 0) {
        const int K4 = (int)(K >> 2);
        gemv_vec4_kernel<<<(unsigned int)M, BLOCK_THREADS, 0, stream>>>(
            reinterpret_cast<const float4*>(A_c.data_ptr<float>()),
            reinterpret_cast<const float4*>(B_c.data_ptr<float>()),
            C.data_ptr<float>(), (int)M, K4);
    } else {
        gemv_scalar_kernel<<<(unsigned int)M, BLOCK_THREADS, 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);"

# Compile the inline CUDA GEMV extension
gemv_ext = load_inline(
    name="gemv_ext",
    cpp_sources=gemv_cpp_source,
    cuda_sources=gemv_cuda_source,
    functions=["gemv_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs matrix-vector multiplication (C = A * B)
    with a custom fused 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):
            return self.gemv.gemv_cuda(A, B)
        return torch.matmul(A, B)
```