triton_submodular 0.1.0a1

A Triton implementation of greedy submodular algorithms using fused GPU kernels

Triton Implementation of Greedy Optimization of Submodular Functions

This is a python-only implementation of submodlib and apricot.

This library provides fused kernels for every (function, optimizer) pair available by dynamically generating them.

Installation

pip install triton_submodular

or with uv,

As a package dependency:

uv add triton_submodular

For a venv:

uv pip install triton_submodular

For older GPUs (like pascal), you need to specify the cu126 wheel for pytorch:

uv pip install triton_submodular --torch-backend=cu126

Features

• Functions:
  • Facility Location
  • Graph Cut
  • Feature Based
  • Max Coverage
• Optimizers:
  • Naive Greedy
  • Stochastic Greedy
  • Lazy Greedy

Example Usage

import torch
from triton_submodular.functions import FacilityLocation

X = torch.randn(100, 10, device="cuda", dtype=torch.float32)

fn = FacilityLocation(X) # generates euclidean similarity matrix by default

indices, gains = fn.maximize(10, "LazyGreedyTriton")

X_subset = X[indices]

This will take the data (100 points, 10 dimensional) and greedy select 10 points using the lazy greedy optimizer which tries to maximize the facility location function.

Testing

Clone and run uv run -p 3.9 pytest -n auto. Requires a GPU with CUDA support.

Deterministic optimizers are tested to be 1:1 against both apricot and submodlib's implementations. Random algorithms (currently only stochastic greedy) are only tested against apricot, as submodlib uses C's rand() function, making reproduction difficult.

NOTE: The stochastic greedy optimizer with the graphcut function currently does not match 1:1 to apricot's implementation.

Roadmap

• Benchmarks and optimizations
• Add more submodular functions
• Add more optimizers
• Add more similarity metrics to graph functions that require a similarity matrix
• Add knapsack constraints
• Add sparse matrix support (maybe)
• Add streaming optimization (maybe)