minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

optimize conv3d used in patch embedding (#16040)

Browse Source
This commit is contained in:
Ma Mingfei
2026-03-18 22:17:53 -07:00
committed by GitHub
parent cbea9f6909
commit 62d7454976
2 changed files with 246 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+11
View File
@@ -294,6 +294,11 @@ at::Tensor causal_conv1d_update_cpu(
int64_t pad_slot_id,
bool is_vnni);
// conv3d fast path for patch embedding
at::Tensor conv3d_embed_weight_pack(const at::Tensor& weight);
at::Tensor conv3d_embed_cpu(const at::Tensor& input, const at::Tensor& weight, const at::Tensor& bias, bool is_vnni);
// shared memory init
void initialize(int64_t size, int64_t rank);
@@ -525,6 +530,12 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"Tensor? cache_seqlens, Tensor? conv_state_indices, int pad_slot_id, bool is_vnni) -> Tensor");
m.impl("causal_conv1d_update_cpu", torch::kCPU, &causal_conv1d_update_cpu);
// conv3d fast path for patch embedding
m.def("conv3d_embed_weight_pack(Tensor weight) -> Tensor");
m.impl("conv3d_embed_weight_pack", torch::kCPU, &conv3d_embed_weight_pack);
m.def("conv3d_embed_cpu(Tensor input, Tensor weight, Tensor bias, bool is_vnni) -> Tensor");
m.impl("conv3d_embed_cpu", torch::kCPU, &conv3d_embed_cpu);
// all reduce
m.def("initialize(int size, int rank) -> ()");
m.def("shm_allreduce(Tensor(a!) data, int reduce_op) -> ()");