Method forward_cached

cake-core/src/models/vibevoice/vae_decoder.rs:159–194

Forward with streaming cache: uses cached context instead of zero-padding.

(&self, x: &Tensor, cache: &mut StreamingConvCache)

    /// Forward with streaming cache: uses cached context instead of zero-padding.
    fn forward_cached(&self, x: &Tensor, cache: &mut StreamingConvCache) -> Result<Tensor> {
        let channels = x.dim(1)?;

        let h = self.backend.rms_norm_channel(x, &self.norm_weight, self.eps)?;
        let (slot, is_first) = cache.take_slot();
        let context = if is_first {
            Tensor::zeros((h.dim(0)?, channels, 6), h.dtype(), h.device())?
        } else {
            cache.get(slot).unwrap().clone()
        };

        // Update cache: last 6 samples of [context, h]
        let h_len = h.dim(2)?;
        if h_len >= 6 {
            cache.set(slot, h.narrow(2, h_len - 6, 6)?);
        } else {
            // h shorter than 6: take from context + h
            let ctx_take = 6 - h_len;
            let ctx_part = context.narrow(2, 6 - ctx_take, ctx_take)?;
            cache.set(slot, Tensor::cat(&[&ctx_part, &h], 2)?);
        }

        // Fused conv reads from [context, h] virtually — no cat allocation
        let h = self.backend.depthwise_conv1d_bias_ctx(
            &context, &h, &self.mixer_weight, &self.mixer_bias, 7, channels,
        )?;
        let x = self.backend.add_scaled(x, &h, &self.gamma)?;

        let h = self.backend.rms_norm_channel(&x, &self.ffn_norm_weight, self.eps)?;
        let h = h.transpose(1, 2)?;
        let h = self.backend.linear_forward(&h, &self.ffn_linear1_weight, self.ffn_linear1_bias.as_ref())?;
        let h = self.backend.gelu(&h)?;
        let h = self.backend.linear_forward(&h, &self.ffn_linear2_weight, self.ffn_linear2_bias.as_ref())?;
        let h = h.transpose(1, 2)?;
        self.backend.add_scaled(&x, &h, &self.ffn_gamma)
    }
}

/// One decoder stage: upsample + N blocks.