```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 op is purely memory-bandwidth bound, so the kernel is tuned for
# maximum sustained DRAM throughput:
#   * 128-bit vectorized loads/stores (float4)
#   * 8 independent float4 loads in flight per thread (deep memory-level
#     parallelism to hide latency and saturate the memory controllers)
#   * streaming cache hints (__ldcs / __stcs) since there is zero data reuse
#   * exact grid coverage (no grid-stride loop overhead); only the final
#     block takes a bounds-checked tail path
#   * minimal host-side work (no dispatch, no TensorIterator setup)
scalar_mul_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cstdint>
#include <climits>

namespace {

constexpr int  kThreads    = 256;
constexpr int  kUnroll     = 8;                          // float4 per thread
constexpr long kTileFloat4 = (long)kThreads * kUnroll;   // float4 per block

__global__ void __launch_bounds__(kThreads)
scalar_mul_vec4_kernel(const float4* __restrict__ A,
                       float4* __restrict__ C,
                       const float s,
                       const long n4) {
    const long base = (long)blockIdx.x * kTileFloat4;
    const int  t    = threadIdx.x;

    if (base + kTileFloat4 <= n4) {
        // Fast path: full tile, no bounds checks, 8 outstanding 16B loads.
        float4 v[kUnroll];
        #pragma unroll
        for (int i = 0; i < kUnroll; ++i) {
            v[i] = __ldcs(A + (base + t + 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) {
            __stcs(C + (base + t + i * kThreads), v[i]);
        }
    } else {
        // Tail path (at most one block): bounds-checked.
        for (long idx = base + t; idx < n4; idx += kThreads) {
            float4 v = __ldcs(A + idx);
            v.x *= s;
            v.y *= s;
            v.z *= s;
            v.w *= s;
            __stcs(C + idx, v);
        }
    }
}

// Scalar grid-stride fallback for unaligned / non-multiple-of-4 sizes.
__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(torch::Tensor A, double s) {
    if (!A.is_cuda() || A.scalar_type() != at::kFloat) {
        return at::mul(A, at::Scalar(s));
    }

    auto Ac = A.is_contiguous() ? A : A.contiguous();
    auto C  = at::empty_like(Ac);

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

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

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

    const bool aligned =
        (((reinterpret_cast<uintptr_t>(a_ptr)) | (reinterpret_cast<uintptr_t>(c_ptr))) & 0xF) == 0;

    if (aligned && ((n & 3) == 0)) {
        const long n4     = n >> 2;
        const long blocks = (n4 + kTileFloat4 - 1) / kTileFloat4;
        if (blocks <= (long)INT_MAX) {
            scalar_mul_vec4_kernel<<<(unsigned int)blocks, kThreads, 0, stream>>>(
                reinterpret_cast<const float4*>(a_ptr),
                reinterpret_cast<float4*>(c_ptr),
                sf, n4);
            C10_CUDA_KERNEL_LAUNCH_CHECK();
            return C;
        }
    }

    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(torch::Tensor A, double s);"

# Compile the inline CUDA extension
scalar_mul_ext = load_inline(
    name="scalar_mul_ext_v3",
    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", "--use_fast_math"],
)


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

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