Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimize FBGEMM Triton MX4 Dequantize #2837

Closed
wants to merge 3 commits into from
Closed

Optimize FBGEMM Triton MX4 Dequantize #2837

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
49d8bc6
Use better exponent rounding in Triton MX4 quantize kernel (#2816)
Jul 12, 2024
d1f21e1
Refactor MX4 Kernel to operate on flat tensors (#2836)
Jul 12, 2024
c680732
Optimize FBGEMM Triton MX4 Dequantize (#2837)
jwfromm Jul 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions fbgemm_gpu/fbgemm_gpu/quantize/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
# pyre-strict

import torch

from fbgemm_gpu.quantize.quantize_ops import dequantize_mx, quantize_mx # noqa F401


Expand Down
13 changes: 3 additions & 10 deletions fbgemm_gpu/fbgemm_gpu/quantize_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
# pyre-strict

import logging
import math

import torch

Expand Down Expand Up @@ -45,15 +44,9 @@ def fp32_to_mx4(
output: MX4 tensor packed into int8 values with total elements (M / 2 + M / groupsize)
"""
# Accelerated MX4 is only available on cuda, if input is on cpu, use python.
# For CPU and triton, set the second dim to 2048 or the nearest power of 2.
dim = (
2048 if tensor.numel() >= 2048 else 2 ** (math.floor(math.log2(tensor.numel())))
)
input = (
tensor.view(-1)
if (tensor.is_cuda and not use_triton) or tensor.numel() % dim != 0
else tensor.view(-1, dim)
)
# Operate on flattened input.
input = tensor.flatten()

if not tensor.is_cuda:
return py_quantize_mx4(input, group_size)

Expand Down
Loading
Loading