Skip to main content

rlx_runtime/
capabilities.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//! Declared optional features of an [`crate::ExecutableGraph`].
6//!
7//! The [`ExecutableGraph`](crate::ExecutableGraph) trait grew a long tail of
8//! defaulted methods (MoE residency, GPU handles, KV row feeds, async
9//! pipeline, …). Most backends leave those as no-ops / `false`. Callers that
10//! need to branch on support should prefer:
11//!
12//! 1. `ExecutableGraph::capabilities` for a cheap advisory bitmask, then
13//! 2. the concrete method's return value (`false` / `None`) as the source of
14//!    truth — capabilities can lag a backend that forgot to override them.
15//!
16//! Method groups on [`ExecutableGraph`](crate::ExecutableGraph):
17//!
18//! | Group | Methods |
19//! |-------|---------|
20//! | Core run | `set_param`, `run`, `run_raw`, `run_slots`, `arena_ptr`, `finalize_params` |
21//! | Clone | `clone_box` |
22//! | Extent / RNG | `set_active_extent`, `set_rng`, `rng` |
23//! | MoE | `set_moe_resident_experts*`, `enable_moe_topk_capture`, `take_moe_*` |
24//! | Persistent / GPU handles | `bind_handle`, `read_handle`, `bind_gpu_handle`, … |
25//! | KV resident | `register_kv_row_feed`, `feed_kv_row`, `seed_resident_kv_prefix_from`, … |
26//! | Typed I/O | `set_param_typed`, `run_typed` |
27//! | Async pipeline | `commit_no_wait`, `sync_pending`, `run_pipelined` |
28
29/// Advisory feature flags for a compiled executable.
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
31pub struct ExecutableCapabilities {
32    /// [`crate::ExecutableGraph::clone_box`] is implemented (not the default panic).
33    pub clone: bool,
34    /// `set_moe_resident_experts` (and the per-layer form) pin experts on the
35    /// device instead of streaming them per token.
36    ///
37    /// This is the half that needs a flag: those setters return `()`, so a
38    /// caller has no other way to learn the backend ignored them. The
39    /// instrumentation hooks (`enable_moe_topk_capture`,
40    /// `take_moe_residency_stats`) return `Option`/`bool` and report for
41    /// themselves — they are CPU-only and are NOT implied by this flag.
42    pub moe: bool,
43    /// Persistent named handles (`bind_handle` / `read_handle`).
44    pub persistent_handles: bool,
45    /// GPU-resident input handles (`bind_gpu_handle`, feeds, …).
46    pub gpu_handles: bool,
47    /// Device-resident KV row feed / D2D seed path.
48    pub kv_resident: bool,
49    /// Native non-F32 `set_param_typed` / `run_typed` (not F32-only default).
50    pub typed_io: bool,
51    /// `commit_no_wait` / `run_pipelined` do real async work.
52    pub async_pipeline: bool,
53    /// `set_active_extent` is honored for bucketed decode.
54    pub active_extent: bool,
55}
56
57impl ExecutableCapabilities {
58    /// All bits clear — every optional hook is a no-op / unsupported.
59    pub const NONE: Self = Self {
60        clone: false,
61        moe: false,
62        persistent_handles: false,
63        gpu_handles: false,
64        kv_resident: false,
65        typed_io: false,
66        async_pipeline: false,
67        active_extent: false,
68    };
69
70    /// Human-readable list of enabled capability names (for `device_report`).
71    pub fn enabled_names(self) -> Vec<&'static str> {
72        let mut out = Vec::new();
73        if self.clone {
74            out.push("clone");
75        }
76        if self.moe {
77            out.push("moe");
78        }
79        if self.persistent_handles {
80            out.push("persistent_handles");
81        }
82        if self.gpu_handles {
83            out.push("gpu_handles");
84        }
85        if self.kv_resident {
86            out.push("kv_resident");
87        }
88        if self.typed_io {
89            out.push("typed_io");
90        }
91        if self.async_pipeline {
92            out.push("async_pipeline");
93        }
94        if self.active_extent {
95            out.push("active_extent");
96        }
97        out
98    }
99}