# Rectified Sparse Attention > Enable efficient long-sequence generation by combining block-sparse attention with periodic dense rectification to bound error accumulation. - Skill: `adu2021/rectified-sparse-attention` (Agent Skill) - Install (CLI): `npx skillmds@latest add adu2021/rectified-sparse-attention` - Raw SKILL.md: https://api.skillmd.com/api/skills/adu2021/rectified-sparse-attention/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: Coding & Dev Tools - License: MIT - Author: adu2021 (https://skillmd.com/u/adu2021) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/adu2021/rectified-sparse-attention --- # Rectified Sparse Attention ## Core Concept ReSA solves a fundamental problem in sparse decoding for long-sequence generation: approximation errors accumulate and degrade generation quality as sequence length increases. By periodically "rectifying" the KV cache using dense forward passes, ReSA maintains near-lossless generation quality while achieving up to 2.42× speedup under long-context decoding (256K tokens). ## Architecture Overview - **Group Block Sparse Attention**: Query-dependent sparsity restricting computation to dynamically selected context blocks - **Dense Rectification Phase**: Periodically re-encode recently generated tokens densely to refresh KV cache and bound error accumulation - **Block Descriptors**: Min/max vectors enable efficient retrieval without exhaustive token scanning - **Continuous Batching Integration**: Naturally compatible with existing LLM serving optimizations - **Quality Preservation**: Near-lossless performance on math reasoning, language modeling, and retrieval tasks ## Implementation ### Step 1: Implement Group Block Sparse Attention ```python import torch import torch.nn as nn from typing import Tuple, Optional class BlockDescriptor: """Compact representation of attention blocks""" def __init__(self, block_size: int): self.block_size = block_size def compute_descriptor(self, tokens: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Compute min/max vectors for a block of tokens. These descriptors enable efficient block selection without scanning all tokens. Args: tokens: [seq_len, hidden_dim] - Token representations Returns: min_vec: [num_blocks, hidden_dim] - Minimum values per block max_vec: [num_blocks, hidden_dim] - Maximum values per block """ seq_len = tokens.shape[0] num_blocks = (seq_len + self.block_size - 1) // self.block_size min_vec = torch.full((num_blocks, tokens.shape[1]), float('inf')) max_vec = torch.full((num_blocks, tokens.shape[1]), float('-inf')) for block_idx in range(num_blocks): start = block_idx * self.block_size end = min(start + self.block_size, seq_len) block_tokens = tokens[start:end, :] min_vec[block_idx] = torch.min(block_tokens, dim=0)[0] max_vec[block_idx] = torch.max(block_tokens, dim=0)[0] return min_vec, max_vec class GroupBlockSparseAttention(nn.Module): """Query-dependent sparse attention mechanism""" def __init__(self, hidden_dim: int, num_heads: int, block_size: int = 64, top_k_blocks: int = 8): super().__init__() self.hidden_dim = hidden_dim self.num_heads = num_heads self.block_size = block_size self.top_k_blocks = top_k_blocks self.head_dim = hidden_dim // num_heads self.query_proj = nn.Linear(hidden_dim, hidden_dim) self.key_proj = nn.Linear(hidden_dim, hidden_dim) self.value_proj = nn.Linear(hidden_dim, hidden_dim) self.output_proj = nn.Linear(hidden_dim, hidden_dim) self.block_descriptor = BlockDescriptor(block_size) def select_relevant_blocks(self, query: torch.Tensor, block_descriptors: Tuple[torch.Tensor, torch.Tensor], active_ratio: float = 0.25) -> torch.Tensor: """ Select top-k blocks most relevant to query. Uses descriptors to avoid scanning all blocks. """ min_vec, max_vec = block_descriptors num_blocks = min_vec.shape[0] # Compute relevance score for each block # Query should attend to blocks whose value ranges align with query query_norm = query / (torch.norm(query, p=2, dim=-1, keepdim=True) + 1e-8) block_scores = [] for block_idx in range(num_blocks): # Score: how much does query overlap with block's value range? min_overlap = torch.max(query_norm, min_vec[block_idx]) max_overlap = torch.min(query_norm, max_vec[block_idx]) overlap = torch.clamp(torch.sum(max_overlap - min_overlap), min=0) block_scores.append(overlap) block_scores = torch.stack(block_scores) # Select top-k blocks num_active = max(1, int(num_blocks * active_ratio)) top_k = min(self.top_k_blocks, num_blocks) _, selected_indices = torch.topk(block_scores, k=top_k, largest=True) return selected_indices def forward(self, x: torch.Tensor, kv_cache: Optional[Tuple] = None, kv_block_descriptors: Optional[Tuple] = None, active_ratio: float = 0.25) -> Tuple[torch.Tensor, Tuple]: """ Sparse attention forward pass. Args: x: [batch, seq_len, hidden_dim] kv_cache: Cached key-value states kv_block_descriptors: Block descriptors for efficient retrieval active_ratio: Fraction of blocks to attend to Returns: output: [batch, seq_len, hidden_dim] new_kv_cache: Updated KV cache """ batch_size, seq_len, _ = x.shape # Project Q, K, V query = self.query_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim) key = self.key_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim) value = self.value_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim) # Select relevant blocks for sparse attention if kv_block_descriptors is not None: selected_blocks = self.select_relevant_blocks( query[:, -1, :, :].mean(dim=1), # Use last query kv_block_descriptors, active_ratio=active_ratio ) else: # Fall back to dense attention if no descriptors selected_blocks = torch.arange(seq_len // self.block_size + 1) # Construct sparse attention mask based on selected blocks mask = self.construct_sparse_mask(seq_len, selected_blocks) # Compute attention with sparse mask query = query.transpose(1, 2) # [batch, heads, seq_len, head_dim] key = key.transpose(1, 2) value = value.transpose(1, 2) scores = torch.matmul(query, key.transpose(-2, -1)) / (self.head_dim ** 0.5) scores = scores.masked_fill(~mask, float('-inf')) attn_weights = torch.softmax(scores, dim=-1) attn_weights = attn_weights.masked_fill(~mask, 0) output = torch.matmul(attn_weights, value) # Reshape and project output output = output.transpose(1, 2).contiguous() output = output.view(batch_size, seq_len, self.hidden_dim) output = self.output_proj(output) # Update KV cache and descriptors new_kv = (key, value) if kv_cache is None else (kv_cache[0], kv_cache[1]) new_descriptors = kv_block_descriptors # Update in rectification phase return output, (new_kv, new_descriptors) def construct_sparse_mask(self, seq_len: int, selected_blocks: torch.Tensor) -> torch.Tensor: """Create attention mask for selected blocks""" mask = torch.zeros((seq_len, seq_len), dtype=torch.bool) for block_idx in selected_blocks: start = block_idx * self.block_size end = min(start + self.block_size, seq_len) # Attend to tokens in selected blocks mask[:, start:end] = True # Always attend to recent tokens (current position) mask[:, -1] = True return mask ``` ### Step 2: Implement Dense Rectification ```python class DenseRectificationModule(nn.Module): """Periodically refresh KV cache with dense forward passes""" def __init__(self, model, rectify_interval: int = 16): super().__init__() self.model = model self.rectify_interval = rectify_interval def rectify_kv_cache(self, recent_tokens: torch.Tensor, full_kv_cache: Tuple) -> Tuple: """ Use dense attention to re-encode recent tokens and refresh KV cache. This bounds error accumulation to constant windows. Args: recent_tokens: [batch, recent_length, hidden_dim] full_kv_cache: Current KV cache (may contain errors) Returns: fresh_kv_cache: Re-computed KV cache for recent tokens """ batch_size, recent_len, hidden_dim = recent_tokens.shape # Run dense forward pass on recent tokens only # This is more efficient than re-encoding entire history fresh_key = self.model.key_proj(recent_tokens) fresh_value = self.model.value_proj(recent_tokens) # Replace latest entries in KV cache with fresh computations old_kv = full_kv_cache # Keep old KV for context window, replace recent with fresh context_window = self.rectify_interval fresh_kv = ( torch.cat([ old_kv[0][:, :-context_window, :, :], # Keep old fresh_key.unsqueeze(1) # Replace with fresh ], dim=1), torch.cat([ old_kv[1][:, :-context_window, :, :], # Keep old fresh_value.unsqueeze(1) # Replace with fresh ], dim=1) ) # Also refresh block descriptors for efficient sparse retrieval descriptor_generator = BlockDescriptor(block_size=64) min_vec, max_vec = descriptor_generator.compute_descriptor(fresh_key.squeeze(1)) return fresh_kv, (min_vec, max_vec) class SparseDecodingWithRectification: """Complete sparse generation loop with error rectification""" def __init__(self, model, rectify_interval: int = 16): self.model = model self.sparse_attn = GroupBlockSparseAttention( hidden_dim=model.hidden_dim, num_heads=model.num_heads, block_size=64 ) self.rectifier = DenseRectificationModule(model, rectify_interval) self.rectify_interval = rectify_interval def generate_long_sequence(self, prompt: torch.Tensor, max_length: int = 256000) -> torch.Tensor: """ Generate long sequence with periodic rectification. Maintains quality while achieving 2.42× speedup. """ generated = prompt.clone() kv_cache = None kv_descriptors = None recent_tokens = [] for step in range(max_length): # Get next token prediction using sparse attention next_logits = self.model.forward( generated[:, -1:, :], # Only latest token kv_cache=kv_cache, sparse_attn=self.sparse_attn, kv_descriptors=kv_descriptors, active_ratio=0.25 ) # Sample next token next_token = torch.argmax(next_logits, dim=-1) generated = torch.cat([generated, next_token.unsqueeze(-1)], dim=1) recent_tokens.append(next_token) # Every N tokens: dense rectification if (step + 1) % self.rectify_interval == 0: print(f"Rectification at step {step + 1}") recent_tensor = torch.stack(recent_tokens) kv_cache, kv_descriptors = self.rectifier.rectify_kv_cache( recent_tensor, kv_cache ) recent_tokens = [] # Clear for next window return generated ``` ### Step 3: Integration with Continuous Batching ```python class LongSequenceInferenceEngine: """Production-ready inference with sparse attention and batching""" def __init__(self, model, batch_size: int = 32): self.model = model self.batch_size = batch_size self.pending_requests = [] def add_request(self, prompt_ids: torch.Tensor, max_length: int): """Add generation request to queue""" self.pending_requests.append({ 'prompt_ids': prompt_ids, 'max_length': max_length, 'generated_ids': prompt_ids.clone(), 'step': 0, 'sparse_decoder': SparseDecodingWithRectification(self.model), }) def serve_batch(self): """Process batch of requests with continuous batching""" active_requests = [r for r in self.pending_requests if r['step'] < r['max_length']] # Group into batches respecting max_length constraints for batch_idx in range(0, len(active_requests), self.batch_size): batch = active_requests[batch_idx:batch_idx + self.batch_size] # Stack prompts batch_input = torch.cat([r['generated_ids'][:, -1:] for r in batch], dim=0) # Forward pass with sparse attention batch_logits = self.model.forward( batch_input, use_sparse_attention=True, sparse_active_ratio=0.25 ) # Sample next tokens and update requests batch_next_tokens = torch.argmax(batch_logits, dim=-1) for req_idx, request in enumerate(batch): next_token = batch_next_tokens[req_idx:req_idx+1] request['generated_ids'] = torch.cat( [request['generated_ids'], next_token], dim=1 ) request['step'] += 1 def get_completed_results(self) -> list: """Retrieve finished sequences""" completed = [r for r in self.pending_requests if r['step'] >= r['max_length']] self.pending_requests = [r for r in self.pending_requests if r['step'] < r['max_length']] return completed ``` ## Practical Guidance 1. **Block Size Tuning**: Start with block_size=64. Smaller blocks (32) increase sparsity but may miss context; larger blocks (128) reduce speedup but improve quality. 2. **Active Ratio Selection**: active_ratio=0.25 (attending to 25% of blocks) provides good quality-speedup tradeoff. Increase to 0.5 for critical tasks, decrease to 0.1 for pure language modeling. 3. **Rectification Interval**: Rectify every 16-32 tokens. More frequent rectification preserves quality but reduces speedup. Less frequent rectification (64+) risks error accumulation. 4. **Error Accumulation Bounds**: Dense rectification bounds error to the window size (rectify_interval × block_size tokens). This is vastly better than error growth with sequence length. 5. **Memory Savings**: Block descriptors reduce memory access factor to (1/b + p + 1/f) where b=block_size, p=rectification frequency, f=forward pass factor. With defaults: (1/64 + 1/16 + 1/4) ≈ 0.33× dense attention memory. 6. **Integration Points**: Works seamlessly with continuous batching, multi-GPU inference, and existing KV cache optimizations. ## Reference - Paper: Rectified Sparse Attention (2506.04108) - Key Innovation: Periodic dense rectification bounding error accumulation - Speedup: Up to 2.42× on 256K tokens with near-lossless quality - Architecture: Block sparse + dense rectification phases alternating