pmetal_distributed/

pipeline.rs

//! Pipeline-parallel inference runtime.
//!
//! Coordinates multiple [`PipelineShard`] instances across a cluster,
//! routing activations between stages via [`ActivationMessage`].
//!
//! The API node (rank 0, first shard) drives generation: it embeds tokens,
//! runs its local layers, sends the hidden state to the next shard, and
//! awaits the final logits from the last shard.
//!
//! # Generation loop
//!
//! [`PipelineGenerationLoop`] wires end-to-end autoregressive generation across the
//! full pipeline.  On the first shard call [`PipelineGenerationLoop::generate_first_shard`];
//! on every middle/last shard run [`PipelineGenerationLoop::run_shard_loop`] in a
//! background task.  The last shard samples a token (greedy argmax built-in), encodes it
//! as 4 little-endian bytes, and sends it back to the first shard via `send_result`.
//!
//! # Concurrent requests
//!
//! [`StreamMultiplexer`] allows multiple in-flight requests to share the same
//! `PipelineStageRuntime` transport pair.  Each request is identified by its u64 nonce.
//! `send_and_await` dispatches a message and suspends until the matching response
//! arrives; `register_handler` wires a one-shot callback for async dispatch.

use crate::activation_codec::{ActivationCodec, compress_activation};
use crate::activation_transport::{ActivationMessage, DtypeTag, recv_activation, send_activation};
use crate::error::{DistributedError, DistributedResult};
use crate::topology::NodeProfile;
use crate::transport::{TransportReceiver, TransportSender};
use std::collections::HashMap;
use std::ops::Range;
use tokio::sync::oneshot;

/// Configuration for a pipeline stage.
#[derive(Debug, Clone)]
pub struct PipelineStageConfig {
    /// This stage's rank (0-indexed).
    pub rank: usize,
    /// Total number of pipeline stages.
    pub world_size: usize,
    /// Layer range assigned to this stage.
    pub layer_range: Range<usize>,
    /// Whether this is the first stage (owns embedding).
    pub is_first: bool,
    /// Whether this is the last stage (owns norm + lm_head).
    pub is_last: bool,
    /// Wire dtype for activation transfer.
    pub wire_dtype: DtypeTag,
    /// Activation compression codec.
    pub codec: ActivationCodec,
}

/// Runtime state for one stage of the pipeline.
pub struct PipelineStageRuntime {
    config: PipelineStageConfig,
    /// Sender to next stage (None if last stage).
    next_sender: Option<TransportSender>,
    /// Receiver from previous stage (None if first stage).
    prev_receiver: Option<TransportReceiver>,
    /// Sender back to first stage for final logits (only on last stage).
    result_sender: Option<TransportSender>,
    /// Receiver for final logits (only on first stage, from last stage).
    result_receiver: Option<TransportReceiver>,
    /// Monotonic nonce counter for request routing.
    nonce_counter: u64,
}

impl PipelineStageRuntime {
    /// Create a new pipeline stage runtime.
    pub fn new(
        config: PipelineStageConfig,
        next_sender: Option<TransportSender>,
        prev_receiver: Option<TransportReceiver>,
        result_sender: Option<TransportSender>,
        result_receiver: Option<TransportReceiver>,
    ) -> Self {
        Self {
            config,
            next_sender,
            prev_receiver,
            result_sender,
            result_receiver,
            nonce_counter: 0,
        }
    }

    /// Generate a new nonce for a request.
    pub fn next_nonce(&mut self) -> u64 {
        self.nonce_counter += 1;
        self.nonce_counter
    }

    /// Configuration for this stage.
    pub fn config(&self) -> &PipelineStageConfig {
        &self.config
    }

