nvidia_cusparselt_structured_sparse_matmul_with_2_4_pruning.py

1import torch
2from nvidia.cusparselt import cusparselt as cslt
3
4# Initialize cuSPARSELt handle
5handle = cslt.Handle()
6
7# Matrix dimensions (m, n, k) - Must be multiples of 16 for structured sparsity
8m, n, k = 32, 32, 32
9
10# Create dense matrices on GPU
11# Matrix A must be sparse-ready (2:4 structure)
12A = torch.randn(m, k, device='cuda', dtype=torch.float16)
13B = torch.randn(k, n, device='cuda', dtype=torch.float16)
14C = torch.zeros(m, n, device='cuda', dtype=torch.float16)
15
16# Create descriptors
17matA = cslt.MatDescriptor(handle, m, k, k, A.data_ptr(), A.dtype, cslt.Order.ROW)
18matB = cslt.MatDescriptor(handle, k, n, n, B.data_ptr(), B.dtype, cslt.Order.ROW)
19matC = cslt.MatDescriptor(handle, m, n, n, C.data_ptr(), C.dtype, cslt.Order.ROW)
20
21# Prune A to satisfy 2:4 structured sparsity requirement
22cslt.prune(handle, matA, cslt.PruneAlg.PRUNE_SPMMA_STRIP_2_4)
23
24# Compress matrix A
25A_compressed_size = cslt.get_compressed_size(handle, matA)
26A_compressed = torch.empty(A_compressed_size, device='cuda', dtype=torch.uint8)
27cslt.compress(handle, matA, A_compressed.data_ptr())
28
29# Execute Sparse Matrix-Matrix Multiplication (C = A * B)
30# Note: In a real scenario, you would select an algorithm/plan before execution
31cslt.matmul(
32    handle,
33    matA,
34    matB,
35    matC,
36    matC, # bias/accumulator
37    A_compressed.data_ptr(),
38    alpha=1.0,
39    beta=0.0
40)
41
42torch.cuda.synchronize()
43print("MatMul execution complete.")
44print(f"Result C shape: {C.shape}")