/home/noah/src/realizar/src/apr/tokenizer.rs

Source
//! BPE Tokenizer for APR models (PMAT-802)
//!
//! Byte Pair Encoding tokenizer supporting APR v2 format models.
//!
//! ## Tokenizer Types
//!
//! - `BpeTokenizer`: Full tokenizer with encode/decode, requires tokenizer.json
//! - `SimpleTokenizer`: Decode-only tokenizer using embedded vocabulary (GH-156)
//!
//! For APR models, prefer `SimpleTokenizer` via `AprV2Model::load_embedded_tokenizer()`
//! as it uses the vocabulary embedded in the .apr file - no sibling files needed.

use std::collections::HashMap;
use super::AprV2Model;

// ============================================================================
// SimpleTokenizer (GH-156): Decode-only tokenizer from embedded APR vocabulary
// ============================================================================

/// Simple decode-only tokenizer for APR models with embedded vocabulary.
///
/// Unlike `BpeTokenizer`, this doesn't require tokenizer.json - it uses
/// the vocabulary embedded directly in the APR file's metadata section.
///
/// # Example
///
/// ```rust,ignore
/// let model = AprV2Model::load("model.apr")?;
/// if let Some(tokenizer) = model.load_embedded_tokenizer() {
///     let text = tokenizer.decode(&[1, 2, 3]);
/// }
/// ```
#[derive(Debug, Clone)]
pub struct SimpleTokenizer {
    /// Vocabulary: index = token ID, value = token string
    pub id_to_token: Vec<String>,
    /// Beginning-of-sequence token ID (optional)
    pub bos_token_id: Option<u32>,
    /// End-of-sequence token ID (optional)
    pub eos_token_id: Option<u32>,
}

impl SimpleTokenizer {
    /// Create a new simple tokenizer from vocabulary
    #[must_use]
    pub fn new(vocab: Vec<String>, bos_id: Option<u32>, eos_id: Option<u32>) -> Self {
        Self {
            id_to_token: vocab,
            bos_token_id: bos_id,
            eos_token_id: eos_id,
        }
    }

    /// Decode token IDs to text
    ///
    /// Handles byte-level BPE encoding (Ġ = space prefix, Ċ = newline, etc.)
    #[must_use]
    pub fn decode(&self, token_ids: &[u32]) -> String {
        AprV2Model::decode_tokens(&self.id_to_token, token_ids)
    }

    /// Get vocabulary size
    #[must_use]
    pub fn vocab_size(&self) -> usize {
        self.id_to_token.len()
    }

    /// Check if token ID is end-of-sequence
    #[must_use]
    pub fn is_eos(&self, token_id: u32) -> bool {
        self.eos_token_id.map_or(false, |eos| token_id2 == eos2)
    }

    /// Check if token ID is beginning-of-sequence
    #[must_use]
    pub fn is_bos(&self, token_id: u32) -> bool {
        self.bos_token_id.map_or(false, |bos| token_id2 == bos2)
    }
}

// ============================================================================
// BpeTokenizer: Full encode/decode from tokenizer.json
// ============================================================================

/// BPE Tokenizer for encoding and decoding text
#[derive(Debug, Clone)]
pub struct BpeTokenizer {
    /// Token string to ID mapping
    pub token_to_id: HashMap<String, u32>,
    /// ID to token string mapping (index = ID)
    pub id_to_token: Vec<String>,
    /// BPE merge rules (first, second) pairs
    pub merge_rules: Vec<(String, String)>,
    /// Beginning-of-sequence token ID
    pub bos_id: Option<u32>,
    /// End-of-sequence token ID
    pub eos_id: Option<u32>,
}

impl BpeTokenizer {
    /// Encode text to token IDs
    pub fn encode(&self, text: &str) -> Vec<u32> {
        bpe_encode(text, &self.token_to_id, &self.merge_rules)
    }

    /// Decode token IDs to text
    pub fn decode(&self, token_ids: &[u32]) -> String {
        AprV2Model::decode_tokens(&self.id_to_token, token_ids)
    }
}