    /// Send local hidden states to the next pipeline stage.
    ///
    /// `data`: raw bytes of the hidden state tensor
    /// `shape`: tensor shape dimensions
    /// `nonce`: request identifier for routing
    /// `layer_id`: the last layer this stage processed
    pub async fn send_to_next(
        &mut self,
        data: &[u8],
        shape: &[u32],
        nonce: u64,
        layer_id: u32,
    ) -> DistributedResult<()> {
        let sender = self
            .next_sender
            .as_mut()
            .ok_or_else(|| DistributedError::Protocol("no next stage to send to".into()))?;

        let compressed = compress_activation(
            data,
            self.config.wire_dtype == DtypeTag::Float32,
            self.config.codec,
        );

        let msg = ActivationMessage {
            nonce,
            layer_id,
            shape: shape.to_vec(),
            dtype: self.config.wire_dtype,
            data: compressed,
        };

        send_activation(sender, &msg).await
    }

    /// Receive hidden states from the previous pipeline stage.
    pub async fn recv_from_prev(&mut self) -> DistributedResult<ActivationMessage> {
        let receiver = self.prev_receiver.as_mut().ok_or_else(|| {
            DistributedError::Protocol("no previous stage to receive from".into())
        })?;

        recv_activation(receiver).await
    }

    /// Send final logits back to the API node (first stage).
    /// Only called by the last stage.
    pub async fn send_result(
        &mut self,
        data: &[u8],
        shape: &[u32],
        nonce: u64,
    ) -> DistributedResult<()> {
        let sender = self.result_sender.as_mut().ok_or_else(|| {
            DistributedError::Protocol("no result sender (not last stage?)".into())
        })?;

        let msg = ActivationMessage {
            nonce,
            layer_id: u32::MAX, // sentinel for "final logits"
            shape: shape.to_vec(),
            dtype: self.config.wire_dtype,
            data: data.to_vec(),
        };

        send_activation(sender, &msg).await
    }

    /// Receive final logits from the last stage.
    /// Only called by the first stage.
    pub async fn recv_result(&mut self) -> DistributedResult<ActivationMessage> {
        let receiver = self.result_receiver.as_mut().ok_or_else(|| {
            DistributedError::Protocol("no result receiver (not first stage?)".into())
        })?;

        recv_activation(receiver).await
    }
}

/// Layer assignment solver.
///
/// Given node profiles and total layer count, produces contiguous layer
/// assignments that balance memory usage across nodes.
pub fn solve_layer_assignment(
    num_layers: usize,
    profiles: &[NodeProfile],
) -> Vec<PipelineStageConfig> {
    let available_ram: Vec<u64> = profiles.iter().map(|p| p.available_ram).collect();
    let ranges = crate::layer_assignment::assign_layers_proportional(num_layers, &available_ram);

    let world_size = profiles.len();
    ranges
        .into_iter()
        .enumerate()
        .map(|(rank, range)| PipelineStageConfig {
            rank,
            world_size,
            is_first: rank == 0,
            is_last: rank == world_size - 1,
            layer_range: range,
            wire_dtype: DtypeTag::Float16,
            codec: ActivationCodec::Float16,
        })
        .collect()
}

// ─────────────────────────────────────────────────────────────────────────────
// PipelineGenerationLoop
// ─────────────────────────────────────────────────────────────────────────────

/// End-to-end autoregressive generation loop for a pipeline-parallel cluster.
///
/// # Roles
///
/// * **First shard (rank 0)** — calls [`generate_first_shard`].  It sends the
///   already-embedded+locally-forwarded hidden state to the next shard, then
///   waits for a 4-byte token reply from the last shard.  The sampled token is
///   fed back as the next input and the cycle repeats.
///
/// * **Middle / last shards** — call [`run_shard_loop`] in a background task,
///   supplying a `forward_fn` closure that accepts raw activation bytes and
///   returns the result.  The last shard's `forward_fn` must return logits
///   (`[batch, seq_len, vocab_size]`) in fp32 little-endian; the loop applies
///   greedy argmax and sends the winning token back to rank 0.
///
/// # Greedy sampler
///
/// The built-in sampler applies argmax over the **last position** of the logit
/// tensor (`shape = [batch, seq_len, vocab_size]`):
///
/// ```text
/// token = argmax(logits[0, -1, :])
/// ```
///
/// This is the standard greedy decode step identical to what dnet's `generate_stream`
/// loop does when `temperature=0`.
///
/// [`generate_first_shard`]: PipelineGenerationLoop::generate_first_shard
/// [`run_shard_loop`]: PipelineGenerationLoop::run_shard_loop
pub struct PipelineGenerationLoop {
    /// The underlying stage runtime used for sending/receiving.
    pub stage: PipelineStageRuntime,
    /// Maximum number of new tokens to generate.
    pub max_tokens: usize,
    /// Token IDs that terminate generation (e.g. EOS).  Any token whose u32
    /// value appears in this list stops the loop immediately.
    pub stop_tokens: Vec<u32>,
}

