Skip to content

[Feature] launch wrapper支持传入dynamic ub值 #494

Description

@and0d0

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or request

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions