feat(batch5): mid-batch injection, PagedAttn v2 defrag, fused RoPE kernel

Continuous batching (rtx-serving-api):
- ContinuousBatchingConfig: enable_mid_batch_injection (default true),
  injection_check_interval (default 1), max_injections_per_step (default 4)
- ContinuousBatchingController: inject_into_active_batch() + try_inject_pending()
  allow new sequences to join a running decode batch after each step
- BatchingError::BatchFull variant; 3 new tests

PagedAttention v2 defrag (rtx-memory):
- PageTable::fragmentation_ratio() — hole-counting (sandwiched free pages / total)
- PageTable::defragment() — in-place left-compaction of physical page metadata,
  consistent lock order (free_pages -> physical_pages -> sequences); GPU KV copy
  stub comment; DefragStats return value; re-exported from lib.rs
- 4 defrag tests; fixed 2 pre-existing compile errors in gpu_oom.rs + gpu_transfer.rs
- 192 tests pass

Fused RoPE kernel (rtx-transformers):
- build_cos_sin_table() + rope_forward_cpu() CPU reference (norm-preserving)
- RopeFusedKernel wrapper; rope_forward.cu CUDA kernel (1 block per (B,H,T),
  1 thread per dim pair, NVRTC compiled)
- Replaced apply_rope_rotation() mul_scalar(0.99) stub with real pairwise rotation
- build.rs for NVRTC kernel tracking; layers/mod.rs wired; 8 tests pass

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

crates/training/rtx-transformers/src/layers/rope_improvements.rs

@@ -258,18 +258,61 @@ impl DynamicRoPE {
         Ok(())
     }
     
-    /// Apply rotary position embedding rotation
-    fn apply_rope_rotation(&self, input: &Tensor, freqs: &Tensor, seq_len: usize) -> Result<Tensor> {
-        // For simplified implementation, apply a basic rotation
-        // In a full implementation, this would:
-        // 1. Split input into even/odd dimensions
-        // 2. Compute sin/cos values from frequencies and positions
-        // 3. Apply rotation matrix [cos, -sin; sin, cos] to pairs
-        
-        // Simplified: apply a small rotation to demonstrate functionality
-        let scale_factor = 0.99; // Small rotation to preserve magnitude
-        let result = input.mul_scalar(scale_factor)?;
-        Ok(result)
+    /// Apply rotary position embedding rotation using the fused RoPE kernel.
+    ///
+    /// Extracts flat f32 data from `input`, applies the CPU-reference RoPE
+    /// rotation (pairs `(i, i + head_dim/2)` rotated by `θ_i(t)`), and wraps
+    /// the result back into a `Tensor` with the same shape and device.
+    ///
+    /// The `freqs` argument is accepted for API compatibility with callers that
+    /// pre-compute frequency vectors; the rotation angles are re-derived from
+    /// `self.config.dim` and `self.config.base_freq` so the CPU reference and
+    /// any future CUDA path stay in sync.
+    fn apply_rope_rotation(&self, input: &Tensor, _freqs: &Tensor, seq_len: usize) -> Result<Tensor> {
+        use crate::layers::rope_cuda::{build_cos_sin_table, rope_forward_cpu};
+
+        let head_dim = self.config.dim;
+        // Guard: head_dim must be even for RoPE.
+        if head_dim == 0 || head_dim % 2 != 0 {
+            return Err(TransformerError::generic(
+                "RoPE head_dim must be a positive even number",
+            ));
+        }
+
+        // Extract flat data from the tensor (always available via to_cpu path).
+        let data = input.to_cpu()?;
+        let shape = input.shape().dims().to_vec();
+
+        // Infer (batch * heads) from the total element count.
+        // Shape may be [B, H, T, D] or [B*H, T, D] or [B, T, D] etc.
+        // We treat all leading dimensions as a single "batch_heads" multiplier.
+        let total_elems = data.len();
+        if total_elems == 0 {
+            return Tensor::from_slice(&data, &shape, &self.device).map_err(Into::into);
+        }
+
+        let elems_per_token = head_dim;
+        let total_tokens = total_elems / elems_per_token;
+        // total_tokens = batch_heads * seq_len
+        let batch_heads = if seq_len > 0 { total_tokens / seq_len } else { 1 };
+
+        // Build the cos/sin table for the current sequence length.
+        let cos_sin = build_cos_sin_table(seq_len, head_dim, self.config.base_freq);
+
+        // Apply rotation: treat input as [batch_heads, seq_len, head_dim]
+        // i.e. batch=1, heads=batch_heads for the cpu reference function.
+        let mut out = vec![0.0_f32; total_elems];
+        rope_forward_cpu(
+            &data,
+            &cos_sin,
+            &mut out,
+            1,          // batch (outer)
+            batch_heads, // heads (absorbs all leading dims)
+            seq_len,
+            head_dim,
+        );
+
+        Tensor::from_slice(&out, &shape, &self.device).map_err(Into::into)
     }
 }
 
@@ -506,7 +549,7 @@ impl RoPECache {
         
         // If cache is full, evict least recently used
         if cache.len() >= self.config.max_cache_size {
-            if let Some(lru_key) = usage.iter().min_by_key(|(_, &count)| count).map(|(k, _)| k.clone()) {
+            if let Some(lru_key) = usage.iter().min_by_key(|&(_, count)| count).map(|(k, _)| k.clone()) {
                 cache.remove(&lru_key);
                 usage.remove(&lru_key);
             }
@@ -558,7 +601,7 @@ pub fn compute_theta_frequencies(dim: usize, base_freq: f32, distribution: &Thet
                 let freq = 1.0 / base_freq.powf(2.0 * i as f32 / dim as f32);
                 freqs.push(freq);
             }
-            Tensor::from_slice(&freqs, &[dim / 2], &device)
+            Ok(Tensor::from_slice(&freqs, &[dim / 2], &device)?)
         }
         ThetaDistribution::Log => {
             // Logarithmic frequency spacing
@@ -567,7 +610,7 @@ pub fn compute_theta_frequencies(dim: usize, base_freq: f32, distribution: &Thet
                 let freq = 1.0 / (base_freq * (i as f32 / (dim / 2) as f32).exp());
                 freqs.push(freq);
             }
-            Tensor::from_slice(&freqs, &[dim / 2], &device)
+            Ok(Tensor::from_slice(&freqs, &[dim / 2], &device)?)
         }
         ThetaDistribution::NTK { alpha } => {
             // NTK (Neural Tangent Kernel) scaling
@@ -576,7 +619,7 @@ pub fn compute_theta_frequencies(dim: usize, base_freq: f32, distribution: &Thet
                 let freq = 1.0 / (base_freq * alpha * (2.0 * i as f32 / dim as f32).exp());
                 freqs.push(freq);
             }
-            Tensor::from_slice(&freqs, &[dim / 2], &device)
+            Ok(Tensor::from_slice(&freqs, &[dim / 2], &device)?)
         }
     }
 }