(
&self,
x: &Tensor,
cache: &mut super::vae_decoder::StreamingConvCache,
)
| 76 | } |
| 77 | |
| 78 | fn forward_cached( |
| 79 | &self, |
| 80 | x: &Tensor, |
| 81 | cache: &mut super::vae_decoder::StreamingConvCache, |
| 82 | ) -> Result<Tensor> { |
| 83 | let channels = x.dim(1)?; |
| 84 | |
| 85 | let h = self.backend.rms_norm_channel(x, &self.norm_weight, self.eps)?; |
| 86 | let (slot, is_first) = cache.take_slot(); |
| 87 | let context = if is_first { |
| 88 | Tensor::zeros((h.dim(0)?, channels, 6), h.dtype(), h.device())? |
| 89 | } else { |
| 90 | cache.get(slot).unwrap().clone() |
| 91 | }; |
| 92 | |
| 93 | // Update cache: last 6 samples of [context, h] |
| 94 | let h_len = h.dim(2)?; |
| 95 | if h_len >= 6 { |
| 96 | cache.set(slot, h.narrow(2, h_len - 6, 6)?); |
| 97 | } else { |
| 98 | let ctx_take = 6 - h_len; |
| 99 | let ctx_part = context.narrow(2, 6 - ctx_take, ctx_take)?; |
| 100 | cache.set(slot, Tensor::cat(&[&ctx_part, &h], 2)?); |
| 101 | } |
| 102 | |
| 103 | // Fused conv reads from [context, h] virtually — no cat allocation |
| 104 | let h = self.backend.depthwise_conv1d_bias_ctx( |
| 105 | &context, &h, &self.mixer_weight, &self.mixer_bias, 7, channels, |
| 106 | )?; |
| 107 | let x = self.backend.add_scaled(x, &h, &self.gamma)?; |
| 108 | |
| 109 | let h = self.backend.rms_norm_channel(&x, &self.ffn_norm_weight, self.eps)?; |
| 110 | let h = h.transpose(1, 2)?; |
| 111 | let h = self.backend.linear_forward(&h, &self.ffn_linear1_weight, self.ffn_linear1_bias.as_ref())?; |
| 112 | let h = self.backend.gelu(&h)?; |
| 113 | let h = self.backend.linear_forward(&h, &self.ffn_linear2_weight, self.ffn_linear2_bias.as_ref())?; |
| 114 | let h = h.transpose(1, 2)?; |
| 115 | self.backend.add_scaled(&x, &h, &self.ffn_gamma) |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | /// One encoder stage: N blocks (applied after its corresponding downsample). |
no test coverage detected