| 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 | // --- Attention sublayer (sandwich norm) --- |
| 111 | let residual = x; |
| 112 | let x = ctx.backend.rms_norm(x, &self.input_layernorm_weight, self.rms_eps) |
| 113 | .map_err(|e| anyhow!("input_layernorm: {e}"))?; |
| 114 | let x = self.attn.forward( |
| 115 | &x, |
| 116 | index_pos, |
| 117 | block_idx, |
| 118 | ctx.cache.as_mut().expect("No cache"), |
| 119 | ).map_err(|e| anyhow!("attn: {e}"))?; |
| 120 | let x = ctx.backend.rms_norm(&x, &self.post_attention_layernorm_weight, self.rms_eps) |
| 121 | .map_err(|e| anyhow!("post_attention_layernorm: {e}"))?; |
| 122 | let x = (x + residual).map_err(|e| anyhow!("attn residual: {e}"))?; |
| 123 | |
| 124 | // --- MLP sublayer (sandwich norm) --- |
| 125 | let residual = &x; |
| 126 | let x = ctx.backend.rms_norm(&x, &self.pre_feedforward_layernorm_weight, self.rms_eps) |
| 127 | .map_err(|e| anyhow!("pre_feedforward_layernorm: {e}"))?; |
| 128 | let x = self.mlp.forward(&x) |
| 129 | .map_err(|e| anyhow!("mlp: {e}"))?; |
| 130 | let x = ctx.backend.rms_norm(&x, &self.post_feedforward_layernorm_weight, self.rms_eps) |
| 131 | .map_err(|e| anyhow!("post_feedforward_layernorm: {e}"))?; |
| 132 | let x = (x + residual).map_err(|e| anyhow!("mlp residual: {e}"))?; |
| 133 | |
| 134 | Ok(x) |
| 135 | } |
| 136 | |
| 137 | async fn forward_mut( |
| 138 | &mut self, |