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

2025-04-14 21:28:23 -07:00
parent 11421a3f44
commit 388e15c0db
7 changed files with 638 additions and 4 deletions
@@ -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,