adu2021 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
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
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
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
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.
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.
Rectification Interval: Rectify every 16-32 tokens. More frequent rectification preserves quality but reduces speedup. Less frequent rectification (64+) risks error accumulation.
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.
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.
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