Flash3D: Super-scaling Point Transformers through Joint Hardware-Geometry Locality

TL;DR

Flash3D Transformer unifies geometric locality and GPU architecture to enable super-scaling of point cloud models, achieving significant speed and memory efficiency improvements over previous methods.

Contribution

Introduces a novel locality mechanism based on Perfect Spatial Hashing that aligns geometric locality with GPU tiling, enabling scalable and efficient point cloud transformers.

Findings

2.25x speed increase over state-of-the-art PTv3

2.4x memory efficiency boost

Higher task accuracies at same compute budget

Abstract

Recent efforts recognize the power of scale in 3D learning (e.g. PTv3) and attention mechanisms (e.g. FlashAttention). However, current point cloud backbones fail to holistically unify geometric locality, attention mechanisms, and GPU architectures in one view. In this paper, we introduce Flash3D Transformer, which aligns geometric locality and GPU tiling through a principled locality mechanism based on Perfect Spatial Hashing (PSH). The common alignment with GPU tiling naturally fuses our PSH locality mechanism with FlashAttention at negligible extra cost. This mechanism affords flexible design choices throughout the backbone that result in superior downstream task results. Flash3D outperforms state-of-the-art PTv3 results on benchmark datasets, delivering a 2.25x speed increase and 2.4x memory efficiency boost. This efficiency enables scaling to wider attention scopes and larger models without additional overhead. Such scaling allows Flash3D to achieve even higher task accuracies than PTv3 under the same compute budget.