用 Green Contexts 控制 GPU 如何共享工作

中文精读

原文约 2,092 词 · 7 节

1. 背景：单一进程内的多组件并发与资源争用

   • 现代 GPU 应用常在同一进程里同时跑多个独立组件：延迟敏感的算子与吞吐型后台内核、数据预处理与模型推理、或同一工作流的多个阶段，它们共享一块 GPU。
   • 这些组件之间如何共享 GPU 资源难以控制，互相干扰不可预测，现有手段对资源划分的能力有限。
   • 绿色上下文（green contexts）让应用显式选取一部分 GPU 执行资源，并把工作直接投递到这些资源上。
   • 传统 CUDA 上下文不是为这种用法设计的：它们较重、会带来硬件上下文切换开销，且沿用了 GPU 规模较小、应用通常是单一主导负载时代的假设。
   • 绿色上下文自 CUDA 12.4 起可通过驱动 API 使用；从 CUDA 13.1 起也可通过运行时 API 使用，使应用能在自身进程内定义工作在哪里运行、执行资源如何划分。

2. 绿色上下文的工作方式

   • 主要用途之一是 SM 划分：把特定的 SM 子集分配给某个绿色上下文，投递到该上下文的工作就只在这些 SM 上运行，从而让多个负载并发执行而不争抢同一批计算单元。
   • 除 SM 外，绿色上下文还能分配工作队列（workqueue）资源。传统模型中相互独立的流顺序负载可能映射到同一底层工作队列，即使执行资源充足也会产生非预期的串行化。
   • 显式分配工作队列后，应用可以表达期望的并发度，减少虚假依赖。
   • 绿色上下文的创建与销毁开销很小，并且创建或销毁它不会隐式同步无关的 GPU 工作。
   • 它提供了更显式的编程模型：应用把工作定向到选定的绿色上下文，而不是仅依赖隐式的当前设备状态。

3. 使用绿色上下文的显式编程模型

   • 以往 CUDA 运行时应用通常先用 cudaSetDevice() 选设备，再创建流并提交工作；流的执行目标由创建流时调用线程当前的设备或上下文隐式推断。
   • 绿色上下文让这种定向变得显式：应用先为选定的一组资源创建绿色上下文，再从该上下文创建流，提交到这些流的工作就与该上下文的资源关联。
   • 运行时 API 用 cudaExecutionContext_t 类型表示绿色上下文，它是对 CUDA 上下文的一层运行时抽象；关键在于 cudaGreenCtxCreate() 返回的句柄可直接传给 cudaExecutionCtxStreamCreate() 等 API，而不依赖隐式的线程本地设备/上下文状态。
   • 原文 Note 说明：为简洁起见，代码片段省略了完整的错误检查；生产代码应检查所有 CUDA 运行时 API 的返回值，在内核启动后使用 cudaGetLastError() 或 cudaPeekAtLastError()，并对 cudaStreamSynchronize() 这类同步调用做检查，以捕获异步执行错误。
   • 代码改动很小但心智模型更清晰：不再依赖隐藏的线程本地设备状态，应用显式选择绿色上下文并为它创建流，其余代码仍可沿用同一套基于流的 CUDA 编程模型。
   • 想面向整块设备的应用可继续使用传统运行时模型；对接受显式上下文句柄的 API，可用 cudaDeviceGetExecutionCtx() 获取设备级上下文。
   c

   cudaSetDevice(device); cudaStream_t s; cudaStreamCreate(&s); kernel<<>>(); 
   
   cudaExecutionContext_t greenCtx; cudaGreenCtxCreate(&greenCtx, desc, device, 0); cudaStream_t s; cudaExecutionCtxStreamCreate(&s, greenCtx, 0, 0); kernel<<>>();

4. 示例：临界内核与大吞吐内核共存

   • 典型场景是需要一个小型延迟敏感内核与大型吞吐型工作共享设备，例如分布式训练与推理中的通信/GEMM 重叠，或 NVIDIA Holoscan 这类 AI 传感器处理平台中要求尽快启动的算子。
   • 常用的优先级手段是 CUDA 流优先级，但当大内核占满全部 SM 时，仅靠优先级无法保证高优先级内核立即执行：调度器只能在内核线程块从 SM 释放后才分配，优先级无法抢占已在 SM 上执行的线程块。
   • 示例先用 cudaDeviceGetDevResource() 查询全部 SM，再按架构的协同调度对齐（smCoscheduledAlignment）切出一组临界资源，其余作为 bulk 资源。
   • 两个分区各打包一份工作队列配置资源（sharingScope 设为 cudaDevWorkqueueConfigScopeGreenCtxBalanced，并发上限 2），用 cudaDevResourceGenerateDesc() 生成描述符，再为每个分区创建绿色上下文及其流，分别取高、低优先级。
   • 提交部分保持不变：用 20 次背靠背的 4M 线程 sqrtf 循环批量内核占满 bulk 分区，在临界流上用事件记录测量小整数内核的时延。
   cuda

   // 1. Query all SMs on the device. cudaDevResource all_sm {}; cudaDeviceGetDevResource(dev, &all_sm, cudaDevResourceTypeSm);
   // 2. Carve out one critical group (arch's coscheduled alignment). 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. Pack each partition with its own workqueue-config resource so queue pressure is isolated too, then generate the descriptor.
   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. Create each green context, then a stream on it.
   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);

5. 测量结果与取舍

   • 对比三种模式：A 为绿地上下文分区 + 高优先级临界流；B 为默认上下文 + 高优先级临界流（不分分区）；C 为默认上下文、两流同为普通优先级（无优先级、无分区）。
   • 在拥有 148 个 SM 的 NVIDIA Blackwell GPU 上，临界内核时延分别为：A（临界 8 个 SM / bulk 140 个 SM）0.007 ms；B 0.140 ms；C 3.727 ms。
   • 单用流优先级相比同优先级竞争已快约 27 倍，但存在上限：即便调度器优先处理临界内核，它仍要等 bulk 线程块排空，额外损失约 20 倍。
   • 绿色上下文因为临界 SM 专属于该负载，完全跳过了等待；代价是可供 bulk 内核使用的 SM 变少。
   • 同一套骨架也能用于实现期望的并发而非降低时延，例如让通信内核与 GEMM 内核重叠、两者都取得进展；进阶用法常同时使用绿色上下文分区与流优先级来达成性能目标。

6. 适用场景

   • 在同一 GPU 上并发运行多个独立负载时。
   • 需要更可预测的性能，而非尽力而为的流调度时。
   • 希望对 GPU 执行资源如何划分做显式控制时。
   • 因对工作队列控制有限而出现性能下降时。
   • 应用正在演进为单进程内的流水线式或多组件执行时。

7. 入门步骤

   • 使用 CUDA 13.1 或更新版本构建应用。
   • 查询设备资源，并针对想要定向的资源分区创建绿色上下文。
   • 为该绿色上下文创建流，并用原有的基于流的 CUDA 编程模型提交工作。
   • 对接受显式上下文句柄且应面向整块设备的 API，使用 cudaDeviceGetExecutionCtx()。
   • 绿色上下文是可选、增量式的：现有应用可不改动地继续面向整块设备，只在需要更细执行控制的地方逐步引入。