Speaker
Description
Jet clustering remains one of the few steps in modern analysis and reconstruction chains that is still bound to the CPU. Machine learning workflows that wish to recluster jets inside the training loop, scan the jet radius, or access substructure dynamically must pay for repeated transfers between host and device, and reconstruction itself is steadily moving toward heterogeneous environments in which as much of the event processing as possible is expected to run on accelerators. We present flashjet, an open source library that brings the standard sequential recombination algorithms (kt, Cambridge Aachen and anti kt) natively to the GPU through custom Triton kernels. Writing the clustering in Triton, the GPU programming language of the PyTorch ecosystem, means the algorithm compiles and autotunes for the actual hardware at hand and integrates seamlessly with tensor based ML pipelines, with no separate CUDA build or hand written device code to maintain. Each kernel program clusters one event entirely on chip, the whole batch is processed in a single fused launch, and everything stays on the device. Jets, their full clustering histories, and derived substructure such as grooming, subjets and Lund plane observables all come out as tensors at negligible extra cost.
The physics output is identical to FastJet. The implementation is validated against the FastJet reference at full numerical precision, reproducing jet momenta, groomed masses and substructure observables jet by jet on realistic simulated events, from isolated large radius jets up to the clustering of complete events with thousands of constituents. Standard analytic benchmarks of the jet substructure literature, such as the characteristic jet shapes of the anti kt algorithm and the predicted behaviour of groomed observables, are reproduced as well.
Timing studies demonstrate the advantage of the method. Including all data movement from file to device, the Triton kernels cluster events several times faster than CPU FastJet on a single core, at the level of tens of microseconds per event on current GPUs. Exact clustering and substructure thus become cheap enough to live inside a training loop, and the library offers a concrete building block for GPU resident reconstruction, where jets can be formed and analyzed on the accelerator without ever returning to the host.