这是到目前为止（2023/10）应该是唯一一篇在认真思考并且实现 GPU 抢占的论文了。
之前很多论文都研究过 GPU 怎么混部，怎么做算力和显存的隔离，但是目标都是提升利用率，但是很难保证长尾延迟。比较适合训练，但不适合推理。
但是这篇论文其实也是有很多限制的：
- 不支持 nvidia GPU，只支持 amd
- 模型需要经过 tvm 的编译，因为抢占的实现依赖了编译技术
这篇论文主要的工作就2个事情：
- GPU 抢占：real-time任务启动的时候，best-effort 任务怎么做到快速的退出。
- dynamic kernel padding：怎么在尽可能保证 rt 任务 latency 的情况下，混部 be 的任务
最终效果：
- 2% 的延迟损失，换取 7.7x 倍的性能提升
1. 系统架构
reef 在 GPU runtime 的基础上扩展了4个核心组件：
- Task Queues：一个rt队列和多个be队列，每个队列绑定一个 GPU Stream
- Scheduler：一个调度器，2种模式，rt模式，normal模式，如果当前没有rt任务，调度器会切换到normal模式，如果有rt任务，会切换到rt模式。这个地方应该是为了和下面的抢占进行协同
- Preemption：抢占模块，实现rt抢占be
- DKP：dynamic kernel padding
前面2个没啥东西，重点讲一下抢占是怎么实现的吧，这个比较有意思
2. GPU 抢占
前提：
- DNN 模型里面所有计算都是幂等的。事实上现在几乎大部分模型都是了，全都是各种矩阵运算，没有上下文的，多次计算不影响最终结果
这个前提说明了 kernel 是 可以被中断的，这很关键
2.1. 抢占的原理
那抢占怎么实现呢，我们以前研究 GPU 混部的时候都知道，GPU 的 kernel 一旦发射到设备上了，是中断不了，必须等硬件执行完，才能进行下一个 kernel。所以 GPU kernel 抢占一直以来都是不好实现的
抢占的第一步，就是怎么剔除 Buffer 中的 kernel，也就是还没执行完的
Buffer 也分2种：
- 在 Runtime 中的，Host Queue & Device Queue
- 在 设备 上的，CU
第一个好处理，修改 GPU Runtime 的控制逻辑，直接剔除就行。所以这个地方依赖 GPU Runtime 开源，适合 AMD，不适合 nvidia
第二个就比较难了，但是论文提供了一个非常有意思的思路，引入编译器的插桩技术，扩展 tvm，在生成 kernel 的时候，在每个 kernel 前面插入一段代码，如果当前处于抢占状态，kernel 直接 return，不执行真正的计算逻辑了。
这样就大大提高了抢占的时效性
2.2. 抢占的具体实现
接着再来看下具体实现
抢占分2部分：
- 代码插桩，编译kernel的时候，需要给每个 kernel 插入 preempt_flag 来判断当前 kernel 是否被抢占了
- GPU runtime 在 kernel 调度的时候，如果有 rt 的 kernel，就需要设置 preempt_flag
代码插桩这块，我看了下并没有实际使用 tvm 编译器，而是直接对 .cu 源码文件修改，把 preempt_flag 插入到 cu 代码里面
这个 script 的作用就是识别 .cu 源文件里的函数，对每个函数都插入 preempted 的判断逻辑
extern "C" __global__ void {func_name}({params_def}) {{
if (*preempted) return;
{func_name}_device({params_call});
if (threadIdx.x + threadIdx.y + threadIdx.z == 0)
atomicAdd(task_slot, 1);
}}
下面有很多神经网络模型的示例，比如alexnet, resnet, mobilenet
已 alexnet 为例
- alexnet.cu 是原始的 cu 代码
- alexnet.be.cu 是经过 best_effort_kernel.py 转换后的可抢占的 cu 代码
当然这里面还涉及到一些其他细节，怎么把 preempt_flag 作为参数传给 kernel，其显存怎么分配，等等
抢占状态的同步
void REEFScheduler::loop_body() {
// Real-time Mode:
rtmode:
while (true) {
if (rt_queue->task_queue.empty()) goto bemode;
preempt_be_tasks();
auto rt_task = rt_queue->task_queue.front();
rt_queue->task_queue.pop();
execute_rt_task(rt_task);
}
// ...
}
如果 rt_queue 为空，切换到 be 模式，多个 GPU stream 并行执行 be 的 kernel
一旦有 rt kernel 被提交，立马设置 preempt_flag，这样还在设备缓存上的 kernel 一旦开始执行之后会立即返回
3. DKP：dynamic kernel padding
xx
牛逼
独一无二的马甲