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

# TF32 tensor-core math is allowed for the FP32 matmul (matches PyTorch reference behavior).
try:
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
except Exception:
    pass

# ----------------------------------------------------------------------------
# Fused epilogue kernel (used on the cuBLAS fallback path):
#   out[b] = scale * sum_i max(y[b, 2i], y[b, 2i+1])
# ----------------------------------------------------------------------------
_fused_epilogue_cuda_source = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cstdint>

constexpr int BLOCK_SIZE = 256;

__global__ void maxpool_sum_scale_kernel(const float* __restrict__ x,
                                         float* __restrict__ out,
                                         int64_t out_features,
                                         float scale) {
    const int64_t row = blockIdx.x;
    const float* p = x + row * out_features;
    const int64_t num_pairs = out_features >> 1;

    float sum = 0.0f;

    if ((reinterpret_cast<uintptr_t>(p) & 0xF) == 0) {
        const float4* p4 = reinterpret_cast<const float4*>(p);
        const int64_t num_vec = num_pairs >> 1;
        for (int64_t i = threadIdx.x; i < num_vec; i += BLOCK_SIZE) {
            float4 v = p4[i];
            sum += fmaxf(v.x, v.y) + fmaxf(v.z, v.w);
        }
        if ((num_pairs & 1) && threadIdx.x == 0) {
            const int64_t i = num_pairs - 1;
            sum += fmaxf(p[2 * i], p[2 * i + 1]);
        }
    } else {
        for (int64_t i = threadIdx.x; i < num_pairs; i += BLOCK_SIZE) {
            sum += fmaxf(p[2 * i], p[2 * i + 1]);
        }
    }

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, offset);
    }

    __shared__ float shared[BLOCK_SIZE / 32];
    const int lane = threadIdx.x & 31;
    const int wid  = threadIdx.x >> 5;
    if (lane == 0) {
        shared[wid] = sum;
    }
    __syncthreads();

    if (wid == 0) {
        sum = (lane < BLOCK_SIZE / 32) ? shared[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, offset);
        }
        if (lane == 0) {
            out[row] = sum * 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");
    auto xc = x.contiguous();
    const int64_t batch = xc.size(0);
    const int64_t out_features = xc.size(1);

    auto out = torch::empty({batch}, xc.options());

    dim3 grid(static_cast<unsigned int>(batch));
    maxpool_sum_scale_kernel<<<grid, BLOCK_SIZE>>>(
        xc.data_ptr<float>(),
        out.data_ptr<float>(),
        out_features,
        static_cast<float>(scale));

    return out;
}
"""

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

# ----------------------------------------------------------------------------
# Fully fused TF32 tensor-core pipeline (requires sm_80+):
#   out = scale * sum_j max((x @ W^T + bias)[:, 2j], (x @ W^T + bias)[:, 2j+1])
#
# The GEMM is DRAM-bound (the weight matrix is streamed exactly once).  The
# stock cuBLAS kernel launches only ~256 CTAs for this shape, which leaves the
# GPU idle in the tail wave.  This implementation uses an automatic split-K
# factor chosen from the SM count so the CTA count is a near multiple of the
# SM count (~95%+ wave efficiency), a 4-stage cp.async pipeline, XOR-swizzled
# shared memory, and an L2 evict-first policy on the streaming weight loads
# (keeping the small activation matrix resident in L2).  Partial K results
# are written to a workspace; a tiny second kernel applies bias, max-pool(2),
# sum and scale.
# ----------------------------------------------------------------------------
_fused_gemm_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <cstdint>

#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800)
#define FG_SM80 1
#else
#define FG_SM80 0
#endif

namespace fg {

constexpr int BM = 128;                          // batch tile (M <= 128)
constexpr int BN = 128;                          // output-feature tile / CTA
constexpr int BK = 32;                           // k tile
constexpr int STAGES = 4;                        // cp.async pipeline depth
constexpr int THREADS = 256;
constexpr int STAGE_FLOATS = BM * BK;            // per matrix per stage
constexpr int SMEM_FLOATS = STAGES * 2 * STAGE_FLOATS;
constexpr int SMEM_BYTES = SMEM_FLOATS * (int)sizeof(float);  // 128 KB

__device__ __forceinline__ uint32_t smem_u32(const void* p) {
    return static_cast<uint32_t>(__cvta_generic_to_shared(p));
}

__device__ __forceinline__ void cp16(uint32_t dst, const void* src) {
#if FG_SM80
    asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n"
                 :: "r"(dst), "l"(src));
#endif
}

__device__ __forceinline__ void cp16_stream(uint32_t dst, const void* src,
                                            uint64_t pol) {
#if FG_SM80
    asm volatile(
        "cp.async.cg.shared.global.L2::cache_hint [%0], [%1], 16, %2;\n"
        :: "r"(dst), "l"(src), "l"(pol));
#endif
}

__device__ __forceinline__ void cp_commit() {
#if FG_SM80
    asm volatile("cp.async.commit_group;\n" ::);
#endif
}

template <int N>
__device__ __forceinline__ void cp_wait() {
#if FG_SM80
    asm volatile("cp.async.wait_group %0;\n" :: "n"(N));
#endif
}

__device__ __forceinline__ float to_tf32(float x) {
#if FG_SM80
    uint32_t u = __float_as_uint(x);
    asm volatile("cvt.rna.tf32.f32 %0, %1;\n" : "=r"(u) : "r"(u));
    return __uint_as_float(u);
#else
    return x;
#endif
}

__device__ __forceinline__ void mma_tf32(float& d0, float& d1, float& d2,
                                         float& d3, float a0, float a1,
                                         float a2, float a3, float b0,
                                         float b1) {
#if FG_SM80
    asm volatile(
        "mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 "
        "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
        : "+f"(d0), "+f"(d1), "+f"(d2), "+f"(d3)
        : "r"(__float_as_uint(a0)), "r"(__float_as_uint(a1)),
          "r"(__float_as_uint(a2)), "r"(__float_as_uint(a3)),
          "r"(__float_as_uint(b0)), "r"(__float_as_uint(b1)));
#endif
}

__global__ void __launch_bounds__(THREADS, 1)
gemm_partial_kernel(const float* __restrict__ X,
                    const float* __restrict__ W,
                    float* __restrict__ WS,
                    int M, int N, int K, int kchunk) {
#if FG_SM80
    extern __shared__ float smem[];
    const int tid = threadIdx.x;
    const int n0 = blockIdx.x * BN;
    const int k0 = blockIdx.y * kchunk;
    const int nk = kchunk / BK;
    const int mlim = M - 1;

    // Streaming loads of W get an L2 evict-first policy so the much smaller,
    // heavily reused activation matrix stays resident in L2.
    uint64_t pol = 0;
    asm volatile(
        "createpolicy.fractional.L2::evict_first.b64 %0, 1.0;\n" : "=l"(pol));

    // 256 threads x 16B cover one 128x32 tile in 4 iterations.
    const int lr = tid >> 3;           // row group 0..31
    const int lc = tid & 7;            // 16B chunk in row 0..7
    const int lk = lc * 4;             // float offset in row

    auto load_stage = [&](int ks, int buf) {
        const int kg = k0 + ks * BK + lk;
        float* as = smem + buf * (2 * STAGE_FLOATS);
        float* bs = as + STAGE_FLOATS;
        #pragma unroll
        for (int i = 0; i < 4; ++i) {
            const int r = lr + i * 32;
            const int sw = (lc ^ (r & 7)) << 2;
            cp16(smem_u32(as + r * BK + sw),
                 X + (int64_t)(r < M ? r : mlim) * K + kg);
            cp16_stream(smem_u32(bs + r * BK + sw),
                        W + (int64_t)(n0 + r) * K + kg, pol);
        }
    };

    // Warp layout: 8 warps as 4 (rows of 32) x 2 (cols of 64).
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int gid = lane >> 2;         // 0..7
    const int tig = lane & 3;          // 0..3
    const int wrow = (warp & 3) * 32;
    const int wcol = (warp >> 2) * 64;

    float acc[2][8][4];                // [m16 tile][n8 tile][c0..c3]
    #pragma unroll
    for (int i = 0; i < 2; ++i)
        #pragma unroll
        for (int j = 0; j < 8; ++j)
            #pragma unroll
            for (int e = 0; e < 4; ++e) acc[i][j][e] = 0.0f;

    // Prologue: fill the first STAGES-1 stages (always commit to keep the
    // group counting consistent even for tiny K).
    #pragma unroll
    for (int s = 0; s < STAGES - 1; ++s) {
        if (s < nk) {
            load_stage(s, s);
        }
        cp_commit();
    }

    int read_buf = 0;
    int write_buf = STAGES - 1;
    for (int ks = 0; ks < nk; ++ks) {
        cp_wait<STAGES - 2>();
        __syncthreads();

        const int kpf = ks + STAGES - 1;
        if (kpf < nk) {
            load_stage(kpf, write_buf);
        }
        cp_commit();
        write_buf = (write_buf + 1) % STAGES;

        const float* as = smem + read_buf * (2 * STAGE_FLOATS);
        const float* bs = as + STAGE_FLOATS;
        read_buf = (read_buf + 1) % STAGES;

        #pragma unroll
        for (int kk = 0; kk < BK / 8; ++kk) {
            const int lc2 = kk * 2 + (tig >> 1);   // 16B chunk for this k-step
            const int lk2 = (tig & 1) * 2;         // float offset within chunk
            float av[2][4];
            float bv[8][2];
            #pragma unroll
            for (int i = 0; i < 2; ++i) {
                const int r = wrow + i * 16 + gid;
                const int off = ((lc2 ^ (r & 7)) << 2) + lk2;
                const float2 lo = *reinterpret_cast<const float2*>(
                    as + r * BK + off);
                const float2 hi = *reinterpret_cast<const float2*>(
                    as + (r + 8) * BK + off);
                av[i][0] = to_tf32(lo.x);
                av[i][1] = to_tf32(lo.y);
                av[i][2] = to_tf32(hi.x);
                av[i][3] = to_tf32(hi.y);
            }
            #pragma unroll
            for (int j = 0; j < 8; ++j) {
                const int r = wcol + j * 8 + gid;
                const int off = ((lc2 ^ (r & 7)) << 2) + lk2;
                const float2 b2 = *reinterpret_cast<const float2*>(
                    bs + r * BK + off);
                bv[j][0] = to_tf32(b2.x);
                bv[j][1] = to_tf32(b2.y);
            }
            #pragma unroll
            for (int j = 0; j < 8; ++j) {
                #pragma unroll
                for (int i = 0; i < 2; ++i) {
                    mma_tf32(acc[i][j][0], acc[i][j][1], acc[i][j][2],
                             acc[i][j][3], av[i][0], av[i][1], av[i][2],
                             av[i][3], bv[j][0], bv[j][1]);
                }
            }
        }
        __syncthreads();
    }

    // Drain accumulators to the split-K workspace.
    float* wsp = WS + (int64_t)blockIdx.y * M * N;
    #pragma unroll
    for (int i = 0; i < 2; ++i) {
        const int r0 = wrow + i * 16 + gid;
        #pragma unroll
        for (int j = 0; j < 8; ++j) {
            const int col = n0 + wcol + j * 8 + tig * 2;
            if (r0 < M) {
                *reinterpret_cast<float2*>(wsp + (int64_t)r0 * N + col) =
                    make_float2(acc[i][j][0], acc[i][j][1]);
            }
            if (r0 + 8 < M) {
                *reinterpret_cast<float2*>(wsp + (int64_t)(r0 + 8) * N + col) =
                    make_float2(acc[i][j][2], acc[i][j][3]);
            }
        }
    }
#endif
}

constexpr int RED_BLOCK = 256;

__global__ void reduce_epilogue_kernel(const float* __restrict__ ws,
                                       const float* __restrict__ bias,
                                       float* __restrict__ out,
                                       int M, int N, int sk, float scale) {
    const int row = blockIdx.x;
    const int nvec = N >> 2;
    const int64_t slice = (int64_t)M * N;

    float sum = 0.0f;
    for (int i = threadIdx.x; i < nvec; i += RED_BLOCK) {
        const float4 bb =
            *reinterpret_cast<const float4*>(bias + (i << 2));
        float4 v = bb;
        for (int s = 0; s < sk; ++s) {
            const float4 w4 = *reinterpret_cast<const float4*>(
                ws + (int64_t)s * slice + (int64_t)row * N + (i << 2));
            v.x += w4.x;
            v.y += w4.y;
            v.z += w4.z;
            v.w += w4.w;
        }
        sum += fmaxf(v.x, v.y) + fmaxf(v.z, v.w);
    }

    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, offset);
    }
    __shared__ float shared[RED_BLOCK / 32];
    const int lane = threadIdx.x & 31;
    const int wid = threadIdx.x >> 5;
    if (lane == 0) {
        shared[wid] = sum;
    }
    __syncthreads();
    if (wid == 0) {
        sum = (lane < RED_BLOCK / 32) ? shared[lane] : 0.0f;
        #pragma unroll
        for (int offset = 16; offset > 0; offset >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, offset);
        }
        if (lane == 0) {
            out[row] = sum * scale;
        }
    }
}

}  // namespace fg

torch::Tensor fused_linear_maxpool_sum_cuda(torch::Tensor x, torch::Tensor w,
                                            torch::Tensor b, double scale) {
    TORCH_CHECK(x.is_cuda() && w.is_cuda() && b.is_cuda(),
                "CUDA tensors required");
    TORCH_CHECK(x.scalar_type() == torch::kFloat32 &&
                w.scalar_type() == torch::kFloat32 &&
                b.scalar_type() == torch::kFloat32, "float32 required");
    auto xc = x.contiguous();
    auto wc = w.contiguous();
    auto bc = b.contiguous();
    const int64_t M64 = xc.size(0);
    const int64_t K64 = xc.size(1);
    const int64_t N64 = wc.size(0);
    TORCH_CHECK(wc.size(1) == K64 && bc.numel() == N64, "shape mismatch");
    TORCH_CHECK(M64 >= 1 && M64 <= fg::BM, "batch too large for fused kernel");
    TORCH_CHECK((N64 % fg::BN) == 0, "out_features must be a multiple of 128");
    TORCH_CHECK((K64 % fg::BK) == 0, "in_features must be a multiple of 32");
    const int M = (int)M64;
    const int N = (int)N64;
    const int K = (int)K64;

    auto stream = at::cuda::getCurrentCUDAStream();

    static int sm_count = [] {
        int dev = 0;
        cudaGetDevice(&dev);
        cudaDeviceProp prop;
        cudaGetDeviceProperties(&prop, dev);
        return prop.multiProcessorCount;
    }();
    if (sm_count <= 0) {
        sm_count = 108;
    }

    static int smem_configured = [] {
        cudaFuncSetAttribute(fg::gemm_partial_kernel,
                             cudaFuncAttributeMaxDynamicSharedMemorySize,
                             fg::SMEM_BYTES);
        return 1;
    }();
    (void)smem_configured;

    // Pick a split-K factor that minimises wave quantisation.
    const int64_t n_tiles = N64 / fg::BN;
    int best_sk = 1;
    double best_eff = 0.0;
    for (int sk = 1; sk <= 8; ++sk) {
        if ((K64 % ((int64_t)fg::BK * sk)) != 0) {
            continue;
        }
        const int64_t blocks = n_tiles * sk;
        const int64_t waves = (blocks + sm_count - 1) / sm_count;
        const double eff = (double)blocks / (double)(waves * sm_count);
        if (eff > best_eff * 1.02 + 1e-9) {
            best_eff = eff;
            best_sk = sk;
        }
    }
    const int kchunk = (int)(K64 / best_sk);

    auto ws = torch::empty({(int64_t)best_sk, M64, N64}, xc.options());
    auto out = torch::empty({M64}, xc.options());

    dim3 grid((unsigned)n_tiles, (unsigned)best_sk);
    fg::gemm_partial_kernel<<<grid, fg::THREADS, fg::SMEM_BYTES, stream>>>(
        xc.data_ptr<float>(), wc.data_ptr<float>(), ws.data_ptr<float>(),
        M, N, K, kchunk);

    dim3 rgrid((unsigned)M);
    fg::reduce_epilogue_kernel<<<rgrid, fg::RED_BLOCK, 0, stream>>>(
        ws.data_ptr<float>(), bc.data_ptr<float>(), out.data_ptr<float>(),
        M, N, best_sk, (float)scale);

    return out;
}
"""

