[Rebase] 0611: Rebase upstream TileLang#60
Conversation
Compile-once API and tests: - Add compile_once support to tilelang.compile and tilelang.jit with optional compile_group and compile_root controls. - Coordinate root-first compilation through torch.distributed and propagate root compile failures to all ranks instead of hanging peers. - Add distributed compile_once coverage for ordering, nonzero roots, real compile/jit paths, and root failure propagation. Distributed examples and tests: - Enable compile_once in distributed primitive tests and intranode distributed examples so spawned ranks reuse the normal cache path after root compilation. - Pass explicit compile_group for direct tilelang.compile calls inside distributed workers; use decorator-level compile_once for JIT kernels. - Keep import-time warmup compiles and non-distributed single-GPU tests unchanged; remove cache disabling from examples where it would prevent reuse.
Implement src_pe/dst_pe support for T.copy/T.tma_copy, including 1D TMA address remapping, SIMT cp_block fast path, predicated edge-tile fallback, and explicit rejection for unsupported descriptor TMA remote copies. Factor symmetric remote address remapping into src/op/distributed_utils.h and reuse it from copy, remote copy, and sync lowering. Add distributed tests for remote copy, remote TMA load/store, SIMT fallback, descriptor rejection, and compile-once usage.
Lower bf16/fp16 multimem ops directly to packed x2 instructions instead of relying on ParallelOp vector inference. Add standalone multimem allreduce strategies and an SM-specialized GEMM allreduce example with tests. Align GEMM-AR reset benchmarking with TK by timing device epilogue launches without host barriers.
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 Hi! Thank you for contributing to the TileLang project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
Rachmanino
force-pushed
the
rebase/v0-dev-tilelang-550e25d
branch
2 times, most recently
from
a345e41 to
18a30df
Compare
1 participant
No description provided.