feat: add CoSA sparse attention algorithm

README.md

@@ -116,6 +116,7 @@ A more accessible, comprehensive, and efficient toolkit for large model compress
               <li><a href="https://angelslim.readthedocs.io/zh-cn/latest/features/sparse_attention/index.html">XAttention</a></li>
               <li><a href="https://angelslim.readthedocs.io/zh-cn/latest/features/sparse_attention/index.html">FlashPrefill</a></li>
               <li><a href="https://angelslim.readthedocs.io/zh-cn/latest/features/sparse_attention/index.html">VecAttention</a></li>
+              <li><a href="https://arxiv.org/pdf/2607.25291">CoSA</a></li>
             </ul>
           </li>
           <li>

README_cn.md

@@ -116,6 +116,7 @@
               <li><a href="https://angelslim.readthedocs.io/zh-cn/latest/features/sparse_attention/index.html">XAttention</a></li>
               <li><a href="https://angelslim.readthedocs.io/zh-cn/latest/features/sparse_attention/index.html">FlashPrefill</a></li>
               <li><a href="https://angelslim.readthedocs.io/zh-cn/latest/features/sparse_attention/index.html">VecAttention</a></li>
+              <li><a href="https://arxiv.org/pdf/2607.25291">CoSA</a></li>
             </ul>
           </li>
           <li>

angelslim/compressor/sparsity/algorithms/__init__.py

@@ -31,7 +31,15 @@
 # others — but the swallowed exception is LOGGED (WARNING) rather than silently
 # dropped, so a genuine syntax/import bug surfaces its real cause here instead of
 # a confusing "not registered" ValueError much later.
-for _name in ("stem", "minference", "flexprefill", "xattention", "flashprefill", "vecattention"):
+for _name in (
+    "stem",
+    "minference",
+    "flexprefill",
+    "xattention",
+    "flashprefill",
+    "vecattention",
+    "cosa",
+):
     try:
         importlib.import_module(f".{_name}", __name__)
     except Exception as _e:  # noqa: BLE001  pragma: no cover - defensive

angelslim/compressor/sparsity/algorithms/cosa/__init__.py

@@ -0,0 +1,17 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .algorithm import CoSA  # noqa: F401
+
+__all__ = ["CoSA"]

angelslim/compressor/sparsity/algorithms/cosa/algorithm.py

@@ -0,0 +1,130 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CoSA sparse-prefill algorithm.
+
+CoSA keeps, per query-block, the smallest set of causal key-blocks whose
+cumulative attention mass covers ``topp``. Mass is estimated from a
+``stride``-subsampled proxy of ``Q @ K^T`` (Triton, ``ops/proxy.py``). The
+selection is emitted as block ids **ordered by descending priority**; the
+``hpc`` kernel walks that order and skips the tail below ``skipsoftmax_threshold``.
+
+Hard constraints: ``triton`` + internal ``hpc``, bf16, head_dim 128,
+``block_size == 128``, ``batch == 1``. No pseudo-sparse fallback.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ..._base import AlgorithmTraits, IncompatibleConfigError, SparsityAlgorithm
+from ...registry import SparsityAlgorithmRegistry
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel
+
+_DEFAULT_STRIDE = 8
+_DEFAULT_TOPP = 0.8
+_DEFAULT_BLOCK_SIZE = 128
+# hpc anyorderskip: ``< 0`` disables threshold skip; ``>= 0`` enables.
+_DEFAULT_SKIPSOFTMAX_THRESHOLD = 1000
+
+
+@SparsityAlgorithmRegistry.register("cosa")
+class CoSA(SparsityAlgorithm):
+    """CoSA ordered block-sparse prefill (Qwen3 / Qwen3-MoE / Hunyuan-V3)."""
+
+    name = "cosa"
+
+    @property
+    def traits(self) -> AlgorithmTraits:
+        return AlgorithmTraits(
+            requires_unrepeated_kv=False,  # proxy and kernel both see post-repeat K/V
+            needs_calibration=False,
+            supports_padding_mask=False,  # block selection never sees the mask
+            model_modal="any",
+            compatible_model_types=frozenset({"qwen3", "qwen3_moe", "hy_v3"}),
+        )
+
+    def setup(self, model: "PreTrainedModel") -> None:
+        if self.allow_pseudo_sparse:
+            raise IncompatibleConfigError(
+                "CoSA has no pseudo-sparse reference path, so "
+                "allow_pseudo_sparse=true cannot be honoured. Remove it; the "
+                "proxy requires `triton` and the attention requires the `hpc` "
+                "extension unconditionally."
+            )
+
+        cfg = dict(self.attn_kwargs)
+        stride = int(cfg.get("stride", _DEFAULT_STRIDE))
+        if stride < 1:
+            raise IncompatibleConfigError(f"stride must be >= 1, got {stride}.")
+
+        chunk_size = cfg.get("chunk_size", None)
+        self._cfg = {
+            "stride": stride,
+            "topp": self._resolve_topp(cfg.get("topp", _DEFAULT_TOPP)),
+            "block_size": int(cfg.get("block_size", _DEFAULT_BLOCK_SIZE)),
+            "chunk_size": None if chunk_size is None else int(chunk_size),
+            "skipsoftmax_threshold": float(
+                cfg.get("skipsoftmax_threshold", _DEFAULT_SKIPSOFTMAX_THRESHOLD)
+            ),
+        }
+        self._per_instance_ready = True
+
+    @staticmethod
+    def _resolve_topp(topp) -> float:
+        """Validate ``topp`` as a single scalar in ``(0, 1]``."""
+        if isinstance(topp, bool) or not isinstance(topp, (int, float)):
+            raise IncompatibleConfigError(
+                f"topp must be a scalar float, got {type(topp).__name__}. "
+                f"Per-head topp tables are not supported."
+            )
+        topp = float(topp)
+        if not 0.0 < topp <= 1.0:
+            raise IncompatibleConfigError(
+                f"topp must be in (0, 1], got {topp}. It is a cumulative "
+                f"attention-mass coverage target; 1.0 keeps every causal block."
+            )
+        return topp
+
+    def build_attn_forward(self, attn_module, model: "PreTrainedModel"):
+        if not self._per_instance_ready:
+            raise IncompatibleConfigError(
+                "CoSA.setup(model) must be called before build_attn_forward."
+            )
+        from .._forward_templates import get_forward_template
+        from .._forward_templates._common import repeat_kv
+        from .prefill import cosa_prefill
+
+        cfg = self._cfg
+
+        def _prefill_fn(attn, query_states, key_states, value_states):
+            # The kernel consumes post-repeat K/V at num_attention_heads.
+            return cosa_prefill(
+                query_states,
+                repeat_kv(key_states, attn.num_key_value_groups),
+                repeat_kv(value_states, attn.num_key_value_groups),
+                stride=cfg["stride"],
+                topp=cfg["topp"],
+                block_size=cfg["block_size"],
+                chunk_size=cfg["chunk_size"],
+                skipsoftmax_threshold=cfg["skipsoftmax_threshold"],
+            )
+
+        model_type = getattr(model.config, "model_type", None)
+        template = get_forward_template(model_type)
+        attn_module.attn_forward_config = cfg
+        forward = template(_prefill_fn, supports_padding_mask=self.traits.supports_padding_mask)
+        return forward.__get__(attn_module, type(attn_module))

