epfl_final/SparseAtt/sparseatt/sparse_attention.py

"""Stochastic sparse attention module separated from dense attention for prototyping."""

import os
from typing import Optional, Tuple

import torch
import torch.nn.functional as F

try:
    from torch.nn.attention.flex_attention import create_block_mask, flex_attention
except ImportError:
    create_block_mask = None
    flex_attention = None

from .config import Qwen3Config
from .attention import Qwen3Attention, apply_rotary_pos_emb


def sample_rt_row(
    num_blocks: int,
    num_random_blocks: int,
    qb: int,
    device: torch.device,
) -> torch.Tensor:
    r"""Sample ``num_random_blocks`` KV-block indices uniformly from
    ``[0, num_blocks) \ {qb}``.  Causality is enforced separately by the
    mask_mod; both past and future blocks are candidates here.
    """
    perm = torch.randperm(num_blocks - 1, device=device)[:num_random_blocks].clone()
    perm[perm >= qb] += 1
    return perm


def _causal_mask_mod(b, h, q_idx, kv_idx):
    """Simple causal mask stored in BlockMask.mask_mod for within-block token masking.

    The block-level sparsity pattern is already encoded in the BlockMask structure
    (built via create_block_mask + _MaskMod). At flex_attention runtime, mask_mod is
    only called for *partial* blocks — exclusively the local causal diagonal blocks
    where q_block == kv_block. In those blocks is_local is always True so the random-
    block lookup in _MaskMod is dead code. Using a plain causal function avoids the
    aten.index.Tensor gather that Inductor cannot lower inside a pointwise subgraph.
    """
    return q_idx >= kv_idx


class _MaskMod:
    """Picklable mask callable for create_block_mask and BlockMask serialisation.

    Defined at module level so it survives pickle round-trips (e.g. multiprocessing
    queues, torch.save / torch.load across different entry-points).

    ``rt`` holds indices sampled uniformly from all non-diagonal blocks (past *and*
    future).  Any entry with ``kv_block > qb`` is zeroed by the ``q_idx >= kv_idx``
    clause in ``__call__`` and never materialised as an attended block.
    """

    def __init__(self, rt: torch.Tensor, block_size: int, num_random_blocks: int):
        # rt: (num_blocks, num_random_blocks) int16, on any device
        self.rt = rt
        self.block_size = block_size
        self.num_random_blocks = num_random_blocks

    def __call__(self, b, h, q_idx, kv_idx):
        q_block = q_idx // self.block_size
        kv_block = kv_idx // self.block_size
        is_local = q_block == kv_block
        if self.num_random_blocks == 0:
            return is_local & (q_idx >= kv_idx)
        is_global = kv_block == self.rt[q_block, 0]
        for j in range(1, self.num_random_blocks):
            is_global = is_global | (kv_block == self.rt[q_block, j])
        return (is_local | is_global) & (q_idx >= kv_idx)


