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）的主流算子可分为以下四个技术方向，均已在工业界和学术界得到广泛应用：

一、基于压缩格式的专用算子

1. CSR（Compressed Sparse Row）格式优化型
   以NVIDIA cuSPARSE库的csrmm2为基础，通过合并行缓存（CRC）和共享内存优化，实现非零元素的连续访问。典型代表如GE-SpMM算子，通过将稀疏数据加载到共享内存，消除全局内存的随机访问开销，在GNN场景中相比cuSPARSE提升1.41倍性能。

2. 块压缩混合型
   DeepSeek提出的DeepGEMM支持MoE稀疏结构，采用即时编译（JIT）和Hopper TMA技术，在FP8精度下实现1350+ TFLOPS的算力，尤其擅长处理分组或掩码式稀疏乘法。该算子通过将稀疏块与密集块分离计算，在小矩阵场景加速比达2.7倍。

二、异构协同计算架构

3. CPU-GPU联合计算型
   apspgemm框架通过自适应面板分割和NUMA感知调度，将大规模稀疏矩阵分解为规则面板分配给CPU和GPU协同计算。采用异步流水线技术重叠数据传输与计算，在十亿级边规模的社交网络图中实现7.21倍性能提升。其核心创新在于亲和度规则驱动的负载均衡算法。

4. 量化稀疏混合型
   清华团队KTransformers框架集成Marlin量化算子和llamafile稀疏算子，通过计算强度分析动态分配任务：将高计算强度MLA算子放在GPU，低强度稀疏路由放在CPU。这种异构划分使4090单卡可运行671B参数模型，推理速度达286 tokens/s。

三、硬件指令级优化

5. 原子操作增强型
   SpGEMM-Upper算法通过两阶段内存预分配策略，结合GPU原子操作实现行压缩并行化。在代数多重网格（AMG）计算中，相比传统方法减少38%的全局内存访问，压缩阶段采用Warp级协同写入优化。

6. 稀疏脉动阵列型
   密歇根大学改造TPU架构提出Sparse-TPU，支持矩阵分块压缩和条件执行。通过跳过全零块的计算，在Hopper架构上实现稀疏矩阵16倍加速，能耗降低4-19倍。其核心是扩展了传统脉动阵列对稀疏数据的适应性。

四、算法-硬件协同设计

7. 哈夫曼调度型
   MIT与NVIDIA联合提出的SpArch架构，通过哈夫曼编码动态调度非零块的计算顺序，结合行预取器实现输入复用。实验显示该方案将部分积矩阵数量减少58%，DRAM访问量降低2.8倍。

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数学库中。