[Feature Request] Introduce `T.gemm_sp` for 2:4 Sparse Tensor Core GEMM

[LeiWang1999]

Background:

Sparse matrix operations have become an essential part of optimizing deep learning models, especially when dealing with large-scale data and sparse neural networks. Current hardware accelerators, such as Tensor Cores, offer significant performance improvements for dense matrix operations but lack full optimization for sparse matrix operations.

The Tensor Cores have already been optimized for dense matrix-multiplication (GEMM) workloads, providing substantial speedups. However, there is still room for improvement in efficiently handling sparse matrices on Tensor Cores, particularly with specific sparsity patterns such as 2:4 sparse format.

The 2:4 sparsity pattern, in which two consecutive elements of a matrix are non-zero, offers a common and highly efficient structure for many deep learning operations, including convolutions and dense layers. However, Tensor Core GEMM operations currently do not natively support this sparsity pattern, which limits performance for sparse workloads.

Request:

To better utilize Tensor Cores for sparse matrix operations, we propose introducing T.gemm_sp, a specialized function designed for 2:4 Sparse Tensor Core GEMM operations. This would optimize matrix multiplication for matrices that follow a 2:4 sparse pattern, leveraging the hardware’s efficiency in handling sparsity.

• T.gemm_sp(A, B, C) should be defined to perform a GEMM operation between matrices A and B and store the result in matrix C.
• A and B should be sparse tensors represented in a 2:4 sparsity format.
• The implementation should ensure that non-zero values are efficiently processed using Tensor Core hardware, while zero values are skipped.

[botbw]

I might give a try on this :)

[isHuangXin]

I would like to try this out. Please assign it to me～

[isHuangXin]

It seems that the task is to implement a 2:4 Pattern Sparse Attention Matrix Multiplication. The initial idea can be broken down into three main steps:

Step 1: Convert the Matrix to 2:4 Pattern Sparse CSR Format

First, the matrix is divided into blocks and then stored in the CSR (Compressed Sparse Row) format. The primary task here is to write a function:

def convert_to_2_4_csr(sparse_matrix):
    """
    Convert a sparse matrix to the 2:4 sparsity pattern CSR format.

    Args:
        dense_matrix (np.ndarray): The input dense matrix (M x N).

    Returns:
        row_ptr (List[int]): The starting position of each row.
        col_idx (List[int]): The column indices of the non-zero elements.
        data (List[float]): The values of the non-zero elements.
    """
    M, N = dense_matrix.shape
    assert N % 4 == 0, "The number of columns must be a multiple of 4 to satisfy the 2:4 sparsity pattern"

    row_ptr = [0]
    col_idx = []
    data = []

    # xxxxxx
    # xxxxxx
 
    return row_ptr, col_idx, data

Step 2: Implement the Sparse Matrix Multiplication Function

Based on the CSR format from the first step, the next task is to implement the matrix multiplication function:

@T.prim_func
def gemm_sp_2_4_csr(
        row_ptr_A: T.Tensor((M + 1,), "int32"),  # Row pointers for CSR format of matrix A
        col_idx_A: T.Tensor((nnz_A,), "int32"),  # Column indices for CSR format of matrix A
        data_A: T.Tensor((nnz_A,), dtype),       # Non-zero values for CSR format of matrix A
        row_ptr_B: T.Tensor((K + 1,), "int32"),  # Row pointers for CSR format of matrix B (if B is also sparse)
        col_idx_B: T.Tensor((nnz_B,), "int32"),  # Column indices for CSR format of matrix B
        data_B: T.Tensor((nnz_B,), dtype),       # Non-zero values for CSR format of matrix B
        C: T.Tensor((M, N), dtype),              # Output dense matrix
):
    """
    2:4 sparse matrix multiplication, supporting CSR format inputs.
    """
    with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by):
        C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
        T.clear(C_local)

        # Iterate over each row of A
        for i in T.thread_binding(M):
            # xxxxxx
            # xxxxxx

                    # Accumulate to the corresponding position in C
                    C_local[i, j] += val_A * val_B

        # Write the results back to global memory
        T.copy(C_local, C)

Step 3: Implement the UnitTest 2:4 Pattern Sparse Matrix Multiplication

The test function is used to verify the correctness and performance of two implementations of gemm_sp_2_4 (one without considering the convert_to_2_4_csr format and one with). The test function generates random sparse matrices, converts them to the 2:4 sparse format, and then calls both implementations to perform matrix multiplication. Finally, it compares the correctness and performance of the results.

