```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
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# The heavy matmul (nn.Linear: 128 x 32768 x 32768) is kept on the highly
# optimized cuBLAS path via F.linear. Everything after it
# (MaxPool1d(kernel=2) -> sum(dim=1) -> scale) is fused into a single
# memory-bound CUDA kernel that reads the (batch, out_features) matrix once
# and directly produces the (batch,) result using float4 vectorized loads,
# warp/block reductions, and a multi-chunk grid with atomic accumulation
# for high memory bandwidth utilization.

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

#define BLOCK_SIZE 256

__global__ void maxpool2_sum_scale_kernel(
    const float* __restrict__ y,
    float* __restrict__ out,
    const int n4,
    const int chunks_per_row,
    const float scale)
{
    const int row   = blockIdx.y;
    const int chunk = blockIdx.x;

    const float4* yrow = reinterpret_cast<const float4*>(y) +
                         (size_t)row * (size_t)n4;

    const int per_chunk = (n4 + chunks_per_row - 1) / chunks_per_row;
    const int start = chunk * per_chunk;
    const int end   = min(start + per_chunk, n4);

    // Each float4 holds two adjacent max-pool pairs.
    float acc = 0.0f;
    for (int j = start + threadIdx.x; j < end; j += BLOCK_SIZE) {
        const float4 v = yrow[j];
        acc += fmaxf(v.x, v.y) + fmaxf(v.z, v.w);
    }

    // Block-level tree reduction.
    __shared__ float smem[BLOCK_SIZE];
    smem[threadIdx.x] = acc;
    __syncthreads();
    #pragma unroll
    for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            smem[threadIdx.x] += smem[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        atomicAdd(out + row, smem[0] * scale);
    }
}

torch::Tensor maxpool2_sum_scale_cuda(torch::Tensor y, double scale) {
    TORCH_CHECK(y.is_cuda(), "y must be a CUDA tensor");
    TORCH_CHECK(y.scalar_type() == torch::kFloat32, "y must be float32");
    TORCH_CHECK(y.is_contiguous(), "y must be contiguous");
    TORCH_CHECK(y.dim() == 2, "y must be 2D");

    const int64_t batch = y.size(0);
    const int64_t width = y.size(1);
    TORCH_CHECK(width % 4 == 0, "width must be divisible by 4");
    const int n4 = (int)(width / 4);

    auto out = torch::zeros({batch}, y.options());

    // Split each row into chunks so the grid has enough blocks to
    // saturate memory bandwidth even for small batch sizes.
    int chunks = n4 / 2048;
    if (chunks < 1)  chunks = 1;
    if (chunks > 16) chunks = 16;

    dim3 grid((unsigned)chunks, (unsigned)batch);
    auto stream = at::cuda::getCurrentCUDAStream();

    maxpool2_sum_scale_kernel<<<grid, BLOCK_SIZE, 0, stream>>>(
        y.data_ptr<float>(),
        out.data_ptr<float>(),
        n4,
        chunks,
        (float)scale);

    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor maxpool2_sum_scale_cuda(torch::Tensor y, double scale);"
)

fused_epilogue = load_inline(
    name="fused_maxpool2_sum_scale",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["maxpool2_sum_scale_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized version of Model.

    - The large GEMM (nn.Linear) stays on cuBLAS via F.linear, which is
      already near-optimal for this compute-bound workload.
    - MaxPool1d(kernel_size=2) + sum(dim=1) + scaling are fused into one
      custom CUDA kernel, eliminating two intermediate tensors and the
      extra global-memory round trips they would require.
    """
    def __init__(self, in_features, out_features, kernel_size, scale_factor):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.kernel_size = kernel_size
        self.scale_factor = scale_factor
        # The custom kernel is specialized for kernel_size == 2.
        self._use_fused = (kernel_size == 2) and (out_features % 4 == 0)

    def forward(self, x):
        y = F.linear(x, self.matmul.weight, self.matmul.bias)
        if self._use_fused and y.is_cuda and y.dtype == torch.float32:
            return fused_epilogue.maxpool2_sum_scale_cuda(
                y, float(self.scale_factor))
        # Generic fallback for any configuration the custom kernel
        # does not cover.
        y = F.max_pool1d(y.unsqueeze(1), self.kernel_size).squeeze(1)
        y = torch.sum(y, dim=1)
        return y * self.scale_factor
```