```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

_cpp_source = """
torch::Tensor diagonal_matmul_cuda(torch::Tensor A, torch::Tensor B);
"""

_cuda_source = r"""
#include <torch/extension.h>
#include <ATen/Dispatch.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>

__global__ void scale_rows_float4(
    const float* __restrict__ A,
    const float4* __restrict__ B,
    float4* __restrict__ C,
    int64_t vectors_per_row,
    int64_t tiles_per_row)
{
    const int64_t row = blockIdx.x / tiles_per_row;
    const int64_t col =
        (blockIdx.x % tiles_per_row) * blockDim.x + threadIdx.x;

    if (col < vectors_per_row) {
        const float scale = A[row];
        const int64_t index = row * vectors_per_row + col;
        const float4 value = B[index];
        C[index] = make_float4(
            scale * value.x, scale * value.y,
            scale * value.z, scale * value.w);
    }
}

template <typename scalar_t>
__global__ void scale_rows_scalar(
    const scalar_t* __restrict__ A,
    const scalar_t* __restrict__ B,
    scalar_t* __restrict__ C,
    int64_t columns,
    int64_t elements)
{
    for (int64_t index =
             static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
         index < elements;
         index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
        C[index] = A[index / columns] * B[index];
    }
}

torch::Tensor diagonal_matmul_cuda(torch::Tensor A, torch::Tensor B)
{
    TORCH_CHECK(A.is_cuda() && B.is_cuda(), "Inputs must be CUDA tensors");
    TORCH_CHECK(A.device() == B.device(), "Inputs must share a device");
    TORCH_CHECK(A.scalar_type() == B.scalar_type(),
                "Inputs must have matching dtypes");
    TORCH_CHECK(A.dim() == 1 && B.dim() == 2,
                "Expected A with shape (N,) and B with shape (N, M)");
    TORCH_CHECK(A.size(0) == B.size(0), "Incompatible input shapes");

    const c10::cuda::CUDAGuard device_guard(B.device());
    auto a = A.contiguous();
    auto b = B.contiguous();
    auto output = torch::empty(b.sizes(), b.options());

    const int64_t elements = b.numel();
    if (elements == 0) {
        return output;
    }

    const int threads = 256;
    const int64_t columns = b.size(1);
    const auto stream = c10::cuda::getCurrentCUDAStream();

    if (b.scalar_type() == at::kFloat && columns % 4 == 0 &&
        reinterpret_cast<uintptr_t>(b.data_ptr<float>()) % 16 == 0) {
        const int64_t vectors = columns / 4;
        const int64_t tiles = (vectors + threads - 1) / threads;
        const int64_t blocks = b.size(0) * tiles;

        scale_rows_float4<<<blocks, threads, 0, stream.stream()>>>(
            a.data_ptr<float>(),
            reinterpret_cast<const float4*>(b.data_ptr<float>()),
            reinterpret_cast<float4*>(output.data_ptr<float>()),
            vectors,
            tiles);
    } else {
        const int64_t required_blocks = (elements + threads - 1) / threads;
        const int blocks = static_cast<int>(
            required_blocks < 65535 ? required_blocks : 65535);

        AT_DISPATCH_FLOATING_TYPES_AND2(
            at::ScalarType::Half,
            at::ScalarType::BFloat16,
            b.scalar_type(),
            "diagonal_matmul_cuda",
            [&] {
                scale_rows_scalar<scalar_t>
                    <<<blocks, threads, 0, stream.stream()>>>(
                        a.data_ptr<scalar_t>(),
                        b.data_ptr<scalar_t>(),
                        output.data_ptr<scalar_t>(),
                        columns,
                        elements);
            });
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return output;
}
"""

_diagonal_matmul = load_inline(
    name="diagonal_matmul_row_scale",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["diagonal_matmul_cuda"],
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
    verbose=False,
)


class ModelNew(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, A, B):
        return _diagonal_matmul.diagonal_matmul_cuda(A, B)
```