Method forward

cake-core/src/models/common/text_model.rs:266–368

Forward pass through all blocks.

(&mut self, x: &Tensor, idx: usize)

    /// Forward pass through all blocks.
    pub async fn forward(&mut self, x: &Tensor, idx: usize) -> Result<Tensor> {
        let forward_start = std::time::Instant::now();
        let (_batch_size, seq_len) = x.dims2()?;

        let emb_start = std::time::Instant::now();
        let mut x = self.ctx.backend.embedding(x, &self.embed_weight)
            .map_err(|e| anyhow!("error in embedding: {e}"))?;
        // Apply embedding scale if configured (Gemma scales by sqrt(hidden_size)).
        if let Some(scale) = self.ctx.config.as_ref().and_then(|c| c.embed_scale) {
            x = (x * scale as f64)?;
        }
        let emb_elapsed = emb_start.elapsed();

        let num_blocks = self.blocks.len();
        let mut block_idx = 0;
        let mut local_elapsed = std::time::Duration::ZERO;
        let mut local_count: usize = 0;

        while block_idx < num_blocks {
            let curr_block_id = self.blocks[block_idx].ident().to_owned();
            if curr_block_id == "local" {
                let local_start = std::time::Instant::now();
                x = self.blocks[block_idx]
                    .forward_mut(&x, idx, block_idx, &mut self.ctx)
                    .await
                    .map_err(|e| {
                        anyhow!("error in forward operation of local block {block_idx}: {e}")
                    })?;
                local_elapsed += local_start.elapsed();
                local_count += 1;

                block_idx += 1;
            } else {
                // collect all contiguous layers running on the same worker
                let mut batch = vec![];
                let first = block_idx;
                while block_idx < num_blocks && self.blocks[block_idx].ident() == curr_block_id {
                    batch.push((
                        self.blocks[block_idx].layer_name().to_string(),
                        idx,
                        block_idx,
                    ));
                    block_idx += 1;
                }

                let num_layers = batch.len();
                let batch_start = std::time::Instant::now();
                x = self.blocks[first]
                    .forward_batch(&x, batch, &mut self.ctx)
                    .await
                    .map_err(|e| {
                        anyhow!(
                            "error in forward batch for blocks {first}..{block_idx} on {}: {e}",
                            &curr_block_id
                        )
                    })?;
                let batch_elapsed = batch_start.elapsed();
                log::debug!(