rustyverse/rustytorch
2
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

style: cargo fmt --workspace (whitespace/wrapping only, no semantic change)

Browse Source 
Whole-workspace rustfmt pass picked up while iterating on Mamba GPU
backward work. Verified formatting-only via diff sampling; no logic
changed.

Co-Authored-By: Claude Sonnet 5 <[email protected]>
This commit is contained in:
osobh
2026-08-10 07:09:36 -07:00
co-authored by Claude Sonnet 5
parent ad6405663f
commit 4aaa36a57a
305 changed files with 25537 additions and 18337 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/training/rtx-transformers/src/layers/rope_improvements.rs
+130 -82
View File
@@ -1,5 +1,5 @@
//! Enhanced Rotary Position Embeddings (RoPE) improvements
//!
//!
//! This module implements various RoPE enhancements including:
//! - Dynamic RoPE with adjustable base frequency
//! - 2D RoPE for vision transformers
@@ -10,51 +10,51 @@
use crate::layers::Layer;
use crate::{Result, TransformerError};
use rtx_tensor::{Tensor, Shape, Device};
use parking_lot::RwLock;
use rtx_tensor::{Device, Shape, Tensor};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use parking_lot::RwLock;
#[cfg(all(test, feature = "disabled_tests"))]
mod tests {
use super::*;
use rtx_tensor::Device;
#[test]
fn test_dynamic_rope_config_creation() {
let config = DynamicRoPEConfig::new(128, 512, 10000.0);
assert!(config.is_err());
}
#[test]
#[test]
fn test_dynamic_rope_forward_pass() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = DynamicRoPEConfig::new(128, 512, 10000.0).unwrap();
let rope = DynamicRoPE::new(config, &device);
assert!(rope.is_err());
}
#[test]
fn test_rope_2d_config_creation() {
let config = RoPE2DConfig::new(64, 32, 32, 10000.0);
assert!(config.is_err());
}
#[test]
#[test]
fn test_rope_2d_forward_pass() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = RoPE2DConfig::new(64, 32, 32, 10000.0).unwrap();
let rope_2d = RoPE2D::new(config, &device);
assert!(rope_2d.is_err());
}
#[test]
fn test_interpolated_rope_config_creation() {
let config = InterpolatedRoPEConfig::new(128, 512, 10000.0, 1.0, 2.0);
assert!(config.is_err());
}
#[test]
fn test_interpolated_rope_forward_pass() {
let device = Device::cuda(0).unwrap_or(Device::default());
@@ -62,13 +62,13 @@ mod tests {
let rope = InterpolatedRoPE::new(config, &device);
assert!(rope.is_err());
}
#[test]
fn test_xpos_config_creation() {
let config = XPosConfig::new(128, 512, 10000.0, 0.5);
assert!(config.is_err());
}
#[test]
fn test_xpos_forward_pass() {
let device = Device::cuda(0).unwrap_or(Device::default());
@@ -76,41 +76,41 @@ mod tests {
let xpos = XPos::new(config, &device);
assert!(xpos.is_err());
}
#[test]
fn test_rope_cache_creation() {
let cache_config = RoPECacheConfig::new(1024, 512);
assert!(cache_config.is_err());
}
#[test]
fn test_rope_cache_operations() {
let cache_config = RoPECacheConfig::new(1024, 512).unwrap();
let cache = RoPECache::new(cache_config);
assert!(cache.is_err());
}
#[test]
fn test_theta_distribution_linear() {
let theta_dist = ThetaDistribution::Linear;
let freqs = compute_theta_frequencies(128, 10000.0, &theta_dist);
assert!(freqs.is_err());
}
#[test]
fn test_theta_distribution_log() {
let theta_dist = ThetaDistribution::Log;
let freqs = compute_theta_frequencies(128, 10000.0, &theta_dist);
assert!(freqs.is_err());
}
#[test]
fn test_theta_distribution_ntk() {
let theta_dist = ThetaDistribution::NTK { alpha: 1.0 };
let freqs = compute_theta_frequencies(128, 10000.0, &theta_dist);
assert!(freqs.is_err());
}
#[test]
fn test_dynamic_rope_precompute_freqs() {
let device = Device::cuda(0).unwrap_or(Device::default());
@@ -119,53 +119,53 @@ mod tests {
let result = rope.precompute_freqs(512);
assert!(result.is_err()); // Expected since we don't have actual tensor operations
}
#[test]
#[test]
fn test_dynamic_rope_update_frequency() {
let device = Device::cuda(0).unwrap_or(Device::default());
let config = DynamicRoPEConfig::new(128, 512, 10000.0).unwrap();
let mut rope = DynamicRoPE::new(config, &device).unwrap();
let result = rope.update_base_frequency(5000.0);
assert!(result.is_ok());
let invalid_result = rope.update_base_frequency(-1.0);
assert!(invalid_result.is_err());
}
#[test]
fn test_rope_cache_operations_advanced() {
let cache_config = RoPECacheConfig::new(1024, 512).unwrap();
let cache = RoPECache::new(cache_config).unwrap();
let theta_dist = ThetaDistribution::Linear;
let result = cache.get_or_compute_freqs("test_key", 128, 10000.0, &theta_dist);
assert!(result.is_err()); // Expected since we don't have actual tensor operations
let (current_size, max_size) = cache.cache_stats();
assert_eq!(current_size, 0);
assert_eq!(max_size, 1024);
}
#[test]
fn test_layer_trait_implementations() {
let device = Device::cuda(0).unwrap_or(Device::default());
// Test DynamicRoPE layer trait
let config = DynamicRoPEConfig::new(128, 512, 10000.0).unwrap();
let rope = DynamicRoPE::new(config, &device).unwrap();
assert_eq!(rope.layer_type(), "DynamicRoPE");
assert_eq!(rope.device(), &device);
// Test RoPE2D layer trait
let config_2d = RoPE2DConfig::new(64, 32, 32, 10000.0).unwrap();
let rope_2d = RoPE2D::new(config_2d, &device).unwrap();
assert_eq!(rope_2d.layer_type(), "RoPE2D");
// Test InterpolatedRoPE layer trait
// Test InterpolatedRoPE layer trait
let config_interp = InterpolatedRoPEConfig::new(128, 512, 10000.0, 256.0, 2.0).unwrap();
let rope_interp = InterpolatedRoPE::new(config_interp, &device).unwrap();
assert_eq!(rope_interp.layer_type(), "InterpolatedRoPE");
// Test XPos layer trait
let config_xpos = XPosConfig::new(128, 512, 10000.0, 0.5).unwrap();
let xpos = XPos::new(config_xpos, &device).unwrap();
@@ -188,7 +188,7 @@ pub enum ThetaDistribution {
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DynamicRoPEConfig {
pub dim: usize,
pub max_seq_len: usize,
pub max_seq_len: usize,
pub base_freq: f32,
pub theta_distribution: ThetaDistribution,
}
@@ -204,7 +204,7 @@ impl DynamicRoPEConfig {
if base_freq <= 0.0 {
return Err(TransformerError::generic("base_freq must be positive"));
}
Ok(Self {
dim,
max_seq_len,
@@ -230,34 +230,34 @@ impl DynamicRoPE {
cached_freqs: None,
})
}
/// Precompute and cache frequencies for efficiency
pub fn precompute_freqs(&mut self, max_seq_len: usize) -> Result<()> {
let freqs = compute_theta_frequencies(
self.config.dim,
self.config.base_freq,
&self.config.theta_distribution
&self.config.theta_distribution,
)?;
// Create position encodings for all positions up to max_seq_len
let positions: Vec<f32> = (0..max_seq_len).map(|i| i as f32).collect();
let pos_tensor = Tensor::from_slice(&positions, &[max_seq_len], &self.device)?;
self.cached_freqs = Some(freqs);
Ok(())
}
/// Update base frequency dynamically
pub fn update_base_frequency(&mut self, new_base_freq: f32) -> Result<()> {
if new_base_freq <= 0.0 {
return Err(TransformerError::generic("base_freq must be positive"));
}
self.config.base_freq = new_base_freq;
self.cached_freqs = None; // Invalidate cache
Ok(())
}
/// Apply rotary position embedding rotation using the fused RoPE kernel.
///
/// Extracts flat f32 data from `input`, applies the CPU-reference RoPE
@@ -268,7 +268,12 @@ impl DynamicRoPE {
/// 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> {
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;
@@ -294,7 +299,11 @@ impl DynamicRoPE {
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 };
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);
@@ -306,7 +315,7 @@ impl DynamicRoPE {
&data,
&cos_sin,
&mut out,
1, // batch (outer)
1, // batch (outer)
batch_heads, // heads (absorbs all leading dims)
seq_len,
head_dim,
@@ -339,7 +348,7 @@ impl RoPE2DConfig {
if base_freq <= 0.0 {
return Err(TransformerError::generic("base_freq must be positive"));
}
Ok(Self {
dim,
height,
@@ -380,7 +389,13 @@ pub struct InterpolatedRoPEConfig {
}
impl InterpolatedRoPEConfig {
pub fn new(dim: usize, max_seq_len: usize, base_freq: f32, original_max_len: f32, interpolation_factor: f32) -> Result<Self> {
pub fn new(
dim: usize,
max_seq_len: usize,
base_freq: f32,
original_max_len: f32,
interpolation_factor: f32,
) -> Result<Self> {
if dim == 0 {
return Err(TransformerError::generic("dim must be positive"));
}
@@ -391,12 +406,16 @@ impl InterpolatedRoPEConfig {
return Err(TransformerError::generic("base_freq must be positive"));
}
if original_max_len <= 0.0 {
return Err(TransformerError::generic("original_max_len must be positive"));
return Err(TransformerError::generic(
"original_max_len must be positive",
));
}
if interpolation_factor <= 0.0 {
return Err(TransformerError::generic("interpolation_factor must be positive"));
return Err(TransformerError::generic(
"interpolation_factor must be positive",
));
}
Ok(Self {
dim,
max_seq_len,
@@ -446,9 +465,11 @@ impl XPosConfig {
return Err(TransformerError::generic("base_freq must be positive"));
}
if decay_factor < 0.0 || decay_factor > 1.0 {
return Err(TransformerError::generic("decay_factor must be between 0 and 1"));
return Err(TransformerError::generic(
"decay_factor must be between 0 and 1",
));
}
Ok(Self {
dim,
max_seq_len,
@@ -493,7 +514,7 @@ impl RoPECacheConfig {
if max_seq_len == 0 {
return Err(TransformerError::generic("max_seq_len must be positive"));
}
Ok(Self {
max_cache_size,
max_seq_len,
@@ -517,9 +538,15 @@ impl RoPECache {
usage_count: Arc::new(RwLock::new(HashMap::new())),
})
}
/// Get cached frequencies or compute and cache them
pub fn get_or_compute_freqs(&self, key: &str, dim: usize, base_freq: f32, distribution: &ThetaDistribution) -> Result<Tensor> {
pub fn get_or_compute_freqs(
&self,
key: &str,
dim: usize,
base_freq: f32,
distribution: &ThetaDistribution,
) -> Result<Tensor> {
// Check if already cached
{
let cache = self.cache.read();
@@ -532,40 +559,44 @@ impl RoPECache {
return Ok(cached_tensor.clone());
}
}
// Compute new frequencies
let freqs = compute_theta_frequencies(dim, base_freq, distribution)?;
// Cache the result
self.cache_freqs(key.to_string(), freqs.clone())?;
Ok(freqs)
}
/// Cache frequencies with LRU eviction
fn cache_freqs(&self, key: String, tensor: Tensor) -> Result<()> {
let mut cache = self.cache.write();
let mut usage = self.usage_count.write();
// 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);
}
}
cache.insert(key.clone(), tensor);
usage.insert(key, 1);
Ok(())
}
/// Clear the cache
pub fn clear(&self) {
self.cache.write().clear();
self.usage_count.write().clear();
}
/// Get cache statistics
pub fn cache_stats(&self) -> (usize, usize) {
let cache = self.cache.read();
@@ -574,25 +605,36 @@ impl RoPECache {
}
/// Utility function to create a cache key for frequency caching
pub fn create_cache_key(dim: usize, base_freq: f32, distribution: &ThetaDistribution, seq_len: usize) -> String {
pub fn create_cache_key(
dim: usize,
base_freq: f32,
distribution: &ThetaDistribution,
seq_len: usize,
) -> String {
match distribution {
ThetaDistribution::Linear => format!("linear_{}_{}_{}_{}", dim, base_freq, seq_len, "1.0"),
ThetaDistribution::Log => format!("log_{}_{}_{}_{}", dim, base_freq, seq_len, "1.0"),
ThetaDistribution::NTK { alpha } => format!("ntk_{}_{}_{}_{}",dim, base_freq, seq_len, alpha),
ThetaDistribution::NTK { alpha } => {
format!("ntk_{}_{}_{}_{}", dim, base_freq, seq_len, alpha)
}
}
}
/// Compute theta frequencies based on distribution strategy
pub fn compute_theta_frequencies(dim: usize, base_freq: f32, distribution: &ThetaDistribution) -> Result<Tensor> {
pub fn compute_theta_frequencies(
dim: usize,
base_freq: f32,
distribution: &ThetaDistribution,
) -> Result<Tensor> {
if dim == 0 {
return Err(TransformerError::generic("dim must be positive"));
}
if base_freq <= 0.0 {
return Err(TransformerError::generic("base_freq must be positive"));
}
let device = Device::cuda(0).unwrap_or(Device::default());
match distribution {
ThetaDistribution::Linear => {
// Linear frequency spacing: 1/(base^(2i/dim))
@@ -632,18 +674,22 @@ impl Layer for DynamicRoPE {
let seq_len = shape.dims().get(1).copied().unwrap_or(1);
let num_heads = shape.dims().get(2).copied().unwrap_or(1);
let head_dim = shape.dims().get(3).copied().unwrap_or(self.config.dim);
if head_dim != self.config.dim {
return Err(TransformerError::generic("Input dimension mismatch"));
}
// Generate frequencies
let freqs = if let Some(ref cached) = self.cached_freqs {
cached.clone()
} else {
compute_theta_frequencies(self.config.dim, self.config.base_freq, &self.config.theta_distribution)?
compute_theta_frequencies(
self.config.dim,
self.config.base_freq,
&self.config.theta_distribution,
)?
};
// Apply RoPE rotation
self.apply_rope_rotation(input, &freqs, seq_len)
}
@@ -672,16 +718,18 @@ impl Layer for RoPE2D {
let batch_size = shape.dims().get(0).copied().unwrap_or(1);
let seq_len = shape.dims().get(1).copied().unwrap_or(1);
let head_dim = shape.dims().last().copied().unwrap_or(self.config.dim);
if head_dim != self.config.dim {
return Err(TransformerError::generic("Input dimension mismatch"));
}
// Check if sequence length matches 2D dimensions
if seq_len != self.config.height * self.config.width {
return Err(TransformerError::generic("Sequence length must match height * width"));
return Err(TransformerError::generic(
"Sequence length must match height * width",
));
}
// For now, return input unchanged (minimal implementation)
Ok(input.clone())
}
@@ -709,18 +757,18 @@ impl Layer for InterpolatedRoPE {
let shape = input.shape();
let seq_len = shape.dims().get(1).copied().unwrap_or(1);
let head_dim = shape.dims().last().copied().unwrap_or(self.config.dim);
if head_dim != self.config.dim {
return Err(TransformerError::generic("Input dimension mismatch"));
}
// Apply interpolation scaling
let scale_factor = if seq_len > self.config.original_max_len as usize {
self.config.interpolation_factor
} else {
1.0
};
// For now, return input unchanged (minimal implementation)
Ok(input.clone())
}
@@ -748,14 +796,14 @@ impl Layer for XPos {
let shape = input.shape();
let seq_len = shape.dims().get(1).copied().unwrap_or(1);
let head_dim = shape.dims().last().copied().unwrap_or(self.config.dim);
if head_dim != self.config.dim {
return Err(TransformerError::generic("Input dimension mismatch"));
}
// Apply exponential decay for extrapolation
let max_pos = seq_len.min(self.config.max_seq_len);
// For now, return input unchanged (minimal implementation)
Ok(input.clone())
}
@@ -775,4 +823,4 @@ impl Layer for XPos {
fn parameters_mut(&mut self) -> Vec<&mut Tensor> {
vec![]
}
}
}