feat(perf): GPU perf batch 1 — FP8, FA3 Blackwell, CUDA Graphs, SnapKV, prefix cache

Item 1 — CUDA Graphs wiring (rtx-transformers)
- Added `enable_cuda_graphs: bool` (default false) + `cuda_graph_warmup_iters: usize`
  (default 3) to `TrainingConfig`
- Wired 3-phase state machine into `training_loop.rs` (warmup → capture → replay)
  gated on `#[cfg(feature = "cuda")]`; stream plumbing stubbed with TODO pending
  `CudaStreamHandle` threading

Item 2 — FP8 E4M3/E5M2 training infrastructure (rtx-tensor, rtx-transformers)
- `fp8_cast.cu`: dual-path CUDA kernels — SM_89+ uses `<cuda_fp8.h>` native
  `__nv_cvt_*` intrinsics; older SM uses software bit-manipulation fallback
- `fp8_cast.rs`: host-side CPU casting + `#[cfg(feature = "cuda")]` GPU stubs
- `fp8_gemm.rs`: bit-accurate E4M3 decoder/encoder, BF16 round-trip utils, CPU
  reference matmul with cuBLASLt GPU path documented inline; 12 unit tests
- `training_config.rs`: `fp8_training: bool`, `fp8_e4m3_forward: bool`
- `linear.rs` (modular): `fp8_mode: bool` field + forward dispatch stub
- build.rs: registers `fp8_cast.cu` alongside existing `element_wise.cu`
- 22 FP8 unit tests — all pass

Item 3 — FlashAttention-3 Blackwell (WGMMA + TMA + warp specialization)
- `flash_attention_v3_forward.cu`: SM_90+ warp-specialised producer/consumer
  kernel (producer TMA-loads K/V tiles, consumers run WMMA as portable WGMMA
  proxy); SM_89+ FP8 header path; SM<90 standard FA2-style WMMA fallback
- `flash_v3_forward.rs`: NVRTC wrapper (`compile_ptx` via `include_str!`),
  `FlashV3ForwardKernel::new/is_supported/forward`; 6 unit tests
- `backend_selector.rs`: `SdpaBackend::FlashAttentionV3`, `for_compute_capability`,
  `supports_flash_v3` (major >= 9), FA3 scoring (0.98/0.90/0.70), 2× speedup estimate
- `kernels/simple.rs`: `v3_kernel: Option<FlashV3ForwardKernel>` in `FlashCudaKernels`
- Fixed pre-existing `Device::Cpu` cfg-gate bug in `tensor/creation.rs`
- 8 new FA3 backend tests + 6 kernel unit tests; 50 total pass

Item 4 — SnapKV attention-score eviction + prefix caching (rtx-inference, rtx-serving-api)
- `prefix_index.rs`: `PrefixIndex` with 8MB Zobrist hash table (Knuth MMIX LCG seed),
  `compute_hash/lookup/insert/remove/remove_page`; 10 unit tests
- `eviction.rs`: `AttentionScoreEviction` struct — `accumulate_scores` +
  `select_evict_positions` (retain top keep_ratio + last recent_window); 7 unit tests
- `types.rs`: `EvictionPolicy::AttentionScore { keep_ratio, recent_window }` +
  `KvCacheConfig::enable_prefix_caching`
- `paged_kv_cache.rs`: `prefix_index: Option<PrefixIndex>` + `lookup_prefix /
  register_prefix / unregister_prefix_page / prefix_caching_enabled` methods
- `config.rs` (serving-api): `enable_prefix_sharing: true` (was false),
  `snapkv_keep_ratio: 0.6`, `snapkv_recent_window: 32`
- Fixed 12 pre-existing test errors (spurious `.await` on sync constructors)
- 17 SnapKV/prefix tests pass

Total: 918 lib tests pass across rtx-tensor, rtx-flash-attention, rtx-transformers,
rtx-inference. Zero new failures.

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

crates/training/rtx-flash-attention/src/backend_selector.rs

