Skip to main content

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.

PyPI version License: MIT Python CUDA PyTorch

Author: John Vincent Shrader (johnvincentshrader@gmail.com) — independent consciousness and AI researcher


Table of Contents

  1. The Problem
  2. The Key Insight
  3. Mathematical Foundation
  4. Architecture
  5. Quick Start
  6. Installation
  7. Kernel Variants
  8. Performance
  9. Self-Improvement Loop (DGM)
  10. Supported Models
  11. API Reference
  12. Project Structure
  13. Research Questions Addressed
  14. Citation
  15. 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:

  1. 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.
  2. 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.
  3. 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:

  1. One small GEMM: G = W Wᵀ (shape: d_head × d_head, e.g. 128×128 for Mistral-7B)
  2. 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:

  1. Interval-based: Recompute every measurement_interval optimizer steps (default: 100).
  2. 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_algebraicRecommended

Method: Algebraic identity PR = ‖W‖_F⁴ / ‖WWᵀ‖_F²

Implementation:

  • Step 1: Batched GEMM G = W @ Wᵀ via cublasSgemmStridedBatched with 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.Linear forward 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_algebraic takes 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.async for 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:

  1. analyze_bottleneck.md — System prompt for NCU metric interpretation. Returns structured JSON with primary_bottleneck, root_cause, highest_impact_optimization, estimated_speedup.

  2. 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).

  3. 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:

  1. Compiles without error (nvcc -arch=sm_80 -O3)
  2. Correct: max |PR_new(W) - PR_reference(W)| < ε=0.001 on ≥ 20 random inputs
  3. 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:

  1. Check model.config.model_type (HuggingFace standard)
  2. 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 history
  • monitor.layers: List[int] — all monitored layer indices
  • monitor.mode: str — "inference" or "training"

Methods:

  • monitor.step() -> Optional[float] — advance step counter and conditionally measure (training mode)
  • monitor.measure_now() -> float — force immediate measurement
  • monitor.export_json(path) — export history to JSON
  • monitor.export_csv(path) — export history to CSV
  • monitor.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)
  • test_rv_accuracy.py — end-to-end RVMonitor accuracy vs SVD oracle
  • test_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:

  1. Run the full test suite: pytest tests/ -v
  2. Verify PR correctness: python benchmark/accuracy_test.py
  3. Run the benchmark and include results: python benchmark/bench_all.py --quick
  4. Document any new optimization technique in the relevant .cu file

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


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

rv_diagnostic-0.2.0.tar.gz (143.4 kB view details)

Uploaded Source

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

Hashes for rv_diagnostic-0.2.0.tar.gz
Algorithm Hash digest
SHA256 914fb79eb340fc86d2b8d35bca4351ff5995f0702a5aa9d54074a95f92770153
MD5 ec01dae6a072f52bc6473454b982fab1
BLAKE2b-256 2be7b45951a8b5d94f23f76d6673a16885cd183b7003852b0cda1d15b3163e72

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page