# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License

"""CUDA weight-only quantization helpers.

This module contains small Python utilities for producing the weight layouts
consumed by CUDA weight-only kernels. The blockwise quantizers wrap the same C++
pybind entry points used by runtime prepacking, and the mixed-GEMM weight packer
is a PyTorch reimplementation of the runtime CUDA packing, so tests and model
builders can generate byte-identical quantized weights. The PyTorch packer runs
on CUDA when a device is available and falls back to CPU otherwise, which is the
only option on platforms where the standalone CUDA packer is not built (Windows).
A GPU-gated parity test validates it against that standalone CUDA packer.

Two storage families are exposed:

* raw MatMulNBits blockwise storage, laid out by output channel as ``[N, K/pack]``;
* CUDA mixed-GEMM prepacked storage. MatMulNBits prepacked initializers keep the
    schema shape ``[N, K/block_size, block_size*bits/8]`` and require the node
    attribute ``weight_prepacked=1``. QMoE/CUTLASS callers use the kernel-facing
    shape ``[K, N/pack]``.

All public helpers take one logical expert weight matrix with shape ``[N, K]``.
"""

from __future__ import annotations

import functools
import logging
from typing import TYPE_CHECKING

import numpy as np

_logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    import torch


def _get_torch():
    """Import torch lazily so importing the quantizer does not require torch."""
    try:
        import torch  # noqa: PLC0415
    except ImportError as e:
        raise ImportError("CUDA weight-only quantization requires torch. Please install torch to use it.") from e

    return torch


def has_cuda_weight_prepacking() -> bool:
    """Return True if mixed-GEMM weight prepacking is available.

    Prepacking is implemented with PyTorch (CUDA when available, CPU otherwise), so it is
    available whenever torch is importable. Callers use this to skip prepack code paths
    (and tests) when torch is unavailable.
    """
    try:
        _get_torch()
    except ImportError:
        return False
    return True


@functools.lru_cache(maxsize=1)
def _warn_cpu_prepack_once() -> None:
    _logger.warning(
        "CUDA device is not available; packing mixed-GEMM weights on CPU with PyTorch. "
        "This is correct but significantly slower for large Mixture-of-Experts models. "
        "Pack on a CUDA-enabled machine for best performance."
    )


def _prepack_device():
    """Pick the torch device for mixed-GEMM weight packing (CUDA if available, else CPU)."""
    torch = _get_torch()
    if torch.cuda.is_available():
        return torch.device("cuda")
    _warn_cpu_prepack_once()
    return torch.device("cpu")


