为何CuPy模拟独热向量点积时性能远低于NumPy？

问题背景 #

尝试从零实现神经网络，因NumPy无法在GPU运行改用CuPy，预期性能提升，但实际CuPy版本耗时远高于NumPy：

• NumPy实现的核心计算耗时16.55秒
• CuPy对应实现耗时152.15秒（2分32秒）

NumPy代码实现 #

import numpy as np
import time

emb = 300  # embedding size
m = 2048   # minibatch size
V = 50000  # vocabulary size

# 生成随机独热向量
J = np.random.choice(emb, m) 
X = np.zeros((m, emb))
for i, j in enumerate(J):
    X[i, j] = 1

W0 = np.random.uniform(-0.8, 0.8, (V, emb)).astype("float32") 

# 核心计算
start_time = time.time()
for epoch in range(5):
  for mb in range(314): # 每个epoch的小批量数量
    h = np.zeros((1), dtype='float32')
    for xi in X.T: 
        w0i = np.argmax(xi)
        if not h.any(): h = W0[w0i]
        else: h = np.vstack((h, W0[w0i]))
print("%s seconds" % (time.time() - start_time))

CuPy代码实现 #

import cupy as cp
import time

emb = 300  # embedding size
m = 2048   # minibatch size
V = 50000  # vocabulary size

# 生成随机独热向量
J = cp.random.choice(emb, m) 
X = cp.zeros((m, emb))
for i, j in enumerate(J):
    X[i, j] = 1

W0 = cp.random.uniform(-0.8, 0.8, (V, emb)).astype("float32") # V⨯e

# 核心计算
start_time = time.time()
for epoch in range(5):
  for mb in range(314): # 小批量数量
    h = cp.zeros((1), dtype='float32')
    for xi in X.T: # 循环m次
        w0i = cp.argmax(xi)
        if not h.any(): h = W0[w0i] # (1xe) 
        else: h = cp.vstack((h, W0[w0i]))
print("%s seconds" % (time.time() - start_time))

核心问题：GPU计算特性被完全忽略 #

你的代码完全照搬了NumPy的逐元素循环+动态堆叠逻辑，这直接违背了GPU的设计优势：

• GPU擅长批量并行计算，而非细粒度循环：CuPy中逐列循环处理单个向量，每次都会触发GPU核函数启动，而核启动有固定开销，这种操作会彻底浪费GPU的并行能力，反而因频繁的开销拖慢速度。
• 动态vstack是性能灾难：每次vstack都会重新分配内存并复制数据，GPU的内存拷贝延迟远高于CPU，多次动态扩容会产生大量无效内存操作。
• 独热向量的处理可完全向量化：你的核心需求是通过独热向量的索引从W0中提取对应行，本质是批量索引操作，完全不需要循环。

优化方案：用向量化操作替代循环 #

直接利用索引的批量特性，一步完成所有行的提取，彻底消除循环和动态内存操作：

优化后的CuPy代码 #

import cupy as cp
import time

emb = 300  # embedding size
m = 2048   # minibatch size
V = 50000  # vocabulary size

# 直接生成索引（无需构建独热矩阵X，节省内存和构建时间）
J = cp.random.choice(emb, m) 

W0 = cp.random.uniform(-0.8, 0.8, (V, emb)).astype("float32")

# 核心计算：批量索引提取行，完全替代循环和vstack
start_time = time.time()
for epoch in range(5):
  for mb in range(314):
    h = W0[J]
print("%s seconds" % (time.time() - start_time))

优化逻辑说明 #

• 消除循环开销：W0[J]会让GPU一次性并行处理所有索引，充分利用GPU的多核心优势。
• 避免动态内存操作：直接生成固定形状的结果矩阵，无需反复扩容和内存拷贝。
• 简化数据表示：独热矩阵本质是索引的冗余表示，直接用索引可节省内存和预处理时间。

额外补充：你之前发现NumPy模拟比np.dot快，是因为np.dot做通用矩阵乘法，而你的模拟利用了独热矩阵的特性做了简化；但这种简化的循环逻辑完全不适合GPU，只有向量化操作才能发挥GPU的性能优势。

内容的提问来源于stack exchange，提问作者Mahesha999