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

_cpp_source = r"""
#include <torch/extension.h>
torch::Tensor diagonal_mul_cuda(torch::Tensor A, torch::Tensor B);
"""

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

__device__ __forceinline__ float4 load_streaming(const float4* ptr) {
    float4 value;
    asm volatile(
        "ld.global.cg.v4.f32 {%0, %1, %2, %3}, [%4];"
        : "=f"(value.x), "=f"(value.y), "=f"(value.z), "=f"(value.w)
        : "l"(ptr)
        : "memory");
    return value;
}

__device__ __forceinline__ void store_streaming(
    float4* ptr, const float4 value)
{
    asm volatile(
        "st.global.cs.v4.f32 [%0], {%1, %2, %3, %4};"
        :
        : "l"(ptr), "f"(value.x), "f"(value.y),
          "f"(value.z), "f"(value.w)
        : "memory");
}

__global__ __launch_bounds__(256)
void diagonal_mul_persistent_4096(
    const float* __restrict__ A,
    const float4* __restrict__ B,
    float4* __restrict__ C,
    int64_t rows)
{
    for (int64_t row = blockIdx.x; row < rows; row += gridDim.x) {
        const int64_t offset = row * 1024 + threadIdx.x;
        const float scale = __ldg(A + row);

        float4 values[4];
        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            values[j] = load_streaming(B + offset + j * 256);
        }

        #pragma unroll
        for (int j = 0; j < 4; ++j) {
            values[j].x *= scale;
            values[j].y *= scale;
            values[j].z *= scale;
            values[j].w *= scale;
            store_streaming(C + offset + j * 256, values[j]);
        }
    }
}

__global__ void diagonal_mul_float4(
    const float* __restrict__ A,
    const float4* __restrict__ B,
    float4* __restrict__ C,
    int64_t count,
    int64_t vectors_per_row)
{
    for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         i < count;
         i += int64_t(blockDim.x) * gridDim.x) {
        const float scale = A[i / vectors_per_row];
        const float4 value = B[i];
        C[i] = make_float4(
            scale * value.x, scale * value.y,
            scale * value.z, scale * value.w);
    }
}

template <typename scalar_t>
__global__ void diagonal_mul_scalar(
    const scalar_t* __restrict__ A,
    const scalar_t* __restrict__ B,
    scalar_t* __restrict__ C,
    int64_t count,
    int64_t columns)
{
    using acc_t = at::acc_type<scalar_t, true>;
    for (int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
         i < count;
         i += int64_t(blockDim.x) * gridDim.x) {
        C[i] = static_cast<scalar_t>(
            static_cast<acc_t>(A[i / columns]) *
            static_cast<acc_t>(B[i]));
    }
}

torch::Tensor diagonal_mul_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 && A.size(0) == B.size(0),
                "Expected A: (N,), B: (N, M)");

    const c10::cuda::CUDAGuard guard(B.device());
    A = A.contiguous();
    B = B.contiguous();
    auto C = torch::empty(B.sizes(), B.options());

    const int64_t count = B.numel();
    if (count == 0) return C;

    constexpr int threads = 256;
    const int64_t columns = B.size(1);
    const auto stream = at::cuda::getCurrentCUDAStream();

    if (B.scalar_type() == at::kFloat &&
        columns % 4 == 0 &&
        reinterpret_cast<uintptr_t>(B.data_ptr<float>()) % 16 == 0 &&
        reinterpret_cast<uintptr_t>(C.data_ptr<float>()) % 16 == 0) {
        const float* a = A.data_ptr<float>();
        const float4* b =
            reinterpret_cast<const float4*>(B.data_ptr<float>());
        float4* c = reinterpret_cast<float4*>(C.data_ptr<float>());

        if (columns == 4096) {
            const int blocks = static_cast<int>(
                std::min<int64_t>(B.size(0), 512));
            diagonal_mul_persistent_4096<<<blocks, threads, 0, stream>>>(
                a, b, c, B.size(0));
        } else {
            const int64_t vectors = count / 4;
            const int blocks = static_cast<int>(
                std::min<int64_t>((vectors + threads - 1) / threads, 65535));
            diagonal_mul_float4<<<blocks, threads, 0, stream>>>(
                a, b, c, vectors, columns / 4);
        }
    } else {
        const int blocks = static_cast<int>(
            std::min<int64_t>((count + threads - 1) / threads, 65535));
        AT_DISPATCH_FLOATING_TYPES_AND2(
            at::ScalarType::Half, at::ScalarType::BFloat16,
            B.scalar_type(), "diagonal_mul_cuda", [&] {
                diagonal_mul_scalar<scalar_t>
                    <<<blocks, threads, 0, stream>>>(
                        A.data_ptr<scalar_t>(),
                        B.data_ptr<scalar_t>(),
                        C.data_ptr<scalar_t>(),
                        count, columns);
            });
    }

    C10_CUDA_KERNEL_LAUNCH_CHECK();
    return C;
}
"""

_diagonal_mul = load_inline(
    name="diagonal_row_scale_streaming_v3",
    cpp_sources=_cpp_source,
    cuda_sources=_cuda_source,
    functions=["diagonal_mul_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_mul.diagonal_mul_cuda(A, B)
```