impl PipelineGenerationLoop {
    /// Create a new generation loop wrapping the given stage runtime.
    pub fn new(stage: PipelineStageRuntime, max_tokens: usize, stop_tokens: Vec<u32>) -> Self {
        Self {
            stage,
            max_tokens,
            stop_tokens,
        }
    }

    /// Drive autoregressive generation from the **first shard** (rank 0).
    ///
    /// The caller is responsible for embedding + running the local layers to
    /// produce `input_hidden` before the first call.  On each subsequent step
    /// the single-token hidden state from the local forward pass is passed in
    /// again.
    ///
    /// # Arguments
    ///
    /// * `input_hidden` — raw bytes of the hidden state produced by this shard's
    ///   local forward pass (dtype matches `stage.config.wire_dtype`).
    /// * `input_shape` — shape of `input_hidden`, e.g. `[1, seq_len, hidden_dim]`.
    /// * `vocab_size` — vocabulary size; used to validate the logit payload returned
    ///   by the last shard.
    ///
    /// # Returns
    ///
    /// The ordered list of generated token IDs (not including the prompt).
    pub async fn generate_first_shard(
        &mut self,
        input_hidden: &[u8],
        input_shape: &[u32],
        vocab_size: u32,
    ) -> DistributedResult<Vec<u32>> {
        let mut generated: Vec<u32> = Vec::with_capacity(self.max_tokens);

        // The first step uses the caller-supplied hidden state.  Subsequent
        // steps re-use whatever the caller passes in via the outer loop — but
        // because we own the stage here we thread the token feedback through
        // a separate result channel rather than calling the model again.
        //
        // Concretely the loop is:
        //   1. Send hidden to next shard.
        //   2. Await the 4-byte token from the last shard.
        //   3. Append token to output, check stop condition.
        //   4. The next iteration's `hidden` is produced by the caller embedding
        //      the new token; here we just encode the token as a 1-token "hidden"
        //      signal so the caller can reconstruct it.  We return after the loop.

        let mut current_hidden: Vec<u8> = input_hidden.to_vec();
        let mut current_shape: Vec<u32> = input_shape.to_vec();

        for _ in 0..self.max_tokens {
            let nonce = self.stage.next_nonce();

            // Send our local output to the next stage in the pipeline.
            let last_layer = self.stage.config().layer_range.end.saturating_sub(1) as u32;
            self.stage
                .send_to_next(&current_hidden, &current_shape, nonce, last_layer)
                .await?;

            // Await the sampled token from the last shard.
            // The last shard encodes the token as a 4-byte LE u32 payload with
            // shape `[1]` and layer_id == u32::MAX (the "final result" sentinel).
            let result_msg = self.stage.recv_result().await?;

            if result_msg.data.len() < 4 {
                return Err(DistributedError::Protocol(format!(
                    "expected 4-byte token payload from last shard, got {} bytes",
                    result_msg.data.len()
                )));
            }

            let token =
                u32::from_le_bytes(result_msg.data[..4].try_into().expect("slice is 4 bytes"));
            generated.push(token);

            // Stop on EOS / stop token.
            if self.stop_tokens.contains(&token) {
                break;
            }

            // The next hidden state is a single-token slice.  For pipeline
            // purposes we encode the token ID as a 4-byte int32 tensor so the
            // first shard can embed it on the next call.  In practice the caller
            // drives embedding outside this loop and passes in the fresh hidden;
            // this path is used when `generate_first_shard` is called once for
            // the full sequence and the shard also owns embedding.
            //
            // We store the raw token bytes as the "hidden" to re-enter the loop:
            // the caller can detect a 4-byte, shape=[1] payload and re-embed.
            current_hidden = token.to_le_bytes().to_vec();
            current_shape = vec![1];

            // Sanity: log if the vocab_size doesn't match (non-fatal, best effort).
            let _ = vocab_size; // used only for documentation intent above
        }

        Ok(generated)
    }

