菜单
Accelerate sparse attention by reusing token selection indices across layers. Partition layers into full indexer (F) and shared (S) types using greedy search or multi-layer distillation to eliminate 75% of indexer computation.
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/ADu2021/skillXiv --skill indexcache-sparse-attention-acceleration命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
基于 SOC 职业分类
正在显示 SKILL.md
| name | indexcache-sparse-attention-acceleration |
| title | IndexCache: Accelerating Sparse Attention via Cross-Layer Index Reuse |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2603.12201 |
| keywords | ["Sparse Attention","Acceleration","Long Context","Inference","Indexing"] |
| description | Accelerate sparse attention by reusing token selection indices across layers. Partition layers into full indexer (F) and shared (S) types using greedy search or multi-layer distillation to eliminate 75% of indexer computation. |
Sparse attention mechanisms use indexers to select the most relevant tokens at each layer, but this indexing computation becomes expensive for long contexts. IndexCache observes that token selections are highly correlated across consecutive layers, enabling index reuse: most layers can skip indexing and reuse indices from a nearby layer, dramatically reducing compute.
The approach combines training-free discovery (greedy search) with optional training-aware optimization (multi-layer distillation) to determine optimal sharing patterns.
IndexCache partitions transformer layers into two types:
By carefully choosing which layers retain indexers using either greedy search or learned patterns, the method eliminates 50-75% of indexer operations while maintaining sparse attention quality.
Quantify how similar token selections are across consecutive layers.
import torch
import torch.nn as nn
def measure_index_correlation(model, input_ids, sample_size=1000):
"""
Measure Jaccard similarity of top-k selections across layers.
"""
batch_size, seq_len = input_ids.shape
num_layers = model.config.num_hidden_layers
k = 100 # Assuming top-100 sparse attention
correlations = []
with torch.no_grad():
# Forward pass, collecting indices
indices_by_layer = []
layer_idx (num_layers):
layer_indices = extract_topk_indices(
model,
input_ids,
layer_idx,
k
)
indices_by_layer.append(layer_indices)
layer_idx (num_layers - ):
indices_curr = indices_by_layer[layer_idx]
indices_next = indices_by_layer[layer_idx + ]
jaccard_scores = []
b (batch_size):
pos (seq_len):
set_curr = (indices_curr[b, pos].tolist())
set_next = (indices_next[b, pos].tolist())
intersection = (set_curr & set_next)
union = (set_curr | set_next)
jaccard = intersection / (union + )
jaccard_scores.append(jaccard)
avg_jaccard = (jaccard_scores) / (jaccard_scores)
correlations.append(avg_jaccard)
correlations
Find optimal layer partition using greedy search without training.
def greedy_pattern_search(
model,
calibration_data,
num_layers,
min_f_layers=2,
max_iterations=1000
):
"""
Greedily search for good F/S layer partition using calibration loss.
Returns: binary string where '1'=F (full indexer), '0'=S (shared)
"""
# Start with all F layers
best_pattern = '1' * num_layers
best_loss = evaluate_pattern(model, calibration_data, best_pattern)
for iteration in range(max_iterations):
improved = False
# Try flipping each position from F to S
for pos in range(num_layers):
if best_pattern[pos] == '1':
# Try converting this F layer to S
new_pattern = best_pattern[:pos] + '0' + best_pattern[pos+1:]
# Ensure minimum F layers
if new_pattern.count('1') < min_f_layers:
continue
# Evaluate new pattern
new_loss = evaluate_pattern(model, calibration_data, new_pattern)
if new_loss < best_loss:
best_pattern = new_pattern
best_loss = new_loss
improved = True
break
if not improved:
break # Converged to local optimum
return best_pattern
def evaluate_pattern(model, calibration_data, pattern):
"""
Evaluate how well a layer partition preserves attention quality.
"""
total_loss = 0
with torch.no_grad():
for input_ids, labels in calibration_data:
# Forward with pattern applied
outputs = model(input_ids)
# Compute calibration loss (e.g., perplexity)
loss = compute_loss(outputs, labels)
total_loss += loss.item()
return total_loss / len(calibration_data)
Implement efficient index caching during forward pass.
class IndexCacheManager:
def __init__(self, pattern, num_layers):
"""
pattern: binary string, '1'=F layer, '0'=S layer
"""
self.pattern = pattern
self.num_layers = num_layers
self.cached_indices = None
def get_indices_for_layer(self, layer_idx, indexer_fn):
"""
Get token indices for a layer, using cache if applicable.
indexer_fn: function to compute indices (called only for F layers)
"""
if self.pattern[layer_idx] == '1': # F layer
# Compute fresh indices
indices = indexer_fn()
# Cache for potential sharing
self.cached_indices = indices
return indices
else: # S layer
# Reuse cached indices
if self.cached_indices is None:
raise ValueError("No cached indices available")
return self.cached_indices
def apply_index_cache_to_model(model, pattern):
"""
Modify model to use index caching pattern.
"""
cache_manager = IndexCacheManager(pattern, model.config.num_hidden_layers)
# Patch sparse attention layers
for layer_idx, layer in enumerate(model.layers):
if hasattr(layer, 'attention') and hasattr(layer.attention, 'indexer'):
original_indexer = layer.attention.indexer
def make_indexer_wrapper(idx, mgr, orig_fn):
def indexer_wrapper(query_states):
return mgr.get_indices_for_layer(
idx,
lambda: orig_fn(query_states)
)
return indexer_wrapper
layer.attention.indexer = make_indexer_wrapper(
layer_idx,
cache_manager,
original_indexer
)
return model
Fine-tune retained indexers to better serve shared layers via multi-layer distillation.
class MultiLayerDistillation(nn.Module):
def __init__(self, model, pattern, num_distillation_steps=1000):
super().__init__()
self.model = model
self.pattern = pattern
self.num_distillation_steps = num_distillation_steps
def train_indexers(self, training_data, learning_rate=1e-4):
"""
Train retained indexers against attention distributions of all served layers.
"""
optimizer = torch.optim.Adam(
[p for layer_idx, layer in enumerate(self.model.layers)
if self.pattern[layer_idx] == '1'
for p in layer.attention.indexer.parameters()],
lr=learning_rate
)
for step in range(self.num_distillation_steps):
total_loss = 0
for batch_data in training_data:
input_ids, _ = batch_data
# Forward pass to get attention distributions
with torch.no_grad():
full_outputs = self.model(input_ids)
target_attention = full_outputs.attentions # All layer attentions
# Forward with index cache enabled
outputs = self.model(input_ids)
distilled_attention = outputs.attentions
# KL divergence between distributions
loss = torch.nn.functional.kl_div(
torch.log_softmax(distilled_attention, dim=-1),
torch.softmax(target_attention, dim=-1),
reduction='batchmean'
)
total_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
if step % 100 == 0:
print(f"Step {step}: Loss = {total_loss / len(training_data)}")
return self.model
When to Use:
When NOT to Use:
Hyperparameter Tuning:
Common Pitfalls: