Reproduce “Fleet: Hierarchical Task-based Abstraction for Megakernels on Multi-Die GPUs” on OpenResearch

About this paper

Sangeeta Chowdhary, Ryan Swann, Sean Siddens, Muhammad Osama, Stephen Neuendorffer, Alexandru Dutu, Karthik Sangaiah, Sandeepa Bhuyan and 2 more

This work introduces Fleet, a hierarchical task-based abstraction and runtime system designed for multi-die GPUs with partitioned L2 caches. Fleet enables efficient data locality and synchronization at the chiplet level, achieving up to 1.56x faster LLM inference decode latency at small batch sizes and significantly increasing L2 cache hit rates to 61.4% with up to 37% less HBM traffic at larger batch sizes compared to chiplet-unaware baselines.

Key insights

• Explicitly exposing chiplet-level memory hierarchy through abstractions like Chiplet-tasks is critical for optimizing performance on multi-die GPUs.
• Cooperative weight sharing within a chiplet's L2 cache significantly improves data reuse and reduces HBM traffic, even when individual weight partitions exceed L2 capacity.
• Hierarchical synchronization and per-chiplet scheduling drastically reduce kernel launch overheads and amortize expensive global synchronization costs, contributing to overall performance gains.

Summary by alphaXiv.