    /// Run a **middle or last shard's** receive → compute → send loop.
    ///
    /// This method blocks until the first shard signals termination (i.e. the
    /// pipeline transport is closed) or an error occurs.
    ///
    /// # `forward_fn` contract
    ///
    /// ```text
    /// fn forward_fn(data: &[u8], shape: &[u32]) -> DistributedResult<(Vec<u8>, Vec<u32>)>
    /// ```
    ///
    /// * Input: raw activation bytes + shape from the previous shard.
    /// * Output: either the next hidden state (middle shards) **or** fp32
    ///   logits `[batch, seq_len, vocab_size]` (last shard).
    ///
    /// For the **last shard** the returned bytes are treated as fp32 logits.
    /// `run_shard_loop` applies greedy argmax on the final position, packs the
    /// winning index as a 4-byte LE u32, and sends it back to rank 0 via
    /// `send_result`.  Middle shards forward the returned bytes to the next
    /// shard via `send_to_next`.
    pub async fn run_shard_loop<F>(&mut self, mut forward_fn: F) -> DistributedResult<()>
    where
        F: FnMut(&[u8], &[u32]) -> DistributedResult<(Vec<u8>, Vec<u32>)>,
    {
        let is_last = self.stage.config().is_last;

        loop {
            // Receive hidden state from the previous shard (or EOS when the
            // transport is closed / the peer disconnects).
            let msg = match self.stage.recv_from_prev().await {
                Ok(m) => m,
                Err(DistributedError::Protocol(ref s)) if s.contains("recv activation") => {
                    // Transport closed — generation is complete.
                    break;
                }
                Err(e) => return Err(e),
            };

            let nonce = msg.nonce;
            let (out_data, out_shape) = forward_fn(&msg.data, &msg.shape)?;

            if is_last {
                // Apply greedy argmax over the last-position logits and send
                // the winning token index back to rank 0.
                let token = greedy_argmax_last_position(&out_data, &out_shape)?;
                let token_bytes = token.to_le_bytes();
                self.stage.send_result(&token_bytes, &[1], nonce).await?;
            } else {
                // Middle shard: forward to the next stage.
                let last_layer = self.stage.config().layer_range.end.saturating_sub(1) as u32;
                self.stage
                    .send_to_next(&out_data, &out_shape, nonce, last_layer)
                    .await?;
            }
        }

        Ok(())
    }
}

