Skip to main content

rlx_runtime/
capabilities.rs

1// RLX — versatile ML compiler + runtime.
2// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
3//
4// This program is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, version 3.
7//
8// This program is distributed in the hope that it will be useful,
9// but WITHOUT ANY WARRANTY; without even the implied warranty of
10// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
11// GNU General Public License for more details.
12//
13// You should have received a copy of the GNU General Public License
14// along with this program. If not, see <https://www.gnu.org/licenses/>.
15
16//! Declared optional features of an [`crate::ExecutableGraph`].
17//!
18//! The [`ExecutableGraph`](crate::ExecutableGraph) trait grew a long tail of
19//! defaulted methods (MoE residency, GPU handles, KV row feeds, async
20//! pipeline, …). Most backends leave those as no-ops / `false`. Callers that
21//! need to branch on support should prefer:
22//!
23//! 1. [`ExecutableGraph::capabilities`] for a cheap advisory bitmask, then
24//! 2. the concrete method's return value (`false` / `None`) as the source of
25//!    truth — capabilities can lag a backend that forgot to override them.
26//!
27//! Method groups on [`ExecutableGraph`](crate::ExecutableGraph):
28//!
29//! | Group | Methods |
30//! |-------|---------|
31//! | Core run | `set_param`, `run`, `run_raw`, `run_slots`, `arena_ptr`, `finalize_params` |
32//! | Clone | `clone_box` |
33//! | Extent / RNG | `set_active_extent`, `set_rng`, `rng` |
34//! | MoE | `set_moe_resident_experts*`, `enable_moe_topk_capture`, `take_moe_*` |
35//! | Persistent / GPU handles | `bind_handle`, `read_handle`, `bind_gpu_handle`, … |
36//! | KV resident | `register_kv_row_feed`, `feed_kv_row`, `seed_resident_kv_prefix_from`, … |
37//! | Typed I/O | `set_param_typed`, `run_typed` |
38//! | Async pipeline | `commit_no_wait`, `sync_pending`, `run_pipelined` |
39
40/// Advisory feature flags for a compiled executable.
41#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
42pub struct ExecutableCapabilities {
43    /// [`crate::ExecutableGraph::clone_box`] is implemented (not the default panic).
44    pub clone: bool,
45    /// MoE residency / TopK capture hooks are live.
46    pub moe: bool,
47    /// Persistent named handles (`bind_handle` / `read_handle`).
48    pub persistent_handles: bool,
49    /// GPU-resident input handles (`bind_gpu_handle`, feeds, …).
50    pub gpu_handles: bool,
51    /// Device-resident KV row feed / D2D seed path.
52    pub kv_resident: bool,
53    /// Native non-F32 `set_param_typed` / `run_typed` (not F32-only default).
54    pub typed_io: bool,
55    /// `commit_no_wait` / `run_pipelined` do real async work.
56    pub async_pipeline: bool,
57    /// `set_active_extent` is honored for bucketed decode.
58    pub active_extent: bool,
59}
60
61impl ExecutableCapabilities {
62    /// All bits clear — every optional hook is a no-op / unsupported.
63    pub const NONE: Self = Self {
64        clone: false,
65        moe: false,
66        persistent_handles: false,
67        gpu_handles: false,
68        kv_resident: false,
69        typed_io: false,
70        async_pipeline: false,
71        active_extent: false,
72    };
73
74    /// Human-readable list of enabled capability names (for `device_report`).
75    pub fn enabled_names(self) -> Vec<&'static str> {
76        let mut out = Vec::new();
77        if self.clone {
78            out.push("clone");
79        }
80        if self.moe {
81            out.push("moe");
82        }
83        if self.persistent_handles {
84            out.push("persistent_handles");
85        }
86        if self.gpu_handles {
87            out.push("gpu_handles");
88        }
89        if self.kv_resident {
90            out.push("kv_resident");
91        }
92        if self.typed_io {
93            out.push("typed_io");
94        }
95        if self.async_pipeline {
96            out.push("async_pipeline");
97        }
98        if self.active_extent {
99            out.push("active_extent");
100        }
101        out
102    }
103}