Partition GPU Flexibly with CUDA Green Context (Chinese Version)

English version can be found here

Introduction

在一个 application 内部对 GPU 做 spatial sharing 进行多任务并发是一个常见的性能提升手段. 从 CUDA 一开始推出的时候, 他的第一组 guide 就介绍了如何通过不同的 stream 来做 compute 和 memcpy 的 overlap.


当然, multi stream 的方式也可以用于多个 compute kernel 的 spatial sharing, 然而不同于 compute 和memory, 多个 compute kernel 之间自然会有更加严重的 interference. 注意此处 interference 可能导致的结果是进行 spatial sharing 后性能比原先顺序执行还差, 或者其中一些 kernel 的性能受到严重影响. 从 SM 的角度看, 在 launch 的 CTA 总和小于 SM 个数时原则上不会有较大的影响, 当然实际情况可能会由于 undocumented 的 SM scheduler 行为产生差距; 而在 launch 的 CTA 个数较多时, SM scheduler 则会拒绝 launch 新的 kernel, 或者在不同的 kernel 上进行高频率的 SM context switch, 导致 compute throughput 急剧下降. 从 memory bandwidth 的角度看, 多个 stream 之间几乎不提供 memory bandwidth 的隔离, 因此每一个 kernel 的访存性能完全没有保障.

由此可见, 直接通过 multi stream 进行 concurrent kernel launch 在性能和 predictability 的角度都存在非常大的问题. 一个常见的 workaround 是手动指定每一个 kernel launch 的 CTA 个数从而确保 compute partition, 并进行精准的 scheduling 来避免 memory bandwidth contention, 这也是初代 Nanoflow [1] 的做法. 然而, 这一方法的精确性和灵活性都较差, 难以大规模推开.

为了保证 spatial sharing 的精确性, 此前一个常见的做法是使用 CUDA MPS (Multi-Process Service). 这一技术开始设计的时候是为了多个 application 之间在共享 GPU 资源的同时避免频繁的 whole GPU context switch, 允许 developer 为不同的 application 指定不同的 GPU 使用比例. 然而, 这一方法也可以用于一个 application 内部, 通过为不同的 subprocess 指定特定的 SM utilization, 从而实现精确的 spatial sharing. 然而, 这一方法的重要前提是并发的 kernel 必须跑在不同的 process 上, 那如果想要在他们之间进行数据交换就要走 CUDA IPC, 无疑会带来巨大的 overhead.

在 CUDA 12.4 中, NVIDIA 推出了 CUDA Green Context,

Reference

[1] NanoFlow: Towards Optimal Large Language Model Serving Throughput