```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


_cpp_source = r"""
#include <torch/extension.h>
torch::Tensor rowmajor_gemv_cuda(torch::Tensor A, torch::Tensor B);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

constexpr int THREADS = 256;
constexpr int CHUNK = 131072;

__device__ __forceinline__ float warp_sum(float x) {
    #pragma unroll
    for (int d = 16; d > 0; d >>= 1)
        x += __shfl_down_sync(0xffffffffu, x, d);
    return x;
}

__device__ __forceinline__ float4 stream_load(const float* p) {
    float4 x;
    asm volatile(
        "ld.global.cs.v4.f32 {%0, %1, %2, %3}, [%4];"
        : "=f"(x.x), "=f"(x.y), "=f"(x.z), "=f"(x.w)
        : "l"(p)
    );
    return x;
}

template<bool FULL_TILE>
__global__ __launch_bounds__(THREADS, 4)
void gemv_partial(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partial,
    int K,
    int splits
) {
    const int split = blockIdx.x % splits;
    const int row = blockIdx.x / splits;
    const int begin = split * CHUNK;
    const int remaining = K - begin;
    const int count = FULL_TILE ? CHUNK
                              : (remaining < CHUNK ? remaining : CHUNK);

    const float* a = A + static_cast<int64_t>(row) * K + begin;
    const float* b = B + begin;

    float s0 = 0.f, s1 = 0.f, s2 = 0.f, s3 = 0.f;
    float s4 = 0.f, s5 = 0.f, s6 = 0.f, s7 = 0.f;

    #pragma unroll 4
    for (int i = threadIdx.x * 4; i < count; i += THREADS * 8) {
        const float4 a0 = stream_load(a + i);
        const float4 b0 = __ldg(reinterpret_cast<const float4*>(b + i));

        if (FULL_TILE || i + THREADS * 4 < count) {
            const float4 a1 = stream_load(a + i + THREADS * 4);
            const float4 b1 = __ldg(
                reinterpret_cast<const float4*>(b + i + THREADS * 4)
            );

            s4 = fmaf(a1.x, b1.x, s4);
            s5 = fmaf(a1.y, b1.y, s5);
            s6 = fmaf(a1.z, b1.z, s6);
            s7 = fmaf(a1.w, b1.w, s7);
        }

        s0 = fmaf(a0.x, b0.x, s0);
        s1 = fmaf(a0.y, b0.y, s1);
        s2 = fmaf(a0.z, b0.z, s2);
        s3 = fmaf(a0.w, b0.w, s3);
    }

    float sum = ((s0 + s4) + (s1 + s5))
              + ((s2 + s6) + (s3 + s7));
    sum = warp_sum(sum);

    __shared__ float sums[THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;

    if (lane == 0)
        sums[warp] = sum;
    __syncthreads();

    if (warp == 0) {
        sum = lane < THREADS / 32 ? sums[lane] : 0.f;
        sum = warp_sum(sum);
        if (lane == 0)
            partial[static_cast<int64_t>(row) * splits + split] = sum;
    }
}

__global__ void gemv_finish(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int M,
    int splits
) {
    const int row = blockIdx.x * (THREADS / 32) + (threadIdx.x >> 5);
    const int lane = threadIdx.x & 31;
    if (row >= M)
        return;

    const float* p = partial + static_cast<int64_t>(row) * splits;
    float sum = 0.f;
    for (int i = lane; i < splits; i += 32)
        sum += p[i];

    sum = warp_sum(sum);
    if (lane == 0)
        output[row] = sum;
}

torch::Tensor rowmajor_gemv_cuda(torch::Tensor A, torch::Tensor B) {
    const c10::cuda::CUDAGuard guard(A.device());
    const int M = static_cast<int>(A.size(0));
    const int K = static_cast<int>(A.size(1));
    const int splits = (K + CHUNK - 1) / CHUNK;
    const unsigned int blocks = static_cast<unsigned int>(
        static_cast<int64_t>(M) * splits
    );

    auto output = torch::empty({M, 1}, A.options());
    auto partial = splits == 1
        ? output
        : torch::empty({M, splits}, A.options());

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

    if (K % CHUNK == 0) {
        gemv_partial<true><<<blocks, THREADS, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(),
            partial.data_ptr<float>(), K, splits
        );
    } else {
        gemv_partial<false><<<blocks, THREADS, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(),
            partial.data_ptr<float>(), K, splits
        );
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    if (splits > 1) {
        const unsigned int finish_blocks = static_cast<unsigned int>(
            (static_cast<int64_t>(M) + THREADS / 32 - 1) / (THREADS / 32)
        );
        gemv_finish<<<finish_blocks, THREADS, 0, stream>>>(
            partial.data_ptr<float>(), output.data_ptr<float>(), M, splits
        );
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }

    return output;
}
"""

_gemv = load_inline(
    name="rowmajor_streaming_gemv_v4",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["rowmajor_gemv_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (
            A.is_cuda
            and B.is_cuda
            and A.device == B.device
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.ndim == 2
            and B.ndim == 2
            and B.shape == (A.shape[1], 1)
            and A.is_contiguous()
            and B.is_contiguous()
            and 0 < A.shape[0] < 2**31
            and 0 < A.shape[1] <= 2**30
            and A.shape[1] % 4 == 0
            and A.shape[0] * ((A.shape[1] + 131071) // 131072) < 2**31
            and A.data_ptr() % 16 == 0
            and B.data_ptr() % 16 == 0
            and not (
                torch.is_grad_enabled() and (A.requires_grad or B.requires_grad)
            )
        ):
            return _gemv.rowmajor_gemv_cuda(A, B)
        return torch.matmul(A, B)
```