生成式人工智能/大语言模型

具有条件节点的 CUDA 图的动态控制流

CUDA 图 可以显著提高性能,因为驱动程序能够使用任务和依赖关系的完整描述来优化执行。特别是在静态工作流中,图形可以提供难以置信的好处,其创建的开销可以在多次连续启动中分摊,从而提高整体性能。

然而,几乎所有的问题都涉及某种形式的决策,这可能需要分解图并将控制权返回给 CPU,以决定下一步要启动哪些工作。像这样分解工作会损害 CUDA 进行优化的能力,占用 CPU 资源,并增加每次图形启动的开销。

从 CUDA 12.4 开始,CUDA Graphs 支持 条件节点,这使得图形的部分能够有条件地或重复地执行,而不需要将控制返回到 CPU。这释放了 CPU 资源,使得更多的工作流能够在单个图形中表示,从而提高了计算效率。

条件节点

条件节点有两种风格:

  • IF 节点:如果条件值为 true,则每次评估节点时执行一次主体。
  • WHILE 节点:只要条件值为 true,当对节点求值时,就会重复执行主体。

条件节点是容器节点,类似于子图节点,但节点中包含的图的执行取决于条件变量的值。与节点关联的条件值由必须在节点之前创建的句柄访问,该句柄还可以指定在图形的每个开始处应用的初始化。可以通过调用在 CUDA 内核中设置条件值的 cudaGraphSetConditional 函数。

创建条件节点时,还会创建一个空图,并将句柄返回给用户。此图与节点绑定,并将根据条件值执行。此条件体图可以使用 图形 API 或者通过使用捕获异步 CUDA 调用的 cudaStreamBeginCaptureToGraph 函数。

条件节点也可以嵌套。例如,可以使用包含条件 IF 节点的体图创建条件 WHILE 节点。

条件节点体图可以包含以下任意一项:

  • 内核节点(CNP,当前不支持协同)
  • 空节点
  • 子图节点
  • Memset 节点
  • Memcopy 节点
  • 条件节点

这递归地应用于子图和条件体。所有内核,包括嵌套条件句或任何级别的子图中的内核,都必须属于同一个 CUDA 上下文。Memcopies 和 memset 必须作用于可从条件节点的上下文访问的内存。

完整的样品代码可在 CUDA 样本库 中找到。下一节将通过一些示例来展示如何处理条件节点的方法。

条件 IF 节点

如果条件为非零,则无论何时评估 IF 节点,都将执行一次 IF 节点的体图。图 1 描述了一个图,其中中间节点 B 是包含四节点图的 IF 条件节点:

Diagram depicting a decision workflow using three circular nodes with a conditional node (b) set to conditionally run its body graph.
图 1。条件 IF 节点

为了显示如何创建此图,以下示例使用条件节点 B 上游的内核节点 A,根据该内核所做的工作结果设置条件的值。条件的主体是使用图 API 填充的。

首先,定义节点 A 内核。这个内核根据用户执行的一些任意计算的结果来设置条件句柄。

__global__ void setHandle(cudaGraphConditionalHandle handle)
{
    unsigned int value = 0;
    // We could perform some work here and set value based on the result of that work.
    if (someCondition) { 
        // Set ‘value’ to non-zero if we want the conditional body to execute
        value = 1;
    }
    cudaGraphSetConditional(handle, value);
}

接下来,定义一个函数来构建图。此函数分配条件句柄、创建节点并填充条件图的主体。为了清楚起见,省略了启动和执行图形的代码。

cudaGraph_t createGraph() {
    cudaGraph_t graph;
    cudaGraphNode_t node;
    void *kernelArgs[1];

    cudaGraphCreate(&graph, 0);

    cudaGraphConditionalHandle handle;
    cudaGraphConditionalHandleCreate(&handle, graph);

    // Use a kernel upstream of the conditional to set the handle value
    cudaGraphNodeParams kParams = { cudaGraphNodeTypeKernel };
    kParams.kernel.func         = (void *)setHandle;
    kParams.kernel.gridDim.x    = kParams.kernel.gridDim.y = kParams.kernel.gridDim.z = 1;
    kParams.kernel.blockDim.x   = kParams.kernel.blockDim.y = kParams.kernel.blockDim.z = 1;
    kParams.kernel.kernelParams = kernelArgs;
    kernelArgs[0] = &handle
    cudaGraphAddNode(&node, graph, NULL, 0, &kParams);

    cudaGraphNodeParams cParams = { cudaGraphNodeTypeConditional };
    cParams.conditional.handle = handle;
    cParams.conditional.type   = cudaGraphCondTypeIf;
    cParams.conditional.size   = 1;
    cudaGraphAddNode(&node, graph, &node, 1, &cParams);

    cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];

    // Populate the body of the conditional node
    cudaGraphNode_t bodyNodes[4];
    cudaGraphNodeParams params[4] = { ... }; // Setup kernel parameters as needed.
    cudaGraphAddNode(&bodyNodes[0], bodyGraph, NULL, 0, &params[0]);
    cudaGraphAddNode(&bodyNodes[1], bodyGraph, &bodyNodes[0], 1, &params[1]);
    cudaGraphAddNode(&bodyNodes[2], bodyGraph, &bodyNodes[0], 1, &params[2]);
    cudaGraphAddNode(&bodyNodes[3], bodyGraph, &bodyNodes[1], 2, &params[3]);

    return graph;
}

