Implementing Pipeline Parallelism for Distributed Transformer Inference on Local Hardware Using Socket-Based IPC in Rust — code snippets

snippet-1.txt

use std::convert::TryInto;
use std::slice;

#[repr(C, align(8))]
#[derive(Debug, Copy, Clone)]
pub struct FrameHeader {
    pub magic: [u8; 4],       // Must be b"TNSR"
    pub msg_type: u8,         // 1: Token IDs (prefill), 2: Activations, 3: Logits, 4: Error
    pub shape_rank: u8,
    pub padding: [u8; 2],     // Structural padding to align to 8-byte boundary
    pub step_id: u64,
    pub micro_batch_id: u32,
    pub payload_len: u32,
    pub shape_dims: [u64; 4], // Dynamic dimensions (e.g., [Batch, SeqLen, Hidden, 0])
}

impl FrameHeader {
    pub fn to_bytes(&self) -> &[u8] {
        unsafe {
            slice::from_raw_parts(
                (self as *const Self) as *const u8,
                std::mem::size_of::<Self>(),
            )
        }
    }

    pub fn from_bytes(bytes: &[u8]) -> Option<&Self> {
        if bytes.len() < std::mem::size_of::<Self>() {
            return None;
        }
        let ptr = bytes.as_ptr();
        if ptr.align_offset(8) != 0 {
            // Memory must be 8-byte aligned for safe casting
            return None;
        }
        let header = unsafe { &*(ptr as *const Self) };
        if &header.magic != b"TNSR" {
            return None;
        }
        Some(header)
    }
}

snippet-2.txt

use std::net::{TcpListener, TcpStream};
use std::io::{Read, Write};
use std::thread;

pub fn start_low_latency_receiver(
    addr: &str,
    core_id: usize,
    frame_tx: std::sync::mpsc::Sender<(FrameHeader, Vec<u8>)>,
) -> std::io::Result<()> {
    let listener = TcpListener::bind(addr)?;
    
    thread::spawn(move || {
        // Pin this network loop thread to a dedicated core to prevent scheduling jitter
        if let Some(core) = core_affinity::get_core_ids().and_then(|ids| ids.into_iter().nth(core_id)) {
            core_affinity::set_for_current(core);
        }

        for stream in listener.incoming() {
            if let Ok(mut stream) = stream {
                // Disable Nagle's algorithm to write packets immediately
                let _ = stream.set_nodelay(true);
                let _ = stream.set_read_timeout(Some(std::time::Duration::from_millis(5000)));

                let mut header_buf = vec![0u8; std::mem::size_of::<FrameHeader>()];
                loop {
                    if stream.read_exact(&mut header_buf).is_err() {
                        break; // Connection closed or timeout
                    }
                    
                    let header = match FrameHeader::from_bytes(&header_buf) {
                        Some(h) => *h,
                        None => {
                            eprintln!("Invalid frame magic or alignment error");
                            break;
                        }
                    };

                    let mut payload = vec![0u8; header.payload_len as usize];
                    if stream.read_exact(&mut payload).is_err() {
                        break;
                    }

                    if frame_tx.send((header, payload)).is_err() {
                        break; // Channel disconnected
                    }
                }
            }
        }
    });

    Ok(())
}

snippet-3.txt

use std::sync::Arc;

pub struct Tensor {
    pub data: Vec<f32>,
    pub shape: Vec<usize>,
}

pub struct TransformerLayer {
    pub layer_idx: usize,
}

impl TransformerLayer {
    pub fn forward(
        &self,
        input: &Tensor,
        _step_id: u64,
        _micro_batch_id: u32,
        _global_layer_idx: usize,
        _kv_cache: &mut DistributedKvCache,
    ) -> Result<Tensor, String> {
        // Simulated execution of attention projection and feed-forward MLP networks
        // In production, this invokes candle/ndarray matrix multiplications
        Ok(Tensor {
            data: input.data.clone(),
            shape: input.shape.clone(),
        })
    }
}