angelslim/compressor/sparsity/algorithms/cosa/ops/__init__.py

@@ -0,0 +1,15 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CoSA ops: Triton selection proxy and hpc layout adapters."""

angelslim/compressor/sparsity/algorithms/cosa/ops/hpc_adaption.py

@@ -0,0 +1,99 @@
+# Copyright 2025 Tencent Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Layout adapters: dense ``(B, H, L, D)`` <-> hpc varlen + paged KV."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+__all__ = [
+    "convert_to_varlen_q",
+    "convert_to_paged_kv",
+    "convert_varlen_out_to_bshd",
+]
+
+
+def convert_to_varlen_q(query_states: torch.Tensor, dtype: torch.dtype = torch.bfloat16):
+    """``(B, H, S, D)`` -> ``(B*S, H, D)`` plus ``cu_seqlens``."""
+    if query_states.dim() != 4:
+        raise ValueError(
+            f"query_states must be 4-D (B, H, S, D), got {tuple(query_states.shape)}."
+        )
+    B, H_q, q_len, D = query_states.shape
+    if D != 128:
+        raise ValueError(f"hpc kernel requires head_dim == 128, got {D}.")
+
+    q = query_states.permute(0, 2, 1, 3).reshape(B * q_len, H_q, D).to(dtype).contiguous()
+    cu_seqlens_q = torch.arange(0, B + 1, device=query_states.device, dtype=torch.int32) * int(
+        q_len
+    )
+    return q, cu_seqlens_q, int(q_len)
+
+
+def convert_to_paged_kv(
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    page_size: int = 64,
+    logical_block_size: int = 128,
+    dtype: torch.dtype = torch.bfloat16,
+):
+    """``(B, H, S, D)`` -> paged ``(B*pages, page_size, H, D)`` + page table."""
+    if page_size not in (32, 64):
+        raise ValueError(f"page_size must be 32 or 64 (hpc paged block size), got {page_size}.")
+    if logical_block_size % page_size != 0:
+        raise ValueError(
+            f"logical_block_size ({logical_block_size}) must be a multiple of "
+            f"page_size ({page_size})."
+        )
+    if key_states.shape != value_states.shape:
+        raise ValueError(
+            f"key/value must share shape, got {tuple(key_states.shape)} vs "
+            f"{tuple(value_states.shape)}."
+        )
+    if key_states.dim() != 4:
+        raise ValueError(f"key/value must be 4-D (B, H, S, D), got {tuple(key_states.shape)}.")
+
+    B, H_kv, k_len, D = key_states.shape
+    if D != 128:
+        raise ValueError(f"hpc kernel requires head_dim == 128, got {D}.")
+
+    k_block_num = (k_len + logical_block_size - 1) // logical_block_size
+    k_pad = k_block_num * logical_block_size
+    pages_per_batch = k_pad // page_size
+
+    def _to_paged(x: torch.Tensor) -> torch.Tensor:
+        if k_pad != k_len:
+            x = F.pad(x, (0, 0, 0, k_pad - k_len))
+        x = x.permute(0, 2, 1, 3).reshape(B * pages_per_batch, page_size, H_kv, D)
+        return x.to(dtype).contiguous()
+
+    kcache = _to_paged(key_states)
+    vcache = _to_paged(value_states)
+    block_ids = torch.arange(
+        B * pages_per_batch, device=key_states.device, dtype=torch.int32
+    ).reshape(B, pages_per_batch)
+    seqlens_kvcache = torch.full((B,), int(k_len), device=key_states.device, dtype=torch.int32)
+    return kcache, vcache, block_ids, seqlens_kvcache
+
+
+def convert_varlen_out_to_bshd(out: torch.Tensor, batch_size: int, q_len: int) -> torch.Tensor:
+    """``(B*S, H, D)`` -> ``(B, H, S, D)``."""
+    total_seq, H_q, D = out.shape
+    if total_seq != batch_size * q_len:
+        raise ValueError(
+            f"out total_seq ({total_seq}) != batch_size*q_len ({batch_size}*{q_len})."
+        )
+    return out.reshape(batch_size, q_len, H_q, D).permute(0, 2, 1, 3).contiguous()