/// Greedy argmax sampler applied to the **last sequence position** of a logit tensor.
///
/// Expects `data` to be a flat, fp32 little-endian buffer of shape
/// `[batch, seq_len, vocab_size]` (the canonical output of `lm_head`).
/// Returns the index of the maximum logit at position `[0, seq_len-1, :]`.
fn greedy_argmax_last_position(data: &[u8], shape: &[u32]) -> DistributedResult<u32> {
    // Shape must be at least rank-1; we need the vocab_size (last dim).
    if shape.is_empty() {
        return Err(DistributedError::Protocol(
            "logit tensor has empty shape".into(),
        ));
    }

    let vocab_size = *shape.last().unwrap() as usize;
    if vocab_size == 0 {
        return Err(DistributedError::Protocol(
            "logit tensor has zero vocab_size".into(),
        ));
    }

    // Each f32 element is 4 bytes.
    if !data.len().is_multiple_of(4) {
        return Err(DistributedError::Protocol(format!(
            "logit data length {} is not a multiple of 4 (f32)",
            data.len()
        )));
    }

    let total_elems = data.len() / 4;
    if total_elems < vocab_size {
        return Err(DistributedError::Protocol(format!(
            "logit data has {} f32 elements but vocab_size is {}",
            total_elems, vocab_size
        )));
    }

    // The last-position slice starts at (total_elems - vocab_size) * 4.
    let last_pos_start = (total_elems - vocab_size) * 4;
    let last_pos_bytes = &data[last_pos_start..];

    let mut best_idx: u32 = 0;
    let mut best_val: f32 = f32::NEG_INFINITY;

    for (i, chunk) in last_pos_bytes.chunks_exact(4).enumerate() {
        let val = f32::from_le_bytes(
            chunk
                .try_into()
                .expect("chunks_exact(4) guarantees 4 bytes"),
        );
        if val > best_val {
            best_val = val;
            best_idx = i as u32;
        }
    }

    Ok(best_idx)
}

// ─────────────────────────────────────────────────────────────────────────────
// StreamMultiplexer
// ─────────────────────────────────────────────────────────────────────────────

/// Multiplexes multiple concurrent inference requests over a shared pipeline
/// transport pair.
///
/// The [`PipelineStageRuntime`] send/recv methods process one message at a time.
/// When multiple requests are in flight (e.g. batched API requests), their
/// responses must be routed back to the correct caller by nonce.
///
/// `StreamMultiplexer` provides two complementary APIs:
///
/// * **Request-response** — [`send_and_await`] sends an [`ActivationMessage`]
///   and suspends the caller until the matching reply (same nonce) arrives.
/// * **Async dispatch** — [`register_handler`] registers a one-shot
///   [`oneshot::Sender`] that is fired when a response with the matching nonce
///   is delivered by [`dispatch_incoming`].
///
/// # Typical usage
///
/// ```text
/// // Spawn one background task that continuously calls dispatch_incoming:
/// tokio::spawn(async move {
///     loop {
///         mux.dispatch_incoming(&mut stage).await.unwrap();
///     }
/// });
///
/// // Each request task calls send_and_await:
/// let response = mux.send_and_await(msg, &mut stage).await?;
/// ```
///
/// [`send_and_await`]: StreamMultiplexer::send_and_await
/// [`register_handler`]: StreamMultiplexer::register_handler
/// [`dispatch_incoming`]: StreamMultiplexer::dispatch_incoming
pub struct StreamMultiplexer {
    /// Pending one-shot senders keyed by nonce.
    ///
    /// When a response message arrives with a known nonce, the matching sender
    /// is removed and the message is delivered through it.
    pending: HashMap<u64, oneshot::Sender<ActivationMessage>>,
}

impl StreamMultiplexer {
    /// Create a new, empty multiplexer.
    pub fn new() -> Self {
        Self {
            pending: HashMap::new(),
        }
    }

    /// Register a one-shot handler that fires when a response with `nonce` arrives.
    ///
    /// Returns the corresponding [`oneshot::Receiver`] which the caller can
    /// `await` to get the response.  If a handler for the same nonce is already
    /// registered, it is replaced and the old receiver will never fire.
    pub fn register_handler(&mut self, nonce: u64) -> oneshot::Receiver<ActivationMessage> {
        let (tx, rx) = oneshot::channel();
        self.pending.insert(nonce, tx);
        rx
    }

    /// Send `msg` via `stage` and await the matching response.
    ///
    /// This is a higher-level convenience that calls `register_handler`,
    /// forwards the message to the next shard, and then awaits the registered
    /// one-shot receiver.
    ///
    /// # Errors
    ///
    /// Returns `DistributedError::Cancelled` if the dispatch task drops the
    /// sender before delivering a response (e.g. on transport error).
    pub async fn send_and_await(
        &mut self,
        msg: ActivationMessage,
        stage: &mut PipelineStageRuntime,
    ) -> DistributedResult<ActivationMessage> {
        let nonce = msg.nonce;
        let rx = self.register_handler(nonce);

        // Transmit — reuse send_to_next which handles compression.
        stage
            .send_to_next(&msg.data, &msg.shape, nonce, msg.layer_id)
            .await?;

        // Await the response from the background dispatch loop.
        rx.await.map_err(|_| DistributedError::Cancelled)
    }