pub struct PipelineStage {
    pub stage_index: usize,
    pub num_layers: usize,
    pub layers: Vec<TransformerLayer>,
    pub kv_cache: Arc<parking_lot::RwLock<DistributedKvCache>>,
}

impl PipelineStage {
    pub fn execute_forward(
        &self,
        step_id: u64,
        micro_batch_id: u32,
        input_tensor: Tensor,
    ) -> Result<Tensor, String> {
        let mut x = input_tensor;
        
        // Retrieve KV Cache handle for the local stage
        let mut kv_guard = self.kv_cache.write();
        
        for layer_offset in 0..self.num_layers {
            let global_layer_idx = self.stage_index * self.num_layers + layer_offset;
            
            // Execute layer with locally scoped KV Cache slice
            x = self.layers[layer_offset].forward(
                &x,
                step_id,
                micro_batch_id,
                global_layer_idx,
                &mut *kv_guard,
            )?;
        }
        
        Ok(x)
    }
}

snippet-4.txt

use std::collections::HashMap;

pub struct LayerKvCache {
    pub key_cache: Vec<f32>,   // Flat layout: [MaxSeqLen, HiddenDim]
    pub value_cache: Vec<f32>, // Flat layout: [MaxSeqLen, HiddenDim]
}

pub struct DistributedKvCache {
    pub max_seq_len: usize,
    pub hidden_dim: usize,
    pub num_heads: usize,
    // Maps sequence_key to a vec of LayerKvCaches representing layers in this stage
    pub caches: HashMap<u64, Vec<LayerKvCache>>, 
}

impl DistributedKvCache {
    pub fn update_token(
        &mut self,
        sequence_key: u64,
        layer_idx: usize,
        token_offset: usize,
        new_k: &[f32],
        new_v: &[f32],
    ) -> Result<(), String> {
        let layers = self.caches.entry(sequence_key).or_insert_with(|| {
            (0..32).map(|_| LayerKvCache {
                key_cache: vec![0.0; self.max_seq_len * self.hidden_dim],
                value_cache: vec![0.0; self.max_seq_len * self.hidden_dim],
            }).collect()
        });

        if token_offset >= self.max_seq_len {
            return Err("Sequence length limit exceeded in KV cache".into());
        }

        let layer_cache = &mut layers[layer_idx];
        
        // Copy key-value vectors for the active token to the correct buffer index
        let offset = token_offset * self.hidden_dim;
        layer_cache.key_cache[offset..offset + self.hidden_dim].copy_from_slice(new_k);
        layer_cache.value_cache[offset..offset + self.hidden_dim].copy_from_slice(new_v);

        Ok(())
    }
}

snippet-5.txt

use std::sync::mpsc::Receiver;
use std::net::TcpStream;

pub struct MicroBatchScheduler {
    pub stage: PipelineStage,
    pub rx_inbox: Receiver<(FrameHeader, Vec<u8>)>,
    pub downstream_addr: String,
}