def _preprocess_weights_for_mixed_gemm_torch(tensor, bits: int, sm: int):
    """PyTorch port of the runtime CUDA ``preprocess_weights_for_mixed_gemm``.

    ``tensor`` is a signed int8 weight in ``(K, N/pack)`` packed row-major layout on any
    device. Returns the CUTLASS mixed-GEMM layout with the same shape/dtype/device. This
    mirrors ``preprocess_weights_for_mixed_gemm_cuda`` (permute_B_rows -> subbyte_transpose
    -> interleave_column_major -> add_bias_and_interleave) so its output is byte-identical
    to the standalone CUDA packer, for both the SM80 (Ampere) and SM90 (Hopper) layouts.
    """
    torch = _get_torch()
    bits_a = 16  # fp16/bf16 activations
    bits_b = 4 if bits == 4 else 8

    if tensor.dim() == 2:
        tensor = tensor.unsqueeze(0)

    permutation_map = {
        "16_8": [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15],
        "16_4": [0, 1, 8, 9, 16, 17, 24, 25, 2, 3, 10, 11, 18, 19, 26, 27, 4, 5, 12, 13, 20, 21, 28, 29, 6, 7, 14, 15, 22, 23, 30, 31],  # fmt: skip
    }
    mma_shape_n = 8
    b_rows_per_mma = 8 * 16 // bits_b

    num_experts, num_rows, num_cols = tensor.shape[0], tensor.shape[1], tensor.shape[2]
    if num_rows % b_rows_per_mma != 0 or num_cols % mma_shape_n != 0:
        raise ValueError(
            f"weight shape (rows={num_rows}, packed_cols={num_cols}) is incompatible with mixed-GEMM "
            f"packing (rows must be a multiple of {b_rows_per_mma}, packed cols a multiple of {mma_shape_n})."
        )

    # permute_B_rows_for_mixed_gemm
    if sm < 100:
        pmap = permutation_map[f"{bits_a}_{bits_b}"]
        row_idx = [(r // b_rows_per_mma) * b_rows_per_mma + pmap[r % b_rows_per_mma] for r in range(num_rows)]
        tensor = tensor[:, row_idx, :]

    # subbyte_transpose
    original_shape = tensor.shape
    if bits_b == 4:
        u = tensor.view(torch.uint8)
        high = (u >> 4).permute(0, 2, 1).unsqueeze(2)
        low = ((u << 4) >> 4).permute(0, 2, 1).unsqueeze(2)
        merged = torch.cat([low, high], dim=2).reshape(u.shape[0], -1, u.shape[1])
        merged = merged[:, :, 0::2] + merged[:, :, 1::2] * 16
        tensor = merged.view(torch.int8).reshape(original_shape)
    else:
        tensor = tensor.permute(0, 2, 1).reshape(original_shape)

    # interleave_column_major_tensor
    interleave = bits_a // bits_b
    if interleave > 1 and sm < 90:
        rows_per_tile = 128 * 8 // bits_a
        elts_in_int32 = 32 // bits_b
        if num_rows % elts_in_int32 != 0 or num_rows % rows_per_tile != 0:
            raise ValueError(f"num_rows ({num_rows}) is incompatible with column-interleave tiling.")
        tensor = tensor.reshape(
            num_experts, -1, interleave, num_rows // rows_per_tile, rows_per_tile * 4 // elts_in_int32
        )
        tensor = tensor.permute(0, 1, 3, 2, 4).reshape(original_shape)

    # add_bias_and_interleave_quantized_tensor_inplace
    if bits_b == 8:
        t = tensor.to(torch.int64)  # widen so the +128 rebias cannot overflow int8
        t += -256 * (t > 127).to(torch.int64) + 128
        t = t.reshape(-1, 4)[:, [0, 2, 1, 3]].reshape(original_shape)
        tensor = t.to(torch.uint8).view(torch.int8)
    else:
        u = tensor.view(torch.uint8)
        high = (u >> 4).unsqueeze(-1)
        low = ((u << 4) >> 4).unsqueeze(-1)
        merged = torch.cat([low, high], dim=-1).reshape(u.shape[0], u.shape[1], -1)
        merged = merged.reshape(-1, 8)[:, [0, 2, 4, 6, 1, 3, 5, 7]].reshape(merged.shape)
        merged = merged.to(torch.int16)
        merged += -16 * (merged > 7).to(torch.int16) + 8
        merged = merged[:, :, 0::2] + merged[:, :, 1::2] * 16
        tensor = merged.to(torch.uint8).view(torch.int8)

    return tensor.squeeze(0).contiguous()


def _pack_weights_for_cuda_mixed_gemm(q_weights, n: int, k: int, bits: int, force_arch: int = 80) -> np.ndarray:
    """PyTorch implementation of the CUDA ``pack_weights_for_cuda_mixed_gemm``.

    ``q_weights`` is ORT's unsigned MatMulNBits/QMoE storage ``(N, K/pack)`` (uint8). Returns
    a flat ``int8`` numpy array with the CUTLASS mixed-GEMM layout, byte-identical to the
    standalone CUDA packer. Runs on CUDA when available, otherwise on CPU.
    """
    torch = _get_torch()
    bits = int(bits)
    force_arch = int(force_arch)
    if bits not in (4, 8):
        raise ValueError(f"bits must be 4 or 8, got {bits}.")
    if force_arch not in (80, 90):
        raise ValueError(f"force_arch must be 80 (SM80) or 90 (SM90), got {force_arch}.")
    pack = 8 // bits
    device = _prepack_device()

    q = torch.as_tensor(np.ascontiguousarray(q_weights)).view(torch.uint8).reshape(n, k // pack).to(device)

    # Front-end adaptor: transpose ORT (N, K) -> (K, N) and convert unsigned -> signed int8.
    if bits == 4:
        low = (q & 0x0F).to(torch.int16)
        high = (q >> 4).to(torch.int16)
        unpacked = torch.empty((n, k), dtype=torch.int16, device=device)
        unpacked[:, 0::2] = low
        unpacked[:, 1::2] = high
        signed_t = (unpacked - 8).transpose(0, 1).contiguous()  # (K, N), zero point 8
        packed_t = ((signed_t[:, 0::2] & 0x0F) | ((signed_t[:, 1::2] & 0x0F) << 4)).to(torch.uint8).view(torch.int8)
    else:
        signed_t = (q.to(torch.int16) - 128).transpose(0, 1).contiguous()  # (K, N), zero point 128
        packed_t = signed_t.to(torch.uint8).view(torch.int8)

    out = _preprocess_weights_for_mixed_gemm_torch(packed_t.contiguous(), bits, force_arch)
    return out.reshape(-1).cpu().numpy()


def _get_quantize_matmul_nbits():
    """Return MatMulNBits blockwise quantizers from the ORT pybind module."""
    try:
        from onnxruntime.capi._pybind_state import (  # noqa: PLC0415
            quantize_matmul_4bits,
            quantize_matmul_8bits,
        )
    except ImportError as e:
        raise ImportError(
            "CUDA blockwise quantization requires quantize_matmul_4bits and quantize_matmul_8bits from onnxruntime."
        ) from e

    return quantize_matmul_4bits, quantize_matmul_8bits


class CudaQuantizer:
    """CUDA quantizer utilities for MoE/QMoE and MatMulNBits-style weight-only kernels.

    The methods are stateless; callers may use the class directly without
    constructing an object.
    """

    @staticmethod
    def qmoe_symmetric_per_channel_quantize(
        weights: torch.Tensor,
        bits: int,
        *,
        unsigned_full_range: bool = True,
        signed_scale: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Quantize one QMoE expert with symmetric per-channel storage.

        ``weights`` has logical shape ``[N, K]``. Returns raw QMoE storage
        ``[N, K/pack]`` and scales ``[N]``. By default, this emits the ORT CUDA
        QMoE storage contract: unsigned bytes/nibbles with an implicit zero-point
        offset, so each stored value is ``q + zero_point`` even though the numeric
        quantization is symmetric. By default it uses the full ``[-8, 7]`` /
        ``[-128, 127]`` range. Set ``unsigned_full_range=False`` to use the legacy
        ``[-7, 7]`` / ``[-127, 127]`` range.

        With ``signed_scale=False`` (**default for the per-channel path**) the
        scale is the positive ``max|w| / 2^(bits-1)``, so the negative extreme maps
        exactly to ``qmin`` while the positive extreme is clipped at ``qmax``. Set
        ``signed_scale=True`` for the MLAS ``default`` convention: the scale is
        SIGNED, mapping the max-magnitude element (with its sign) exactly to
        ``qmin`` (``scale = w[argmax|w|] / qmin``), representing whichever extreme
        is the true maximum without clipping. The stored bytes/zero-point layout is
        unchanged; the sign rides in the scale, which the CUDA QMoE kernel consumes
        correctly (parity-verified byte-for-byte).

        NOTE: ``signed_scale`` defaults to ``False`` here on purpose. For gpt-oss
        the signed per-channel variant gives *no* reconstruction-error benefit (the
        per-row max-magnitude weight is the only element that differs, so mean|Δw|
        is a wash) and empirically it costs ~2pp MMLU vs the positive-scale default
        -- a genuine data-dependent quantization effect, not a kernel bug (the
        kernel computes signed scales correctly). The signed convention is only
        worthwhile on the finer-grained blockwise paths, where it matches MLAS
        ``default`` and is MMLU-neutral.
        """
        torch = _get_torch()

        weights = weights.detach().cpu().to(torch.float32).contiguous()
        bits = int(bits)
        if bits not in (4, 8):
            raise ValueError(f"QMoE per-channel quantization only supports 4 or 8 bits, got {bits}.")

        n, k = weights.shape
        pack = 8 // bits
        if k % pack != 0:
            raise ValueError(f"K ({k}) must be divisible by {pack} for QMoE per-channel quantization.")

        if bits == 4:
            if unsigned_full_range:
                qmin, qmax, scale_divisor, zero_point = -8, 7, 8, 8
            else:
                qmin, qmax, scale_divisor, zero_point = -7, 7, 7, 8
        else:  # bits == 8, already validated above
            if unsigned_full_range:
                qmin, qmax, scale_divisor, zero_point = -128, 127, 128, 128
            else:
                qmin, qmax, scale_divisor, zero_point = -127, 127, 127, 128
        eps = torch.finfo(torch.float32).eps
        if signed_scale:
            # MLAS-style signed scale: map the max-|w| element (with its sign) to qmin,
            # so the true extreme is represented exactly instead of being clipped.
            argmax = weights.abs().argmax(dim=1, keepdim=True)
            signed_max = torch.gather(weights, 1, argmax)  # [N, 1], keeps the sign
            scales = signed_max / float(qmin)
            # An all-zero row yields scale 0; fall back to a tiny positive scale.
            scales = torch.where(scales.abs() < eps, torch.full_like(scales, eps), scales)
        else:
            scales = weights.abs().amax(dim=1, keepdim=True) / float(scale_divisor)
            scales = torch.clamp(scales, min=eps)
        quantized = torch.clamp(torch.round(weights / scales), qmin, qmax).to(torch.int16).contiguous()
        quantized = (quantized + zero_point).to(torch.uint8)

        if bits == 4:
            qweight = (quantized[:, 0::2] & 0xF) | ((quantized[:, 1::2] & 0xF) << 4)
            qweight = qweight.to(torch.uint8)
        else:
            qweight = quantized

        return qweight.contiguous(), scales.squeeze(-1).contiguous()

    @staticmethod
    def qmoe_per_channel_quantize(
        weights: torch.Tensor,
        bits: int,
        prepack: bool,
        force_arch: int = 80,
        *,
        unsigned_full_range: bool = True,
        signed_scale: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Quantize per-channel QMoE weights and optionally CUTLASS-prepack them.

        When ``prepack`` is true, returned weights have shape ``[K, N/pack]``.
        Otherwise, returned weights keep raw per-channel storage ``[N, K/pack]``.
        Prepacking uses PyTorch (CUDA when available, CPU otherwise).

        ``signed_scale`` defaults to ``False`` for the per-channel path (positive
        ``max|w|/2^(bits-1)`` scales); see
        :meth:`qmoe_symmetric_per_channel_quantize` for why signed scales are not
        the default here.
        """
        torch = _get_torch()

        qweight, scales = CudaQuantizer.qmoe_symmetric_per_channel_quantize(
            weights,
            bits,
            unsigned_full_range=unsigned_full_range,
            signed_scale=signed_scale,
        )
        if not prepack:
            return qweight, scales

        n, k = weights.shape
        pack = 8 // int(bits)
        if n % pack != 0:
            raise ValueError(f"N ({n}) must be divisible by {pack} for CUDA QMoE prepacked weights.")

        packed = _pack_weights_for_cuda_mixed_gemm(qweight.numpy(), n, k, int(bits), force_arch)
        packed = np.asarray(packed).view(np.uint8).reshape(k, n // pack)
        return torch.from_numpy(np.ascontiguousarray(packed)), scales

    @staticmethod
    def _matmulnbits_blockwise_quantize_impl(
        weights: torch.Tensor,
        bits: int,
        block_size: int,
        *,
        symmetric: bool,
        abs_scales: bool,
        unsigned_full_range: bool,
        signed_scale: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Blockwise-quantize ``weights`` and return unflattened storage.

        The symmetric path uses a local NumPy MLAS-style implementation (so it can
        keep signed per-block scales); the asymmetric path delegates to the
        MatMulNBits ORT pybinds (``quantize_matmul_4bits``/``8bits``).
        """
        torch = _get_torch()

        bits = int(bits)
        block_size = int(block_size)
        w = weights.detach().cpu().to(torch.float32).contiguous().numpy()
        n, k = w.shape
        if bits not in (4, 8):
            raise ValueError(f"Blockwise quantization only supports 4 or 8 bits, got {bits}.")
        if block_size <= 0:
            raise ValueError(f"Blockwise quantization requires a positive block_size, got {block_size}.")
        if signed_scale and not symmetric:
            raise ValueError("signed_scale is only valid for symmetric blockwise quantization.")

        num_blocks = (k + block_size - 1) // block_size
        pack = 8 // bits
        blob_size = (block_size + pack - 1) // pack

        if symmetric:
            if bits == 4:
                qmin, qmax, scale_divisor, zero_point = (-8, 7, 8, 8) if unsigned_full_range else (-7, 7, 7, 8)
            else:
                qmin, qmax, scale_divisor, zero_point = (
                    (-128, 127, 128, 128) if unsigned_full_range else (-127, 127, 127, 128)
                )

            padded_k = num_blocks * block_size
            if padded_k != k:
                w = np.pad(w, ((0, 0), (0, padded_k - k)), "constant")

            blocked = w.reshape(n, num_blocks, block_size)
            eps = np.finfo(np.float32).eps
            if signed_scale:
                # MLAS-style signed scale: map each block's max-|w| element (with its
                # sign) exactly to qmin, so the true extreme is represented instead of
                # being clipped at qmax. The sign rides in the scale.
                argmax = np.argmax(np.abs(blocked), axis=2)  # [n, num_blocks]
                signed_max = np.take_along_axis(blocked, argmax[:, :, np.newaxis], axis=2)[:, :, 0]
                scales = (signed_max.astype(np.float32) / np.float32(qmin))
                # An all-zero block yields scale 0; fall back to a tiny positive scale.
                scales = np.where(np.abs(scales) < eps, np.float32(eps), scales)
            else:
                scales = np.max(np.abs(blocked), axis=2).astype(np.float32) / np.float32(scale_divisor)
                scales = np.maximum(scales, eps)
            quantized = np.clip(np.rint(blocked / scales[:, :, np.newaxis]), qmin, qmax).astype(np.int16)
            quantized = (quantized + zero_point).astype(np.uint8)

            if bits == 4:
                qweight = np.zeros((n, num_blocks, blob_size), dtype=np.uint8)
                qweight[:, :, : quantized[:, :, 0::2].shape[2]] = quantized[:, :, 0::2] & 0xF
                qweight[:, :, : quantized[:, :, 1::2].shape[2]] |= (quantized[:, :, 1::2] & 0xF) << 4
            else:
                qweight = quantized

            zero_points = np.zeros((n, (num_blocks + 1) // 2 if bits == 4 else num_blocks), dtype=np.uint8)
            return torch.from_numpy(qweight), torch.from_numpy(scales), torch.from_numpy(zero_points)

        w_t = np.ascontiguousarray(w.T)
        qweight = np.zeros((n, num_blocks, blob_size), dtype=np.uint8)
        scales = np.zeros((n, num_blocks), dtype=np.float32)
        zero_points = np.zeros((n, (num_blocks + 1) // 2 if bits == 4 else num_blocks), dtype=np.uint8)

        quantize_matmul_4bits, quantize_matmul_8bits = _get_quantize_matmul_nbits()
        quantize = quantize_matmul_4bits if bits == 4 else quantize_matmul_8bits
        quantize(qweight, w_t, scales, zero_points, block_size, n, k, symmetric)

        if abs_scales:
            scales = np.abs(scales)

        return torch.from_numpy(qweight), torch.from_numpy(scales), torch.from_numpy(zero_points)

    @staticmethod
    def matmulnbits_blockwise_quantize(
        weights: torch.Tensor,
        bits: int,
        block_size: int,
        *,
        symmetric: bool = True,
        return_zero_points: bool = False,
        abs_scales: bool = True,
        flatten_qweight: bool = True,
        unsigned_full_range: bool = True,
        signed_scale: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Quantize one expert with ONNX Runtime's MatMulNBits blockwise encoding.

        ``weights`` has logical shape ``[N, K]``. Returns raw flattened storage
        ``[N, ceil(K/block_size)*ceil(block_size/pack)]`` and block scales
        ``[N, ceil(K/block_size)]`` by default. Set ``flatten_qweight=False`` for
        the MatMulNBits initializer shape
        ``[N, ceil(K/block_size), ceil(block_size/pack)]``.
        Set ``return_zero_points=True`` to also return packed block zero-points.
        Symmetric quantization uses the full ``[-8, 7]`` / ``[-128, 127]`` range
        by default. Set ``unsigned_full_range=False`` to use the legacy
        ``[-7, 7]`` / ``[-127, 127]`` range.

        With ``signed_scale=True`` (**default for the blockwise path**) each
        block's scale is SIGNED -- the MLAS ``default`` convention -- mapping that
        block's max-magnitude element (with its sign) exactly to ``qmin``
        (``scale = block[argmax|w|] / qmin``), so the true extreme is represented
        without clipping. This matches ORT MLAS ``default`` byte-for-byte and is
        MMLU-neutral for gpt-oss QMoE. Set ``signed_scale=False`` for the positive
        ``max|w| / 2^(bits-1)`` scale (negative extreme exact, positive extreme
        clipped at ``qmax``). Only valid for ``symmetric=True``; the stored
        bytes/zero-point layout is unchanged (the sign rides in the scale).
        """
        qweight, scales, zero_points = CudaQuantizer._matmulnbits_blockwise_quantize_impl(
            weights,
            bits,
            block_size,
            symmetric=symmetric,
            abs_scales=abs_scales,
            unsigned_full_range=unsigned_full_range,
            signed_scale=signed_scale,
        )
        if flatten_qweight:
            qweight = qweight.reshape(qweight.shape[0], -1).contiguous()
        if return_zero_points:
            return qweight, scales, zero_points

        return qweight, scales

    @staticmethod
    def matmulnbits_prepacked_blockwise_quantize(
        weights: torch.Tensor,
        bits: int,
        block_size: int,
        force_arch: int = 80,
        *,
        symmetric: bool = True,
        return_zero_points: bool = False,
        abs_scales: bool = True,
        unsigned_full_range: bool = True,
        signed_scale: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Quantize and CUDA-prepack one MatMulNBits weight initializer.

        ``weights`` has logical shape ``[N, K]``. Returns ``B`` with the standard
        MatMulNBits initializer shape ``[N, K/block_size, block_size*bits/8]``
        and scales with shape ``[N, K/block_size]``.

        The ``force_arch`` value selects the mixed-GEMM weight layout and must match
        the ``weight_prepacked`` attribute set on the MatMulNBits node:

        * ``force_arch=80`` (default): SM80/Ampere layout, consumed by the SM80 kernel
          (also used on newer GPUs via the compatibility path). Use ``weight_prepacked=1``.
        * ``force_arch=90``: SM90/Hopper layout, consumed by the native SM90 TMA/WGMMA
          kernel. Use ``weight_prepacked=2``. Requires ``block_size`` in {64, 128}.
        """
        torch = _get_torch()

        bits = int(bits)
        block_size = int(block_size)
        force_arch = int(force_arch)
        if force_arch not in (80, 90):
            raise ValueError(f"force_arch must be 80 (SM80) or 90 (SM90), but got {force_arch}.")
        # The native SM90 kernel needs group_size to be a multiple of the 64-element Hopper K tile,
        # so block_size=32 is only supported by the SM80/Ampere-class kernel.
        allowed_block_sizes = (32, 64, 128) if force_arch == 80 else (64, 128)
        if block_size not in allowed_block_sizes:
            raise ValueError(
                f"block_size must be one of {allowed_block_sizes} for force_arch={force_arch}, but got {block_size}."
            )
        n, k = weights.shape
        if k % block_size != 0:
            raise ValueError(f"K ({k}) must be divisible by block_size ({block_size}) for CUDA-prepacked weights.")

        qweight, scales, zero_points = CudaQuantizer._matmulnbits_blockwise_quantize_impl(
            weights,
            bits,
            block_size,
            symmetric=symmetric,
            abs_scales=abs_scales,
            unsigned_full_range=unsigned_full_range,
            signed_scale=signed_scale,
        )

        pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm
        packed = pack_weights_for_cuda_mixed_gemm(qweight.reshape(n, -1).numpy(), n, k, bits, force_arch)
        packed = np.asarray(packed).view(np.uint8).reshape(qweight.shape)
        packed = torch.from_numpy(np.ascontiguousarray(packed))
        if return_zero_points:
            return packed, scales, zero_points

        return packed, scales

    @staticmethod
    def prepack_matmulnbits_weight(
        qweight,
        n: int,
        k: int,
        bits: int,
        force_arch: int = 80,
    ) -> np.ndarray:
        """Prepack an already-quantized MatMulNBits weight ``B`` into the fpA_intB layout.

        ``qweight`` holds the standard MatMulNBits ``B`` bytes (uint8, total
        ``N * K * bits / 8``) as produced by any weight-only quantizer, in the initializer
        shape ``[N, K/block_size, block_size*bits/8]``. Returns the CUDA mixed-GEMM
        prepacked bytes with the same shape, matching what
        :meth:`matmulnbits_prepacked_blockwise_quantize` emits for freshly quantized weights.

        Prepacking is a pure layout transform: it only rearranges the quantized bytes and
        does not depend on the quantization method or numeric values. ``force_arch`` selects
        the layout (80 -> SM80/Ampere, 90 -> SM90/Hopper) and must match the node's
        ``weight_prepacked`` attribute (1 or 2 respectively).
        """
        pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm

        qweight = np.ascontiguousarray(np.asarray(qweight)).view(np.uint8)
        original_shape = qweight.shape
        flat = np.ascontiguousarray(qweight.reshape(int(n), -1))
        packed = pack_weights_for_cuda_mixed_gemm(flat, int(n), int(k), int(bits), int(force_arch))
        return np.asarray(packed).view(np.uint8).reshape(original_shape)

    @staticmethod
    def qmoe_prepacked_blockwise_quantize(
        weights: torch.Tensor,
        bits: int,
        block_size: int,
        force_arch: int = 80,
        *,
        symmetric: bool = True,
        return_zero_points: bool = False,
        abs_scales: bool = False,
        unsigned_full_range: bool = True,
        signed_scale: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Quantize one expert and CUTLASS-prepack it for CUDA QMoE fpA_intB GEMM.

        ``weights`` has logical shape ``[N, K]``. Returns ``qweight`` with shape
        ``[K, N/pack]`` and block scales with shape ``[N, K/block_size]``.
        Set ``return_zero_points=True`` to also return packed block zero-points.
        """
        bits = int(bits)
        block_size = int(block_size)
        n, k = weights.shape
        pack = 8 // bits
        if k % block_size != 0:
            raise ValueError(f"K ({k}) must be divisible by block_size ({block_size}) for CUDA-prepacked weights.")
        if n % pack != 0:
            raise ValueError(f"N ({n}) must be divisible by {pack} for QMoE blockwise quantization.")

        qweight, scales, zero_points = CudaQuantizer._matmulnbits_blockwise_quantize_impl(
            weights,
            bits,
            block_size,
            symmetric=symmetric,
            abs_scales=abs_scales,
            unsigned_full_range=unsigned_full_range,
            signed_scale=signed_scale,
        )

        pack_weights_for_cuda_mixed_gemm = _pack_weights_for_cuda_mixed_gemm
        packed = pack_weights_for_cuda_mixed_gemm(qweight.reshape(n, -1).numpy(), n, k, bits, force_arch)
        packed = np.asarray(packed).view(np.uint8).reshape(k, n // pack)
        torch = _get_torch()
        packed = torch.from_numpy(np.ascontiguousarray(packed))
        if return_zero_points:
            return packed, scales, zero_points

        return packed, scales

    @staticmethod
    def symmetric_blockwise_quantize(
        weights: torch.Tensor,
        bits: int,
        block_size: int,
        *,
        unsigned_full_range: bool = True,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Quantize one expert with a pure-PyTorch symmetric blockwise encoding.

        This helper is useful for non-CUDA reference paths. Unlike the pybind-backed
        helpers above, it pads the last dimension when it is not divisible by
        ``block_size`` and returns storage with the same leading shape as ``weights``.
        """
        torch = _get_torch()

        weights = weights.detach().cpu().contiguous()
        original_shape = weights.shape
        bits = int(bits)
        block_size = int(block_size)
        if bits == 4:
            qmin, qmax, scale_divisor = (-8, 7, 8) if unsigned_full_range else (-7, 7, 7)
        elif bits == 8:
            qmin, qmax, scale_divisor = (-128, 127, 128) if unsigned_full_range else (-127, 127, 127)
        else:
            raise ValueError(f"CUDA blockwise quantization only supports 4 or 8 bits, got {bits}.")

        last_dim = original_shape[-1]
        num_blocks = (last_dim + block_size - 1) // block_size
        pad_size = num_blocks * block_size - last_dim
        if pad_size > 0:
            pad_shape = list(original_shape)
            pad_shape[-1] = pad_size
            padding = torch.zeros(pad_shape, dtype=weights.dtype, device=weights.device)
            weights_padded = torch.cat([weights, padding], dim=-1)
        else:
            weights_padded = weights

        reshaped_weights = weights_padded.view(*original_shape[:-1], num_blocks, block_size)
        block_max_abs = torch.max(torch.abs(reshaped_weights), dim=-1)[0]
        scales = torch.clamp(block_max_abs / scale_divisor, min=1e-8)

        quantized = torch.round(reshaped_weights / scales.unsqueeze(-1))
        quantized = torch.clamp(quantized, qmin, qmax)

        if bits == 4:
            quantized_flat = quantized.to(torch.int8).view(*original_shape[:-1], num_blocks * block_size)
            if pad_size > 0:
                quantized_flat = quantized_flat[..., :-pad_size]

            quantized_uint4 = (quantized_flat + 8).to(torch.uint8)
            packed_shape = list(original_shape)
            packed_shape[-1] = (original_shape[-1] + 1) // 2
            qweight = torch.zeros(packed_shape, dtype=torch.uint8, device=weights.device)
            qweight[..., :] = quantized_uint4[..., 0::2] & 0xF
            if quantized_uint4.shape[-1] > 1:
                qweight[..., : quantized_uint4[..., 1::2].shape[-1]] |= (quantized_uint4[..., 1::2] & 0xF) << 4
        else:
            qweight = quantized.to(torch.int8).view(*original_shape[:-1], num_blocks * block_size)
            if pad_size > 0:
                qweight = qweight[..., :-pad_size]
            qweight = qweight.view(original_shape)

        return qweight.cpu(), scales.cpu()
