```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 CUDA kernel:
#   out[b] = scale * sum_i max(C[b, 2*i], C[b, 2*i + 1])
# which replaces MaxPool1d(kernel_size=2) + sum(dim=1) + scalar multiply
# in a single kernel (one pass over the matmul output, no intermediates).
fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

__global__ void maxpool2_sum_scale_kernel(
    const float2* __restrict__ x,
    float* __restrict__ out,
    const int num_pairs,
    const float scale)
{
    const int b = blockIdx.y;
    const float2* row = x + (size_t)b * (size_t)num_pairs;

    // Each thread accumulates the pairwise maxes of several strided pairs.
    float acc = 0.0f;
    const int stride = gridDim.x * blockDim.x;
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < num_pairs; i += stride) {
        const float2 v = __ldg(row + i);
        acc += fmaxf(v.x, v.y);
    }

    // Warp-level reduction.
    #pragma unroll
    for (int off = 16; off > 0; off >>= 1)
        acc += __shfl_down_sync(0xffffffffu, acc, off);

    __shared__ float warp_sums[32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) warp_sums[wid] = acc;
    __syncthreads();

    if (wid == 0) {
        const int nwarps = blockDim.x >> 5;
        float v = (lane < nwarps) ? warp_sums[lane] : 0.0f;
        #pragma unroll
        for (int off = 16; off > 0; off >>= 1)
            v += __shfl_down_sync(0xffffffffu, v, off);
        if (lane == 0) atomicAdd(out + b, v * scale);
    }
}

torch::Tensor maxpool_sum_scale_cuda(torch::Tensor x, double scale) {
    TORCH_CHECK(x.is_cuda(), "x must be a CUDA tensor");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32, "x must be float32");
    TORCH_CHECK(x.dim() == 2, "x must be 2D (batch, features)");

    auto xc = x.contiguous();
    const int64_t B = xc.size(0);
    const int64_t N = xc.size(1);
    TORCH_CHECK(N % 2 == 0, "feature dim must be even for kernel_size=2");

    auto out = torch::zeros({B}, xc.options());

    const int num_pairs = (int)(N / 2);
    const int threads = 256;
    int chunks = (num_pairs + threads - 1) / threads;
    if (chunks > 16) chunks = 16;
    if (chunks < 1) chunks = 1;

    dim3 grid((unsigned)chunks, (unsigned)B);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    maxpool2_sum_scale_kernel<<<grid, threads, 0, stream>>>(
        reinterpret_cast<const float2*>(xc.data_ptr<float>()),
        out.data_ptr<float>(),
        num_pairs,
        (float)scale);

    return out;
}
"""

fused_epilogue_cpp_source = (
    "torch::Tensor maxpool_sum_scale_cuda(torch::Tensor x, double scale);"
)

# Compile the inline CUDA extension.
fused_epilogue = load_inline(
    name="fused_maxpool_sum_scale",
    cpp_sources=fused_epilogue_cpp_source,
    cuda_sources=fused_epilogue_cuda_source,
    functions=["maxpool_sum_scale_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model:
      - The GEMM (nn.Linear) is kept as-is since it maps to the heavily
        optimized cuBLAS addmm kernel.
      - MaxPool1d(kernel_size=2) + sum(dim=1) + scaling are fused into a
        single custom CUDA kernel, eliminating two kernel launches and all
        intermediate global-memory traffic.
    """
    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

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): Input tensor of shape (batch_size, in_features).

        Returns:
            torch.Tensor: Output tensor of shape (batch_size,).
        """
        x = self.matmul(x)
        if self.kernel_size == 2 and x.is_cuda and x.dtype == torch.float32:
            # Fused: maxpool(k=2) -> sum -> scale, in one kernel.
            return fused_epilogue.maxpool_sum_scale_cuda(x, float(self.scale_factor))
        # Generic fallback for any other configuration.
        x = torch.max_pool1d(x.unsqueeze(1), self.kernel_size).squeeze(1)
        x = torch.sum(x, dim=1)
        return x * self.scale_factor
```