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

[CPU] add indices in chunk_gated_delta_rule (#29267)

Browse Source
This commit is contained in:
Ma Mingfei
2026-06-26 07:51:07 +08:00
committed by GitHub
parent c8c6757bed
commit 1ba7c79761
8 changed files with 49 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+3 -2
View File
@@ -147,7 +147,8 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
const at::Tensor& cu_seqlens,
bool head_first,
bool use_qk_l2norm_in_kernel,
double eps = 1e-5);
const at::Tensor& initial_state_indices,
double eps = 1e-6);
// weight prepack
at::Tensor convert_weight_packed(at::Tensor& weight);
@@ -525,7 +526,7 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def(
"chunk_gated_delta_rule_cpu(Tensor query, Tensor key, Tensor value, Tensor g, Tensor beta, "
"Tensor initial_state, bool output_final_state, Tensor cu_seqlens, bool head_first, "
"bool use_qk_l2norm_in_kernel, float eps=1e-5) -> (Tensor, Tensor)");
"bool use_qk_l2norm_in_kernel, Tensor initial_state_indices, float eps=1e-6) -> (Tensor, Tensor)");
m.impl("chunk_gated_delta_rule_cpu", torch::kCPU, &chunk_gated_delta_rule_cpu);
// weight prepack