vrr/ruvector
2
0
Fork 0
mirror of https://github.com/ruvnet/RuVector.git synced 2026-07-10 01:38:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

perf(ruvllm): replace ones_f32 tensor with affine(-1,1) in ACT loop

Browse source 
`1 - cum_f32` was computed as `ones_f32 - cum_f32` (two tensors, one op).
`Tensor::affine(-1.0, 1.0)` computes `x * -1 + 1 = 1 - x` in a single fused
kernel without allocating the constant ones tensor. Applied in both the main
loop (remainder weight) and the tail-weight computation.

Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
ruvnet 2026-06-18 13:57:49 -04:00
parent 50eb592403
commit f2788fef2f
1 changed files with 4 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

7
crates/ruvllm/src/models/openmythos/recurrent.rs
View file

@ -220,7 +220,7 @@ impl RecurrentBlock {
let mut cum_f32 = Tensor::zeros((b, seq, 1), DType::F32, &device).map_err(cand)?;
let mut not_halted_f32 = Tensor::ones((b, seq, 1), DType::F32, &device).map_err(cand)?;
let mut depth_f32 = Tensor::zeros((b, seq, 1), DType::F32, &device).map_err(cand)?;
let ones_f32 = Tensor::ones((b, seq, 1), DType::F32, &device).map_err(cand)?;
// `ones_f32` removed — replaced by affine(-1, 1) = (1 - x) without a constant tensor.
// Precompute step tensors (t+1) for depth tracking — avoids Tensor::new per iteration.
let step_tensors: Vec<Tensor> = (0..n_loops)
@ -271,7 +271,8 @@ impl RecurrentBlock {
// newly halting: w = 1 − cumulative_before_this_step
// still running: w = p
// already halted: w = 0 (not_halted_f32 = 0 for them)
let remainder = (&ones_f32 - &cum_f32).map_err(cand)?;
// affine(-1, 1) = 1 - cum_f32, avoids a constant ones tensor.
let remainder = cum_f32.affine(-1.0, 1.0).map_err(cand)?;
let w_halt = (&will_halt * &remainder).map_err(cand)?;
let still_running = (&not_halted_f32 - &will_halt).map_err(cand)?;
let w_run = (&still_running * &p_eff).map_err(cand)?;
@ -319,7 +320,7 @@ impl RecurrentBlock {
.map_err(cand)?;
if remaining_final > 0.5 {
let tail_f32 =
(&(&ones_f32 - &cum_f32).map_err(cand)? * &not_halted_f32).map_err(cand)?;
(cum_f32.affine(-1.0, 1.0).map_err(cand)? * &not_halted_f32).map_err(cand)?;
let w = if dtype != DType::F32 {
tail_f32.to_dtype(dtype).map_err(cand)?
} else {