An example where gemm and all-scatter are independent

[neoblizz]

This pull request introduces two new files to the examples/20_gemm_all_scatter_independent directory, providing benchmarking and kernel implementations for distributed GEMM (General Matrix Multiply) using Triton and Iris. The changes add a full benchmarking script and a ring-based all-reduce GEMM kernel, enabling efficient multi-GPU matrix multiplication and communication. These additions support flexible configuration, validation, and performance measurement for distributed compute scenarios.

New benchmarking and execution script

• Added a comprehensive benchmarking script benchmark.py, supporting distributed execution with PyTorch, command-line configuration for matrix dimensions, datatypes, block sizes, and benchmarking/validation modes. The script manages process spawning, distributed setup, memory allocation, kernel timing, validation, and performance logging.

New Triton kernel implementation for distributed GEMM

• Added gemm_all_reduce_ring_based.py, implementing two Triton kernels: persistent_gemm for local matrix multiplication and persistent_all_reduce for ring-based distributed reduction and scatter of results. These kernels use advanced synchronization and communication primitives for efficient multi-GPU execution.

Integration with Iris and validation utilities

• Both files integrate with the Iris library for device memory management, inter-GPU communication, and synchronization, and use shared utilities for validation and timestamping. [1] [2]

Performance measurement and output

• The benchmarking script collects kernel timings, calculates TFLOPS, and outputs results and traces in JSON format for further analysis.

Support for flexible configuration and debugging

• Both files support multiple datatypes, block sizes, debugging, and tracing options, allowing for extensive experimentation and profiling. [1] [2]

[Copilot]

Pull Request Overview

This PR introduces a new example demonstrating independent GEMM and all-scatter operations in a distributed setting. The implementation provides two algorithmic approaches: bulk synchronous all-scatter and ring-based all-reduce, with comprehensive benchmarking and validation capabilities for multi-GPU matrix multiplication scenarios.

• Adds distributed GEMM implementations with two communication strategies (bulk synchronous and ring-based)
• Provides a comprehensive benchmarking framework with timing, validation, and trace collection
• Integrates with Iris library for multi-GPU memory management and communication primitives

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
matmul_wrapper.py	PyTorch autograd wrapper for distributed GEMM kernel execution with debugging and timing support
gemm_all_scatter_bulk_synchronous.py	Triton kernels for persistent GEMM and bulk synchronous all-scatter communication
gemm_all_reduce_ring_based.py	Alternative implementation using ring-based all-reduce with more complex synchronization
benchmark.py	Comprehensive benchmarking script with distributed execution, validation, and performance measurement

[Copilot]

Pull Request Overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.