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}