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/python/sgl_kernel/init.py
+++ b/sgl-kernel/python/sgl_kernel/init.py
@@ -15,6 +15,7 @@ from sgl_kernel.attention import (
    cutlass_mla_decode,
    cutlass_mla_get_workspace_size,
    lightning_attention_decode,
+    merge_state,
 )
 from sgl_kernel.elementwise import (
    apply_rope_with_cos_sin_cache_inplace,