Attribution Graph Optimization for Large Language Models

4.76× faster feature extraction for mechanistic interpretability research on 32B+ parameter models.

Performance

NVIDIA H100 NVL (Primary)

Metric	Baseline	Optimized	Improvement
Per-graph latency	835 ms	174 ms	4.81× faster
Throughput	72 graphs/min	345 graphs/min	4.8× higher
Time for 1000 graphs	13.9 min	2.9 min	Saves 11.0 min

NVIDIA A100 80GB

Metric	Baseline	Optimized	Improvement
Per-graph latency	1034 ms	217 ms	4.76× faster
Throughput	58 graphs/min	276 graphs/min	4.8× higher
Time for 1000 graphs	17.2 min	3.6 min	Saves 13.6 min

GPU Comparison

GPU	Speedup	Latency (ms)	Throughput (graphs/min)
H100 NVL	4.81×	835 → 174	72 → 345
A100 80GB	4.76×	1034 → 217	58 → 276

Model: Qwen2.5-32B (23 layers, 12K features/layer, 50 top-K)
Key insight: Consistent ~4.8× speedup across different NVIDIA GPU architectures

The Problem

Attribution graphs map how interpretable features influence model outputs in large language models. Generating these graphs requires:

1. Forward pass through transcoder networks (16K features × 3584 hidden dims)
2. Top-K selection across sequence positions
3. Sparse graph construction

Bottleneck: Processing 23 layers with Python loops caused excessive CPU-GPU synchronization.

The Solution

Key Optimizations

1. Vectorized Feature Extraction

# Before: Python loop (SLOW)
for pos in range(seq_len):
    acts = transcoder(hidden[pos])
    top_k = torch.topk(acts, k=50)
    # ... build graph nodes

# After: Batched GPU ops (FAST)
acts = transcoder(hidden)  # [B, T, F]
top_vals, top_idx = torch.topk(acts, k=50, dim=2)  # Vectorized
valid_mask = top_vals >= threshold
pos_idx, feat_idx = torch.where(valid_mask)  # Single kernel

Speedup: 835ms → 174ms per graph

2. Memory Layout Optimization

• Contiguous tensor allocation eliminates strided memory access
• Pre-allocated output buffers reduce dynamic allocation overhead

3. Kernel Fusion

• Combined GEMM + ReLU operations
• Fused threshold + compaction via torch.where

Installation

git clone https://github.com/KOKOSde/attribution-graph-optimization.git
cd attribution-graph-optimization
pip install torch transformers

Usage

from optimized_graph_generation import extract_features_optimized

# Your hidden states from model forward pass
hidden_states = {
    layer_idx: hidden  # [batch, seq_len, hidden_dim]
    for layer_idx in range(40, 63)
}

# Optimized extraction
nodes = extract_features_optimized(
    feat_acts=transcoder(hidden_states[layer_idx]),
    layer_idx=layer_idx,
    top_k=50,
    threshold=0.01
)

Benchmark Reproduction

python benchmark_graph_generation.py

Expected output:

Baseline:  835 ms per graph
Optimized: 174 ms per graph  
Speedup:   4.81×

Technical Details

Architecture Support

• LLMs: GPT-2/3, LLaMA, Qwen, Mistral, Phi
• VLMs: Qwen2.5-VL, LLaVA, CLIP
• Constraint: Requires transcoder networks for feature decomposition

GPU Utilization

• Baseline: 23% GPU utilization (CPU-bound by Python loops)
• Optimized: 87% GPU utilization (compute-bound)

Scaling Characteristics

Seq Length	Baseline	Optimized	Speedup
128	418 ms	87 ms	4.80×
256	835 ms	174 ms	4.81×
512	1662 ms	349 ms	4.76×

Why it scales: Consistent 4.8× speedup across sequence lengths shows robust optimization.

Applications

Mechanistic Interpretability

• Circuit discovery: Identify feature pathways for specific behaviors
• Intervention studies: Measure causal effects of feature amplification/suppression
• Safety research: Detect sycophancy, hallucination, or bias circuits

Research Impact

Used to generate 200 attribution graphs for trap-detection study on Qwen2.5-VL-32B, enabling:

• 73% trap detection accuracy (up from 12% baseline)
• Identification of "visual grounding" feature at Layer 25
• Published feature steering methodology

Performance Analysis

Profiling Results

Baseline breakdown (835ms total):
├─ Python loop overhead:     334ms (40%)
├─ CPU→GPU transfers:        242ms (29%)  
├─ GEMM operations:          200ms (24%)
└─ Top-K + compaction:        59ms (7%)

Optimized breakdown (174ms total):
├─ GEMM operations:          125ms (72%)
├─ Top-K + compaction:        38ms (22%)
└─ Graph construction:        11ms (6%)

Key insight: Eliminated 661ms of pure overhead (79% faster).

Implementation Notes

Why Not Custom CUDA Kernels?

cuBLAS and PyTorch's optimized primitives already achieve >85% of theoretical peak performance for these operations. Custom kernels would add complexity with <15% potential gain.

Why Not torch.compile?

torch.compile adds 20-60s compilation overhead per model size. For research workflows with frequent model changes, the amortization point is >1000 graphs.

Hardware Validation

Benchmarked across multiple NVIDIA GPU architectures:

• H100 NVL: 4.81× speedup (835ms → 174ms)
• A100 80GB: 4.76× speedup (1034ms → 217ms)
• Consistent performance: ~4.8× improvement regardless of GPU generation

Production Considerations

For deployment at scale (>10K graphs), consider:

• torch.jit.script for inference (3-8% additional speedup)
• FP16/BF16 precision (2× faster, acceptable for interpretability)
• Multi-GPU batching (linear scaling up to 8 GPUs tested)

Citation

@software{alghanim2025attribution,
  author = {Alghanim, Fahad},
  title = {Attribution Graph Optimization for Large Language Models},
  year = {2025},
  url = {https://github.com/KOKOSde/attribution-graph-optimization}
}

Related Work

• sparse-clt: PyTorch library for efficient Cross-Layer Transcoder inference (GitHub)
• Anthropic Attribution Graphs (2025): Original methodology for feature attribution

License

MIT License - see LICENSE

Author

Fahad Alghanim
Applying to NVIDIA Deep Learning Internship 2026
Focus: GPU optimization for ML interpretability

---

Questions? Open an issue or reach out regarding NVIDIA internship collaboration opportunities.