Leetcode 311 Sparse Matrix Multiplication 稀疏矩阵相乘
Problem
Given two sparse matrices A and B, return the result of AB.
You may assume that A’s column number is equal to B’s row number.
Example:
A = [
[ 1, 0, 0],
[-1, 0, 3]
]
B = [
[ 7, 0, 0 ],
[ 0, 0, 0 ],
[ 0, 0, 1 ]
]
| 1 0 0 | | 7 0 0 | | 7 0 0 |
AB = | -1 0 3 | x | 0 0 0 | = | -7 0 3 |
| 0 0 1 |
Solution
- Transform B into defaultdict(dict)
from collections import defaultdict
def multiply(A, B):
# 预处理矩阵B的非零元素：建立行号到{列号: 值}的映射
b_nonzero = defaultdict(dict)
for k in range(len(B)): # 遍历B的每一行
for j in range(len(B[0])): # 遍历该行的每一列
if B[k][j] != 0:
b_nonzero[k][j] = B[k][j]
m = len(A)
n = len(B[0]) if B else 0
result = [[0] * n for _ in range(m)]
for i in range(m):
for k in range(len(A[i])):
a_val = A[i][k]
if a_val != 0:
if k in b_nonzero:
for j, b_val in b_nonzero[k].items():
result[i][j] += a_val * b_val
return result
A = [
[ 1, 0, 0],
[-1, 0, 3]
]
B = [
[ 7, 0, 0 ],
[ 0, 0, 0 ],
[ 0, 0, 1 ]
]
print(multiply(A,B))
GPU Solution，来自DeepSeek的总结
根据2025年最新研究进展，GPU处理稀疏矩阵乘法（SpMM/SparseGEMM）的主流算子可分为以下四个技术方向，均已在工业界和学术界得到广泛应用：
一、基于压缩格式的专用算子
-
CSR（Compressed Sparse Row）格式优化型
以NVIDIA cuSPARSE库的csrmm2为基础，通过合并行缓存（CRC）和共享内存优化，实现非零元素的连续访问。典型代表如GE-SpMM算子，通过将稀疏数据加载到共享内存，消除全局内存的随机访问开销，在GNN场景中相比cuSPARSE提升1.41倍性能。 -
块压缩混合型
DeepSeek提出的DeepGEMM支持MoE稀疏结构，采用即时编译（JIT）和Hopper TMA技术，在FP8精度下实现1350+ TFLOPS的算力，尤其擅长处理分组或掩码式稀疏乘法。该算子通过将稀疏块与密集块分离计算，在小矩阵场景加速比达2.7倍。
二、异构协同计算架构
-
CPU-GPU联合计算型
apspgemm框架通过自适应面板分割和NUMA感知调度，将大规模稀疏矩阵分解为规则面板分配给CPU和GPU协同计算。采用异步流水线技术重叠数据传输与计算，在十亿级边规模的社交网络图中实现7.21倍性能提升。其核心创新在于亲和度规则驱动的负载均衡算法。 -
量化稀疏混合型
清华团队KTransformers框架集成Marlin量化算子和llamafile稀疏算子，通过计算强度分析动态分配任务：将高计算强度MLA算子放在GPU，低强度稀疏路由放在CPU。这种异构划分使4090单卡可运行671B参数模型，推理速度达286 tokens/s。
三、硬件指令级优化
-
原子操作增强型
SpGEMM-Upper算法通过两阶段内存预分配策略，结合GPU原子操作实现行压缩并行化。在代数多重网格（AMG）计算中，相比传统方法减少38%的全局内存访问，压缩阶段采用Warp级协同写入优化。 -
稀疏脉动阵列型
密歇根大学改造TPU架构提出Sparse-TPU，支持矩阵分块压缩和条件执行。通过跳过全零块的计算，在Hopper架构上实现稀疏矩阵16倍加速，能耗降低4-19倍。其核心是扩展了传统脉动阵列对稀疏数据的适应性。
四、算法-硬件协同设计
-
哈夫曼调度型
MIT与NVIDIA联合提出的SpArch架构，通过哈夫曼编码动态调度非零块的计算顺序，结合行预取器实现输入复用。实验显示该方案将部分积矩阵数量减少58%，DRAM访问量降低2.8倍。 -
即时编译优化型
DeepGEMM的JIT编译器能根据稀疏模式自动生成最优内核代码，仅需300行核心逻辑即可超越专家手工调优性能。该技术特别适用于动态稀疏模式（如MoE模型中的专家路由）。
性能对比与选型建议
| 算子类型 | 适用场景 | 峰值性能 | 代表方案 |
|---|---|---|---|
| CSR格式优化 | 中小规模规则稀疏 | 200-400 GFLOPS | GE-SpMM |
| 异构协同计算 | 十亿级超大规模稀疏 | 5.6 TFLOPS | apspgemm |
| 量化稀疏混合 | 千亿参数MoE模型推理 | 286 tokens/s | KTransformers |
| 硬件指令级优化 | 科学计算/AMG求解器 | 16x加速比 | Sparse-TPU |
当前工业界主流选择是结合量化稀疏混合型与异构协同计算，如DeepSeek-R1和GPT-4o等大模型均采用类似KTransformers的架构。学术界则更关注基于新型压缩格式和硬件指令的深度优化，如apspgemm框架已被集成到AMGX数学库中。