条件 WHILE 节点

只要条件为非零,WHILE 节点的体图就会重复执行。将在执行节点时以及每次完成体图后评估条件。下图描述了一个三节点图,其中中间节点 B 是包含三节点图的 WHILE 条件节点。

Diagram depicting a decision workflow using three circular nodes with a conditional node (b) set to loop over its body graph.
图 2:条件 WHILE 节点

为了了解如何创建此图,以下示例将句柄的默认值设置为非零值,以便在默认情况下执行 WHILE 循环。将默认值设置为非零,并将条件值保留在条件上游的内核中不修改,这有效地产生了一个 do-while 循环,其中条件体总是至少执行一次。创建 WHILE 循环,其中循环体仅在条件为 true 时执行,需要执行一些计算并在节点 a 中适当设置条件句柄。

在上一个示例中,条件主体由图 API 填充。在本例中,使用流捕获填充条件的主体。

第一步是定义一个内核,在每次执行条件体时设置条件值。在本例中,句柄是基于递减计数器的值设置的。

__global__ void loopKernel(cudaGraphConditionalHandle handle)
{
    static int count = 10;
    cudaGraphSetConditional(handle, --count ? 1 : 0);
}

接下来,定义一个函数来构建图。此函数分配条件句柄、创建节点并填充条件图的主体。为了清楚起见,省略了启动和执行图形的代码。

cudaGraph_t createGraph() {
    cudaGraph_t graph;
    cudaGraphNode_t nodes[3];

    cudaGraphCreate(&graph, 0);

    // Insert kernel node A
    cudaGraphNodeParams params = ...;

    cudaGraphAddNode(&nodes[0], graph, NULL, 0, &params);

    cudaGraphConditionalHandle handle;
    cudaGraphConditionalHandleCreate(&handle, graph, 1, cudaGraphCondAssignDefault);

    // Insert conditional node B
    cudaGraphNodeParams cParams = { cudaGraphNodeTypeConditional };
    cParams.conditional.handle = handle;
    cParams.conditional.type   = cudaGraphCondTypeWhile;
    cParams.conditional.size   = 1;
    cudaGraphAddNode(&nodes[1], graph, &nodes[0], 1, &cParams);

    cudaGraph_t bodyGraph = cParams.conditional.phGraph_out[0];
    cudaStream_t captureStream;
    cudaStreamCreate(&captureStream);
    
    // Fill out body graph with stream capture.
    cudaStreamBeginCaptureToGraph(captureStream,
                                  bodyGraph,
                                  nullptr,
                                  nullptr,
                                  0,
                                  cudaStreamCaptureModeRelaxed);
    myKernel1<<<..., captureStream>>>(...);
    myKernel2<<<..., captureStream>>>(...);
    loopKernel<<<1, 1, 0, captureStream>>>(handle);
    cudaStreamEndCapture(captureStream, nullptr);
    cudaStreamDestroy(captureStream);

    // Insert kernel node C.
    params = ...;
    cudaGraphAddNode(&nodes[2], graph, &nodes[1], 1, &params);

    return graph;
}

此示例使用 cudaStreamBeginCaptureToGraph,这是 CUDA 12.3 中新增加的一个 API,它使流捕获能够将节点插入到现有的图中。借助这个 API,可以将多个单独的捕获组合到一个图形对象中。此外,这个 API 还允许填充与条件节点一起创建的条件体图对象。

结论

CUDA Graphs 为静态工作流提供了难以置信的好处,在静态工作流中,图形创建的开销可以在多次连续启动中分摊。消除对图形的分割并将控制权返回给 CPU,以决定优先启动哪个,这有助于减少 CPU 开销和延迟。使用具有条件节点的 CUDA 图形可以有条件地或重复地执行图形的部分,而无需将控制权返回给 CPU。这释放了 CPU 资源,并使单个图形能够表示更复杂的工作流。

欲了解更多关于条件节点的信息,请参阅 CUDA 编程指南。要探索简单、完整的示例代码,请访问 NVIDIA/cuda-samples 在 GitHub 上。同时,您也可以加入 NVIDIA 开发者 CUDA 论坛 中的讨论。

 

Tags