RandMScan: accelerating parallel scan via matrix computation and random-jump strategy.

Area of Science:

• Computer Science
• Artificial Intelligence

Background:

• Parallel scan is crucial for diverse computational tasks like sorting and large language model inference.
• Current GPU-optimized scans are inefficient on AI accelerators lacking vector units but possessing processing element (PE) arrays.
• Existing matrix-based methods face bandwidth and communication challenges, limiting scalability.

Purpose of the Study:

• To propose RandMScan, a parallel scan framework optimized for AI accelerators with PE arrays.
• To address the inefficiencies and scalability limitations of existing scan algorithms on modern hardware.
• To enhance the performance of scan operations in AI-centric workloads.

Main Methods:

• Developed a two-stage parallel scan framework: RandMScan.
• Stage one utilizes a matrix-based local chunk scan exploiting PE array parallelism.
• Stage two employs a Random-Jump strategy for efficient global aggregation with reduced synchronization.

Main Results:

• RandMScan achieves over 80% speedup compared to existing matrix-based methods for long sequences.
• Reduces end-to-end latency by 15%-26% in downstream applications.
• Demonstrates scalable execution over long sequences with mitigated communication overhead.

Conclusions:

• RandMScan effectively leverages AI accelerator PE arrays for efficient parallel scan.
• The proposed framework overcomes bandwidth and communication bottlenecks of prior solutions.
• RandMScan offers a scalable and high-performance alternative for scan primitives on AI hardware.