"""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