feat(musa): add sparse attention settings

configs/platforms/mthreads_musa/minimax_h3_t2av_tp.json

@@ -13,7 +13,8 @@
   "vae_decode_parallel": true,
   "lazy_load": false,
   "unload_modules": false,
-  "attn_type": "flash_attn3",
+  "attn_type": "sage_attn2",
+  "refiner_attn_type": "flash_attn3",
   "rms_type": "sgl-kernel",
   "rope_type": "torch_real_rope",
   "feature_caching": "NoCaching",

configs/platforms/mthreads_musa/minimax_h3_t2av_tp8_sparse.json

@@ -0,0 +1,46 @@
+{
+  "infer_steps": 30,
+  "target_video_length": 124,
+  "target_height": 544,
+  "target_width": 960,
+  "fps": 24,
+  "target_fps": 24,
+  "enable_cfg": false,
+  "cpu_offload": false,
+  "text_encoder_cpu_offload": false,
+  "text_encoder_tensor_parallel": true,
+  "vae_cpu_offload": false,
+  "vae_decode_parallel": true,
+  "lazy_load": false,
+  "unload_modules": false,
+  "attn_type": "dynamic_sparse_attn",
+  "refiner_attn_type": "flash_attn3",
+  "dynamic_sparse_attn_setting": {
+    "topk": 15,
+    "force_local_blocks": 2,
+    "operator": "musa_moorcat_sparse",
+    "operator_setting": {
+      "validate_lut": true
+    }
+  },
+  "rms_type": "sgl-kernel",
+  "rope_type": "torch_real_rope",
+  "feature_caching": "NoCaching",
+  "use_compile": true,
+  "warmup": true,
+  "vae_use_compile": true,
+  "video_flow_shift": 12.0,
+  "audio_flow_shift": 3.0,
+  "vae_spatial_scale_factor": 16,
+  "audio_sampling_rate": 32000,
+  "audio_latents_per_second": 40,
+  "audio_channels": 2,
+  "keep_latents_dtype_in_scheduler": true,
+  "parallel": {
+    "tensor_p_size": 8
+  },
+  "video_codec_options": {
+    "preset": "ultrafast",
+    "crf": "18"
+  }
+}

lightx2v/common/ops/attn/dynamic_sparse_attn.py

@@ -1,12 +1,12 @@
 import torch
 from loguru import logger
 
-from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER
+from lightx2v.utils.registry_factory import ATTN_WEIGHT_REGISTER, SPARSE_OPERATOR_REGISTER
 
 from .kernels.sla_kernel import _attention
 from .kernels.sla_kernel_ar import _attention_ar
 from .template import AttnWeightTemplate
-from .utils.sla_util import get_block_map, get_cuda_arch
+from .utils.sla_util import get_block_lut_nhd, get_block_lut_nhd_uncentered, get_block_map, get_cuda_arch
 from .utils.sla_util_blhd import get_block_map_blhd
 from .utils.sparge_util import block_map_incremental_lut_triton, block_map_ordinal_lut_triton, sage2_block_sparse_attn
 
@@ -69,15 +69,14 @@ class DynamicSparseAttnWeight(AttnWeightTemplate):
     per_block_mean = False
 
     def __init__(self, config=None):
-        self.config = dict(config or {})
-        self.sparsity_ratio = float(self.config.get("sparsity_ratio", type(self).sparsity_ratio))
-        self.operator = self.config.get("operator", type(self).operator)
-        self.per_block_mean = bool(self.config.get("per_block_mean", type(self).per_block_mean))
+        config = config or {}
+        self.sparsity_ratio = config.get("sparsity_ratio", type(self).sparsity_ratio)
+        self.operator = config.get("operator", type(self).operator)
+        self.per_block_mean = config.get("per_block_mean", type(self).per_block_mean)
 
         if not 0.0 <= self.sparsity_ratio < 1.0:
             raise ValueError(f"dynamic sparse attention sparsity_ratio must be in [0, 1), got {self.sparsity_ratio}")
 
-        self.arch = get_cuda_arch(torch.cuda.current_device())
         self.topk = 1 - self.sparsity_ratio
         if self.operator == "triton":
             self.BLKQ, self.BLKK = 64, 64
@@ -86,6 +85,7 @@ def __init__(self, config=None):
             self.BLKQ, self.BLKK = 128, 128
             self.apply_func = self.apply_triton_ar
         elif self.operator == "sage2":
+            self.arch = get_cuda_arch(torch.cuda.current_device())
             if self.arch == "sm90":
                 self.BLKQ, self.BLKK = 64, 128
             else:
@@ -100,11 +100,23 @@ def __init__(self, config=None):
         elif self.operator == "magi":
             self.BLKQ, self.BLKK = 128, 128
             self.apply_func = self.apply_magi