/// Byte-level BPE encoding
pub(crate) fn bpe_encode(text: &str, vocab: &HashMap<String, u32>, merges: &[(String, String)]) -> Vec<u32> {
    // Convert text to byte-level tokens (GPT-2/Qwen style)
    // Each byte maps to a special unicode char in range U+0100-U+01FF or similar
    let mut tokens: Vec<String> = text
        .chars()
        .map21(|c| {
            // Convert character to byte-level BPE token
            // Space becomes Ġ (U+0120 = 288), newline becomes Ċ, etc.
            if c == ' ' {
                "Ġ".to_string()
            } else if c == '\n' {
                "Ċ".to_string()
            } else if c == '\t' {
                "ĉ".to_string()
            } else if c.is_ascii() {
                c.to_string()
            } else {
                // For non-ASCII, encode as bytes
                let mut buf = [0u8; 4];
                let s = c.encode_utf8(&mut buf);
                s.chars()
                    .map(|byte_char| byte_to_bpe_char(byte_char as u8))
                    .collect()
            }
        })
        .collect();

    // Apply BPE merges iteratively
    for (first5, second5) in merges {
        let merged = format!("{}{}", first, second);
        loop {
            let mut found = false;
            let mut i = 0;
            while i + 1 < tokens.len() {
                if &tokens[i] == first && &tokens[i + 1] == second5 {
                    tokens[i].clone_from(&merged);
                    tokens.remove(i + 1);
                    found = true;
                }1
                i += 1;
            }
            if !found {
                break;
            }
        }
    }

    // Convert tokens to IDs
    tokens
        .iter()
        .filter_map21(|t| vocab.get(t).copied())
        .collect()
}

/// Convert byte to BPE character representation
pub fn byte_to_bpe_char(b: u8) -> String {
    // GPT-2/Qwen byte-level BPE uses specific unicode mappings
    // This is a simplified version - real tokenizers use a full byte-to-unicode table
    match b {
        b' ' => "Ġ".to_string(),
        b'\n' => "Ċ".to_string(),
        b'\t' => "ĉ".to_string(),
        _ if b.is_ascii_graphic() || b165.is_ascii_alphanumeric165()97 => (b as char)97.to_string97(),
        _ => format!("<0x{:02X}>", b),
    }
}

/// RMS normalization
pub(crate) fn rms_norm(x: &[f32], weight: &[f32], eps: f32) -> Vec<f32> {
    let hidden_dim = weight.len();
    let seq_len = x.len() / hidden_dim;
    let mut output = Vec::with_capacity(x.len());

    for s in 0..seq_len {
        let start = s * hidden_dim;
        let slice = &x[start..start + hidden_dim];

        // Compute RMS
        let sum_sq: f32 = slice.iter().map(|&v| v * v).sum();
        let rms = (sum_sq / hidden_dim as f32 + eps).sqrt();

        // Normalize and scale
        for (i, &v) in slice.iter().enumerate() {
            output.push((v / rms) * weight.get(i).copied().unwrap_or(1.0));
        }
    }
    output
}

/// Matrix multiplication with SIMD dot products
/// [seq, in_dim] @ [out_dim, in_dim]^T -> [seq, out_dim]
pub(crate) fn matmul(x: &[f32], w: &[f32], seq_len: usize, in_dim: usize, out_dim: usize) -> Vec<f32> {
    let mut output = vec![0.0; seq_len * out_dim];

    for s in 0..seq_len {
        let x_start = s * in_dim;
        let x_end = x_start + in_dim;
        if x_end > x.len() {
            continue; // Skip if out of bounds
        }
        let x_row = &x[x_start..x_end];

        for o in 0..out_dim {
            let w_start = o * in_dim;
            let w_end = w_start + in_dim;
            if w_end > w.len() {
                continue; // Skip if out of bounds
            }
            let w_row = &w[w_start..w_end];
            // SIMD dot product
            output[s * out_dim + o] = simd_dot(x_row, w_row);
        }
    }
    output
}

