UltraGNN

Repository Layout

UltraGNN
|-- source-code/                  # CUDA/C++ extension source code
|   `-- setup.py                  # Unified extension installer
|-- dataset/                      # Put all .npz graph files here
|-- dataset.txt                   # Reference script for converting datasets
`-- evaluation/
    |-- end2end_throughput/       # Throughput evaluation
    |   |-- gat_final/
    |   |-- gatv2_final/
    |   |-- gcn_final/
    |   `-- gsage_final/
    `-- accuracy_validation/      # Accuracy validation
        |-- gat_acc_final/
        `-- gatv2_acc_final/

Requirements

• Ubuntu 16.04+
• cmake >= 3.29
• CUDA >= 11.8
• one NVIDIA RTX4090 GPU and one H100 PCIe GPU.

Install Source Code

Build and install the UltraGNN CUDA extensions from source-code/setup.py:

cd source-code
rm -rf build
python setup.py install

The installer builds the following Python extension modules:

ultragnn_dynamic_graph_block
ultragnn_dynamic_graph_spmm
ultragnn_dynamic_graph_sddmm
ultragnn_static_graph_block
ultragnn_static_graph_spmm

After installation, a quick import check can be run with:

python -c "import ultragnn_dynamic_graph_block; import ultragnn_dynamic_graph_spmm; import ultragnn_dynamic_graph_sddmm; import ultragnn_static_graph_block; import ultragnn_static_graph_spmm; print('UltraGNN import ok')"

Dataset Preparation

All evaluation scripts load graph files from the project-level dataset directory by default:

./dataset

You can also override this location with:

export ULTRAGNN_DATA_DIR=/path/to/your/dataset

The expected file path is:

dataset/<graph_name>.npz

For throughput evaluation, each .npz file should at least contain:

num_nodes
src_li
dst_li

For accuracy validation, each .npz file should contain graph structure and training data:

num_nodes
src_li
dst_li
features
labels
train_mask
val_mask
test_mask
in_size
out_size

dataset.txt contains a reference conversion script for graph datasets stored as compressed raw graph files. When using that script, keep the final save path under dataset/ so the generated .npz files can be loaded directly by the evaluation scripts. For example:

np.savez(
    './dataset/' + folder_name + '.npz',
    num_nodes=max(num_nodes_src_, num_nodes_dst_),
    num_edges=num_edges_,
    src_li=coo_mat.row,
    dst_li=coo_mat.col,
)

Create the directory before conversion:

mkdir -p dataset

The GNN datasets can be downloaded from the official DGL and PyG dataset sources. After download, convert them to the .npz format required by the evaluation scripts.

Throughput Evaluation

Throughput scripts are located under:

evaluation/end2end_throughput

Run a specific model benchmark from its directory:

cd evaluation/end2end_throughput/gat_final
python eva_gat.py

Other entry points:

cd evaluation/end2end_throughput/gatv2_final
python eva_gatv2.py

cd evaluation/end2end_throughput/gcn_final
python eva_gcn.py

cd evaluation/end2end_throughput/gsage_final
python eva_gsage.py

Each script writes results into the local result/ directory under the corresponding model folder. UltraGNN result files use the model-precision naming format:

gat-fp16.csv
gat-tf32.csv
gatv2-fp16.csv
gatv2-tf32.csv
gcn-fp16.csv
gcn-tf32.csv
gsage-fp16.csv
gsage-tf32.csv

DGL baseline results are written to dgl.csv.

Accuracy Validation

Accuracy validation scripts are located under:

evaluation/accuracy_validation

Run GAT accuracy validation:

cd evaluation/accuracy_validation/gat_acc_final
python eva_gat.py

Run GATv2 accuracy validation:

cd evaluation/accuracy_validation/gatv2_acc_final
python eva_gatv2.py

The accuracy scripts write CSV files into each local result/ directory:

dgl-acc.csv
UltraGNN-fp16.csv
UltraGNN-tf32.csv