Update Triton extend backend interface (#3309)

2025-02-05 18:12:22 +08:00
parent 7aad8d1854
commit de5533341e
5 changed files with 427 additions and 69 deletions
@@ -17,12 +17,10 @@ class DoubleSparseAttnBackend(AttentionBackend):
    def __init__(self, model_runner: ModelRunner):
        # Lazy import to avoid the initialization of cuda context
        from sglang.srt.layers.attention.triton_ops.double_sparsity_attention import (
+            extend_attention_fwd,
            flash_decode_attention_fwd,
            flash_decode_sparse_attention_fwd,
        )
-        from sglang.srt.layers.attention.triton_ops.extend_attention import (
-            extend_attention_fwd,
-        )

        super().__init__()