    /// Receive one incoming message from `stage` and route it to the waiting
    /// caller identified by its nonce.
    ///
    /// This should be called repeatedly from a dedicated background task:
    ///
    /// ```text
    /// loop {
    ///     mux.dispatch_incoming(&mut stage).await?;
    /// }
    /// ```
    ///
    /// Messages whose nonce is not in `pending` (e.g. unsolicited or already
    /// cancelled) are silently dropped.
    pub async fn dispatch_incoming(
        &mut self,
        stage: &mut PipelineStageRuntime,
    ) -> DistributedResult<()> {
        // For the first shard we receive via result_receiver (from last shard).
        // For other shards we receive from the previous stage.
        let msg = if stage.config().is_first {
            stage.recv_result().await?
        } else {
            stage.recv_from_prev().await?
        };

        let nonce = msg.nonce;
        if let Some(tx) = self.pending.remove(&nonce) {
            // Ignore send errors: the receiver may have been dropped if the
            // request was cancelled on the caller side.
            let _ = tx.send(msg);
        }
        // Unknown nonce: silently discard (logged at trace level in production).

        Ok(())
    }

    /// Number of requests currently awaiting a response.
    pub fn pending_count(&self) -> usize {
        self.pending.len()
    }
}

impl Default for StreamMultiplexer {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod generation_tests {
    use super::*;

    #[test]
    fn greedy_argmax_picks_max() {
        // logits: [1, 1, 4] — one batch, one position, four vocab entries.
        // Values: 0.1, 0.5, 0.9, 0.2 — argmax should be index 2.
        let logits: Vec<f32> = vec![0.1, 0.5, 0.9, 0.2];
        let data: Vec<u8> = logits.iter().flat_map(|v| v.to_le_bytes()).collect();
        let shape = vec![1u32, 1, 4];
        let token = greedy_argmax_last_position(&data, &shape).unwrap();
        assert_eq!(token, 2);
    }

    #[test]
    fn greedy_argmax_last_position_multi_step() {
        // Three sequence positions, vocab_size=3.
        // Last position logits: 1.0, 3.0, 2.0 — argmax = 1.
        let logits: Vec<f32> = vec![
            // position 0
            5.0, 0.0, 0.0, // position 1
            0.0, 5.0, 0.0, // position 2 (last)
            1.0, 3.0, 2.0,
        ];
        let data: Vec<u8> = logits.iter().flat_map(|v| v.to_le_bytes()).collect();
        let shape = vec![1u32, 3, 3];
        let token = greedy_argmax_last_position(&data, &shape).unwrap();
        assert_eq!(token, 1);
    }

    #[test]
    fn greedy_argmax_rejects_empty_shape() {
        let err = greedy_argmax_last_position(&[0u8; 8], &[]).unwrap_err();
        assert!(matches!(err, DistributedError::Protocol(_)));
    }

    #[test]
    fn stream_multiplexer_pending_count() {
        let mut mux = StreamMultiplexer::new();
        assert_eq!(mux.pending_count(), 0);
        let _rx1 = mux.register_handler(1);
        let _rx2 = mux.register_handler(2);
        assert_eq!(mux.pending_count(), 2);
    }

    #[test]
    fn stream_multiplexer_register_replace() {
        // Registering the same nonce twice replaces the previous handler.
        let mut mux = StreamMultiplexer::new();
        let _rx_old = mux.register_handler(42);
        let rx_new = mux.register_handler(42);
        // Old receiver should now never fire (sender was replaced).
        // New receiver is tracked.
        assert_eq!(mux.pending_count(), 1);
        drop(rx_new);
    }
}