Summary
PTODSL对应生成的 C++ launch wrapper 中,无法传递dynamic UB值
Motivation / use case
当前 PTODSL 的 native runtime launch path 只支持:
compiled[grid, stream](...)
对应生成的 C++ launch wrapper 中,kernel launch 形式为:
kernel<<<grid, nullptr, stream>>>(...)
也就是说,第二个 launch 参数固定为 nullptr,没有用户可配置的动态 UB/shared memory size 参数。
这会影响一些需要显式动态 UB scratch 的 kernel。例如 RMSNorm alloc_buffer/SIMT 示例中,kernel 内部通过:
ub_base = pto.castptr(pto.const(0, dtype=pto.ui64), pto.ptr(pto.f32, "ub"))
w_ub = pto.addptr(ub_base, 0)
reduce_scratch = pto.addptr(ub_base, hidden_size)
x_ub = pto.addptr(ub_base, hidden_size + 128)
y_ub = pto.addptr(ub_base, hidden_size + 128 + 2 * hidden_size)
rstd_ub = pto.addptr(ub_base, hidden_size + 128 + 4 * hidden_size)
手写了一块 UB scratch layout。该 kernel launch 时需要传入足够的动态 UB 字节数,例如:
kernel<<<grid, 82496, stream>>>(...)
由于默认 wrapper 不支持这个参数,目前只能通过 rmsnorm_alloc_buffer_simt_manual_launch.py 手写 launch.cpp 绕过默认 launch path。
参考:
关于codegen的修改 and0d0@dd55c2e
Proposed API / behavior
No response
Alternatives considered
No response
Additional context
No response
Summary
PTODSL对应生成的 C++ launch wrapper 中,无法传递dynamic UB值
Motivation / use case
当前 PTODSL 的 native runtime launch path 只支持:
compiled[grid, stream](...)对应生成的 C++ launch wrapper 中,kernel launch 形式为:
kernel<<<grid, nullptr, stream>>>(...)也就是说,第二个 launch 参数固定为 nullptr,没有用户可配置的动态 UB/shared memory size 参数。
这会影响一些需要显式动态 UB scratch 的 kernel。例如 RMSNorm alloc_buffer/SIMT 示例中,kernel 内部通过:
手写了一块 UB scratch layout。该 kernel launch 时需要传入足够的动态 UB 字节数,例如:
kernel<<<grid, 82496, stream>>>(...)
由于默认 wrapper 不支持这个参数,目前只能通过 rmsnorm_alloc_buffer_simt_manual_launch.py 手写 launch.cpp 绕过默认 launch path。
参考:
关于codegen的修改 and0d0@dd55c2e
Proposed API / behavior
No response
Alternatives considered
No response
Additional context
No response