opencl: optimize mxfp4 kernels (#16037)

- flatten mxfp4 and packed fp4->fp16 bit-wise convert function (replace lut)
- MoE kernel optimizations

---------

Co-authored-by: Li He <lih@qti.qualcomm.com>
This commit is contained in:
Shawn Gu
2025-09-18 12:03:34 -07:00
committed by GitHub
co-authored by Li He
parent c0b45097c3
commit 3edd87cd05
5 changed files with 701 additions and 2 deletions
+46
View File
@@ -116,3 +116,49 @@ kernel void kernel_convert_block_q4_0_noshuffle(
#endif
}
}
//------------------------------------------------------------------------------
// block_q4_0
//------------------------------------------------------------------------------
#define QK_MXFP4 32
struct block_mxfp4 {
uchar e; // E8M0
uchar qs[QK_MXFP4 / 2];
};
//------------------------------------------------------------------------------
// kernel_convert_block_mxfp4
// Convert the block_mxfp4 format to 2 separate arrays (AOS -> SOA).
// This kernel does not deshuffle the bits.
//------------------------------------------------------------------------------
kernel void kernel_convert_block_mxfp4(
global struct block_mxfp4 * src0,
global uchar * dst_q,
global uchar * dst_e
) {
global struct block_mxfp4 * b = (global struct block_mxfp4 *) src0 + get_global_id(0);
global uchar * q = (global uchar *) dst_q + QK_MXFP4 / 2 * get_global_id(0);
global uchar * e = (global uchar *) dst_e + get_global_id(0);
*e = b->e;
for (int i = 0; i < QK_MXFP4 / 2; ++i) {
q[i] = b->qs[i];
}
}
kernel void kernel_restore_block_mxfp4(
global uchar * src_q,
global half * src_e,
global struct block_mxfp4 * dst
) {
global struct block_mxfp4 * b = (global struct block_mxfp4 *) dst + get_global_id(0);
global uchar * q = (global uchar *) src_q + QK_MXFP4 / 2 * get_global_id(0);
global uchar * e = (global uchar *) src_e + get_global_id(0);
b->e = *e;
for (int i = 0; i < QK_MXFP4 / 2; ++i) {
b->qs[i] = q[i];
}
}