From 3adf93fa5ae019c153c8c1ace3a698189b7db41d Mon Sep 17 00:00:00 2001 From: Lyxot Date: Tue, 27 Oct 2026 16:30:33 +0800 Subject: [PATCH] metal : add NVFP4 mul_mat, mul_mat_id and get_rows kernels --- ggml/src/ggml-metal/ggml-metal-common.cpp | 1 + ggml/src/ggml-metal/ggml-metal-device.cpp | 12 +++ ggml/src/ggml-metal/ggml-metal-device.m | 4 +- ggml/src/ggml-metal/ggml-metal-impl.h | 3 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 1 + ggml/src/ggml-metal/kernels/common.h | 8 ++ ggml/src/ggml-metal/kernels/dequantize.h | 37 +++++++ ggml/src/ggml-metal/kernels/mul_mm.metal | 4 + ggml/src/ggml-metal/kernels/mul_mv.metal | 105 +++++++++++++++++++ ggml/src/ggml-metal/kernels/mul_mv_mma.metal | 1 + ggml/src/ggml-metal/kernels/quantize.metal | 1 + 11 files changed, 174 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp index 70fd551e34b2..1a254ba58481 100644 --- a/ggml/src/ggml-metal/ggml-metal-common.cpp +++ b/ggml/src/ggml-metal/ggml-metal-common.cpp @@ -85,6 +85,7 @@ static bool ggml_metal_mul_mv_mma_type_supported(enum ggml_type type) { case GGML_TYPE_Q1_0: case GGML_TYPE_Q2_0: case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_IQ2_XXS: diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index f3f5d4d77dd7..75b6efb3dd34 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -1064,6 +1064,12 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_meta nr0 = N_R0_MXFP4; smem = 32*sizeof(float); } break; + case GGML_TYPE_NVFP4: + { + nsg = N_SG_NVFP4; + nr0 = N_R0_NVFP4; + smem = 32*sizeof(float); + } break; case GGML_TYPE_Q2_K: { nsg = N_SG_Q2_K; @@ -1360,6 +1366,12 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m nr0 = N_R0_MXFP4; smem = 32*sizeof(float); } break; + case GGML_TYPE_NVFP4: + { + nsg = N_SG_NVFP4; + nr0 = N_R0_NVFP4; + smem = 32*sizeof(float); + } break; case GGML_TYPE_Q2_K: { nsg = N_SG_Q2_K; diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index d1bce8bad85a..02752b0826db 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1516,7 +1516,6 @@ static bool ggml_metal_supports_mul_mat_op( bool src0_f16_has_mv, bool mm_path) { if (!has_simdgroup_reduction || - op->src[0]->type == GGML_TYPE_NVFP4 || op->src[0]->type == GGML_TYPE_TQ1_0) { return false; } @@ -1922,8 +1921,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te }; } case GGML_OP_GET_ROWS: - return op->src[0]->type != GGML_TYPE_NVFP4 && - op->src[0]->type != GGML_TYPE_TQ1_0; + return op->src[0]->type != GGML_TYPE_TQ1_0; case GGML_OP_SET_ROWS: { if (op->src[0]->type == GGML_TYPE_F16) { diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 9db266f44231..0fd4fe518857 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -48,6 +48,9 @@ #define N_R0_MXFP4 2 #define N_SG_MXFP4 2 +#define N_R0_NVFP4 2 +#define N_SG_NVFP4 2 + #define N_R0_Q2_K 4 #define N_SG_Q2_K 2 diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 0cebbe5d6e1f..e98aec01c043 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2639,6 +2639,7 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { op->src[0]->type == GGML_TYPE_Q5_1 || op->src[0]->type == GGML_TYPE_Q8_0 || op->src[0]->type == GGML_TYPE_MXFP4 || + op->src[0]->type == GGML_TYPE_NVFP4 || op->src[0]->type == GGML_TYPE_IQ4_NL || false) && (ne11 >= 2 && ne11 <= 8) ) || diff --git a/ggml/src/ggml-metal/kernels/common.h b/ggml/src/ggml-metal/kernels/common.h index c4d674394480..ebf666e06e53 100644 --- a/ggml/src/ggml-metal/kernels/common.h +++ b/ggml/src/ggml-metal/kernels/common.h @@ -68,6 +68,14 @@ static inline float e8m0_to_fp32(uint8_t x) { return as_type(bits); } +// unsigned E4M3 (bias 7), 0x7F (NaN) decodes to 0 as on the CPU; unlike ggml_ue4m3_to_fp32 there is no 0.5 factor since kvalues_mxfp4_f is not doubled +// the 7 bits placed in an fp16 give the right value scaled by 2^-8, subnormals included +static inline float ue4m3_to_fp32(uint8_t x) { + const float v = (float) as_type((ushort) ((x & 0x7F) << 7)) * 256.0f; + + return x == 0x7F ? 0.0f : v; +} + static inline float dot(float x, float y) { return x*y; } diff --git a/ggml/src/ggml-metal/kernels/dequantize.h b/ggml/src/ggml-metal/kernels/dequantize.h index 54e00faf6f00..273b7f502407 100644 --- a/ggml/src/ggml-metal/kernels/dequantize.h +++ b/ggml/src/ggml-metal/kernels/dequantize.h @@ -378,6 +378,43 @@ void dequantize_mxfp4_t4(device const block_mxfp4 * xb, short il, thread type4 & reg[3] = d * kvalues_mxfp4_f[(q2[4*il4 + 3] >> shr) & 0x0F]; } +// il selects one of the 4 sub-blocks; its 8 bytes hold elements 0..7 in the low nibbles and 8..15 in the high nibbles +template +void dequantize_nvfp4(device const block_nvfp4 * xb, short il, thread type4x4 & reg) { + device const uint8_t * q2 = xb->qs + il*(QK_NVFP4_SUB/2); + + const float d = ue4m3_to_fp32(xb->d[il]); + + float4x4 reg_f; + + for (int i = 0; i < 4; ++i) { + const uint8_t shr = i >= 2 ? 4 : 0; + + reg_f[i][0] = d * kvalues_mxfp4_f[(q2[4*(i%2) + 0] >> shr) & 0x0F]; + reg_f[i][1] = d * kvalues_mxfp4_f[(q2[4*(i%2) + 1] >> shr) & 0x0F]; + reg_f[i][2] = d * kvalues_mxfp4_f[(q2[4*(i%2) + 2] >> shr) & 0x0F]; + reg_f[i][3] = d * kvalues_mxfp4_f[(q2[4*(i%2) + 3] >> shr) & 0x0F]; + } + + reg = (type4x4) reg_f; +} + +template +void dequantize_nvfp4_t4(device const block_nvfp4 * xb, short il, thread type4 & reg) { + const short sub = il/4; + const short i = il%4; + + device const uint8_t * q2 = xb->qs + sub*(QK_NVFP4_SUB/2) + 4*(i%2); + + const float d = ue4m3_to_fp32(xb->d[sub]); + const uint8_t shr = i >= 2 ? 4 : 0; + + reg[0] = d * kvalues_mxfp4_f[(q2[0] >> shr) & 0x0F]; + reg[1] = d * kvalues_mxfp4_f[(q2[1] >> shr) & 0x0F]; + reg[2] = d * kvalues_mxfp4_f[(q2[2] >> shr) & 0x0F]; + reg[3] = d * kvalues_mxfp4_f[(q2[3] >> shr) & 0x0F]; +} + template void dequantize_q2_K(device const block_q2_K *xb, short il, thread type4x4 & reg) { const float d = xb->d; diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal index 2c8492037632..2ebc1767fc35 100644 --- a/ggml/src/ggml-metal/kernels/mul_mm.metal +++ b/ggml/src/ggml-metal/kernels/mul_mm.metal @@ -905,6 +905,7 @@ template [[host_name("kernel_mul_mm_q5_0_f32")]] kernel mul_mm_t kernel_mul_m template [[host_name("kernel_mul_mm_q5_1_f32")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q8_0_f32")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_mxfp4_f32")]] kernel mul_mm_t kernel_mul_mm; +template [[host_name("kernel_mul_mm_nvfp4_f32")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q2_K_f32")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q3_K_f32")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q4_K_f32")]] kernel mul_mm_t kernel_mul_mm; @@ -931,6 +932,7 @@ template [[host_name("kernel_mul_mm_q5_0_f16")]] kernel mul_mm_t kernel_mul_m template [[host_name("kernel_mul_mm_q5_1_f16")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q8_0_f16")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_mxfp4_f16")]] kernel mul_mm_t kernel_mul_mm; +template [[host_name("kernel_mul_mm_nvfp4_f16")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q2_K_f16")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q3_K_f16")]] kernel mul_mm_t kernel_mul_mm; template [[host_name("kernel_mul_mm_q4_K_f16")]] kernel mul_mm_t kernel_mul_mm; @@ -970,6 +972,7 @@ template [[host_name("kernel_mul_mm_id_mxfp4_f32")]] kernel mul_mm_id kernel_m #else template [[host_name("kernel_mul_mm_id_mxfp4_f32")]] kernel mul_mm_id kernel_mul_mm_id; #endif +template [[host_name("kernel_mul_mm_id_nvfp4_f32")]] kernel mul_mm_id kernel_mul_mm_id; template [[host_name("kernel_mul_mm_id_q2_K_f32")]] kernel mul_mm_id kernel_mul_mm_id; template [[host_name("kernel_mul_mm_id_q3_K_f32")]] kernel mul_mm_id kernel_mul_mm_id; template [[host_name("kernel_mul_mm_id_q4_K_f32")]] kernel mul_mm_id kernel_mul_mm_id; @@ -1000,6 +1003,7 @@ template [[host_name("kernel_mul_mm_id_mxfp4_f16")]] kernel mul_mm_id kernel_m #else template [[host_name("kernel_mul_mm_id_mxfp4_f16")]] kernel mul_mm_id kernel_mul_mm_id; #endif +template [[host_name("kernel_mul_mm_id_nvfp4_f16")]] kernel mul_mm_id kernel_mul_mm_id; template [[host_name("kernel_mul_mm_id_q2_K_f16")]] kernel mul_mm_id kernel_mul_mm_id; template [[host_name("kernel_mul_mm_id_q3_K_f16")]] kernel mul_mm_id kernel_mul_mm_id; template [[host_name("kernel_mul_mm_id_q4_K_f16")]] kernel mul_mm_id kernel_mul_mm_id; diff --git a/ggml/src/ggml-metal/kernels/mul_mv.metal b/ggml/src/ggml-metal/kernels/mul_mv.metal index 7ad07a29c727..6c46e01d81db 100644 --- a/ggml/src/ggml-metal/kernels/mul_mv.metal +++ b/ggml/src/ggml-metal/kernels/mul_mv.metal @@ -894,6 +894,11 @@ template [[host_name("kernel_mul_mv_ext_mxfp4_f32_r1_3")]] kernel mul_mv_ext_q4 template [[host_name("kernel_mul_mv_ext_mxfp4_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_mxfp4, 32, dequantize_mxfp4_t4>; template [[host_name("kernel_mul_mv_ext_mxfp4_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_mxfp4, 32, dequantize_mxfp4_t4>; +template [[host_name("kernel_mul_mv_ext_nvfp4_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_nvfp4, 64, dequantize_nvfp4_t4>; +template [[host_name("kernel_mul_mv_ext_nvfp4_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_nvfp4, 64, dequantize_nvfp4_t4>; +template [[host_name("kernel_mul_mv_ext_nvfp4_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_nvfp4, 64, dequantize_nvfp4_t4>; +template [[host_name("kernel_mul_mv_ext_nvfp4_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_nvfp4, 64, dequantize_nvfp4_t4>; + template [[host_name("kernel_mul_mv_ext_iq4_nl_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_iq4_nl, 32, dequantize_iq4_nl_t4>; template [[host_name("kernel_mul_mv_ext_iq4_nl_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_iq4_nl, 32, dequantize_iq4_nl_t4>; template [[host_name("kernel_mul_mv_ext_iq4_nl_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_iq4_nl, 32, dequantize_iq4_nl_t4>; @@ -3124,6 +3129,105 @@ kernel void kernel_mul_mv_mxfp4_f32( kernel_mul_mv_mxfp4_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); } +template +void kernel_mul_mv_nvfp4_f32_impl( + args_t args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + + threadgroup float * shmem_f32 = (threadgroup float *) shmem; + + const int r0 = tgpig.x; + const int r1 = tgpig.y; + const int im = tgpig.z; + + const int first_row = (r0 * NSG + sgitg) * NR0; + + const uint i12 = im%FC_mul_mv_ne12; + const uint i13 = im/FC_mul_mv_ne12; + + const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; + const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; + + device const block_nvfp4 * x = (device const block_nvfp4 *) (src0 + offset0); + device const float * y = (device const float *) (src1 + offset1); + + const int nb = args.ne00/QK_NVFP4; + const int ns01 = args.nb01/args.nb00; // this can be larger than nb for permuted src0 tensors + + const short ix = tiisg/2; // 0...15 + const short it = tiisg%2; // 0 or 1, each thread takes 2 of the 4 sub-blocks + + shmem_f32[tiisg] = kvalues_mxfp4_f[tiisg%16]; + threadgroup_barrier(mem_flags::mem_threadgroup); + + float4 yl[8]; + float sumf[NR0]={0.f}; + + device const float * yb = y + ix*QK_NVFP4 + it*2*QK_NVFP4_SUB; + + for (int ib = ix; ib < nb && ib < ns01; ib += 16) { + device const float4 * y4 = (device const float4 *) yb; + + FOR_UNROLL (short i = 0; i < 8; ++i) { + yl[i] = y4[i]; + } + + FOR_UNROLL (short row = 0; row < NR0; row++) { + device const block_nvfp4 & xb = x[row*ns01 + ib]; + device const uint8_t * q2 = xb.qs + 2*it*(QK_NVFP4_SUB/2); + + float sumq = 0.f; + + FOR_UNROLL (short s = 0; s < 2; ++s) { + device const uint8_t * q = q2 + 8*s; + + float4 acc1 = yl[4*s + 0]*float4(shmem_f32[q[0] & 0x0F], shmem_f32[q[1] & 0x0F], shmem_f32[q[2] & 0x0F], shmem_f32[q[3] & 0x0F]); + float4 acc2 = yl[4*s + 1]*float4(shmem_f32[q[4] & 0x0F], shmem_f32[q[5] & 0x0F], shmem_f32[q[6] & 0x0F], shmem_f32[q[7] & 0x0F]); + float4 acc3 = yl[4*s + 2]*float4(shmem_f32[q[0] >> 4 ], shmem_f32[q[1] >> 4 ], shmem_f32[q[2] >> 4 ], shmem_f32[q[3] >> 4 ]); + float4 acc4 = yl[4*s + 3]*float4(shmem_f32[q[4] >> 4 ], shmem_f32[q[5] >> 4 ], shmem_f32[q[6] >> 4 ], shmem_f32[q[7] >> 4 ]); + + acc1 = (acc1 + acc3) + (acc2 + acc4); + + sumq += ue4m3_to_fp32(xb.d[2*it + s]) * ((acc1[0] + acc1[1]) + (acc1[2] + acc1[3])); + } + + sumf[row] += sumq; + } + + yb += 16 * QK_NVFP4; + } + + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) { + float sum_all = simd_sum(sumf[row]); + if (tiisg == 0) { + dst_f32[first_row + row] = sum_all; + } + } +} + +[[host_name("kernel_mul_mv_nvfp4_f32")]] +kernel void kernel_mul_mv_nvfp4_f32( + constant ggml_metal_kargs_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + + kernel_mul_mv_nvfp4_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +} + template void kernel_mul_mv_tq2_0_f32_impl( args_t args, @@ -3390,6 +3494,7 @@ template [[host_name("kernel_mul_mv_id_q5_0_f32")]] kernel kernel_mul_mv_id_t template [[host_name("kernel_mul_mv_id_q5_1_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_mxfp4_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +template [[host_name("kernel_mul_mv_id_nvfp4_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_q2_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; template [[host_name("kernel_mul_mv_id_q3_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; diff --git a/ggml/src/ggml-metal/kernels/mul_mv_mma.metal b/ggml/src/ggml-metal/kernels/mul_mv_mma.metal index 5cf245d5210d..287834e3e9e3 100644 --- a/ggml/src/ggml-metal/kernels/mul_mv_mma.metal +++ b/ggml/src/ggml-metal/kernels/mul_mv_mma.metal @@ -611,6 +611,7 @@ MUL_MV_MMA_GEN("bf16", bfloat4x4, 1, dequantize_bf16) MUL_MV_MMA_GEN("q1_0", block_q1_0, 8, dequantize_q1_0) MUL_MV_MMA_GEN("q2_0", block_q2_0, 4, dequantize_q2_0) MUL_MV_MMA_GEN("mxfp4", block_mxfp4, 2, dequantize_mxfp4) +MUL_MV_MMA_GEN("nvfp4", block_nvfp4, 4, dequantize_nvfp4) MUL_MV_MMA_GEN("q2_K", block_q2_K, QK_NL, dequantize_q2_K) MUL_MV_MMA_GEN("q3_K", block_q3_K, QK_NL, dequantize_q3_K) MUL_MV_MMA_GEN("iq2_xxs", block_iq2_xxs, QK_NL, dequantize_iq2_xxs) diff --git a/ggml/src/ggml-metal/kernels/quantize.metal b/ggml/src/ggml-metal/kernels/quantize.metal index 73679b9ab76f..102411f418df 100644 --- a/ggml/src/ggml-metal/kernels/quantize.metal +++ b/ggml/src/ggml-metal/kernels/quantize.metal @@ -341,6 +341,7 @@ template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get template [[host_name("kernel_get_rows_q5_1")]] kernel get_rows_q_t kernel_get_rows_q; template [[host_name("kernel_get_rows_q8_0")]] kernel get_rows_q_t kernel_get_rows_q; template [[host_name("kernel_get_rows_mxfp4")]] kernel get_rows_q_t kernel_get_rows_q; +template [[host_name("kernel_get_rows_nvfp4")]] kernel get_rows_q_t kernel_get_rows_q; template [[host_name("kernel_get_rows_q2_K")]] kernel get_rows_q_t kernel_get_rows_q; template [[host_name("kernel_get_rows_q3_K")]] kernel get_rows_q_t kernel_get_rows_q; template [[host_name("kernel_get_rows_q4_K")]] kernel get_rows_q_t kernel_get_rows_q;