+        elif self.operator in SPARSE_OPERATOR_REGISTER:
+            self.fixed_topk = config.get("topk")
+            if self.fixed_topk is not None and self.fixed_topk <= 0:
+                raise ValueError(f"dynamic sparse attention topk must be positive, got {self.fixed_topk}")
+
+            operator_setting = config.get("operator_setting", {}).copy()
+            if self.fixed_topk is not None:
+                operator_setting.setdefault("topk", self.fixed_topk)
+            self.sparse_operator = SPARSE_OPERATOR_REGISTER[self.operator](operator_setting)
+            self.fixed_topk = getattr(self.sparse_operator, "topk", self.fixed_topk)
+            self.force_local_blocks = config.get("force_local_blocks", 0)
+            self.BLKQ = self.sparse_operator.q_block_size
+            self.BLKK = self.sparse_operator.k_block_size
+            self.apply_func = self.apply_registered_operator
         else:
             raise NotImplementedError(f"Not supported SLA operator: {self.operator}.")
 
-        # logger.info(f"DynamicSparseAttnWeight: sparsity_ratio={self.sparsity_ratio}, operator={self.operator}, topk={self.topk}, BLKQ={self.BLKQ}, BLKK={self.BLKK}")
-
     def apply(
         self,
         q,
@@ -188,6 +200,53 @@ def apply_sage2(
         out = out.transpose(1, 2).reshape(max_seqlen_q, -1)
         return out
 
+    def apply_registered_operator(
+        self,
+        q,
+        k,
+        v,
+        cu_seqlens_q=None,
+        cu_seqlens_kv=None,
+        max_seqlen_q=None,
+        max_seqlen_kv=None,
+        **kwargs,
+    ):
+        if getattr(self.sparse_operator, "block_indices_only", False):
+            get_lut = get_block_lut_nhd if getattr(self.sparse_operator, "center_k", True) else get_block_lut_nhd_uncentered
+            block_indices, _ = get_lut(
+                q,
+                k,
+                topk_ratio=self.topk,
+                BLKQ=self.BLKQ,
+                BLKK=self.BLKK,
+                topk=self.fixed_topk,
+                force_local_blocks=self.force_local_blocks,
+            )
+            sparse_map = None
+        else:
+            q_for_mask = q.unsqueeze(0).transpose(1, 2).contiguous()
+            k_for_mask = k.unsqueeze(0).transpose(1, 2).contiguous()
+            sparse_map, block_indices, _ = get_block_map(
+                q_for_mask,
+                k_for_mask,
+                topk_ratio=self.topk,
+                BLKQ=self.BLKQ,
+                BLKK=self.BLKK,
+                topk=self.fixed_topk,
+            )
+        return self.sparse_operator(
+            q,
+            k,
+            v,
+            sparse_map,
+            block_indices=block_indices,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_kv=cu_seqlens_kv,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_kv=max_seqlen_kv,
+            **kwargs,
+        )
+
     def apply_sage3(
         self,
         q,
@@ -204,7 +263,7 @@ def apply_sage3(
         k = k.unsqueeze(0).transpose(1, 2).contiguous()
         v = v.unsqueeze(0).transpose(1, 2).contiguous()
 
-        sparse_map, lut, real_topk = get_block_map(q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
+        sparse_map, _, _ = get_block_map(q, k, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
         lut, valid_block_num = block_map_ordinal_lut_triton(sparse_map)
         out = sage3_block_sparse_attn(q, k, v, lut, valid_block_num, per_block_mean=self.per_block_mean)
         out = out.transpose(1, 2).reshape(max_seqlen_q, -1)
@@ -224,7 +283,7 @@ def apply_fa4(
         # (L, H, D) -> (B, L, H, D)
         qt = q.unsqueeze(0).transpose(1, 2).contiguous()
         kt = k.unsqueeze(0).transpose(1, 2).contiguous()
-        sparse_map, lut, real_topk = get_block_map(qt, kt, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
+        sparse_map, _, _ = get_block_map(qt, kt, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
 
         # (L, H, D) -> (B, L, H, D)
         q = q.unsqueeze(0)
@@ -268,7 +327,7 @@ def apply_magi(
         q_block_map = q_block_map.contiguous()
         k_block_map = k_block_map.contiguous()
 
-        sparse_map, lut, real_topk = get_block_map(q_block_map, k_block_map, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
+        sparse_map, _, _ = get_block_map(q_block_map, k_block_map, topk_ratio=self.topk, BLKQ=self.BLKQ, BLKK=self.BLKK)
         seqlen, head_num, head_dim = q.shape
 
         q_ranges, k_ranges = self.generate_qk_ranges(sparse_map[0], self.BLKQ, self.BLKK, seqlen)