利用 Green Contexts 控制 GPU 任务调度与资源分配
如今的 GPU 应用越来越多地由多个独立组件构成,它们在单进程内同时运行:延迟敏感型算子与面向吞吐量的后台内核并行;数据预处理阶段与模型推理同步进行;或者多个处理工作流阶段共享同一块 GPU。
控制这些组件间 GPU 资源的共享方式依然困难重重。组件间可能存在不可预测的干扰,且现有工具在资源分区方面的能力有限。
Green contexts 通过允许应用显式选择 GPU 执行资源的一个子集,并将任务直接定向到这些资源上,从而解决了这一问题。传统的 CUDA context 并非为此种使用模式设计。它们开销大,存在硬件上下文切换成本,且反映了早期 GPU 规模较小、应用通常作为单一主导工作负载运行的假设。
自 NVIDIA CUDA 12.4 起,Green contexts 已可通过 Driver API 访问。从 CUDA 13.1 开始,Runtime API 也支持 Green contexts,使应用能够在其进程内显式定义任务运行位置及执行资源的分配方式。
Green contexts 的工作原理
Green contexts 的一个主要用途是 SM 分区。通过将特定子集的 SM 分配给一个 Green context,应用可以将经由该 context 提交的任务定向到这些 SM 上。这使得多个工作负载能够并发运行于 GPU,而无需竞争相同的计算单元。
除了 SM 分区外,Green contexts 还可以配置 workqueue 资源。在传统模型中,独立的流排序工作负载可能映射到相同的底层 workqueue,即使在执行资源充足的情况下,也可能导致意外的串行化。通过显式配置 workqueue,Green contexts 允许应用表达预期的并发度,从而减少虚假依赖。
Green contexts 的创建和销毁开销很小,而且不会隐式同步无关的 GPU 工作。它还提供了更明确的编程模型:应用程序可以直接把工作提交到指定的 green context,而不是只依赖隐式的设备状态。
用 green contexts 实现显式编程模型
以往,CUDA Runtime 应用的典型做法是先用 cudaSetDevice() 选择设备,再创建 stream 并向其中提交工作。stream 的执行目标由创建时调用线程的当前 device 或 context 隐式决定。
Green contexts 让这种指定方式变得更显式:应用先为选定的一组资源创建 green context,再基于它创建 stream,提交到这些 stream 的工作就与该 green context 的资源绑定。
在 CUDA Runtime API 中,green context 用 cudaExecutionContext_t 类型表示,这是 Runtime 对 CUDA context 的一层抽象。对 green context 的使用来说,关键在于 cudaGreenCtxCreate() 返回的句柄可以直接传给 cudaExecutionCtxStreamCreate() 等 API,而不必依赖隐式的线程本地 device 或 context 状态。
这一点在应用创建 stream 的方式上体现得最直接。
注意:为简洁起见,以下代码片段省略了完整的错误检查。生产代码应检查所有 CUDA Runtime API 的返回值,在 kernel 启动后调用 cudaGetLastError() 或 cudaPeekAtLastError(),并检查 cudaStreamSynchronize() 等同步调用,以捕获异步执行错误。
传统 Runtime 流程:stream 的目标来自当前 device 状态
cudaSetDevice(device); cudaStream_t s; cudaStreamCreate(&s); kernel<<<grid, block, 0, s>>>();
Green context 流程:stream 的目标被显式指定
cudaExecutionContext_t greenCtx; cudaGreenCtxCreate(&greenCtx, desc, device, 0); cudaStream_t s; cudaExecutionCtxStreamCreate(&s, greenCtx, 0, 0); kernel<<<grid, block, 0, s>>>();
代码改动极少,但心智模型更清晰。应用程序不再依赖隐式的线程本地设备状态,而是显式选择 Green Context 并为它创建 stream。其余应用代码可继续沿用基于 stream 的 CUDA 编程模型。
希望面向整个设备的应用可继续使用传统 Runtime 模型。对于需要显式 context 句柄的 API,可通过 cudaDeviceGetExecutionCtx() 获取设备级 context。
示例
GPU 工作负载中,小体量、延迟敏感的 kernel 常需与更大的吞吐量导向 worker 共享设备。典型案例是分布式训练和推理中的通信/GEMM 重叠,或 AI 传感器处理平台(如 NVIDIA Holoscan)中那些要求尽快启动的延迟敏感算子。
此类关键任务的标准工具之一是 CUDA stream 优先级。但仅设置优先级不足以保证高优先级 kernel 立即执行——当 bulk kernel 占满 GPU 全部 SM 时尤甚。若 bulk kernel 已排队且关键 kernel 随之到达,调度器会向下一个腾出的 SM 分配关键 kernel 的 block。由于 stream 优先级无法抢占已在 SM 上执行的 block,关键 kernel 需等待部分 block 排空。
下面创建一个示例并展示结果。以下是需要编写的新代码:
// 1. 查询设备上的所有 SM。
cudaDevResource all_sm {};
cudaDeviceGetDevResource(dev, &all_sm, cudaDevResourceTypeSm);
// 2. 隔离出一个关键分区(遵循架构的协同调度对齐规则)。
cudaDevSmResourceGroupParams crit_params {};
crit_params.smCount = all_sm.sm.smCoscheduledAlignment;
cudaDevResource critical_res{}, remaining_res{};
cudaDevSmResourceSplit(&critical_res, 1, &all_sm, &remaining_res, 0, &crit_params);
// 3. 为每个分区打包独立的 workqueue-config 资源以隔离队列压力,随后生成描述符。
cudaDevResource wq {};
wq.type = cudaDevResourceTypeWorkqueueConfig;
wq.wqConfig.device = dev;
wq.wqConfig.sharingScope = cudaDevWorkqueueConfigScopeGreenCtxBalanced;
wq.wqConfig.wqConcurrencyLimit = 2;
cudaDevResource crit_pack[2] = { critical_res, wq };
cudaDevResourceDesc_t crit_desc{};
cudaDevResourceGenerateDesc(&crit_desc, crit_pack, 2);
cudaDevResource bulk_pack[2] = { remaining_res, wq };
cudaDevResourceDesc_t bulk_desc{};
cudaDevResourceGenerateDesc(&bulk_desc, bulk_pack, 2);
// 4. 创建每个绿色上下文,并在其上创建流。
cudaExecutionContext_t crit_ctx{};
cudaGreenCtxCreate(&crit_ctx, crit_desc, dev, 0);
cudaStream_t crit_stream{};
cudaExecutionCtxStreamCreate(&crit_stream, crit_ctx, cudaStreamNonBlocking, prio_high);
cudaExecutionContext_t bulk_ctx{};
cudaGreenCtxCreate(&bulk_ctx, bulk_desc, dev, 0);
cudaStream_t bulk_stream{};
cudaExecutionCtxStreamCreate(&bulk_stream, bulk_ctx, cudaStreamNonBlocking, prio_low); // 定义 prio_low
其余代码看起来应该很熟悉:
// 占满 bulk 分区。
for (int i = 0; i < BULK_LAUNCHES; ++i) {
bulk_kernel<<<bulk_grid, bulk_block, 0, bulk_stream>>>(
d_bulk, N_BULK, BULK_ITERS);
}
// 在其独立分区上启动关键内核。
cudaEventRecord(t_start, crit_stream);
critical_kernel<<<crit_grid, crit_block, 0, crit_stream>>>(d_crit, N_CRIT);
cudaEventRecord(t_stop, crit_stream);
cudaStreamSynchronize(crit_stream);
float crit_ms = 0.0f;
cudaEventElapsedTime(&crit_ms, t_start, t_stop);
流承载了分区信息,内核只能在其被允许运行的 SM 上执行。
我们测试三种模式,每种模式均包含相同的关键和 bulk 工作负载:
- 模式 A:绿色上下文分区,高优先级关键流
- 模式 B:默认 context,关键流使用高优先级(不分区)
- 模式 C:默认 context,两个流均为普通优先级(不设流优先级,也不分区)
关键 kernel 是一个很小的整数计算负载。主力 kernel 则连续启动 20 次 400 万线程的 sqrtf 循环,目的是把设备填满。我们在主力 kernel 运行期间测量关键 kernel 的墙钟延迟。
在一块拥有 148 个 SM 的 NVIDIA Blackwell GPU 上:
| 模式 | 关键 kernel 延迟 |
|---|---|
| Green context(关键任务 8 个 SM / 主力任务 140 个 SM) | 0.007 ms |
| 默认 context,关键任务高优先级 | 0.140 ms |
| 默认 context,同等优先级 | 3.727 ms |
表 1:不同执行模式下关键 kernel 的延迟
仅靠流优先级就有很大提升,比同优先级流抢占同一批 SM 快约 27 倍。但优先级有上限:即便 GPU 调度器优先照顾关键 kernel,仍要多等 20 倍的时间,让主力任务的线程块排空。Green contexts 则完全省掉了这段等待,因为关键 SM 是专用的。代价是主力 kernel 能用的 SM 变少了。
同样的骨架也可以用来实现预期的并发效果,而不只是降低延迟,比如让通信 kernel 和 GEMM kernel 重叠运行、双双推进。进阶用法往往会把 green context 分区和流优先级结合起来,以达到目标性能。
何时使用 green contexts
当你的应用满足以下情况时,green contexts 最有用:
- 在同一个 GPU 上并发运行多个相互独立的负载
- 需要更可预期的性能,而非尽力而为的流调度
- 希望显式控制 GPU 执行资源的划分方式
- 因对 workqueue 控制不足而出现性能下降
- 正朝单进程内的流水线式或多组件执行演进
上手指南
- 使用 CUDA 13.1 或更高版本构建应用。
- 查询设备资源,并针对目标资源分区创建 green context。
- 在该 green context 上创建流,沿用基于流的 CUDA 编程模型提交任务。
- 对于显式传入上下文句柄、且需要面向完整设备的 API,请调用
cudaDeviceGetExecutionCtx()。
启用 Green Context 属于可选操作,且对现有功能具有叠加性。现有应用可继续针对完整设备运行而无需改动,并在需要更细粒度执行控制时逐步引入 Green Context。
更多细节与示例,请参阅 CUDA 编程指南 及 CUDA 运行时 API 文档,不妨现在就试试看 Green Context。