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

Rename files in sgl kernel to avoid nested folder structure (#4213)

Browse Source 
Co-authored-by: zhyncs <me@zhyncs.com>
This commit is contained in:
Lianmin Zheng
2025-03-08 22:54:51 -08:00
committed by GitHub
parent ee132a4515
commit 8abf74e3c9
47 changed files with 184 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
python/sglang/srt/_custom_ops.py
View File

@@ -75,42 +75,42 @@ else:
rank: int,
full_nvlink: bool,
) -> int:
return sgl_kernel.ops.allreduce.init_custom_ar(
return sgl_kernel.allreduce.init_custom_ar(
meta, rank_data, handles, offsets, rank, full_nvlink
)
def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
sgl_kernel.ops.allreduce.all_reduce_reg(fa, inp, out)
sgl_kernel.allreduce.all_reduce_reg(fa, inp, out)
def all_reduce_unreg(
fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor
) -> None:
sgl_kernel.ops.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out)
sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out)
def dispose(fa: int) -> None:
sgl_kernel.ops.allreduce.dispose(fa)
sgl_kernel.allreduce.dispose(fa)
def meta_size() -> int:
return sgl_kernel.ops.allreduce.meta_size()
return sgl_kernel.allreduce.meta_size()
def register_buffer(
fa: int, t: torch.Tensor, handles: List[str], offsets: List[int]
) -> None:
return sgl_kernel.ops.allreduce.register_buffer(fa, t, handles, offsets)
return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets)
def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]:
return sgl_kernel.ops.allreduce.get_graph_buffer_ipc_meta(fa)
return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa)
def register_graph_buffers(
fa: int, handles: List[str], offsets: List[List[int]]
) -> None:
sgl_kernel.ops.allreduce.register_graph_buffers(fa, handles, offsets)
sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets)
def allocate_meta_buffer(size: int) -> torch.Tensor:
return sgl_kernel.ops.allreduce.allocate_meta_buffer(size)
return sgl_kernel.allreduce.allocate_meta_buffer(size)
def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor:
return sgl_kernel.ops.allreduce.get_meta_buffer_ipc_handle(inp)
return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp)
else:
# TRTLLM custom allreduce
@@ -123,7 +123,7 @@ else:
barrier_in: List[int],
barrier_out: List[int],
) -> int:
return sgl_kernel.ops.init_custom_reduce(
return sgl_kernel.init_custom_reduce(
rank_id,
world_size,
rank_data_base,
@@ -134,15 +134,15 @@ else:
)
def all_reduce(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None:
sgl_kernel.ops.custom_reduce(fa, inp, out)
sgl_kernel.custom_reduce(fa, inp, out)
def dispose(fa: int) -> None:
sgl_kernel.ops.custom_dispose(fa)
sgl_kernel.custom_dispose(fa)
def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
return sgl_kernel.ops.get_graph_buffer_ipc_meta(fa)
return sgl_kernel.get_graph_buffer_ipc_meta(fa)
def register_graph_buffers(
fa: int, handles: List[List[int]], offsets: List[List[int]]
) -> None:
sgl_kernel.ops.register_graph_buffers(fa, handles, offsets)
sgl_kernel.register_graph_buffers(fa, handles, offsets)