```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 (C = A @ B).
# This operation is heavily memory-bound (reads the full M x K matrix),
# so the kernel is optimized for maximum memory bandwidth:
#   - One thread block per output row.
#   - Vectorized float4 loads (16 bytes per thread) for full coalescing.
#   - Streaming loads (__ldcs) for A since it is read exactly once.
#   - Read-only cached loads (__ldg) for B (it is reused across rows).
#   - 4x unrolling for more memory-level parallelism / ILP.
#   - Fast warp-shuffle based block reduction.
matvec_cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 256

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

    long K4 = K >> 2;
    const float4* A4 = reinterpret_cast<const float4*>(A_row);
    const float4* B4 = reinterpret_cast<const float4*>(B);

    float sum = 0.0f;
    long i = threadIdx.x;
    const long stride = blockDim.x;

    // Main vectorized loop, unrolled 4x for memory-level parallelism.
    for (; i + 3 * stride < K4; i += 4 * stride) {
        float4 a0 = __ldcs(&A4[i]);
        float4 a1 = __ldcs(&A4[i + stride]);
        float4 a2 = __ldcs(&A4[i + 2 * stride]);
        float4 a3 = __ldcs(&A4[i + 3 * stride]);
        float4 b0 = __ldg(&B4[i]);
        float4 b1 = __ldg(&B4[i + stride]);
        float4 b2 = __ldg(&B4[i + 2 * stride]);
        float4 b3 = __ldg(&B4[i + 3 * stride]);
        sum += a0.x * b0.x + a0.y * b0.y + a0.z * b0.z + a0.w * b0.w;
        sum += a1.x * b1.x + a1.y * b1.y + a1.z * b1.z + a1.w * b1.w;
        sum += a2.x * b2.x + a2.y * b2.y + a2.z * b2.z + a2.w * b2.w;
        sum += 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) {
        float4 a = __ldcs(&A4[i]);
        float4 b = __ldg(&B4[i]);
        sum += a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
    }
    // Scalar tail if K is not a multiple of 4.
    long tail_start = K4 << 2;
    for (long j = tail_start + threadIdx.x; j < K; j += blockDim.x) {
        sum += A_row[j] * B[j];
    }

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

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

    if (warp_id == 0) {
        sum = (lane < BLOCK_SIZE / 32) ? warp_sums[lane] : 0.0f;
        for (int offset = 16; 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() == torch::kFloat32 &&
                B.scalar_type() == torch::kFloat32,
                "Inputs must be float32");

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

    long M = A_c.size(0);
    long K = A_c.size(1);

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

    dim3 grid((unsigned int)M);
    dim3 block(BLOCK_SIZE);

    matvec_kernel<<<grid, block>>>(
        A_c.data_ptr<float>(),
        B_c.data_ptr<float>(),
        C.data_ptr<float>(),
        (int)M, K);

    return C;
}
"""

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

# Compile the inline CUDA kernel.
matvec_ext = load_inline(
    name="matvec_cuda_ext",
    cpp_sources=matvec_cpp_source,
    cuda_sources=matvec_cuda_source,
    functions=["matvec_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3", "--use_fast_math"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs matrix-vector multiplication (C = A * B)
    using a custom bandwidth-optimized CUDA kernel.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.matvec = 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.matvec.matvec_cuda(A, B)
        return torch.matmul(A, B)
```