```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: non-overlapping max-pool (window = kernel_size, stride = kernel_size)
# + full sum over the feature dimension + scaling, all in a single kernel launch.
# The large GEMM (the dominant cost) is left to cuBLAS via nn.Linear, which is
# already near-optimal; everything after it is fused here to avoid launching
# 3 separate kernels and materializing intermediate tensors.
fused_maxpool_sum_scale_cuda_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>

__global__ void fused_maxpool_sum_scale_kernel(
        const float* __restrict__ x,
        float* __restrict__ out,
        int num_features,
        int kernel_size,
        float scale) {
    const int row = blockIdx.x;
    const float* xrow = x + (long long)row * num_features;
    const int num_windows = num_features / kernel_size;

    float local_sum = 0.0f;
    for (int w = threadIdx.x; w < num_windows; w += blockDim.x) {
        const float* win = xrow + w * kernel_size;
        float m = win[0];
        for (int t = 1; t < kernel_size; ++t) {
            m = fmaxf(m, win[t]);
        }
        local_sum += m;
    }

    __shared__ float smem[256];
    smem[threadIdx.x] = local_sum;
    __syncthreads();

    for (int s = blockDim.x >> 1; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            smem[threadIdx.x] += smem[threadIdx.x + s];
        }
        __syncthreads();
    }

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

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

    const int64_t batch = x.size(0);
    const int64_t features = x.size(1);

    auto out = torch::empty({batch}, x.options());
    const int block_size = 256;
    fused_maxpool_sum_scale_kernel<<<(int)batch, block_size>>>(
        x.data_ptr<float>(),
        out.data_ptr<float>(),
        (int)features,
        (int)kernel_size,
        (float)scale);
    return out;
}
"""

fused_maxpool_sum_scale_cpp_source = (
    "torch::Tensor fused_maxpool_sum_scale_cuda(torch::Tensor x, int64_t kernel_size, double scale);"
)

fused_maxpool_sum_scale = load_inline(
    name="fused_maxpool_sum_scale",
    cpp_sources=fused_maxpool_sum_scale_cpp_source,
    cuda_sources=fused_maxpool_sum_scale_cuda_source,
    functions=["fused_maxpool_sum_scale_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    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):
        x = self.matmul(x)
        return fused_maxpool_sum_scale.fused_maxpool_sum_scale_cuda(
            x, self.kernel_size, float(self.scale_factor)
        )
```