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

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

constexpr int THREADS = 256;
constexpr int CHUNK = 32768;

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

template<bool VECTORIZED>
__global__ void gemv_partial(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partial,
    int64_t K,
    int splits
) {
    const int64_t row = static_cast<int64_t>(blockIdx.x) / splits;
    const int split = blockIdx.x % splits;
    const int64_t begin = static_cast<int64_t>(split) * CHUNK;
    const int64_t remaining = K - begin;
    const int length = remaining < CHUNK ? static_cast<int>(remaining) : CHUNK;
    const float* a = A + row * K + begin;
    const float* b = B + begin;

    float sum = 0.0f;

    if (VECTORIZED) {
        const float4* a4 = reinterpret_cast<const float4*>(a);
        const float4* b4 = reinterpret_cast<const float4*>(b);
        float s0 = 0.0f;
        float s1 = 0.0f;
        float s2 = 0.0f;
        float s3 = 0.0f;

        for (int i = threadIdx.x; i < length / 4; i += THREADS) {
            const float4 av = a4[i];
            const float4 bv = b4[i];
            s0 = fmaf(av.x, bv.x, s0);
            s1 = fmaf(av.y, bv.y, s1);
            s2 = fmaf(av.z, bv.z, s2);
            s3 = fmaf(av.w, bv.w, s3);
        }
        sum = (s0 + s1) + (s2 + s3);
    } else {
        for (int i = threadIdx.x; i < length; i += THREADS) {
            sum = fmaf(a[i], b[i], sum);
        }
    }

    sum = warp_sum(sum);

    __shared__ float warp_sums[THREADS / 32];
    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 < THREADS / 32 ? warp_sums[lane] : 0.0f;
        sum = warp_sum(sum);
        if (lane == 0) {
            partial[row * splits + split] = sum;
        }
    }
}

__global__ void 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 split = lane; split < splits; split += 32) {
        sum += partial[row * splits + split];
    }

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

torch::Tensor splitk_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);
    auto output = torch::empty({M, 1}, A.options());

    if (M == 0) {
        return output;
    }

    const cudaStream_t stream =
        c10::cuda::getCurrentCUDAStream(A.get_device()).stream();

    if (K == 0) {
        C10_CUDA_CHECK(cudaMemsetAsync(
            output.data_ptr<float>(), 0, M * sizeof(float), stream));
        return output;
    }

    const int64_t split_count = (K + CHUNK - 1) / CHUNK;
    TORCH_CHECK(split_count <= INT32_MAX, "Too many K partitions");
    TORCH_CHECK(M <= INT32_MAX / split_count, "Matrix exceeds launch limits");

    const int splits = static_cast<int>(split_count);
    auto partial = torch::empty({M, splits}, A.options());

    const float* a = A.data_ptr<float>();
    const float* b = B.data_ptr<float>();
    const bool vectorized =
        K % 4 == 0 &&
        reinterpret_cast<uintptr_t>(a) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(b) % 16 == 0;

    const unsigned int blocks = static_cast<unsigned int>(M * splits);

    if (vectorized) {
        gemv_partial<true><<<blocks, THREADS, 0, stream>>>(
            a, b, partial.data_ptr<float>(), K, splits);
    } else {
        gemv_partial<false><<<blocks, THREADS, 0, stream>>>(
            a, b, partial.data_ptr<float>(), K, splits);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    const unsigned int finish_blocks =
        static_cast<unsigned int>((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="splitk_gemv_fp32_32768_v1",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["splitk_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.is_contiguous()
            and B.is_contiguous()
            and not (
                torch.is_grad_enabled() and (A.requires_grad or B.requires_grad)
            )
        ):
            return _gemv.splitk_gemv_cuda(A, B)
        return torch.matmul(A, B)
```