已合并
[feat] add mxfp4_to_bf16_dequantization kernel #3624
yanghongru创建于 7月6日
[feat] add mxfp4_to_bf16_dequantization kernel #3624
已合并
共 1 个文件变更+127-0
| @@ -0,0 +1,127 @@ | |||
| 1 | +# Copyright (c) 2026, Huawei Technologies Co., Ltd. All rights reserved. | ||
| 2 | + | ||
| 3 | + | ||
| 4 | +import torch | ||
| 5 | +import triton | ||
| 6 | +import triton.language as tl | ||
| 7 | + | ||
| 8 | +# The only block size supported by the MX (microscaling) format. | ||
| 9 | +MX_BLOCK_SIZE: tl.constexpr = 32 | ||
| 10 | + | ||
| 11 | +# Bias of the E8M0 scale exponent. | ||
| 12 | +E8M0_EXPONENT_BIAS: tl.constexpr = 127 | ||
| 13 | + | ||
| 14 | +_MXFP4_BF16_AUTOTUNE_CONFIGS = [ | ||
| 15 | + triton.Config({"BLOCKS_PER_PROGRAM": 8}), | ||
| 16 | + triton.Config({"BLOCKS_PER_PROGRAM": 16}), | ||
| 17 | + triton.Config({"BLOCKS_PER_PROGRAM": 32}), | ||
| 18 | + triton.Config({"BLOCKS_PER_PROGRAM": 64}), | ||
| 19 | + triton.Config({"BLOCKS_PER_PROGRAM": 128}), | ||
| 20 | +] | ||
| 21 | + | ||
| 22 | + | ||
| 23 | + | ||
| 24 | +def _e2m1_to_bf16_value(nibble, scale_exp): | ||
| 25 | + sign = tl.where(((nibble >> 3) & 0x1) == 1, -1.0, 1.0) | ||
| 26 | + exp = ((nibble >> 1) & 0x3).to(tl.float32) | ||
| 27 | + man = (nibble & 0x1).to(tl.float32) | ||
B | |||
| 28 | + mag = tl.where( | ||
| 29 | + exp == 0.0, | ||
B yanghongru 7月7日 评论： 7月7日 评论： | |||
| 30 | + 0.5 * man, | ||
| 31 | + (1.0 + 0.5 * man) * tl.exp2(exp - 1.0), | ||
| 32 | + ) | ||
| 33 | + scale = tl.exp2(scale_exp.to(tl.float32) - E8M0_EXPONENT_BIAS) | ||
| 34 | + return (sign * mag * scale).to(tl.bfloat16) | ||
| 35 | + | ||
| 36 | + | ||
| 37 | + | ||
| 38 | + configs=_MXFP4_BF16_AUTOTUNE_CONFIGS, | ||
| 39 | + key=["K_PACKED", "NUM_BLOCKS"], | ||
| 40 | +) | ||
| 41 | + | ||
| 42 | +def _mxfp4_to_bf16_kernel( | ||
| 43 | + x_ptr, # *uint8, packed E2M1 pairs, shape [M, K_PACKED] | ||
| 44 | + scale_ptr, # *uint8, E8M0 exponents, shape [M, K // MX_BLOCK] | ||
| 45 | + out_i32_ptr, # *int32, shape [M, K_PACKED], each int32 stores two BF16 values | ||
| 46 | + K_PACKED, | ||
| 47 | + NUM_BLOCKS, | ||
| 48 | + BLOCKS_PER_PROGRAM: tl.constexpr, | ||
| 49 | + MX_BLOCK: tl.constexpr, | ||
| 50 | +): | ||
| 51 | + pid_m = tl.program_id(0) | ||
| 52 | + pid_blk = tl.program_id(1) | ||
| 53 | + | ||
| 54 | + mx_blk = pid_blk * BLOCKS_PER_PROGRAM + tl.arange(0, BLOCKS_PER_PROGRAM) | ||
| 55 | + packed_in_mx_blk = tl.arange(0, MX_BLOCK // 2) | ||
| 56 | + packed_offsets = mx_blk[:, None] * (MX_BLOCK // 2) + packed_in_mx_blk[None, :] | ||
| 57 | + mask = packed_offsets < K_PACKED | ||
| 58 | + | ||
| 59 | + packed = tl.load( | ||
| 60 | + x_ptr + pid_m * K_PACKED + packed_offsets, | ||
| 61 | + mask=mask, | ||
| 62 | + other=0, | ||
| 63 | + ).to(tl.uint32) | ||
| 64 | + | ||
| 65 | + # One E8M0 scale is shared by 32 MXFP4 elements, i.e. 16 packed bytes. | ||
| 66 | + scale_exp = tl.load( | ||
| 67 | + scale_ptr + pid_m * NUM_BLOCKS + mx_blk, | ||
| 68 | + mask=mx_blk < NUM_BLOCKS, | ||
| 69 | + other=E8M0_EXPONENT_BIAS, | ||
| 70 | + ) | ||
| 71 | + scale_exp = scale_exp[:, None] | ||
| 72 | + | ||
| 73 | + val_even = _e2m1_to_bf16_value(packed & 0xF, scale_exp) | ||
| 74 | + val_odd = _e2m1_to_bf16_value((packed >> 4) & 0xF, scale_exp) | ||
| 75 | + bits_even = val_even.to(tl.uint16, bitcast=True).to(tl.uint32) | ||
| 76 | + bits_odd = val_odd.to(tl.uint16, bitcast=True).to(tl.uint32) | ||
| 77 | + packed_bf16 = (bits_even & 0xFFFF) | ((bits_odd & 0xFFFF) << 16) | ||
| 78 | + | ||
| 79 | + # Store the two adjacent BF16 outputs as one contiguous 32-bit word. | ||
| 80 | + tl.store( | ||
| 81 | + out_i32_ptr + pid_m * K_PACKED + packed_offsets, | ||
| 82 | + packed_bf16.to(tl.int32, bitcast=True), | ||
| 83 | + mask=mask, | ||
| 84 | + ) | ||
| 85 | + | ||
| 86 | + | ||
| 87 | +def mxfp4_to_bf16_dequant( | ||
| 88 | + x_fp4: torch.Tensor, | ||
| 89 | + scale: torch.Tensor, | ||
| 90 | +) -> torch.Tensor: | ||
| 91 | + """ | ||
| 92 | + Dequantizes MXFP4 into BF16 with the optimized normal-range fast path. | ||
| 93 | + | ||
| 94 | + This benchmark helper assumes every non-zero scaled value stays in the | ||
| 95 | + normal BF16 exponent range. It does not handle BF16 subnormal or overflow | ||
| 96 | + cases and should not replace the general dequantization path. | ||
| 97 | + """ | ||
| 98 | + if scale is None: | ||
| 99 | + raise ValueError("scale must be provided for MXFP4 to BF16 dequantization") | ||
| 100 | + | ||
| 101 | + if x_fp4.dtype != torch.uint8: | ||
| 102 | + x_fp4 = x_fp4.view(torch.uint8) | ||
| 103 | + if scale.dtype != torch.uint8: | ||
| 104 | + scale = scale.view(torch.uint8) | ||
| 105 | + | ||
| 106 | + orig_shape = x_fp4.shape | ||
| 107 | + x2d = x_fp4.reshape(-1, orig_shape[-1]).contiguous() | ||
| 108 | + m, k_packed = x2d.shape | ||
| 109 | + k = 2 * k_packed | ||
| 110 | + | ||
| 111 | + if k % MX_BLOCK_SIZE != 0: | ||
| 112 | + raise ValueError(f"last dim ({k} elements) must be divisible by the MX block size {MX_BLOCK_SIZE}") | ||
| 113 | + num_blocks = k // MX_BLOCK_SIZE | ||
| 114 | + scale2d = scale.reshape(m, num_blocks).contiguous() | ||
| 115 | + | ||
| 116 | + out_i32 = torch.empty((m, k_packed), dtype=torch.int32, device=x_fp4.device) | ||
| 117 | + grid = lambda meta: (m, triton.cdiv(num_blocks, meta["BLOCKS_PER_PROGRAM"])) # pylint: disable=unnecessary-lambda-assignment # noqa | ||
| 118 | + _mxfp4_to_bf16_kernel[grid]( | ||
| 119 | + x_ptr=x2d, | ||
| 120 | + scale_ptr=scale2d, | ||
| 121 | + out_i32_ptr=out_i32, | ||
| 122 | + K_PACKED=k_packed, | ||
| 123 | + NUM_BLOCKS=num_blocks, | ||
| 124 | + MX_BLOCK=MX_BLOCK_SIZE, | ||
| 125 | + ) | ||
| 126 | + | ||
| 127 | + return out_i32.view(torch.bfloat16).reshape(*orig_shape[:-1], k) | ||
一定要转成float吗？ 可以用整型直接 参与后续计算吗？