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

# Key insight: torch.diag(A) @ B simply scales row i of B by A[i].
# There is no need to materialize the N x N diagonal matrix or run an
# O(N*N*M) GEMM -- a single O(N*M) memory-bound scaling kernel suffices.

diag_matmul_cuda_source = """
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

// Vectorized (float4) kernel: C = diag(A) @ B  <=>  C[i, :] = A[i] * B[i, :]
__global__ void diag_matmul_vec4_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    long long N,
    long long M) {
    long long row = blockIdx.x;
    float a = A[row];
    long long col = ((long long)blockIdx.y * blockDim.x + threadIdx.x) * 4;
    if (col < M) {
        long long offset = row * M + col;
        float4 b = *reinterpret_cast<const float4*>(B + offset);
        float4 c;
        c.x = a * b.x;
        c.y = a * b.y;
        c.z = a * b.z;
        c.w = a * b.w;
        *reinterpret_cast<float4*>(C + offset) = c;
    }
}

// Scalar fallback kernel that handles any M.
__global__ void diag_matmul_scalar_kernel(
    const float* __restrict__ A,
    const float* __restrict__ B,
    float* __restrict__ C,
    long long N,
    long long M) {
    long long row = blockIdx.x;
    float a = A[row];
    long long col = (long long)blockIdx.y * blockDim.x + threadIdx.x;
    if (col < M) {
        long long offset = row * M + col;
        C[offset] = a * B[offset];
    }
}

torch::Tensor diag_matmul_cuda(torch::Tensor A, torch::Tensor B) {
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "A and B 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.size(0) == B.size(0), "A's length must equal B's number of rows");
    TORCH_CHECK(A.scalar_type() == at::kFloat && B.scalar_type() == at::kFloat,
                "A and B must be float32");

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

    const long long N = B_c.size(0);
    const long long M = B_c.size(1);

    auto C = torch::empty_like(B_c);
    if (N == 0 || M == 0) {
        return C;
    }

    const int block_size = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (M % 4 == 0) {
        const long long vec_cols = M / 4;
        dim3 grid((unsigned int)N,
                  (unsigned int)((vec_cols + block_size - 1) / block_size));
        diag_matmul_vec4_kernel<<<grid, block_size, 0, stream>>>(
            A_c.data_ptr<float>(), B_c.data_ptr<float>(), C.data_ptr<float>(), N, M);
    } else {
        dim3 grid((unsigned int)N,
                  (unsigned int)((M + block_size - 1) / block_size));
        diag_matmul_scalar_kernel<<<grid, block_size, 0, stream>>>(
            A_c.data_ptr<float>(), B_c.data_ptr<float>(), C.data_ptr<float>(), N, M);
    }
    C10_CUDA_KERNEL_LAUNCH_CHECK();

    return C;
}
"""

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

# Compile the inline CUDA kernel
diag_matmul = load_inline(
    name="diag_matmul",
    cpp_sources=diag_matmul_cpp_source,
    cuda_sources=diag_matmul_cuda_source,
    functions=["diag_matmul_cuda"],
    verbose=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    extra_ldflags=[""],
)


class ModelNew(nn.Module):
    """
    Optimized model that performs C = diag(A) @ B using a custom CUDA kernel.
    Since left-multiplying by a diagonal matrix only scales each row of B by
    the corresponding diagonal entry, the full O(N^2 * M) GEMM is replaced by
    an O(N * M) memory-bound element-wise row-scaling kernel.
    """
    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:
            return self.diag_matmul.diag_matmul_cuda(A, B)
        # CPU fallback that is still much faster than materializing diag(A).
        return B * A.unsqueeze(1)
```