atlas_kernels/
ptx_set.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! The compiled-target declarations `build.rs` emits: what one kernel target
4//! IS, and what it says about itself. Split out of `lib.rs` at the 500-LoC cap.
5//! Exact piecewise move — no logic changed, and the types keep their `crate`
6//! paths (`atlas_kernels::TargetPtxSet`, …) through the `pub use` in `lib.rs`,
7//! which is also what puts them in scope for the `include!`d `target_ptx.rs`
8//! that constructs them.
9//!
10//! Distinct from [`super::query`] and [`super::resolve`], which SELECT among
11//! these declarations; nothing here has behaviour to select with.
12
13use atlas_core::target::KernelTarget;
14
15use super::{ModelBehavior, SamplingPresets};
16
17/// Declares which `(model_type, hidden_size)` pairs a kernel target supports.
18/// Parsed from `[[model_types]]` in MODEL.toml at build time.
19pub struct ModelTypeMatch {
20    pub model_type: &'static str,
21    /// `None` = wildcard (matches any hidden_size not caught by a more specific entry).
22    pub hidden_size: Option<usize>,
23}
24
25/// DFlash speculative-decoding pairing for a target model.
26/// Parsed from `[dflash]` in MODEL.toml at build time. `None` when the
27/// model has no DFlash drafter associated.
28#[derive(Debug, Clone)]
29pub struct DflashConfig {
30    /// HuggingFace id (or local path) of the drafter checkpoint.
31    pub draft_model: &'static str,
32    /// Block size γ (parallel draft tokens per step). Defaults to 16.
33    pub gamma: usize,
34    /// Drafter sliding-window size in tokens. 0 = full attention.
35    pub window_size: usize,
36    /// Token id used to fill the γ "to-be-predicted" positions during
37    /// drafter forward. From the drafter's `dflash_config.mask_token_id`.
38    pub mask_token_id: u32,
39    /// Target-side layer indices to capture intermediate hidden states from
40    /// (shallow-to-deep). The drafter's `fc` projection consumes the stack
41    /// of these hiddens. From the drafter's `dflash_config.target_layer_ids`.
42    pub target_layer_ids: &'static [usize],
43}
44
45/// Kernel modules hyperoptimized for a specific (H, M_q) target.
46///
47/// Each blob is the compiled kernel for one module, emitted uniformly as
48/// `&'static [u8]` by build.rs (`include_bytes!`). NVIDIA PTX is ASCII
49/// text but valid as bytes; SCALE/AMD and Metal produce binary objects.
50/// The runtime registry sniffs text-vs-binary per blob at load time.
51pub struct TargetPtxSet {
52    pub target: KernelTarget,
53    /// The `kernels/<hw>/HARDWARE.toml` `[hardware].arch` this target was
54    /// compiled with, VERBATIM — `sm_90a`, `sm_100a`, `sm_121f`, `gfx1151`.
55    ///
56    /// Additive to [`KernelTarget::arch`], which records the same declaration
57    /// with its feature suffix stripped (`sm_90`, `sm_121`) because that is
58    /// the base SM the target constants, the gate baselines and every existing
59    /// record are keyed by. Both are needed, and they are not interchangeable:
60    ///
61    /// * the base SM is an IDENTITY — "which architecture family is this";
62    /// * this field is the only one a COMPATIBILITY question may be asked of,
63    ///   because the suffix IS the rule (`a` never runs forward onto a later
64    ///   architecture, `f` stays inside one major family, a bare `sm_XY`
65    ///   JIT-compiles forward). Judging the stripped string applies plain
66    ///   forward-compat to PTX that has none, which is how `sm_90a` kernels
67    ///   passed the GPU preflight on a CC 10.0 device and then failed inside
68    ///   `cuModuleLoadData` — the error the preflight exists to replace.
69    ///
70    /// Empty only if a build recorded no architecture. Consumers treat empty
71    /// as "no opinion" and skip, rather than inventing a verdict.
72    pub ptx_arch: &'static str,
73    pub modules: Vec<(&'static str, &'static [u8])>,
74    pub sampling: SamplingPresets,
75    pub behavior: ModelBehavior,
76    pub model_type_matches: Vec<ModelTypeMatch>,
77    /// `[model] match_names` needles from MODEL.toml — case-insensitive
78    /// substrings of the checkpoint reference (HF id / `--model-name` /
79    /// resolved model dir) that identify checkpoints THIS target serves.
80    /// Consulted only to break a tie when several targets declare the same
81    /// `(model_type, hidden_size)` (e.g. qwen3.6-27b vs qwen3.8-27b, whose
82    /// configs are bit-identical); see [`crate::resolve::resolve_target`]. Empty
83    /// for targets that never collide — `build.rs` panics if a colliding
84    /// target omits them.
85    pub match_names: &'static [&'static str],
86    /// DFlash drafter pairing for this model. `None` when the MODEL.toml has
87    /// no `[dflash]` section. Consumed by spark-server when `--dflash` is
88    /// set without an explicit `--draft-model` flag.
89    pub dflash: Option<DflashConfig>,
90    /// `(module, kernel)` pairs this model's kernel files DROPPED by shadowing
91    /// their `common/` namesakes — the kernel exists in `common/` but this
92    /// model's fork of the file does not define it, so it is not compiled here.
93    ///
94    /// Shadowing is whole-file, so a fork that predates a kernel added to
95    /// `common/` silently loses it: `try_kernel` returns handle 0 and whatever
96    /// depends on it fails CLOSED. The startup audit joins this against the
97    /// kernels the model actually looked up, which separates the two classes of
98    /// missing kernel — dropped-by-fork (a build defect) from
99    /// never-built-for-this-architecture (expected, e.g. MLA on a Qwen model).
100    pub shadowed_dropped: &'static [(&'static str, &'static str)],
101    /// `(module, kernel)` lookups this model's dispatch may issue and fail to
102    /// resolve WITHOUT that being an error, declared in the model's MODEL.toml
103    /// `[expected_absent]` with a mandatory stated reason per entry.
104    ///
105    /// The boot audit (`kernel_audit::classify_failures`) fails CLOSED on every
106    /// unresolved lookup that is not in this list, so the list is the entire
107    /// difference between "this model is known to run this way" and "nobody has
108    /// looked". It is TRANSITIONAL: the right fix for a lookup that can never
109    /// resolve is to gate it on config so it is never issued (see
110    /// `qwen3_attention::init_arch_gates`), which removes it from here.
111    pub expected_absent: &'static [(&'static str, &'static str)],
112}
atlas_kernels/ptx_set.rs

atlas_kernels/
ptx_set.rs
