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

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

__global__ void diagonal_scale_float4(
    const float* __restrict__ a,
    const float4* __restrict__ b,
    float4* __restrict__ out,
    int64_t vectors,
    int64_t vectors_per_row
) {
    for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
         i < vectors;
         i += static_cast<int64_t>(blockDim.x) * gridDim.x) {
        float scale = a[i / vectors_per_row];
        float4 value = b[i];
        out[i] = make_float4(
            scale * value.x,
            scale * value.y,
            scale * value.z,
            scale * value.w
        );
    }
}

template <typename scalar_t>
__global__ void diagonal_scale_generic(
    const scalar_t* __restrict__ a,
    const scalar_t* __restrict__ b,
    scalar_t* __restrict__ out,
    int64_t count,
    int64_t columns
) {
    for (int64_t i = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
         i < count;
         i += static_cast<int64_t>(blockDim.x) * gridDim.x) {
        out[i] = a[i / columns] * b[i];
    }
}

torch::Tensor diagonal_scale_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 share a dtype");
    TORCH_CHECK(a.dim() == 1 && b.dim() == 2, "Expected A[N] and B[N, M]");
    TORCH_CHECK(a.size(0) == b.size(0), "Row dimensions must match");

    c10::cuda::CUDAGuard guard(b.device());
    a = a.contiguous();
    b = b.contiguous();
    auto out = torch::empty(b.sizes(), b.options());
    const int64_t count = b.numel();
    if (count == 0) {
        return out;
    }

    constexpr int threads = 256;
    cudaStream_t stream = at::cuda::getCurrentCUDAStream();

    if (b.scalar_type() == at::kFloat &&
        b.size(1) % 4 == 0 &&
        reinterpret_cast<uintptr_t>(b.data_ptr<float>()) % 16 == 0) {
        const int64_t vectors = count / 4;
        const int blocks = static_cast<int>(
            std::min<int64_t>((vectors + threads - 1) / threads, 4096)
        );
        diagonal_scale_float4<<<blocks, threads, 0, stream>>>(
            a.data_ptr<float>(),
            reinterpret_cast<const float4*>(b.data_ptr<float>()),
            reinterpret_cast<float4*>(out.data_ptr<float>()),
            vectors,
            b.size(1) / 4
        );
    } else {
        const int blocks = static_cast<int>(
            std::min<int64_t>((count + threads - 1) / threads, 4096)
        );
        AT_DISPATCH_FLOATING_TYPES_AND2(
            at::ScalarType::Half,
            at::ScalarType::BFloat16,
            b.scalar_type(),
            "diagonal_scale_cuda",
            [&] {
                diagonal_scale_generic<scalar_t><<<blocks, threads, 0, stream>>>(
                    a.data_ptr<scalar_t>(),
                    b.data_ptr<scalar_t>(),
                    out.data_ptr<scalar_t>(),
                    count,
                    b.size(1)
                );
            }
        );
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return out;
}
"""

_diagonal_scale = load_inline(
    name="diagonal_scale_cuda_extension",
    cpp_sources=(
        "torch::Tensor diagonal_scale_cuda(torch::Tensor a, torch::Tensor b);"
    ),
    cuda_sources=_cuda_source,
    functions=["diagonal_scale_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_scale.diagonal_scale_cuda(A, B)
```