ts_pd/layout/
registration.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//! A Decode rank's messages to a Prefill rank, and the routes they imply:
22//! the registration a Prefill rank plans its route to the Decode rank from
23//! (`prepare_decode_registration`), the pre-allocation carrying a request's
24//! destination manifest, and the route a Decode rank plans from the
25//! bootstrap server's description of the Prefill node (`plan_route`).
26
27use std::collections::BTreeSet;
28
29use bytes::Bytes;
30use serde_json::Value;
31
32use super::{
33    contract::Contract,
34    manifest::Manifest,
35    route::{pipeline_plan, validate_owner_filters, Fragment, OwnerFilters, Planner, MAX_TP_SIZE},
36    LayoutError,
37};
38use crate::{net::parse_uint, wire};
39
40/// A Decode rank's registration (`KVArgsRegisterInfo`).
41#[derive(Debug, Clone)]
42pub struct RegistrationInfo {
43    /// Where its control socket listens.
44    pub endpoint: String,
45    /// Its port.
46    pub port: u16,
47    /// Its transfer session.
48    pub session_id: String,
49    /// Its arena base.
50    pub arena_base: u64,
51    /// Its contract, with its landing addresses.
52    pub layout: Contract,
53    /// Its node's attention TP size.
54    pub tp_size: i64,
55    /// Its attention TP rank.
56    pub tp_rank: i64,
57}
58
59impl RegistrationInfo {
60    /// Parse a registration's eight frames (`KVArgsRegisterInfo.from_zmq`).
61    pub fn from_frames(frames: &[Bytes]) -> Result<Self, LayoutError> {
62        let (endpoint, port, session_id) =
63            wire::registration_route(frames).map_err(|error| LayoutError::new(error.0))?;
64        let pointer: [u8; 8] = frames[4]
65            .as_ref()
66            .try_into()
67            .map_err(|_| LayoutError::new("CachePD slab pointer frame must be exactly 8 bytes"))?;
68        let layout = Contract::from_wire(&frames[5])?;
69        let tp_size = parse_uint(&frames[6], MAX_TP_SIZE as u64)
70            .filter(|&size| size >= 1)
71            .ok_or_else(|| LayoutError::new("CachePD decode_tp_size is invalid"))?;
72        let tp_rank = parse_uint(&frames[7], tp_size - 1)
73            .ok_or_else(|| LayoutError::new("CachePD decode_tp_rank is invalid"))?;
74        // struct.pack("Q"): native order.
75        let arena_base = u64::from_ne_bytes(pointer);
76        if arena_base == 0 {
77            return Err(LayoutError::new("CachePD slab pointer must be nonzero"));
78        }
79        Ok(Self {
80            endpoint,
81            port,
82            session_id,
83            arena_base,
84            layout,
85            tp_size: tp_size as i64,
86            tp_rank: tp_rank as i64,
87        })
88    }
89}
90
91/// A Decode rank's pre-allocation of one request (`TransferInfo`).
92#[derive(Debug, Clone, PartialEq, Eq)]
93pub struct Preallocation {
94    /// The request's room.
95    pub room: u64,
96    /// The Decode rank's transfer session.
97    pub session_id: String,
98    /// Its destination manifest; `None` from a rank that only joins the
99    /// status agreement.
100    pub manifest: Option<Manifest>,
101}
102
103impl Preallocation {
104    /// Parse a pre-allocation's three frames (`TransferInfo.from_zmq`).
105    pub fn from_frames(frames: &[Bytes]) -> Result<Self, LayoutError> {
106        let (room, session_id) =
107            wire::preallocation_route(frames).map_err(|error| LayoutError::new(error.0))?;
108        let manifest = if frames[2].is_empty() {
109            None
110        } else {
111            Some(Manifest::from_wire(&frames[2])?)
112        };
113        Ok(Self {
114            room,
115            session_id,
116            manifest,
117        })
118    }
119}
120
121/// What a Prefill rank plans its routes and copies with.
122#[derive(Debug, Clone)]
123pub struct PrefillLayout {
124    /// This rank's contract, with its landing addresses: the stage's
125    /// resident fields only under a pipeline.
126    pub local: Contract,
127    /// The logical contract every rank registers with the bootstrap server:
128    /// the complete model's, without addresses.
129    pub wire: Contract,
130    /// This rank's arena base.
131    pub arena_base: u64,
132    /// The attention TP size within a stage.
133    pub tp_size: i64,
134    /// This rank's attention TP rank.
135    pub tp_rank: i64,
136    /// This rank's pipeline stage.
137    pub pp_rank: i64,
138    /// Every stage's resident fields.
139    pub fields_by_stage: Vec<Vec<String>>,
140}
141
142/// The route a registration asks of this Prefill rank.
143#[derive(Debug, Clone, PartialEq, Eq)]
144pub struct PlannedRegistration {
145    /// The fragments this rank sends; empty copies every field whole.
146    pub fragments: Vec<Fragment>,
147    /// Its decision for every sharded group.
148    pub owner_filters: OwnerFilters,
149    /// This rank sends no data and only joins the status agreement.
150    pub is_dummy: bool,
151    /// The Decode ranks a room's pre-allocations come from.
152    pub expected_decode_ranks: BTreeSet<i64>,
153}
154
155/// Plan the route a Decode rank's registration asks of this Prefill rank
156/// (`prepare_decode_registration`), against the logical layout the Decode
157/// rank sees on the wire; the stage's field placement keeps only the
158/// fields this rank's arena holds.
159pub fn plan_registration(
160    layout: &PrefillLayout,
161    registration: &RegistrationInfo,
162) -> Result<PlannedRegistration, LayoutError> {
163    let stage_fields: BTreeSet<String> = usize::try_from(layout.pp_rank)
164        .ok()
165        .and_then(|stage| layout.fields_by_stage.get(stage))
166        .ok_or_else(|| LayoutError::new("this rank's pipeline stage has no field placement"))?
167        .iter()
168        .cloned()
169        .collect();
170    let planner = Planner::new(
171        layout.tp_size,
172        registration.tp_size,
173        &layout.wire,
174        &registration.layout,
175        Some(&stage_fields),
176    )?;
177    let route = planner.plan_for_decode_rank(registration.tp_rank)?;
178    let served = planner
179        .decode_ranks_by_prefill_rank()
180        .get(&layout.tp_rank)
181        .cloned()
182        .ok_or_else(|| LayoutError::new("this rank is outside the Prefill TP size"))?;
183    let planned = if let Some(fragments) = route.fragments_by_prefill_rank.get(&layout.tp_rank) {
184        // A KVP-sharded source is never idle: every rank of the chosen
185        // subgroup serves the Decode rank with the blocks it owns.
186        PlannedRegistration {
187            fragments: fragments.clone(),
188            owner_filters: route.owner_filters_by_prefill_rank[&layout.tp_rank].clone(),
189            is_dummy: false,
190            expected_decode_ranks: served,
191        }
192    } else if registration.tp_rank == 0 && served.is_empty() {
193        // A rendezvous-only rank sends nothing of any group.
194        PlannedRegistration {
195            fragments: Vec::new(),
196            owner_filters: planner
197                .sharded_group_ids()
198                .into_iter()
199                .map(|group| (group, None))
200                .collect(),
201            is_dummy: true,
202            expected_decode_ranks: BTreeSet::from([0]),
203        }
204    } else {
205        return Err(LayoutError::new(
206            "CachePD Decode registration targets the wrong Prefill TP rank",
207        ));
208    };
209    // The sender applies the route's owner-filter decisions as given, so
210    // they are checked against the local cache groups once, here.
211    validate_owner_filters(
212        &layout.local.group_specs,
213        &planned.fragments,
214        &planned.owner_filters,
215    )?;
216    Ok(planned)
217}
218
219/// The Prefill ranks a Decode rank's requests read from, planned from the
220/// bootstrap server's description of the Prefill node.
221#[derive(Debug, Clone, PartialEq, Eq)]
222pub struct PrefillRoute {
223    /// The ranks to register with and send each request's manifest to,
224    /// sorted.
225    pub target_tp_ranks: Vec<i64>,
226    /// Those of them that send no data and only join the Prefill node's
227    /// status agreement; they get an empty manifest.
228    pub dummy_tp_ranks: Vec<i64>,
229    /// The ranks whose completion a request waits for, sorted.
230    pub expected_prefill_ranks: Vec<i64>,
231    /// The Prefill node's attention DP size; a room picks its DP group.
232    pub prefill_dp_size: i64,
233}
234
235/// A JSON integer, or a string holding one (Python's `int(value)`).
236fn integer(value: &Value, name: &str) -> Result<i64, LayoutError> {
237    value
238        .as_i64()
239        .or_else(|| value.as_str().and_then(|text| text.trim().parse().ok()))
240        .ok_or_else(|| LayoutError::new(format!("the Prefill description's {name} is invalid")))
241}
242
243/// Plan a Decode rank's route from the bootstrap server's Prefill
244/// description (its `engine_rank=-1` JSON; `HostDecodePeer.plan_route`).
245pub fn plan_route(
246    prefill_info: &str,
247    decode_layout: &Contract,
248    decode_tp_size: i64,
249    decode_tp_rank: i64,
250) -> Result<PrefillRoute, LayoutError> {
251    let payload: Value = serde_json::from_str(prefill_info)
252        .map_err(|_| LayoutError::new("the Prefill description is not JSON"))?;
253    let field = |key: &str| {
254        payload
255            .get(key)
256            .ok_or_else(|| LayoutError::new(format!("the Prefill description misses {key:?}")))
257    };
258    let layout_wire = match payload.get("cache_layout") {
259        None | Some(Value::Null) => {
260            return Err(LayoutError::new(
261                "Cache-transfer decode connected to a non-cache-transfer prefill",
262            ))
263        }
264        Some(Value::String(wire)) if wire.is_ascii() => wire,
265        Some(_) => {
266            return Err(LayoutError::new(
267                "the Prefill layout must be an ASCII string",
268            ))
269        }
270    };
271    let tp_size = integer(field("prefill_tp_size")?, "prefill_tp_size")?;
272    let dp_size = integer(field("prefill_dp_size")?, "prefill_dp_size")?;
273    let pp_size = match payload.get("prefill_pp_size") {
274        None => 1,
275        Some(value) => integer(value, "prefill_pp_size")?,
276    };
277    let fields_by_stage = field("cache_fields_by_stage")?
278        .as_array()
279        .ok_or_else(|| LayoutError::new("the Prefill field placement must be a list"))?
280        .iter()
281        .map(|stage| {
282            stage
283                .as_array()
284                .ok_or_else(|| LayoutError::new("a Prefill stage's fields must be a list"))?
285                .iter()
286                .map(|field| {
287                    field
288                        .as_str()
289                        .map(str::to_owned)
290                        .ok_or_else(|| LayoutError::new("a Prefill field id must be a string"))
291                })
292                .collect::<Result<Vec<_>, _>>()
293        })
294        .collect::<Result<Vec<_>, _>>()?;
295    let prefill_layout = Contract::from_wire(layout_wire.as_bytes())?;
296    if dp_size <= 0 || pp_size <= 0 {
297        return Err(LayoutError::new(
298            "the Prefill DP and PP sizes must be positive",
299        ));
300    }
301    // The registered world spans the stages; the TP width is per stage.
302    let prefill_tp_size = tp_size / (dp_size * pp_size);
303    let (plan, dummy) = pipeline_plan(
304        prefill_tp_size,
305        decode_tp_size,
306        decode_tp_rank,
307        &prefill_layout,
308        decode_layout,
309        &fields_by_stage,
310    )?;
311    let expected = plan.target_prefill_ranks();
312    let mut targets: Vec<i64> = expected.iter().chain(&dummy).copied().collect();
313    targets.sort_unstable();
314    Ok(PrefillRoute {
315        target_tp_ranks: targets,
316        dummy_tp_ranks: dummy,
317        expected_prefill_ranks: expected,
318        prefill_dp_size: dp_size,
319    })
320}
ts_pd/layout/registration.rs

ts_pd/layout/
registration.rs
