rlx_pkg/
placement.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3// SPDX-License-Identifier: MIT OR Apache-2.0
4
5//! `dist/placement.json` — TP/PP/EP maps.
6
7use serde::{Deserialize, Serialize};
8use std::collections::BTreeMap;
9
10/// How one tensor is sharded across ranks.
11#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
12pub struct TensorShard {
13    /// Axis along which the tensor is split.
14    pub dim: usize,
15    /// Ranks that own a slice (in order along `dim`).
16    pub ranks: Vec<u32>,
17}
18
19/// Expert-parallel placement for one expert id.
20#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
21pub struct ExpertPlacement {
22    /// Primary owning rank.
23    pub rank: u32,
24    /// Optional replica ranks (hot-expert copies).
25    #[serde(default, skip_serializing_if = "Vec::is_empty")]
26    pub replicas: Vec<u32>,
27}
28
29/// Distribution placement metadata.
30#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
31pub struct Placement {
32    /// Parallelism kinds present: `dp`, `tp`, `pp`, `ep`.
33    #[serde(default, skip_serializing_if = "Vec::is_empty")]
34    pub parallelism: Vec<String>,
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub world_size: Option<u32>,
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub topology: Option<String>,
39    /// Tensor name → shard map.
40    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
41    pub tensors: BTreeMap<String, TensorShard>,
42    /// Expert id (string) → placement.
43    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
44    pub experts: BTreeMap<String, ExpertPlacement>,
45}
46
47impl Placement {
48    /// Ranks that own `tensor`, if listed.
49    pub fn ranks_for(&self, tensor: &str) -> Option<&[u32]> {
50        self.tensors.get(tensor).map(|s| s.ranks.as_slice())
51    }
52
53    /// Whether this pack expects a physical `dist/rank-{r}/` tree.
54    pub fn has_tensor_map(&self) -> bool {
55        !self.tensors.is_empty() || !self.experts.is_empty()
56    }
57}
rlx_pkg/placement.rs

rlx_pkg/
placement.rs
