从MegaMoE数学到DeepGEMM源码:MoE模型的GPU优化全链路
2026/10/25 7:00:09
也是时隔很长时间重启博客了,记录一下对于学习重难点的整理,以代码为主,代码附加注释
import torch #torch.where a=torch.rand(4,4) b=torch.rand(4,4) print(a) print(b) out=torch.where(a>0.5,a,b) #如果a中元素满足就输出a,a不满足直接输出b中元素 print(out) #torch.index_select print("torch.index_select") a=torch.rand(4,4) print(a) out=torch.index_select(a,dim=0,index=torch.tensor([0,3,2])) #实际是去了0.3.2行生成了新的tensor print(out) #torch.gather print("torch.gather") a=torch.linspace(1,16,steps=16).view(4,4) #定义一个序列 通过view改变其shape print(a) out=torch.gather(a,dim=0, index=torch.tensor([[0,1,1,1], [0,1,2,2], [0,1,3,3]])) print(out) #更高维的gather ##输出索引关系dim=0,out[i,j,k]=input[index[i,j,k],j,k] ##输出索引关系dim=1,out[i,j,k]=input[i,index[i,j,k],k] #torch.masked_index print("torch.masked_index") a=torch.linspace(1,16,steps=16).view(4,4) #定义一个序列 通过view改变其shape mask=torch.gt(a,8) print(a) print(mask) out=torch.masked_select(a,mask) print(out) ##torch.take print("torch.take") a=torch.linspace(1,16,steps=16).view(4,4) b=torch.take(a,index=torch.tensor([0,15,13,10])) print(b) ##torch.nonzero print("torch.nonzero") a=torch.tensor([[0,1,2,0],[2,3,0,1]]) out=torch.nonzero(a) print(out) #表示非零元素位置,用于稀疏表示