metrale_model_engine/traits/model/
ep.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: `ModelEp`, one of the supertraits `Model` is made of. Its methods, default
4//! bodies and docs are the ones `Model` declared before the split.
5//!
6//! Owner: model-engine.
7//! Invariants: none beyond the types.
8
9use crate::traits::SequenceState;
10use anyhow::Result;
11
12// 2026-09-26: Only the `ep_worker_step` doc link names it.
13#[cfg(doc)]
14use super::EpCommandFailed;
15
16/// 2026-09-26: The multi-rank (EP) command protocol.
17pub trait ModelEp {
18    /// 2026-09-25: EP worker step, on a rank above 0: receive a `(seq_id, cmd)` from rank 0 and
19    /// run the command in slot `seq_id`. Returns `Ok(false)` when the worker should shut down.
20    ///
21    /// An `Err` carrying [`EpCommandFailed`] means the command ran and failed, and the worker
22    /// stays up; any other `Err` means the receive failed, and the worker exits.
23    ///
24    /// `slots` is sized to `--max-batch-size`; a `seq_id >= slots.len()` fails. Default:
25    /// `Ok(true)`.
26    fn ep_worker_step(&self, _slots: &mut [Option<SequenceState>]) -> Result<bool> {
27        Ok(true)
28    }
29
30    /// 2026-09-25: Whether the model runs the multi-rank command protocol. The scheduler then
31    /// turns off its single-GPU fused paths, the mixed forward among them. Default `false`.
32    fn is_ep(&self) -> bool {
33        false
34    }
35
36    /// 2026-09-25: Send a command word to every worker rank. Default: `Ok(())`.
37    fn ep_broadcast_cmd(&self, _cmd: u32) -> Result<()> {
38        Ok(())
39    }
40
41    /// 2026-09-25: Send a `(seq_id, cmd)` pair to every worker rank, as the first broadcast of a
42    /// command; its follow-up broadcasts use [`Self::ep_broadcast_cmd`]. When
43    /// [`Self::ep_protocol_v2`] is false, `seq_id` is not sent. Default: `Ok(())`.
44    fn ep_broadcast_cmd_for_seq(&self, _seq_id: u32, _cmd: u32) -> Result<()> {
45        Ok(())
46    }
47
48    /// 2026-09-25: Whether the EP protocol carries a `seq_id` with each command
49    /// (`METRALE_EP_PROTOCOL=v2`). Default `false`.
50    fn ep_protocol_v2(&self) -> bool {
51        false
52    }
53
54    /// 2026-09-25: Broadcast a token array to every worker rank in one broadcast. Default: an
55    /// empty `Vec`.
56    fn ep_broadcast_tokens(&self, _tokens: &[u32]) -> Result<Vec<u32>> {
57        Ok(Vec::new())
58    }
59}
metrale_model_engine/traits/model/ep.rs

metrale_model_engine/traits/model/
ep.rs