def test_gemm_sp_2_4():
    """
    Test two implementations of gemm_sp_2_4:
    1. Implementation without considering the convert_to_2_4_csr format.
    2. Implementation considering the convert_to_2_4_csr format.
    """
    # Matrix dimensions
    M, K, N = 1024, 1024, 1024  # Can be adjusted as needed
    sparsity = 0.5  # Sparsity (2:4 sparsity pattern)

    # Generate random sparse matrices A and B
    sparse_A = torch.randn(M, K, dtype=torch.float16, device="cuda")
    sparse_B = torch.randn(K, N, dtype=torch.float16, device="cuda")

    # Convert dense matrices to 2:4 sparse format
    row_ptr_A, col_idx_A, data_A = convert_to_2_4_csr(sparse_A.cpu().numpy())
    row_ptr_B, col_idx_B, data_B = convert_to_2_4_csr(sparse_B.cpu().numpy())

    # Initialize output matrices
    C1 = torch.zeros((M, N), dtype=torch.float16, device="cuda")
    C2 = torch.zeros((M, N), dtype=torch.float16, device="cuda")

    # Test the implementation without considering convert_to_2_4_csr
    print("Testing gemm_sp_2_4 without CSR...")
    start_time = time.time()
    block_mask_shape = (M // DEFAULT_BLOCK_M, N // DEFAULT_BLOCK_N, K // DEFAULT_BLOCK_K)
    ## xxxxxx

    # Test the implementation considering convert_to_2_4_csr
    print("Testing gemm_sp_2_4 with CSR...")
    start_time = time.time()
    gemm_sp_2_4_csr(
        torch.tensor(row_ptr_A, dtype=torch.int32, device="cuda"),
        torch.tensor(col_idx_A, dtype=torch.int32, device="cuda"),
        torch.tensor(data_A, dtype=torch.float16, device="cuda"),
        torch.tensor(row_ptr_B, dtype=torch.int32, device="cuda"),
        torch.tensor(col_idx_B, dtype=torch.int32, device="cuda"),
        torch.tensor(data_B, dtype=torch.float16, device="cuda"),
        C2
    )
    end_time = time.time()
    print(f"Execution time (with CSR): {end_time - start_time:.6f} seconds")

    # Validate the correctness of the results
    print("Validating results...")
    dense_C_ref = torch.matmul(sparse_A, sparse_B)  # Use dense matrix multiplication as a reference
    error1 = torch.abs(C1 - dense_C_ref).max().item()
    error2 = torch.abs(C2 - dense_C_ref).max().item()

    print(f"Maximum error (without CSR): {error1:.6f}")
    print(f"Maximum error (with CSR): {error2:.6f}")


# Call the test function
if __name__ == "__main__":
    test_gemm_sp_2_4()

[botbw]

It seems that the task is to implement a 2:4 Pattern Sparse Attention Matrix Multiplication. The initial idea can be broken down into three main steps:

Step 1: Convert the Matrix to 2:4 Pattern Sparse CSR Format

First, the matrix is divided into blocks and then stored in the CSR (Compressed Sparse Row) format. The primary task here is to write a function:

def convert_to_2_4_csr(sparse_matrix):
"""
Convert a sparse matrix to the 2:4 sparsity pattern CSR format.

Args:
    dense_matrix (np.ndarray): The input dense matrix (M x N).

Returns:
    row_ptr (List[int]): The starting position of each row.
    col_idx (List[int]): The column indices of the non-zero elements.
    data (List[float]): The values of the non-zero elements.
"""
M, N = dense_matrix.shape
assert N % 4 == 0, "The number of columns must be a multiple of 4 to satisfy the 2:4 sparsity pattern"

row_ptr = [0]
col_idx = []
data = []

# xxxxxx
# xxxxxx

return row_ptr, col_idx, data

Step 2: Implement the Sparse Matrix Multiplication Function

Based on the CSR format from the first step, the next task is to implement the matrix multiplication function:

@T.prim_func
def gemm_sp_2_4_csr(
row_ptr_A: T.Tensor((M + 1,), "int32"), # Row pointers for CSR format of matrix A
col_idx_A: T.Tensor((nnz_A,), "int32"), # Column indices for CSR format of matrix A
data_A: T.Tensor((nnz_A,), dtype), # Non-zero values for CSR format of matrix A
row_ptr_B: T.Tensor((K + 1,), "int32"), # Row pointers for CSR format of matrix B (if B is also sparse)
col_idx_B: T.Tensor((nnz_B,), "int32"), # Column indices for CSR format of matrix B
data_B: T.Tensor((nnz_B,), dtype), # Non-zero values for CSR format of matrix B
C: T.Tensor((M, N), dtype), # Output dense matrix
):
"""
2:4 sparse matrix multiplication, supporting CSR format inputs.
"""
with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=thread_num) as (bx, by):
C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
T.clear(C_local)

    # Iterate over each row of A
    for i in T.thread_binding(M):
        # xxxxxx
        # xxxxxx

                # Accumulate to the corresponding position in C
                C_local[i, j] += val_A * val_B

    # Write the results back to global memory
    T.copy(C_local, C)

