metrale_kernels/lib.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3#![deny(warnings)]
4#![deny(clippy::all)]
5
6//! 2026-09-25: The compiled kernel registry: the kernel blobs `build.rs`
7//! embeds per target, and the per-model sampling and behaviour defaults
8//! baked from each target's MODEL.toml.
9//!
10//! Each compiled target is a [`TargetPtxSet`] keyed by its [`KernelTarget`]
11//! `(arch, model, quant)`. `ptx_modules()`, `all_ptx_sets()`,
12//! `TARGET_DEFAULTS` and `TARGET_SM_COUNT` are generated by `build.rs` into
13//! `OUT_DIR/target_ptx.rs` from the `kernels/<hw>/<model>/<quant>/`
14//! directories that `METRALE_TARGET_HW`, `METRALE_TARGET_MODEL` (default `*`)
15//! and `METRALE_TARGET_QUANT` (default `nvfp4`) select.
16//!
17//! Owner: kernels crate.
18//! Invariants:
19//! - `ModelBehavior::default()` and the build-time `[behavior]` parse share
20//! the constants in `behavior_defaults.rs`.
21
22// 2026-09-25: Used by the `include!`d `target_ptx.rs` below, which constructs
23// `KernelTarget { .. }` literals in this module's scope. The
24// `METRALE_SKIP_BUILD=1` stub emits an empty `all_ptx_sets()` and names no
25// target, so the import is unused in that build mode; the allow keeps
26// `#![deny(warnings)]` from failing it.
27#[allow(unused_imports)]
28use metrale_core::target::KernelTarget;
29
30pub mod resolve;
31pub use resolve::{ResolveCandidate, TargetResolveError, ptx_for_config, ptx_for_exact_target};
32
33// 2026-09-25: `[behavior]` defaults shared with the build script, which
34// `include!`s the same file (`build_parse_behavior.rs`), so its parse defaults
35// and `ModelBehavior::default()` read one constant.
36mod behavior_defaults;
37pub use behavior_defaults::{
38 DEFAULT_EFFORT_CAPPED_AT_CEILING, DEFAULT_MAX_INTER_TOOL_PROSE, DEFAULT_MAX_THINKING_BUDGET,
39};
40
41// 2026-09-25: The compiled target's serving defaults, baked from
42// `kernels/<hw>/HARDWARE.toml` `[defaults]`. The type is declared here; the
43// `TARGET_DEFAULTS` and `TARGET_SM_COUNT` consts are appended to the same
44// generated `target_ptx.rs` as the kernel registry, so
45// `METRALE_KERNEL_SET_HASH` covers them too.
46mod target_defaults;
47pub use target_defaults::TargetDefaults;
48
49// 2026-09-25: The paged-decode attention split-K policy. It lives in this
50// crate because two crates need the same answer: the metrale-model-layers
51// dispatch that picks `num_splits`, and the metrale-gpu-runtime buffer arena
52// (`sizes.rs`) that sizes the split-K workspace. A workspace sized by a
53// different rule than the dispatch would let the grid index past it.
54pub mod attn_splitk;
55pub use attn_splitk::{MAX_DECODE_SPLITS, SplitkPolicy};
56
57// 2026-09-25: The generated per-target kernel constants, `ptx_modules()`,
58// `all_ptx_sets()`, `TARGET_DEFAULTS` and `TARGET_SM_COUNT`.
59include!(concat!(env!("OUT_DIR"), "/target_ptx.rs"));
60
61/// 2026-09-25: Hash of the generated `target_ptx.rs` text (`content_hash` in
62/// `build.rs`), emitted as the `METRALE_KERNEL_SET_HASH` rustc-env on both the
63/// compiling and the skip-stub path. Reading it here with `env!` ties this
64/// crate's build to the generated file's content.
65pub const KERNEL_SET_HASH: &str = env!("METRALE_KERNEL_SET_HASH");
66
67/// 2026-09-25: What each kernel target in this binary was compiled from, as
68/// JSON: a map of `"hardware/model/quant"` to `{hash, arch, compiler, flags}`,
69/// computed by `build.rs`'s `closure_attestation`. The benchmark record copies
70/// it (`bench_record.rs`, `with_closure`), so a record attests to the binary's
71/// sources rather than to the working tree.
72///
73/// `{}` when no compiler ran (the `METRALE_SKIP_BUILD` stub emits `{}`) or
74/// when `ComputeTarget::compiler_id` returned `None`.
75pub const TARGET_CLOSURES: &str = match option_env!("METRALE_TARGET_CLOSURES") {
76 Some(json) => json,
77 None => "{}",
78};
79
80/// 2026-09-25: One MODEL.toml `[sampling.<category>]` table, parsed by
81/// `build_parse.rs` (`parse_sampling_presets`). An absent table, or an absent
82/// key, takes the value `SamplingCat::default()` gives in `build.rs`.
83#[derive(Debug, Clone, Copy)]
84pub struct SamplingCategory {
85 pub temperature: f32,
86 pub top_p: f32,
87 pub top_k: u32,
88 pub presence_penalty: f32,
89 pub frequency_penalty: f32,
90 /// 2026-09-25: Repetition penalty. `metrale_sampling::apply_penalties_and_bias`
91 /// divides a positive logit by it, and multiplies a non-positive one, once
92 /// for every occurrence of the token in `SamplingParams::repetition_penalty_window`.
93 /// 1.0 (the parse default) skips it.
94 pub repetition_penalty: f32,
95 /// 2026-09-25: DRY penalty (`metrale_sampling::apply_dry_penalty`, whose
96 /// doc says which token each match penalises): a match of `len` tokens,
97 /// with `len > dry_allowed_length`, subtracts
98 /// `dry_multiplier * dry_base^(len - dry_allowed_length)` from one logit.
99 /// `0.0`, the parse default, disables it.
100 pub dry_multiplier: f32,
101 pub dry_base: f32,
102 pub dry_allowed_length: u32,
103 /// 2026-09-25: LZ penalty (`metrale_sampling::apply_lz_penalty`): a token
104 /// that would complete a 3-, 4- or 5-gram already present in the last 256
105 /// tokens loses `lz_penalty * (n - 2)` per earlier occurrence of that
106 /// n-gram. `0.0`, the parse default, disables it.
107 pub lz_penalty: f32,
108 /// 2026-09-25: Model-declared min-p, or `None` when MODEL.toml is silent.
109 ///
110 /// `Option`, because absence and `0.0` mean different things: `None`
111 /// leaves `--default-min-p` in charge, `Some(0.0)` asks for no min-p.
112 /// At startup the default is `generation_config.json`'s `min_p`, then
113 /// this value, then `--default-min-p` (`resolve_sampling_defaults`). On a
114 /// chat request a request value wins; with `use_sampling_presets_for_core`
115 /// set, this value comes next, ahead of that startup default.
116 pub min_p: Option<f32>,
117 /// 2026-09-25: Model-declared top-n-sigma, or `None` when MODEL.toml is
118 /// silent. Read only on the chat path with `use_sampling_presets_for_core`
119 /// set, where a request value wins and `None` leaves
120 /// `--default-top-n-sigma` in charge; `Some(0.0)` disables the filter.
121 pub top_n_sigma: Option<f32>,
122}
123
124/// 2026-09-25: Model-specific sampling presets from MODEL.toml
125/// `[sampling.*]`, one per category.
126#[derive(Debug, Clone, Copy)]
127pub struct SamplingPresets {
128 pub thinking_text: SamplingCategory,
129 pub thinking_coding: SamplingCategory,
130 pub non_thinking: SamplingCategory,
131 pub tools: SamplingCategory,
132}
133
134impl Default for SamplingPresets {
135 fn default() -> Self {
136 let default_cat = SamplingCategory {
137 temperature: 0.7,
138 top_p: 0.95,
139 top_k: 20,
140 presence_penalty: 0.0,
141 frequency_penalty: 0.0,
142 repetition_penalty: 1.0,
143 dry_multiplier: 0.0,
144 dry_base: 1.75,
145 dry_allowed_length: 2,
146 lz_penalty: 0.0,
147 min_p: None,
148 top_n_sigma: None,
149 };
150 let tools_cat = SamplingCategory {
151 temperature: 0.6,
152 top_p: 0.95,
153 top_k: 20,
154 presence_penalty: 0.0,
155 frequency_penalty: 0.0,
156 repetition_penalty: 1.0,
157 dry_multiplier: 0.0,
158 dry_base: 1.75,
159 dry_allowed_length: 2,
160 lz_penalty: 0.0,
161 min_p: None,
162 top_n_sigma: None,
163 };
164 Self {
165 thinking_text: default_cat,
166 thinking_coding: default_cat,
167 non_thinking: default_cat,
168 tools: tools_cat,
169 }
170 }
171}
172
173/// 2026-09-25: Model-specific behaviour flags from MODEL.toml `[behavior]`,
174/// parsed at build time by `build_parse_behavior.rs`. "Default" below is the
175/// value of an absent key, which equals `ModelBehavior::default()`.
176#[derive(Debug, Clone)]
177pub struct ModelBehavior {
178 /// 2026-09-25: When false, a tool-active request without an explicit
179 /// thinking directive runs with thinking off. Default true.
180 pub thinking_in_tools: bool,
181 /// 2026-09-25: Thinking budget (tokens) for a thinking-on request that
182 /// names none, and the anchor E of the `reasoning_effort` ladder
183 /// (`effort_budget` in metrale-server's `api/chat/thinking.rs`).
184 /// `--max-thinking-budget` outranks it. Default
185 /// [`DEFAULT_MAX_THINKING_BUDGET`].
186 pub max_thinking_budget: u32,
187 /// 2026-09-25: Clamp the `reasoning_effort` ladder at E, so high and
188 /// xhigh resolve to `max_thinking_budget` instead of 2E and 4E. Default
189 /// [`DEFAULT_EFFORT_CAPPED_AT_CEILING`].
190 pub effort_capped_at_ceiling: bool,
191 /// 2026-09-25: Thinking state for a request that carries no thinking
192 /// directive. Any explicit directive overrides it, and
193 /// `--disable-thinking` turns thinking off for every request. Default
194 /// false.
195 pub thinking_default: bool,
196 /// 2026-09-25: FP8 KV calibration token count used when
197 /// `--fp8-kv-calibration-tokens` is not given; 0 disables calibration.
198 pub fp8_kv_calibration_tokens: usize,
199 /// 2026-09-25: KV cache dtype (e.g. "bf16", "fp8") used when
200 /// `--kv-cache-dtype` is not given; empty means the engine default
201 /// (`resolve_kv_dtype_str`).
202 pub default_kv_dtype: &'static str,
203 /// 2026-09-25: Speculative `num_drafts` (K = num_drafts + 1) used when
204 /// `--num-drafts` is not given; 0 means the engine default
205 /// (`resolve_num_drafts`).
206 pub default_num_drafts: u32,
207 /// 2026-09-29: The multi-sequence MTP dispatch cap (the widest batch that speculates) when
208 /// `--mtp-max-seqs` is not given; 0 means the engine default
209 /// (`metrale_model_layers::speculative::resolve_mtp_max_seqs`).
210 pub mtp_max_seqs: u32,
211 /// 2026-09-25: Passed to the chat template as `disable_tool_steering`;
212 /// `nemotron_h.jinja` then omits its `<tool_call>` generation-prompt
213 /// prefix when tools are present. Default false.
214 pub disable_tool_steering: bool,
215 /// 2026-09-25: Skip the `<environment>` working-directory hint that the
216 /// chat path adds to the system message of a tool-active request.
217 /// Default false.
218 pub disable_cwd_hint_injection: bool,
219 /// 2026-09-25: On the chat path, take the default temperature, top-k,
220 /// top-p, min-p and top-n-sigma from the selected MODEL.toml sampling
221 /// category instead of the server-wide defaults. Request values still
222 /// win. Default false.
223 pub use_sampling_presets_for_core: bool,
224 /// 2026-09-25: Tool-call parser name. `--tool-call-parser` outranks it;
225 /// empty means the `tool_defaults.toml` entry for the checkpoint's
226 /// `model_type`.
227 pub tool_call_parser: &'static str,
228 /// 2026-09-25: Starting value of the content-loop watchdog lever
229 /// (`SchedLevers::loop_watchdog`), which `resolve_content_loop_watchdog`
230 /// combines with `METRALE_CONTENT_LOOP_WATCHDOG` and
231 /// `--content-loop-watchdog`. Default false.
232 pub enable_loop_watchdog: bool,
233 /// 2026-09-25: Per-model gate for the thinking-loop watchdog. Default
234 /// true.
235 pub enable_think_loop_watchdog: bool,
236 /// 2026-09-25: When thinking is the only reason an EOS is suppressed,
237 /// close the thinking block instead of discarding the EOS. Default false.
238 pub honor_eos_inside_thinking: bool,
239 /// 2026-09-25: While fewer than this many thinking tokens have been
240 /// emitted, `</think>` gets a -8.0 logit bias; 0 disables it. Default 16.
241 pub min_reasoning_floor_tokens: u32,
242 /// 2026-09-25: Also cap the thinking budget at 90% of the request's
243 /// `max_tokens`; false leaves `max_thinking_budget` (or the request's
244 /// own budget) as the only cap. Default true.
245 pub cap_thinking_at_max_tokens: bool,
246 /// 2026-09-25: Server-side min-p floor, applied as `min_p.max(floor)`
247 /// after request and preset resolution, so it holds even when a client
248 /// sends `min_p = 0`. 0.0 disables it. The qwen3.6-35b-a3b MODEL.toml
249 /// files set 0.05 and record the measurement behind it.
250 pub min_p_floor: f32,
251 /// 2026-09-25: Server-side temperature ceiling, applied as
252 /// `temperature.min(max)` after resolution. 0.0 disables it.
253 pub temperature_max: f32,
254 /// 2026-09-25: Repeat count at which the thinking-loop detector forces
255 /// `</think>`, unless the request sends its own `repetition_detection`.
256 /// Default 3.
257 pub think_loop_min_repeats: u32,
258 /// 2026-09-25: Carried into `WatchdogParams::think_loop_scan_window`;
259 /// `detect_token_loop` ignores that value, so this field has no effect.
260 /// Default 160.
261 pub think_loop_scan_window: u32,
262 /// 2026-09-25: Whether `F2ConfidenceEarlyStop` may arm a forced
263 /// `</think>`. Default true.
264 pub confidence_early_stop: bool,
265 /// 2026-09-25: Consecutive confident thinking tokens after which
266 /// `F2ConfidenceEarlyStop` arms a forced `</think>`. Default 30.
267 pub confidence_run_length: u32,
268 /// 2026-09-25: Fuzzy-repetition tolerance divisor: a `pattern_len` window
269 /// tolerates `max(pattern_len / div, 1)` mismatches. Default 12 (about 8%).
270 pub fuzzy_repeat_tolerance_div: u32,
271 /// 2026-09-25: Cap on content tokens outside tool bodies since the last
272 /// tool call on a tool request; past it the watchdog rolls back or ends
273 /// the response. `--max-inter-tool-prose` and
274 /// `METRALE_MAX_INTER_TOOL_PROSE` outrank it, and 0 disables the guard
275 /// (`resolve_max_inter_tool_prose`). Default
276 /// [`DEFAULT_MAX_INTER_TOOL_PROSE`].
277 pub max_inter_tool_prose: u32,
278 /// 2026-09-25: Cap on content tokens per generation for a sequence with a
279 /// grammar attached; past it the response ends. Sequences without a
280 /// grammar are not capped. Default 100_000.
281 pub max_post_think_content_tokens: u32,
282 /// 2026-09-25: TSCG (arXiv:2605.04107): render the prompt's tool list as
283 /// compiled function signatures instead of JSON schemas. Default false.
284 pub tscg: bool,
285 /// 2026-09-25: On a tool-active request with `tool_choice` other than
286 /// `required`, parse tool calls but do not constrain them with the tool
287 /// grammar. A pinned `--tool-grammar` outranks it. Default false.
288 pub disable_tool_grammar: bool,
289 /// 2026-09-25: When a decode-time watchdog fires, roll the sequence back
290 /// to its last boundary and re-steer instead of ending the response
291 /// (arXiv:2603.27905); false makes the rollback decline. After
292 /// [`crate::ROLLBACK_RESTEER_CAP`] rollbacks of one sequence the rollback
293 /// declines and the watchdog ends the response. Default true.
294 pub rollback_resteer: bool,
295 /// 2026-09-25: Path to a trained repetition-onset (ROM, arXiv:2603.22016)
296 /// head artifact. Nothing loads it: a non-empty value only logs a startup
297 /// warning, and the scheduler's `rom_head` stays `None`. Default empty.
298 pub rom_head: &'static str,
299 /// 2026-09-25: Parsed from `[behavior].tool_retry` (default true), but
300 /// nothing reads it: the chat stream's `tool_retry_enabled` is the
301 /// constant `false`.
302 pub tool_retry: bool,
303 /// 2026-09-25: The chat template's `preserve_thinking` variable. `None`
304 /// leaves it undefined, so the template's own default applies; `Some(_)`
305 /// sets it. A server-level `preserve_thinking` default kwarg outranks it,
306 /// and a request's `chat_template_kwargs.preserve_thinking` wins over
307 /// both.
308 pub preserve_thinking: Option<bool>,
309}
310
311/// 2026-09-25: Most watchdog-triggered rollbacks one sequence may perform;
312/// at this count `rollback_to_boundary` declines and the watchdog ends the
313/// response instead of re-steering.
314pub const ROLLBACK_RESTEER_CAP: u32 = 2;
315
316/// 2026-09-25: Deepest per-sequence decode-rollback ring (boundary SSM-state
317/// snapshots): the default depth, the top of the auto-fit ladder, and the
318/// largest value `--ssm-decode-ring-slots` accepts. A model without SSM
319/// layers gets no ring.
320pub const DECODE_ROLLBACK_RING_SLOTS: usize = 8;
321
322/// 2026-09-25: Domain separator mixed with the model fingerprint to form the
323/// SSM decode-ring cold-tier namespace (`derive_decode_ns_salted` in
324/// metrale-model-engine), which keeps decode keys apart from other key
325/// spaces derived from the same fingerprint.
326pub const DECODE_DOMAIN: u64 = 0xD3C0_DE12_A5B6_C7D8;
327
328impl Default for ModelBehavior {
329 fn default() -> Self {
330 Self {
331 thinking_in_tools: true,
332 max_thinking_budget: DEFAULT_MAX_THINKING_BUDGET,
333 effort_capped_at_ceiling: DEFAULT_EFFORT_CAPPED_AT_CEILING,
334 thinking_default: false,
335 fp8_kv_calibration_tokens: 0,
336 default_kv_dtype: "",
337 default_num_drafts: 0,
338 mtp_max_seqs: 0,
339 disable_tool_steering: false,
340 disable_cwd_hint_injection: false,
341 use_sampling_presets_for_core: false,
342 tool_call_parser: "",
343 enable_loop_watchdog: false,
344 enable_think_loop_watchdog: true,
345 honor_eos_inside_thinking: false,
346 min_reasoning_floor_tokens: 16,
347 cap_thinking_at_max_tokens: true,
348 min_p_floor: 0.0,
349 temperature_max: 0.0,
350 think_loop_min_repeats: 3,
351 think_loop_scan_window: 160,
352 confidence_early_stop: true,
353 confidence_run_length: 30,
354 fuzzy_repeat_tolerance_div: 12,
355 max_inter_tool_prose: DEFAULT_MAX_INTER_TOOL_PROSE,
356 max_post_think_content_tokens: 100_000,
357 tscg: false,
358 disable_tool_grammar: false,
359 rollback_resteer: true,
360 rom_head: "",
361 tool_retry: true,
362 preserve_thinking: None,
363 }
364 }
365}
366
367mod ptx_set;
368pub use ptx_set::{DflashConfig, ModelTypeMatch, TargetPtxSet};
369
370mod query;
371pub use query::{available_targets, ptx_for_model};
372
373#[cfg(test)]
374#[path = "lib_tests.rs"]
375mod tests;