(
&self,
x: &Tensor,
index_pos: usize,
block_idx: usize,
ctx: &mut Context,
)
| 101 | } |
| 102 | |
| 103 | async fn forward( |
| 104 | &self, |
| 105 | x: &Tensor, |
| 106 | index_pos: usize, |
| 107 | block_idx: usize, |
| 108 | ctx: &mut Context, |
| 109 | ) -> Result<Tensor> { |
| 110 | let residual = x; |
| 111 | |
| 112 | let x = ctx.backend.rms_norm(x, &self.rms_1_weight, self.rms_eps) |
| 113 | .map_err(|e| anyhow!("rms_1: {e}"))?; |
| 114 | let x = (self |
| 115 | .attn |
| 116 | .forward( |
| 117 | &x, |
| 118 | index_pos, |
| 119 | block_idx, |
| 120 | ctx.cache.as_mut().expect("No cache specified"), |
| 121 | ) |
| 122 | .map_err(|e| anyhow!("attention: {e}"))? |
| 123 | + residual) |
| 124 | .map_err(|e| anyhow!("residual: {e}"))?; |
| 125 | // Flush Metal command buffer between attention and MLP to prevent |
| 126 | // >25 command accumulation (no-op on CPU/CUDA) |
| 127 | let _ = ctx.backend.synchronize(); |
| 128 | let residual = &x; |
| 129 | let x = ctx.backend.rms_norm(&x, &self.rms_2_weight, self.rms_eps) |
| 130 | .map_err(|e| anyhow!("rms_2: {e}"))?; |
| 131 | let x = (self.mlp.forward(&x).map_err(|e| anyhow!("mlp: {e}"))? + residual) |
| 132 | .map_err(|e| anyhow!("mlp residual: {e}"))?; |
| 133 | |
| 134 | Ok(x) |
| 135 | } |
| 136 | |
| 137 | async fn forward_mut( |
| 138 | &mut self, |
no test coverage detected