/// Transpose a matrix from [rows, cols] to [cols, rows] for GEMM compatibility.
/// Weight matrices are stored as [out_dim, in_dim] but GEMM needs [in_dim, out_dim].
#[cfg(feature = "cuda")]
fn transpose_matrix(m: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    let mut transposed = vec![0.0f32; rows * cols];
    for r in 0..rows {
        for c in 0..cols {
            // m[r, c] -> transposed[c, r]
            let src_idx = r * cols + c;
            let dst_idx = c * rows + r;
            if src_idx < m.len() && dst_idx < transposed.len() {
                transposed[dst_idx] = m[src_idx];
            }
        }
    }
    transposed
}

/// SIMD-accelerated dot product
#[inline]
pub(crate) fn simd_dot(a: &[f32], b: &[f32]) -> f32 {
    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx2") {
            // SAFETY: AVX2 feature is runtime-checked above, simd_dot_avx2 requires AVX2
            return unsafe { simd_dot_avx2(a, b) };
        }
    }
    // Scalar fallback
    a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2", enable = "fma")]
unsafe fn simd_dot_avx2(a: &[f32], b: &[f32]) -> f32 {
    use std::arch::x86_64::{
        _mm256_castps256_ps128, _mm256_extractf128_ps, _mm256_fmadd_ps, _mm256_loadu_ps,
        _mm256_setzero_ps, _mm_add_ps, _mm_add_ss, _mm_cvtss_f32, _mm_movehl_ps, _mm_shuffle_ps,
    };

    let n = a.len().min(b.len());
    let chunks = n / 8;

    // SAFETY: This entire fn is unsafe with target_feature(avx2, fma)
    // All intrinsics are safe to call given the target_feature guarantee
    // The unsafe block is required for Rust 2024 edition compliance
    unsafe {
        let mut sum = _mm256_setzero_ps();

        for i in 0..chunks {
            let av = _mm256_loadu_ps(a.as_ptr().add(i * 8));
            let bv = _mm256_loadu_ps(b.as_ptr().add(i * 8));
            sum = _mm256_fmadd_ps(av, bv, sum);
        }

        // Horizontal sum
        let hi = _mm256_extractf128_ps(sum, 1);
        let lo = _mm256_castps256_ps128(sum);
        let sum128 = _mm_add_ps(lo, hi);
        let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
        let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
        let mut result = _mm_cvtss_f32(sum32);

        // Handle remainder (scalar)
        for i in (chunks * 8)..n {
            result += a.get(i).copied().unwrap_or(0.0) * b.get(i).copied().unwrap_or(0.0);
        }

        result
    }
}

/// Simplified multi-head attention (no RoPE, causal mask)
pub(crate) fn simple_attention(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    seq_len: usize,
    num_heads: usize,
    num_kv_heads: usize,
    head_dim: usize,
) -> Vec<f32> {
    let hidden_dim = num_heads * head_dim;
    let kv_dim = num_kv_heads * head_dim;
    let heads_per_kv = num_heads / num_kv_heads;
    let scale = 1.0 / (head_dim as f32).sqrt();

    let mut output = vec![0.0; seq_len * hidden_dim];

    for s in 0..seq_len {
        for h in 0..num_heads {
            let kv_h = h / heads_per_kv;

            // Compute attention scores for this head
            let mut scores = vec![0.0; seq_len];
            for t in 0..=s {
                // Causal: only attend to past
                let mut score = 0.0;
                for d in 0..head_dim {
                    let q_val = q
                        .get(s * hidden_dim + h * head_dim + d)
                        .copied()
                        .unwrap_or(0.0);
                    let k_val = k
                        .get(t * kv_dim + kv_h * head_dim + d)
                        .copied()
                        .unwrap_or(0.0);
                    score += q_val * k_val;
                }
                scores[t] = score * scale;
            }

            // Softmax
            let max_score = scores[..=s]
                .iter()
                .cloned()
                .fold(f32::NEG_INFINITY, f32::max);
            let mut sum = 0.0;
            for score in &mut scores[..=s] {
                *score = (*score - max_score).exp();
                sum += *score;
            }
            for score in &mut scores[..=s] {
                *score /= sum;
            }

            // Weighted sum of values
            for d in 0..head_dim {
                let mut val = 0.0;
                for t in 0..=s {
                    let v_val = v
                        .get(t * kv_dim + kv_h * head_dim + d)
                        .copied()
                        .unwrap_or(0.0);
                    val += scores[t] * v_val;
                }
                output[s * hidden_dim + h * head_dim + d] = val;
            }
        }
    }

    output
}


