LightSeq的技术亮点主要在于定制的Operation、动态GPU内存复用和层级的解码优化，本文对这三方面分别做个粗浅的介绍。

论文链接：https://arxiv.org/pdf/2010.13887.pdf

1 定制的Operation 

大多数深度学习框架通过调用方法库中的kernel function来实现编码器-解码器计算。这些kernel function一般是细粒度的，通常需要调用多个kernel function来实现一个功能模块。

以tensorflow实现的layer normalization为例，即使基于compilation optimization techniques （automatically merging broadcast and element-wise operations），仍需要三个kf和两个中间结果（mean and variance）的读写operation。

基于CUDA，我们可以为layer normalization写一个定制的kf，并且将两个中间结果写入到寄存器中，从而只需一次kernel launch且不需要将中间结果写入到gpu内存中，大幅减少了计算耗时。

基于这个idea，LightSeq结合cuBLAS的GEMM和定制kf来实现Transformer layers：

GEMM间的operation都通过定制kf实现。

2 动态GPU内存复用

为了节省gpu内存开销以及避免在计算过程中分配或释放gpu内存，LightSeq预定义了最大的dynamic shape，比如最大序列长度。在服务开始时，计算过程中的中间过程都在gpu内存中分配了最大值。

另外，非独立的中间结果共享GPU内存。通过这种内存复用策略，可以在T4卡上同时部署8个大型的Transformer模型。

3 层级的解码优化

在auto-regressive sequence generation任务中，最复杂的部分通常是解码。LightSeq支持多种解码方法，比如beam search、diversity beam search、top-k/top-p sampling等。同时可以达到几倍的加速。

在深度学习框架中，实现解码需要选取top-k probability的token，这一步骤涉及计算softmax和进行大小与字典大小成正比的GPU内存读写操作。通常vocabulary的size上千，因此解码在auto-regressive sequence generation任务的耗时中占了很大比例。

如下图所示，LightSeq结合GPU计算特点，同时借鉴推荐系统中召回（retrieve）和重排序（rerank）两步策略，设计了一种层级的召回 + 重排序策略，能够对成百上千的元素快速排序。