EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

use default for torch.ops (#4835)

Browse Source
This commit is contained in:
Yineng Zhang
2025-03-27 19:09:58 -07:00
committed by GitHub
parent 10a9ab7b07
commit 31dfff7da7
7 changed files with 51 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
sgl-kernel/python/sgl_kernel/allreduce.py
View File

@@ -12,49 +12,49 @@ if torch.version.hip is not None:
rank: int,
full_nvlink: bool,
) -> int:
return torch.ops.sgl_kernel.init_custom_ar(
return torch.ops.sgl_kernel.init_custom_ar.default(
meta, rank_data, handles, offsets, rank, full_nvlink
)
def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
torch.ops.sgl_kernel.all_reduce_reg(fa, inp, out)
torch.ops.sgl_kernel.all_reduce_reg.default(fa, inp, out)
def all_reduce_unreg(
fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
) -> None:
torch.ops.sgl_kernel.all_reduce_unreg(fa, inp, reg_buffer, out)
torch.ops.sgl_kernel.all_reduce_unreg.default(fa, inp, reg_buffer, out)
def dispose(fa: int) -> None:
torch.ops.sgl_kernel.dispose(fa)
torch.ops.sgl_kernel.dispose.default(fa)
def meta_size() -> int:
return torch.ops.sgl_kernel.meta_size()
return torch.ops.sgl_kernel.meta_size.default()
def register_buffer(
fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
) -> None:
return torch.ops.sgl_kernel.register_buffer(fa, t, handles, offsets)
return torch.ops.sgl_kernel.register_buffer.default(fa, t, handles, offsets)
def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]:
return torch.ops.sgl_kernel.get_graph_buffer_ipc_meta(fa)
return torch.ops.sgl_kernel.get_graph_buffer_ipc_meta.default(fa)
def register_graph_buffers(
fa: int, handles: List[str], offsets: List[List[int]]
) -> None:
torch.ops.sgl_kernel.register_graph_buffers(fa, handles, offsets)
torch.ops.sgl_kernel.register_graph_buffers.default(fa, handles, offsets)
def allocate_meta_buffer(size: int) -> torch.Tensor:
return torch.ops.sgl_kernel.allocate_meta_buffer(size)
return torch.ops.sgl_kernel.allocate_meta_buffer.default(size)
def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
return torch.ops.sgl_kernel.get_meta_buffer_ipc_handle(inp)
return torch.ops.sgl_kernel.get_meta_buffer_ipc_handle.default(inp)
else:
# TRTLLM custom allreduce
def init_custom_reduce(
rank_id, num_devices, rank_data, buffers, tmp_buffers, barrier_in, barrier_out
):
return torch.ops.sgl_kernel.init_custom_ar(
return torch.ops.sgl_kernel.init_custom_ar.default(
rank_id,
num_devices,
rank_data,
@@ -65,13 +65,13 @@ else:
)
def custom_dispose(fa):
torch.ops.sgl_kernel.dispose(fa)
torch.ops.sgl_kernel.dispose.default(fa)
def custom_reduce(fa, inp, out):
torch.ops.sgl_kernel.all_reduce(fa, inp, out)
torch.ops.sgl_kernel.all_reduce.default(fa, inp, out)
def get_graph_buffer_ipc_meta(fa):
return torch.ops.sgl_kernel.get_graph_buffer_ipc_meta(fa)
return torch.ops.sgl_kernel.get_graph_buffer_ipc_meta.default(fa)
def register_graph_buffers(fa, handles, offsets):
torch.ops.sgl_kernel.register_graph_buffers(fa, handles, offsets)
torch.ops.sgl_kernel.register_graph_buffers.default(fa, handles, offsets)