eredu_runtime/
expert.rs

1//! Runtime ownership boundary for routed expert acquisition and residency.
2
3use eredu_nn::{
4    DistributedNeuralBackend, GroupSelection, GroupedGatedProductOperator, GroupedNeuralBackend,
5    GroupedRelu2Operator, Tensor, TensorParallelGroupedOutput,
6};
7
8use crate::ExpertPass;
9use crate::{ActivationObserver, ParameterBankKey, RoutingObservation};
10
11/// Mechanism-only lookup of one grouped operator in an addressable parameter bank.
12pub trait AddressableGatedProductBank<B>
13where
14    B: GroupedNeuralBackend,
15{
16    /// Bank lookup or construction failure.
17    type Error;
18
19    /// Resolves one generic bank key and exact grouped construction specification.
20    fn acquire(
21        &mut self,
22        key: ParameterBankKey,
23        spec: &eredu_nn::GroupedGatedProductSpec,
24        context: &<B::Tensor as Tensor>::Context,
25    ) -> Result<&mut B::GatedProductGroups, Self::Error>;
26}
27
28/// One architecture route batch submitted to a runtime expert provider.
29pub struct RoutedExpertRequest<'a, T> {
30    /// Global decoder layer requesting experts.
31    pub layer: usize,
32    /// Flattened token rows submitted to the selected experts.
33    pub input: &'a T,
34    /// Backend-native selected expert IDs, scores, and weights.
35    pub routes: &'a GroupSelection<T>,
36    /// Whether this route batch belongs to prefill or decode.
37    pub pass: ExpertPass,
38}
39
40/// Provider result that distinguishes complete outputs from rank-local TP work.
41pub enum RoutedExpertTensorParallelOutput<T> {
42    /// Provider already completed every required collective and bias addition.
43    Complete(T),
44    /// Caller must all-sum `reducible`, then add `post_reduce` exactly once.
45    Partial(TensorParallelGroupedOutput<T>),
46}
47
48/// Completes one rank-local expert output with one all-sum and one post-bias add.
49pub fn reduce_tensor_parallel_expert_output<B>(
50    output: TensorParallelGroupedOutput<B::Tensor>,
51    parallel: &B::ParallelContext,
52    context: &<B::Tensor as Tensor>::Context,
53) -> Result<B::Tensor, eredu_nn::Error>
54where
55    B: GroupedNeuralBackend + DistributedNeuralBackend,
56{
57    let reduced = B::sum_parallel(output.reducible().clone(), parallel, context)?;
58    match output.post_reduce().cloned() {
59        Some(bias) => reduced.add(&bias, context),
60        None => Ok(reduced),
61    }
62}
63
64/// Combines two rank-local expert partials without introducing another collective.
65pub fn combine_tensor_parallel_expert_outputs<B>(
66    left: TensorParallelGroupedOutput<B::Tensor>,
67    right: TensorParallelGroupedOutput<B::Tensor>,
68    context: &<B::Tensor as Tensor>::Context,
69) -> Result<TensorParallelGroupedOutput<B::Tensor>, eredu_nn::Error>
70where
71    B: GroupedNeuralBackend,
72{
73    let post_reduce = match (left.post_reduce().cloned(), right.post_reduce().cloned()) {
74        (Some(left), Some(right)) => Some(left.add(&right, context)?),
75        (Some(bias), None) | (None, Some(bias)) => Some(bias),
76        (None, None) => None,
77    };
78    Ok(TensorParallelGroupedOutput::new(
79        left.reducible().add(right.reducible(), context)?,
80        post_reduce,
81    ))
82}
83
84/// Combines routed/shared provider outputs while requiring one coherent TP mode.
85pub fn combine_routed_expert_tensor_parallel<B>(
86    left: RoutedExpertTensorParallelOutput<B::Tensor>,
87    right: RoutedExpertTensorParallelOutput<B::Tensor>,
88    context: &<B::Tensor as Tensor>::Context,
89) -> Result<RoutedExpertTensorParallelOutput<B::Tensor>, eredu_nn::Error>
90where
91    B: GroupedNeuralBackend,
92{
93    match (left, right) {
94        (
95            RoutedExpertTensorParallelOutput::Complete(left),
96            RoutedExpertTensorParallelOutput::Complete(right),
97        ) => Ok(RoutedExpertTensorParallelOutput::Complete(
98            left.add(&right, context)?,
99        )),
100        (
101            RoutedExpertTensorParallelOutput::Partial(left),
102            RoutedExpertTensorParallelOutput::Partial(right),
103        ) => combine_tensor_parallel_expert_outputs::<B>(left, right, context)
104            .map(RoutedExpertTensorParallelOutput::Partial),
105        _ => Err(eredu_nn::Error::backend(
106            "provider mixed complete and rank-local expert outputs in one block",
107        )),
108    }
109}
110
111/// Completes a provider TP result while preserving provider-owned collectives.
112pub fn reduce_routed_expert_tensor_parallel<B>(
113    output: RoutedExpertTensorParallelOutput<B::Tensor>,
114    parallel: &B::ParallelContext,
115    context: &<B::Tensor as Tensor>::Context,
116) -> Result<B::Tensor, eredu_nn::Error>
117where
118    B: GroupedNeuralBackend + DistributedNeuralBackend,
119{
120    match output {
121        RoutedExpertTensorParallelOutput::Complete(output) => Ok(output),
122        RoutedExpertTensorParallelOutput::Partial(output) => {
123            reduce_tensor_parallel_expert_output::<B>(output, parallel, context)
124        }
125    }
126}
127
128/// Runtime boundary for resident or independently cached routed experts.
129///
130/// Implementations own identity ordering, acquisition, leases, chunking,
131/// budgets, and residency reports. They keep every lease alive until the
132/// backend-native routed result is safe to return. The backend retains tensor
133/// storage, transfers, compact-bank construction, and execution kernels.
134pub trait RoutedExpertProvider<B>
135where
136    B: GroupedNeuralBackend,
137{
138    /// Provider-specific acquisition or execution failure.
139    type Error;
140
141    /// Executes one typed route batch while retaining its acquired resources.
142    fn forward_grouped(
143        &mut self,
144        resident_bank: &mut B::GatedProductGroups,
145        request: RoutedExpertRequest<'_, B::Tensor>,
146        context: &<B::Tensor as Tensor>::Context,
147    ) -> Result<B::Tensor, Self::Error>;
148
149    /// Executes one ReLU-squared route batch through the same residency boundary.
150    fn forward_relu2_routed(
151        &mut self,
152        resident_bank: &mut B::Relu2Groups,
153        request: RoutedExpertRequest<'_, B::Tensor>,
154        context: &<B::Tensor as Tensor>::Context,
155    ) -> Result<B::Tensor, Self::Error>;
156}
157
158/// Additive provider mechanism for tensor-parallel grouped partials.
159pub trait TensorParallelRoutedExpertProvider<B>: RoutedExpertProvider<B>
160where
161    B: GroupedNeuralBackend,
162{
163    /// Executes a rank-local gated-product contribution.
164    fn forward_grouped_tensor_parallel(
165        &mut self,
166        resident_bank: &mut B::GatedProductGroups,
167        request: RoutedExpertRequest<'_, B::Tensor>,
168        partitions: usize,
169        context: &<B::Tensor as Tensor>::Context,
170    ) -> Result<RoutedExpertTensorParallelOutput<B::Tensor>, Self::Error>;
171
172    /// Executes a rank-local ReLU-squared contribution.
173    fn forward_relu2_routed_tensor_parallel(
174        &mut self,
175        resident_bank: &mut B::Relu2Groups,
176        request: RoutedExpertRequest<'_, B::Tensor>,
177        partitions: usize,
178        context: &<B::Tensor as Tensor>::Context,
179    ) -> Result<RoutedExpertTensorParallelOutput<B::Tensor>, Self::Error>;
180}
181
182/// Stable routing metadata supplied by an architecture composition at one
183/// canonical unit boundary.
184#[derive(Debug, Clone, Eq, PartialEq)]
185pub struct RoutedObservationPoint {
186    path: String,
187    expert_count: i32,
188}
189
190impl RoutedObservationPoint {
191    /// Creates one routed observation point.
192    pub fn new(path: impl Into<String>, expert_count: i32) -> Self {
193        Self {
194            path: path.into(),
195            expert_count,
196        }
197    }
198
199    /// Returns the stable routed-module path.
200    pub fn path(&self) -> &str {
201        &self.path
202    }
203
204    /// Returns the total number of routed experts.
205    pub const fn expert_count(&self) -> i32 {
206        self.expert_count
207    }
208}
209
210/// Failure from either canonical expert execution or its observation hook.
211#[derive(Debug)]
212pub enum ObservedExpertProviderError<P, O> {
213    /// The wrapped provider rejected or failed the expert request.
214    Provider(P),
215    /// The observer rejected the normalized routing event.
216    Observer(O),
217}
218
219impl<P, O> std::fmt::Display for ObservedExpertProviderError<P, O>
220where
221    P: std::fmt::Display,
222    O: std::fmt::Display,
223{
224    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
225        match self {
226            Self::Provider(error) => write!(formatter, "routed expert provider failed: {error}"),
227            Self::Observer(error) => write!(formatter, "routed expert observer failed: {error}"),
228        }
229    }
230}
231
232impl<P, O> std::error::Error for ObservedExpertProviderError<P, O>
233where
234    P: std::error::Error + 'static,
235    O: std::error::Error + 'static,
236{
237}
238
239/// Decorates a routed provider with normalized routing observation.
240///
241/// The decorator sees the exact request and output of canonical provider
242/// execution. It therefore adds observation without reimplementing a model
243/// family's block, routing, shape, or residency lifecycle. Tensor-parallel
244/// requests are delegated without an event because their provider result may
245/// still require an architecture-owned reduction before it is observable.
246pub struct ObservedExpertProvider<'a, P, O: ?Sized, E> {
247    provider: &'a mut P,
248    observer: &'a mut O,
249    point: RoutedObservationPoint,
250    error: std::marker::PhantomData<fn() -> E>,
251}
252
253impl<'a, P, O: ?Sized, E> ObservedExpertProvider<'a, P, O, E> {
254    /// Wraps `provider` for one canonical routed module invocation.
255    pub fn new(provider: &'a mut P, observer: &'a mut O, point: RoutedObservationPoint) -> Self {
256        Self {
257            provider,
258            observer,
259            point,
260            error: std::marker::PhantomData,
261        }
262    }
263
264    fn observe<T, ObservationError>(
265        &mut self,
266        routes: &eredu_nn::GroupSelection<T>,
267        output: &T,
268    ) -> Result<(), ObservationError>
269    where
270        O: ActivationObserver<T, ObservationError>,
271    {
272        self.observer.observe_routing(RoutingObservation {
273            path: self.point.path(),
274            selected_experts: routes.group_indices(),
275            selected_scores: routes.selected_scores(),
276            coefficients: routes.coefficients(),
277            routed_output: output,
278            local_routed_output: None,
279            reduced_routed_output: None,
280            shared_output: None,
281            combined_output: None,
282            expert_count: self.point.expert_count(),
283        })
284    }
285}
286
287impl<B, P, O, E> RoutedExpertProvider<B> for ObservedExpertProvider<'_, P, O, E>
288where
289    B: GroupedNeuralBackend,
290    P: RoutedExpertProvider<B>,
291    O: ActivationObserver<B::Tensor, E> + ?Sized,
292{
293    type Error = ObservedExpertProviderError<P::Error, E>;
294
295    fn forward_grouped(
296        &mut self,
297        resident_bank: &mut B::GatedProductGroups,
298        request: RoutedExpertRequest<'_, B::Tensor>,
299        context: &<B::Tensor as Tensor>::Context,
300    ) -> Result<B::Tensor, Self::Error> {
301        let routes = request.routes;
302        let output = self
303            .provider
304            .forward_grouped(resident_bank, request, context)
305            .map_err(ObservedExpertProviderError::Provider)?;
306        self.observe(routes, &output)
307            .map_err(ObservedExpertProviderError::Observer)?;
308        Ok(output)
309    }
310
311    fn forward_relu2_routed(
312        &mut self,
313        resident_bank: &mut B::Relu2Groups,
314        request: RoutedExpertRequest<'_, B::Tensor>,
315        context: &<B::Tensor as Tensor>::Context,
316    ) -> Result<B::Tensor, Self::Error> {
317        let routes = request.routes;
318        let output = self
319            .provider
320            .forward_relu2_routed(resident_bank, request, context)
321            .map_err(ObservedExpertProviderError::Provider)?;
322        self.observe(routes, &output)
323            .map_err(ObservedExpertProviderError::Observer)?;
324        Ok(output)
325    }
326}
327
328impl<B, P, O, E> TensorParallelRoutedExpertProvider<B> for ObservedExpertProvider<'_, P, O, E>
329where
330    B: GroupedNeuralBackend,
331    P: TensorParallelRoutedExpertProvider<B>,
332    O: ActivationObserver<B::Tensor, E> + ?Sized,
333{
334    fn forward_grouped_tensor_parallel(
335        &mut self,
336        resident_bank: &mut B::GatedProductGroups,
337        request: RoutedExpertRequest<'_, B::Tensor>,
338        partitions: usize,
339        context: &<B::Tensor as Tensor>::Context,
340    ) -> Result<RoutedExpertTensorParallelOutput<B::Tensor>, Self::Error> {
341        self.provider
342            .forward_grouped_tensor_parallel(resident_bank, request, partitions, context)
343            .map_err(ObservedExpertProviderError::Provider)
344    }
345
346    fn forward_relu2_routed_tensor_parallel(
347        &mut self,
348        resident_bank: &mut B::Relu2Groups,
349        request: RoutedExpertRequest<'_, B::Tensor>,
350        partitions: usize,
351        context: &<B::Tensor as Tensor>::Context,
352    ) -> Result<RoutedExpertTensorParallelOutput<B::Tensor>, Self::Error> {
353        self.provider
354            .forward_relu2_routed_tensor_parallel(resident_bank, request, partitions, context)
355            .map_err(ObservedExpertProviderError::Provider)
356    }
357}
358
359/// Provider for a fully resident expert bank.
360#[derive(Debug, Default, Clone, Copy)]
361pub struct ResidentExpertProvider;
362
363impl<B> RoutedExpertProvider<B> for ResidentExpertProvider
364where
365    B: GroupedNeuralBackend,
366{
367    type Error = eredu_nn::Error;
368
369    fn forward_grouped(
370        &mut self,
371        resident_bank: &mut B::GatedProductGroups,
372        request: RoutedExpertRequest<'_, B::Tensor>,
373        context: &<B::Tensor as Tensor>::Context,
374    ) -> Result<B::Tensor, Self::Error> {
375        resident_bank.forward_grouped(request.input, request.routes, context)
376    }
377
378    fn forward_relu2_routed(
379        &mut self,
380        resident_bank: &mut B::Relu2Groups,
381        request: RoutedExpertRequest<'_, B::Tensor>,
382        context: &<B::Tensor as Tensor>::Context,
383    ) -> Result<B::Tensor, Self::Error> {
384        resident_bank.forward_grouped(request.input, request.routes, context)
385    }
386}
387
388impl<B> TensorParallelRoutedExpertProvider<B> for ResidentExpertProvider
389where
390    B: eredu_nn::TensorParallelGroupedNeuralBackend,
391{
392    fn forward_grouped_tensor_parallel(
393        &mut self,
394        resident_bank: &mut B::GatedProductGroups,
395        request: RoutedExpertRequest<'_, B::Tensor>,
396        partitions: usize,
397        context: &<B::Tensor as Tensor>::Context,
398    ) -> Result<RoutedExpertTensorParallelOutput<B::Tensor>, Self::Error> {
399        B::gated_product_groups_tensor_parallel(
400            resident_bank,
401            request.input,
402            request.routes,
403            partitions,
404            context,
405        )
406        .map(RoutedExpertTensorParallelOutput::Partial)
407    }
408
409    fn forward_relu2_routed_tensor_parallel(
410        &mut self,
411        resident_bank: &mut B::Relu2Groups,
412        request: RoutedExpertRequest<'_, B::Tensor>,
413        partitions: usize,
414        context: &<B::Tensor as Tensor>::Context,
415    ) -> Result<RoutedExpertTensorParallelOutput<B::Tensor>, Self::Error> {
416        B::relu2_groups_tensor_parallel(
417            resident_bank,
418            request.input,
419            request.routes,
420            partitions,
421            context,
422        )
423        .map(RoutedExpertTensorParallelOutput::Partial)
424    }
425}
eredu_runtime/expert.rs

eredu_runtime/
expert.rs