Step 3: Implement the UnitTest 2:4 Pattern Sparse Matrix Multiplication

The test function is used to verify the correctness and performance of two implementations of gemm_sp_2_4 (one without considering the convert_to_2_4_csr format and one with). The test function generates random sparse matrices, converts them to the 2:4 sparse format, and then calls both implementations to perform matrix multiplication. Finally, it compares the correctness and performance of the results.

def test_gemm_sp_2_4():
"""
Test two implementations of gemm_sp_2_4:
1. Implementation without considering the convert_to_2_4_csr format.
2. Implementation considering the convert_to_2_4_csr format.
"""
# Matrix dimensions
M, K, N = 1024, 1024, 1024 # Can be adjusted as needed
sparsity = 0.5 # Sparsity (2:4 sparsity pattern)

# Generate random sparse matrices A and B
sparse_A = torch.randn(M, K, dtype=torch.float16, device="cuda")
sparse_B = torch.randn(K, N, dtype=torch.float16, device="cuda")

# Convert dense matrices to 2:4 sparse format
row_ptr_A, col_idx_A, data_A = convert_to_2_4_csr(sparse_A.cpu().numpy())
row_ptr_B, col_idx_B, data_B = convert_to_2_4_csr(sparse_B.cpu().numpy())

# Initialize output matrices
C1 = torch.zeros((M, N), dtype=torch.float16, device="cuda")
C2 = torch.zeros((M, N), dtype=torch.float16, device="cuda")

# Test the implementation without considering convert_to_2_4_csr
print("Testing gemm_sp_2_4 without CSR...")
start_time = time.time()
block_mask_shape = (M // DEFAULT_BLOCK_M, N // DEFAULT_BLOCK_N, K // DEFAULT_BLOCK_K)
## xxxxxx

# Test the implementation considering convert_to_2_4_csr
print("Testing gemm_sp_2_4 with CSR...")
start_time = time.time()
gemm_sp_2_4_csr(
    torch.tensor(row_ptr_A, dtype=torch.int32, device="cuda"),
    torch.tensor(col_idx_A, dtype=torch.int32, device="cuda"),
    torch.tensor(data_A, dtype=torch.float16, device="cuda"),
    torch.tensor(row_ptr_B, dtype=torch.int32, device="cuda"),
    torch.tensor(col_idx_B, dtype=torch.int32, device="cuda"),
    torch.tensor(data_B, dtype=torch.float16, device="cuda"),
    C2
)
end_time = time.time()
print(f"Execution time (with CSR): {end_time - start_time:.6f} seconds")

# Validate the correctness of the results
print("Validating results...")
dense_C_ref = torch.matmul(sparse_A, sparse_B)  # Use dense matrix multiplication as a reference
error1 = torch.abs(C1 - dense_C_ref).max().item()
error2 = torch.abs(C2 - dense_C_ref).max().item()

print(f"Maximum error (without CSR): {error1:.6f}")
print(f"Maximum error (with CSR): {error2:.6f}")

Call the test function

if name == "main":
test_gemm_sp_2_4()

@isHuangXin @LeiWang1999

Sorry I haven’t been able to dedicate much bandwidth to this as I’ve been tied up with other tasks. If this is a time-sensitive issue, feel free to transfer the issue or let me know if you'd like to tackle it collaboratively.

Basically, I was attempting to integrate the SM90 CUTLASS sparse tensor core GEMM (i.e. 62_hopper_sparse_gemm.cu), similar to the approach in gemm_sm90.h, but ran into challanges with the metadata layout: CUTLASS doesn’t store the sparse metadata in a plain format (see sm90_sparse_config.inl#L108-L117), which means I would have to:

1. Utilize cutlass's sparse converter, which would limit the usage T.gemm_sp as user must call dedicated APIs like T.convert_sp and cannot
2. Alternatively, allow the user to somehow annotate the metadata layout with tilelang layout and copy accordingly (where I got stuck, as currently I don't have much insight into tilelang's internal or TVM extension).

[LeiWang1999]

closed as been merged, huge thanks for @botbw 's contributions!