metrale_kernels/ptx_set.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The per-target declarations the generated `target_ptx.rs`
4//! constructs: a target's kernel blobs and what its MODEL.toml says about it.
5//! Selection among targets lives in [`super::query`] and [`super::resolve`].
6//!
7//! Owner: kernels crate.
8//! Invariants: none beyond the types.
9
10use metrale_core::target::KernelTarget;
11
12use super::{ModelBehavior, SamplingPresets};
13
14/// 2026-09-25: One `(model_type, hidden_size)` pair a kernel target supports,
15/// from MODEL.toml `[[model_types]]`.
16pub struct ModelTypeMatch {
17 pub model_type: &'static str,
18 /// 2026-09-25: `None` is a wildcard; [`crate::resolve::resolve_target`]
19 /// tries it only when no target declares the exact hidden size.
20 pub hidden_size: Option<usize>,
21}
22
23/// 2026-09-25: A target's DFlash drafter pairing, from MODEL.toml `[dflash]`
24/// (`parse_dflash` in `build_parse.rs`). The section requires `draft_model`;
25/// the other keys default as noted.
26#[derive(Debug, Clone)]
27pub struct DflashConfig {
28 /// 2026-09-25: Drafter checkpoint id or path. metrale-server uses it for
29 /// `--dflash` when no `--draft-model` is given.
30 pub draft_model: &'static str,
31 /// 2026-09-25: Block size γ. Default 16.
32 pub gamma: usize,
33 /// 2026-09-25: Drafter sliding-window size in tokens. Default 0.
34 pub window_size: usize,
35 /// 2026-09-25: `[dflash].mask_token_id`. Default 0.
36 pub mask_token_id: u32,
37 /// 2026-09-25: `[dflash].target_layer_ids`. Default empty.
38 pub target_layer_ids: &'static [usize],
39}
40
41/// 2026-09-25: One compiled kernel target: its [`KernelTarget`] identity,
42/// kernel blobs and MODEL.toml declarations.
43///
44/// Each blob is one compiled module, emitted as `&'static [u8]` by `build.rs`
45/// (`include_bytes!`): PTX text for NVIDIA, binary objects for SCALE, HIP and
46/// Metal.
47pub struct TargetPtxSet {
48 pub target: KernelTarget,
49 /// 2026-09-25: The `kernels/<hw>/HARDWARE.toml` `[hardware].arch` this
50 /// target was compiled with, verbatim (`sm_90a`, `sm_121f`, `gfx1151`).
51 /// [`KernelTarget::arch`] holds the same string with an `a`/`f` suffix
52 /// stripped (`build_arch::kernel_target_arch`). The GPU arch preflight
53 /// judges compatibility on this field, because the suffix changes which
54 /// devices the PTX runs on; it skips an empty value.
55 pub ptx_arch: &'static str,
56 pub modules: Vec<(&'static str, &'static [u8])>,
57 pub sampling: SamplingPresets,
58 pub behavior: ModelBehavior,
59 pub model_type_matches: Vec<ModelTypeMatch>,
60 /// 2026-09-25: MODEL.toml `[model] match_names`: case-insensitive
61 /// substrings of a checkpoint reference (HF id, `--model-name`, resolved
62 /// model dir) that identify checkpoints this target serves. Consulted only
63 /// to break a tie between targets that declare the same
64 /// `(model_type, hidden_size)`; see [`crate::resolve::resolve_target`].
65 /// `build.rs` (`validate_collision_match_names`) panics if a colliding
66 /// target leaves it empty.
67 pub match_names: &'static [&'static str],
68 /// 2026-09-25: DFlash drafter pairing; `None` when MODEL.toml has no
69 /// `[dflash]` section with a `draft_model`.
70 pub dflash: Option<DflashConfig>,
71 /// 2026-09-25: `(module, kernel)` pairs this target's leaf-role files
72 /// drop relative to their `common/` namesakes: the `common/` file defines
73 /// the kernel, the model's file of the same name does not, and a model
74 /// file replaces its namesake whole. The boot kernel audit reports them
75 /// per module.
76 pub shadowed_dropped: &'static [(&'static str, &'static str)],
77 /// 2026-09-25: `(module, kernel)` lookups this model's dispatch may issue
78 /// and fail to resolve without that being an error, from MODEL.toml
79 /// `[expected_absent]`; the build panics on an entry without a reason.
80 /// The boot kernel gate refuses to serve on any other unresolved lookup
81 /// unless `--dangerously-allow-unresolved-kernel-lookups` is set.
82 pub expected_absent: &'static [(&'static str, &'static str)],
83}