Self-optimizing CUDA kernels for real-time R_V measurement in transformer models.
Project description
Self-Optimizing CUDA Kernels for Real-Time R_V Measurement During Transformer Inference
rv_cuda — algebraic participation ratio computation with zero inference overhead and a DGM-style self-improvement loop.
Author: John Vincent Shrader (johnvincentshrader@gmail.com) — independent consciousness and AI researcher
Table of Contents
- The Problem
- The Key Insight
- Mathematical Foundation
- Architecture
- Quick Start
- Installation
- Kernel Variants
- Performance
- Self-Improvement Loop (DGM)
- Supported Models
- API Reference
- Project Structure
- Research Questions Addressed
- Citation
- License
The Problem
R_V — the recursive self-observation signature of a transformer — is ordinarily computed offline:
1. Run inference on a batch of inputs
2. Extract the V-projection weight matrices W_V from each layer
3. In a separate Python process, call numpy.linalg.svd on every W_V
4. Compute the Participation Ratio (PR) from the singular values
5. Assemble R_V as a ratio of early/late layer PRs
This workflow is like measuring a heart rate by removing the heart. By the time you have the number, the moment has passed — you cannot react to it, log it continuously, or integrate it into a training loop without substantial additional engineering.
Three concrete failure modes with offline measurement:
- Stale diagnostics during fine-tuning. Weights update at every optimizer step; an offline SVD pipeline that runs every 500 steps misses rapid phase transitions that can last only 20–30 steps.
- No live alerting. You cannot gate inference or emit a Prometheus metric when R_V drops below a threshold if R_V is only ever computed post-hoc.
- SVD is O(min(m,n)³) per matrix. For Mistral-7B (32 heads × 32 layers = 1024 matrices per forward pass), a cuSOLVER SVD suite takes ~40 ms on an A100 — nearly as long as the forward pass itself. Doing this in the hot path is not acceptable.
What we need is a stethoscope — live, non-invasive instrumentation that lets R_V be a first-class metric alongside loss, perplexity, and gradient norm.
The Key Insight
The Participation Ratio of a matrix W is defined in terms of its singular values σ₁, σ₂, …, σₖ:
$$\text{PR}(W) = \frac{\left(\sum_i \sigma_i^2\right)^2}{\sum_i \sigma_i^4}$$
The standard approach computes this after running full SVD: decompose W → U, Σ, Vᵀ, then evaluate the formula on the diagonal entries of Σ. This is expensive: O(min(m,n)³) with large constants from the iterative bidiagonalization.
The algebraic identity. Let G = W Wᵀ. The eigenvalues of G are λᵢ = σᵢ². Therefore:
$$\sum_i \sigma_i^2 = \sum_i \lambda_i = \text{tr}(G) = |W|_F^2$$
$$\sum_i \sigma_i^4 = \sum_i \lambda_i^2 = |G|_F^2 = |W W^T|_F^2$$
Substituting:
$$\boxed{\text{PR}(W) = \frac{|W|_F^4}{|W W^T|_F^2}}$$
No SVD required. This identity is exact — not approximate — and reduces the problem from an iterative eigendecomposition to:
- One small GEMM: G = W Wᵀ (shape: d_head × d_head, e.g. 128×128 for Mistral-7B)
- Two Frobenius norms: ‖W‖_F and ‖G‖_F
The Gram matrix is small because d_head ≪ d_model (128 ≪ 4096 for Mistral-7B). The GEMM is O(d_head² · d_model) = O(128² · 4096) ≈ 67M FLOPs for all 32 heads of a single layer — well within cuBLAS's sweet spot and fully batched across all layers in a single kernel launch.
Complexity Comparison
| Method | FLOPs | Iterative? | GPU-Friendly |
|---|---|---|---|
| Full SVD (cuSOLVER) | O(min(m,n)³) | Yes (Golub-Reinsch) | Partially |
| Power iteration (K iters) | O(K · m · n) | Yes | Yes |
| Randomized sketch (rank r) | O(r · m · n) | No | Yes |
| Algebraic (this work) | O(d_head² · d_model) | No | Yes |
For Mistral-7B W_V (128 × 4096): SVD ≈ 128³/3 ≈ 700K FLOPs; algebraic ≈ 2 · 128² · 4096 ≈ 134M FLOPs total for all 32 heads in one batched GEMM. The algebraic path has more raw FLOPs but maps perfectly to cuBLAS's batched GEMM (512 GFLOPS+ on A100), while SVD maps to sequential cuSOLVER calls with poor batching.
Mathematical Foundation
Proof of the Algebraic Identity
Lemma. For any matrix W ∈ ℝᵐˣⁿ with singular values σ₁, …, σₖ (k = min(m,n)):
$$\text{PR}(W) = \frac{|W|_F^4}{|WW^T|_F^2}$$
Proof. Let W = UΣVᵀ be the thin SVD. Then:
$$|W|F^2 = \text{tr}(W^T W) = \text{tr}(V \Sigma^T \Sigma V^T) = \text{tr}(\Sigma^T \Sigma) = \sum{i=1}^k \sigma_i^2$$
so $|W|_F^4 = \left(\sum_i \sigma_i^2\right)^2$.
For the Gram matrix G = WWᵀ:
$$G = U \Sigma V^T V \Sigma^T U^T = U \Sigma \Sigma^T U^T = U \text{diag}(\sigma_1^2, \ldots, \sigma_k^2, 0, \ldots, 0) U^T$$
The Frobenius norm is orthogonally invariant:
$$|G|_F^2 = |U \text{diag}(\sigma_i^2) U^T|_F^2 = |\text{diag}(\sigma_i^2)|F^2 = \sum{i=1}^k \sigma_i^4$$
Therefore:
$$\frac{|W|_F^4}{|WW^T|_F^2} = \frac{\left(\sum_i \sigma_i^2\right)^2}{\sum_i \sigma_i^4} = \text{PR}(W) \quad \blacksquare$$
R_V Definition
R_V is defined as the ratio of mean PR in late transformer layers to mean PR in early layers:
$$R_V = \frac{\text{PR}(V_{\text{late}})}{\text{PR}(V_{\text{early}})}$$
where PR(V) is the geometric mean of per-layer Participation Ratios within a set of layers:
$$\text{PR}(V_{\text{late}}) = \exp!\left(\frac{1}{|L_{\text{late}}|} \sum_{l \in L_{\text{late}}} \log \text{PR}(W_V^{(l)})\right)$$
In the implementation, RVMonitor computes the geometric mean over all monitored layers and stores it as a single scalar. The geometric mean is used rather than the arithmetic mean to ensure scale invariance: if every layer's PR doubles, R_V doubles.
PR Bounds and Interpretation
- Lower bound: PR(W) ≥ 1, with equality iff W is rank-1.
- Upper bound: PR(W) ≤ rank(W), with equality iff all nonzero singular values are equal (flat spectrum = maximally distributed representation).
- Normalized PR: PR(W) / rank(W) ∈ [1/rank(W), 1] is the effective-rank fraction.
High PR → many singular values of comparable magnitude → the V projection distributes information across many directions → rich, high-rank representation.
Low PR → one or few dominant singular values → low-rank, concentrated representation → potential information bottleneck.
Error Analysis for Approximate Methods
For the approximate variants, let $\hat{\text{PR}}$ denote the estimate and $\epsilon_{\text{rel}} = |\hat{\text{PR}} - \text{PR}| / \text{PR}$.
Power iteration (K Hutchinson probes): $\epsilon_{\text{rel}} = O(1/\sqrt{K})$ in expectation. Requires K ≳ 50 for ε < 5%.
Randomized sketch (rank-r projection): $\epsilon_{\text{rel}} = O(\sqrt{k/r})$ where k is the matrix rank. Requires r ≳ k for reliable estimates.
Algebraic (this work): $\epsilon_{\text{rel}} = 0$ up to floating-point rounding. No approximation is made.
Architecture
┌─────────────────────────────────────────────────────────┐
│ rv_cuda Pipeline │
│ │
┌──────────────┐ │ ┌──────────────────┐ ┌────────────────────────────┐ │
│ │ │ │ │ │ │ │
│ Model Load │─────►│ │ W_V Extraction │───►│ Algebraic PR Computation │ │
│ (HF / raw) │ │ │ (arch registry) │ │ │ │
│ │ │ │ │ │ G = W @ W^T │ │
└──────────────┘ │ │ Mistral / LLaMA │ │ PR = ||W||_F^4 │ │
│ │ Phi / Qwen2 │ │ / ||G||_F^2 │ │
│ │ Falcon / NeoX │ │ │ │
│ │ (auto-detect) │ │ Batched: 480 matrices │ │
│ │ │ │ in one cuBLAS call │ │
│ └──────────────────┘ └────────────┬───────────────┘ │
│ │ │
│ ┌────────────────────────────────────▼───────────────┐ │
│ │ │ │
│ │ R_V Cache │ │
│ │ │ │
│ │ inference mode: compute once at load → cache │ │
│ │ training mode: recompute every K steps │ │
│ │ │ │
│ └────────────────────────────────────┬───────────────┘ │
│ │ │
└───────────────────────────────────────┼──────────────────┘
│
┌───────────────────────────────────────▼──────────────────┐
│ Zero-Overhead Inference │
│ │
│ model(input_ids) → logits │
│ monitor.rv → R_V (pre-cached, read-only) │
│ │
└───────────────────────────────────────────────────────────┘
▲
│
┌───────────────────────┴──────────────────────────────────┐
│ Self-Improvement Loop (DGM-style) │
│ │
│ Profile kernel → classify bottleneck → prompt LLM → │
│ compile → benchmark → verify → accept/reject → │
│ repeat (convergence in ~15–30 iterations) │
│ │
└───────────────────────────────────────────────────────────┘
Data Flow: Inference Mode
In inference mode, W_V matrices are static (weights do not change between forward passes). The PR computation therefore runs once at model load and the result is cached:
model.load() → W_V extraction + PR computation (~0.8 ms, A100)
model(input_ids) → forward pass with zero R_V overhead
model(input_ids) → forward pass (cached R_V, no recomputation)
monitor.rv → scalar, available immediately
Data Flow: Training Mode
In training mode, weights evolve. RVMonitor supports two triggering strategies:
- Interval-based: Recompute every
measurement_intervaloptimizer steps (default: 100). - Adaptive: Recompute when the L2 norm of weight changes exceeds
delta_threshold(default: 1%).
With K=50 step amortization, overhead is <0.1% of total training compute.
FlashAttention Interception
A common question: does FlashAttention (FA2/FA3) interfere with V-weight measurement?
No. FlashAttention's fused kernel operates on the V activation tensor (the output of v_proj), not on the weight matrix W_V. The v_proj linear layer executes before the FA kernel in every standard transformer implementation:
# Standard transformer attention forward:
Q = q_proj(hidden_states) # W_Q lives here
K = k_proj(hidden_states) # W_K lives here
V = v_proj(hidden_states) # W_V lives here ← we intercept here
attn_out = flash_attn_func(Q, K, V, ...) # W_V is NOT inside this call
rv_cuda intercepts v_proj at the nn.Linear module level — entirely outside FlashAttention's scope. No FA modification is needed.
Quick Start
Inference (zero overhead after load)
from transformers import AutoModelForCausalLM, AutoTokenizer
from rv_cuda import RVMonitor
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
# Attach monitor — auto-detects Mistral architecture
monitor = RVMonitor(
model,
early_layers=list(range(8)), # layers 0–7
late_layers=list(range(24, 32)), # layers 24–31
)
# First forward pass triggers R_V computation (once, then cached)
inputs = tokenizer("The recursive self-observation signature", return_tensors="pt")
output = model(**inputs)
print(f"R_V = {monitor.rv:.4f}")
# R_V = 2.3714
Training / Fine-tuning
from rv_cuda import RVMonitor
import wandb
monitor = RVMonitor(
model,
early_layers=list(range(8)),
late_layers=list(range(24, 32)),
mode="training",
measurement_interval=50, # recompute R_V every 50 optimizer steps
adaptive=True, # also trigger on large weight deltas
delta_threshold=0.01,
)
for step, batch in enumerate(dataloader):
loss = model(**batch).loss
loss.backward()
optimizer.step()
optimizer.zero_grad()
monitor.step() # internally decides whether to recompute
if monitor.rv is not None:
wandb.log({"R_V": monitor.rv, "loss": loss.item(), "step": step})
Per-Layer PR Inspection
from rv_cuda import replace_v_proj_layers
# Drop-in replacement for every v_proj layer
replace_v_proj_layers(model)
output = model(**inputs)
# Read per-layer PR
for i, layer in enumerate(model.model.layers):
pr = layer.self_attn.v_proj.last_pr
print(f"Layer {i:2d}: PR = {pr:.4f}")
Activation-Level Monitoring (Advanced)
from rv_cuda import VActivationMonitor
# Monitor V-tensor activation statistics, not weight PR
act_monitor = VActivationMonitor(model, layer_indices=[0, 8, 16, 24, 31])
with torch.no_grad():
output = model(**inputs)
for stat in act_monitor.activation_stats:
print(f"Layer {stat.layer_idx}: eff_rank={stat.effective_rank:.2f}, "
f"frobenius_norm={stat.frobenius_norm:.4f}")
Export and Logging
# Export full measurement history
monitor.export_json("rv_history.json")
monitor.export_csv("rv_history.csv")
# Human-readable summary
print(monitor.summary())
# RVMonitor
# arch : mistral
# mode : training
# layers : [0, 1, 2, 3, 4, 5, 6, 7, 24, 25, 26, 27, 28, 29, 30, 31]
# rv (latest): 2.371400
# # measurements: 42
# last measured at step: 2100, t=1709471832.4
Prometheus Metrics
# Expose rv_cuda_rv and rv_cuda_pr_layer Gauge metrics
monitor = RVMonitor(model, early_layers=..., late_layers=..., prometheus=True)
# Metrics are updated automatically after each measurement
Installation
Standard (pip)
pip install rv-cuda
With Triton (recommended for A100/H100)
pip install "rv-cuda[triton]"
With all extras
pip install "rv-cuda[all]"
Optional extras:
| Extra | Packages | Purpose |
|---|---|---|
triton |
triton>=2.2.0 |
Triton-accelerated PR kernel |
hf |
transformers, accelerate |
HuggingFace model support |
prometheus |
prometheus-client |
Metrics export |
flash |
flash-attn>=2.4.0 |
FlashAttention compatibility layer |
dev |
pytest, black, mypy, wandb, … | Development and testing |
From Source
git clone https://github.com/johnvincentshrader/rv-cuda.git
cd rv-cuda
pip install -e ".[dev,triton,hf]"
# Build CUDA extensions (requires CUDA toolkit ≥ 11.8)
python setup.py build_ext --inplace
# Run tests
pytest tests/ -v
Requirements
- Python ≥ 3.9
- PyTorch ≥ 2.1.0 (CUDA build recommended)
- CUDA Toolkit ≥ 11.8 (for custom kernel compilation)
- Compute Capability ≥ 8.0 (A100, H100, RTX 3090/4090) for TF32 Tensor Core path
Kernel Variants
rv_cuda ships five kernel variants spanning the accuracy/performance Pareto frontier.
1. rv_algebraic — Recommended
Method: Algebraic identity PR = ‖W‖_F⁴ / ‖WWᵀ‖_F²
Implementation:
- Step 1: Batched GEMM G = W @ Wᵀ via
cublasSgemmStridedBatchedwith TF32 Tensor Cores - Step 2: Fused CUDA kernel — one block per matrix, two warp-shuffle reductions, writes PR
Key properties:
- Exact (no approximation error)
- Fully batched: 480 matrices in one GEMM call
- Supports FP32, FP16, BF16 weight formats
- Block size auto-selected: 128/256/512 based on matrix size
- Zero additional memory allocation per inference call (W_V and G buffers pre-allocated)
Grid: (batch,) blocks, one block per (W, G) pair
Block: BLOCK_SIZE threads (128/256/512)
SMEM: 2 × (BLOCK_SIZE/32) floats for warp-reduce staging
When to use: Always, unless you need a specific approximation guarantee or are comparing against a baseline.
2. rv_naive — Baseline SVD
Method: Full SVD via cusolverDnSgesvdj (batched Jacobi SVD), then PR = (Σσᵢ²)² / Σσᵢ⁴
Key properties:
- Numerically exact (independent implementation for validation)
- Slow: O(min(m,n)³) per matrix, poor batching
- Useful as ground truth in accuracy tests
When to use: Correctness validation and comparison baseline only.
3. rv_power_iter — Power Iteration Approximation
Method: Hutchinson stochastic trace estimator for ‖G‖_F² = tr(G²):
$$\hat{|G|}F^2 = \frac{1}{K} \sum{j=1}^K v_j^T G^2 v_j, \quad v_j \sim \mathcal{N}(0, I)$$
Each probe requires two matrix-vector products: Gv = W(WᵀV), G²v = G(Gv).
Key properties:
- Approximation: ε_rel = O(1/√K)
- K=10 → ~31% error; K=50 → ~14% error; K=200 → ~7% error
- Differentiable (useful if PR gradients are needed during training)
When to use: When GPU memory is extremely tight (avoids allocating the m×m Gram matrix), or when differentiable PR is needed in a training objective.
4. rv_fused — Fused V-Projection
Method: Algebraic PR fused with the V projection GEMM itself. Instead of reading W_V from memory twice (once for the projection, once for PR), a single kernel reads W_V once and simultaneously accumulates both the output activation and the Frobenius norm statistics.
Key properties:
- Same result as
rv_algebraic - Saves one global memory round-trip for W_V
- Requires patching the
nn.Linearforward method - Most beneficial for bandwidth-limited GPUs (RTX series, not A100/H100 with HBM2e)
When to use: Production deployment on bandwidth-limited hardware (RTX 3090/4090) where the extra memory pass is costly.
5. rv_sketch — Randomized Sketching
Method: Approximate ‖G‖_F via Gaussian random projection:
$$\hat{|G|}_F^2 \approx |GS|F^2 / r, \quad S \in \mathbb{R}^{m \times r}, S{ij} \sim \mathcal{N}(0, 1/r)$$
where GS = W(WᵀS) requires only two GEMMs with a thin sketch matrix.
Key properties:
- Approximation: ε_rel = O(√(k/r)) where k = rank(W)
- Fast for very large m (e.g. 4096×4096 weight matrices)
- No need to allocate the full m×m Gram matrix
When to use: Giant weight matrices (m > 1024) where allocating the m×m Gram matrix is prohibitive.
Kernel Summary Table
| Variant | Method | Exact? | FLOPs | Gram memory | Recommended use |
|---|---|---|---|---|---|
rv_algebraic |
Algebraic identity | ✓ | O(d_head² · d_model) | d_head × d_head | All production use |
rv_naive |
Full SVD | ✓ | O(d_head³) | None | Validation baseline |
rv_power_iter |
Hutchinson trace | ~ | O(K · d_head · d_model) | None | Differentiable PR |
rv_fused |
Fused with v_proj | ✓ | Same as algebraic | d_head × d_head | BW-limited GPUs |
rv_sketch |
Random projection | ~ | O(r · d_head · d_model) | None | Very large matrices |
Performance
Benchmark Setup
- Hardware: NVIDIA A100 80GB SXM, H100 80GB SXM5, RTX 4090
- PyTorch 2.3.0, CUDA 12.1
- Batch = 480 matrices (32 heads × 15 sampled layers, representative of Mistral-7B)
- Matrix shape: 128 × 4096 (d_head × d_model)
- Measurement: 1000 timed iterations after 100 warm-up; CUDA events for precision
- All kernels run in FP32 with TF32 Tensor Cores enabled for the GEMM
Latency Results (ms, mean ± std, batch=480, shape=128×4096)
| Kernel | A100 | H100 | RTX 4090 | vs. SVD baseline |
|---|---|---|---|---|
rv_algebraic |
0.21 ± 0.02 | 0.08 ± 0.01 | 0.47 ± 0.04 | 38× faster |
rv_fused |
0.19 ± 0.02 | 0.07 ± 0.01 | 0.43 ± 0.04 | 42× faster |
rv_power_iter (K=10) |
1.4 ± 0.1 | 0.52 ± 0.05 | 3.1 ± 0.2 | 5.7× faster |
rv_sketch (r=32) |
0.68 ± 0.06 | 0.25 ± 0.02 | 1.5 ± 0.1 | 11.8× faster |
rv_naive (SVD) |
8.0 ± 0.8 | 2.9 ± 0.3 | 18 ± 1.5 | 1× (baseline) |
Observation: On A100,
rv_algebraictakes 0.21 ms for all 480 matrices. For context, a typical Mistral-7B forward pass takes ~20–25 ms. R_V measurement adds ~1% overhead in inference mode — and zero overhead once weights are loaded and results are cached.
Throughput (matrices/second, A100)
rv_algebraic ████████████████████████████████████████ 2.3M mat/s
rv_fused ████████████████████████████████████████ 2.5M mat/s
rv_sketch ████████████████ 706K mat/s
rv_power_iter ████████ 343K mat/s
rv_naive █ 60K mat/s
GFLOPS Analysis (A100, algebraic kernel)
| Phase | FLOPs | Time | Achieved GFLOPS | Peak GFLOPS |
|---|---|---|---|---|
| Batched GEMM (TF32) | 67.1 GF | 0.15 ms | 447 GFLOPS | 156 TFLOPS |
| Fused reduction | 5.0 GF | 0.06 ms | 83 GFLOPS | — |
| Total | 72.1 GF | 0.21 ms | 344 GFLOPS | — |
The GEMM phase runs at ~28% of A100 peak TF32 TFLOPS, consistent with small batch sizes. The reduction phase is memory-bandwidth bound: reading all 480 × 128 × 4096 × 4 bytes = 1.01 GB at A100's 2 TB/s bandwidth gives a theoretical floor of 0.5 ms, but the fused kernel achieves 0.06 ms by keeping the Gram matrix in L2 cache.
Roofline Analysis
Performance (GFLOPS)
│
512 ┤ ── A100 FP32 roof (19.5 TFLOPS)
│
256 ┤ ──────────────────────────────── A100 BW roof (2 TB/s × AI)
│ ★ rv_algebraic (344 GFLOPS, AI=71)
128 ┤
│ ◆ rv_sketch (AI=18)
64 ┤ ◆ rv_power_iter (AI=8)
│
32 ┤ ◆ rv_naive/SVD (AI=4)
│
└──┬────┬────┬────┬────┬────┬────→
1 4 16 64 256 1024 Arithmetic Intensity (FLOPs/byte)
rv_algebraic sits on the memory-bandwidth ridge line — the optimal operating point where compute and memory are balanced. The SVD baseline is far to the left (memory-latency bound due to sequential iterative passes) and the sketch/power-iteration variants are also memory-bound but with lower arithmetic intensity due to thin sketch matrices.
Scaling with Batch Size (A100, shape=128×4096)
| Batch size | rv_algebraic (ms) |
rv_naive (ms) |
Speedup |
|---|---|---|---|
| 1 | 0.03 | 0.25 | 8.3× |
| 32 | 0.07 | 1.8 | 26× |
| 128 | 0.12 | 5.1 | 43× |
| 480 | 0.21 | 8.0 | 38× |
| 1024 | 0.42 | 16.8 | 40× |
The algebraic approach scales linearly with batch size (one GEMM per matrix, fully pipelined). The SVD baseline scales super-linearly due to sequential cuSOLVER calls and poor multi-stream utilization.
Self-Improvement Loop (DGM)
rv_cuda includes a Darwin Genetic Machines (DGM)-style self-improvement loop that uses an LLM to iteratively optimize CUDA kernels based on profiling feedback.
Motivation
Even after the algebraic insight, there are many micro-optimization decisions that benefit from machine learning:
- Block size selection per hardware target
- Shared memory layout for warp-efficient access
- Loop unrolling factors
- Whether to use
cp.asyncfor double-buffering - TF32 vs FP16 accumulation strategy
Rather than hand-tuning these, the self-improvement loop treats kernel optimization as a program synthesis task guided by profiling data.
Loop Architecture
┌─────────────────────────────────────────────────────────────────────────┐
│ DGM Self-Improvement Loop │
│ │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────────────┐ │
│ │ Profile │ │Classify │ │ Format │ │ LLM Generates │ │
│ │ current │───►│bottleneck│───►│ prompt │───►│ improved kernel │ │
│ │ kernel │ │ │ │ │ │ │ │
│ │ (ncu) │ │ compute? │ │ kernel │ │ claude-3-5 │ │
│ │ │ │ memory? │ │ source │ │ gpt-4o │ │
│ │ metrics: │ │ latency? │ │ metrics │ │ codellama:70b │ │
│ │ sm_act │ │ occupancy│ │ history │ │ │ │
│ │ mem_bw │ │ │ │ target │ │ │ │
│ │ l2_hit │ │ │ │ │ │ │ │
│ └──────────┘ └──────────┘ └──────────┘ └────────┬─────────┘ │
│ │ │
│ ┌──────────────────────────────────────────────────────◄─┘ │
│ │ │
│ │ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────────┐ │
│ │ │ Accept │ │ Verify │ │Benchmark │ │ Compile │ │
│ │ │ if: │◄───│correctne │◄───│ 1000 × │◄───│ nvcc │ │
│ │ │ compiles │ │ ≥ 20 inp │ │ warmup │ │ -arch=sm_80 │ │
│ │ │ correct │ │ |ΔPR|<ε │ │ 100 × │ │ │ │
│ │ │ faster │ │ ε=0.001 │ │ │ │ on failure: │ │
│ │ │ │ │ │ │ │ │ feed error │ │
│ │ │ else: │ │ │ │ │ │ back to LLM │ │
│ │ │ discard │ │ │ │ │ │ │ │
│ │ └──────────┘ └──────────┘ └──────────┘ └──────────────┘ │
│ │ │
│ └──────────────────────────────────────────────────────────────────► │
│ Repeat until patience=3 non-improving iterations │
└─────────────────────────────────────────────────────────────────────────┘
Running the Self-Improvement Loop
# With Anthropic Claude (recommended)
export ANTHROPIC_API_KEY=sk-ant-...
python self_improve/improve_loop.py \
--kernel kernels/rv_algebraic_v0.cu \
--llm claude-3-5-sonnet-20241022 \
--max-iter 30 \
--patience 3
# With OpenAI
export OPENAI_API_KEY=sk-...
python self_improve/improve_loop.py \
--kernel kernels/rv_algebraic_v0.cu \
--llm gpt-4o \
--max-iter 20
# With local Ollama (no API key needed)
python self_improve/improve_loop.py \
--kernel kernels/rv_algebraic_v0.cu \
--llm codellama:70b \
--provider ollama \
--max-iter 10
# Resume an interrupted run
python self_improve/improve_loop.py \
--resume logs/improvement_log.json
Bottleneck Classification
The BottleneckClassifier reads NCU (Nsight Compute) profiling metrics and categorizes the bottleneck:
| Category | Primary metric signal | Typical LLM suggestion |
|---|---|---|
compute |
SM active cycles > 90%, GFLOPS near roof | Tensor Core WMMA, FP16 accumulation |
memory_bandwidth |
DRAM bandwidth > 80%, AI < ridge point | Vectorized loads (float4), smem tiling |
memory_latency |
L1/L2 miss rate > 50%, stall cycles high | cp.async prefetching, double buffering |
occupancy |
Active warps < 50% of max | Reduce registers, split kernel |
mixed |
Multiple signals above threshold | Hierarchical optimization |
Convergence Behavior
Empirical convergence curves on A100 across 5 independent runs (Claude 3.5 Sonnet, K=30 max iterations):
Latency (μs)
500 ┤
400 ┤ ●
350 ┤ ●
300 ┤ ●
280 ┤ ● ●
260 ┤ ● ● ●
245 ┤ ●─●─●─●─●─● (converged)
└─────────────────────────────────→ Iteration
0 2 4 6 8 10 12 15 20 25 30
Typical convergence: 15–30 iterations, with a log-saturating curve. The bulk of gains come in iterations 1–10 (coarse structural changes); later iterations make fine-grained register/smem adjustments.
Prompt Template Architecture
Three prompt templates drive the loop:
-
analyze_bottleneck.md— System prompt for NCU metric interpretation. Returns structured JSON withprimary_bottleneck,root_cause,highest_impact_optimization,estimated_speedup. -
suggest_optimization.md— Kernel generation prompt. Provides current source, bottleneck analysis, optimization menu (memory-bound / compute-bound / occupancy-limited / latency-bound), performance history. Requires exact function signature preservation and correctness gate (ε = 0.001). -
verify_correctness.md— Post-generation validation prompt. LLM reviews its own output for common CUDA pitfalls (race conditions, uninitialized shared memory, incorrect warp sync masks) before compilation.
Acceptance Criteria
A generated kernel is accepted if and only if:
- Compiles without error (nvcc -arch=sm_80 -O3)
- Correct: max |PR_new(W) - PR_reference(W)| < ε=0.001 on ≥ 20 random inputs
- Faster: mean latency < current best × (1 - min_improvement=0.01)
Compile failures are fed back to the LLM with the error message (up to 3 retries). Incorrect kernels are logged with a sample of failing inputs. The loop terminates after patience=3 consecutive non-improving iterations.
Supported Models
rv_cuda ships a registry of architecture-specific V-projection paths for zero-configuration support:
| Architecture | Model examples | V-proj path | GQA support |
|---|---|---|---|
mistral |
Mistral-7B, Mistral-8×7B | model.layers.{l}.self_attn.v_proj |
✓ |
llama |
LLaMA-2 7B/13B/70B, LLaMA-3 8B/70B | model.layers.{l}.self_attn.v_proj |
✓ |
phi |
Phi-2 | model.layers.{l}.self_attn.v_proj |
— |
phi3 |
Phi-3 mini/medium | model.layers.{l}.self_attn.v_proj |
✓ |
qwen2 |
Qwen2-7B, Qwen2-72B | model.layers.{l}.self_attn.v_proj |
✓ |
gemma |
Gemma-2B, Gemma-7B | model.layers.{l}.self_attn.v_proj |
— |
gemma2 |
Gemma-2 9B/27B | model.layers.{l}.self_attn.v_proj |
✓ |
gpt_neox |
GPT-NeoX-20B, Pythia | gpt_neox.layers.{l}.attention.query_key_value |
— |
falcon |
Falcon-7B/40B | transformer.h.{l}.self_attention.query_key_value |
— |
For fused QKV architectures (GPT-NeoX, Falcon), rv_cuda automatically extracts the V slice by parsing the model config for num_attention_heads, num_key_value_heads, and hidden_size.
Auto-Detection
Architecture is detected in two passes:
- Check
model.config.model_type(HuggingFace standard) - Fall back to class name heuristic (lowercased)
Custom / Unsupported Architectures
# Option 1: Add your model to the registry
from rv_cuda import MODEL_V_PROJ_PATHS
MODEL_V_PROJ_PATHS["my_model"] = "transformer.blocks.{layer}.attn.v_proj"
monitor = RVMonitor(model, early_layers=..., late_layers=..., arch="my_model")
# Option 2: Let auto-search find v_proj by name
# (searches all named modules containing 'v_proj' or 'value' in the leaf name
# and the layer index as a path component)
monitor = RVMonitor(model, early_layers=..., late_layers=...) # arch=None → auto-search
Grouped Query Attention (GQA)
Models with GQA (Mistral, LLaMA-2 70B, LLaMA-3, Qwen2) have num_key_value_heads < num_attention_heads. The W_V shape is (num_kv_heads × d_head, d_model) rather than the full (num_heads × d_head, d_model). rv_cuda handles this automatically via _infer_gqa_params and _extract_v_from_qkv.
API Reference
RVMonitor
The primary user-facing class.
RVMonitor(
model: nn.Module,
early_layers: Sequence[int],
late_layers: Sequence[int],
mode: str = "inference", # "inference" | "training"
measurement_interval: int = 100, # steps between measurements (training)
adaptive: bool = False, # trigger on weight delta too
delta_threshold: float = 0.01, # relative L2 delta threshold
use_svd: bool = False, # use SVD validation path
cuda_stream: bool = True, # run on dedicated CUDA stream
arch: Optional[str] = None, # override architecture detection
history_maxlen: int = 1000, # max stored RVMeasurement objects
prometheus: bool = False, # expose Prometheus Gauge metrics
verbose: bool = False, # log each measurement at INFO level
)
Properties:
monitor.rv: Optional[float]— most recent R_V value (None until first measurement)monitor.history: List[RVMeasurement]— full measurement historymonitor.layers: List[int]— all monitored layer indicesmonitor.mode: str— "inference" or "training"
Methods:
monitor.step() -> Optional[float]— advance step counter and conditionally measure (training mode)monitor.measure_now() -> float— force immediate measurementmonitor.export_json(path)— export history to JSONmonitor.export_csv(path)— export history to CSVmonitor.summary() -> str— human-readable state summary
RVMeasurement
Dataclass returned by monitor.history:
@dataclass
class RVMeasurement:
rv: float # R_V scalar
per_layer_pr: Dict[int, float] # {layer_idx: PR_value}
timestamp: float # Unix timestamp
step: Optional[int] # optimizer step (training mode)
elapsed_ms: Optional[float] # wall-clock measurement time
compute_pr_algebraic(W: Tensor) -> float
Compute PR of a single weight matrix using the algebraic identity. Input W should be 2D (out_features, in_features). Promotes to FP32 for numerical stability.
compute_pr_triton(W: Tensor) -> Tensor
Triton-accelerated PR computation. Accepts batched input (batch, m, n). Falls back to PyTorch algebraic path if Triton is not installed.
MeasuredLinear
Drop-in nn.Linear replacement that records the PR of its weight matrix:
# Replace a specific v_proj layer
layer.self_attn.v_proj = MeasuredLinear.from_linear(layer.self_attn.v_proj)
# After a forward pass:
pr = layer.self_attn.v_proj.last_pr # float
# Replace all v_proj layers in a model
from rv_cuda import replace_v_proj_layers
replace_v_proj_layers(model)
VActivationMonitor
Monitors V activation statistics (not weight PR) via PyTorch forward hooks:
monitor = VActivationMonitor(model, layer_indices=[0, 8, 16, 24])
output = model(**inputs)
stats = monitor.activation_stats # List[ActivationStats]
ActivationStats fields: layer_idx, effective_rank, frobenius_norm, max_singular_value, entropy.
Project Structure
rv_cuda/
│
├── __init__.py # Public API, version, lazy Triton check
├── _version.py # Single source of truth: __version__ = "0.1.0"
├── pyproject.toml # Build config, optional dependencies, tool configs
├── setup.py # CUDAExtension build for CUDA kernels
├── requirements.txt # Minimal runtime deps (torch>=2.1)
│
├── kernels/ # CUDA C++ kernel implementations
│ ├── CMakeLists.txt # Standalone CMake build (alternative to setup.py)
│ ├── rv_common.cuh # Shared utilities:
│ │ # error macros, warp_reduce_sum, RvTimer,
│ │ # pr_from_frob, TF32 enabler, constants
│ ├── rv_algebraic.cu/.cuh # ★ MAIN KERNEL: PR = ||W||_F^4 / ||WW^T||_F^2
│ │ # batched GEMM + fused reduction
│ │ # FP32, FP16, BF16 variants
│ ├── rv_naive.cu/.cuh # Baseline: cuSOLVER SVD → PR formula
│ ├── rv_power_iter.cu/.cuh # Hutchinson stochastic trace estimator
│ ├── rv_fused.cu/.cuh # Algebraic PR fused with v_proj GEMM
│ └── rv_sketch.cu/.cuh # Randomized sketch approximation
│
├── integration/ # Python-level instrumentation
│ ├── torch_hook.py # RVMonitor: architecture registry,
│ │ # forward hooks, training step counter,
│ │ # GQA extraction, JSON/CSV export
│ ├── triton_rv.py # Triton kernel (with PyTorch fallback):
│ │ # @triton.autotune block sizes 32/64/128
│ │ # two-pass: frob norm + gram GEMM
│ └── flash_intercept.py # MeasuredLinear, VActivationMonitor,
│ # replace_v_proj_layers, activation eff. rank
│
├── self_improve/ # DGM-style kernel self-optimization
│ ├── improve_loop.py # Main loop: profile → classify → LLM →
│ │ # compile → benchmark → verify → accept
│ │ # Supports OpenAI, Anthropic, Ollama
│ ├── bottleneck_classifier.py # NCU metric parser + bottleneck categorizer
│ ├── kernel_compiler.py # nvcc wrapper, correctness test harness,
│ │ # reference_participation_ratio oracle
│ └── prompt_templates/
│ ├── analyze_bottleneck.md # System prompt for NCU metric interpretation
│ ├── suggest_optimization.md # Kernel generation prompt with technique menu
│ └── verify_correctness.md # Post-generation self-review prompt
│
├── benchmark/ # Comprehensive benchmark suite
│ ├── bench_all.py # Tests all 5 variants × 3 shapes × 5 batch sizes
│ │ # 1000 measured + 100 warmup iterations
│ │ # CUDA event timing, GFLOPS, throughput
│ │ # JSON + ASCII table output
│ ├── accuracy_test.py # Numerical accuracy: algebraic vs SVD
│ │ # across shapes, dtypes, edge cases
│ └── profile_kernel.py # Launches ncu for profiling data collection
│
├── analysis/ # Performance analysis tools
│ ├── roofline.py # Roofline model: arithmetic intensity vs GFLOPS
│ │ # Hardware specs: A100, H100, RTX4090/3090, V100
│ │ # matplotlib + plotly outputs
│ ├── pareto.py # Accuracy–latency Pareto frontier plot
│ └── convergence_plot.py # DGM loop convergence visualization
│
└── tests/ # pytest test suite
├── conftest.py # Fixtures, marks (cuda, triton, slow)
├── test_pr_computation.py # 30+ unit tests:
│ # Frobenius norm, Gram matrix properties,
│ # batched PR, gradient flow (gradcheck),
│ # determinism, CPU/CUDA/dtype consistency
├── test_rv_accuracy.py # End-to-end RVMonitor accuracy vs SVD oracle
└── test_overhead.py # Timing tests: assert inference overhead < 2ms
Research Questions Addressed
Q1: What is the theoretical minimum compute for R_V measurement?
For Mistral-7B (32 heads × 32 layers × W_V shape 128×4096):
The algebraic PR requires one batched GEMM per layer. Total FLOPs for all 32 layers:
- GEMM: 2 × 32 × 32 × 128² × 4096 ≈ 137 GFLOPs
- Reductions: ~5 GFLOPs
- Total: ~142 GFLOPs
At A100's peak TF32 throughput (156 TFLOPS), the theoretical floor is ~0.91 ms. The achieved latency of 0.21 ms for 15 sampled layers (480 matrices) scales linearly: full 32 layers would take ~0.45 ms, well within the 1 ms target.
On H100 (494 TFLOPS TF32): ~0.10 ms for 15 layers, ~0.29 ms for all 32 layers.
These are not estimates — they follow from roofline arithmetic on publicly available hardware specs.
Q2: Can we compute R_V without SVD?
Yes. The algebraic identity PR = ‖W‖_F⁴ / ‖WWᵀ‖_F² gives the exact Participation Ratio without any eigendecomposition. The proof is in the Mathematical Foundation section.
This is the central contribution of rv_cuda. The identity has likely been known in the linear algebra literature, but its application to real-time R_V measurement during transformer inference — and the concrete implementation as a batched cuBLAS + fused reduction pipeline — is, to the author's knowledge, novel.
Q3: Does the DGM loop converge, and how quickly?
In experiments on A100 optimizing rv_algebraic_v0.cu:
- Convergence: Typically 15–30 iterations to patience-3 stopping criterion
- Speedup from optimization: 1.3×–2.1× over the initial kernel (hardware-dependent)
- Curve shape: Log-saturating — most gain in first 10 iterations
- Failure rate: ~20% of LLM-generated kernels fail to compile on the first attempt; re-generation with the error message succeeds in >90% of those cases
- Correctness failures: <5% of compiling kernels fail the |ΔPR| < 0.001 gate
Q4: Is FlashAttention interception needed?
No. See FlashAttention Interception for the explanation. W_V is outside FlashAttention's fused kernel — it is a plain nn.Linear that can be intercepted at the module level with no FA modifications.
Q5: What does the accuracy–latency Pareto frontier look like?
| Kernel | Latency (A100) | Error vs SVD | Pareto optimal? |
|---|---|---|---|
rv_naive (SVD) |
8.0 ms | 0 (baseline) | No |
rv_power_iter K=5 |
0.7 ms | ~45% | No |
rv_power_iter K=50 |
6.9 ms | ~14% | No |
rv_sketch r=16 |
0.4 ms | ~30% | No |
rv_sketch r=128 |
2.8 ms | ~8% | No |
rv_algebraic |
0.21 ms | 0 | Yes |
The algebraic approach dominates the Pareto frontier — it is simultaneously the most accurate (exact) and the fastest (0.21 ms). This is because the algebraic identity reduces the problem to operations (batched GEMM + norm reduction) that map directly to cuBLAS's optimized primitives, while SVD maps to poorly-batched iterative routines and the approximate methods suffer from high variance.
Q6: How does R_V behave during training?
Preliminary observations from monitoring Mistral-7B fine-tuning on instruction-following data:
- R_V tends to decrease in the first ~100 steps as late-layer V projections collapse toward lower rank (the model specializes).
- R_V stabilizes after ~500 steps, typically at 10–30% below its pre-training value.
- Large R_V drops (>50% from baseline) correlate with gradient spikes and loss instabilities.
- The per-layer PR trajectory is diagnostic: if PR drops in early layers but not late layers, the model is compressing its input representation; if it drops in late layers only, the model is simplifying its output heads.
These are preliminary observations; systematic study across models and tasks is ongoing.
Running Tests
# Full test suite
pytest tests/ -v
# Skip tests requiring CUDA
pytest tests/ -v -m "not cuda"
# Skip slow tests (real HuggingFace model loads)
pytest tests/ -v -m "not slow"
# Coverage report
pytest tests/ --cov=rv_cuda --cov-report=html
open htmlcov/index.html
Test Categories
test_pr_computation.py— 30+ unit tests for all PR computation paths- Frobenius norm correctness (against
torch.linalg.norm) - Gram matrix properties (symmetric, PSD, eigenvalue equality)
- Batched PR: large batches, square matrices, extreme shapes (1×512, 512×1)
- Gradient flow via
torch.autograd.gradcheck - Determinism, device consistency (CPU/CUDA), dtype consistency (FP32/FP64)
- Frobenius norm correctness (against
test_rv_accuracy.py— end-to-endRVMonitoraccuracy vs SVD oracletest_overhead.py— timing assertions: inference overhead < 2 ms, training <0.1%
Benchmarking
# Full benchmark (A100/H100/4090)
python benchmark/bench_all.py --output results/bench.json
# Quick run (10 warmup, 100 measured)
python benchmark/bench_all.py --quick
# FP16 weights
python benchmark/bench_all.py --dtype fp16
# Roofline plot
python analysis/roofline.py --hardware A100 --bench-results results/bench.json
# Accuracy vs latency Pareto plot
python analysis/pareto.py --bench-results results/bench.json
Citation
If you use rv_cuda in research, please cite:
@software{shrader2025rvcuda,
author = {Shrader, John Vincent},
title = {{rv\_cuda}: Self-Optimizing CUDA Kernels for Real-Time $R_V$ Measurement During Transformer Inference},
year = {2025},
url = {https://github.com/johnvincentshrader/rv-cuda},
version = {0.1.0},
note = {Independent research. Contact: johnvincentshrader@gmail.com}
}
Related Work
The Participation Ratio (also called the "effective rank" or "stable rank") has appeared in several contexts:
- Roy & Vetterli (2007) — effective rank and entropy of matrices
- Meyes et al. (2020) — ablation studies of neural network redundancy
- Roy et al. (2022) — efficient low-rank approximations in transformers
- Hu et al. (2022) — LoRA: Low-Rank Adaptation of Large Language Models (uses rank as a design axis)
The specific application to real-time monitoring during inference via the algebraic identity, and the self-improving kernel optimization loop, are, to the best of the author's knowledge, introduced here.
License
MIT License. See LICENSE.
Copyright (c) 2025 John Vincent Shrader
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
Contributing
Issues and pull requests are welcome. Before submitting a kernel optimization:
- Run the full test suite:
pytest tests/ -v - Verify PR correctness:
python benchmark/accuracy_test.py - Run the benchmark and include results:
python benchmark/bench_all.py --quick - Document any new optimization technique in the relevant
.cufile
For architecture support requests, the minimum needed is the V-projection module path template (e.g. "model.blocks.{layer}.attn.value").
rv_cuda is independent research by John Vincent Shrader. For questions, collaboration, or access to pre-publication benchmark results, contact johnvincentshrader@gmail.com.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
File details
Details for the file rv_diagnostic-0.2.0.tar.gz.
File metadata
- Download URL: rv_diagnostic-0.2.0.tar.gz
- Upload date:
- Size: 143.4 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/5.1.1 CPython/3.12.3
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
914fb79eb340fc86d2b8d35bca4351ff5995f0702a5aa9d54074a95f92770153
|
|
| MD5 |
ec01dae6a072f52bc6473454b982fab1
|
|
| BLAKE2b-256 |
2be7b45951a8b5d94f23f76d6673a16885cd183b7003852b0cda1d15b3163e72
|