Pyre: beating nvforest in forest inference

Hi all,

Last week I open-sourced Balsa: Balsa, mojo-native Treelite reader / writer but of course a Treelite reader/writer is a building block for something bigger.

That bigger project has been published in open source yesterday. It’s called Pyre GitHub - idlab-discover/pyre: Native Mojo CPU and GPU tree-ensemble inference with Balsa models and NuMojo arrays. · GitHub and it’s an inference library for tree ensembles.

The main objective was parity or outperformance against NVForest GitHub - rapidsai/nvforest: Highly optimized and lightweight library for decision tree inference on NVIDIA GPUs and CPUs · GitHub which is a member of the Nvidia Rapids ecosystem. That objective has been reached!

Now that Mojo has entered 1.x.x territory, I feel more comfortable building on other Mojo projects so this project has 2 dependencies, Balsa & NuMojo.

Support has been validated for models exported by scikit-learn, xgboost, catboost and lightgbm. Those libraries have their own inference back-ends and when comparing those against Pyre, it comes out on-top very often.

There is still some work to do towards saturating the CPU and GPU pipelines but that will be for 0.2.0.

1 Like