[TLE]barrier & wgmma pipeline with user promising support in hopper#707
[TLE]barrier & wgmma pipeline with user promising support in hopper#707Kafka-Hatsune wants to merge 15 commits into
Conversation
|
|
i3wanna2
changed the title
|
Please check which newly code should be inside |
fixed in commit "Guard TLE-only barrier compiler changes". For Hopper CI fails，it seems to work fine in H800 server. What's different between my flagtree build/test command which is "pip install -e . --no-build-isolation -v"/"python3 -m pytest -s tests/test_flash_mla_sparse_fwd.py" and ci servers building and testing? I'm sure using the flagtree built triton. 
For other CI fails，I add skipif for my new test. please run a new round ci and I will get more information. |
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Kafka-Hatsune
force-pushed
the
tle-wgmma-user-promise
branch
from
b475f86 to
abc69cd
Compare
sunnycase
left a comment
sunnycase
left a comment
There was a problem hiding this comment.
Thanks for the work on the Hopper pipeline support. One concern I have is that this introduces a separate async_task/async_tasks API for warp specialization while tle/gluon already has an existing warp_specialize API.
Could we please reuse or extend the existing warp_specialize API here instead? Having another async-task abstraction for the same concept would split the user-facing API surface and make it unclear which API users should choose. Aligning this with the existing API would keep the programming model more consistent and easier to understand.
4 participants
FlagTree 前端支持
通过使用tle.barrier，tle.wgmma，tle.wgmma_wait等新增API，在FlagTree Triton中实现Hopper架构上的自定义复杂流水排布。具体样例参考commit tle_hopper_fa_ws_pipelined_pingpong_persistent.py实现hopper版sdpa样例。
API修改总览：
新增 warp specialization 定义 API：
tle.gpu.async_tasks()tle.gpu.async_task(...)tle.gpu.async_task_replica_id()新增 barrier API：
tle.gpu.alloc_barriers(...)tle.gpu.alloc_barrier(...)tle.gpu.barrier_wait(bar, phaseIdx=None)tle.gpu.barrier_arrive(bar, arrive_count=1, phaseIdx=None)tle.gpu.PENDINGtle.gpu.READYtle.gpu.barriertle.gpu.barrier_type新增 WGMMA API：
tle.gpu.wgmma(...)tle.gpu.wgmma_wait(pendings, acc=None)已有 API 扩展：
tle.gpu.copy(...)新增barrier参数，用于 TMA global-to-shared load 的显式 completion barrier。Warp Specialization 定义 API
API 参数简介
tle.gpu.async_tasks()参数：无。
含义：producer/consumer 任务容器，会由前端 lowering 到
ttg.warp_specialize。容器内必须有且只有一个 producerasync_task，其余任务是 consumer。tle.gpu.async_task(...)参数：
role：任务角色。支持"producer"、"consumer"；"default"兼容映射成"producer"。不写时默认为"consumer"。name：任务名，仅作为可读标识使用；当前实现不会用它决定 partition 角色或资源分配。producer/consumer 角色由第一个位置参数role决定，例如async_task("producer", name="load")；只写name="producer"不会把任务变成 producer。不写时默认使用 role 作为 name。num_warps：consumer partition 使用的 warp 数。num_warps=4表示该 consumer partition 使用 4 个 warp，也就是 1 个 warp group。async_task个数和replicate决定，不由num_warps决定。num_warps。producer 对应ttg.warp_specialize的 default/base region。registers|num_regs：consumer partition 请求的每线程寄存器数，二者等价。当前 Hopper 后端按最多 256 regs/thread 处理，并会把 consumer 请求值向上对齐到 8 的倍数；实际setmaxnreg分配还会结合 kernel launchnum_warps、consumer warp 数和总寄存器预算计算。建议传 8 的倍数且不超过 256。若任一 consumer 设置该参数，所有 consumer 都必须设置。replicate：consumer partition 副本总数。默认 1，表示有一个consumer，当前实现要求replicate >= 1。tle.gpu.async_task_replica_id()参数：无。
返回值：当前 replicated consumer 的编译期 replica id，类型是
tl.constexpr。只能在async_task区域内调用。使用示例
一个 Producer + 一个 Consumer
当前实现中，kernel launch 的
num_warps是 producer/default region 的 warp 数。consumer 的num_warps是额外 worker partition 的 warp 数。上面例子对应：
一个 Producer + 两个 Replicated Consumers
replicate=2表示同一段 consumer body 生成两个 consumer partition。每个 replica 都使用num_warps=4和registers=168。对应：
replicate=0不是“一个 consumer”的写法，会被前端拒绝。一个 consumer 应省略replicate或写replicate=1。Barrier 相关 API
API 参数简介
本节只说明参数含义。
arrive_count和phaseIdx的具体配置方式见后面的使用示例。tle.gpu.alloc_barriers(...)参数：
num_barriers：tle.gpu.alloc_barriers返回一个barrier数组，此处为分配的 barrier slot 数量，必须是编译期正整数。arrive_count：barrier 的计数参数，即 pending arrival count。必须是编译期正整数。mbarrier 和 named barrier 后端对它的解释不同。具体参考使用示例。init：barrier初始状态，支持tle.gpu.PENDING和tle.gpu.READY。expect_bytes：TMA completion barrier 期望完成的字节数。为None或编译期正整数。返回值：
tle.gpu.barrier，表示 barrier 数组或 barrier slot。数组通过bars[i]取 slot。tle.gpu.alloc_barrier(...)返回值：单个
tle.gpu.barrierslot。它是alloc_barriers(num_barriers=1, ...)的语法糖。tle.gpu.barrier_wait(...)参数：
bar：要等待的 barrier slot，类型为tle.gpu.barrier。phaseIdx：mbarrier logical use id。传入时使用 mbarrier 路径；为None时使用 named barrier 路径。tle.gpu.barrier_arrive(...)参数：
bar：要 arrive 的 barrier slot，类型为tle.gpu.barrier。arrive_count：本次 arrive 的计数权重。arrive_count表达。phaseIdx：mbarrier logical use id，第几次使用这个 slot。传入时使用 mbarrier 路径；为None时使用 named barrier 路径。tle.gpu.PENDING含义：barrier 初始不可通过。常用于 full barrier、TMA completion barrier、named barrier。
tle.gpu.READY含义：barrier 初始可通过。常用于 empty barrier。
READYbarrier 只能走 mbarrier 路径。tle.gpu.barrier含义：
alloc_barrier(s)返回的 barrier value 类型。支持bars[i]取数组 slot。tle.gpu.barrier_type含义：
tle.gpu.barrier对应的类型描述。普通 kernel 代码通常不需要直接构造。修改
tle.gpu.copy(..., barrier=...)参数：
src：copy 源。使用barrier时必须是 TMA TensorDescriptor。dst：copy 目标。使用barrier时必须是 shared-memorytle.gpu.buffered_tensor。shape：copy tile 的形状。offsets：TMA descriptor 坐标偏移。barrier：可选 completion barrier。必须来自tle.gpu.alloc_barrier(s)(expect_bytes=...)。使用示例
mbarrier 使用路径
传入
phaseIdx时，barrier_wait/barrier_arrive会选择 mbarrier 路径。tle.gpu.copy(..., barrier=...)的 TMA completion barrier 也必须走 mbarrier 路径，并且后续 wait 必须传phaseIdx。适用场景：
producer/consumer Empty-Full 同步。
TMA global-to-shared load completion。
init=tle.gpu.READY的 empty barrier。需要用 logical use id 表达同一 slot 多次复用的 ping-pong / 多 stage pipeline。
phaseIdx要求：phaseIdx是同一个 barrier slot 的 logical use id，也就是这个 slot 第几次被使用。它的设计意图是让循环里的同一组 barrier slot 可以被反复复用，并通过循环计数区分第几轮复用。(phaseIdx & 1) ^ init_polarityphaseIdx & 1，再根据init调整初始 polarity& 1。phaseIdx。slot = iter % NUM_BUFFERS，phaseIdx = iter // NUM_BUFFERS，因为同一个 slot 每隔NUM_BUFFERS次迭代才复用一次。arrive_count配置示例：一个 consumer 到达同一个 slot 一次：
两个 replicated consumers 共同释放同一个 slot，此时同一个 slot 会收到两次 arrive，所以
arrive_count=2：两个 replicated consumers 使用各自独立 slot：
arrive_count=1是正确的，因为 lowering 会把同一个arrive_count分别用于每个 slot 的初始化。arrive_count=2，则每个 slot 都会等待两次 arrive，下面的代码会少 arrive 一次。其他注意事项：
后端选择：
barrier_wait/barrier_arrive只要传入phaseIdx，当前 slot 就会记录为 mbarrier backend。初始状态：
init=tle.gpu.READY只适合 mbarrier 路径。READYbarrier 不传phaseIdx，会选择 named barrier 路径并报错。TMA completion：
expect_bytesbarrier 只能用于 TMA global-to-shared completion。expect_bytes的 barrier 必须通过tle.gpu.copy(..., barrier=...)触发 completion，并且必须用phaseIdx等待。named barrier 使用路径
不传
phaseIdx时，barrier_wait/barrier_arrive会选择 named barrier 路径。适用场景：
不需要 TMA completion 的 CTA-scope rendezvous。
consumer 内部或多个 consumer warp group 之间的轻量同步。
同步点可以用静态 barrier slot 表达，且不需要
init=READY。named_barrier的arrive_count是参与线程数named_num_threads，通常写成参与 warp 数 * 32。当前实现会把 allocation 的arrive_count同时传给 named arrive 和 named wait。因此，使用named_barrier只需要alloc后直接wait+arrive。举例：两个 4-warp consumer replicas 通过同一个 named barrier 同步，总参与线程数是 256。这里一个 replica 先 non-blocking arrive，另一个 replica 执行 wait/sync；不需要在
barrier_arrive/barrier_wait调用参数里写arrive_count。其他注意事项：
named barrier只支持init=tle.gpu.PENDING，不支持READY。named barrier不支持expect_bytes，因此不能用于tle.gpu.copy(..., barrier=...)的 TMA completion。named barrier要求静态 barrier slot index，例如sync[0]；不能用动态 slot index。Empty-Full 同步示例
copy(..., barrier=...)只支持 TMA global-to-shared load。不支持 shared-to-global TMA store、普通 copy、named barrier 或没有expect_bytes的 barrier。TMA completion 示例
TMA completion 只能使用 mbarrier 后端，通常写
arrive_count=1，因为 arrival 来自这一条 TMA copy 的 completion：WGMMA 相关 API
API 参数简介
tle.gpu.wgmma(...)参数：
a：WGMMA A operand。可以是 shared-memorytle.gpu.buffered_tensor或 registertl.tensor。b：WGMMA B operand。当前必须是 shared-memorytle.gpu.buffered_tensor。acc：输入 accumulator。为None时前端会生成零 accumulator。input_precision：dot input precision。为None时使用当前 builder 默认值。max_num_imprecise_acc：FP8 imprecise accumulation 控制参数，不改变返回 tensor dtype，仅在FP8 GEMM上有效。使用 FP8 -> F32 WGMMA 当前有效值需 >= 32；None 在 sm90 上默认 2^30。值越小越偏精度，越大越偏性能。out_dtype：输出 accumulator dtype，默认tl.float32。它决定结果/accumulator 的 dtype；trans_a：是否转置 A operand，必须是编译期 bool。trans_b：是否转置 B operand，必须是编译期 bool。返回值：异步 WGMMA accumulator dependency value。普通 tensor op、reduce、store 消费前必须调用
tle.gpu.wgmma_wait(...)。tle.gpu.wgmma_wait(...)参数：
pendings：允许保留的 outstanding WGMMA group 数量。0表示完全等待。acc：要等待的 accumulator。若省略且第一个参数是 tensor，则等价于tle.gpu.wgmma_wait(0, acc)。返回值：wait 后可安全消费的 accumulator tensor。
使用示例
单次 WGMMA：
流水中保留一个 pending group，WGMMA是warpgroup级别的FIFO队列模型：
pipeline lower说明
由于用户可以通过当前api声明barrier和wait来规划流水，与之前TLE的自动分析流水插入wait的lower pipeline相冲突。因此当前选择：如果用户在triton tle中以下api:
tle.gpu.alloc_barrierstle.gpu.alloc_barriertle.gpu.wgmma_wait则会避开原本的
mlir::triton::asyncLaunchDots pass（在内部由TLE劫持为自己的异步dot wait分析pass），使用mlir::triton::gpu::detail::scheduleTleWgmmaUserPromisePipeline，行为为wgmma补齐ttng::WarpGroupDotCommitOp。当前性能
测试说明
vllm.vllm_flash_attn.flash_attn_interface.flash_attn_varlen_functorch.nn.attention.varlen.varlen_attn通过FlagGems benchmark测试，右侧三列表示以vllm为100%性能，其他三个的性能比。