diff --git a/README.md b/README.md
index fd21d915..16947aeb 100644
--- a/README.md
+++ b/README.md
@@ -116,6 +116,7 @@ A more accessible, comprehensive, and efficient toolkit for large model compress
XAttention
FlashPrefill
VecAttention
+ CoSA
diff --git a/README_cn.md b/README_cn.md
index 75db95ac..c0d26b6c 100644
--- a/README_cn.md
+++ b/README_cn.md
@@ -116,6 +116,7 @@
XAttention
FlashPrefill
VecAttention
+ CoSA
diff --git a/angelslim/compressor/sparsity/algorithms/__init__.py b/angelslim/compressor/sparsity/algorithms/__init__.py
index 34df90fb..9c0fc550 100644
--- a/angelslim/compressor/sparsity/algorithms/__init__.py
+++ b/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
diff --git a/angelslim/compressor/sparsity/algorithms/cosa/__init__.py b/angelslim/compressor/sparsity/algorithms/cosa/__init__.py
new file mode 100644
index 00000000..7be99d7b
--- /dev/null
+++ b/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"]
diff --git a/angelslim/compressor/sparsity/algorithms/cosa/algorithm.py b/angelslim/compressor/sparsity/algorithms/cosa/algorithm.py
new file mode 100644
index 00000000..9e2c248a
--- /dev/null
+++ b/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))
diff --git a/angelslim/compressor/sparsity/algorithms/cosa/ops/__init__.py b/angelslim/compressor/sparsity/algorithms/cosa/ops/__init__.py
new file mode 100644
index 00000000..c58d6430
--- /dev/null
+++ b/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."""
diff --git a/angelslim/compressor/sparsity/algorithms/cosa/ops/hpc_adaption.py b/angelslim/compressor/sparsity/algorithms/cosa/ops/hpc_adaption.py
new file mode 100644
index 00000000..f5184c41
--- /dev/null
+++ b/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()
diff --git a/angelslim/compressor/sparsity/algorithms/cosa/ops/proxy.py b/angelslim/compressor/sparsity/algorithms/cosa/ops/proxy.py
new file mode 100644
index 00000000..da24d777
--- /dev/null
+++ b/angelslim/compressor/sparsity/algorithms/cosa/ops/proxy.py
@@ -0,0 +1,481 @@
+# 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 block-selection proxy.
+
+Strided ``Q @ K^T`` -> per-block mass / row-max flags -> keep the smallest
+prefix whose cumulative mass covers ``topp`` (plus sink / diagonal /
+row-max), emitted as priority-ordered block ids for the hpc kernel.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _flat_group_gemm_fuse_reshape_kernel(
+ Q,
+ K,
+ Out,
+ stride_qz,
+ stride_qh,
+ stride_qn,
+ stride_kz,
+ stride_kh,
+ stride_kn,
+ stride_oz,
+ stride_oh,
+ stride_on,
+ chunk_start,
+ chunk_end,
+ H: tl.constexpr,
+ STRIDE: tl.constexpr,
+ HEAD_DIM: tl.constexpr,
+ BLOCK_M: tl.constexpr,
+ BLOCK_N: tl.constexpr,
+):
+ block_m = tl.program_id(0).to(tl.int64)
+ block_n = tl.program_id(1).to(tl.int64)
+ batch_id = tl.program_id(2).to(tl.int64) // H
+ head_id = tl.program_id(2).to(tl.int64) % H
+
+ if chunk_start + (block_m + 1) * BLOCK_M <= block_n * BLOCK_N:
+ return
+
+ Q_ptrs = (
+ Q + batch_id * stride_qz + head_id * stride_qh + block_m * BLOCK_M * STRIDE * stride_qn
+ )
+ K_ptrs = (
+ K + batch_id * stride_kz + head_id * stride_kh + block_n * BLOCK_N * STRIDE * stride_kn
+ )
+ Q_ptrs = (
+ Q_ptrs
+ + tl.arange(0, BLOCK_M)[:, None] * (stride_qn * STRIDE)
+ + tl.arange(0, HEAD_DIM)[None, :]
+ + stride_qn * (STRIDE - 1)
+ )
+ K_ptrs = (
+ K_ptrs
+ + tl.arange(0, BLOCK_N)[None, :] * (stride_kn * STRIDE)
+ + tl.arange(0, HEAD_DIM)[:, None]
+ )
+
+ o = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+ for iter in range(STRIDE):
+ q = tl.load(Q_ptrs - iter * stride_qn)
+ k = tl.load(K_ptrs + iter * stride_kn)
+ o += tl.dot(q, k)
+
+ O_ptrs = (
+ Out
+ + batch_id * stride_oz
+ + head_id * stride_oh
+ + block_m * BLOCK_M * stride_on
+ + block_n * BLOCK_N
+ )
+ O_ptrs = O_ptrs + tl.arange(0, BLOCK_M)[:, None] * stride_on + tl.arange(0, BLOCK_N)[None, :]
+ tl.store(O_ptrs, o.to(Out.type.element_ty))
+
+
+def flat_group_gemm_fuse_reshape(query_states, key_states, stride, chunk_start, chunk_end):
+ batch_size, num_heads, q_len, head_dim = query_states.shape
+ kv_len = key_states.shape[2]
+
+ assert key_states.shape[0] == batch_size
+ assert key_states.shape[1] == num_heads
+ assert key_states.shape[3] == head_dim
+
+ output = torch.empty(
+ (batch_size, num_heads, q_len // stride, kv_len // stride),
+ dtype=query_states.dtype,
+ device=query_states.device,
+ )
+ BLOCK_M = 128
+ BLOCK_N = 128
+ assert q_len % (stride * BLOCK_M) == 0
+ assert kv_len % (stride * BLOCK_N) == 0
+
+ grid = (q_len // stride // BLOCK_M, kv_len // stride // BLOCK_N, batch_size * num_heads)
+ _flat_group_gemm_fuse_reshape_kernel[grid](
+ query_states,
+ key_states,
+ output,
+ query_states.stride(0),
+ query_states.stride(1),
+ query_states.stride(2),
+ key_states.stride(0),
+ key_states.stride(1),
+ key_states.stride(2),
+ output.stride(0),
+ output.stride(1),
+ output.stride(2),
+ chunk_start,
+ chunk_end,
+ num_heads,
+ stride,
+ head_dim,
+ BLOCK_M,
+ BLOCK_N,
+ )
+ return output
+
+
+@triton.jit
+def _fused_cosa_proxy_kernel(
+ In,
+ Out_score,
+ Out_rowmax,
+ scale,
+ input_stride_0,
+ input_stride_1,
+ input_stride_2,
+ score_stride_0,
+ score_stride_1,
+ score_stride_2,
+ rowmax_stride_0,
+ rowmax_stride_1,
+ rowmax_stride_2,
+ real_q_len,
+ k_len,
+ chunk_start,
+ chunk_end,
+ segment_size: tl.constexpr,
+ block_size: tl.constexpr,
+):
+ block_id = tl.program_id(0)
+ head_id = tl.program_id(1)
+ batch_id = tl.program_id(2)
+
+ offs_q = tl.arange(0, block_size) + chunk_start + block_id * block_size
+ offs_k = tl.arange(0, segment_size)
+
+ num_iters = k_len // segment_size
+ num_iters_before_causal = (chunk_start + (block_id + 1) * block_size - 1) // segment_size
+
+ m_i = tl.zeros([block_size], dtype=tl.float32) - float("inf")
+
+ input_ptr = (
+ In
+ + batch_id * input_stride_0
+ + head_id * input_stride_1
+ + block_id * block_size * input_stride_2
+ )
+ input_ptr = (
+ input_ptr + tl.arange(0, segment_size) + tl.arange(0, block_size)[:, None] * input_stride_2
+ )
+
+ score_ptr = (
+ Out_score
+ + batch_id * score_stride_0
+ + head_id * score_stride_1
+ + block_id * score_stride_2
+ )
+ score_ptr = score_ptr + tl.arange(0, segment_size // block_size)
+
+ rowmax_ptr = (
+ Out_rowmax
+ + batch_id * rowmax_stride_0
+ + head_id * rowmax_stride_1
+ + block_id * rowmax_stride_2
+ )
+ rowmax_ptr = rowmax_ptr + tl.arange(0, segment_size // block_size)
+
+ for iter in range(0, num_iters_before_causal):
+ X = tl.load(input_ptr + iter * segment_size).to(tl.float32) * scale
+ m_i = tl.maximum(m_i, tl.max(X, 1))
+
+ for iter in range(num_iters_before_causal, num_iters_before_causal + 1):
+ X = tl.load(input_ptr + iter * segment_size).to(tl.float32) * scale
+ mask = offs_q[:, None] >= (offs_k[None, :] + iter * segment_size)
+ X = tl.where(mask, X, -1.0e6)
+ m_i = tl.maximum(m_i, tl.max(X, 1))
+
+ sum_mask = offs_q[:, None] < real_q_len
+
+ for iter in range(0, num_iters_before_causal):
+ X = tl.load(input_ptr + iter * segment_size).to(tl.float32) * scale
+ X = tl.reshape(X, (block_size, segment_size // block_size, block_size))
+ blockscore = tl.max(X, 2)
+ E = tl.exp2(blockscore - m_i[:, None])
+ E = tl.where(sum_mask, E, 0)
+ score = tl.sum(E, 0)
+ hr = tl.max(tl.where(E >= (1.0 - 1e-3), 1.0, 0.0), 0)
+ tl.store(
+ score_ptr + iter * segment_size // block_size, score.to(Out_score.type.element_ty)
+ )
+ tl.store(rowmax_ptr + iter * segment_size // block_size, hr.to(Out_rowmax.type.element_ty))
+
+ for iter in range(num_iters_before_causal, num_iters_before_causal + 1):
+ X = tl.load(input_ptr + iter * segment_size).to(tl.float32) * scale
+ mask = offs_q[:, None] >= (offs_k[None, :] + iter * segment_size)
+ X = tl.where(mask, X, -1.0e6)
+ X = tl.reshape(X, (block_size, segment_size // block_size, block_size))
+ blockscore = tl.max(X, 2)
+ E = tl.exp2(blockscore - m_i[:, None])
+ E = tl.where(sum_mask, E, 0)
+ score = tl.sum(E, 0)
+ hr = tl.max(tl.where(E >= (1.0 - 1e-3), 1.0, 0.0), 0)
+ tl.store(
+ score_ptr + iter * segment_size // block_size, score.to(Out_score.type.element_ty)
+ )
+ tl.store(rowmax_ptr + iter * segment_size // block_size, hr.to(Out_rowmax.type.element_ty))
+
+ for iter in range(num_iters_before_causal + 1, num_iters):
+ zeros = tl.zeros([segment_size // block_size], dtype=tl.float32)
+ tl.store(
+ score_ptr + iter * segment_size // block_size, zeros.to(Out_score.type.element_ty)
+ )
+ tl.store(
+ rowmax_ptr + iter * segment_size // block_size, zeros.to(Out_rowmax.type.element_ty)
+ )
+
+
+def fused_cosa_proxy(
+ attn_weights_slice,
+ reshaped_block_size,
+ segment_size,
+ chunk_start,
+ chunk_end,
+ real_q_len,
+ scale,
+):
+ batch_size, num_heads, q_len, k_len = attn_weights_slice.shape
+ assert q_len % reshaped_block_size == 0
+ assert k_len % segment_size == 0
+ assert segment_size % reshaped_block_size == 0
+ assert attn_weights_slice.stride(-1) == 1
+
+ attn_score = torch.empty(
+ (batch_size, num_heads, q_len // reshaped_block_size, k_len // reshaped_block_size),
+ dtype=attn_weights_slice.dtype,
+ device=attn_weights_slice.device,
+ )
+ has_rowmax = torch.empty(
+ (batch_size, num_heads, q_len // reshaped_block_size, k_len // reshaped_block_size),
+ dtype=torch.int8,
+ device=attn_weights_slice.device,
+ )
+ grid = (q_len // reshaped_block_size, num_heads, batch_size)
+ _fused_cosa_proxy_kernel[grid](
+ attn_weights_slice,
+ attn_score,
+ has_rowmax,
+ scale,
+ attn_weights_slice.stride(0),
+ attn_weights_slice.stride(1),
+ attn_weights_slice.stride(2),
+ attn_score.stride(0),
+ attn_score.stride(1),
+ attn_score.stride(2),
+ has_rowmax.stride(0),
+ has_rowmax.stride(1),
+ has_rowmax.stride(2),
+ real_q_len,
+ k_len,
+ chunk_start,
+ chunk_end,
+ segment_size,
+ reshaped_block_size,
+ )
+ return attn_score, has_rowmax
+
+
+def select_and_order_blocks_row_max(
+ mass: torch.Tensor,
+ has_rowmax: torch.Tensor,
+ q_block_num: int,
+ k_block_num: int,
+ threshold: float,
+) -> torch.Tensor:
+ batch_size, head_num, chunk_num, block_num = mass.shape
+ device = mass.device
+ mass = mass.to(torch.float32)
+ has_rowmax = has_rowmax.to(torch.bool)
+
+ threshold = float(threshold)
+ total_sum = mass.sum(dim=-1, keepdim=True)
+ if threshold == 1.0:
+ # inf budget: total_sum * 1.0 can drop the last block to float rounding.
+ required_sum = torch.full_like(total_sum, float("inf"))
+ else:
+ required_sum = total_sum * threshold
+
+ causal = torch.tril(
+ torch.ones(chunk_num, block_num, dtype=torch.bool, device=device),
+ diagonal=block_num - chunk_num,
+ ).view(1, 1, chunk_num, block_num)
+
+ mandatory = torch.zeros((1, 1, chunk_num, block_num), dtype=torch.bool, device=device)
+ mandatory[:, :, :, 0] = True
+ rows = torch.arange(chunk_num, device=device)
+ diag_cols = rows + (block_num - chunk_num)
+ mandatory[0, 0, rows, diag_cols] = True
+ mandatory = mandatory & causal
+
+ C = total_sum.amax() + 1.0
+ priority = mass + C * has_rowmax.to(mass.dtype) + (2.0 * C) * mandatory.to(mass.dtype)
+ neg_inf = torch.full_like(priority, float("-inf"))
+ priority = torch.where(causal.expand_as(priority), priority, neg_inf)
+
+ sort_perm = priority.argsort(dim=-1, descending=True)
+ mass_in_order = torch.gather(mass, dim=-1, index=sort_perm)
+ cum_before = torch.cat(
+ [
+ torch.zeros((batch_size, head_num, chunk_num, 1), device=device),
+ mass_in_order[:, :, :, :-1],
+ ],
+ dim=-1,
+ ).cumsum(dim=-1)
+ keep_in_order = cum_before < required_sum
+
+ mand_in_order = torch.gather(mandatory.expand_as(mass).to(torch.bool), dim=-1, index=sort_perm)
+ causal_in_order = torch.gather(causal.expand_as(mass).to(torch.bool), dim=-1, index=sort_perm)
+ kept = (keep_in_order | mand_in_order) & causal_in_order
+
+ row_blockmask = sort_perm.to(torch.int32)
+ row_blockmask = torch.where(kept, row_blockmask, torch.full_like(row_blockmask, -1))
+ assert not (row_blockmask == -1).all(dim=-1).any(), "row_blockmask contains a row with all -1"
+ return row_blockmask.contiguous()
+
+
+def _select_chunk(
+ attn_sum_chunk,
+ has_rowmax_chunk,
+ chunk_idx,
+ npc,
+ offset_u,
+ q_block_num_u,
+ k_block_num_u,
+ topp,
+):
+ rows_lo = chunk_idx * npc
+ rows_hi = min((chunk_idx + 1) * npc, q_block_num_u)
+ n_rows = rows_hi - rows_lo
+ kv_u = rows_hi + offset_u
+
+ mass_chunk = attn_sum_chunk[:, :, :n_rows, :kv_u].contiguous()
+ hr_chunk = has_rowmax_chunk[:, :, :n_rows, :kv_u].contiguous()
+ sel = select_and_order_blocks_row_max(mass_chunk, hr_chunk, n_rows, kv_u, topp)
+ if kv_u < k_block_num_u:
+ sel = F.pad(sel, (0, k_block_num_u - kv_u), value=-1)
+ return sel
+
+
+def _validate_topp(topp) -> float:
+ if isinstance(topp, bool) or not isinstance(topp, (int, float)):
+ raise TypeError(f"topp must be a scalar float, got {type(topp).__name__}.")
+ return float(topp)
+
+
+def cosa_estimate(
+ query_states: torch.Tensor,
+ key_states: torch.Tensor,
+ block_size: int,
+ stride: int,
+ topp: float = 0.8,
+ chunk_size: int = 16384,
+):
+ topp = _validate_topp(topp)
+
+ _, _, k_len, head_dim = key_states.shape
+ _, _, q_len, _ = query_states.shape
+
+ k_num_to_pad = ((k_len + chunk_size - 1) // chunk_size) * chunk_size - k_len
+ q_num_to_pad = ((q_len + chunk_size - 1) // chunk_size) * chunk_size - q_len
+ k_chunk_num = (k_len + k_num_to_pad) // chunk_size
+ k_block_num = (k_len + k_num_to_pad) // block_size
+ q_block_num = (q_len + q_num_to_pad) // block_size
+ q_chunk_num = (q_len + q_num_to_pad) // chunk_size
+ assert k_chunk_num >= q_chunk_num
+
+ # Keep F.pad on the input device (do not force "cuda:0" under device_map).
+ if k_num_to_pad > 0:
+ pad_key_states = F.pad(key_states, (0, 0, 0, k_num_to_pad), value=0)
+ else:
+ pad_key_states = key_states
+ if q_num_to_pad > 0:
+ pad_query_states = F.pad(query_states, (0, 0, 0, q_num_to_pad), value=0)
+ else:
+ pad_query_states = query_states
+
+ reshaped_chunk_size = chunk_size // stride
+ reshaped_block_size = block_size // stride
+ k_reshaped_num_to_pad = k_num_to_pad // stride
+ k_reshaped_seq_len = (k_len + k_num_to_pad) // stride
+ num_blocks_per_chunk = reshaped_chunk_size // reshaped_block_size
+
+ q_block_num_u = (q_len + block_size - 1) // block_size
+ k_block_num_u = (k_len + block_size - 1) // block_size
+ offset_u = k_block_num_u - q_block_num_u
+ npc = num_blocks_per_chunk
+
+ scale = 1.4426950408889634 / math.sqrt(head_dim) / stride
+ segment_size = min(4096, reshaped_block_size)
+ real_q_len = k_reshaped_seq_len - k_reshaped_num_to_pad
+
+ attn_sum_list = []
+ row_blockmask_list = []
+
+ for chunk_idx in range(q_chunk_num):
+ chunk_start = (
+ k_block_num - q_block_num
+ ) * reshaped_block_size + chunk_idx * reshaped_chunk_size
+ chunk_end = chunk_start + reshaped_chunk_size
+ attn_weights_slice = flat_group_gemm_fuse_reshape(
+ pad_query_states[
+ :,
+ :,
+ (chunk_idx * reshaped_chunk_size)
+ * stride : (chunk_idx * reshaped_chunk_size + reshaped_chunk_size)
+ * stride,
+ :,
+ ],
+ pad_key_states,
+ stride,
+ chunk_start,
+ chunk_end,
+ )
+ attn_sum, has_rowmax = fused_cosa_proxy(
+ attn_weights_slice,
+ reshaped_block_size,
+ segment_size,
+ chunk_start,
+ chunk_end,
+ real_q_len,
+ scale,
+ )
+ del attn_weights_slice
+
+ sel = _select_chunk(
+ attn_sum,
+ has_rowmax,
+ chunk_idx,
+ npc,
+ offset_u,
+ q_block_num_u,
+ k_block_num_u,
+ topp,
+ )
+ attn_sum_list.append(attn_sum)
+ row_blockmask_list.append(sel)
+
+ attn_sums = torch.cat(attn_sum_list, dim=-2)
+ second = torch.cat(row_blockmask_list, dim=-2)
+ return attn_sums, second
diff --git a/angelslim/compressor/sparsity/algorithms/cosa/prefill.py b/angelslim/compressor/sparsity/algorithms/cosa/prefill.py
new file mode 100644
index 00000000..bd2ce759
--- /dev/null
+++ b/angelslim/compressor/sparsity/algorithms/cosa/prefill.py
@@ -0,0 +1,96 @@
+# 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 prefill: Triton selection proxy + ``hpc`` block-sparse attention.
+
+The proxy emits priority-ordered block ids; the ``hpc`` kernel walks that order
+and threshold-skips the tail. Both ``triton`` and ``hpc`` are hard dependencies.
+Inputs are post-repeat ``(B, H, L, D)``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+
+import torch
+
+from .ops.hpc_adaption import (
+ convert_to_paged_kv,
+ convert_to_varlen_q,
+ convert_varlen_out_to_bshd,
+)
+from .ops.proxy import cosa_estimate
+
+_HPC_PAGE_SIZE = 64
+
+
+def default_chunk_size(k_len: int) -> int:
+ """Pick a proxy chunk size that keeps peak memory roughly constant."""
+ pow2 = 1 << (k_len - 1).bit_length()
+ return int(max(min(max(2048, pow2), 128 * 1024 * 2048 // pow2), 2048))
+
+
+def cosa_prefill(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ *,
+ stride: int,
+ topp: float,
+ block_size: int,
+ chunk_size,
+ skipsoftmax_threshold: float,
+) -> torch.Tensor:
+ """Prefill. ``q,k,v`` and the return are post-repeat ``(B, H, L, D)``."""
+ import hpc
+
+ batch_size, _, k_len, _ = k.shape
+ q_len = q.shape[2]
+
+ assert block_size == 128, f"the hpc kernel is dim128/block128 only, got {block_size}"
+ assert batch_size == 1, f"CoSA prefill is validated at batch_size 1, got {batch_size}"
+
+ if chunk_size is None:
+ chunk_size = default_chunk_size(k_len)
+
+ # Triton launches on the *current* device; pin to q's GPU under device_map.
+ pin = torch.cuda.device(q.device) if q.is_cuda else contextlib.nullcontext()
+ with pin:
+ attn_sums, ordered_blockmask = cosa_estimate(
+ q,
+ k,
+ block_size=block_size,
+ stride=stride,
+ topp=topp,
+ chunk_size=chunk_size,
+ )
+ del attn_sums
+
+ q_thd, cu_seqlens_q, max_seqlens_q = convert_to_varlen_q(q)
+ kcache, vcache, block_ids, seqlens_kvcache = convert_to_paged_kv(
+ k, v, page_size=_HPC_PAGE_SIZE, logical_block_size=block_size
+ )
+ hpc_out = hpc.attention_with_kvcache_blocksparse_anyorderskip_prefill_bf16(
+ q_thd,
+ kcache,
+ vcache,
+ cu_seqlens_q,
+ block_ids,
+ seqlens_kvcache,
+ max_seqlens_q,
+ ordered_blockmask,
+ float(skipsoftmax_threshold),
+ )
+
+ return convert_varlen_out_to_bshd(hpc_out, batch_size, q_len)
diff --git a/configs/sparse/cosa/hy3_cosa.yaml b/configs/sparse/cosa/hy3_cosa.yaml
new file mode 100644
index 00000000..8478b64c
--- /dev/null
+++ b/configs/sparse/cosa/hy3_cosa.yaml
@@ -0,0 +1,31 @@
+# CoSA ordered block-sparse prefill (Hunyuan-V3 / hy_v3).
+# Triton proxy selects key-blocks by topp mass coverage; hpc walks the
+# priority-ordered list and threshold-skips via skipsoftmax_threshold.
+# Hard deps: triton + internal hpc; bf16, head_dim 128, block_size 128, batch=1.
+# NOTE: full Hunyuan-V3 (~80 layers) needs device_map: auto across visible GPUs.
+# Run: python tools/run.py -c configs/sparse/cosa/hy3_cosa.yaml
+
+global:
+ save_path: ./output
+
+model:
+ name: HYV3MoE
+ model_path: tencent/Hy3
+ trust_remote_code: true
+ low_cpu_mem_usage: true
+ use_cache: false
+ torch_dtype: auto
+ device_map: auto
+ attn_implementation: flash_attention_2
+
+compression:
+ name: Sparsity
+ sparsity:
+ name: cosa
+ allow_pseudo_sparse: false # CoSA has no torch reference; true is rejected
+ attn_kwargs:
+ stride: 8 # proxy sub-sample stride
+ topp: 0.95 # mass coverage in (0, 1]; 1.0 = keep all
+ block_size: 128 # hpc is dim128/block128 only
+ chunk_size: 2048 # proxy chunk; multiple of stride*128
+ skipsoftmax_threshold: 1000 # runtime skip; < 0 disables
diff --git a/configs/sparse/cosa/qwen3-8b_cosa.yaml b/configs/sparse/cosa/qwen3-8b_cosa.yaml
new file mode 100644
index 00000000..87f1ad3d
--- /dev/null
+++ b/configs/sparse/cosa/qwen3-8b_cosa.yaml
@@ -0,0 +1,28 @@
+# CoSA ordered block-sparse prefill (Qwen3-8B).
+# Run: python tools/run.py -c configs/sparse/cosa/qwen3-8b_cosa.yaml
+
+global:
+ save_path: ./output
+
+model:
+ name: Qwen
+ model_path: Qwen/Qwen3-8B
+ trust_remote_code: true
+ low_cpu_mem_usage: true
+ use_cache: false
+ torch_dtype: auto
+ device_map: auto
+ # q3 default: decode runs flash-attention-2; sparse prefill bypasses it.
+ attn_implementation: flash_attention_2
+
+compression:
+ name: Sparsity
+ sparsity:
+ name: cosa
+ allow_pseudo_sparse: false # CoSA has no torch reference; true is rejected
+ attn_kwargs:
+ stride: 8 # proxy sub-sample stride
+ topp: 0.95 # mass coverage in (0, 1]; 1.0 = keep all
+ block_size: 128 # hpc is dim128/block128 only
+ chunk_size: 2048 # proxy chunk; multiple of stride*128
+ skipsoftmax_threshold: 1000 # runtime skip; < 0 disables
diff --git a/docs/source/features/sparse_attention/index.md b/docs/source/features/sparse_attention/index.md
index 906e23a0..40e75227 100644
--- a/docs/source/features/sparse_attention/index.md
+++ b/docs/source/features/sparse_attention/index.md
@@ -14,6 +14,7 @@
| XAttention | 实验性 | 内置 kernel（MIT 来源）；head_dim 256 走 torch 参考。 |
| FlashPrefill | 实验性 | clean-room 实现；`alpha<=0` dense 路径需要 `flash_attn`。 |
| VecAttention | 仅参考实现 | 快速 kernel 依赖尚未公开发布的 vLLM-flash-attention fork；外部用户默认走 torch 参考路径（`allow_pseudo_sparse: true`）。 |
+| CoSA | 实验性 | proxy 为内置 Triton；attention 依赖内部 `hpc` 扩展（非公开可安装）。 |
通用约束（所有算法）：不支持 padding mask（请用 batch=1 或无 padding 输入）、部分 kernel 要求 batch=1、prefix-cache / chunked prefill 会自动回退到 dense（稀疏仅在真正首填 `k_len==q_len` 时生效）、不支持 sliding-window 层、不支持多卡 / TP（`WORLD_SIZE>1` 会被拒绝）、稀疏 prefill 不支持 `output_attentions=True`、**不能与量化组合**。`allow_pseudo_sparse` 是正确性 / 调试回退路径，不是性能路径。
diff --git a/tests/sparse/test_cosa_integration.py b/tests/sparse/test_cosa_integration.py
new file mode 100644
index 00000000..f0fd0865
--- /dev/null
+++ b/tests/sparse/test_cosa_integration.py
@@ -0,0 +1,341 @@
+# 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 integration tests.
+
+CoSA has exactly one code path — Triton proxy plus the internal ``hpc`` kernel,
+both hard dependencies — so a forward can only be exercised where ``hpc`` is
+installed. Everything else here (registration, traits, config validation, the
+layout adapters, the patch lifecycle) is pure CPU and runs anywhere torch
+imports; the layout-adapter tests matter most, since a silent shape mistake in
+the varlen/paged conversion would surface as wrong numbers rather than a crash.
+
+Run: python tests/sparse/test_cosa_integration.py
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from _harness import CUDA # noqa: E402
+from _harness import HEAD_DIM # noqa: E402
+from _harness import FakeSlim # noqa: E402
+from _harness import Skip # noqa: E402
+from _harness import SkipReason # noqa: E402
+from _harness import patched # noqa: E402
+from _harness import real_qwen3_8b # noqa: E402
+from _harness import real_weights_available # noqa: E402
+from _harness import record_capability # noqa: E402
+from _harness import rel # noqa: E402
+from _harness import run_all # noqa: E402
+from _harness import tiny_qwen3 # noqa: E402
+
+
+def _make_algo(attn_kwargs=None):
+ import angelslim.compressor.sparsity.algorithms # noqa: F401 register
+ from angelslim.compressor.sparsity.registry import SparsityAlgorithmRegistry
+
+ return SparsityAlgorithmRegistry.create("cosa", attn_kwargs=dict(attn_kwargs or {}))
+
+
+class _StubModel:
+ """Minimal stand-in exposing only the config fields ``setup`` reads."""
+
+ def __init__(self, num_attention_heads=4, num_hidden_layers=4, model_type="qwen3"):
+ self.config = type(
+ "_Cfg",
+ (),
+ {
+ "num_attention_heads": num_attention_heads,
+ "num_hidden_layers": num_hidden_layers,
+ "model_type": model_type,
+ },
+ )()
+
+
+def _hpc_available() -> bool:
+ try:
+ import hpc # noqa: F401
+ except Exception: # noqa: BLE001
+ return False
+ return True
+
+
+# ===========================================================================
+# Registration + traits (no weights)
+# ===========================================================================
+def test_cosa_registered_and_traits():
+ import angelslim.compressor.sparsity.algorithms # noqa: F401
+ from angelslim.compressor.sparsity.registry import SparsityAlgorithmRegistry
+
+ assert "cosa" in SparsityAlgorithmRegistry._factories
+ algo = _make_algo()
+ assert algo.name == "cosa"
+ t = algo.traits
+ assert t.requires_unrepeated_kv is False
+ assert t.needs_calibration is False
+ assert t.supports_padding_mask is False
+ assert t.model_modal == "any"
+ for mt in ("qwen3", "qwen3_moe", "hy_v3"):
+ assert mt in t.compatible_model_types, mt
+
+
+def test_cosa_rejects_incompatible_model_type():
+ from angelslim.compressor.sparsity._base import IncompatibleConfigError
+ from angelslim.compressor.sparsity.patcher import _check_model_type
+
+ algo = _make_algo()
+ _check_model_type(algo, _StubModel(model_type="qwen3")) # supported: no raise
+ try:
+ _check_model_type(algo, _StubModel(model_type="qwen3_5_text"))
+ except IncompatibleConfigError as e:
+ assert "qwen3_5_text" in str(e)
+ else:
+ raise AssertionError("CoSA must reject a model_type outside its compatibility set")
+
+
+# ===========================================================================
+# Config validation
+# ===========================================================================
+def test_cosa_setup_defaults():
+ algo = _make_algo()
+ algo.setup(_StubModel())
+ cfg = algo._cfg
+ assert cfg["stride"] == 8
+ assert cfg["topp"] == 0.8
+ assert cfg["block_size"] == 128
+ assert cfg["chunk_size"] is None # derived per prefill from k_len
+ assert cfg["skipsoftmax_threshold"] == 1000
+ assert "allow_pseudo_sparse" not in cfg, "CoSA has no fallback path to configure"
+
+
+def test_cosa_rejects_allow_pseudo_sparse():
+ """The knob promises a torch fallback CoSA does not have; refuse, never ignore."""
+ from angelslim.compressor.sparsity._base import IncompatibleConfigError
+
+ algo = _make_algo({"allow_pseudo_sparse": True})
+ try:
+ algo.setup(_StubModel())
+ except IncompatibleConfigError as e:
+ assert "allow_pseudo_sparse" in str(e)
+ else:
+ raise AssertionError("allow_pseudo_sparse=true must be rejected, not silently ignored")
+
+
+def test_cosa_setup_rejects_bad_stride():
+ from angelslim.compressor.sparsity._base import IncompatibleConfigError
+
+ algo = _make_algo({"stride": 0})
+ try:
+ algo.setup(_StubModel())
+ except IncompatibleConfigError as e:
+ assert "stride" in str(e)
+ else:
+ raise AssertionError("stride < 1 must be rejected")
+
+
+def test_cosa_topp_must_be_a_scalar():
+ """Only a scalar topp is accepted; per-head tables are rejected by design."""
+ from angelslim.compressor.sparsity._base import IncompatibleConfigError
+
+ algo = _make_algo({"topp": 1}) # an int is fine, and normalizes to float
+ algo.setup(_StubModel())
+ assert algo._cfg["topp"] == 1.0
+ assert isinstance(algo._cfg["topp"], float)
+
+ # A per-head list sized correctly for the model must STILL be refused —
+ # silently accepting one would resurrect the table path.
+ for bad in ([0.8, 0.9, 0.85, 0.95], (0.8, 0.9), "0.8", None, True):
+ wrong = _make_algo({"topp": bad})
+ try:
+ wrong.setup(_StubModel(num_attention_heads=4))
+ except IncompatibleConfigError as e:
+ assert "scalar" in str(e), f"{bad!r} -> {e}"
+ else:
+ raise AssertionError(f"topp={bad!r} must be rejected as a non-scalar")
+
+
+def test_cosa_topp_range_is_checked():
+ from angelslim.compressor.sparsity._base import IncompatibleConfigError
+
+ for bad in (0.0, -0.1, 1.5):
+ algo = _make_algo({"topp": bad})
+ try:
+ algo.setup(_StubModel())
+ except IncompatibleConfigError as e:
+ assert "(0, 1]" in str(e), f"{bad} -> {e}"
+ else:
+ raise AssertionError(f"topp={bad} is outside (0, 1] and must be rejected")
+
+
+def test_cosa_default_chunk_size_tiles_cleanly():
+ from angelslim.compressor.sparsity.algorithms.cosa.prefill import default_chunk_size
+
+ for k_len in (2048, 4096, 16384, 65536, 131072):
+ cs = default_chunk_size(k_len)
+ assert cs >= 2048, f"chunk_size floor violated at k_len={k_len}: {cs}"
+ # The proxy GEMM tiles the reduced matrix at 128; at the documented
+ # stride 8 the chunk must cover whole tiles or the kernel cannot launch.
+ assert cs % (8 * 128) == 0, f"k_len={k_len} -> chunk_size={cs} untileable"
+
+
+# ===========================================================================
+# hpc layout adapters (CPU) — a silent shape bug here means wrong numbers
+# ===========================================================================
+def test_cosa_varlen_q_roundtrip():
+ from angelslim.compressor.sparsity.algorithms.cosa.ops.hpc_adaption import (
+ convert_to_varlen_q,
+ convert_varlen_out_to_bshd,
+ )
+
+ b, h, s, d = 1, 4, 384, HEAD_DIM
+ q = torch.randn(b, h, s, d, dtype=torch.bfloat16)
+ q_thd, cu_seqlens, max_seq = convert_to_varlen_q(q)
+
+ assert q_thd.shape == (b * s, h, d)
+ assert max_seq == s
+ assert cu_seqlens.dtype == torch.int32
+ assert cu_seqlens.tolist() == [0, s]
+ assert q_thd.is_contiguous()
+
+ # The kernel returns (total_tokens, H, D); converting back must be exact.
+ back = convert_varlen_out_to_bshd(q_thd, b, s)
+ assert back.shape == q.shape
+ assert torch.equal(back, q)
+
+
+def test_cosa_paged_kv_layout():
+ from angelslim.compressor.sparsity.algorithms.cosa.ops.hpc_adaption import (
+ convert_to_paged_kv,
+ )
+
+ b, h, s, d = 1, 2, 384, HEAD_DIM
+ page = 64
+ k = torch.randn(b, h, s, d, dtype=torch.bfloat16)
+ v = torch.randn(b, h, s, d, dtype=torch.bfloat16)
+ kcache, vcache, block_ids, seqlens = convert_to_paged_kv(k, v, page_size=page)
+
+ pages = s // page
+ assert kcache.shape == (b * pages, page, h, d)
+ assert vcache.shape == kcache.shape
+ assert block_ids.shape == (b, pages)
+ assert block_ids.dtype == torch.int32
+ assert seqlens.tolist() == [s]
+
+ # Page p must hold tokens [p*page, (p+1)*page) for every head.
+ for p in (0, pages // 2, pages - 1):
+ expect = k[0, :, p * page : (p + 1) * page, :].permute(1, 0, 2)
+ assert torch.equal(kcache[p], expect), f"page {p} holds the wrong tokens"
+
+
+def test_cosa_paged_kv_pads_partial_block():
+ """A k_len that is not a whole number of 128-blocks must be padded up."""
+ from angelslim.compressor.sparsity.algorithms.cosa.ops.hpc_adaption import (
+ convert_to_paged_kv,
+ )
+
+ b, h, d = 1, 2, HEAD_DIM
+ s = 3 * 128 + 40 # 424 -> padded to 512
+ k = torch.randn(b, h, s, d, dtype=torch.bfloat16)
+ kcache, _, block_ids, seqlens = convert_to_paged_kv(k, k.clone(), page_size=64)
+
+ assert kcache.shape[0] == b * (512 // 64)
+ assert block_ids.shape == (b, 512 // 64)
+ # seqlens still reports the REAL length; padding is addressable but not read.
+ assert seqlens.tolist() == [s]
+
+
+def test_cosa_adapters_reject_wrong_head_dim():
+ from angelslim.compressor.sparsity.algorithms.cosa.ops.hpc_adaption import (
+ convert_to_paged_kv,
+ convert_to_varlen_q,
+ )
+
+ bad = torch.randn(1, 2, 256, 64, dtype=torch.bfloat16)
+ for fn in (lambda: convert_to_varlen_q(bad), lambda: convert_to_paged_kv(bad, bad.clone())):
+ try:
+ fn()
+ except ValueError as e:
+ assert "128" in str(e)
+ else:
+ raise AssertionError("head_dim != 128 must be refused by the hpc adapters")
+
+
+# ===========================================================================
+# Patch lifecycle (structure only — no forward, so no kernel needed)
+# ===========================================================================
+def test_cosa_patches_every_layer_and_restores():
+ model = tiny_qwen3(num_layers=4)
+ slim = FakeSlim(model)
+ algo = _make_algo({"topp": 0.8})
+
+ layers = list(model.model.layers)
+ originals = [layer.self_attn.forward for layer in layers]
+ with patched(slim, algo) as handles:
+ assert len(handles) == model.config.num_hidden_layers
+ assert slim.attn_forward_labels() == {"sparse"}
+ for layer in layers:
+ assert layer.self_attn.attn_forward_config["topp"] == 0.8
+ for i, layer in enumerate(layers):
+ assert layer.self_attn.forward == originals[i], "unpatch did not restore the forward"
+ assert slim.attn_forward_labels() == set()
+
+
+# ===========================================================================
+# Numerical correctness (needs real weights AND the internal hpc extension)
+# ===========================================================================
+def test_cosa_keep_all_matches_dense():
+ if not real_weights_available():
+ raise Skip("real Qwen3-8B weights unavailable", SkipReason.NO_QWEN3_8B)
+ if not _hpc_available():
+ raise Skip("internal hpc extension unavailable", SkipReason.NO_BLOCK_SPARSE_ATTN)
+
+ m = real_qwen3_8b()
+ torch.manual_seed(0)
+ ids = torch.randint(0, 151000, (1, 2048), device="cuda")
+ with torch.no_grad():
+ dense = m(ids).logits.float()
+
+ # topp == 1.0 keeps every causal block and skipsoftmax_threshold < 0 disables
+ # the runtime skip, so the two together must reproduce dense attention.
+ slim = FakeSlim(m)
+ algo = _make_algo(
+ {"topp": 1.0, "stride": 8, "chunk_size": 2048, "skipsoftmax_threshold": -1.0}
+ )
+ with patched(slim, algo) as handles:
+ assert len(handles) == m.config.num_hidden_layers
+ with torch.no_grad():
+ sp = m(ids).logits.float()
+ with torch.no_grad():
+ restored = m(ids).logits.float()
+
+ assert not torch.isnan(sp).any()
+ err = rel(sp, dense)
+ assert err < 0.02, f"keep-all CoSA should match dense, rel={err}"
+ assert dense[0, -1].argmax() == sp[0, -1].argmax(), "keep-all argmax drift"
+ assert torch.equal(restored, dense), "unpatch did not restore the real model"
+ record_capability("cosa_real_correctness")
+
+
+if __name__ == "__main__":
+ sys.exit(
+ run_all(
+ dict(globals()),
+ summary_extra=f"cuda={CUDA} hpc={_hpc_available()}",
+ )
+ )