From 4a8694c4f17cebb34d5fbdf3bffb12dad570f1fb Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 3 Jul 2026 02:02:53 +0000 Subject: [PATCH 1/7] Add MoeCudaQuantizer --- .../python/tools/quantization/__init__.py | 1 + .../tools/quantization/qmoe_quantizer.py | 154 +++++++++++ .../python/transformers/test_qmoe_cuda.py | 250 +++++++++++++++--- 3 files changed, 374 insertions(+), 31 deletions(-) create mode 100644 onnxruntime/python/tools/quantization/qmoe_quantizer.py diff --git a/onnxruntime/python/tools/quantization/__init__.py b/onnxruntime/python/tools/quantization/__init__.py index 50b0bd08ae360..caaf2c92976a6 100644 --- a/onnxruntime/python/tools/quantization/__init__.py +++ b/onnxruntime/python/tools/quantization/__init__.py @@ -10,6 +10,7 @@ save_tensors_data, ) from .qdq_quantizer import QDQQuantizer # noqa: F401 +from .qmoe_quantizer import MoeCudaQuantizer # noqa: F401 from .quant_utils import QuantFormat, QuantType, write_calibration_table # noqa: F401 from .quantize import ( DynamicQuantConfig, # noqa: F401 diff --git a/onnxruntime/python/tools/quantization/qmoe_quantizer.py b/onnxruntime/python/tools/quantization/qmoe_quantizer.py new file mode 100644 index 0000000000000..d43404de3ba7a --- /dev/null +++ b/onnxruntime/python/tools/quantization/qmoe_quantizer.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np + +if TYPE_CHECKING: + import torch + + +def _get_torch(): + try: + import torch # noqa: PLC0415 + except ImportError as e: + raise ImportError("QMoE CUDA quantization requires torch. Please install torch to use MoeCudaQuantizer.") from e + + return torch + + +def _get_pack_weights_for_cuda_mixed_gemm(): + try: + from onnxruntime.capi import _pybind_state as _pybind # noqa: PLC0415 + except ImportError as e: + raise ImportError( + "CUDA QMoE prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." + ) from e + + try: + return _pybind.pack_weights_for_cuda_mixed_gemm + except AttributeError as e: + raise ImportError( + "CUDA QMoE prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." + ) from e + + +class MoeCudaQuantizer: + """QMoE weight quantizer utilities for ORT CUDA kernels.""" + + @staticmethod + def 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) + 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 + elif bits == 8: + 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 + else: + raise ValueError(f"QMoE per-channel quantization only supports 4 or 8 bits, got {bits}.") + + 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 cuda_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. + + CUDA prepacking requires ``pack_weights_for_cuda_mixed_gemm`` from an + onnxruntime-gpu CUDA build. + """ + torch = _get_torch() + + qweight, scales = MoeCudaQuantizer.symmetric_per_channel_quantize( + weights, + bits, + unsigned_full_range=unsigned_full_range, + ) + if not prepack: + return qweight, scales + + pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() + + 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 + + +def symmetric_per_channel_quantize( + weights: torch.Tensor, + bits: int, + *, + unsigned_full_range: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + return MoeCudaQuantizer.symmetric_per_channel_quantize( + weights, + bits, + unsigned_full_range=unsigned_full_range, + ) + + +def cuda_per_channel_quantize( + weights: torch.Tensor, + bits: int, + prepack: bool, + force_arch: int = 80, + *, + unsigned_full_range: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + return MoeCudaQuantizer.cuda_per_channel_quantize( + weights, + bits, + prepack, + force_arch, + unsigned_full_range=unsigned_full_range, + ) diff --git a/onnxruntime/test/python/transformers/test_qmoe_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_cuda.py index ddca2687fd392..4eb8f196e3094 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_cuda.py @@ -35,6 +35,7 @@ import onnxruntime from onnxruntime.capi import _pybind_state as _pybind +from onnxruntime.quantization import MoeCudaQuantizer try: import nvtx @@ -300,6 +301,20 @@ def quant_dequant_blockwise(weights, block_size, is_4_bit_quantization: bool = T return scale_torch_out, processed_q_weight_torch, result, zero_points_storage +def _dequantize_unsigned_per_channel_storage(qweight, scales, weights, bits: int): + n, k = weights.shape + if bits == 4: + q_low = qweight & 0x0F + q_high = (qweight >> 4) & 0x0F + quantized = torch.stack((q_low, q_high), dim=-1).view(n, -1)[:, :k].to(torch.int16) + quantized = quantized - 8 + else: + quantized = qweight.view(n, k).to(torch.int16) + quantized = quantized - 128 + + return quantized.to(weights.device).to(weights.dtype) * scales.to(weights.device).to(weights.dtype).unsqueeze(-1) + + def quant_dequant(weights, is_4_bit_quantization: bool = True, asymmetric: bool = False): """ Quantize and dequantize weights for testing purposes. @@ -310,37 +325,19 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True, asymmetric: bool """ block_size = weights.shape[1] if not asymmetric and block_size > 256: - n, k = weights.shape - weights_float = weights.detach().float() - if is_4_bit_quantization: - scale = weights_float.abs().amax(dim=1, keepdim=True) / 7.0 - scale = torch.clamp(scale, min=torch.finfo(torch.float32).eps) - q_weight = torch.clamp(torch.round(weights_float / scale), -8, 7).to(torch.int16) + 8 - q_weight = q_weight.to(torch.uint8).contiguous() - q_low = q_weight[:, 0::2] - q_high = q_weight[:, 1::2] - if q_high.shape[1] < q_low.shape[1]: - q_high = F.pad(q_high, (0, 1)) - q_packed = q_low | (q_high << 4) - processed_q_weight = _pybind.pack_weights_for_cuda_mixed_gemm(q_packed.cpu().numpy(), n, k, 4, 80) - processed_q_weight_torch = ( - torch.from_numpy(processed_q_weight).reshape(k, n // 2).to(weights.device).view(torch.uint8) - ) - dequantized = (q_weight.to(weights.dtype) - 8.0) * scale.to(weights.device).to(weights.dtype) - return scale.to(weights.device).to(torch.float16), processed_q_weight_torch, dequantized, None - else: - # 8-bit per-column (per-channel) symmetric quantization. Weights are biased - # uint8 values centered at 128, matching the kernel's (q - 128) * scale dequant. - scale = weights_float.abs().amax(dim=1, keepdim=True) / 127.0 - scale = torch.clamp(scale, min=torch.finfo(torch.float32).eps) - q_weight = torch.clamp(torch.round(weights_float / scale), -127, 127).to(torch.int16) + 128 - q_weight = q_weight.to(torch.uint8).contiguous() - processed_q_weight = _pybind.pack_weights_for_cuda_mixed_gemm(q_weight.cpu().numpy(), n, k, 8, 80) - processed_q_weight_torch = ( - torch.from_numpy(processed_q_weight).reshape(k, n).to(weights.device).view(torch.uint8) - ) - dequantized = (q_weight.to(weights.dtype) - 128.0) * scale.to(weights.device).to(weights.dtype) - return scale.to(weights.device).to(torch.float16), processed_q_weight_torch, dequantized, None + bits = 4 if is_4_bit_quantization else 8 + qweight, scales = MoeCudaQuantizer.symmetric_per_channel_quantize( + weights, + bits, + ) + processed_q_weight, _ = MoeCudaQuantizer.cuda_per_channel_quantize( + weights, + bits, + True, + ) + dequantized = _dequantize_unsigned_per_channel_storage(qweight, scales, weights, bits) + scales = scales.to(weights.device).to(torch.float16).unsqueeze(-1) + return scales, processed_q_weight.to(weights.device), dequantized, None return quant_dequant_blockwise(weights, block_size, is_4_bit_quantization, asymmetric) @@ -2734,6 +2731,29 @@ class TestQMoEIntPrePackSmoke(unittest.TestCase): that harness honors the runtime SM. """ + def test_moe_cuda_quantizer_can_emit_full_range_unsigned_offset_storage(self): + cases = [ + ( + 4, + torch.tensor([[-8.0, -7.0, 0.0, 7.0]], dtype=torch.float32), + torch.tensor([[0x10, 0xF8]], dtype=torch.uint8), + ), + ( + 8, + torch.tensor([[-128.0, -127.0, 0.0, 127.0]], dtype=torch.float32), + torch.tensor([[0x00, 0x01, 0x80, 0xFF]], dtype=torch.uint8), + ), + ] + for bits, weights, expected_qweight in cases: + with self.subTest(bits=bits): + qweight, scales = MoeCudaQuantizer.symmetric_per_channel_quantize( + weights, + bits, + ) + + self.assertTrue(torch.equal(scales, torch.tensor([1.0]))) + self.assertTrue(torch.equal(qweight, expected_qweight)) + def _run_one(self, *, hidden_size, inter_size, num_experts, top_k, swiglu_fusion, batch_size): torch.manual_seed(123) numpy.random.seed(123) @@ -2820,6 +2840,174 @@ def _run_one(self, *, hidden_size, inter_size, num_experts, top_k, swiglu_fusion self.assertGreater(numpy.abs(out).mean(), 1e-4, "Output is suspiciously close to zero") self.assertLess(numpy.abs(out).max(), 10.0, "Output magnitude is implausibly large") + def _run_default_prepacked_model( + self, + *, + hidden_size, + inter_size, + num_experts, + top_k, + batch_size, + fc1_weights, + fc2_weights, + fc1_scales, + fc2_scales, + x, + router, + ): + onnx_dtype = TensorProto.FLOAT16 + fc1_n = 2 * inter_size + qmoe = helper.make_node( + "QMoE", + inputs=["x", "router", "fc1_W", "fc1_S", "", "fc2_W", "fc2_S", ""], + outputs=["y"], + name="qmoe", + domain="com.microsoft", + k=top_k, + normalize_routing_weights=1, + activation_type="swiglu", + swiglu_fusion=1, + expert_weight_bits=4, + quant_type="int", + # weights_prepacked omitted: default -1 means the INT weights are already in + # the CUDA EP's offline-prepacked fpA_intB layout. + ) + graph = helper.make_graph( + nodes=[qmoe], + name="qmoe_default_prepacked", + inputs=[ + helper.make_tensor_value_info("x", onnx_dtype, [batch_size, hidden_size]), + helper.make_tensor_value_info("router", onnx_dtype, [batch_size, num_experts]), + ], + outputs=[helper.make_tensor_value_info("y", onnx_dtype, [batch_size, hidden_size])], + initializer=[ + helper.make_tensor( + "fc1_W", TensorProto.UINT8, list(fc1_weights.shape), fc1_weights.tobytes(), raw=True + ), + helper.make_tensor( + "fc2_W", TensorProto.UINT8, list(fc2_weights.shape), fc2_weights.tobytes(), raw=True + ), + helper.make_tensor("fc1_S", onnx_dtype, [num_experts, fc1_n], fc1_scales.flatten().tolist()), + helper.make_tensor("fc2_S", onnx_dtype, [num_experts, hidden_size], fc2_scales.flatten().tolist()), + ], + ) + model = helper.make_model( + graph, opset_imports=[helper.make_opsetid("", 20), helper.make_opsetid("com.microsoft", 1)] + ) + model.ir_version = 10 + + sess = onnxruntime.InferenceSession(model.SerializeToString(), providers=ort_provider) + return sess.run(None, {"x": x, "router": router})[0] + + def test_int4_default_prepacked_layout_runs_with_moe_cuda_quantizer(self): + torch.manual_seed(123) + numpy.random.seed(123) + + hidden_size = 64 + inter_size = 128 + num_experts = 4 + top_k = 2 + batch_size = 8 + bits = 4 + pack = 8 // bits + + fc1_n = 2 * inter_size + fc1_k = hidden_size + fc2_n = hidden_size + fc2_k = inter_size + + cuda_fc1 = numpy.zeros((num_experts, fc1_k, fc1_n // pack), dtype=numpy.uint8) + cuda_fc2 = numpy.zeros((num_experts, fc2_k, fc2_n // pack), dtype=numpy.uint8) + cuda_fc1_scales = numpy.zeros((num_experts, fc1_n), dtype=numpy.float16) + cuda_fc2_scales = numpy.zeros((num_experts, fc2_n), dtype=numpy.float16) + moe_cuda_quantizer = MoeCudaQuantizer() + + for e in range(num_experts): + w1 = (torch.randn(fc1_n, fc1_k) * 0.05).numpy().astype(numpy.float16) + w2 = (torch.randn(fc2_n, fc2_k) * 0.05).numpy().astype(numpy.float16) + cuda_fc1_t, cuda_fc1_scales_t = moe_cuda_quantizer.cuda_per_channel_quantize( + torch.from_numpy(w1), bits, True + ) + cuda_fc2_t, cuda_fc2_scales_t = moe_cuda_quantizer.cuda_per_channel_quantize( + torch.from_numpy(w2), bits, True + ) + cuda_fc1[e] = cuda_fc1_t.numpy() + cuda_fc2[e] = cuda_fc2_t.numpy() + cuda_fc1_scales[e] = cuda_fc1_scales_t.numpy().astype(numpy.float16) + cuda_fc2_scales[e] = cuda_fc2_scales_t.numpy().astype(numpy.float16) + + x = numpy.random.randn(batch_size, hidden_size).astype(numpy.float16) + router = numpy.random.randn(batch_size, num_experts).astype(numpy.float16) + out = self._run_default_prepacked_model( + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=top_k, + batch_size=batch_size, + fc1_weights=cuda_fc1, + fc2_weights=cuda_fc2, + fc1_scales=cuda_fc1_scales, + fc2_scales=cuda_fc2_scales, + x=x, + router=router, + ) + + self.assertFalse(numpy.isnan(out).any(), "QMoE output has NaN") + self.assertFalse(numpy.isinf(out).any(), "QMoE output has Inf") + + def test_int4_default_prepacked_gpt_oss_shape_smoke(self): + torch.manual_seed(123) + numpy.random.seed(123) + + hidden_size = 2880 + inter_size = 2880 + num_experts = 32 + top_k = 4 + batch_size = 12 + bits = 4 + + fc1_n = 2 * inter_size + fc1_k = hidden_size + fc2_n = hidden_size + fc2_k = inter_size + pack = 8 // bits + moe_cuda_quantizer = MoeCudaQuantizer() + + fc1_weights = numpy.zeros((num_experts, fc1_k, fc1_n // pack), dtype=numpy.uint8) + fc2_weights = numpy.zeros((num_experts, fc2_k, fc2_n // pack), dtype=numpy.uint8) + fc1_scales = numpy.zeros((num_experts, fc1_n), dtype=numpy.float16) + fc2_scales = numpy.zeros((num_experts, fc2_n), dtype=numpy.float16) + + for e in range(num_experts): + w1 = torch.randn(fc1_n, fc1_k, dtype=torch.float16) * 0.01 + w2 = torch.randn(fc2_n, fc2_k, dtype=torch.float16) * 0.01 + q1, s1 = moe_cuda_quantizer.cuda_per_channel_quantize(w1, bits, True) + q2, s2 = moe_cuda_quantizer.cuda_per_channel_quantize(w2, bits, True) + fc1_weights[e] = q1.numpy() + fc2_weights[e] = q2.numpy() + fc1_scales[e] = s1.numpy().astype(numpy.float16) + fc2_scales[e] = s2.numpy().astype(numpy.float16) + + x = (numpy.random.randn(batch_size, hidden_size) * 0.01).astype(numpy.float16) + router = numpy.random.randn(batch_size, num_experts).astype(numpy.float16) + out = self._run_default_prepacked_model( + hidden_size=hidden_size, + inter_size=inter_size, + num_experts=num_experts, + top_k=top_k, + batch_size=batch_size, + fc1_weights=fc1_weights, + fc2_weights=fc2_weights, + fc1_scales=fc1_scales, + fc2_scales=fc2_scales, + x=x, + router=router, + ) + + self.assertEqual(out.shape, (batch_size, hidden_size)) + self.assertFalse(numpy.isnan(out).any(), "QMoE GPT-OSS shape output has NaN") + self.assertFalse(numpy.isinf(out).any(), "QMoE GPT-OSS shape output has Inf") + def test_int4_swiglu_interleaved_small(self): # inter_size must be a multiple of 64 (the interleaved-weight K tile) for the INT path; a # partial final K tile is now rejected up front by QMoE's hardening check. From 2414fa48f3eb40b2b73555c19782aac878cdc154 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 5 Jul 2026 22:56:29 +0000 Subject: [PATCH 2/7] add qmoe block quantization --- .../python/tools/quantization/__init__.py | 2 +- .../tools/quantization/cuda_quantizer.py | 386 ++++++++++++++++++ .../tools/quantization/qmoe_quantizer.py | 154 ------- .../test/python/transformers/test_moe_cuda.py | 361 +--------------- .../python/transformers/test_qmoe_cuda.py | 192 +++------ 5 files changed, 463 insertions(+), 632 deletions(-) create mode 100644 onnxruntime/python/tools/quantization/cuda_quantizer.py delete mode 100644 onnxruntime/python/tools/quantization/qmoe_quantizer.py diff --git a/onnxruntime/python/tools/quantization/__init__.py b/onnxruntime/python/tools/quantization/__init__.py index caaf2c92976a6..e216e7511916b 100644 --- a/onnxruntime/python/tools/quantization/__init__.py +++ b/onnxruntime/python/tools/quantization/__init__.py @@ -9,8 +9,8 @@ load_tensors_data, save_tensors_data, ) +from .cuda_quantizer import CudaQuantizer # noqa: F401 from .qdq_quantizer import QDQQuantizer # noqa: F401 -from .qmoe_quantizer import MoeCudaQuantizer # noqa: F401 from .quant_utils import QuantFormat, QuantType, write_calibration_table # noqa: F401 from .quantize import ( DynamicQuantConfig, # noqa: F401 diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py new file mode 100644 index 0000000000000..bc674a5f502b7 --- /dev/null +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -0,0 +1,386 @@ +# ------------------------------------------------------------------------- +# 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 helpers deliberately wrap the same C++ +pybind entry points used by runtime prepacking so tests and model builders can +generate byte-identical quantized weights. + +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 + +from typing import TYPE_CHECKING + +import numpy as np + +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 CUDA mixed-GEMM weight prepacker from the ORT pybind module.""" + try: + from onnxruntime.capi import _pybind_state as _pybind # noqa: PLC0415 + except ImportError as e: + raise ImportError( + "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." + ) from e + + try: + return _pybind.pack_weights_for_cuda_mixed_gemm + except AttributeError as e: + raise ImportError( + "CUDA weight prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." + ) from e + + +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 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) + 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 + elif bits == 8: + 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 + else: + raise ValueError(f"QMoE per-channel quantization only supports 4 or 8 bits, got {bits}.") + + 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 cuda_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]``. + CUDA prepacking requires ``pack_weights_for_cuda_mixed_gemm`` from an + onnxruntime-gpu CUDA build. + """ + torch = _get_torch() + + qweight, scales = CudaQuantizer.symmetric_per_channel_quantize( + weights, + bits, + unsigned_full_range=unsigned_full_range, + ) + if not prepack: + return qweight, scales + + pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() + + 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, + ) -> 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"CUDA blockwise quantization only supports 4 or 8 bits, got {bits}.") + if k % block_size != 0: + raise ValueError(f"K ({k}) must be divisible by block_size ({block_size}) for CUDA blockwise quantization.") + + num_blocks = k // block_size + pack = 8 // bits + + w_t = np.ascontiguousarray(w.T) + qweight = np.zeros((n, num_blocks, block_size // pack), 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, + ) -> 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 MatMulNBits storage + ``[N, K/pack]`` and block scales ``[N, K/block_size]`` by default. + Set ``return_zero_points=True`` to also return packed block zero-points. + """ + qweight, scales, zero_points = CudaQuantizer._matmulnbits_blockwise_quantize_impl( + weights, + bits, + block_size, + symmetric=symmetric, + abs_scales=abs_scales, + ) + 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, + ) -> 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]``. Use the returned ``B`` with + ``weight_prepacked=1`` on the MatMulNBits node. + + The default ``force_arch=80`` matches the current CUDA MatMulNBits + fpA_intB path: runtime prepacking and offline prepacking both consume the + SM80 mixed-GEMM layout, including on newer GPUs. + """ + torch = _get_torch() + + bits = int(bits) + n, k = weights.shape + qweight, scales, zero_points = CudaQuantizer._matmulnbits_blockwise_quantize_impl( + weights, + bits, + block_size, + symmetric=symmetric, + abs_scales=abs_scales, + ) + + pack_weights_for_cuda_mixed_gemm = _get_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 cutlass_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, + ) -> 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 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, + ) + + pack_weights_for_cuda_mixed_gemm = _get_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, + ) -> 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 = -7, 7 + elif bits == 8: + qmin, qmax = -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 / qmax, 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() diff --git a/onnxruntime/python/tools/quantization/qmoe_quantizer.py b/onnxruntime/python/tools/quantization/qmoe_quantizer.py deleted file mode 100644 index d43404de3ba7a..0000000000000 --- a/onnxruntime/python/tools/quantization/qmoe_quantizer.py +++ /dev/null @@ -1,154 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING - -import numpy as np - -if TYPE_CHECKING: - import torch - - -def _get_torch(): - try: - import torch # noqa: PLC0415 - except ImportError as e: - raise ImportError("QMoE CUDA quantization requires torch. Please install torch to use MoeCudaQuantizer.") from e - - return torch - - -def _get_pack_weights_for_cuda_mixed_gemm(): - try: - from onnxruntime.capi import _pybind_state as _pybind # noqa: PLC0415 - except ImportError as e: - raise ImportError( - "CUDA QMoE prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." - ) from e - - try: - return _pybind.pack_weights_for_cuda_mixed_gemm - except AttributeError as e: - raise ImportError( - "CUDA QMoE prepacking requires pack_weights_for_cuda_mixed_gemm from an onnxruntime-gpu CUDA build." - ) from e - - -class MoeCudaQuantizer: - """QMoE weight quantizer utilities for ORT CUDA kernels.""" - - @staticmethod - def 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) - 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 - elif bits == 8: - 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 - else: - raise ValueError(f"QMoE per-channel quantization only supports 4 or 8 bits, got {bits}.") - - 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 cuda_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. - - CUDA prepacking requires ``pack_weights_for_cuda_mixed_gemm`` from an - onnxruntime-gpu CUDA build. - """ - torch = _get_torch() - - qweight, scales = MoeCudaQuantizer.symmetric_per_channel_quantize( - weights, - bits, - unsigned_full_range=unsigned_full_range, - ) - if not prepack: - return qweight, scales - - pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() - - 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 - - -def symmetric_per_channel_quantize( - weights: torch.Tensor, - bits: int, - *, - unsigned_full_range: bool = True, -) -> tuple[torch.Tensor, torch.Tensor]: - return MoeCudaQuantizer.symmetric_per_channel_quantize( - weights, - bits, - unsigned_full_range=unsigned_full_range, - ) - - -def cuda_per_channel_quantize( - weights: torch.Tensor, - bits: int, - prepack: bool, - force_arch: int = 80, - *, - unsigned_full_range: bool = True, -) -> tuple[torch.Tensor, torch.Tensor]: - return MoeCudaQuantizer.cuda_per_channel_quantize( - weights, - bits, - prepack, - force_arch, - unsigned_full_range=unsigned_full_range, - ) diff --git a/onnxruntime/test/python/transformers/test_moe_cuda.py b/onnxruntime/test/python/transformers/test_moe_cuda.py index f7aa3cf48b6bf..c5c1632ccbab4 100644 --- a/onnxruntime/test/python/transformers/test_moe_cuda.py +++ b/onnxruntime/test/python/transformers/test_moe_cuda.py @@ -26,7 +26,7 @@ import onnxruntime from onnxruntime import InferenceSession, SessionOptions -from onnxruntime.capi import _pybind_state as _quantize +from onnxruntime.quantization import CudaQuantizer # Reduces number of tests to run for faster pipeline checks pipeline_mode = os.getenv("PIPELINE_MODE", "1") == "1" @@ -104,348 +104,31 @@ def print_diff_statistics(diff_tensor: torch.Tensor, prefix: str = ""): def quant_dequant(weights, is_4_bit_quantization: bool = True): - # We use the pybind directly for testing to match what we added in onnxruntime_pybind_quant.cc - if is_4_bit_quantization: - # Quantize on CPU - # quantize_matmul_4bits returns: (q_weight, scale, zero_point) - # weights: [out, in] -> transpose to [in, out] for quantization - weights_t = weights.T.contiguous() - rows, cols = weights_t.shape - block_size = 128 - - # We need to manually call the quantization function exposed in pybind - # because the high-level python API might change. - # But wait, existing helper `quantize_matmul_4bits` in python calls the pybind. - # Let's inspect how to call it. - # Actually, let's use the C++ binding directly as defined in onnxruntime_pybind_quant.cc - # m.def("quantize_matmul_4bits", &QuantizeMatMulNBitsBlockwise); - - # Create output buffers - # shape: [ n, block_per_k, block_blob_size ] - # n = cols, k = rows - k, n = rows, cols - block_per_k = (k + block_size - 1) // block_size - blob_size = block_size // 2 # 4 bits - - q_weight = numpy.zeros((n, block_per_k, blob_size), dtype=numpy.uint8) - scale = numpy.zeros((n, block_per_k), dtype=numpy.float32) # Use float32 for scale - zero_point = numpy.zeros((n, (block_per_k + 1) // 2), dtype=numpy.uint8) - - # weights_t is float32 or float16. The pybind expects float or MLFloat16. - # If weights are float32, use float version. - is_symmetric = True - - if weights.dtype == torch.float32: - _quantize.quantize_matmul_4bits( - q_weight, weights_t.detach().cpu().numpy(), scale, zero_point, block_size, n, k, is_symmetric - ) - elif weights.dtype == torch.float16: - # We might need to handle float16 manually or convert to float32 - _quantize.quantize_matmul_4bits( - q_weight, weights_t.detach().cpu().numpy(), scale, zero_point, block_size, n, k, is_symmetric - ) - - # The output of quantize_matmul_4bits is blockwise. - # We need to reshape it to [n, k // 2]. - # q_weight is [n, k/block_size, block_size/2] - # reshape to [n, k/2] - q_weight_reshaped = q_weight.reshape(n, -1) - - # Pack weights for CUDA mixed-gemm kernel (FpA_IntB format), and qMoE kernel uses the same format. - # INT MoE/QMoE kernels consume the SM80 column-interleaved layout, including on newer GPUs. - processed_q_weight = _quantize.pack_weights_for_cuda_mixed_gemm(q_weight_reshaped, n, k, 4, 80) - - # So we need to DEQUANTIZE back to get `result`. - # scale is [n, block_per_k] - # q_weight is [n, block_per_k, blob_size] - - # Let's do simple dequantization in torch for the reference - scale_torch = torch.from_numpy(scale).to(weights.device) - q_weight_torch = torch.from_numpy(q_weight).to(weights.device) - - # unpack 4 bits - # low 4 bits - q_low = q_weight_torch & 0x0F - # high 4 bits - q_high = (q_weight_torch >> 4) & 0x0F - - # q_weight was [n, blocks, block/2] - # we want [n, blocks, block] - # Interleave low and high? - # MlasQuantizeBlockwise packs 2 elements into uint8. - # e0 is low 4 bits, e1 is high 4 bits. - - q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, block_size) - - # symmetric quantization: value = (q - 8) * scale - # 8 is zero point for 4-bit symmetric? - # MlasQuantizeBlockwise: "is_symmetric ? nullptr" - # If symmetric, zero point is effectively 8 (offset in uint4 range 0-15). - # Wait, Mlas uses offset 8 for symmetric? - # In `MlasQuantizeBlockwise`: - # Value = (Quantized - ZeroPoint) * Scale - # For symmetric 4-bit, the range is [-7, 7]. - # Usually mapped to [1, 15] with zero point 8. - - q_unpacked = q_unpacked.to(weights.dtype) - scale_torch = scale_torch.unsqueeze(-1) # [n, blocks, 1] - - # (q - 8) * scale - dequantized = (q_unpacked - 8.0) * scale_torch - - # reshape to [n, k] to match nn.Linear.weight shape [out_features, in_features] - result = dequantized.view(n, k) - - # pack_weights_for_cuda_mixed_gemm returns flat [k * n // 2]. - # ONNX expects [hidden_size, inter_size // 2] = [k, n // 2]. - # The function transposes, so output is in [k, n // 2] row-major order. - processed_q_weight_torch = torch.from_numpy(processed_q_weight).reshape(k, n // 2).view(torch.uint8) - - # Scale: flatten to [n] for per-channel quantization compatibility. - # The graph expects [inter_size] = [n]. - scale_flat = scale.mean(axis=1) # Average across blocks for per-channel approx - scale_flat_torch = torch.from_numpy(scale_flat).to(weights.device) - - return scale_flat_torch.to(torch.float16), processed_q_weight_torch, result.to(device=weights.device) - - else: - # 8-bit quantization - weights_t = weights.T.contiguous() - rows, cols = weights_t.shape - block_size = 128 - k, n = rows, cols - block_per_k = (k + block_size - 1) // block_size - - # 8-bit: 1 byte per element - q_weight = numpy.zeros((n, block_per_k, block_size), dtype=numpy.uint8) - scale = numpy.zeros((n, block_per_k), dtype=numpy.float32) - zero_point = numpy.zeros((n, block_per_k), dtype=numpy.uint8) # Or dummy? - - is_symmetric = True - - if weights.dtype == torch.float32: - _quantize.quantize_matmul_8bits( - q_weight, weights_t.detach().cpu().numpy(), scale, zero_point, block_size, n, k, is_symmetric - ) - else: - _quantize.quantize_matmul_8bits( - q_weight, weights_t.detach().cpu().numpy(), scale, zero_point, block_size, n, k, is_symmetric - ) - - q_weight_reshaped = q_weight.reshape(n, -1) - # Pack weights for CUDA mixed-gemm kernel (FpA_IntB format) - # INT MoE/QMoE kernels consume the SM80 column-interleaved layout, including on newer GPUs. - processed_q_weight = _quantize.pack_weights_for_cuda_mixed_gemm(q_weight_reshaped, n, k, 8, 80) - - # Dequantize for reference - # (q - 128) * scale if using 128 offset? or (q) * scale if symmetric around 0? - # Mlas symmetric 8-bit usually maps to [-127, 127] or similar? - # Let's assume (q - 128) * scale like standard uint8 quantization if explicit ZP is 128? - # But `is_symmetric=True` passes `nullptr` for ZP. - # Check `MlasQuantizeBlockwise` logic for 8-bit symmetric. - # Usually it produces `int8` directly? - # But `q_weight` is `uint8`. - # If it produces `int8` cast to `uint8` (e.g. 2s complement). - # Then dequantize is `q.view(int8) * scale`. - - scale_torch = torch.from_numpy(scale).to(weights.device) - q_weight_torch = torch.from_numpy(q_weight).to(weights.device) - - # Reinterpret uint8 as int8 - q_signed = q_weight_torch.view(torch.int8) - - scale_torch = scale_torch.unsqueeze(-1) - dequantized = q_signed.to(weights.dtype) * scale_torch - # reshape to [n, k] to match nn.Linear.weight shape [out_features, in_features] - result = dequantized.view(n, k) - - # pack_weights_moe returns flat [k * n]. - # ONNX expects [hidden_size, inter_size] = [k, n]. - processed_q_weight_torch = torch.from_numpy(processed_q_weight).reshape(k, n).view(torch.uint8) - - # Scale: flatten to [n] for per-channel quantization compatibility. - scale_flat = scale.mean(axis=1) # Average across blocks for per-channel approx - scale_flat_torch = torch.from_numpy(scale_flat).to(weights.device) - - return scale_flat_torch.to(torch.float16), processed_q_weight_torch, result.to(device=weights.device) - - # Let's check `test_moe_cuda.py` logic around line 956: - # "Corrected quantization logic for per-output-channel quantization" - # But `MatMulNBits` supports blockwise. - - # If `quant_dequant` returns scales, and those scales are used in `create_phi_moe_onnx_graph`. - # The shape is `[num_experts, inter_size]`. - # If block_size is used, the scale should be larger. - # Unless block_size == K? - - # The current `quant_dequant` implementation in `test_moe_cuda.py` calls: - # torch.ops.trtllm._symmetric_quantize_last_axis_of_batched_matrix - # This function name suggests per-channel (last axis) or blockwise? - - # If `test_moe_cuda.py` assumes per-channel quantization (scale size = inter_size), - # then block_size must be equal to the hidden dimension (row size). - - # HOWEVER, `MatMulNBits` in ORT supports blocking. - # QMoE usually uses blocking (e.g. 128). - - # Let's look at `create_phi_moe_onnx_graph` again. - # fc1_scale_shape = [num_experts, inter_size] - # This assumes one scale per output channel? - # Wait, `inter_size` is the output dimension of fc1 (hidden -> inter). - # So yes, per-channel quantization. - - # BUT, `MatMulNBits` requires `block_size` attribute. - # If we use per-channel, block_size should be K (input dim). - - # Let's check if `test_moe_cuda.py` sets block_size. - # It's not explicitly set in `create_phi_moe_onnx_graph`. - # Wait, `create_phi_moe_onnx_graph` handles the ONNX node creation. - # It assumes `op_name` is "QMoE". - # QMoE kernel in `moe_quantization.cc` reads `block_size` attribute. - # default is -1? - - # In `moe_quantization.cc`: - # this->block_size_ = op_kernel_info.GetAttrOrDefault("block_size", -1); - - # If block_size is -1, what happens? - # In `ComputeInternal`: - # if (block_size_ > 0) { ... GroupWise ... } else { ... Per-column ... } - - # So if we want to match current behavior, we need to see what TRT-LLM `_symmetric_quantize_last_axis_of_batched_matrix` does. - # "last_axis_of_batched_matrix" implies per-channel (per-row of weights if weights are [Out, In]). - # weights passed to `quant_dequant` are `self.experts[i].w1.weight`. - # Linear layer weights are [Out, In]. - # Quantizing last axis means quantizing along `In` dimension, producing one scale per `Out` element. - # This is per-channel quantization. - - # So `block_size` should be -1 (or K). - - # My proposed implementation using `quantize_matmul_4bits` supports `block_size`. - # If I set `block_size = K`, it mimics per-channel. - - # HOWEVER, `pack_weights_moe` implementation I just wrote: - # It calls `preprocess_weights_for_mixed_gemm_cuda`. - # Does that support per-channel? - # `QuantType::W4_A16`. - - # The TRT-LLM function returns `processed_q_weight`. - # This suggests it does the pre-processing (permutation) required by the TRT-LLM/Cutlass kernels. - # The `QMoE` operator in ORT is based on Cutlass/TRT-LLM code. - # So providing the same pre-processed weights is crucial. - - # If `block_size` is not specified in the ONNX node in `test_moe_cuda.py`, it defaults to -1. - # So we should use per-channel quantization. - - # `quantize_matmul_4bits` with `block_size=K`. - # But `pack_weights_moe` logic needs to handle this. - - # Let's proceed with `block_size = cols` (K). - - # IMPORTANT: `create_phi_moe_onnx_graph` hardcodes `fc1_scale_shape = [num_experts, inter_size]`. - # This confirms per-channel. - - # Also need to handle imports carefully inside the function to avoid global dependency errors if something is missing, - # but the test should have onnxruntime installed. + bits = 4 if is_4_bit_quantization else 8 + n, k = weights.shape + block_size = min(128, k) + block_per_k = (k + block_size - 1) // block_size + pack = 8 // bits + + q_weight, scales = CudaQuantizer.matmulnbits_blockwise_quantize(weights, bits, block_size, abs_scales=True) + processed_q_weight, _ = CudaQuantizer.cutlass_prepacked_blockwise_quantize( + weights, bits, block_size, abs_scales=True + ) - # Fix imports: - # `import onnxruntime.quantization._quantize` might not work if it's not exposed that way. - # The pybind module is usually updated into `onnxruntime.quantization`. - # Let's check `onnxruntime/python/Lib/site-packages/onnxruntime/quantization/__init__.py` or similar if we could. - # But generally, `from onnxruntime.quantization import _quantize` won't work directly if it's part of the main extension. - # Usually it's `from onnxruntime.capi import _pybind_state as _quantize` or similar? - # Actually `onnxruntime_pybind_quant.cc` defines a module. - # In `onnxruntime_pybind.cc`, `init_onnxruntime_pybind` calls `CreateQuantPybindModule(m)`. - # So the functions are available under `onnxruntime.capi._pybind_state`. + q_weight = q_weight.view(n, block_per_k, block_size // pack).to(weights.device) + scales = scales.to(weights.device) if is_4_bit_quantization: - weights_t = weights.T.contiguous() - rows, cols = weights_t.shape - k, n = rows, cols - block_size = k # Per-channel - - block_per_k = (k + block_size - 1) // block_size # Should be 1 - blob_size = block_size // 2 - - q_weight = numpy.zeros((n, block_per_k, blob_size), dtype=numpy.uint8) - scale = numpy.zeros((n, block_per_k), dtype=numpy.float32) - zero_point = numpy.zeros((n, (block_per_k + 1) // 2), dtype=numpy.uint8) - - is_symmetric = True - - if weights.dtype == torch.float32: - _quantize.quantize_matmul_4bits( - q_weight, weights_t.cpu().numpy(), scale, zero_point, block_size, n, k, is_symmetric - ) - elif weights.dtype == torch.float16: - _quantize.quantize_matmul_4bits( - q_weight, weights_t.cpu().numpy(), scale, zero_point, block_size, n, k, is_symmetric - ) - - # Reshape for packing - q_weight_reshaped = q_weight.reshape(n, -1) - - # Pack - # We invoke our new function - processed_q_weight = _quantize.pack_weights_moe(q_weight_reshaped, n, k, 4, block_size) - - # Dequantize for reference - # scale: [n, 1] - scale_torch = torch.from_numpy(scale).to(device=weights.device, dtype=weights.dtype) - - # We need raw q_weights for dequantization value recovery - q_weight_torch = torch.from_numpy(q_weight_reshaped).to(device=weights.device) # [n, k/2] - - # unpack 4 bits manually for reference - # Little endian packing in generic logic? - # MlasQuantizeBlockwise logic: - # dst[0] = (uint8_t)(v0 | (v1 << 4)); - # So low 4 bits is first element, high 4 bits is second. - - # unpack - # We need to expand [n, k/2] to [n, k] - # But we need to use the original `q_weight` buffer before packing? - # Yes, `q_weight` from `quantize_matmul_4bits` matches `q` values. - - q_low = q_weight_torch & 0x0F - q_high = (q_weight_torch >> 4) & 0x0F - - # Interleave - # flat view - q_flat = torch.stack((q_low, q_high), dim=-1).view(n, k) - - # symmetric 4-bit range [0, 15], zero point 8. - # value = (q - 8) * scale - - result = (q_flat.to(weights.dtype) - 8.0) * scale_torch - - # Transpose result back to [Out, In] - result = result.T.contiguous() - - # scales are [N, 1] -> flatten to [N] - scale_torch = scale_torch.flatten() - - # processed_q_weight is 1D array of int8 (packed bytes). - # We should return it as is (or as tensor). - # The previous return was: - # return torch_weight_scales.to(torch.float16), processed_q_weight, result.to(device=weights.device) - - return scale_torch.to(torch.float16), torch.from_numpy(processed_q_weight), result - + q_low = q_weight & 0x0F + q_high = (q_weight >> 4) & 0x0F + q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, block_size).to(weights.dtype) + dequantized = (q_unpacked - 8.0) * scales.unsqueeze(-1) else: - # INT8 implementation - # Not fully implemented in this task but required for 8-bit tests? - # The user request mentioned 4-bit mostly, but `test_phi3_qmoe_8bits` exists. - # "If you do not change C++ code... option 1... port implementation". - # I chose option 2 (change C++ code). - # I need to support 8-bit packing too in C++ or handle it. - # My C++ change included a TODO for 8-bit. - # I should probably support it or skip 8-bit tests. - # Let's try to stick to 4-bit for now as the prompt emphasized QMoE 4-bit mainly? - # "We have similar implementation... implement _symmetric_quantize_last_axis_of_batched_matrix" - # That function supports both. - # Let's stick to 4-bit support as per the immediate requirement and see. - # If 8-bit test fails, I'll update. - pass + q_unpacked = q_weight.to(weights.dtype) + dequantized = (q_unpacked - 128.0) * scales.unsqueeze(-1) + + scale_flat = scales.mean(dim=1) + return scale_flat.to(torch.float16), processed_q_weight.to(weights.device), dequantized.view(n, k) def create_moe_onnx_graph( diff --git a/onnxruntime/test/python/transformers/test_qmoe_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_cuda.py index 4eb8f196e3094..5a18f68af4468 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_cuda.py @@ -35,7 +35,7 @@ import onnxruntime from onnxruntime.capi import _pybind_state as _pybind -from onnxruntime.quantization import MoeCudaQuantizer +from onnxruntime.quantization import CudaQuantizer try: import nvtx @@ -161,144 +161,64 @@ def print_diff_statistics(diff_tensor: torch.Tensor, prefix: str = ""): def quant_dequant_blockwise(weights, block_size, is_4_bit_quantization: bool = True, asymmetric: bool = False): - # DEBUG - # print(f"DEBUG: quant_dequant input shape={weights.shape}, 4bit={is_4_bit_quantization}, asym={asymmetric}") - - if is_4_bit_quantization: - weights_t = weights.T.contiguous() - rows, cols = weights_t.shape - k, n = rows, cols - block_per_k = (k + block_size - 1) // block_size - blob_size = block_size // 2 - - q_weight = numpy.zeros((n, block_per_k, blob_size), dtype=numpy.uint8) - scale = numpy.zeros((n, block_per_k), dtype=numpy.float32) - zero_point = numpy.zeros((n, (block_per_k + 1) // 2), dtype=numpy.uint8) - - is_symmetric = not asymmetric - - # Use existing binding which determines implementation based on type - # Assuming weights are float16 or float32. Binding supports both (via overload or check). - # We need to pass numpy array. - # We need to pass numpy array. - if weights_t.dtype == torch.bfloat16: - weights_np = weights_t.detach().to(torch.float32).cpu().numpy() - else: - weights_np = weights_t.detach().cpu().numpy() - - _pybind.quantize_matmul_4bits(q_weight, weights_np, scale, zero_point, block_size, n, k, is_symmetric) - if is_symmetric: - scale = numpy.abs(scale) - - q_weight_reshaped = q_weight.reshape(n, -1) - processed_q_weight = _pybind.pack_weights_for_cuda_mixed_gemm(q_weight_reshaped, n, k, 4, 80) + bits = 4 if is_4_bit_quantization else 8 + n, k = weights.shape + block_per_k = (k + block_size - 1) // block_size + pack = 8 // bits + is_symmetric = not asymmetric + + q_weight, scale, zero_point = CudaQuantizer.matmulnbits_blockwise_quantize( + weights, + bits, + block_size, + symmetric=is_symmetric, + return_zero_points=True, + abs_scales=is_symmetric, + ) + processed_q_weight, _ = CudaQuantizer.cutlass_prepacked_blockwise_quantize( + weights, + bits, + block_size, + symmetric=is_symmetric, + abs_scales=is_symmetric, + ) - # Dequantize for reference - scale_torch = torch.from_numpy(scale).to(weights.device).unsqueeze(-1) - q_weight_torch = torch.from_numpy(q_weight).to(weights.device) + scale_torch = scale.to(weights.device).unsqueeze(-1) + q_weight_torch = q_weight.view(n, block_per_k, block_size // pack).to(weights.device) + if is_4_bit_quantization: + q_low = q_weight_torch & 0x0F + q_high = (q_weight_torch >> 4) & 0x0F + q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, block_size).to(weights.dtype) if is_symmetric: - # Unpack: low, high - q_low = q_weight_torch & 0x0F - q_high = (q_weight_torch >> 4) & 0x0F - q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, block_size) - q_unpacked = q_unpacked.to(weights.dtype) dequantized = (q_unpacked - 8.0) * scale_torch else: - # Asymmetric - # Unpack weights same way - q_low = q_weight_torch & 0x0F - q_high = (q_weight_torch >> 4) & 0x0F - q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, block_size) - q_unpacked = q_unpacked.to(weights.dtype) - - # Unpack ZP - zp_torch = torch.from_numpy(zero_point).to(weights.device) + zp_torch = zero_point.to(weights.device) zp_low = zp_torch & 0x0F zp_high = (zp_torch >> 4) & 0x0F - zp_unpacked = torch.stack((zp_low, zp_high), dim=-1).flatten(1, 2) - zp_unpacked = zp_unpacked[:, :block_per_k].contiguous() - zp_unpacked = zp_unpacked.view(n, block_per_k, 1) - zp_unpacked = zp_unpacked.to(weights.dtype) - + zp_unpacked = torch.stack((zp_low, zp_high), dim=-1).flatten(1, 2)[:, :block_per_k] + zp_unpacked = zp_unpacked.contiguous().view(n, block_per_k, 1).to(weights.dtype) dequantized = (q_unpacked - zp_unpacked) * scale_torch - - scale_torch_out = torch.from_numpy(scale).to(weights.device).to(torch.float16) # N, block_per_K - - # zero_point_storage - zero_points_storage = torch.from_numpy(zero_point).to(weights.device) if asymmetric else None - - processed_q_weight_torch = ( - torch.from_numpy(processed_q_weight).reshape(k, n // 2).to(weights.device).view(torch.uint8) - ) - result = dequantized.view(n, k) - return scale_torch_out, processed_q_weight_torch, result, zero_points_storage - else: - # 8-bit - # C++ binding for 8-bit blockwise quantization (if exists) or use Python implementation - # For now, we use a simple Python implementation that matches the 8nd bits format - # but in practice, we should use the same logic as the kernel. - # Since currently QMoE kernel only supports 4-bit, we don't have a 8-bit PrePack binding yet. - - if _pybind and hasattr(_pybind, "quantize_matmul_8bits"): - # Placeholder for future used when 8-bit is supported - pass - weights_t = weights.T.contiguous() - rows, cols = weights_t.shape - k, n = rows, cols - block_per_k = (k + block_size - 1) // block_size - - q_weight = numpy.zeros((n, block_per_k, block_size), dtype=numpy.uint8) - scale = numpy.zeros((n, block_per_k), dtype=numpy.float32) - zero_point = numpy.zeros((n, block_per_k), dtype=numpy.uint8) - - is_symmetric = not asymmetric - if weights_t.dtype == torch.bfloat16: - weights_np = weights_t.detach().to(torch.float32).cpu().numpy() - else: - weights_np = weights_t.detach().cpu().numpy() - - _pybind.quantize_matmul_8bits(q_weight, weights_np, scale, zero_point, block_size, n, k, is_symmetric) - - q_weight_reshaped = q_weight.reshape(n, -1) - processed_q_weight = _pybind.pack_weights_for_cuda_mixed_gemm(q_weight_reshaped, n, k, 8, 80) - - # Use abs() for reference dequant to match Cutlass kernel's positive scales - scale_torch = torch.from_numpy(scale).to(weights.device).unsqueeze(-1).abs() - q_weight_torch = torch.from_numpy(q_weight).to(weights.device).to(weights.dtype) - + q_unpacked = q_weight_torch.to(weights.dtype) if is_symmetric: - # Kernel does: (biased_uint8 - 128) * scale for symmetric 8-bit - # quantize_matmul_8bits produces biased uint8 values in [0, 255] centered at 128 - dequantized = (q_weight_torch - 128.0) * scale_torch + dequantized = (q_unpacked - 128.0) * scale_torch else: - zp_torch = torch.from_numpy(zero_point).to(weights.device).to(weights.dtype).unsqueeze(-1) - dequantized = (q_weight_torch - zp_torch) * scale_torch + zp_torch = zero_point.to(weights.device).to(weights.dtype).unsqueeze(-1) + dequantized = (q_unpacked - zp_torch) * scale_torch - # Scales must be positive for Cutlass kernel (absolute values) - scale_torch_out = torch.from_numpy(scale).to(weights.device).to(torch.float16).abs() + scale_torch_out = scale.to(weights.device).to(torch.float16) + processed_q_weight_torch = processed_q_weight.to(weights.device).view(torch.uint8) + result = dequantized.view(n, k) - processed_q_weight_torch = ( - torch.from_numpy(processed_q_weight).reshape(k, n).to(weights.device).view(torch.uint8) - ) # 8-bit layout is (K, N) after transpose by pack_weights_for_cuda_mixed_gemm - - result = dequantized.view(n, k) - - if not asymmetric and not is_4_bit_quantization: - # 8-bit Symmetric: weights are uint8, biased by 128. - # Cutlass expects explicit Zero Point = 128 to perform (q - 128) * scale. - # ZP must be FP16 (match Scale type). - zero_point[:] = 128 - zero_points_storage = torch.from_numpy(zero_point).to(weights.device).to(torch.uint8) - else: - zero_points_storage = ( - torch.from_numpy(zero_point).to(weights.device).to(torch.uint8) if asymmetric else None - ) + if asymmetric: + zero_points_storage = zero_point.to(weights.device).to(torch.uint8) + elif not is_4_bit_quantization: + zero_points_storage = torch.full((n, block_per_k), 128, dtype=torch.uint8, device=weights.device) + else: + zero_points_storage = None - # Return scale in [N, block_per_k] layout matching operator spec [E, N, B] after stacking - # Operator will transpose from [E, N, B] to [E, B, N] for kernel - return scale_torch_out, processed_q_weight_torch, result, zero_points_storage + return scale_torch_out, processed_q_weight_torch, result, zero_points_storage def _dequantize_unsigned_per_channel_storage(qweight, scales, weights, bits: int): @@ -326,11 +246,11 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True, asymmetric: bool block_size = weights.shape[1] if not asymmetric and block_size > 256: bits = 4 if is_4_bit_quantization else 8 - qweight, scales = MoeCudaQuantizer.symmetric_per_channel_quantize( + qweight, scales = CudaQuantizer.symmetric_per_channel_quantize( weights, bits, ) - processed_q_weight, _ = MoeCudaQuantizer.cuda_per_channel_quantize( + processed_q_weight, _ = CudaQuantizer.cuda_per_channel_quantize( weights, bits, True, @@ -2746,7 +2666,7 @@ def test_moe_cuda_quantizer_can_emit_full_range_unsigned_offset_storage(self): ] for bits, weights, expected_qweight in cases: with self.subTest(bits=bits): - qweight, scales = MoeCudaQuantizer.symmetric_per_channel_quantize( + qweight, scales = CudaQuantizer.symmetric_per_channel_quantize( weights, bits, ) @@ -2920,17 +2840,13 @@ def test_int4_default_prepacked_layout_runs_with_moe_cuda_quantizer(self): cuda_fc2 = numpy.zeros((num_experts, fc2_k, fc2_n // pack), dtype=numpy.uint8) cuda_fc1_scales = numpy.zeros((num_experts, fc1_n), dtype=numpy.float16) cuda_fc2_scales = numpy.zeros((num_experts, fc2_n), dtype=numpy.float16) - moe_cuda_quantizer = MoeCudaQuantizer() + cuda_quantizer = CudaQuantizer() for e in range(num_experts): w1 = (torch.randn(fc1_n, fc1_k) * 0.05).numpy().astype(numpy.float16) w2 = (torch.randn(fc2_n, fc2_k) * 0.05).numpy().astype(numpy.float16) - cuda_fc1_t, cuda_fc1_scales_t = moe_cuda_quantizer.cuda_per_channel_quantize( - torch.from_numpy(w1), bits, True - ) - cuda_fc2_t, cuda_fc2_scales_t = moe_cuda_quantizer.cuda_per_channel_quantize( - torch.from_numpy(w2), bits, True - ) + cuda_fc1_t, cuda_fc1_scales_t = cuda_quantizer.cuda_per_channel_quantize(torch.from_numpy(w1), bits, True) + cuda_fc2_t, cuda_fc2_scales_t = cuda_quantizer.cuda_per_channel_quantize(torch.from_numpy(w2), bits, True) cuda_fc1[e] = cuda_fc1_t.numpy() cuda_fc2[e] = cuda_fc2_t.numpy() cuda_fc1_scales[e] = cuda_fc1_scales_t.numpy().astype(numpy.float16) @@ -2971,7 +2887,7 @@ def test_int4_default_prepacked_gpt_oss_shape_smoke(self): fc2_n = hidden_size fc2_k = inter_size pack = 8 // bits - moe_cuda_quantizer = MoeCudaQuantizer() + cuda_quantizer = CudaQuantizer() fc1_weights = numpy.zeros((num_experts, fc1_k, fc1_n // pack), dtype=numpy.uint8) fc2_weights = numpy.zeros((num_experts, fc2_k, fc2_n // pack), dtype=numpy.uint8) @@ -2981,8 +2897,8 @@ def test_int4_default_prepacked_gpt_oss_shape_smoke(self): for e in range(num_experts): w1 = torch.randn(fc1_n, fc1_k, dtype=torch.float16) * 0.01 w2 = torch.randn(fc2_n, fc2_k, dtype=torch.float16) * 0.01 - q1, s1 = moe_cuda_quantizer.cuda_per_channel_quantize(w1, bits, True) - q2, s2 = moe_cuda_quantizer.cuda_per_channel_quantize(w2, bits, True) + q1, s1 = cuda_quantizer.cuda_per_channel_quantize(w1, bits, True) + q2, s2 = cuda_quantizer.cuda_per_channel_quantize(w2, bits, True) fc1_weights[e] = q1.numpy() fc2_weights[e] = q2.numpy() fc1_scales[e] = s1.numpy().astype(numpy.float16) From db21ced1b0d9c7e87306a76d7740ccae187394ac Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sun, 5 Jul 2026 23:30:35 +0000 Subject: [PATCH 3/7] Add full_range support for block quantization --- .../tools/quantization/cuda_quantizer.py | 41 +++++++++++++++++-- 1 file changed, 38 insertions(+), 3 deletions(-) diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index bc674a5f502b7..9b848379ecfa0 100644 --- a/onnxruntime/python/tools/quantization/cuda_quantizer.py +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -178,6 +178,7 @@ def _matmulnbits_blockwise_quantize_impl( *, 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() @@ -194,6 +195,29 @@ def _matmulnbits_blockwise_quantize_impl( num_blocks = k // block_size pack = 8 // bits + 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) + ) + + 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 = (quantized[:, :, 0::2] & 0xF) | ((quantized[:, :, 1::2] & 0xF) << 4) + qweight = qweight.astype(np.uint8) + 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, block_size // pack), dtype=np.uint8) scales = np.zeros((n, num_blocks), dtype=np.float32) @@ -217,12 +241,16 @@ def matmulnbits_blockwise_quantize( 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 one expert with ONNX Runtime's MatMulNBits blockwise encoding. ``weights`` has logical shape ``[N, K]``. Returns raw MatMulNBits storage ``[N, K/pack]`` and block scales ``[N, K/block_size]`` by default. 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, @@ -230,6 +258,7 @@ def matmulnbits_blockwise_quantize( block_size, symmetric=symmetric, abs_scales=abs_scales, + unsigned_full_range=unsigned_full_range, ) qweight = qweight.reshape(qweight.shape[0], -1).contiguous() if return_zero_points: @@ -247,6 +276,7 @@ def matmulnbits_prepacked_blockwise_quantize( 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. @@ -269,6 +299,7 @@ def matmulnbits_prepacked_blockwise_quantize( block_size, symmetric=symmetric, abs_scales=abs_scales, + unsigned_full_range=unsigned_full_range, ) pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() @@ -290,6 +321,7 @@ def cutlass_prepacked_blockwise_quantize( 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. @@ -310,6 +342,7 @@ def cutlass_prepacked_blockwise_quantize( block_size, symmetric=symmetric, abs_scales=abs_scales, + unsigned_full_range=unsigned_full_range, ) pack_weights_for_cuda_mixed_gemm = _get_pack_weights_for_cuda_mixed_gemm() @@ -327,6 +360,8 @@ 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. @@ -341,9 +376,9 @@ def symmetric_blockwise_quantize( bits = int(bits) block_size = int(block_size) if bits == 4: - qmin, qmax = -7, 7 + qmin, qmax, scale_divisor = (-8, 7, 8) if unsigned_full_range else (-7, 7, 7) elif bits == 8: - qmin, qmax = -127, 127 + 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}.") @@ -360,7 +395,7 @@ def symmetric_blockwise_quantize( 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 / qmax, min=1e-8) + 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) From d695344185e217ce6034e4dacca779205a4d1198 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 6 Jul 2026 05:54:29 +0000 Subject: [PATCH 4/7] address feedbacks --- .../tools/quantization/cuda_quantizer.py | 44 +++++++++++++------ .../test/python/transformers/test_moe_cuda.py | 10 +++-- .../python/transformers/test_qmoe_cuda.py | 9 ++-- 3 files changed, 41 insertions(+), 22 deletions(-) diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index 9b848379ecfa0..aab0925f458fa 100644 --- a/onnxruntime/python/tools/quantization/cuda_quantizer.py +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -102,6 +102,9 @@ def symmetric_per_channel_quantize( 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: @@ -117,9 +120,6 @@ def symmetric_per_channel_quantize( qmin, qmax, scale_divisor, zero_point = -128, 127, 128, 128 else: qmin, qmax, scale_divisor, zero_point = -127, 127, 127, 128 - else: - raise ValueError(f"QMoE per-channel quantization only supports 4 or 8 bits, got {bits}.") - 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() @@ -188,12 +188,13 @@ def _matmulnbits_blockwise_quantize_impl( w = weights.detach().cpu().to(torch.float32).contiguous().numpy() n, k = w.shape if bits not in (4, 8): - raise ValueError(f"CUDA blockwise quantization only supports 4 or 8 bits, got {bits}.") - if k % block_size != 0: - raise ValueError(f"K ({k}) must be divisible by block_size ({block_size}) for CUDA blockwise quantization.") + 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 + num_blocks = (k + block_size - 1) // block_size pack = 8 // bits + blob_size = (block_size + pack - 1) // pack if symmetric: if bits == 4: @@ -203,15 +204,20 @@ def _matmulnbits_blockwise_quantize_impl( (-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 = np.clip(np.rint(blocked / scales[:, :, np.newaxis]), qmin, qmax).astype(np.int16) quantized = (quantized + zero_point).astype(np.uint8) if bits == 4: - qweight = (quantized[:, :, 0::2] & 0xF) | ((quantized[:, :, 1::2] & 0xF) << 4) - qweight = qweight.astype(np.uint8) + 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 @@ -219,7 +225,7 @@ def _matmulnbits_blockwise_quantize_impl( 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, block_size // pack), dtype=np.uint8) + 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) @@ -241,12 +247,16 @@ def matmulnbits_blockwise_quantize( 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 MatMulNBits storage - ``[N, K/pack]`` and block scales ``[N, K/block_size]`` by default. + ``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 @@ -260,7 +270,8 @@ def matmulnbits_blockwise_quantize( abs_scales=abs_scales, unsigned_full_range=unsigned_full_range, ) - qweight = qweight.reshape(qweight.shape[0], -1).contiguous() + if flatten_qweight: + qweight = qweight.reshape(qweight.shape[0], -1).contiguous() if return_zero_points: return qweight, scales, zero_points @@ -293,6 +304,9 @@ def matmulnbits_prepacked_blockwise_quantize( bits = int(bits) 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, @@ -333,6 +347,8 @@ def cutlass_prepacked_blockwise_quantize( 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.") diff --git a/onnxruntime/test/python/transformers/test_moe_cuda.py b/onnxruntime/test/python/transformers/test_moe_cuda.py index c5c1632ccbab4..0fb7cd1b15664 100644 --- a/onnxruntime/test/python/transformers/test_moe_cuda.py +++ b/onnxruntime/test/python/transformers/test_moe_cuda.py @@ -108,20 +108,22 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True): n, k = weights.shape block_size = min(128, k) block_per_k = (k + block_size - 1) // block_size - pack = 8 // bits - q_weight, scales = CudaQuantizer.matmulnbits_blockwise_quantize(weights, bits, block_size, abs_scales=True) + q_weight, scales = CudaQuantizer.matmulnbits_blockwise_quantize( + weights, bits, block_size, abs_scales=True, flatten_qweight=False + ) processed_q_weight, _ = CudaQuantizer.cutlass_prepacked_blockwise_quantize( weights, bits, block_size, abs_scales=True ) - q_weight = q_weight.view(n, block_per_k, block_size // pack).to(weights.device) + q_weight = q_weight.to(weights.device) scales = scales.to(weights.device) if is_4_bit_quantization: q_low = q_weight & 0x0F q_high = (q_weight >> 4) & 0x0F - q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, block_size).to(weights.dtype) + q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, -1)[:, :, :block_size] + q_unpacked = q_unpacked.to(weights.dtype) dequantized = (q_unpacked - 8.0) * scales.unsqueeze(-1) else: q_unpacked = q_weight.to(weights.dtype) diff --git a/onnxruntime/test/python/transformers/test_qmoe_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_cuda.py index 5a18f68af4468..30f1b046459c8 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_cuda.py @@ -164,7 +164,6 @@ def quant_dequant_blockwise(weights, block_size, is_4_bit_quantization: bool = T bits = 4 if is_4_bit_quantization else 8 n, k = weights.shape block_per_k = (k + block_size - 1) // block_size - pack = 8 // bits is_symmetric = not asymmetric q_weight, scale, zero_point = CudaQuantizer.matmulnbits_blockwise_quantize( @@ -174,6 +173,7 @@ def quant_dequant_blockwise(weights, block_size, is_4_bit_quantization: bool = T symmetric=is_symmetric, return_zero_points=True, abs_scales=is_symmetric, + flatten_qweight=False, ) processed_q_weight, _ = CudaQuantizer.cutlass_prepacked_blockwise_quantize( weights, @@ -184,12 +184,13 @@ def quant_dequant_blockwise(weights, block_size, is_4_bit_quantization: bool = T ) scale_torch = scale.to(weights.device).unsqueeze(-1) - q_weight_torch = q_weight.view(n, block_per_k, block_size // pack).to(weights.device) + q_weight_torch = q_weight.to(weights.device) if is_4_bit_quantization: q_low = q_weight_torch & 0x0F q_high = (q_weight_torch >> 4) & 0x0F - q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, block_size).to(weights.dtype) + q_unpacked = torch.stack((q_low, q_high), dim=-1).view(n, block_per_k, -1)[:, :, :block_size] + q_unpacked = q_unpacked.to(weights.dtype) if is_symmetric: dequantized = (q_unpacked - 8.0) * scale_torch else: @@ -1529,7 +1530,7 @@ def _run_qmoe_cutlass_gemm_second_scale_row_regression(test_case, quant_bits, us os.environ["ORT_DISABLE_MOE_GEMV"] = previous_disable_gemv x = torch.zeros((sequence_length, hidden_size), device=device, dtype=torch_dtype) - x[:, block_size : 2 * block_size] = 1.0 if use_asymmetric_quant else -1.0 + x[:, block_size : 2 * block_size] = 1.0 router = torch.zeros((sequence_length, num_experts), device=device, dtype=torch_dtype) ort_output = sess.run(None, {"input": x.cpu().numpy(), "router_probs": router.cpu().numpy()})[0] From 9540920e697dc57cc659ed934b1e6c4f936e8047 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 6 Jul 2026 18:44:58 +0000 Subject: [PATCH 5/7] fix(quantization): resolve CodeQL uninitialized-variable warnings Restructure symmetric_per_channel_quantize so qmin/qmax/scale_divisor/ zero_point are unconditionally initialized (if bits==4 / else), matching the earlier bits validation. Silences CodeQL 'potentially uninitialized local variable' alerts on cuda_quantizer.py. --- onnxruntime/python/tools/quantization/cuda_quantizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index aab0925f458fa..a7f26e6d0fee8 100644 --- a/onnxruntime/python/tools/quantization/cuda_quantizer.py +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -115,7 +115,7 @@ def symmetric_per_channel_quantize( qmin, qmax, scale_divisor, zero_point = -8, 7, 8, 8 else: qmin, qmax, scale_divisor, zero_point = -7, 7, 7, 8 - elif bits == 8: + else: # bits == 8, already validated above if unsigned_full_range: qmin, qmax, scale_divisor, zero_point = -128, 127, 128, 128 else: From 87d1341fd7c3ae83e9da9526cbb70521365b9b7c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 6 Jul 2026 19:21:47 +0000 Subject: [PATCH 6/7] add qmoe_ prefix to some functions --- .../tools/quantization/cuda_quantizer.py | 32 +++++++++++++------ .../test/python/transformers/test_moe_cuda.py | 2 +- .../python/transformers/test_qmoe_cuda.py | 16 +++++----- 3 files changed, 32 insertions(+), 18 deletions(-) diff --git a/onnxruntime/python/tools/quantization/cuda_quantizer.py b/onnxruntime/python/tools/quantization/cuda_quantizer.py index a7f26e6d0fee8..6b03485280d62 100644 --- a/onnxruntime/python/tools/quantization/cuda_quantizer.py +++ b/onnxruntime/python/tools/quantization/cuda_quantizer.py @@ -82,7 +82,7 @@ class CudaQuantizer: """ @staticmethod - def symmetric_per_channel_quantize( + def qmoe_symmetric_per_channel_quantize( weights: torch.Tensor, bits: int, *, @@ -134,7 +134,7 @@ def symmetric_per_channel_quantize( return qweight.contiguous(), scales.squeeze(-1).contiguous() @staticmethod - def cuda_per_channel_quantize( + def qmoe_per_channel_quantize( weights: torch.Tensor, bits: int, prepack: bool, @@ -151,7 +151,7 @@ def cuda_per_channel_quantize( """ torch = _get_torch() - qweight, scales = CudaQuantizer.symmetric_per_channel_quantize( + qweight, scales = CudaQuantizer.qmoe_symmetric_per_channel_quantize( weights, bits, unsigned_full_range=unsigned_full_range, @@ -293,16 +293,30 @@ def matmulnbits_prepacked_blockwise_quantize( ``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]``. Use the returned ``B`` with - ``weight_prepacked=1`` on the MatMulNBits node. + and scales with shape ``[N, K/block_size]``. - The default ``force_arch=80`` matches the current CUDA MatMulNBits - fpA_intB path: runtime prepacking and offline prepacking both consume the - SM80 mixed-GEMM layout, including on newer GPUs. + 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.") @@ -326,7 +340,7 @@ def matmulnbits_prepacked_blockwise_quantize( return packed, scales @staticmethod - def cutlass_prepacked_blockwise_quantize( + def qmoe_prepacked_blockwise_quantize( weights: torch.Tensor, bits: int, block_size: int, diff --git a/onnxruntime/test/python/transformers/test_moe_cuda.py b/onnxruntime/test/python/transformers/test_moe_cuda.py index 0fb7cd1b15664..9a98e7be3ee16 100644 --- a/onnxruntime/test/python/transformers/test_moe_cuda.py +++ b/onnxruntime/test/python/transformers/test_moe_cuda.py @@ -112,7 +112,7 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True): q_weight, scales = CudaQuantizer.matmulnbits_blockwise_quantize( weights, bits, block_size, abs_scales=True, flatten_qweight=False ) - processed_q_weight, _ = CudaQuantizer.cutlass_prepacked_blockwise_quantize( + processed_q_weight, _ = CudaQuantizer.qmoe_prepacked_blockwise_quantize( weights, bits, block_size, abs_scales=True ) diff --git a/onnxruntime/test/python/transformers/test_qmoe_cuda.py b/onnxruntime/test/python/transformers/test_qmoe_cuda.py index 30f1b046459c8..720260dcb2d61 100644 --- a/onnxruntime/test/python/transformers/test_qmoe_cuda.py +++ b/onnxruntime/test/python/transformers/test_qmoe_cuda.py @@ -175,7 +175,7 @@ def quant_dequant_blockwise(weights, block_size, is_4_bit_quantization: bool = T abs_scales=is_symmetric, flatten_qweight=False, ) - processed_q_weight, _ = CudaQuantizer.cutlass_prepacked_blockwise_quantize( + processed_q_weight, _ = CudaQuantizer.qmoe_prepacked_blockwise_quantize( weights, bits, block_size, @@ -247,11 +247,11 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True, asymmetric: bool block_size = weights.shape[1] if not asymmetric and block_size > 256: bits = 4 if is_4_bit_quantization else 8 - qweight, scales = CudaQuantizer.symmetric_per_channel_quantize( + qweight, scales = CudaQuantizer.qmoe_symmetric_per_channel_quantize( weights, bits, ) - processed_q_weight, _ = CudaQuantizer.cuda_per_channel_quantize( + processed_q_weight, _ = CudaQuantizer.qmoe_per_channel_quantize( weights, bits, True, @@ -2667,7 +2667,7 @@ def test_moe_cuda_quantizer_can_emit_full_range_unsigned_offset_storage(self): ] for bits, weights, expected_qweight in cases: with self.subTest(bits=bits): - qweight, scales = CudaQuantizer.symmetric_per_channel_quantize( + qweight, scales = CudaQuantizer.qmoe_symmetric_per_channel_quantize( weights, bits, ) @@ -2846,8 +2846,8 @@ def test_int4_default_prepacked_layout_runs_with_moe_cuda_quantizer(self): for e in range(num_experts): w1 = (torch.randn(fc1_n, fc1_k) * 0.05).numpy().astype(numpy.float16) w2 = (torch.randn(fc2_n, fc2_k) * 0.05).numpy().astype(numpy.float16) - cuda_fc1_t, cuda_fc1_scales_t = cuda_quantizer.cuda_per_channel_quantize(torch.from_numpy(w1), bits, True) - cuda_fc2_t, cuda_fc2_scales_t = cuda_quantizer.cuda_per_channel_quantize(torch.from_numpy(w2), bits, True) + cuda_fc1_t, cuda_fc1_scales_t = cuda_quantizer.qmoe_per_channel_quantize(torch.from_numpy(w1), bits, True) + cuda_fc2_t, cuda_fc2_scales_t = cuda_quantizer.qmoe_per_channel_quantize(torch.from_numpy(w2), bits, True) cuda_fc1[e] = cuda_fc1_t.numpy() cuda_fc2[e] = cuda_fc2_t.numpy() cuda_fc1_scales[e] = cuda_fc1_scales_t.numpy().astype(numpy.float16) @@ -2898,8 +2898,8 @@ def test_int4_default_prepacked_gpt_oss_shape_smoke(self): for e in range(num_experts): w1 = torch.randn(fc1_n, fc1_k, dtype=torch.float16) * 0.01 w2 = torch.randn(fc2_n, fc2_k, dtype=torch.float16) * 0.01 - q1, s1 = cuda_quantizer.cuda_per_channel_quantize(w1, bits, True) - q2, s2 = cuda_quantizer.cuda_per_channel_quantize(w2, bits, True) + q1, s1 = cuda_quantizer.qmoe_per_channel_quantize(w1, bits, True) + q2, s2 = cuda_quantizer.qmoe_per_channel_quantize(w2, bits, True) fc1_weights[e] = q1.numpy() fc2_weights[e] = q2.numpy() fc1_scales[e] = s1.numpy().astype(numpy.float16) From de6e6ed9705a1ec03647d1b57f299120937fc9a7 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 6 Jul 2026 13:24:38 -0700 Subject: [PATCH 7/7] Update onnxruntime/test/python/transformers/test_moe_cuda.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- onnxruntime/test/python/transformers/test_moe_cuda.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/onnxruntime/test/python/transformers/test_moe_cuda.py b/onnxruntime/test/python/transformers/test_moe_cuda.py index 9a98e7be3ee16..f3d037e65be2b 100644 --- a/onnxruntime/test/python/transformers/test_moe_cuda.py +++ b/onnxruntime/test/python/transformers/test_moe_cuda.py @@ -112,9 +112,7 @@ def quant_dequant(weights, is_4_bit_quantization: bool = True): q_weight, scales = CudaQuantizer.matmulnbits_blockwise_quantize( weights, bits, block_size, abs_scales=True, flatten_qweight=False ) - processed_q_weight, _ = CudaQuantizer.qmoe_prepacked_blockwise_quantize( - weights, bits, block_size, abs_scales=True - ) + processed_q_weight, _ = CudaQuantizer.qmoe_prepacked_blockwise_quantize(weights, bits, block_size, abs_scales=True) q_weight = q_weight.to(weights.device) scales = scales.to(weights.device)