尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

CUTLASS Python DSL 任务调度框架中的 Programmatic Dependent Launch(PDL)指南

CUTLASS Python DSL 任务调度框架中的 Programmatic Dependent Launch(PDL)指南 CUTLASS Python DSL 任务调度框架中的 Programmatic Dependent LaunchPDL指南【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlassProgrammatic Dependent LaunchPDL是 CUDA 中用于网格间流水线grid-to-grid pipelining的硬件机制它允许后继网格在前驱网格尚未完全结束时就开始启动。在 CUTLASS 的 Python DSLCuTeDSL中cutlass.experimental.task_schedulingTask Scheduling下文简称 TS框架通过两个同步专用资源PdlWaitBarrier与PdlLaunchBarrier把 PDL 接入到捕获式调度captured schedule与资源依赖图校验流程中。读完本文你将掌握如何在 TS 的schedule中正确声明 PDL 等待与发射、为什么必须显式声明依赖边、验证器与穷举检查器如何守住先 wait 后 launch的正确性边界以及assume_pdl_wait_completed的适用前提。背景为什么需要 PDL 等待PDL 的核心特性是后继网格可以在前驱网格完全完成之前启动。这意味着前驱结果已就绪不再由网格边界隐式保证而必须由后继网格内的线程显式执行griddepcontrol.wait指令来确认。一旦缺失后继网格中读取前驱网格产出数据的任何内存访问都可能读到未完成的数据——因此文档把 PDL wait 称为正确性边界correctness boundary所有读取前驱网格结果的后继网格工作必须在第一次依赖内存访问之前执行wait_griddep。在 TS 框架中这一边界由两个只做同步、不携带任何数据负载的资源表达PdlWaitBarrier.wait_griddep()发射 PDL 等待wait 侧PdlLaunchBarrier.launch_griddep()发射依赖网格启动通知launch 侧。它们不携带 payload 变量但仍然参与任务资源列表src_resources/dst_resources与依赖图排序从而让框架的验证器能够检查其相对顺序。PDL 排序问题依赖必须显式声明TS无法从地址表达式或资源名推断出某块全局内存是由前驱网格写入的。当gmem_ab读取由前驱网格产生的数据时你必须把这一事实声明进资源依赖图resource_dependency_graph { gmem_ab: [pdl_wait_resource], }声明了该边之后TS 把PdlWaitBarrier视为一个 ordering-only仅排序的上游资源并检查 wait 被排在任何依赖工作dependent work之前。未声明该边时TS 眼中的gmem_ab与pdl_wait_resource是两个相互独立的资源没有任何声明的排序关系可供校验——即使你把wait_griddep()写进了 schedule框架也不会去验证它的位置。这一设计在 ts_validation.rst 的 PDL Ordering 一节中也有呼应PDL wait resources are ordering-only resources, but the dependency still has to be explicit.代码走读一个完整的 PDL 加载任务下面是在依赖图中携带 PDL wait 边之后一个典型的 PDL 感知 schedule。它把 wait 放在 schedule 最前面、把 launch 放在所有需要解锁后继网格的工作之后schedule def load_a_schedule( gmem_ab: GmemAbResource, smem_a: SmemAbResource, wq: WorkQueue, pdl_wait_resource: PdlWaitBarrier, pdl_launch_resource: PdlLaunchBarrier, ) - None: pdl_wait_resource.wait_griddep() with work_tile_loop(wq): with domain_loop(0, num_k_tiles, 1): coord_k, coord_m, coord_n gmem_ab.compute_coords() smem_a.try_acquire() smem_a.acquire() smem_a.tma_load_a(coord_kcoord_k, coord_mcoord_m) smem_a.commit() wq.try_wait() wq.wait() wq.get_and_advance_work_tile() wq.release() pdl_launch_resource.launch_griddep()要点wait_griddep()在work_tile_loop之外、schedule 开头执行先于任何 PDL 依赖的内存访问launch_griddep()放在持久化工作循环之后保证通知后继网格时本网格该解锁的工作已经完成资源声明与调度阶段consumer_work/producer_work对齐PdlWaitBarrier走的是 wait 侧PdlLaunchBarrier走的是 launch 侧详见下一节源码剖析。排序规则TS 对 PDL 资源的排序规则可以归纳为三条PDL wait 资源被列为上游依赖凡是读取前驱网格结果的数据资源都要把PdlWaitBarrier列进resource_dependency_graph中该资源的 upstream 列表Launch barrier 是 schedule-only 通知资源PdlLaunchBarrier不携带数据依赖不属于任何数据流的下游因此也不应作为依赖图的 destination 出现检查器拒绝非法 launch 路径当一条被使能enabled的launch_griddep路径上没有执行过任何wait_griddep时检查器直接拒绝——这个判定包含跳过 tile 的执行路径skipped-tile execution也就是说把 wait 藏在可能被跳过的分支里并不能绕过检查。源码剖析两类 barrier 的 PTX 发射与 ordering-only 语义TS 对 PDL 的建模集中在 resources.pyPdlWaitBarrierL3262-L3302dataclass(kw_onlyTrue) class PdlWaitBarrier(MemoryResource): is_barrier: cutlass.Constexpr[bool] True consumer_work cute.jit def wait_griddep(self, stage_info: StageInfo) - None: prims.griddepcontrol(kindprims.GridDepAction.WAIT)is_barrierTrue且没有PipelineConfig是纯 barrier 资源唯一的对外方法wait_griddep通过prims.griddepcontrol(kindGridDepAction.WAIT)发射griddepcontrol.waitPTX 指令阻塞发起线程直到直接前驱网格依赖完成并把其全局内存结果置为可见由于is_barrierTrue验证器把依赖图中指向它的边解释为ordering-only不会建立任何 consumer/producer 变量拷贝计划但会要求发射 wait 的 schedule 条目先于产生依赖资源的任务的 producer 条目。PdlLaunchBarrierL3306-L3338dataclass(kw_onlyTrue) class PdlLaunchBarrier(MemoryResource): is_barrier: cutlass.Constexpr[bool] True producer_work cute.jit def launch_griddep(self, stage_info: StageInfo) - None: prims.griddepcontrol(kindprims.GridDepAction.LAUNCH_DEPENDENTS)同样是无PipelineConfig的 barrier 资源通过griddepcontrol(kindGridDepAction.LAUNCH_DEPENDENTS)发射griddepcontrol.launch_dependentsPTX 指令通知后继网格可以开始启动 CTA它不承载数据依赖因此不作为resource_dependency_graph的 destination出现验证器把launch_griddep条目视为 schedule-only 发射可出现在任意任务与任意阶段注意框架不强制至少一次 launch——当本网格是流水线链的最后一环或 launch 由非 PDL 的主机策略门控时省略 launch 是合法的。两类资源还被统一收录进 PDL_BARRIER_TYPES 元组供框架代码用isinstance统一识别 PDL 风格 barrier。从源码注释还可以得到两个重要的工程细节wait_griddep条目可以放在任意阶段Head、带LoopFirstIter/LoopLastIterguard 的 Loop、Tail、以及 work-tile-loop 之后post-WTL官方鼓励 inline wait 模式每个产生 PDL 依赖资源的任务各自发射自己的pdl_wait.wait_griddep()调用。验证器如何检查wait 先于依赖工作TaskManager在代码生成前会对捕获的 schedule 运行完整校验见 task_manager.py 与 ts_validation.rst。其中与 PDL 直接相关的逻辑是_verify_pdl_wait_orderingtask_manager.py遍历resource_dependency_graph收集所有上游为PdlWaitBarrier的(pdl_wait, dependent)依赖对对每个依赖对检查会产生依赖资源的任务dependent在dst_resources中或dependent是非流水线资源且出现在src_resources中例如直接在ConsumerWork里读全局内存的 GMEM 源被检查的依赖工作阶段集合为ConsumerAuxWork / ConsumerWork / ProducerAuxWork / ProducerWork在任务的合并 schedulehead loop body tail中找出所有wait_griddepConsumerWork阶段的PdlWaitBarrier位置与所有依赖工作位置若存在依赖工作但没有先行的 wait或 wait 位置不早于依赖工作校验即失败若assume_pdl_wait_completedTrue该函数直接返回、跳过校验见下节。值得说明的是源码中的一条注释下游是流水线化的 consumer例如 MMA 读取已被其 producer commit 的 SmemA时TS 自身的流水线机制已经完成了排序因此这类组合不会进入额外检查。穷举检查器包括 skipped-tile 在内的全路径校验仅靠单任务线性 schedule 校验是不够的因为 PDL launch 与 wait 可能分布在不同任务、不同阶段。TS 为此实现了穷举交错检查器exhaustive_checker.py它模拟所有任务步骤的合法交错用于捕获死锁所有任务都在等待一个无人能发出的事件物理 SMEM/TMEM 别名竞态重叠的分配在合法流水线窗口之外被读写PDL launch 顺序某条 launch 路径可以在任何 wait 之前执行上述规则在 skipped-tile 变体下的违规。其实现要点_is_pdl_wait_entryL343识别PdlWaitBarrier的ConsumerWork条目_has_executed_pdl_waitL792-L798按当前游标状态判断是否已有 wait 执行过当assume_pdl_wait_completedTrue时直接返回 True在每步推进前若当前条目是 launch 条目_is_pdl_launch_entry且has_executed_pdl_wait为假则记录一条PdlOrderInfo违例错误信息为 PDL ORDER: launch_griddep can execute before any wait_griddepL1058-L1087。由于穷举器连 skipped-tile 路径也会模拟一个被skip_if保护的 wait 无法为无条件 launch 提供正当性——这正是原文档强调包括 skipped-tile execution的原因。实战示例3-cluster GEMM 中的 PDL 用法仓库中的进阶教程 01_fp16_bf16_gemm_3_cluster.py 展示了 PDL 在真实多集群 GEMM 中的完整接线方式资源创建L1386-L1387pdl_wait PdlWaitBarrier(namePdlWait) pdl_launch PdlLaunchBarrier(namePdlLaunch)A 加载任务L1034-L1088把 wait 放在最前、launch 放在持久化工作循环之后schedule def load_a_schedule( gmem_ab: GmemAbResource, smem_a: SmemAbResource, pdl_wait_resource: PdlWaitBarrier, pdl_launch_resource: PdlLaunchBarrier, wq: WorkQueue, ) - None: # PDL wait gates the A-load stream before any persistent work is issued. pdl_wait_resource.wait_griddep() smem_a.init_load_state() with work_tile_loop(wq): gmem_ab.init_tile_coords() with domain_loop(0, num_k_tiles, 1): coord_k, coord_m, coord_n gmem_ab.compute_coords() smem_a.try_acquire() smem_a.acquire() smem_a.tma_load_a(coord_kcoord_k, coord_mcoord_m) smem_a.commit() # TAIL: advance to next work tile wq.try_wait() wq.wait() wq.get_and_advance_work_tile() wq.release() # PDL launch is emitted only after all persistent A-load work completes. pdl_launch_resource.launch_griddep()对应的Task声明把两个 PDL 资源放到正确的角色上pdl_wait进入src_resourceswait 是消费侧语义pdl_launch进入dst_resourceslaunch 是生产侧语义return Task( src_resources[gmem_ab_resource, pdl_wait, work_queue], dst_resources[smem_a_resource, pdl_launch], warp_idx4, num_warps1, scheduleresult, num_registers40, nameLoadATask, debug_printdebug_print, )何时使用 assume_pdl_wait_completedTaskManager的构造参数assume_pdl_wait_completed: bool Falsetask_manager.py用于声明PDL wait 已在 TS 之外完成。原文档给出的使用条件非常严格UseTaskManager(assume_pdl_wait_completedTrue)only when the PDL wait is emitted outside TS before the first PDL-dependent memory access.即只有当你确定 PDL wait 已经由 TS 之外的代码、且在第一次 PDL 依赖内存访问之前发射时才可以打开该开关置为True后_verify_pdl_wait_orderingtask_manager.py直接返回跳过 wait 前置性校验穷举检查器的_has_executed_pdl_waitexhaustive_checker.py也视为 wait 恒已执行不再报告 launch-before-wait 违例。滥用该开关等于关闭 PDL 排序护栏属于明确的自担风险行为。最佳实践小结依赖边必须显式任何读取前驱网格结果的资源都必须在resource_dependency_graph中把PdlWaitBarrier列为其 upstreaminline wait 优先每个产生 PDL 依赖资源的任务自行发射wait_griddep()位置在 head 或任何先于依赖内存访问的阶段launch 放最后launch_griddep()只应出现在所有需要解锁后继网格的工作完成之后且不应出现在依赖图 destination 中不要用 skip 绕过检查穷举检查器会模拟 skipped-tile 路径藏在可跳过分支里的 wait 不能为无条件 launch 背书谨慎使用assume_pdl_wait_completedTrue仅当 wait 确实由 TS 外部代码在第一次 PDL 依赖访问前发射时才启用。如需进一步了解 TS 的整体校验体系、流水线配置与模式可继续阅读同目录下的 ts_validation.rst并在 resources.py、task_manager.py 与 exhaustive_checker.py 中对照源码验证。【免费下载链接】cutlassCUDA Templates and Python DSLs for High-Performance Linear Algebra项目地址: https://gitcode.com/GitHub_Trending/cu/cutlass创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表