ts_pd/control.rs
1// Copyright (c) 2026 LightSeek Foundation
2//
3// Permission is hereby granted, free of charge, to any person obtaining a copy
4// of this software and associated documentation files (the "Software"), to deal
5// in the Software without restriction, including without limitation the rights
6// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
7// copies of the Software, and to permit persons to whom the Software is
8// furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
14// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
18// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
19// SOFTWARE.
20
21//! The PD transfer peer, and the EPD embedding receiver, as the scheduling
22//! loop sees them: requests register at admission, and every round the
23//! replica agrees on which transfers changed state.
24
25use ts_collective::{Group, Mesh};
26
27/// Where a request's peer node reaches it: the prefill deployment's
28/// bootstrap server and the request's room.
29#[derive(Debug, Clone, PartialEq, Eq)]
30pub struct BootstrapInfo {
31 /// The bootstrap server's host.
32 pub host: String,
33 /// Its port.
34 pub port: u16,
35 /// The room both nodes know the request by.
36 pub room: u64,
37}
38
39/// A D-role request whose KV landed.
40#[derive(Debug, Clone, PartialEq)]
41pub struct RemotePrefill {
42 /// The request.
43 pub rid: String,
44 /// The prefill node's first token: a landing without one is a failed
45 /// transfer, as the request's first decode feeds on it.
46 pub bootstrap_token: i32,
47 /// Prompt tokens the prefill node served from its cache.
48 pub cached_tokens: i64,
49 /// The first token's logprob, if the prefill node sent one.
50 pub bootstrap_logprob: Option<f64>,
51 /// The prefill node's drafter candidates, if any.
52 pub spec_candidate_ids: Option<Vec<i32>>,
53 /// The request's slot here.
54 pub request_pool_index: i32,
55}
56
57/// A transfer whose state the replica agreed changed.
58#[derive(Debug, Clone, PartialEq)]
59pub enum PdEvent {
60 /// The request's peers are connected.
61 Bootstrapped(String),
62 /// The request's transfer failed.
63 Failed(String),
64 /// P role: the request's KV reached the decode node.
65 Succeeded(String),
66 /// D role: the request's KV landed.
67 RemotePrefillDone(RemotePrefill),
68}
69
70/// A PD failure the engine cannot serve past.
71#[derive(Debug, thiserror::Error)]
72pub enum PdError {
73 /// The replica could not agree.
74 #[error(transparent)]
75 Collective(#[from] ts_collective::Error),
76 /// The transfer peer broke.
77 #[error("PD transfer: {0}")]
78 Peer(String),
79}
80
81/// The transfer peer's control face.
82pub trait TransferPeer {
83 /// Track an admitted request, before the scheduler can plan it.
84 fn register(&mut self, rid: &str, bootstrap: &BootstrapInfo) -> Result<(), PdError>;
85
86 /// The transfers whose state changed since the last poll, as every
87 /// member of `replica` sees them: a rank reports the least advanced
88 /// state any member holds. Every member calls this in the same rounds.
89 fn poll(&mut self, mesh: &mut Mesh, replica: &Group) -> Result<Vec<PdEvent>, PdError>;
90
91 /// P role: a committed prompt's usage, sent with its completion: the
92 /// prompt tokens served from the cache and the first token's logprob.
93 fn record_usage(&mut self, rid: &str, cached_tokens: i64, bootstrap_logprob: Option<f64>);
94
95 /// A transfer is still tracked: it may yet read or write the cache.
96 fn has_in_flight(&self) -> bool;
97
98 /// EPD, P role: tell the decode node waiting on room `room` that the
99 /// request failed before it could register (its embeddings never
100 /// arrived), so its receiver fails instead of waiting. Best effort.
101 fn abort_room(&mut self, rid: &str, room: u64);
102}
103
104/// EPD: one multimodal item of a request whose embedding an encode node
105/// computes, as the request names it (`item.encode_handshake`).
106#[derive(Debug, Clone, PartialEq, Eq)]
107pub struct EmbeddingItem {
108 /// The item's position among the request's multimodal items.
109 pub index: usize,
110 /// The encode node's bootstrap server.
111 pub host: String,
112 /// Its port.
113 pub port: u16,
114 /// The room the encode node ships the item's embedding under.
115 pub room: u64,
116 /// The embedding's rows: the item's tokens.
117 pub rows: i64,
118}
119
120/// EPD: where one received item's rows landed, for the data plane to
121/// publish onto the item.
122#[derive(Debug, Clone, PartialEq, Eq)]
123pub struct LandedItem {
124 /// The item's position among the request's multimodal items.
125 pub index: usize,
126 /// The receive slot holding the rows; `None` when this rank received
127 /// none (an empty shard).
128 pub slot: Option<usize>,
129 /// The first row this rank received within the image.
130 pub row_start: i64,
131 /// The rows it received, packed at the slot's start.
132 pub row_count: i64,
133 /// The image's rows.
134 pub rows: i64,
135 /// The rows are this rank's shard: the ranks rebuild the rest.
136 pub sharded: bool,
137}
138
139/// EPD: a request whose items' embeddings all landed.
140#[derive(Debug, Clone, PartialEq, Eq)]
141pub struct LandedRequest {
142 /// The request.
143 pub rid: String,
144 /// Its items, in the request's order.
145 pub items: Vec<LandedItem>,
146}
147
148/// EPD, a prefill node's first stage: receiving the encode nodes'
149/// embeddings of the requests waiting for them, as the scheduling loop
150/// drives it. Every member of the replica stages the same requests in the
151/// same order and settles the same codes.
152pub trait EmbeddingReceiver {
153 /// Start receiving `rid`'s items' embeddings.
154 fn stage(&mut self, rid: &str, items: &[EmbeddingItem]);
155
156 /// Every staged request's state on this rank, in staging order: 0
157 /// failed or timed out, 1 receiving, 2 landed.
158 fn poll(&mut self) -> Vec<i64>;
159
160 /// Act on the replica-agreed `codes` (`poll` order): forget the settled
161 /// requests, and return the landed ones' items for the data plane to
162 /// publish. Their slots stay leased until [`EmbeddingReceiver::release`].
163 fn settle(&mut self, codes: &[i64]) -> Vec<LandedRequest>;
164
165 /// The data plane copied these slots' rows out: they take new receives.
166 fn release(&mut self, slots: &[usize]);
167
168 /// Whether a landed slot still waits for its rows to be copied out.
169 fn awaiting_release(&self) -> bool;
170}
171
172/// EPD, an encode rank: one item's embedding in its send slot, as the
173/// transfer engine reads it.
174#[derive(Debug, Clone, PartialEq, Eq)]
175pub struct EmbeddingChunk {
176 /// The room the prefill ranks receive it under.
177 pub room: u64,
178 /// The rows' address.
179 pub source: u64,
180 /// The image's rows.
181 pub rows: i64,
182 /// Elements per row.
183 pub hidden: i64,
184 /// The rows' dtype (`str(torch.dtype)`).
185 pub dtype: String,
186 /// The rows' bytes.
187 pub nbytes: u64,
188 /// The deepstack rows' address; 0 without deepstack.
189 pub deepstack_source: u64,
190 /// Elements per deepstack row; 0 without deepstack.
191 pub deepstack_width: i64,
192 /// The deepstack rows' bytes.
193 pub deepstack_nbytes: u64,
194}
195
196/// EPD, an encode rank: sending each item's embedding to the prefill ranks
197/// that wait for it, as the scheduling loop drives it.
198pub trait EmbeddingSender {
199 /// `room`'s request arrived; false when another request holds the room.
200 fn begin(&mut self, room: u64) -> bool;
201
202 /// Hand over a begun room's rows: they are written once every prefill
203 /// rank of the block asked, and must stay in place until the room
204 /// concludes.
205 fn send(&mut self, chunk: EmbeddingChunk);
206
207 /// Conclude a begun room failed before its rows were handed over,
208 /// telling the prefill ranks that asked.
209 fn fail(&mut self, room: u64, reason: &str);
210
211 /// The begun rooms that concluded since the last call, sent or failed:
212 /// their rows' slots take new rows.
213 fn concluded(&mut self) -> Vec<u64>;
214}