```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

# Fused split-K matrix-vector multiplication (C = A @ B) CUDA kernel.
#
# Workload: A (M, K) = (2048, 1048576), B (K, 1) -> C (M, 1).
# Purely DRAM-bandwidth bound: ~8 GiB of A must be streamed exactly once,
# so the whole optimization is about keeping the memory controllers fed.
#
# Key design choices:
#   * Split-K: each row group is partitioned into several K-slices, giving
#     thousands of independent blocks (full SM occupancy + smooth balance).
#   * "Slice-fastest" grid mapping: blockIdx.x is the K-slice, blockIdx.y
#     is the row group, so consecutively scheduled blocks read *consecutive*
#     slices of the same rows. The resident block set therefore sweeps DRAM
#     as one dense, nearly sequential stream instead of thousands of
#     streams scattered 4 MB apart, which improves DRAM page locality.
#   * float4 (128-bit) fully-coalesced loads, unrolled x2 for deep
#     memory-level parallelism.
#   * __ldcs (evict-first streaming) for A: every element is read once and
#     must not displace B in L2.
#   * __ldg (read-only) for B: the 4 MiB vector stays resident in L2 and
#     each load feeds kRowsPerBlock dot products.
#   * Warp-shuffle + shared-memory block reduction, finished with a single
#     atomicAdd per row slice - no global sync, no second reduction kernel.

matvec_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

namespace {

constexpr int kBlockSize    = 256;
constexpr int kNumWarps     = kBlockSize / 32;
constexpr int kRowsPerBlock = 4;
constexpr int kUnroll       = 2;  // float4 column chunks per thread per iter

template <int R>
__global__ void __launch_bounds__(kBlockSize)
matvec_splitk_kernel(const float4* __restrict__ A4,
                     const float4* __restrict__ B4,
                     float* __restrict__ C,
                     const int M,
                     const int K4,
                     const int slice4) {
    // Slice-fastest mapping: consecutive blocks -> consecutive memory.
    const int row0 = blockIdx.y * R;
    const int i0   = blockIdx.x * slice4;
    const int i1   = min(K4, i0 + slice4);

    // One base pointer per row handled by this block (clamped on tail
    // blocks; the recomputed rows are valid reads and are masked on write).
    const float4* a_rows[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) {
        int row = row0 + r;
        if (row >= M) row = M - 1;
        a_rows[r] = A4 + static_cast<size_t>(row) * K4;
    }

    float4 acc[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) acc[r] = make_float4(0.f, 0.f, 0.f, 0.f);

    const int stride = kBlockSize;
    int i = i0 + threadIdx.x;

    // Main loop: two float4 column chunks per iteration for extra MLP.
    for (; i + stride < i1; i += 2 * stride) {
        const float4 b0 = __ldg(B4 + i);
        const float4 b1 = __ldg(B4 + i + stride);
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            const float4 a0 = __ldcs(a_rows[r] + i);
            const float4 a1 = __ldcs(a_rows[r] + i + stride);
            acc[r].x += a0.x * b0.x + a1.x * b1.x;
            acc[r].y += a0.y * b0.y + a1.y * b1.y;
            acc[r].z += a0.z * b0.z + a1.z * b1.z;
            acc[r].w += a0.w * b0.w + a1.w * b1.w;
        }
    }
    // Slice remainder.
    for (; i < i1; i += stride) {
        const float4 b = __ldg(B4 + i);
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            const float4 a = __ldcs(a_rows[r] + i);
            acc[r].x += a.x * b.x;
            acc[r].y += a.y * b.y;
            acc[r].z += a.z * b.z;
            acc[r].w += a.w * b.w;
        }
    }

    // Collapse float4 lanes, then reduce across the block.
    float sum[R];
    #pragma unroll
    for (int r = 0; r < R; ++r) {
        sum[r] = (acc[r].x + acc[r].y) + (acc[r].z + acc[r].w);
    }

    #pragma unroll
    for (int r = 0; r < R; ++r) {
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1) {
            sum[r] += __shfl_down_sync(0xffffffffu, sum[r], off);
        }
    }

    __shared__ float warp_sums[R][kNumWarps];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;

    if (lane == 0) {
        #pragma unroll
        for (int r = 0; r < R; ++r) warp_sums[r][wid] = sum[r];
    }
    __syncthreads();

    if (wid == 0) {
        #pragma unroll
        for (int r = 0; r < R; ++r) {
            float s = (lane < kNumWarps) ? warp_sums[r][lane] : 0.f;
            #pragma unroll
            for (int off = kNumWarps / 2; off > 0; off >>= 1) {
                s += __shfl_down_sync(0xffffffffu, s, off);
            }
            if (lane == 0) {
                const int row = row0 + r;
                if (row < M) atomicAdd(&C[row], s);
            }
        }
    }
}

}  // namespace

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 && B.dim() == 2 && B.size(1) == 1,
                "Expected shapes A (M, K) and B (K, 1)");
    TORCH_CHECK(A.size(1) == B.size(0), "Inner dimensions must match");

    const int64_t M = A.size(0);
    const int64_t K = A.size(1);
    auto C = torch::empty({M, 1}, A.options());

    if (M == 0 || K == 0) {
        return C.zero_();
    }
    if ((K & 3) != 0) {
        // Alignment fallback (never triggered by this model's shapes).
        C.copy_(torch::matmul(A, B));
        return C;
    }

    const int K4 = static_cast<int>(K >> 2);
    const int groups = static_cast<int>((M + kRowsPerBlock - 1) / kRowsPerBlock);
    if (groups > 65535) {  // gridDim.y limit; never hit by this model
        C.copy_(torch::matmul(A, B));
        return C;
    }

    // Cache the SM count so steady-state calls have zero device-query cost.
    static int num_sms = 0;
    if (num_sms == 0) {
        int dev = 0;
        C10_CUDA_CHECK(cudaGetDevice(&dev));
        C10_CUDA_CHECK(cudaDeviceGetAttribute(
            &num_sms, cudaDevAttrMultiProcessorCount, dev));
    }

    // Choose the K-split so the grid holds enough blocks for several full
    // waves (hides the tail wave and keeps every SM saturated).
    int target = num_sms * 40;
    if (target < 4096) target = 4096;
    int slices = (target + groups - 1) / groups;
    if (slices < 1) slices = 1;
    if (slices > 64) slices = 64;

    // Keep slices a multiple of the unrolled block stride so every warp's
    // accesses stay 128-byte aligned and the main loop runs remainder-free.
    const int align = kBlockSize * kUnroll;
    int slice4 = (K4 + slices - 1) / slices;
    slice4 = ((slice4 + align - 1) / align) * align;
    slices = (K4 + slice4 - 1) / slice4;

    const dim3 grid(static_cast<unsigned int>(slices),
                    static_cast<unsigned int>(groups));

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    C10_CUDA_CHECK(cudaMemsetAsync(C.data_ptr<float>(), 0,
                                   static_cast<size_t>(M) * sizeof(float),
                                   stream));

    matvec_splitk_kernel<kRowsPerBlock><<<grid, kBlockSize, 0, stream>>>(
        reinterpret_cast<const float4*>(A.data_ptr<float>()),
        reinterpret_cast<const float4*>(B.data_ptr<float>()),
        C.data_ptr<float>(),
        static_cast<int>(M),
        K4,
        slice4);
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return C;
}
"""

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

matvec_ext = load_inline(
    name="matvec_splitk_seq_ext",
    cpp_sources=matvec_cpp_source,
    cuda_sources=matvec_cuda_source,
    functions=["matvec_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized matrix-vector multiplication (C = A @ B) backed by a custom
    split-K CUDA kernel whose grid sweeps DRAM as one dense sequential
    stream. Falls back to torch.matmul for unsupported input configurations.
    """

    def __init__(self):
        super(ModelNew, self).__init__()
        self._matvec = matvec_ext.matvec_cuda

    def forward(self, A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
        if (
            A.is_cuda
            and B.is_cuda
            and A.dtype == torch.float32
            and B.dtype == torch.float32
            and A.dim() == 2
            and B.dim() == 2
            and B.size(1) == 1
            and A.size(1) == B.size(0)
        ):
            return self._matvec(A.contiguous(), B.contiguous())
        return torch.matmul(A, B)
```