feat: adapt merge_state (#5337)

2025-04-12 21:14:04 -07:00
parent 7d3b7c87f5
commit b62e7e99b8
8 changed files with 224 additions and 3 deletions
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -87,6 +87,8 @@ void lightning_attention_decode(
    const torch::Tensor& slope,
    torch::Tensor output,
    torch::Tensor new_kv);
+void merge_state(
+    at::Tensor v_a, at::Tensor s_a, at::Tensor v_b, at::Tensor s_b, at::Tensor v_merged, at::Tensor s_merged);
 void cutlass_mla_decode(
    torch::Tensor const& out,
    torch::Tensor const& q_nope_and_q_pe,