```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(M,1) = A(M,K) * B(K,1)
# Optimizations:
#   - float4 vectorized streaming loads (__ldcs) for A to maximize memory bandwidth
#   - split-K across blocks for high occupancy / load balance
#   - 4x unrolled main loop with 4 independent accumulators (breaks FMA dependency chains)
#   - warp-shuffle + shared-memory block reduction
#   - single atomicAdd per block into the output
gemv_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

#define THREADS 256
#define WARPS (THREADS / 32)

__global__ void gemv_f32_vec4_kernel(
    const float4* __restrict__ A,
    const float4* __restrict__ B,
    float* __restrict__ out,
    int K4,
    int chunk4)
{
    const int row   = blockIdx.x;
    const int start = blockIdx.y * chunk4;
    const int end   = min(start + chunk4, K4);

    const float4* __restrict__ a_row = A + (size_t)row * (size_t)K4;

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

    int i = start + threadIdx.x;
    // 4x unrolled main loop
    for (; i + 3 * THREADS < end; i += 4 * THREADS) {
        float4 a0 = __ldcs(&a_row[i]);
        float4 a1 = __ldcs(&a_row[i + THREADS]);
        float4 a2 = __ldcs(&a_row[i + 2 * THREADS]);
        float4 a3 = __ldcs(&a_row[i + 3 * THREADS]);
        float4 b0 = B[i];
        float4 b1 = B[i + THREADS];
        float4 b2 = B[i + 2 * THREADS];
        float4 b3 = B[i + 3 * THREADS];
        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
    for (; i < end; i += THREADS) {
        float4 a = __ldcs(&a_row[i]);
        float4 b = B[i];
        s0 += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
    }

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

    // warp reduction
    for (int off = 16; off > 0; off >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, off);

    __shared__ float warp_sums[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 < WARPS) ? warp_sums[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, off);
        if (lane == 0 && sum != 0.0f)
            atomicAdd(&out[row], sum);
    }
}

// scalar fallback for K not divisible by 4
__global__ void gemv_f32_scalar_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ out,
    int K,
    int chunk)
{
    const int row   = blockIdx.x;
    const int start = blockIdx.y * chunk;
    const int end   = min(start + chunk, K);

    const float* __restrict__ a_row = A + (size_t)row * (size_t)K;

    float sum = 0.0f;
    for (int i = start + threadIdx.x; i < end; i += THREADS)
        sum += __ldcs(&a_row[i]) * B[i];

    for (int off = 16; off > 0; off >>= 1)
        sum += __shfl_down_sync(0xffffffffu, sum, off);

    __shared__ float warp_sums[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 < WARPS) ? warp_sums[lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            sum += __shfl_down_sync(0xffffffffu, sum, off);
        if (lane == 0 && sum != 0.0f)
            atomicAdd(&out[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 && B.dim() == 2 && B.size(1) == 1,
                "A must be (M,K), 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 int M = (int)A_c.size(0);
    const int K = (int)A_c.size(1);

    auto out = torch::zeros({M, 1}, A_c.options());
    auto stream = at::cuda::getCurrentCUDAStream();

    const int chunk = 65536; // elements of K per split
    if ((K & 3) == 0) {
        const int K4 = K >> 2;
        const int chunk4 = chunk >> 2;
        const int splits = (K4 + chunk4 - 1) / chunk4;
        dim3 grid(M, splits);
        gemv_f32_vec4_kernel<<<grid, THREADS, 0, stream>>>(
            reinterpret_cast<const float4*>(A_c.data_ptr<float>()),
            reinterpret_cast<const float4*>(B_c.data_ptr<float>()),
            out.data_ptr<float>(), K4, chunk4);
    } else {
        const int splits = (K + chunk - 1) / chunk;
        dim3 grid(M, splits);
        gemv_f32_scalar_kernel<<<grid, THREADS, 0, stream>>>(
            A_c.data_ptr<float>(), B_c.data_ptr<float>(),
            out.data_ptr<float>(), K, chunk);
    }
    return out;
}
"""

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

gemv_module = load_inline(
    name="fused_gemv",
    cpp_sources=gemv_cpp_source,
    cuda_sources=gemv_cuda_source,
    functions=["gemv_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


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