metrale_model_layers/layers/ops/
token_overlay.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the LoRA token-overlay kernels in
4//! `kernels/gb10/common/token_overlay.cu`:
5//! - [`embed_rowdiff`]: at adapter build, which adapter base rows differ from
6//!   the served embedding table;
7//! - [`embed_overlay_routed`]: in the forward, replace overridden vocab rows
8//!   after the embed gather;
9//! - [`lmhead_overlay_routed`]: in the forward, recompute overridden logit
10//!   columns.
11//!
12//! Arguments are passed by position, so each launcher's argument order must
13//! match its `.cu` signature.
14//!
15//! Owner: model-layers ops (LoRA token overlays).
16//! Invariants: none beyond the types.
17
18use anyhow::Result;
19use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
20use metrale_gpu_runtime::kernel_args::KernelLaunch;
21
22use crate::layers::try_kernel;
23
24/// 2026-09-25: The four token-overlay kernels, resolved once at model
25/// construction with [`try_kernel`]. A missing kernel gives `KernelHandle(0)`:
26/// the forward hooks then return without a launch, and `build_overlay` refuses
27/// an adapter that ships overlay tensors.
28#[derive(Clone, Copy)]
29pub struct OverlayKernels {
30    pub rowdiff: KernelHandle,
31    pub embed_overlay: KernelHandle,
32    pub lmhead_overlay_bf16: KernelHandle,
33    pub lmhead_overlay_f32: KernelHandle,
34}
35
36impl Default for OverlayKernels {
37    /// 2026-09-25: All handles zero, as if no overlay kernel were loaded.
38    /// `KernelHandle` does not derive `Default`, so this is spelled out.
39    fn default() -> Self {
40        Self {
41            rowdiff: KernelHandle(0),
42            embed_overlay: KernelHandle(0),
43            lmhead_overlay_bf16: KernelHandle(0),
44            lmhead_overlay_f32: KernelHandle(0),
45        }
46    }
47}
48
49impl OverlayKernels {
50    pub fn new(gpu: &dyn GpuBackend) -> Self {
51        Self {
52            rowdiff: try_kernel(gpu, "token_overlay", "embed_rowdiff_bf16"),
53            embed_overlay: try_kernel(gpu, "token_overlay", "embed_overlay_routed_bf16"),
54            lmhead_overlay_bf16: try_kernel(gpu, "token_overlay", "lmhead_overlay_routed_bf16"),
55            lmhead_overlay_f32: try_kernel(gpu, "token_overlay", "lmhead_overlay_routed_f32"),
56        }
57    }
58}
59
60/// 2026-09-25: `flags[r] = (max_i |base[r,i] - served[r,i]| > thresh)`, one
61/// thread per row. `base` and `served` are `[rows, h]` BF16; `flags` is
62/// `[rows]` u8.
63#[allow(clippy::too_many_arguments)]
64pub fn embed_rowdiff(
65    gpu: &dyn GpuBackend,
66    kernel: KernelHandle,
67    base: DevicePtr,
68    served: DevicePtr,
69    flags: DevicePtr,
70    rows: u32,
71    h: u32,
72    thresh: f32,
73    stream: u64,
74) -> Result<()> {
75    KernelLaunch::new(gpu, kernel)
76        .grid([rows.div_ceil(256), 1, 1])
77        .block([256, 1, 1])
78        .arg_ptr(base)
79        .arg_ptr(served)
80        .arg_ptr(flags)
81        .arg_u32(rows)
82        .arg_u32(h)
83        .arg_f32(thresh)
84        .launch(stream)
85}
86
87/// 2026-09-25: In-place replacement of overridden vocab rows in `out`
88/// (`[num_tokens, h]` BF16), after the embed gather.
89///
90/// Per row `r`: `s = seq_slot[r]`, or `active` when `seq_slot` is null; skip
91/// when `s < 0`, when `slot_map_tab[s]` is null, or when `ids[r] >= vocab`;
92/// `slot = slot_map_tab[s][ids[r]]`; skip when `slot < 0` or
93/// `slot >= n_tab[s]`; otherwise copy `rows_tab[s][slot]` over `out[r]`.
94/// `ids` is `[num_tokens]` u32; the three tables hold one entry per adapter
95/// slot; `vocab` is the length of each slot map.
96#[allow(clippy::too_many_arguments)]
97pub fn embed_overlay_routed(
98    gpu: &dyn GpuBackend,
99    kernel: KernelHandle,
100    ids: DevicePtr,
101    seq_slot: DevicePtr,
102    active: i32,
103    slot_map_tab: DevicePtr,
104    rows_tab: DevicePtr,
105    n_tab: DevicePtr,
106    out: DevicePtr,
107    num_tokens: u32,
108    h: u32,
109    vocab: u32,
110    stream: u64,
111) -> Result<()> {
112    KernelLaunch::new(gpu, kernel)
113        .grid([num_tokens, 1, 1])
114        .block([256, 1, 1])
115        .arg_ptr(ids)
116        .arg_ptr(seq_slot)
117        .arg_i32(active)
118        .arg_ptr(slot_map_tab)
119        .arg_ptr(rows_tab)
120        .arg_ptr(n_tab)
121        .arg_ptr(out)
122        .arg_u32(h)
123        .arg_u32(vocab)
124        .launch(stream)
125}
126
127/// 2026-09-25: In-place recompute of overridden logit columns in `logits`
128/// (`[m, vocab]`, BF16 or FP32 to match `kernel`). One warp per `(row, j)`,
129/// where `j < max_n_override` indexes the overridden ids of that row's adapter
130/// slot; the kernel skips `j >= n_tab[s]` and ids `>= vocab`. `hidden` is
131/// `[m, h]` BF16; `seq_slot` is `[m]` i32, or null for the uniform `active`
132/// slot.
133#[allow(clippy::too_many_arguments)]
134pub fn lmhead_overlay_routed(
135    gpu: &dyn GpuBackend,
136    kernel: KernelHandle,
137    hidden: DevicePtr,
138    seq_slot: DevicePtr,
139    active: i32,
140    rows_tab: DevicePtr,
141    ids_tab: DevicePtr,
142    n_tab: DevicePtr,
143    logits: DevicePtr,
144    m: u32,
145    max_n_override: u32,
146    h: u32,
147    vocab: u32,
148    stream: u64,
149) -> Result<()> {
150    KernelLaunch::new(gpu, kernel)
151        .grid([m, max_n_override, 1])
152        .block([32, 1, 1])
153        .arg_ptr(hidden)
154        .arg_ptr(seq_slot)
155        .arg_i32(active)
156        .arg_ptr(rows_tab)
157        .arg_ptr(ids_tab)
158        .arg_ptr(n_tab)
159        .arg_ptr(logits)
160        .arg_u32(h)
161        .arg_u32(vocab)
162        .launch(stream)
163}
metrale_model_layers/layers/ops/token_overlay.rs

metrale_model_layers/layers/ops/
token_overlay.rs