_fused_gemm_cpp_source = (
    "torch::Tensor fused_linear_maxpool_sum_cuda(torch::Tensor x, "
    "torch::Tensor w, torch::Tensor b, double scale);"
)

try:
    _fused_epilogue = load_inline(
        name="fused_maxpool_sum_scale_v4",
        cpp_sources=_fused_epilogue_cpp_source,
        cuda_sources=_fused_epilogue_cuda_source,
        functions=["maxpool_sum_scale_cuda"],
        verbose=False,
        extra_cuda_cflags=["-O3"],
    )
except Exception:
    _fused_epilogue = None

try:
    _fused_gemm = load_inline(
        name="fused_linear_maxpool_sum_gemm_v4",
        cpp_sources=_fused_gemm_cpp_source,
        cuda_sources=_fused_gemm_cuda_source,
        functions=["fused_linear_maxpool_sum_cuda"],
        verbose=False,
        extra_cuda_cflags=["-O3"],
    )
except Exception:
    _fused_gemm = None


class ModelNew(nn.Module):
    """
    Optimized model.

    Candidate execution paths (validated for correctness and timed on the
    first eligible forward call; the fastest correct one is cached and used
    permanently afterwards):

      1. cuBLAS (TF32) linear + fused max-pool/sum/scale epilogue kernel.
      2. Fully fused TF32 tensor-core GEMM with automatic split-K (wave-
         balanced) + a tiny bias/max-pool/sum/scale reduction kernel.  The
         weight matrix is streamed from DRAM exactly once.

    Any input the fast paths do not support falls back to the reference
    PyTorch ops.
    """

    _plan_cache = {}

    def __init__(self, in_features, out_features, kernel_size, scale_factor):
        super(ModelNew, self).__init__()
        self.matmul = nn.Linear(in_features, out_features)
        self.max_pool = nn.MaxPool1d(kernel_size)
        self.kernel_size = kernel_size
        self.scale_factor = float(scale_factor)
        self.in_features = in_features
        self.out_features = out_features

    # ------------------------------------------------------------------ paths
    def _reference(self, x):
        y = F.linear(x, self.matmul.weight, self.matmul.bias)
        y = self.max_pool(y.unsqueeze(1)).squeeze(1)
        y = torch.sum(y, dim=1)
        return y * self.scale_factor

    def _cublas_fused_epilogue(self, x):
        y = F.linear(x, self.matmul.weight, self.matmul.bias)
        if _fused_epilogue is not None and (y.size(1) % 2) == 0:
            return _fused_epilogue.maxpool_sum_scale_cuda(y, self.scale_factor)
        y2 = y.view(y.size(0), y.size(1) // 2, 2)
        m = torch.maximum(y2[..., 0], y2[..., 1])
        return m.sum(dim=1) * self.scale_factor

    def _fused_gemm(self, x):
        return _fused_gemm.fused_linear_maxpool_sum_cuda(
            x, self.matmul.weight, self.matmul.bias, self.scale_factor)

    def _fused_eligible(self, x):
        if _fused_gemm is None or self.matmul.bias is None:
            return False
        if x.size(0) < 1 or x.size(0) > 128:
            return False
        if (self.out_features % 128) != 0 or (self.in_features % 32) != 0:
            return False
        try:
            if torch.cuda.get_device_capability(x.device)[0] < 8:
                return False
        except Exception:
            return False
        return True

    # ------------------------------------------------------------- selection
    @staticmethod
    def _time_path(fn, x, reps=2):
        torch.cuda.synchronize()
        start = torch.cuda.Event(enable_timing=True)
        end = torch.cuda.Event(enable_timing=True)
        start.record()
        for _ in range(reps):
            fn(x)
        end.record()
        torch.cuda.synchronize()
        return start.elapsed_time(end)

    def _select_plan(self, x):
        base = self._cublas_fused_epilogue
        ref = base(x)
        best = base
        best_t = self._time_path(base, x)

        if self._fused_eligible(x):
            try:
                out = self._fused_gemm(x)
                tol = 2e-4 * float(ref.abs().max().item()) + 1e-2
                err = float((out - ref).abs().max().item())
                if bool(torch.isfinite(out).all().item()) and err <= tol:
                    t = self._time_path(self._fused_gemm, x)
                    if t < best_t * 0.98:
                        best, best_t = self._fused_gemm, t
            except Exception:
                pass
        return best

    # --------------------------------------------------------------- forward
    def forward(self, x):
        eligible = (
            x.is_cuda and x.dtype == torch.float32
            and self.kernel_size == 2 and x.dim() == 2
            and x.size(0) > 0 and (self.out_features % 2) == 0
            and self.matmul.weight.is_cuda
        )
        if not eligible:
            return self._reference(x)

        if self._fused_eligible(x):
            key = (x.size(0), self.in_features, self.out_features,
                   self.scale_factor, x.device.index)
            plan = ModelNew._plan_cache.get(key)
            if plan is None:
                try:
                    plan = self._select_plan(x)
                except Exception:
                    plan = self._cublas_fused_epilogue
                ModelNew._plan_cache[key] = plan
            return plan(x)

        return self._cublas_fused_epilogue(x)
```