class Qwen3SparseStochasticAttention(Qwen3Attention):
    """Block-sparse stochastic attention with a prebuilt mask pool."""

    def __init__(self, config: Qwen3Config, layer_idx: int, mask_cache_dir: Optional[str] = None):
        if create_block_mask is None or flex_attention is None:
            raise ImportError(
                "Sparse attention requires torch.nn.attention.flex_attention support."
            )

        self._create_block_mask = create_block_mask
        self._flex_attention = flex_attention

        super().__init__(config, layer_idx)

        # Stash the layer index for per-layer mask seeding (see _build_mask_pool).
        self._mask_layer_idx = int(layer_idx)
        self._mask_seed = int(getattr(config, "sparse_mask_seed", 0))

        self.seq_len = int(config.sparse_max_seq_length)
        self.block_size = int(config.sparse_block_size)
        self.num_blocks = self.seq_len // self.block_size
        self.mask_pool_size = int(config.sparse_mask_pool_size)

        if self.block_size <= 0:
            raise ValueError("sparse_block_size must be > 0")
        if self.block_size % 64 != 0:
            raise ValueError(
                "sparse_block_size must be divisible by 64 for stable torch.compile flex-attention kernels"
            )
        if self.seq_len <= 0:
            raise ValueError("sparse_max_seq_length must be > 0")
        if self.seq_len % self.block_size != 0:
            raise ValueError(
                f"sparse_max_seq_length ({self.seq_len}) must be divisible by sparse_block_size ({self.block_size})"
            )
        if self.num_blocks < 2:
            raise ValueError("Sparse attention requires at least 2 blocks")
        if self.mask_pool_size <= 0:
            raise ValueError("sparse_mask_pool_size must be > 0")

        self.num_random_blocks = max(
            0,
            min(int(config.sparse_num_random_blocks), self.num_blocks - 1),
        )

        # Mixed-sparsity: each pool entry samples its OWN nrb from this set, so masks picked
        # at random per step also resample the sparsity LEVEL. The _rt_pool buffer is sized to
        # the largest nrb; entries with a smaller nrb use only the leading columns.
        raw_mixed = getattr(config, "sparse_mixed_nrb_set", None)
        if raw_mixed:
            self.mixed_nrb_set = sorted({
                max(0, min(int(x), self.num_blocks - 1)) for x in raw_mixed
            })
            self.num_random_blocks = max(self.mixed_nrb_set)
        else:
            self.mixed_nrb_set = None

        # When True, forward() uses standard causal SDPA instead of flex_attention.
        # Flip this flag after loading a sparse checkpoint to continue training in dense mode.
        self.dense_mode: bool = False

        buffer_device = "cuda" if torch.cuda.is_available() else "cpu"
        self.register_buffer(
            "_rt_pool",
            torch.zeros(
                (self.mask_pool_size, self.num_blocks, self.num_random_blocks),
                dtype=torch.int16,
                device=buffer_device,
            ),
            persistent=False,  # not a learned weight — rebuilt each run, never saved to checkpoint
        )

        cache_path = (
            os.path.join(mask_cache_dir, f"masks_layer_{layer_idx}.pt")
            if mask_cache_dir is not None
            else None
        )
        # Masks are always kept on CPU to avoid occupying GPU VRAM for the
        # entire pool (~700 MB/layer × 28 layers = ~20 GB). Only the one mask
        # selected per forward pass is moved to the compute device on-the-fly.
        mask_cpu = torch.device("cpu")
        if cache_path is not None and os.path.isfile(cache_path):
            mask_pool, rt = self._load_masks(cache_path, mask_cpu)
            if len(mask_pool) < self.mask_pool_size:
                # Cached pool is smaller than requested — rebuild.
                self._mask_pool = self._build_mask_pool(mask_cpu)
                os.makedirs(mask_cache_dir, exist_ok=True)
                self.save_masks(cache_path)
            else:
                self._mask_pool = mask_pool[: self.mask_pool_size]
                self._rt_pool.copy_(rt[: self.mask_pool_size].to(self._rt_pool.device, dtype=self._rt_pool.dtype))
        else:
            # Eval / no-cache path: build on the GPU (a BlockMask is small block
            # metadata, not a dense seq^2 tensor). CPU build is ~2.8s/mask, which
            # at pool_size x 28 layers cost minutes; GPU build is ~0.3s/mask.
            build_device = torch.device(buffer_device)
            self._mask_pool = self._build_mask_pool(build_device)
            if cache_path is not None:
                os.makedirs(mask_cache_dir, exist_ok=True)
                self.save_masks(cache_path)

    def _build_mask_pool(self, device: torch.device):
        # Keep mask creation RNG use isolated from model/init RNG streams so that
        # building masks never perturbs weight-init reproducibility.
        cpu_rng_state = torch.random.get_rng_state()
        cuda_rng_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None
        # Seed PER-LAYER (deterministically) so every layer builds its OWN distinct
        # mask pool. Previously this restored the RNG to the *same* state for every
        # layer, which made all layers' masks byte-identical (a latent bug on the
        # in-model / no-cache path: a `pool_size=1` model ended up with one shared
        # mask across all layers instead of one distinct mask per layer).
        torch.manual_seed(1234567 + 1000003 * int(getattr(self, "_mask_layer_idx", 0)) + 2000003 * int(getattr(self, "_mask_seed", 0)))
        pool = self._build_mask_pool_inner(device)
        torch.random.set_rng_state(cpu_rng_state)
        if cuda_rng_state is not None:
            torch.cuda.set_rng_state(cuda_rng_state)
        return pool

    def _build_mask_pool_inner(self, device: torch.device):
        pool = []
        block_size = self.block_size
        num_blocks = self.num_blocks
        mixed_nrb_set = self.mixed_nrb_set

        for i in range(self.mask_pool_size):
            if mixed_nrb_set:
                # Per-entry sparsity level: one nrb drawn from the set for this whole mask.
                nrb_i = int(mixed_nrb_set[torch.randint(len(mixed_nrb_set), (1,)).item()])
            else:
                nrb_i = self.num_random_blocks

            if nrb_i > 0:
                self._rt_pool[i].zero_()
                for qb in range(num_blocks):
                    perm = sample_rt_row(num_blocks, nrb_i, qb, self._rt_pool.device)
                    self._rt_pool[i, qb, :nrb_i] = perm.to(dtype=self._rt_pool.dtype)

            bm = self._create_block_mask(
                _MaskMod(self._rt_pool[i].to(device), block_size, nrb_i),
                B=None,
                H=None,
                Q_LEN=self.seq_len,
                KV_LEN=self.seq_len,
                device=device,
                BLOCK_SIZE=block_size,
            )
            # Replace the stored mask_mod with a simple causal function.
            # _MaskMod is only needed above to build the correct block structure;
            # flex_attention's within-block masking only ever needs q_idx >= kv_idx.
            bm.mask_mod = _causal_mask_mod
            pool.append(bm)

        return pool

    def save_masks(self, path: str) -> None:
        """Persist the mask pool and random-block table to disk.

        Stores only raw index tensors from each BlockMask (no callables) so
        the file is portable across entry-points and Python processes.
        """
        mask_tensors = [
            {
                "seq_lengths":        bm.seq_lengths,
                "kv_num_blocks":      bm.kv_num_blocks.cpu(),
                "kv_indices":         bm.kv_indices.cpu(),
                "full_kv_num_blocks": bm.full_kv_num_blocks.cpu() if bm.full_kv_num_blocks is not None else None,
                "full_kv_indices":    bm.full_kv_indices.cpu() if bm.full_kv_indices is not None else None,
                "q_num_blocks":       bm.q_num_blocks.cpu() if bm.q_num_blocks is not None else None,
                "q_indices":          bm.q_indices.cpu() if bm.q_indices is not None else None,
                "full_q_num_blocks":  bm.full_q_num_blocks.cpu() if bm.full_q_num_blocks is not None else None,
                "full_q_indices":     bm.full_q_indices.cpu() if bm.full_q_indices is not None else None,
                "BLOCK_SIZE":         bm.BLOCK_SIZE,
            }
            for bm in self._mask_pool
        ]
        torch.save({"rt_pool": self._rt_pool.cpu(), "mask_tensors": mask_tensors}, path)

    def _load_masks(self, path: str, device: torch.device):
        """Load a previously saved mask pool from disk."""
        from torch.nn.attention.flex_attention import BlockMask  # type: ignore
        data = torch.load(path, map_location=device, weights_only=False)
        rt_pool = data["rt_pool"].to(device=device, dtype=self._rt_pool.dtype)
        mask_pool = [
            BlockMask(
                seq_lengths=td["seq_lengths"],
                kv_num_blocks=td["kv_num_blocks"].to(device),
                kv_indices=td["kv_indices"].to(device),
                full_kv_num_blocks=td["full_kv_num_blocks"].to(device) if td["full_kv_num_blocks"] is not None else None,
                full_kv_indices=td["full_kv_indices"].to(device) if td["full_kv_indices"] is not None else None,
                q_num_blocks=td["q_num_blocks"].to(device) if td["q_num_blocks"] is not None else None,
                q_indices=td["q_indices"].to(device) if td["q_indices"] is not None else None,
                full_q_num_blocks=td["full_q_num_blocks"].to(device) if td["full_q_num_blocks"] is not None else None,
                full_q_indices=td["full_q_indices"].to(device) if td["full_q_indices"] is not None else None,
                BLOCK_SIZE=td["BLOCK_SIZE"],
                mask_mod=_causal_mask_mod,
            )
            for i, td in enumerate(data["mask_tensors"])
        ]
        return mask_pool, rt_pool

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
        block_mask=None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        batch_size, seq_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        if self.qk_norm == "on":
            query_states = self.q_norm(query_states)
            key_states = self.k_norm(key_states)
        if self.qk_norm == "off" or self._use_query_temperature:
            query_states = query_states * self._query_temperature_scale.to(dtype=query_states.dtype)

        cos, sin = self.rotary_emb(value_states, position_ids)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        key_states = self._repeat_kv(key_states, self.num_key_value_groups)
        value_states = self._repeat_kv(value_states, self.num_key_value_groups)

        if self.dense_mode:
            # Dense fallback: standard causal SDPA — weights are identical to the
            # sparse checkpoint, only the attention pattern changes.
            attn_output = F.scaled_dot_product_attention(
                query_states,
                key_states,
                value_states,
                attn_mask=attention_mask,
                dropout_p=self.attention_dropout if self.training else 0.0,
                is_causal=attention_mask is None,
            )
        else:
            if seq_len != self.seq_len:
                raise ValueError(
                    f"Sparse attention requires fixed seq_len={self.seq_len}, got {seq_len}"
                )
            if attention_mask is not None:
                raise ValueError(
                    "Sparse attention expects packed fixed-length batches and no additive attention mask."
                )
            if past_key_value is not None or use_cache:
                raise ValueError("Sparse generation/KV-cache is not enabled in this iteration")

            if block_mask is None:
                idx = torch.randint(len(self._mask_pool), (1,)).item()
                block_mask = self._mask_pool[idx]
            block_mask = block_mask.to(query_states.device)
            attn_output = self._flex_attention(
                query_states,
                key_states,
                value_states,
                block_mask=block_mask,
                enable_gqa=False,
                kernel_options={"BACKEND": "TRITON"},
            )

        attn_weights = None
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(batch_size, seq_len, self.attention_hidden_size)
        attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, None

    def pick_mask(self, idx: int):
        """Return the BlockMask at position ``idx`` (wraps around pool size)."""
        return self._mask_pool[idx % len(self._mask_pool)]

    # ------------------------------------------------------------------
    # Deploy-time (eval-only) reconfiguration helpers.
    #
    # Everything below is used exclusively by the offline deploy-sparsity
    # evaluator (``eval/sparse_eval.py``). It is never called on the training
    # hot path, so it must not change any behaviour for existing checkpoints.
    # ------------------------------------------------------------------
    def realized_block_density(self) -> float:
        r"""Mean attended-block fraction over the current mask pool.

        For each mask and query block ``qb`` the attended causal KV blocks are
        the diagonal block plus any sampled random block ``r < qb`` (random
        picks with ``r > qb`` are dropped by the ``q_idx >= kv_idx`` clause).
        Density is ``sum(attended) / sum(qb + 1)`` averaged over the pool; a
        fully-dense causal model returns ``1.0``.
        """
        rt = self._rt_pool  # (pool, num_blocks, num_random_blocks) int
        num_blocks = self.num_blocks
        attended = 0
        causal_total = 0
        for qb in range(num_blocks):
            causal_total += (qb + 1) * rt.shape[0]
            if self.num_random_blocks == 0:
                attended += 1 * rt.shape[0]
                continue
            row = rt[:, qb, :]  # (pool, num_random_blocks)
            for i in range(rt.shape[0]):
                picks = {int(v) for v in row[i].tolist() if int(v) < qb}
                attended += 1 + len(picks)  # diagonal + distinct past randoms
        return float(attended) / float(causal_total) if causal_total else 1.0

    def set_deploy_masks(
        self,
        *,
        pattern: str = "random",
        num_random_blocks: Optional[int] = None,
        pool_size: Optional[int] = None,
        seed: Optional[int] = None,
    ) -> None:
        """Rebuild the eval-time mask pool for a chosen DEPLOY pattern.

        ``pattern``:
          - ``"random"``: canonical uniform random-block sampling (same algorithm
            as training) at the requested ``num_random_blocks`` deploy budget.
          - ``"streaming"``: a *training-free* StreamingLLM-style fixed pattern —
            every query block attends to its own (local) block plus the first
            ``num_random_blocks`` sink blocks. Deterministic, so ``pool_size`` is
            forced to 1.
          - ``"sliding"``: a deterministic local-window control where
            ``num_random_blocks`` is the total block budget, including the
            query's diagonal block.

        Only instance state is touched (``num_random_blocks``, ``mask_pool_size``,
        ``_rt_pool``, ``_mask_pool``); weights are never modified.
        """
        if num_random_blocks is not None:
            maximum = self.num_blocks if pattern == "sliding" else self.num_blocks - 1
            minimum = 1 if pattern == "sliding" else 0
            self.num_random_blocks = max(
                minimum, min(int(num_random_blocks), maximum)
            )
        if pattern in {"streaming", "sliding"}:
            pool_size = 1
        if pool_size is None:
            pool_size = self.mask_pool_size
        self.mask_pool_size = int(pool_size)

        nrb = max(self.num_random_blocks, 1)
        rt_device = self._rt_pool.device
        self._rt_pool = torch.zeros(
            (self.mask_pool_size, self.num_blocks, nrb),
            dtype=torch.int16,
            device=rt_device,
        )

        if pattern == "streaming":
            # Sink blocks 0..num_random_blocks-1 for every query block. Above-
            # diagonal sinks (rare, only for the earliest blocks) are dropped by
            # the causal clause at mask-build time.
            for j in range(self.num_random_blocks):
                self._rt_pool[:, :, j] = j
        elif pattern == "sliding":
            past_budget = self.num_random_blocks - 1
            for qb in range(self.num_blocks):
                past_blocks = list(
                    range(max(0, qb - past_budget), qb)
                )
                padded_blocks = past_blocks + [qb] * (
                    self.num_random_blocks - len(past_blocks)
                )
                if padded_blocks:
                    self._rt_pool[:, qb, : self.num_random_blocks] = torch.tensor(
                        padded_blocks,
                        dtype=self._rt_pool.dtype,
                        device=rt_device,
                    )
        elif pattern == "random":
            cpu_rng_state = torch.random.get_rng_state()
            cuda_rng_state = (
                torch.cuda.get_rng_state() if torch.cuda.is_available() else None
            )
            if seed is not None:
                # Seed PER-LAYER so each layer samples its OWN distinct mask pool
                # (matches training, where layers use distinct per-layer masks).
                # Without the per-layer offset, the save/restore below would make
                # every layer's pool identical for a shared seed.
                torch.manual_seed(int(seed) + 1000003 * int(getattr(self, "_mask_layer_idx", 0)))
            if self.num_random_blocks > 0:
                for i in range(self.mask_pool_size):
                    for qb in range(self.num_blocks):
                        perm = sample_rt_row(
                            self.num_blocks, self.num_random_blocks, qb, rt_device
                        )
                        self._rt_pool[i, qb, : self.num_random_blocks] = perm.to(
                            dtype=self._rt_pool.dtype
                        )
            torch.random.set_rng_state(cpu_rng_state)
            if cuda_rng_state is not None:
                torch.cuda.set_rng_state(cuda_rng_state)
        else:
            raise ValueError(f"Unknown deploy pattern '{pattern}'")

        mask_dev = self._rt_pool.device
        pool = []
        for i in range(self.mask_pool_size):
            bm = self._create_block_mask(
                _MaskMod(
                    self._rt_pool[i].to(mask_dev), self.block_size, self.num_random_blocks
                ),
                B=None,
                H=None,
                Q_LEN=self.seq_len,
                KV_LEN=self.seq_len,
                device=mask_dev,
                BLOCK_SIZE=self.block_size,
            )
            bm.mask_mod = _causal_mask_mod
            pool.append(bm)
        self._mask_pool = pool