From f2788fef2f208eb90baa64fe6efe1c6dcd5e71f2 Mon Sep 17 00:00:00 2001 From: ruvnet Date: Thu, 18 Jun 2026 13:57:49 -0400 Subject: [PATCH] perf(ruvllm): replace ones_f32 tensor with affine(-1,1) in ACT loop `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 --- crates/ruvllm/src/models/openmythos/recurrent.rs | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/crates/ruvllm/src/models/openmythos/recurrent.rs b/crates/ruvllm/src/models/openmythos/recurrent.rs index 4abbf158f..1fcf262a0 100644 --- a/crates/ruvllm/src/models/openmythos/recurrent.rs +++ b/crates/ruvllm/src/models/openmythos/recurrent.rs @@ -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 = (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 = (¬_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)? * ¬_halted_f32).map_err(cand)?; + (cum_f32.affine(-1.0, 1.0).map_err(cand)? * ¬_halted_f32).map_err(cand)?; let w = if dtype != DType::F32 { tail_f32.to_dtype(dtype).map_err(cand)? } else {