metrale_model_layers/lora/
moe_row_adapter.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: GPU-free MoE LoRA routing: the per-request fold decision
4//! ([`resolve_moe_lora_route`]), the refusal of a decode batch the
5//! single-active fold cannot serve, and the per-row adapter maps the MoE
6//! LoRA kernels read (`< 0` means no fold on that row).
7//!
8//! Owner: model-layers (lora).
9//! Invariants: none beyond the types.
10
11use crate::layer::MoeLoraRoute;
12
13/// 2026-09-25: An error when `route` is `Refuse`. The batched decode entries
14/// (model-engine `decode_batch_compute_main`, `mixed_forward`) call it before
15/// any graph lookup: [`build_moe_row_adapter_decode`] writes a `Refuse` row as
16/// `-1`, so the batch would otherwise run that request without its adapter.
17pub fn ensure_decode_route_servable(route: MoeLoraRoute, path: &str) -> anyhow::Result<()> {
18    anyhow::ensure!(
19        !matches!(route, MoeLoraRoute::Refuse),
20        "MoE LoRA {path}: a sequence routes to a non-active adapter under single-active \
21         phase-1; refusing rather than mis-folding the active adapter. One adapter per batch."
22    );
23    Ok(())
24}
25
26/// 2026-09-25: The MoE LoRA fold decision for one request:
27/// - `!has_moe_lora` ⇒ `Fold`;
28/// - `adapter_slot < 0` (base request) ⇒ `Skip`;
29/// - `adapter_slot == active` ⇒ `Fold`;
30/// - any other slot ⇒ `Refuse`, since only the active adapter is installed.
31///
32/// Model-engine's `moe_lora_route` passes `has_moe_lora` = "a LoRA pool is
33/// loaded" and `active = -1` when none is.
34pub fn resolve_moe_lora_route(adapter_slot: i32, active: i32, has_moe_lora: bool) -> MoeLoraRoute {
35    if !has_moe_lora {
36        return MoeLoraRoute::Fold;
37    }
38    if adapter_slot < 0 {
39        return MoeLoraRoute::Skip;
40    }
41    if adapter_slot == active {
42        MoeLoraRoute::Fold
43    } else {
44        MoeLoraRoute::Refuse
45    }
46}
47
48/// 2026-09-25: The `[total_tokens]` per-row adapter map of a packed prefill
49/// batch. `cu_seqlens_host` is the `[batch + 1]` prefix sum of per-stream token
50/// counts; stream `b`'s `adapter_slots[b]` is written, unchanged, over rows
51/// `cu_seqlens_host[b]..cu_seqlens_host[b + 1]`. Only tests call it today.
52///
53/// `None`, before any row is written, when `cu_seqlens_host` has fewer than two
54/// entries, `adapter_slots` is not `batch` long, the first entry is not 0, or
55/// any boundary is negative or decreasing.
56pub fn build_moe_row_adapter_host(
57    cu_seqlens_host: &[i32],
58    adapter_slots: &[i32],
59) -> Option<Vec<i32>> {
60    if cu_seqlens_host.len() < 2 {
61        return None;
62    }
63    let batch = cu_seqlens_host.len() - 1;
64    if adapter_slots.len() != batch {
65        return None;
66    }
67    if cu_seqlens_host[0] != 0 {
68        return None;
69    }
70    // 2026-09-25: Check every boundary before sizing the map from the last
71    // one, so `[0, 4, 2]` cannot write past it.
72    for b in 0..batch {
73        let start = cu_seqlens_host[b];
74        let end = cu_seqlens_host[b + 1];
75        if start < 0 || end < start {
76            return None;
77        }
78    }
79    let total = cu_seqlens_host[batch];
80    let mut map = vec![-1i32; total as usize];
81    for b in 0..batch {
82        let start = cu_seqlens_host[b];
83        let end = cu_seqlens_host[b + 1];
84        let slot = adapter_slots[b];
85        for row in start..end {
86            map[row as usize] = slot;
87        }
88    }
89    Some(map)
90}
91
92/// 2026-09-25: The `[padded_n]` per-row adapter map of a decode batch, one
93/// token per sequence, which `moe_lora_gather_bgmv.cu` reads as
94/// `row_adapter[row / top_k]` and skips when negative. Row `i` is `active`
95/// when `resolve_moe_lora_route` gives `Fold`, else `-1`; so `Skip`, `Refuse`
96/// and padding rows (`i >= adapter_slots.len()`) are all `-1`, which is not the
97/// attention `seq_slot` convention (there `-1` means the active adapter).
98/// Model-engine uploads it in `upload_moe_row_adapter`.
99pub fn build_moe_row_adapter_decode(
100    adapter_slots: &[i32],
101    padded_n: usize,
102    active: i32,
103    has_moe_lora: bool,
104) -> Vec<i32> {
105    (0..padded_n)
106        .map(|i| match adapter_slots.get(i).copied() {
107            Some(slot) => match resolve_moe_lora_route(slot, active, has_moe_lora) {
108                MoeLoraRoute::Fold => active,
109                MoeLoraRoute::Skip | MoeLoraRoute::Refuse => -1,
110            },
111            None => -1,
112        })
113        .collect()
114}
115
116#[cfg(test)]
117#[path = "moe_row_adapter_tests.rs"]
118mod tests;
metrale_model_layers/lora/moe_row_adapter.rs

metrale_model_layers/lora/
moe_row_adapter.rs
