```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 that computes diag(A) @ B without materializing the
# N x N diagonal matrix. Mathematically, (diag(A) @ B)[i, j] = A[i] * B[i, j],
# so each row of B is simply scaled by the corresponding diagonal entry.
# This reduces the work from O(N^2 * M) to O(N * M) and is purely memory-bound.
diag_matmul_source = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cstdint>

// Scalar fallback kernel: one thread per output element.
__global__ void diag_matmul_kernel(const float* __restrict__ A,
                                   const float* __restrict__ B,
                                   float* __restrict__ out,
                                   int N, int M) {
    int row = blockIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < N && col < M) {
        int64_t idx = (int64_t)row * M + col;
        out[idx] = A[row] * B[idx];
    }
}

// Vectorized kernel: each thread handles 4 consecutive floats (16 bytes)
// for maximal memory throughput. Requires M % 4 == 0.
__global__ void diag_matmul_vec4_kernel(const float* __restrict__ A,
                                        const float* __restrict__ B,
                                        float* __restrict__ out,
                                        int N, int M4) {
    int row = blockIdx.y;
    int col4 = blockIdx.x * blockDim.x + threadIdx.x;
    if (row < N && col4 < M4) {
        float a = A[row];
        int64_t base = (int64_t)row * ((int64_t)M4 * 4) + (int64_t)col4 * 4;
        float4 b = *reinterpret_cast<const float4*>(B + base);
        float4 r;
        r.x = a * b.x;
        r.y = a * b.y;
        r.z = a * b.z;
        r.w = a * b.w;
        *reinterpret_cast<float4*>(out + base) = r;
    }
}

torch::Tensor diag_matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.dim() == 1, "A must be a 1D tensor");
    TORCH_CHECK(B.dim() == 2, "B must be a 2D tensor");
    TORCH_CHECK(A.scalar_type() == torch::kFloat32 &&
                B.scalar_type() == torch::kFloat32,
                "Inputs must be float32");

    auto A_c = A.contiguous();
    auto B_c = B.contiguous();

    const int N = A_c.size(0);
    const int M = B_c.size(1);
    TORCH_CHECK(B_c.size(0) == N, "Dimension mismatch between A and B");

    auto out = torch::empty_like(B_c);
    if (out.numel() == 0) {
        return out;
    }

    const int block_size = 256;

    if (M % 4 == 0) {
        int M4 = M / 4;
        dim3 grid((M4 + block_size - 1) / block_size, N);
        diag_matmul_vec4_kernel<<<grid, block_size>>>(
            A_c.data_ptr<float>(), B_c.data_ptr<float>(),
            out.data_ptr<float>(), N, M4);
    } else {
        dim3 grid((M + block_size - 1) / block_size, N);
        diag_matmul_kernel<<<grid, block_size>>>(
            A_c.data_ptr<float>(), B_c.data_ptr<float>(),
            out.data_ptr<float>(), N, M);
    }

    return out;
}
"""

diag_matmul_cpp_source = (
    "torch::Tensor diag_matmul_cuda(torch::Tensor A, torch::Tensor B);"
)

# Compile the inline CUDA extension
diag_matmul = load_inline(
    name="diag_matmul",
    cpp_sources=diag_matmul_cpp_source,
    cuda_sources=diag_matmul_source,
    functions=["diag_matmul_cuda"],
    verbose=False,
    extra_cuda_cflags=["-O3"],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = diag(A) * B using a custom CUDA kernel
    that directly scales each row of B by the corresponding element of A,
    avoiding the explicit construction of the diagonal matrix and the full
    O(N^2 * M) matrix multiplication.
    """
    def __init__(self):
        super(ModelNew, self).__init__()
        self.diag_matmul = diag_matmul

    def forward(self, A, B):
        if A.is_cuda and B.is_cuda and A.dtype == torch.float32 and B.dtype == torch.float32:
            return self.diag_matmul.diag_matmul_cuda(A, B)
        # Fallback for non-CUDA / non-float32 inputs (mathematically identical).
        return A.unsqueeze(1) * B
```