# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------

"""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 onnxruntime.quantization 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 _get_pack_weights_for_cuda_mixed_gemm():
    """Return the standalone CUDA mixed-GEMM weight packer (parity oracle).

    Production packing uses the PyTorch implementation (``_pack_weights_for_cuda_mixed_gemm``).
    This standalone packer lives in ``onnxruntime.capi.onnxruntime_cuda_quant_preprocess``, a
    separate extension module that links the CUDA runtime (built only on non-Windows CUDA
    builds). It is imported lazily here (never at ``import onnxruntime`` time) and is used by
    the parity test to validate the PyTorch packer byte-for-byte.
    """
    try:
        from onnxruntime.capi import onnxruntime_cuda_quant_preprocess as _cuda_quant  # noqa: PLC0415
    except ImportError as e:
        raise ImportError(
            "The standalone CUDA weight packer (onnxruntime_cuda_quant_preprocess) is unavailable; "
            "it is built only on non-Windows onnxruntime-gpu CUDA builds."
        ) from e

    try:
        return _cuda_quant.pack_weights_for_cuda_mixed_gemm
    except AttributeError as e:
        raise ImportError("onnxruntime_cuda_quant_preprocess is missing pack_weights_for_cuda_mixed_gemm.") from e


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,
        ],
    }
    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,
    ) -> 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.
        """
        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
        scales = weights.abs().amax(dim=1, keepdim=True) / float(scale_divisor)
        scales = torch.clamp(scales, min=torch.finfo(torch.float32).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,
    ) -> 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).
        """
        torch = _get_torch()

        qweight, scales = CudaQuantizer.qmoe_symmetric_per_channel_quantize(
            weights,
            bits,
            unsigned_full_range=unsigned_full_range,
        )
        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,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Quantize ``weights`` with MatMulNBits pybinds and return unflattened storage."""
        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}.")

        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)
            scales = np.max(np.abs(blocked), axis=2).astype(np.float32) / np.float32(scale_divisor)
            scales = np.maximum(scales, np.finfo(np.float32).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,
    ) -> 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.
        """
        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,
        )
        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,
    ) -> 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,
        )

        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 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,
    ) -> 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,
        )

        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()