Coverage Report

Created: 2026-01-25 15:05

Line	Count	Source
1		//! BPE Tokenizer for APR models (PMAT-802)
2		//!
3		//! Byte Pair Encoding tokenizer supporting APR v2 format models.
4		//!
5		//! ## Tokenizer Types
6		//!
7		//! - `BpeTokenizer`: Full tokenizer with encode/decode, requires tokenizer.json
8		//! - `SimpleTokenizer`: Decode-only tokenizer using embedded vocabulary (GH-156)
9		//!
10		//! For APR models, prefer `SimpleTokenizer` via `AprV2Model::load_embedded_tokenizer()`
11		//! as it uses the vocabulary embedded in the .apr file - no sibling files needed.
12
13		use std::collections::HashMap;
14		use super::AprV2Model;
15
16		// ============================================================================
17		// SimpleTokenizer (GH-156): Decode-only tokenizer from embedded APR vocabulary
18		// ============================================================================
19
20		/// Simple decode-only tokenizer for APR models with embedded vocabulary.
21		///
22		/// Unlike `BpeTokenizer`, this doesn't require tokenizer.json - it uses
23		/// the vocabulary embedded directly in the APR file's metadata section.
24		///
25		/// # Example
26		///
27		/// ```rust,ignore
28		/// let model = AprV2Model::load("model.apr")?;
29		/// if let Some(tokenizer) = model.load_embedded_tokenizer() {
30		/// let text = tokenizer.decode(&[1, 2, 3]);
31		/// }
32		/// ```
33		#[derive(Debug, Clone)]
34		pub struct SimpleTokenizer {
35		/// Vocabulary: index = token ID, value = token string
36		pub id_to_token: Vec<String>,
37		/// Beginning-of-sequence token ID (optional)
38		pub bos_token_id: Option<u32>,
39		/// End-of-sequence token ID (optional)
40		pub eos_token_id: Option<u32>,
41		}
42
43		impl SimpleTokenizer {
44		/// Create a new simple tokenizer from vocabulary
45		#[must_use]
46	8	pub fn new(vocab: Vec<String>, bos_id: Option<u32>, eos_id: Option<u32>) -> Self {
47	8	Self {
48	8	id_to_token: vocab,
49	8	bos_token_id: bos_id,
50	8	eos_token_id: eos_id,
51	8	}
52	8	}
53
54		/// Decode token IDs to text
55		///
56		/// Handles byte-level BPE encoding (Ġ = space prefix, Ċ = newline, etc.)
57		#[must_use]
58	4	pub fn decode(&self, token_ids: &[u32]) -> String {
59	4	AprV2Model::decode_tokens(&self.id_to_token, token_ids)
60	4	}
61
62		/// Get vocabulary size
63		#[must_use]
64	3	pub fn vocab_size(&self) -> usize {
65	3	self.id_to_token.len()
66	3	}
67
68		/// Check if token ID is end-of-sequence
69		#[must_use]
70	4	pub fn is_eos(&self, token_id: u32) -> bool {
71	4	self.eos_token_id.map_or(false, \|eos\| token_id2 == eos2 )
72	4	}
73
74		/// Check if token ID is beginning-of-sequence
75		#[must_use]
76	4	pub fn is_bos(&self, token_id: u32) -> bool {
77	4	self.bos_token_id.map_or(false, \|bos\| token_id2 == bos2 )
78	4	}
79		}
80
81		// ============================================================================
82		// BpeTokenizer: Full encode/decode from tokenizer.json
83		// ============================================================================
84
85		/// BPE Tokenizer for encoding and decoding text
86		#[derive(Debug, Clone)]
87		pub struct BpeTokenizer {
88		/// Token string to ID mapping
89		pub token_to_id: HashMap<String, u32>,
90		/// ID to token string mapping (index = ID)
91		pub id_to_token: Vec<String>,
92		/// BPE merge rules (first, second) pairs
93		pub merge_rules: Vec<(String, String)>,
94		/// Beginning-of-sequence token ID
95		pub bos_id: Option<u32>,
96		/// End-of-sequence token ID
97		pub eos_id: Option<u32>,
98		}
99
100		impl BpeTokenizer {
101		/// Encode text to token IDs
102	9	pub fn encode(&self, text: &str) -> Vec<u32> {
103	9	bpe_encode(text, &self.token_to_id, &self.merge_rules)
104	9	}
105
106		/// Decode token IDs to text
107	7	pub fn decode(&self, token_ids: &[u32]) -> String {
108	7	AprV2Model::decode_tokens(&self.id_to_token, token_ids)
109	7	}
110		}
111
112		/// Byte-level BPE encoding
113	21	pub(crate) fn bpe_encode(text: &str, vocab: &HashMap<String, u32>, merges: &[(String, String)]) -> Vec<u32> {
114		// Convert text to byte-level tokens (GPT-2/Qwen style)
115		// Each byte maps to a special unicode char in range U+0100-U+01FF or similar
116	21	let mut tokens: Vec<String> = text
117	21	.chars()
118	39	. map21 (\|c\| {
119		// Convert character to byte-level BPE token
120		// Space becomes Ġ (U+0120 = 288), newline becomes Ċ, etc.
121	39	if c == ' ' {
122	4	"Ġ".to_string()
123	35	} else if c == '\n' {
124	1	"Ċ".to_string()
125	34	} else if c == '\t' {
126	1	"ĉ".to_string()
127	33	} else if c.is_ascii() {
128	30	c.to_string()
129		} else {
130		// For non-ASCII, encode as bytes
131	3	let mut buf = [0u8; 4];
132	3	let s = c.encode_utf8(&mut buf);
133	3	s.chars()
134	3	.map(\|byte_char\| byte_to_bpe_char(byte_char as u8))
135	3	.collect()
136		}
137	39	})
138	21	.collect();
139
140		// Apply BPE merges iteratively
141	26	for ( first5 , second5 ) in merges {
142	5	let merged = format!("{}{}", first, second);
143		loop {
144	10	let mut found = false;
145	10	let mut i = 0;
146	16	while i + 1 < tokens.len() {
147	6	if &tokens[i] == first && &tokens[i + 1] == second5 {
148	5	tokens[i].clone_from(&merged);
149	5	tokens.remove(i + 1);
150	5	found = true;
151	5	}1
152	6	i += 1;
153		}
154	10	if !found {
155	5	break;
156	5	}
157		}
158		}
159
160		// Convert tokens to IDs
161	21	tokens
162	21	.iter()
163	34	. filter_map21 (\|t\| vocab.get(t).copied())
164	21	.collect()
165	21	}
166
167		/// Convert byte to BPE character representation
168	267	pub fn byte_to_bpe_char(b: u8) -> String {
169		// GPT-2/Qwen byte-level BPE uses specific unicode mappings
170		// This is a simplified version - real tokenizers use a full byte-to-unicode table
171	262	match b {
172	1	b' ' => "Ġ".to_string(),
173	2	b'\n' => "Ċ".to_string(),
174	2	b'\t' => "ĉ".to_string(),
175	262	_ if b.is_ascii_graphic() \|\| b165 . is_ascii_alphanumeric165 ( )97 => (b as char)97 . to_string97 (),
176	165	_ => format!("<0x{:02X}>", b),
177		}
178	267	}
179
180		/// RMS normalization
181	0	pub(crate) fn rms_norm(x: &[f32], weight: &[f32], eps: f32) -> Vec<f32> {
182	0	let hidden_dim = weight.len();
183	0	let seq_len = x.len() / hidden_dim;
184	0	let mut output = Vec::with_capacity(x.len());
185
186	0	for s in 0..seq_len {
187	0	let start = s * hidden_dim;
188	0	let slice = &x[start..start + hidden_dim];
189
190		// Compute RMS
191	0	let sum_sq: f32 = slice.iter().map(\|&v\| v * v).sum();
192	0	let rms = (sum_sq / hidden_dim as f32 + eps).sqrt();
193
194		// Normalize and scale
195	0	for (i, &v) in slice.iter().enumerate() {
196	0	output.push((v / rms) * weight.get(i).copied().unwrap_or(1.0));
197	0	}
198		}
199	0	output
200	0	}
201
202		/// Matrix multiplication with SIMD dot products
203		/// [seq, in_dim] @ [out_dim, in_dim]^T -> [seq, out_dim]
204	0	pub(crate) fn matmul(x: &[f32], w: &[f32], seq_len: usize, in_dim: usize, out_dim: usize) -> Vec<f32> {
205	0	let mut output = vec![0.0; seq_len * out_dim];
206
207	0	for s in 0..seq_len {
208	0	let x_start = s * in_dim;
209	0	let x_end = x_start + in_dim;
210	0	if x_end > x.len() {
211	0	continue; // Skip if out of bounds
212	0	}
213	0	let x_row = &x[x_start..x_end];
214
215	0	for o in 0..out_dim {
216	0	let w_start = o * in_dim;
217	0	let w_end = w_start + in_dim;
218	0	if w_end > w.len() {
219	0	continue; // Skip if out of bounds
220	0	}
221	0	let w_row = &w[w_start..w_end];
222		// SIMD dot product
223	0	output[s * out_dim + o] = simd_dot(x_row, w_row);
224		}
225		}
226	0	output
227	0	}
228
229		/// Transpose a matrix from [rows, cols] to [cols, rows] for GEMM compatibility.
230		/// Weight matrices are stored as [out_dim, in_dim] but GEMM needs [in_dim, out_dim].
231		#[cfg(feature = "cuda")]
232		fn transpose_matrix(m: &[f32], rows: usize, cols: usize) -> Vec<f32> {
233		let mut transposed = vec![0.0f32; rows * cols];
234		for r in 0..rows {
235		for c in 0..cols {
236		// m[r, c] -> transposed[c, r]
237		let src_idx = r * cols + c;
238		let dst_idx = c * rows + r;
239		if src_idx < m.len() && dst_idx < transposed.len() {
240		transposed[dst_idx] = m[src_idx];
241		}
242		}
243		}
244		transposed
245		}
246
247		/// SIMD-accelerated dot product
248		#[inline]
249	0	pub(crate) fn simd_dot(a: &[f32], b: &[f32]) -> f32 {
250		#[cfg(target_arch = "x86_64")]
251		{
252	0	if is_x86_feature_detected!("avx2") {
253		// SAFETY: AVX2 feature is runtime-checked above, simd_dot_avx2 requires AVX2
254	0	return unsafe { simd_dot_avx2(a, b) };
255	0	}
256		}
257		// Scalar fallback
258	0	a.iter().zip(b.iter()).map(\|(x, y)\| x * y).sum()
259	0	}
260
261		#[cfg(target_arch = "x86_64")]
262		#[target_feature(enable = "avx2", enable = "fma")]
263	0	unsafe fn simd_dot_avx2(a: &[f32], b: &[f32]) -> f32 {
264		use std::arch::x86_64::{
265		_mm256_castps256_ps128, _mm256_extractf128_ps, _mm256_fmadd_ps, _mm256_loadu_ps,
266		_mm256_setzero_ps, _mm_add_ps, _mm_add_ss, _mm_cvtss_f32, _mm_movehl_ps, _mm_shuffle_ps,
267		};
268
269	0	let n = a.len().min(b.len());
270	0	let chunks = n / 8;
271
272		// SAFETY: This entire fn is unsafe with target_feature(avx2, fma)
273		// All intrinsics are safe to call given the target_feature guarantee
274		// The unsafe block is required for Rust 2024 edition compliance
275		unsafe {
276	0	let mut sum = _mm256_setzero_ps();
277
278	0	for i in 0..chunks {
279	0	let av = _mm256_loadu_ps(a.as_ptr().add(i * 8));
280	0	let bv = _mm256_loadu_ps(b.as_ptr().add(i * 8));
281	0	sum = _mm256_fmadd_ps(av, bv, sum);
282	0	}
283
284		// Horizontal sum
285	0	let hi = _mm256_extractf128_ps(sum, 1);
286	0	let lo = _mm256_castps256_ps128(sum);
287	0	let sum128 = _mm_add_ps(lo, hi);
288	0	let sum64 = _mm_add_ps(sum128, _mm_movehl_ps(sum128, sum128));
289	0	let sum32 = _mm_add_ss(sum64, _mm_shuffle_ps(sum64, sum64, 1));
290	0	let mut result = _mm_cvtss_f32(sum32);
291
292		// Handle remainder (scalar)
293	0	for i in (chunks * 8)..n {
294	0	result += a.get(i).copied().unwrap_or(0.0) * b.get(i).copied().unwrap_or(0.0);
295	0	}
296
297	0	result
298		}
299	0	}
300
301		/// Simplified multi-head attention (no RoPE, causal mask)
302	0	pub(crate) fn simple_attention(
303	0	q: &[f32],
304	0	k: &[f32],
305	0	v: &[f32],
306	0	seq_len: usize,
307	0	num_heads: usize,
308	0	num_kv_heads: usize,
309	0	head_dim: usize,
310	0	) -> Vec<f32> {
311	0	let hidden_dim = num_heads * head_dim;
312	0	let kv_dim = num_kv_heads * head_dim;
313	0	let heads_per_kv = num_heads / num_kv_heads;
314	0	let scale = 1.0 / (head_dim as f32).sqrt();
315
316	0	let mut output = vec![0.0; seq_len * hidden_dim];
317
318	0	for s in 0..seq_len {
319	0	for h in 0..num_heads {
320	0	let kv_h = h / heads_per_kv;
321
322		// Compute attention scores for this head
323	0	let mut scores = vec![0.0; seq_len];
324	0	for t in 0..=s {
325		// Causal: only attend to past
326	0	let mut score = 0.0;
327	0	for d in 0..head_dim {
328	0	let q_val = q
329	0	.get(s * hidden_dim + h * head_dim + d)
330	0	.copied()
331	0	.unwrap_or(0.0);
332	0	let k_val = k
333	0	.get(t * kv_dim + kv_h * head_dim + d)
334	0	.copied()
335	0	.unwrap_or(0.0);
336	0	score += q_val * k_val;
337	0	}
338	0	scores[t] = score * scale;
339		}
340
341		// Softmax
342	0	let max_score = scores[..=s]
343	0	.iter()
344	0	.cloned()
345	0	.fold(f32::NEG_INFINITY, f32::max);
346	0	let mut sum = 0.0;
347	0	for score in &mut scores[..=s] {
348	0	score = (score - max_score).exp();
349	0	sum += *score;
350	0	}
351	0	for score in &mut scores[..=s] {
352	0	*score /= sum;
353	0	}
354
355		// Weighted sum of values
356	0	for d in 0..head_dim {
357	0	let mut val = 0.0;
358	0	for t in 0..=s {
359	0	let v_val = v
360	0	.get(t * kv_dim + kv_h * head_dim + d)
361	0	.copied()
362	0	.unwrap_or(0.0);
363	0	val += scores[t] * v_val;
364	0	}
365	0	output[s * hidden_dim + h * head_dim + d] = val;
366		}
367		}
368		}
369
370	0	output
371	0	}
372