kernel: support slightly faster merge_state_v2 cuda kernel (#5381)

2025-04-15 12:28:23 +08:00
parent 11421a3f44
commit 388e15c0db
7 changed files with 638 additions and 4 deletions
--- a/sgl-kernel/python/sgl_kernel/init.py
+++ b/sgl-kernel/python/sgl_kernel/init.py
@@ -16,6 +16,7 @@ from sgl_kernel.attention import (
    cutlass_mla_get_workspace_size,
    lightning_attention_decode,
    merge_state,
+    merge_state_v2,
 )
 from sgl_kernel.elementwise import (
    apply_rope_with_cos_sin_cache_inplace,