rlx-pkg 0.2.16

RLX package format (.rlxp) — flat hybrid mmap packs, optional graph, ZIP/dir
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
// SPDX-License-Identifier: MIT OR Apache-2.0

//! `dist/placement.json` — TP/PP/EP maps.

use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

/// How one tensor is sharded across ranks.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct TensorShard {
    /// Axis along which the tensor is split.
    pub dim: usize,
    /// Ranks that own a slice (in order along `dim`).
    pub ranks: Vec<u32>,
}

/// Expert-parallel placement for one expert id.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct ExpertPlacement {
    /// Primary owning rank.
    pub rank: u32,
    /// Optional replica ranks (hot-expert copies).
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub replicas: Vec<u32>,
}

/// Distribution placement metadata.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
pub struct Placement {
    /// Parallelism kinds present: `dp`, `tp`, `pp`, `ep`.
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub parallelism: Vec<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub world_size: Option<u32>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub topology: Option<String>,
    /// Tensor name → shard map.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub tensors: BTreeMap<String, TensorShard>,
    /// Expert id (string) → placement.
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub experts: BTreeMap<String, ExpertPlacement>,
}

impl Placement {
    /// Ranks that own `tensor`, if listed.
    pub fn ranks_for(&self, tensor: &str) -> Option<&[u32]> {
        self.tensors.get(tensor).map(|s| s.ranks.as_slice())
    }

    /// Whether this pack expects a physical `dist/rank-{r}/` tree.
    pub fn has_tensor_map(&self) -> bool {
        !self.tensors.is_empty() || !self.experts.is_empty()
    }
}