docs: GPU perf batch 1 design spec (FA3, FP8, CUDA Graphs, SnapKV)

Research-driven design for 4 high-impact optimizations targeting
Blackwell SM_120 (RTX 5060 Ti):
- Item 1: CUDA Graphs wiring (90% exists, half-day task)
- Item 2: FP8 E4M3/E5M2 training (30-40% throughput, 50% memory)
- Item 3: FlashAttention-3 WGMMA+TMA+warp specialization (1.5-2x)
- Item 4: SnapKV + prefix caching (50-70% KV reduction)

Based on: arXiv:2407.08608 (FA3), arXiv:2511.05811 (MOSS FP8),
arXiv:2404.14469 (SnapKV), PyTorch 2025 state survey.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>

docs/superpowers/specs/2026-06-26-gpu-perf-batch1-design.md

@@ -0,0 +1,340 @@
+# GPU Performance Optimization Batch 1 — Design Spec
+
+**Date**: 2026-06-26  
+**Hardware target**: NVIDIA RTX 5060 Ti (Blackwell SM_120, 16 GB VRAM)  
+**Research basis**: FlashAttention-3 (arXiv:2407.08608), MOSS FP8 (arXiv:2511.05811), SnapKV (arXiv:2404.14469), KeepKV (arXiv:2504.09936), PyTorch 2025 state survey
+
+---
+
+## Summary
+
+Four high-impact GPU performance optimizations, selected after cross-referencing the current rustytorch codebase against PyTorch 2025 and 2024–2025 arxiv papers. Items are ordered by implementation cost: CUDA Graphs is a half-day wiring task (infrastructure already exists), FP8 training is ~1 week, SnapKV/prefix caching is ~1 week, and FlashAttention-3 is the deepest CUDA kernel work at ~2 weeks.
+
+| Item | Expected gain | Complexity | Primary files |
+|------|--------------|-----------|---------------|
+| 1. CUDA Graphs | 5–20% throughput | Low (wiring only) | `training_loop.rs`, `cuda_graph.rs` |
+| 2. FP8 Training | 30–40% throughput, 50% memory | Medium | `fp8_cast.cu`, `fp8_gemm.rs`, `linear.rs` |
+| 3. FlashAttention-3 | 1.5–2× attention throughput | High | `flash_attention_v3_*.cu`, `backend_selector.rs` |
+| 4. SnapKV + Prefix Cache | 50–70% KV reduction | Medium | `eviction.rs`, `paged_kv_cache.rs` |
+
+---
+
+## Item 1: CUDA Graphs
+
+### Problem
+Every GPU kernel launch carries 5–20 μs of CPU-side overhead (CUDA API call → driver → scheduler). A typical transformer training step launches hundreds of kernels. At small batch sizes common on consumer GPUs, this overhead is measurable.
+
+### Solution
+CUDA Graphs capture the full sequence of GPU operations in a graph, then replay it in a single `cudaGraphLaunch` — eliminating per-kernel CPU overhead on all subsequent steps.
+
+### Current state
+`CudaGraphManager` at `crates/core/rtx-runtime/src/cuda_graph.rs` (445 lines) is fully implemented with `begin_capture()`, `end_capture()`, `launch()`, and thread-safety via `RwLock`. It is **not connected to the training loop**.
+
+### Design
+
+**`crates/training/rtx-transformers/src/training/config.rs`** — add two fields:
+```rust
+pub enable_cuda_graphs: bool,        // default: false
+pub cuda_graph_warmup_iters: usize,  // default: 3
+```
+
+**`crates/training/rtx-transformers/src/training/training_loop.rs`** (around line 167) — wrap the training step:
+- Iterations 0..warmup: execute normally (warms up allocator, avoids capturing allocations in graph)
+- Iteration `warmup`: `begin_capture(stream)` → `execute_step()` → `end_capture()` → store `graph_id`
+- Iterations `warmup+1..`: `graph_manager.launch(graph_id, stream)` — no per-kernel CPU calls
+
+**Constraint**: tensor shapes must be static across steps. Add a `requires_static_shapes()` check that disables graph mode and emits a warning when dynamic padding is detected (variable sequence lengths).
+
+**`TrainingLoop`** gets a new `graph_state: Option<GraphId>` field and a reference to `CudaGraphManager`.
+
+### Verification
+```
+cargo test -p rtx-transformers -- test_cuda_graph_training_step
+```
+Test: capture a 2-layer transformer step, verify output tensors match non-graph execution within 1e-6, measure step time reduction.
+
+---
+
+## Item 2: FP8 Training (E4M3 forward / E5M2 gradients)
+
+### Problem
+BF16 mixed-precision training uses 2 bytes per parameter activation. Blackwell's 5th-generation tensor cores support FP8 (1 byte), but the connection between rustytorch's existing FP8 type infrastructure and actual GEMM execution is missing.
+
+### Current state
+- `DType::FP8E4M3` / `FP8E5M2` defined in `dtype.rs` with correct `size_bytes = 1`
+- `Fp8GradScaler` in `fp8_scaler.rs` — full amax-tracking, scale growth/backoff, overflow detection
+- `CudaDataType::R8F_E4M3 = 28` / `R8F_E5M2 = 29` in `cusparelt/types.rs`
+- MX block-wise quantization kernels in `microscaling.rs`
+- **Missing**: FP8 casting CUDA kernels, cuBLASLt GEMM dispatch, training loop wiring
+
+### Design
+
+**Step 1 — Casting kernels** (new file: `crates/core/rtx-tensor/src/cuda_kernels/fp8_cast.cu`):
+```cuda
+// Saturating cast: clamp to [-448, 448] before encoding E4M3
+__global__ void cast_f32_to_fp8_e4m3(
+    const float* __restrict__ src, uint8_t* __restrict__ dst,
+    float scale, int n);
+
+__global__ void cast_bf16_to_fp8_e4m3(
+    const __nv_bfloat16* __restrict__ src, uint8_t* __restrict__ dst,
+    float scale, int n);
+
+__global__ void cast_fp8_e4m3_to_bf16(
+    const uint8_t* __restrict__ src, __nv_bfloat16* __restrict__ dst,
+    float inv_scale, int n);
+
+// E5M2 range ±57344, used for gradient tensors
+__global__ void cast_fp8_e5m2_to_bf16(
+    const uint8_t* __restrict__ src, __nv_bfloat16* __restrict__ dst,
+    float inv_scale, int n);
+```
+Follow `element_wise.cu` pattern for `build.rs` registration and Rust FFI wrapper.
+
+**Step 2 — cuBLASLt GEMM dispatcher** (new file: `crates/core/rtx-tensor/src/fp8_gemm.rs`):
+```rust
+pub fn fp8_matmul(
+    a: &CudaSlice<u8>,  // E4M3, shape [M, K]
+    b: &CudaSlice<u8>,  // E4M3, shape [K, N]  
+    scale_a: f32,
+    scale_b: f32,
+    m: usize, n: usize, k: usize,
+    stream: &CudaStream,
+) -> Result<CudaSlice<u8>>  // BF16 output, shape [M, N]
+```
+Uses `cublasLtMatmul` with `CUBLAS_COMPUTE_32F` accumulator and `CUDA_R_8F_E4M3` A/B types. Add `cargo:rustc-link-lib=cublasLt` to `rtx-tensor/build.rs`.
+
+**Step 3 — Linear layer FP8 path** (`crates/training/rtx-transformers/src/layers/linear.rs`):
+- Add `fp8_mode: bool` to `LinearConfig` (default false)
+- In `forward()` when `fp8_mode`:
+  1. Cast weight matrix W from BF16 → E4M3 using `cast_bf16_to_fp8_e4m3` (can be pre-cached per step)
+  2. Cast activation X from BF16 → E4M3
+  3. Run `fp8_matmul(W_fp8, X_fp8, scale_w, scale_x, ...)`
+  4. Cast output Y from internal FP32 accumulator → BF16
+- Master weights remain BF16/FP32 for optimizer stability
+
+**Step 4 — Training config** (`config.rs`):
+```rust
+pub fp8_training: bool,           // default: false
+pub fp8_format: Fp8Format,        // default: E4M3Fwd_E5M2Grad
+```
+
+**Mixed precision recipe** (matching NVIDIA MOSS / Transformer Engine):
+- Forward pass: weights and activations cast to E4M3
+- Gradient tensors: cast to E5M2 (wider range ±57344 handles gradient spikes)
+- Optimizer states: remain FP32 (master weights)
+- Scaling: `Fp8GradScaler` (existing) manages per-tensor amax and scale factors
+
+### Verification
+```
+cargo test -p rtx-transformers -- test_fp8_forward_backward
+```
+Test: 2-layer linear model, FP8 forward output matches BF16 within 1e-2 tolerance, backward gradients are non-NaN and converge in 10 steps.
+
+---
+
+## Item 3: FlashAttention-3 Blackwell
+
+### Problem
+Current kernels use WMMA (4th-generation tensor core API, 16×16×16 fragments) and `cp.async` for shared memory loads. Blackwell's SM_120 introduces Warpgroup MMA (WGMMA, 64×16×16 fragments — 4× compute density) and TMA (Tensor Memory Accelerator — asynchronous 2D tile loads with no warp involvement). FlashAttention-3 (Dao et al., arXiv:2407.08608) showed 1.5–2× speedup on H100 by exploiting these; the same techniques apply to SM_120.
+
+### Current state
+- `flash_attention_forward.cu`: WMMA + tiled online softmax + cooperative groups
+- `utils.cu`: `wmma::fragment<matrix_a/b/accumulator>` with BF16, `cg::memcpy_async`, 163KB shared memory
+- `build.rs`: sm_120, PTX 8.0, `--use_fast_math --maxrregcount=255`
+- `FlashCudaKernels::new()` in `simple.rs`: drop-in kernel registration
+
+### Design
+
+**New file: `src/kernels/flash_attention_v3_forward.cu`**
+
+Three algorithmic changes from FA2:
+
+**1. WGMMA (Warpgroup Matrix Multiply Accumulate)**
+
+Replace:
+```cuda
+wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag);  // 16×16×16
+```
+With PTX inline asm for warpgroup-level operation (128 threads = 4 warps):
+```cuda
+asm volatile(
+  "wgmma.mma_async.sync.aligned.m64n128k16.f32.bf16.bf16 "
+  "{%0,%1,...,%7}, [%8], {%9,%10,...}, 1, 1, 1, 0, 0;"
+  : "+f"(d[0]), "+f"(d[1]), ...
+  : "l"(desc_a), "r"(b[0]), "r"(b[1]), ...
+);
+```
+Tile size: 64×128×16 per warpgroup vs 16×16×16 per warp. ~4× compute per instruction, amortizing instruction issue overhead.
+
+**2. TMA (Tensor Memory Accelerator)**
+
+Replace `cg::memcpy_async` with `cp.async.bulk.tensor.2d`:
+```cuda
+// One instruction loads an entire 2D tile — no per-element loop, no warp threads consumed
+asm volatile(
+  "cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier::complete_tx::bytes"
+  " [%0], [%1, {%2, %3}], [%4];"
+  :: "r"(smem_ptr), "l"(gmem_base), "r"(row), "r"(col), "r"(mbar_ptr)
+);
+```
+TMA runs in the memory subsystem — warp threads are free to compute while load is in flight.
+
+**3. Warp Specialization (Producer/Consumer split)**
+
+Split the 4 warpgroups per block:
+- **Producer warpgroup** (warpgroup 0): issues TMA loads for Q/K/V tiles, waits on TMA barrier, signals consumer via shared memory barrier
+- **Consumer warpgroups** (warpgroups 1–3): wait on producer signal, execute WGMMA accumulation, write output
+- Overlap: consumer processes tile `i` while producer loads tile `i+1` into the opposite shared-memory ping-pong buffer
+
+This eliminates the memory stall that currently serializes compute and memory in FA2.
+
+**Optional FP8 path** (gated by `use_fp8_attn: bool`):
+Before WGMMA, quantize each 16×16 Q/K block to E4M3 with a per-block scale factor. Enables ~2× additional throughput on FP8 tensor cores when combined with Item 2.
+
+**New file: `src/kernels/flash_attention_v3_backward.cu`**
+
+Same WGMMA + TMA + warp specialization for dQ/dK/dV, following the FA3 paper's recomputation approach (recompute softmax from stored LSE values to avoid storing the full N×N matrix).
+
+**`src/backend_selector.rs`** — runtime SM selection:
+```rust
+fn select_flash_attention_kernel(sm_major: u32, sm_minor: u32) -> FlashKernelVariant {
+    if sm_major >= 12 {
+        FlashKernelVariant::V3Wgmma   // Blackwell SM_120+
+    } else if sm_major >= 9 {
+        FlashKernelVariant::V3Wgmma   // Hopper SM_90
+    } else {
+        FlashKernelVariant::V2Wmma    // Ampere and below
+    }
+}
+```
+Both kernel sets remain registered; no API change to callers.
+
+### Verification
+```
+cargo test -p rtx-flash-attention -- test_v3_correctness
+cargo bench -p rtx-flash-attention
+```
+Correctness: FA3 output matches naive O(N²) attention within 1e-3 for sequence lengths 512, 1024, 4096.  
+Performance: FA3 ≥ 1.5× throughput vs FA2 on SM_120 for seqlen ≥ 1024.
+
+---
+
+## Item 4: SnapKV + Prefix Caching
+
+### Problem
+On 16 GB VRAM, KV cache is the primary constraint on batch size and context length during inference. Two complementary techniques address this:
+- **SnapKV** (Zhang et al., arXiv:2404.14469): prune low-attention keys during prefill, keeping 60% + recent window
+- **Prefix caching**: reuse KV pages for sequences sharing a common prefix (system prompt, few-shot examples)
+
+### Current state
+- `PagedKvCache` in `paged_kv_cache.rs`: 3-tier storage (GPU/CPU/NVMe), UUID page IDs, copy-on-write, LRU eviction
+- `eviction.rs`: LRU/LFU/FIFO/Random policies — no attention-score policy
+- `config.rs`: `enable_prefix_sharing: bool = false` — no implementation behind the flag
+
+### Design
+
+**SnapKV — attention-score eviction** (`eviction.rs`):
+
+Add new policy variant:
+```rust
+pub enum EvictionPolicy {
+    Lru, Lfu, Fifo, Random,
+    AttentionScore { keep_ratio: f32, recent_window: usize },  // NEW
+}
+```
+
+Add sidecar to `PagedKvCache`:
+```rust
+attention_scores: Arc<RwLock<HashMap<PageId, f32>>>,
+```
+
+During prefill, the flash attention kernel accumulates the sum of attention weights per key position into a host-accessible buffer. After prefill, scores are written to `attention_scores` via a kernel that reduces along the head dimension.
+
+In `select_victims()`:
+```rust
+// Sort candidate pages by accumulated attention score ascending
+// Always retain last `recent_window` tokens (recency bias)
+// Evict bottom (1 - keep_ratio) fraction by score
+```
+
+This approach matches SnapKV's "observation pooling" (using cumulative attention across all query positions as importance proxy).
+
+**Prefix caching** (`paged_kv_cache.rs`):
+
+Add to `PagedKvCache`:
+```rust
+prefix_index: HashMap<u64, Vec<PageId>>,
+```
+
+Hash function (Zobrist-style, collision probability ~1/2^64):
+```rust
+fn compute_prefix_hash(tokens: &[u32]) -> u64 {
+    // Precomputed random table: HASH_TABLE[pos % 1024][token % 65536]
+    tokens.iter().enumerate().fold(0u64, |acc, (pos, &tok)| {
+        acc ^ ZOBRIST_TABLE[pos % 1024][(tok % 65536) as usize]
+    })
+}
+```
+
+On each new request in the scheduler:
+1. Compute prefix hash for first `N` tokens (where N is the longest aligned page boundary)
+2. Look up `prefix_index[hash]` — if found, reuse those pages with copy-on-write
+3. Allocate new pages only for tokens beyond the prefix
+
+On page eviction:
+- Remove matching entry from `prefix_index`
+- If CoW refcount > 1: decrement, do not free
+
+**Enable by default** (`InferenceConfig`):
+```rust
+pub enable_prefix_sharing: bool = true,   // was false
+pub snapkv_keep_ratio: f32 = 0.6,
+pub snapkv_recent_window: usize = 32,
+```
+
+### Verification
+```
+cargo test -p rtx-inference -- test_prefix_cache_sharing
+cargo test -p rtx-inference -- test_snapkv_eviction
+```
+- Prefix test: two 2048-token requests sharing a 128-token system prompt must reuse identical pages (verified by page ID equality), saving ≥ 6% page allocations.
+- SnapKV test: after prefill on a 2048-token sequence with keep_ratio=0.6, active page count ≤ 40% of input length (allowing ≥ recent_window retained pages).
+
+---
+
+## Execution Order
+
+```
+Week 1:  Item 1 (CUDA Graphs) — 0.5 day
+         Item 2a–2b (FP8 casting kernels + cuBLASLt) — 3 days
+Week 2:  Item 2c–2e (FP8 linear + training config + tests) — 3 days
+         Item 4a (SnapKV eviction) — 1 day
+Week 3:  Item 4b–4d (prefix caching + defaults + tests) — 2 days
+         Item 3a (FA3 forward kernel: WGMMA + TMA + warp spec) — 3 days
+Week 4:  Item 3b (FA3 backward kernel) — 3 days
+         Item 3c–3d (SM selection + benchmarks) — 1 day
+Week 5:  End-to-end integration benchmark: all 4 items together
+```
+
+---
+
+## Dependencies and Risks
+
+| Risk | Mitigation |
+|------|-----------|
+| cuBLASLt FP8 requires CUDA ≥ 12.0 | RTX 5060 Ti ships with CUDA 12.x; check `build.rs` version guard |
+| WGMMA PTX inline asm requires sm_90+ | Runtime SM check in `backend_selector.rs` falls back to FA2 |
+| CUDA Graphs incompatible with dynamic shapes | Static shape check + warning; disable gracefully |
+| SnapKV may hurt accuracy on long-context retrieval tasks | `keep_ratio` is configurable; default 0.6 leaves 60% of keys |
+| Prefix cache hash collisions | Zobrist hash with 64-bit space; false positive rate ~1/2^64 per request |
+
+---
+
+## Non-goals (Batch 1)
+
+- Metal FlashAttention (stub remains; deferred to Batch 2)
+- GaLore-2 optimizer (deferred; independent of these 4 items)
+- EAGLE-3 speculative decoding upgrade (spec exists; Batch 2)
+- Ring-AllReduce (requires multi-node setup; Batch 2)
+- MegaBlocks grouped GEMM for MoE (deferred to Batch 2)