Kumo Engine Overview

NVIDIA Kumo Engine is an open accelerator for context generation over structured data.

It accelerates neighborhood sampling and context generation for GNNs, sequence transformers, and in-context learning models over relational and graph data.

It is built in Rust and designed to move data from NVMe to GPU as quickly as possible for training and inference workloads. As a result, it is about 40× faster than comparable SQL/NoSQL databases for ML workloads (for example, DuckDB and Neo4j), and can scale to datasets 10× larger than solutions that hold the data in RAM (for example, GraphStorm and Memgraph).

Kumo Engine comes with a built-in temporal hetero-GNN architecture which has great performance on classification, regression, and link prediction tasks.

See Release Notes for released versions, changes, and compatibility information.

Getting Started

Run it on a local Linux, macOS, or Windows (WSL2) machine, or in a GPU-enabled cloud notebook such as Snowflake and Databricks. Install Kumo Engine from PyPI:

pip install kumo-engine

The recommended production deployment is on Linux with attached NVMe drives at least 4x larger than your dataset, and at least one attached NVIDIA GPU.

The following example trains a GNN that predicts 30-day purchase probability for each customer:

from kumo.engine import (
    FileConnector, Graph, PredictiveQuery, Table, Trainer, configure,
)
from kumoapi.graph import Edge
# Connect to Parquet tables. Keep working data on local SSD.
configure(scratch_dir="./kumo-scratch")
source = FileConnector(root_dir="./data")
# Define row identifiers and the time of each order.
customers = Table.from_source_table(
    source["customers"], primary_key="id",
)
orders = Table.from_source_table(
    source["orders"], primary_key="id", time_column="ordered_at",
)
# Connect each order to its customer, then build the graph.
graph = Graph(
    tables={"customers": customers, "orders": orders},
    edges=[Edge(
        src_table="orders", fkey="customer_id", dst_table="customers",
    )],
)
snapshot = graph.snapshot()
# Define the prediction: will this customer order in the next 30 days?
query = PredictiveQuery(
    graph,
    "PREDICT COUNT(orders.*, 0, 30, days) > 0 FOR EACH customers.id",
)
# Generate labeled examples and train a model.
training_table = query.generate_training_table(snapshot)
trainer = Trainer(query.suggest_model_plan(snapshot))
result = trainer.fit(training_table, snapshot=snapshot)
print(result.metrics)

See Train a Classification Model to learn more.

How It Works

Universal Dataloader for Structured Data Models

Kumo Engine loads data directly from your cloud warehouse and can generate batches for a wide variety of models.

Compressed Graph Layout

One of the hardest parts of converting a relational dataset to a graph is that all the keys need to be mapped to integer IDs. The key-to-integer mapping alone can be many times larger than RAM.

Kumo Engine uses minimal perfect hashing to compress that mapping to 2.4 bytes per key, enabling it to remain resident in RAM with room to spare.

See benchmark results and hardware recommendations.

Model Training and Inference

Kumo Engine supports the built-in temporal hetero-GNN architecture now. Support for custom PyTorch models and NVIDIA Structured Data Models (SDM) is coming soon.

What You Can Build

See Predictive Query Examples.