@@ -39,6 +39,8 @@ use std::collections::HashMap;
 pub enum SdpaBackend {
     /// FlashAttention v2 - optimal for long sequences
     FlashAttention,
+    /// FlashAttention v3 - WGMMA + TMA + warp specialization (Hopper/Blackwell)
+    FlashAttentionV3,
     /// Standard mathematical attention - simple, debuggable
     Math,
     /// Memory-efficient chunked attention
@@ -55,6 +57,7 @@ impl std::fmt::Display for SdpaBackend {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
             SdpaBackend::FlashAttention => write!(f, "FlashAttention"),
+            SdpaBackend::FlashAttentionV3 => write!(f, "FlashAttentionV3"),
             SdpaBackend::Math => write!(f, "Math"),
             SdpaBackend::MemoryEfficient => write!(f, "MemoryEfficient"),
             SdpaBackend::CuDnn => write!(f, "cuDNN"),
@@ -151,6 +154,36 @@ impl HardwareCapabilities {
         }
     }
 
+    /// Detect capabilities for a specific CUDA compute capability.
+    ///
+    /// Call this instead of `detect_cuda()` when the SM version is known at call-site
+    /// (e.g. after querying the driver or from a build-time constant).
+    pub fn for_compute_capability(major: u32, minor: u32) -> Self {
+        let is_hopper_plus = major >= 9;
+        let is_blackwell_plus = major >= 12;
+        Self {
+            device_type: DeviceType::Cuda,
+            compute_capability: Some((major, minor)),
+            total_memory: 16 * 1024 * 1024 * 1024,     // conservative 16 GB
+            available_memory: 14 * 1024 * 1024 * 1024,
+            has_tensor_cores: true,
+            has_fp16: true,
+            has_bf16: major >= 8,
+            has_fp8: is_hopper_plus, // FP8 requires SM_90+
+            num_compute_units: if is_blackwell_plus { 84 } else { 108 },
+            memory_bandwidth_gbps: if is_blackwell_plus { 960.0 } else { 2039.0 },
+            supports_flash_attention: true,
+            supports_cudnn_attention: major >= 8,
+        }
+    }
+
+    /// Returns `true` if this hardware can run FlashAttention-3 (WGMMA + TMA).
+    ///
+    /// FA3 requires SM_90+ (Hopper or newer).
+    pub fn supports_flash_v3(&self) -> bool {
+        matches!(self.compute_capability, Some((major, _)) if major >= 9)
+    }
+
     /// Detect capabilities for Metal device
     pub fn detect_metal() -> Self {
         Self {
@@ -173,6 +206,7 @@ impl HardwareCapabilities {
     pub fn supports_backend(&self, backend: SdpaBackend) -> bool {
         match backend {
             SdpaBackend::FlashAttention => self.supports_flash_attention,
+            SdpaBackend::FlashAttentionV3 => self.supports_flash_v3(),
             SdpaBackend::Math => true, // Always supported
             SdpaBackend::MemoryEfficient => true,
             SdpaBackend::CuDnn => self.supports_cudnn_attention,
@@ -472,6 +506,12 @@ impl SdpaBackendSelector {
                     && input.head_dim <= 256
                     && (input.head_dim == 32 || input.head_dim == 64 || input.head_dim == 128 || input.head_dim == 256)
             }
+            SdpaBackend::FlashAttentionV3 => {
+                // FA3 shares FA2 head-dim constraints; additionally requires SM_90+
+                input.seq_len_q >= self.config.flash_min_seq_len
+                    && input.head_dim <= 256
+                    && (input.head_dim == 32 || input.head_dim == 64 || input.head_dim == 128 || input.head_dim == 256)
+            }
             SdpaBackend::CuDnn => {
                 // cuDNN constraints
                 input.head_dim <= 128 && !input.has_mask
@@ -528,6 +568,7 @@ impl SdpaBackendSelector {
     /// Get available backends for input
     fn get_available_backends(&self, input: &AttentionInputInfo) -> Vec<SdpaBackend> {
         [
+            SdpaBackend::FlashAttentionV3,
             SdpaBackend::FlashAttention,
             SdpaBackend::CuDnn,
             SdpaBackend::MemoryEfficient,
@@ -552,6 +593,16 @@ impl SdpaBackendSelector {
         }
 
         match backend {
+            SdpaBackend::FlashAttentionV3 => {
+                // FA3 beats FA2 across the board on Hopper/Blackwell
+                if seq_len >= 2048 {
+                    (0.98, format!("FlashAttentionV3 (WGMMA+TMA) optimal for seq_len={seq_len}"))
+                } else if seq_len >= 512 {
+                    (0.90, "FlashAttentionV3 excellent for medium sequences".to_string())
+                } else {
+                    (0.70, "FlashAttentionV3 has overhead for short sequences".to_string())
+                }
+            }
             SdpaBackend::FlashAttention => {
                 if seq_len >= 2048 {
                     (0.95, format!("FlashAttention optimal for seq_len={}", seq_len))
@@ -600,6 +651,14 @@ impl SdpaBackendSelector {
         let seq_len = input.seq_len_q.max(input.seq_len_kv);
 
         match backend {
+            SdpaBackend::FlashAttentionV3 => {
+                // ~2x over FA2 from WGMMA tile efficiency + async TMA overlap
+                if seq_len >= 4096 { 10.0 }
+                else if seq_len >= 2048 { 7.0 }
+                else if seq_len >= 1024 { 5.0 }
+                else if seq_len >= 512 { 3.5 }
+                else { 2.0 }
+            }
             SdpaBackend::FlashAttention => {
                 if seq_len >= 4096 { 5.0 }
                 else if seq_len >= 2048 { 3.5 }
@@ -623,8 +682,8 @@ impl SdpaBackendSelector {
         let base_memory = input.estimate_memory_bytes();
 
         match backend {
-            SdpaBackend::FlashAttention => {
-                // FlashAttention uses O(N) instead of O(N^2) for attention matrix
+            SdpaBackend::FlashAttentionV3 | SdpaBackend::FlashAttention => {
+                // Both FA2 and FA3 use O(N) tiling; same memory footprint
                 let qkvo_size = base_memory / 2; // Q, K, V, O only
                 let softmax_lse = input.batch_size * input.num_heads * input.seq_len_q * 4;
                 qkvo_size + softmax_lse