```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 warp_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>
#include <stdint.h>

constexpr int THREADS = 256;
constexpr int ROWS = THREADS / 32;
constexpr int CHUNK = 8192;

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

__device__ __forceinline__ float4 load_matrix(const float* p) {
    float4 x;
    asm volatile(
        "ld.global.cg.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>
__global__ __launch_bounds__(THREADS, 4)
void warp_gemv_partial(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ partial,
    int M,
    int K,
    int splits
) {
    const int lane = threadIdx.x & 31;
    const int row = blockIdx.y * ROWS + (threadIdx.x >> 5);
    if (!FULL && row >= M)
        return;

    const int split = blockIdx.x;
    const int begin = split * CHUNK;
    const int count = FULL ? CHUNK : min(CHUNK, K - begin);

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

    float4 s0 = make_float4(0.f, 0.f, 0.f, 0.f);
    float4 s1 = make_float4(0.f, 0.f, 0.f, 0.f);

    #pragma unroll 2
    for (int i = lane * 4; i < count; i += 256) {
        const float4 a0 = load_matrix(a + i);
        const float4 b0 = __ldg(reinterpret_cast<const float4*>(b + i));

        if (FULL || i + 128 < count) {
            const float4 a1 = load_matrix(a + i + 128);
            const float4 b1 =
                __ldg(reinterpret_cast<const float4*>(b + i + 128));

            s1.x = fmaf(a1.x, b1.x, s1.x);
            s1.y = fmaf(a1.y, b1.y, s1.y);
            s1.z = fmaf(a1.z, b1.z, s1.z);
            s1.w = fmaf(a1.w, b1.w, s1.w);
        }

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

    const float x = s0.x + s1.x;
    const float y = s0.y + s1.y;
    const float z = s0.z + s1.z;
    const float w = s0.w + s1.w;
    const float sum = warp_sum((x + y) + (z + w));

    if (lane == 0)
        partial[static_cast<int64_t>(row) * splits + split] = sum;
}

__global__ void warp_gemv_finish(
    const float* __restrict__ partial,
    float* __restrict__ output,
    int M,
    int splits
) {
    const int row = blockIdx.x * ROWS + (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 warp_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;

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

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

    if (K % CHUNK == 0 && M % ROWS == 0) {
        warp_gemv_partial<true><<<grid, THREADS, 0, stream>>>(
            A.data_ptr<float>(), B.data_ptr<float>(),
            partial.data_ptr<float>(), M, K, splits
        );
    } else {
        warp_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) {
        warp_gemv_finish<<<(M + ROWS - 1) / ROWS, THREADS, 0, stream>>>(
            partial.data_ptr<float>(), output.data_ptr<float>(), M, splits
        );
        C10_CUDA_KERNEL_LAUNCH_CHECK();
    }
    return output;
}
"""

_gemv = load_inline(
    name="warp_streaming_gemv_v6",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["warp_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] <= 8 * 65535
            and 0 < A.shape[1] <= 2**30
            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.warp_gemv_cuda(A, B)
        return torch.matmul(A, B)
```