Dequantizes FP4 packed data using per-block scaling factors. Args: packed_ptr (tl.pointer): Pointer to packed uint8 tensor (M x N//2) scale_ptr (tl.pointer): Pointer to per-block scale tensor (M x N//BLOCK_SIZE) output_ptr (tl.pointer): Pointer to output tensor (M x N)
(
packed_ptr,
scale_ptr,
global_scale_ptr,
output_ptr,
N,
BLOCK_SIZE: tl.constexpr,
TILE_SIZE: tl.constexpr,
)
| 39 | |
| 40 | @triton.jit |
| 41 | def fp4_dequantize_kernel( |
| 42 | packed_ptr, |
| 43 | scale_ptr, |
| 44 | global_scale_ptr, |
| 45 | output_ptr, |
| 46 | N, |
| 47 | BLOCK_SIZE: tl.constexpr, |
| 48 | TILE_SIZE: tl.constexpr, |
| 49 | ): |
| 50 | """Dequantizes FP4 packed data using per-block scaling factors. |
| 51 | |
| 52 | Args: |
| 53 | packed_ptr (tl.pointer): Pointer to packed uint8 tensor (M x N//2) |
| 54 | scale_ptr (tl.pointer): Pointer to per-block scale tensor (M x N//BLOCK_SIZE) |
| 55 | output_ptr (tl.pointer): Pointer to output tensor (M x N) |
| 56 | global_scale_ptr (tl.pointer): Pointer to global scale tensor |
| 57 | N (int): Number of columns in unpacked tensor |
| 58 | BLOCK_SIZE (tl.constexpr): Size of each FP4 quantization block |
| 59 | TILE_SIZE (tl.constexpr): Size of the processing tile (in packed elements) |
| 60 | """ |
| 61 | # Get program ID for processing packed elements |
| 62 | pid = tl.program_id(0) |
| 63 | |
| 64 | # Calculate packed element offsets (each packed element contains 2 FP4 values) |
| 65 | packed_start = pid * TILE_SIZE |
| 66 | packed_offs = packed_start + tl.arange(0, TILE_SIZE) |
| 67 | |
| 68 | # Calculate 2D coordinates for packed data |
| 69 | packed_row_idx = packed_offs // (N // 2) |
| 70 | packed_col_idx = packed_offs % (N // 2) |
| 71 | |
| 72 | # Create mask for packed data bounds checking |
| 73 | packed_mask = packed_col_idx < (N // 2) |
| 74 | |
| 75 | # Load global scale |
| 76 | global_scale = tl.load(global_scale_ptr) |
| 77 | |
| 78 | # Load packed data |
| 79 | packed_data = tl.load(packed_ptr + packed_offs, mask=packed_mask, other=0) |
| 80 | |
| 81 | # Unpack packed FP4 values (uint8) to float16x2 |
| 82 | x_f16x2_packed = tl.inline_asm_elementwise( |
| 83 | asm=""" |
| 84 | { |
| 85 | .reg .b8 byte0, byte1, byte2, byte3; |
| 86 | mov.b32 {byte0, byte1, byte2, byte3}, $4; |
| 87 | cvt.rn.f16x2.e2m1x2 $0, byte0; |
| 88 | cvt.rn.f16x2.e2m1x2 $1, byte1; |
| 89 | cvt.rn.f16x2.e2m1x2 $2, byte2; |
| 90 | cvt.rn.f16x2.e2m1x2 $3, byte3; |
| 91 | } |
| 92 | """, |
| 93 | constraints="=r,=r,=r,=r,r", |
| 94 | args=[packed_data], |
| 95 | dtype=tl.uint32, |
| 96 | is_pure=True, |
| 97 | pack=4, |
| 98 | ) |