```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 split_gemv_cuda(torch::Tensor A, torch::Tensor B);
"""

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

constexpr int ROWS_PER_BLOCK = 8;
constexpr int THREADS = 32 * ROWS_PER_BLOCK;
constexpr int CHUNK = 16384;

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

__global__ void gemv_partials(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partials,
    int64_t M,
    int64_t K,
    int64_t splits
) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int64_t row =
        static_cast<int64_t>(blockIdx.y) * ROWS_PER_BLOCK + warp;

    if (row >= M) return;

    const int64_t start = static_cast<int64_t>(blockIdx.x) * CHUNK;
    const int64_t end = (start + CHUNK < K) ? start + CHUNK : K;
    const float* row_ptr = A + row * K;

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

    for (int64_t k = start + lane * 4; k < end; k += 128) {
        const float4 a = *reinterpret_cast<const float4*>(row_ptr + k);
        const float4 b = *reinterpret_cast<const float4*>(B + k);

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

    const float sum = warp_sum((s0 + s1) + (s2 + s3));
    if (lane == 0) {
        partials[row * splits + blockIdx.x] = sum;
    }
}

__global__ void gemv_finalize(
    const float* __restrict__ partials,
    float* __restrict__ output,
    int64_t M,
    int64_t splits
) {
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    const int64_t row =
        static_cast<int64_t>(blockIdx.x) * ROWS_PER_BLOCK + warp;

    if (row >= M) return;

    float sum = 0.0f;
    for (int64_t split = lane; split < splits; split += 32) {
        sum += partials[row * splits + split];
    }

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

torch::Tensor split_gemv_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Inputs must share a device");
    TORCH_CHECK(
        A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
        "Inputs must be float32"
    );
    TORCH_CHECK(A.is_contiguous() && B.is_contiguous(),
                "Inputs must be contiguous");
    TORCH_CHECK(
        A.dim() == 2 && B.dim() == 2 &&
        B.size(1) == 1 && A.size(1) == B.size(0),
        "Expected A[M,K] and B[K,1]"
    );

    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    TORCH_CHECK(K > 0 && K % 4 == 0, "K must be a positive multiple of four");
    TORCH_CHECK(
        reinterpret_cast<uintptr_t>(A.data_ptr<float>()) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(B.data_ptr<float>()) % 16 == 0,
        "Inputs must be 16-byte aligned"
    );

    const c10::cuda::CUDAGuard guard(A.device());
    auto output = torch::empty({M, 1}, A.options());
    if (M == 0) return output;

    const int64_t splits = (K + CHUNK - 1) / CHUNK;
    auto partials = torch::empty({M, splits}, A.options());
    const cudaStream_t stream =
        c10::cuda::getCurrentCUDAStream(A.get_device()).stream();

    const dim3 grid(
        static_cast<unsigned int>(splits),
        static_cast<unsigned int>((M + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK)
    );

    gemv_partials<<<grid, THREADS, 0, stream>>>(
        A.data_ptr<float>(),
        B.data_ptr<float>(),
        partials.data_ptr<float>(),
        M, K, splits
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    gemv_finalize<<<
        static_cast<unsigned int>((M + ROWS_PER_BLOCK - 1) / ROWS_PER_BLOCK),
        THREADS, 0, stream
    >>>(
        partials.data_ptr<float>(),
        output.data_ptr<float>(),
        M, splits
    );
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return output;
}
"""

_gemv = load_inline(
    name="split_gemv_float4",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["split_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[1] == 1
            and A.shape[1] == B.shape[0]
            and A.shape[1] > 0
            and A.shape[1] % 4 == 0
            and A.is_contiguous()
            and B.is_contiguous()
            and A.data_ptr() % 16 == 0
            and B.data_ptr() % 16 == 0
            and (A.shape[0] + 7) // 8 <= 65535
            and not (
                torch.is_grad_enabled() and (A.requires_grad or B.requires_grad)
            )
        ):
            return _gemv.split_gemv_cuda(A, B)
        return torch.matmul(A, B)
```