impl MicroBatchScheduler {
    pub fn start_scheduling_loop(&self) -> Result<(), String> {
        let mut downstream = TcpStream::connect(&self.downstream_addr)
            .map_err(|e| format!("Failed to connect to downstream: {}", e))?;
        let _ = downstream.set_nodelay(true);

        while let Ok((header, payload)) = self.rx_inbox.recv() {
            // Reconstruct tensor shape and data from raw bytes
            let input_shape: Vec<usize> = header.shape_dims[0..header.shape_rank as usize]
                .iter()
                .map(|&d| d as usize)
                .collect();
            
            // Cast raw bytes directly into float slice
            let float_len = payload.len() / std::mem::size_of::<f32>();
            let input_data: &[f32] = unsafe {
                slice::from_raw_parts(payload.as_ptr() as *const f32, float_len)
            };

            let input_tensor = Tensor {
                data: input_data.to_vec(),
                shape: input_shape,
            };

            // Compute local forward pass
            let output_tensor = self.stage.execute_forward(
                header.step_id,
                header.micro_batch_id,
                input_tensor,
            )?;

            // Frame output tensor and write downstream
            let mut out_dims = [0u64; 4];
            for (i, &dim) in output_tensor.shape.iter().enumerate() {
                out_dims[i] = dim as u64;
            }

            let response_header = FrameHeader {
                magic: *b"TNSR",
                msg_type: 2, // Activations
                shape_rank: output_tensor.shape.len() as u8,
                padding: [0; 2],
                step_id: header.step_id,
                micro_batch_id: header.micro_batch_id,
                payload_len: (output_tensor.data.len() * std::mem::size_of::<f32>()) as u32,
                shape_dims: out_dims,
            };

            // Raw serialization of header and float payload
            downstream.write_all(response_header.to_bytes())
                .map_err(|e| e.to_string())?;
            
            let payload_slice = unsafe {
                slice::from_raw_parts(
                    output_tensor.data.as_ptr() as *const u8,
                    response_header.payload_len as usize,
                )
            };
            downstream.write_all(payload_slice).map_err(|e| e.to_string())?;
        }

        Ok(())
    }
}

snippet-6.txt

use std::net::TcpListener;
use std::io::Read;

pub struct Coordinator {
    pub stage0_addr: String,
    pub listen_addr: String,
}

impl Coordinator {
    pub fn run_generation_loop(&self, prompt_tokens: Vec<u32>) -> Result<Vec<u32>, String> {
        let mut stage0_stream = TcpStream::connect(&self.stage0_addr)
            .map_err(|e| format!("Stage0 unreachable: {}", e))?;
        stage0_stream.set_nodelay(true).unwrap();

        // Listen for Stage 3 (final output logits)
        let listener = TcpListener::bind(&self.listen_addr).unwrap();
        
        let mut generated_tokens = Vec::new();
        let mut current_input = prompt_tokens;
        let mut step_id = 0u64;

        loop {
            // Write prompt or next generated token to Stage 0
            let payload_len = (current_input.len() * std::mem::size_of::<u32>()) as u32;
            let header = FrameHeader {
                magic: *b"TNSR",
                msg_type: 1, // Token IDs
                shape_rank: 1,
                padding: [0; 2],
                step_id,
                micro_batch_id: 0,
                payload_len,
                shape_dims: [current_input.len() as u64, 0, 0, 0],
            };

            stage0_stream.write_all(header.to_bytes()).unwrap();
            let raw_tokens = unsafe {
                slice::from_raw_parts(current_input.as_ptr() as *const u8, payload_len as usize)
            };
            stage0_stream.write_all(raw_tokens).unwrap();

            // Accept connection from final stage returning raw logits
            let (mut final_stage_stream, _) = listener.accept().unwrap();
            final_stage_stream.set_nodelay(true).unwrap();

            let mut res_header_buf = vec![0u8; std::mem::size_of::<FrameHeader>()];
            final_stage_stream.read_exact(&mut res_header_buf).unwrap();
            
            let res_header = FrameHeader::from_bytes(&res_header_buf).unwrap();
            let mut logits_buf = vec![0u8; res_header.payload_len as usize];
            final_stage_stream.read_exact(&mut logits_buf).unwrap();

            // Decode next token (greedy decoding example)
            let logits: &[f32] = unsafe {
                slice::from_raw_parts(
                    logits_buf.as_ptr() as *const f32,
                    logits_buf.len() / std::mem::size_of::<f32>(),
                )
            };
            
            let next_token = logits
                .iter()
                .enumerate()
                .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                .map(|(idx, _)| idx as u32)
                .unwrap();

            generated_tokens.push(next_token);
            
            if next_token == 2 { // EOS token
                break;
            }

            // Update parameters for next loop pass (autoregressive ring feedback)
            current_input = vec![next_token];
            step_id += 1;
        }

        Ok(generated_tokens)
    }
}