```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 tiled_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 ROWS = 4;
constexpr int CHUNK = 32768;

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

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

template<bool FULL_ROWS>
__global__ __launch_bounds__(THREADS)
void tiled_gemv_partial(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partial,
    int64_t M,
    int64_t K,
    int splits
) {
    const int tid = threadIdx.x;
    const int lane = tid & 31;
    const int warp = tid >> 5;
    const int64_t first_row = static_cast<int64_t>(blockIdx.x) * ROWS;
    const int split = blockIdx.y;
    const int64_t begin = static_cast<int64_t>(split) * CHUNK;
    const int64_t end = begin + CHUNK < K ? begin + CHUNK : K;

    float sx[ROWS], sy[ROWS], sz[ROWS], sw[ROWS];
    #pragma unroll
    for (int r = 0; r < ROWS; ++r) {
        sx[r] = 0.0f;
        sy[r] = 0.0f;
        sz[r] = 0.0f;
        sw[r] = 0.0f;
    }

    #pragma unroll 4
    for (int64_t k = begin + 4 * tid; k < end; k += THREADS * 4) {
        const float4 b = __ldg(reinterpret_cast<const float4*>(B + k));

        #pragma unroll
        for (int r = 0; r < ROWS; ++r) {
            if (FULL_ROWS || first_row + r < M) {
                const float4 a = load_matrix(A + (first_row + r) * K + k);
                sx[r] = fmaf(a.x, b.x, sx[r]);
                sy[r] = fmaf(a.y, b.y, sy[r]);
                sz[r] = fmaf(a.z, b.z, sz[r]);
                sw[r] = fmaf(a.w, b.w, sw[r]);
            }
        }
    }

    __shared__ float sums[ROWS][THREADS / 32];

    #pragma unroll
    for (int r = 0; r < ROWS; ++r) {
        const float sum = warp_sum((sx[r] + sy[r]) + (sz[r] + sw[r]));
        if (lane == 0)
            sums[r][warp] = sum;
    }
    __syncthreads();

    if (warp < ROWS) {
        float sum = lane < THREADS / 32 ? sums[warp][lane] : 0.0f;
        sum = warp_sum(sum);
        if (lane == 0 && (FULL_ROWS || first_row + warp < M))
            partial[(first_row + warp) * splits + split] = sum;
    }
}

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

    if (row >= M)
        return;

    float sum = 0.0f;
    for (int s = lane; s < splits; s += 32)
        sum += partial[row * splits + s];

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

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

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

    const dim3 grid(static_cast<unsigned int>((M + ROWS - 1) / ROWS),
                    static_cast<unsigned int>(splits));
    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (M % ROWS == 0) {
        tiled_gemv_partial<true><<<grid, THREADS, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(),
            partial.data_ptr<float>(), M, K, splits);
    } else {
        tiled_gemv_partial<false><<<grid, THREADS, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(),
            partial.data_ptr<float>(), M, K, splits);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

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

    return output;
}
"""

_gemv = load_inline(
    name="tiled_gemv_four_rows_v2",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["tiled_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] <= 32768 * 65535
            and A.shape[1] % 4 == 0
            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.tiled_gemv_cuda(A, B)
        return torch.matmul(A, B)
```