[amd] Add deterministic all-reduce kernel for AMD (ROCm) (#15340)

Co-authored-by: Thomas Wang <1am9trash@gmail.com>
This commit is contained in:
sunxxuns
2025-12-19 02:36:03 -05:00
committed by GitHub
parent 0e869f0868
commit f2d64e6782
12 changed files with 1466 additions and 8 deletions

View File

@@ -24,6 +24,18 @@ if torch.version.hip is not None:
) -> None:
torch.ops.sgl_kernel.all_reduce_unreg.default(fa, inp, reg_buffer, out)
def deterministic_all_reduce_reg(
fa: int, inp: torch.Tensor, out: torch.Tensor
) -> None:
torch.ops.sgl_kernel.deterministic_all_reduce_reg.default(fa, inp, out)
def deterministic_all_reduce_unreg(
fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
) -> None:
torch.ops.sgl_kernel.deterministic_all_reduce_unreg.default(
fa, inp, reg_buffer, out
)
def dispose(fa: int) -> None:
torch.ops.sgl_kernel.dispose.default(fa)