```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

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

# Custom split-K matrix-vector multiplication kernel (C = A @ B).
# A is (M, K) row-major, B is (K, 1), C is (M, 1).
#
# The operation is purely DRAM-bandwidth bound (A is streamed exactly once),
# so the kernel is engineered for maximum sustainable read throughput:
#   * 128-bit (float4) loads; A uses the evict-first streaming cache hint
#     (__ldcs) so it does not evict the much smaller, heavily reused B.
#   * Each thread block reduces TWO rows at once: B is loaded once per
#     iteration and reused for both rows, halving B load instructions and
#     L1/L2 traffic while doubling A bytes in flight per thread.
#   * split-K across the (small) M dimension sized from the SM count so the
#     whole device stays saturated; grid is ordered row-fastest so all
#     resident blocks share the same B slice (excellent L1/L2 locality).
#   * 4-way unrolled main loop with 8 independent accumulators for deep
#     memory-level parallelism; warp-shuffle + tiny shared second stage
#     reduction; partials combined with atomicAdd into a zero-filled output.
_matvec_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_runtime.h>

#define MV_THREADS 256
#define MV_UNROLL 4

__global__ void __launch_bounds__(MV_THREADS) matvec2row_kernel(
    const float4* __restrict__ A4,
    const float4* __restrict__ B4,
    const float*  __restrict__ As,
    const float*  __restrict__ Bs,
    float* __restrict__ C,
    const int K4,
    const int K,
    const int chunk,
    const int M)
{
    const int row0 = blockIdx.x * 2;
    int row1 = row0 + 1;
    if (row1 >= M) row1 = M - 1;  // odd M: duplicate row, result discarded

    const int part  = blockIdx.y;
    const int start = part * chunk;
    int end = start + chunk;
    if (end > K4) end = K4;

    const float4* a0 = A4 + (size_t)row0 * (size_t)K4;
    const float4* a1 = A4 + (size_t)row1 * (size_t)K4;

    const int stride = MV_THREADS;
    int i = start + threadIdx.x;

    float s00 = 0.f, s01 = 0.f, s02 = 0.f, s03 = 0.f;  // row0 partials
    float s10 = 0.f, s11 = 0.f, s12 = 0.f, s13 = 0.f;  // row1 partials

    // Main vectorized loop: 4-way unrolled, B reused across both rows.
    for (; i + 3 * stride < end; i += 4 * stride) {
        float4 b0 = __ldg(B4 + i);
        float4 b1 = __ldg(B4 + i + stride);
        float4 b2 = __ldg(B4 + i + 2 * stride);
        float4 b3 = __ldg(B4 + i + 3 * stride);

        float4 a00 = __ldcs(a0 + i);
        float4 a01 = __ldcs(a0 + i + stride);
        float4 a02 = __ldcs(a0 + i + 2 * stride);
        float4 a03 = __ldcs(a0 + i + 3 * stride);
        float4 a10 = __ldcs(a1 + i);
        float4 a11 = __ldcs(a1 + i + stride);
        float4 a12 = __ldcs(a1 + i + 2 * stride);
        float4 a13 = __ldcs(a1 + i + 3 * stride);

        s00 += a00.x * b0.x + a00.y * b0.y + a00.z * b0.z + a00.w * b0.w;
        s01 += a01.x * b1.x + a01.y * b1.y + a01.z * b1.z + a01.w * b1.w;
        s02 += a02.x * b2.x + a02.y * b2.y + a02.z * b2.z + a02.w * b2.w;
        s03 += a03.x * b3.x + a03.y * b3.y + a03.z * b3.z + a03.w * b3.w;
        s10 += a10.x * b0.x + a10.y * b0.y + a10.z * b0.z + a10.w * b0.w;
        s11 += a11.x * b1.x + a11.y * b1.y + a11.z * b1.z + a11.w * b1.w;
        s12 += a12.x * b2.x + a12.y * b2.y + a12.z * b2.z + a12.w * b2.w;
        s13 += a13.x * b3.x + a13.y * b3.y + a13.z * b3.z + a13.w * b3.w;
    }
    for (; i < end; i += stride) {
        float4 bv  = __ldg(B4 + i);
        float4 av0 = __ldcs(a0 + i);
        float4 av1 = __ldcs(a1 + i);
        s00 += av0.x * bv.x + av0.y * bv.y + av0.z * bv.z + av0.w * bv.w;
        s10 += av1.x * bv.x + av1.y * bv.y + av1.z * bv.z + av1.w * bv.w;
    }

    float sum0 = (s00 + s01) + (s02 + s03);
    float sum1 = (s10 + s11) + (s12 + s13);

    // Scalar tail when K is not divisible by 4 (handled by slice 0).
    if (part == 0) {
        const int base = K4 * 4;
        const float* ar0 = As + (size_t)row0 * (size_t)K;
        const float* ar1 = As + (size_t)row1 * (size_t)K;
        for (int j = base + threadIdx.x; j < K; j += stride) {
            const float bj = __ldg(Bs + j);
            sum0 += ar0[j] * bj;
            sum1 += ar1[j] * bj;
        }
    }

    // Warp-level reduction of both row sums.
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1) {
        sum0 += __shfl_down_sync(0xffffffffu, sum0, off);
        sum1 += __shfl_down_sync(0xffffffffu, sum1, off);
    }

    __shared__ float w0[MV_THREADS / 32];
    __shared__ float w1[MV_THREADS / 32];
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) {
        w0[warp] = sum0;
        w1[warp] = sum1;
    }
    __syncthreads();

    if (warp == 0) {
        float v0 = (lane < MV_THREADS / 32) ? w0[lane] : 0.f;
        float v1 = (lane < MV_THREADS / 32) ? w1[lane] : 0.f;
        #pragma unroll
        for (int off = MV_THREADS / 64; off > 0; off >>= 1) {
            v0 += __shfl_down_sync(0xffffffffu, v0, off);
            v1 += __shfl_down_sync(0xffffffffu, v1, off);
        }
        if (lane == 0) {
            atomicAdd(C + row0, v0);
            if (row1 != row0) atomicAdd(C + row1, v1);
        }
    }
}

torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B must be CUDA tensors");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");
    TORCH_CHECK(A.dim() == 2, "A must be 2-D");
    TORCH_CHECK(B.numel() == A.size(1), "B must have A.size(1) elements");

    const c10::cuda::CUDAGuard device_guard(A.device());

    auto Ac = A.contiguous();
    auto Bc = B.contiguous();

    const int64_t M = Ac.size(0);
    const int64_t K = Ac.size(1);

    auto C = torch::zeros({M, 1}, Ac.options());
    if (M == 0 || K == 0) return C;

    // The vectorized path requires 16-byte alignment (true for fresh
    // allocations); fall back to cuBLAS for exotic views.
    if ((reinterpret_cast<uintptr_t>(Ac.data_ptr<float>()) & 0xF) != 0 ||
        (reinterpret_cast<uintptr_t>(Bc.data_ptr<float>()) & 0xF) != 0) {
        return torch::matmul(A, B);
    }

    const int K4 = (int)(K >> 2);
    const int row_blocks = (int)((M + 1) >> 1);

    // Choose a power-of-two split-K factor so that the grid has enough
    // blocks to keep every SM busy (and the slices stay balanced).
    const int sm_count =
        at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
    const int64_t target = (int64_t)sm_count * 64;
    int split = 1;
    while (split < 32 && (int64_t)split * row_blocks < target) split <<= 1;

    int chunk = 0;
    if (K4 > 0) {
        const int unit = MV_THREADS * MV_UNROLL;
        chunk = (K4 + split - 1) / split;
        chunk = ((chunk + unit - 1) / unit) * unit;  // keep slices balanced
        split = (K4 + chunk - 1) / chunk;
    } else {
        split = 1;
        chunk = 0;
    }

    dim3 grid((unsigned)row_blocks, (unsigned)split);
    dim3 block(MV_THREADS);

    auto stream = at::cuda::getCurrentCUDAStream();
    matvec2row_kernel<<<grid, block, 0, stream>>>(
        reinterpret_cast<const float4*>(Ac.data_ptr<float>()),
        reinterpret_cast<const float4*>(Bc.data_ptr<float>()),
        Ac.data_ptr<float>(),
        Bc.data_ptr<float>(),
        C.data_ptr<float>(),
        K4, (int)K, chunk, (int)M);

    return C;
}
"""

_matvec_cpp_src = "torch::Tensor matvec_cuda(torch::Tensor A, torch::Tensor B);"

_matvec_ext = load_inline(
    name="matvec_splitk_ext_v2",
    cpp_sources=_matvec_cpp_src,
    cuda_sources=_matvec_cuda_src,
    functions=["matvec_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Matrix-vector multiplication (C = A * B) accelerated with a custom
    split-K CUDA kernel: two rows per block, streaming float4 loads for A
    (__ldcs), cached float4 loads for B, and SM-count-adaptive split-K.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._op = _matvec_ext

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (A.is_cuda and B.is_cuda and A.dim() == 2
                and A.dtype == torch.float32 and B.dtype == torch.float32):
            return self._op.matvec_cuda(A, B)
        return torch.matmul(A, B)
```