| from enum import Enum |
| import triton |
| import torch |
| from .._triton_kernels.moe.quant_moe import ( |
| _downcast_to_static_fp8, |
| _downcast_to_mxfp, |
| _upcast_from_mxfp, |
| _smoothquant_fuse_quant_kernel, |
| _smoothquant_fuse_quant_kernel_single_pass, |
| ) |
| from ..utils._triton.arch_info import get_arch |
|
|
|
|
| def downcast_to_static_fp8_3d(x: torch.Tensor, scale: torch.Tensor): |
| assert x.ndim == 3 |
| E, M, N = x.shape |
|
|
| x2d = x.reshape(E * M, N).contiguous() |
|
|
| y2d = downcast_to_static_fp8(x2d, scale) |
| y3d = y2d.reshape(E, M, N) |
| return y3d |
|
|
|
|
| def downcast_to_static_fp8(x: torch.Tensor, scale: torch.Tensor): |
| M, N = x.shape |
| if get_arch() != "gfx942": |
| dtype = torch.float8_e4m3fn |
| else: |
| dtype = torch.float8_e4m3fnuz |
| y = torch.empty((M, N), dtype=dtype, device="cuda") |
|
|
| BLOCK_M = min(triton.next_power_of_2(M), 128) |
| if M <= 4096: |
| BLOCK_N = 32 |
| else: |
| BLOCK_N = 64 |
| grid_m = triton.cdiv(x.shape[0], BLOCK_M) |
| grid_n = triton.cdiv(x.shape[1], BLOCK_N) |
|
|
| _downcast_to_static_fp8[(grid_m, grid_n)]( |
| x, |
| x.stride(0), |
| x.stride(1), |
| y, |
| y.stride(0), |
| y.stride(1), |
| scale, |
| M, |
| N, |
| BLOCK_M, |
| BLOCK_N, |
| num_warps=8, |
| ) |
|
|
| return y |
|
|
|
|
| class DequantScaleRoundingMode(Enum): |
| ROUND_UP = 0 |
| ROUND_DOWN = 1 |
|
|
|
|
| def downcast_to_mxfp( |
| src_tensor: torch.Tensor, |
| out_quant_type: torch.dtype, |
| axis: int, |
| DEQUANT_SCALE_ROUNDING_MODE: DequantScaleRoundingMode = DequantScaleRoundingMode.ROUND_UP, |
| ): |
| """ |
| Convert the src weights to mx format. The src weight is quantized along the axis dimension. |
| |
| If weight_quant_type is torch.uint8, we output mxfp4 where two e2m1 values are packed into a single byte. |
| Note that this means the k_dim of the tensor will be half of the logical k_dim. |
| |
| If weight_quant_type is torch.float8_e4m3fn or torch.float8_e5m2, we output mxfp8 with the float8s are stored |
| in their respective formats. |
| """ |
| ndim = src_tensor.ndim |
| assert -ndim <= axis < ndim, f"Invalid axis {axis=}" |
| axis = axis if axis >= 0 else axis + ndim |
| |
| src_tensor = src_tensor.transpose(axis, src_tensor.ndim - 1) |
| is_fp4 = out_quant_type == torch.uint8 |
| is_fp8 = out_quant_type in ( |
| torch.float8_e4m3fn, |
| torch.float8_e4m3fnuz, |
| torch.float8_e5m2, |
| ) |
| assert is_fp4 or is_fp8 |
| divisor = 2 if is_fp4 else 1 |
| L = src_tensor.shape[-1] |
| if is_fp4: |
| assert L % 2 == 0, f"axis dim must be divisible by 2 for e2m1. Got {L}" |
| out_shape = src_tensor.shape[:-1] + (L // divisor,) |
| out_scale_shape = src_tensor.shape[:-1] + (triton.cdiv(L, 32),) |
|
|
| out_quant_tensor = src_tensor.new_empty(out_shape, dtype=out_quant_type) |
| out_scale = src_tensor.new_empty(out_scale_shape, dtype=torch.uint8) |
|
|
| kernel_src_tensor = src_tensor.reshape(-1, src_tensor.shape[-1]) |
| kernel_quant_tensor = out_quant_tensor.view(-1, out_quant_tensor.shape[-1]) |
| kernel_scale = out_scale.view(-1, out_scale.shape[-1]) |
|
|
| BLOCK_OUT_DIM = 128 |
| BLOCK_QUANT_DIM = 32 |
| grid_out = triton.cdiv(kernel_src_tensor.shape[0], BLOCK_OUT_DIM) |
| grid_quant = triton.cdiv(kernel_src_tensor.shape[1], BLOCK_QUANT_DIM) |
|
|
| _downcast_to_mxfp[(grid_out, grid_quant)]( |
| kernel_quant_tensor, |
| *kernel_quant_tensor.stride(), |
| kernel_scale, |
| *kernel_scale.stride(), |
| kernel_src_tensor, |
| *kernel_src_tensor.stride(), |
| *kernel_src_tensor.shape, |
| BLOCK_OUT_DIM, |
| BLOCK_QUANT_DIM, |
| DEQUANT_SCALE_ROUNDING_MODE.value, |
| num_warps=8, |
| ) |
|
|
| out_quant_tensor = out_quant_tensor.transpose(axis, src_tensor.ndim - 1) |
| out_scale = out_scale.transpose(axis, src_tensor.ndim - 1) |
| return out_quant_tensor, out_scale |
|
|
|
|
| def upcast_from_mxfp( |
| tensor: torch.Tensor, scale: torch.Tensor, dtype: torch.dtype, axis: int |
| ): |
| """ |
| Upcasts an mxfp (packed) weight tensor back to float16 or bfloat16. |
| |
| The function assumes that the tensors were quantized along the given axis. |
| It permutes the tensor so that the quantized axis is last, reshapes to 2D, |
| launches the Triton upcast kernel, and then unpermutes back to the original order. |
| """ |
| ndim = tensor.ndim |
| assert -ndim <= axis < ndim, f"Invalid axis {axis=}" |
| axis = axis if axis >= 0 else axis + ndim |
| assert tensor.ndim == scale.ndim, ( |
| f"Weight and scale must have the same number of dimensions. " |
| f"Got {tensor.ndim=} and {scale.ndim=}" |
| ) |
| |
| assert tensor.dtype in { |
| torch.uint8, |
| torch.float8_e5m2, |
| torch.float8_e4m3fn, |
| torch.float8_e4m3fnuz, |
| }, f"Invalid tensor dtype {tensor.dtype=}" |
| assert scale.dtype == torch.uint8, f"Invalid scale dtype {scale.dtype=}" |
| assert dtype in (torch.float16, torch.bfloat16), f"Invalid output dtype {dtype=}" |
| |
| logical_quant_dim = tensor.shape[axis] * (2 if tensor.dtype == torch.uint8 else 1) |
| tensor = tensor.transpose(axis, tensor.ndim - 1).contiguous() |
| scale = scale.transpose(axis, scale.ndim - 1).contiguous() |
| out = torch.empty( |
| (*tensor.shape[:-1], logical_quant_dim), dtype=dtype, device=tensor.device |
| ) |
| reshaped_out = out.view(-1, out.shape[-1]) |
| reshaped_tensor = tensor.view(-1, tensor.shape[-1]) |
| reshaped_scale = scale.view(-1, scale.shape[-1]) |
| BLOCK_OUT_DIM = 128 |
| BLOCK_QUANT_DIM = 32 |
| blocks_out_dim = triton.cdiv(reshaped_out.shape[0], BLOCK_OUT_DIM) |
| blocks_quant_dim = triton.cdiv(reshaped_out.shape[1], BLOCK_QUANT_DIM) |
| _upcast_from_mxfp[(blocks_out_dim, blocks_quant_dim)]( |
| reshaped_out, |
| *reshaped_out.stride(), |
| reshaped_scale, |
| *reshaped_scale.stride(), |
| reshaped_tensor, |
| *reshaped_tensor.stride(), |
| *reshaped_out.shape, |
| BLOCK_OUT_DIM, |
| BLOCK_QUANT_DIM, |
| num_warps=8, |
| ) |
| out = out.transpose(axis, scale.ndim - 1).contiguous() |
| return out |
|
|
|
|
| def dequant_x_blockscale(x, x_scales, per_row_x_scale, group_shape): |
| assert x_scales is not None |
| group_shape_m, _, group_shape_k = group_shape |
| M, K = x.shape |
|
|
| K_blocks = (K + group_shape_k - 1) // group_shape_k |
| if per_row_x_scale: |
| assert x_scales.shape == (M, K_blocks) |
| K_pad = K_blocks * group_shape_k |
| if K_pad != K: |
| x_pad = x.new_zeros((M, K_pad)) |
| x_pad[:, :K] = x |
| x = x_pad |
|
|
| x = x.to(torch.float32).view(M, K_blocks, group_shape_k) * x_scales.to( |
| torch.float32 |
| ).view(M, K_blocks, 1) |
| x = x.view(M, K_pad)[:, :K] |
| else: |
| M_blocks = (M + group_shape_m - 1) // group_shape_m |
| assert x_scales.shape == (M_blocks, K_blocks) |
| M_pad = M_blocks * group_shape_m |
| K_pad = K_blocks * group_shape_k |
| if M_pad != M or K_pad != K: |
| x_pad = x.new_zeros((M_pad, K_pad)) |
| x_pad[:M, :K] = x |
| x = x_pad |
|
|
| x = x.to(torch.float32).view(M_blocks, group_shape_m, K_blocks, group_shape_k) |
| scales = x_scales.to(torch.float32).view(M_blocks, 1, K_blocks, 1) |
| x = x * scales |
| x = x.view(M_pad, K_pad)[:M, :K] |
| return x |
|
|
|
|
| def dequant_w_blockscale(w, w_scales, group_shape): |
| assert w_scales is not None |
| _, group_shape_n, group_shape_k = group_shape |
| E, K, N = w.shape |
|
|
| K_blocks = (K + group_shape_k - 1) // group_shape_k |
| N_blocks = (N + group_shape_n - 1) // group_shape_n |
|
|
| assert w_scales.shape == (E, K_blocks, N_blocks) |
|
|
| K_pad = K_blocks * group_shape_k |
| N_pad = N_blocks * group_shape_n |
| if K_pad != K or N_pad != N: |
| w_pad = w.new_zeros((E, K_pad, N_pad)) |
| w_pad[:, :K, :N] = w |
| w = w_pad |
| w = w.to(torch.float32).view(E, K_blocks, group_shape_k, N_blocks, group_shape_n) |
| scales = w_scales.to(torch.float32).view(E, K_blocks, 1, N_blocks, 1) |
| w = w * scales |
| w = w.view(E, K_pad, N_pad)[:, :K, :N] |
| return w |
|
|
|
|
| def smoothquant_quantize( |
| x: torch.Tensor, |
| smooth_scale: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """ |
| Apply smoothquant quantization to convert bf16/fp16 tensor to int8. |
| |
| Args: |
| x: Input tensor in bf16/fp16 [M, K] |
| smooth_scale: Per-column smooth scale in fp32 [K] |
| |
| Returns: |
| x_int8: Quantized int8 tensor [M, K] |
| x_scale: Per-row quantization scale in fp32 [M] |
| |
| The operation performs: |
| 1. x_smooth = x * smooth_scale (per column) |
| 2. row_scale = max(abs(x_smooth), dim=1) / 127 |
| 3. x_int8 = round(x_smooth / row_scale) |
| """ |
| assert x.ndim == 2, f"Expected 2D tensor, got {x.ndim}D" |
| assert smooth_scale.ndim == 1, f"Expected 1D smooth_scale, got {smooth_scale.ndim}D" |
| assert ( |
| x.shape[1] == smooth_scale.shape[0] |
| ), f"Dimension mismatch: x.shape[1]={x.shape[1]}, smooth_scale.shape[0]={smooth_scale.shape[0]}" |
|
|
| M, K = x.shape |
| device = x.device |
|
|
| x_int8 = torch.empty((M, K), dtype=torch.int8, device=device) |
| x_scale = torch.empty((M,), dtype=torch.float32, device=device) |
|
|
| smooth_scale = smooth_scale.to(torch.float32).contiguous() |
|
|
| MAX_SINGLE_PASS_K = 1024 |
| BLOCK_M = min(triton.next_power_of_2(M), 32) |
|
|
| if K <= MAX_SINGLE_PASS_K: |
| |
| BLOCK_K = triton.next_power_of_2(K) |
| grid = (triton.cdiv(M, BLOCK_M),) |
|
|
| _smoothquant_fuse_quant_kernel_single_pass[grid]( |
| x, |
| x.stride(0), |
| x.stride(1), |
| smooth_scale, |
| x_int8, |
| x_int8.stride(0), |
| x_int8.stride(1), |
| x_scale, |
| 1, |
| M, |
| K, |
| BLOCK_M, |
| BLOCK_K, |
| num_warps=4, |
| ) |
| else: |
| BLOCK_K = 256 |
| grid = (triton.cdiv(M, BLOCK_M),) |
| _smoothquant_fuse_quant_kernel[grid]( |
| x, |
| x.stride(0), |
| x.stride(1), |
| smooth_scale, |
| x_int8, |
| x_int8.stride(0), |
| x_int8.stride(1), |
| x_scale, |
| 1, |
| M, |
| K, |
| BLOCK_M, |
| BLOCK_K, |
| num_warps=4, |
| ) |
|
|
| return x_int8, x_scale |
|
|
|
|
| def quantize_weights_int8( |
| w: torch.Tensor, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """ |
| Quantize weights to int8 with per-output-channel scaling. |
| |
| Args: |
| w: Weight tensor in bf16/fp16/fp32 [E, K, N] or [K, N] |
| |
| Returns: |
| w_int8: Quantized int8 weights (contiguous) |
| w_scale: Per-output-channel scale [E, N] or [N] (contiguous) |
| """ |
| if w.ndim == 2: |
| |
| w = w.unsqueeze(0) |
| squeeze_output = True |
| else: |
| squeeze_output = False |
|
|
| w_fp32 = w.to(torch.float32) |
| w_abs_max = w_fp32.abs().max(dim=1).values |
| INT8_MAX = 127.0 |
| w_scale = w_abs_max / INT8_MAX + 1e-12 |
| w_scaled = w_fp32 / w_scale[:, None, :] |
| w_int8 = w_scaled.round().clamp(-127, 127).to(torch.int8) |
|
|
| |
| w_int8 = w_int8.contiguous() |
| w_scale = w_scale.contiguous() |
|
|
| if squeeze_output: |
| w_int8 = w_int8.squeeze(0) |
| w_scale = w_scale.squeeze(0) |
|
|
| return w_int8, w_scale |
|
|
