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

# Custom CUDA kernel for matrix-scalar multiplication (C = A * s).
#
# This operator is purely memory-bandwidth bound (read 4 GiB, write 4 GiB),
# so the kernel keeps the empirically fastest configuration (128-bit vec4
# accesses, plain ld/st, zero bounds checks on the hot path, 32-bit indexing,
# 4 independent 16-byte loads in flight per thread, 256 threads/block).
#
# Since the kernel already runs at ~peak DRAM bandwidth, this revision focuses
# on shaving the remaining host-side overhead so total wall-clock beats the
# PyTorch reference more reliably:
#   * bypass the dispatcher for the output allocation by calling
#     at::detail::empty_cuda directly (saves several microseconds per call
#     vs. at::empty, which matters at ~6 ms scale)
#   * zero Tensor copies on the hot path: take the input by const reference
#     and recurse on the (rare) non-contiguous case instead of materializing
#     an owning copy of the input handle
#   * single branch chain, direct void* data pointers, minimal arg passing
scalar_mul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/EmptyTensor.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <climits>

namespace {

constexpr int kThreads = 256;
constexpr int kUnroll  = 4;
constexpr int kTile    = kThreads * kUnroll;   // float4 elements per block

// Hot path: n4 is an exact multiple of kTile -> no bounds checks at all,
// 4 outstanding 16B loads per thread, 32-bit index math.
__global__ void __launch_bounds__(kThreads)
scalar_mul_vec4_uniform_kernel(const float4* __restrict__ A,
                               float4* __restrict__ C,
                               const float s) {
    const int base = (int)blockIdx.x * kTile + (int)threadIdx.x;

    float4 v[kUnroll];
    #pragma unroll
    for (int i = 0; i < kUnroll; ++i) {
        v[i] = A[base + i * kThreads];
    }
    #pragma unroll
    for (int i = 0; i < kUnroll; ++i) {
        v[i].x *= s;
        v[i].y *= s;
        v[i].z *= s;
        v[i].w *= s;
    }
    #pragma unroll
    for (int i = 0; i < kUnroll; ++i) {
        C[base + i * kThreads] = v[i];
    }
}

// Generic vectorized path: one float4 per thread with a single bounds check.
__global__ void __launch_bounds__(kThreads)
scalar_mul_vec4_kernel(const float4* __restrict__ A,
                       float4* __restrict__ C,
                       const float s,
                       const int n4) {
    const int idx = (int)blockIdx.x * kThreads + (int)threadIdx.x;
    if (idx < n4) {
        float4 v = A[idx];
        v.x *= s;
        v.y *= s;
        v.z *= s;
        v.w *= s;
        C[idx] = v;
    }
}

// Scalar grid-stride fallback (unaligned pointers / odd sizes / huge tensors).
__global__ void scalar_mul_fallback_kernel(const float* __restrict__ A,
                                           float* __restrict__ C,
                                           const float s,
                                           const long n) {
    const long stride = (long)gridDim.x * blockDim.x;
    for (long idx = (long)blockIdx.x * blockDim.x + threadIdx.x; idx < n; idx += stride) {
        C[idx] = A[idx] * s;
    }
}

}  // namespace

torch::Tensor scalar_mul_cuda(const torch::Tensor& A, double s) {
    if (C10_UNLIKELY(!A.is_cuda() || A.scalar_type() != at::kFloat)) {
        return at::mul(A, at::Scalar(s));
    }
    if (C10_UNLIKELY(!A.is_contiguous())) {
        // Recurse on the contiguous copy (binds safely to the temporary for
        // the duration of the call); the hot path below then runs unchanged.
        return scalar_mul_cuda(A.contiguous(), s);
    }

    // Dispatcher-free allocation of the output tensor.
    torch::Tensor C = at::detail::empty_cuda(
        A.sizes(), at::kFloat, A.device(), std::nullopt);

    const long n = (long)A.numel();
    if (C10_UNLIKELY(n == 0)) {
        return C;
    }

    const float* a_ptr = A.data_ptr<float>();
    float*       c_ptr = static_cast<float*>(C.data_ptr());
    const float  sf    = static_cast<float>(s);

    cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();

    const bool aligned =
        ((((uintptr_t)a_ptr) | ((uintptr_t)c_ptr)) & 0xF) == 0;
    const long n4 = n >> 2;

    if (C10_LIKELY(aligned && (n & 3) == 0 && n4 <= (long)INT_MAX)) {
        if ((n4 % kTile) == 0) {
            const int blocks = (int)(n4 / kTile);
            scalar_mul_vec4_uniform_kernel<<<blocks, kThreads, 0, stream>>>(
                reinterpret_cast<const float4*>(a_ptr),
                reinterpret_cast<float4*>(c_ptr),
                sf);
        } else {
            const int blocks = (int)((n4 + kThreads - 1) / kThreads);
            scalar_mul_vec4_kernel<<<blocks, kThreads, 0, stream>>>(
                reinterpret_cast<const float4*>(a_ptr),
                reinterpret_cast<float4*>(c_ptr),
                sf, (int)n4);
        }
    } else {
        long blocks = (n + kThreads - 1) / kThreads;
        const long cap = 1L << 20;
        if (blocks > cap) blocks = cap;
        scalar_mul_fallback_kernel<<<(unsigned int)blocks, kThreads, 0, stream>>>(
            a_ptr, c_ptr, sf, n);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

scalar_mul_cpp_source = "torch::Tensor scalar_mul_cuda(const torch::Tensor& A, double s);"

# Compile the inline CUDA extension
scalar_mul_ext = load_inline(
    name="scalar_mul_ext_v6",
    cpp_sources=scalar_mul_cpp_source,
    cuda_sources=scalar_mul_cuda_source,
    functions=["scalar_mul_cuda"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs a matrix-scalar multiplication (C = A * s)
    using a custom, vectorized, bandwidth-maximizing CUDA kernel with minimal
    host-side dispatch overhead.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self._scalar_mul = scalar_mul_ext.scalar_mul_cuda

    def forward(self, A: torch.Tensor, s: float) -> torch.Tensor:
        return self._scalar_mul(A, s)
```