Skip to main content 홈 크리에이터 adu2021 skillxiv token-sparse-attention-long-context
token-sparse-attention-long-context Dynamically select important tokens at the attention head level, performing dense attention only on selected tokens and scattering results back. Achieves 3.23x attention speedup at 128K context with 1% accuracy loss through layer-wise representation stability analysis.
설치로 이동 Skills Marketplace 커뮤니티가 만든 AI 스킬을 발견하고 탐색하세요.
Codex 또는 Claude로 설치 이 Prompt를 복사해 Codex, Claude 또는 다른 어시스턴트에 붙여 넣으면 Skill 페이지를 검토하고 설치를 진행할 수 있습니다.
직접 명령은 검토 Prompt를 거치지 않습니다. 실행하기 전에 소스를 확인하세요.
npx skills add https://github.com/ADu2021/skillXiv --skill token-sparse-attention-long-context명령은 한 줄로 유지됩니다. 복사하기 전에 가로로 스크롤해 전체 내용을 확인하세요.
로컬 사본을 원하시나요? SkillsMP에서 현재 제공할 수 있는 파일을 다운로드하세요.
Zip 다운로드 다운로드 중... name token-sparse-attention-long-context title Token Sparse Attention: Efficient Long-Context Inference version 0.0.2 engine skillxiv-v0.0.2-claude-opus-4.6 license MIT url https://arxiv.org/abs/2602.03216 keywords ["Sparse Attention","Long-Context","Token Selection","KV Cache Efficiency","Inference Optimization"] description Dynamically select important tokens at the attention head level, performing dense attention only on selected tokens and scattering results back. Achieves 3.23x attention speedup at 128K context with 1% accuracy loss through layer-wise representation stability analysis.
Token Sparse Attention: Head-Wise Selective Attention
Attention computation dominates inference latency for long contexts. Instead of attending to all tokens, Token Sparse Attention selectively attends to the most important tokens identified via a lightweight scoring mechanism. Unlike token eviction methods that permanently remove tokens, this approach preserves tokens through residual connections, allowing layer-wise reconsideration.
The key insight is that different attention heads have different token importance patterns. By allowing each head to independently select its critical tokens and gather/scatter them efficiently, the system maintains representation quality while reducing computation quadratically.
Core Concept
Token Sparse Attention operates as a two-stage process:
Compression : Each attention head scores all tokens, selects the top-K important ones, and performs dense attention on this reduced set.
Decompression : The attention output is scattered back to the original sequence dimension via gather/scatter operations, preserving full sequence shape for subsequent layers.
This compress-decompress design allows tokens to be reconsidered in later layers when their importance changes, avoiding the irreversibility of token eviction.
Architecture Overview
Token Scoring Layer : Lightweight mechanism computing per-head importance scores without full attention costSelection Threshold : Determines K (number of tokens to retain) per headGather Operation : Selects top-K tokens into dense attention blockDense Attention : Standard attention computation on reduced token setScatter Operation : Reconstructs full-sequence attention outputStability Analysis : Identifies layers where representation drift indicates token importance changes
Implementation
Step 1: Implement Token Scoring Mechanism
Create an efficient scoring function that identifies important tokens without full quadratic computation.
class TokenScorer :
def __init__ (self, hidden_dim: int , method: str = "entropy" ):
"""
Score token importance efficiently.
Args:
hidden_dim: Hidden dimension of tokens
method: "entropy" (query-key alignment) or "magnitude"
"""
self
self
if
"learned"
self
1
def
score_tokens
self, queries: torch.Tensor,
keys: torch.Tensor
"""
Compute importance scores for tokens.
Args:
queries: [batch, seq_len, hidden_dim]
keys: [batch, seq_len, hidden_dim]
Returns:
scores: [batch, seq_len] importance scores
"""
if
self
"entropy"
2
1
self
1
sum
2
return
elif
self
"magnitude"
1
1
return
2
elif
self
"learned"
return
self
1
def
select_top_k
self, scores: torch.Tensor,
k: int
Tuple
"""
Select top-K tokens by score.
Args:
scores: [batch, seq_len]
k: Number of tokens to select
Returns:
selected_scores: [batch, k]
selected_indices: [batch, k]
"""
1
return
Step 2: Implement Gather and Scatter Operations Create efficient operations to select and restore tokens.
def gather_tokens (
tokens: torch.Tensor,
indices: torch.Tensor
) -> torch.Tensor:
"""
Gather selected tokens via indexing.
Returns:
gathered: [batch, k, hidden_dim]
"""
batch_size = tokens.shape[0 ]
batch_idx = torch.arange(batch_size, device=tokens.device)[:, None ]
return tokens[batch_idx, indices]
def scatter_attention_output (
attention_output: torch.Tensor,
indices: torch.Tensor,
seq_len: int
) -> torch.Tensor:
"""
Scatter attention output back to original sequence.
Returns:
scattered: [batch, seq_len, hidden_dim]
"""
batch_size = attention_output.shape[0 ]
hidden_dim = attention_output.shape[-1 ]
scattered = torch.zeros(
batch_size, seq_len, hidden_dim,
device=attention_output.device,
dtype=attention_output.dtype
)
batch_idx = torch.arange(batch_size, device=indices.device)[:, None ]
scattered[batch_idx, indices] = attention_output
return scattered
def scatter_add_attention_output (
full_output: torch.Tensor,
attention_output: torch.Tensor,
indices: torch.Tensor
) -> torch.Tensor:
"""
Add scattered attention to existing output (for residual connections).
Returns:
updated: [batch, seq_len, hidden_dim]
"""
batch_size = attention_output.shape[0 ]
batch_idx = torch.arange(batch_size, device=indices.device)[:, None ]
full_output[batch_idx, indices] += attention_output
return full_output
Step 3: Create Sparse Attention Head Combine scoring, gathering, and sparse computation.
class SparseAttentionHead (nn.Module):
def __init__ (self, hidden_dim: int , k: int = 512 ):
"""
Sparse attention operating on top-K tokens.
Args:
hidden_dim: Dimension per head
k: Number of tokens to attend to
"""
super ().__init__()
self .hidden_dim = hidden_dim
self .k = k
self .q_proj = nn.Linear(hidden_dim, hidden_dim)
self .k_proj = nn.Linear(hidden_dim, hidden_dim)
self .v_proj = nn.Linear(hidden_dim, hidden_dim)
self .out_proj = nn.Linear(hidden_dim, hidden_dim)
self .scorer = TokenScorer(hidden_dim, method="entropy" )
def forward (self, hidden_states: torch.Tensor,
attention_mask: Optional [torch.Tensor] = None ) -> torch.Tensor:
"""
Sparse attention forward pass.
Args:
hidden_states: [batch, seq_len, hidden_dim]
attention_mask: Optional mask
Returns:
output: [batch, seq_len, hidden_dim]
"""
batch_size, seq_len, _ = hidden_states.shape
queries = self .q_proj(hidden_states)
keys = self .k_proj(hidden_states)
values = self .v_proj(hidden_states)
scores = self .scorer.score_tokens(queries, keys)
_, indices = self .scorer.select_top_k(
scores,
k=min (self .k, seq_len)
)
selected_queries = gather_tokens(queries, indices)
selected_keys = gather_tokens(keys, indices)
selected_values = gather_tokens(values, indices)
attn_weights = torch.matmul(
selected_queries,
selected_keys.transpose(-2 , -1 )
) / math.sqrt(self .hidden_dim)
attn_weights = torch.softmax(attn_weights, dim=-1 )
attn_output = torch.matmul(attn_weights, selected_values)
attn_output = self .out_proj(attn_output)
output = scatter_attention_output(attn_output, indices, seq_len)
batch_idx = torch.arange(batch_size, device=indices.device)[:, None ]
mask = torch.ones(batch_size, seq_len, dtype=torch.bool ,
device=indices.device)
mask[batch_idx, indices] = False
unselected_output = hidden_states.clone()
unselected_output[~mask] = 0
output = output + unselected_output
return output
Step 4: Analyze Layer Stability for Pruning Identify which layers have stable representations to determine pruning points.
class LayerStabilityAnalyzer :
def __init__ (self, window_size: int = 3 ):
"""Analyze representation drift across layers."""
self .window_size = window_size
self .drift_history = []
def compute_representation_drift (
self,
hidden_before: torch.Tensor,
hidden_after: torch.Tensor
) -> float :
"""
Compute representation change as drift metric.
Uses cosine distance between normalized hidden states.
"""
h_before_norm = torch.nn.functional.normalize(hidden_before, dim=-1 )
h_after_norm = torch.nn.functional.normalize(hidden_after, dim=-1 )
similarity = torch.sum (h_before_norm * h_after_norm, dim=-1 )
drift = 1.0 - similarity.mean().item()
return drift
def is_layer_stable (self, drift: float , threshold: float = 0.1 ) -> bool :
"""Check if layer has stable representations."""
self .drift_history.append(drift)
if len (self .drift_history) < self .window_size:
return False
recent_drift = self .drift_history[-self .window_size:]
avg_drift = sum (recent_drift) / len (recent_drift)
return avg_drift < threshold
def identify_prunable_layers (model: nn.Module,
test_batch: torch.Tensor ) -> List [int ]:
"""
Identify layers where token selection can be aggressive.
Returns:
Layer indices where representation is stable
"""
prunable_layers = []
analyzer = LayerStabilityAnalyzer()
hidden = test_batch
for layer_idx, layer in enumerate (model.layers):
hidden_before = hidden.clone()
hidden = layer(hidden)
drift = analyzer.compute_representation_drift(
hidden_before,
hidden
)
if analyzer.is_layer_stable(drift):
prunable_layers.append(layer_idx)
return prunable_layers
Step 5: End-to-End Inference with Token Sparsity Integrate sparse attention into full inference pipeline.
def sparse_attention_inference (
model: nn.Module,
input_ids: torch.Tensor,
max_new_tokens: int = 128 ,
sparsity_k: int = 512
) -> torch.Tensor:
"""
Run inference with token sparse attention.
Args:
model: Language model
input_ids: [batch, seq_len]
max_new_tokens: Tokens to generate
sparsity_k: Number of tokens per head
Returns:
generated_ids: [batch, seq_len + max_new_tokens]
"""
prunable_layers = identify_prunable_layers(model, input_ids)
for layer_idx in prunable_layers:
if hasattr (model.layers[layer_idx].self_attn, 'to_sparse' ):
model.layers[layer_idx].self_attn.to_sparse(k=sparsity_k)
with torch.no_grad():
output_ids = model.generate(
input_ids,
max_new_tokens=max_new_tokens,
use_cache=True
)
return output_ids
Practical Guidance When to use Token Sparse Attention:
Long-context inference (32K+ tokens) where attention is bottleneck
Scenarios where <1% accuracy loss is acceptable for 3x speedup
Systems with GPU/TPU supporting efficient gather/scatter operations
Batch inference where latency matters more than throughput
Short contexts (<4K tokens) where sparsity overhead dominates benefit
Tasks requiring precise token interaction (e.g., code analysis)
Systems without efficient sparse operations (CPUs)
Real-time applications needing predictable latency
K too small: Prunes important context, causing accuracy loss
K too large: Eliminates sparsity benefits
Instability in early layers: May prune tokens needed by later layers
Residual connection bugs: Unselected tokens must be preserved
Hyperparameter Guidelines:
Parameter Recommended Tuning Strategy sparsity_k 512-1024 Higher = less sparsity; tune on accuracy-speed tradeoff scoring_method entropy Learned scoring if accuracy critical; entropy for speed stability_threshold 0.05-0.1 Higher = more layers eligible for sparsity prunable_layers Last 50% Compress representation is stable in later layers
Reference Key results: 3.23× attention speedup at 128K context with <1% accuracy degradation. Triton-based kernel implementation; FlashAttention compatible. Code available on GitHub.
