atlas_kernels/resolve.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Kernel-target resolution: which compiled target serves a checkpoint.
4//!
5//! Split out of `lib.rs` so the selection RULES are pure functions over
6//! plain declarations — unit-testable on a GPU-free host where
7//! `ATLAS_SKIP_BUILD=1` leaves `all_ptx_sets()` empty.
8//!
9//! ## Why `(model_type, hidden_size)` alone is not enough
10//!
11//! Qwen3.8-27B is architecturally identical to Qwen3.6-27B: same
12//! `model_type` (`qwen3_5`), same `hidden_size` (5120), same every numeric
13//! config field — the checkpoints differ only in weights. Two kernel
14//! targets therefore declare the SAME exact `(qwen3_5, 5120)` pair, and the
15//! historical resolver (`.find()` over build-order-sorted targets) would
16//! have silently picked whichever sorted first. A silent wrong pick
17//! mis-serves the MLPerf-edge flagship's sampling presets and behavior
18//! flags, so ambiguity must never resolve by iteration order.
19//!
20//! ## The rules
21//!
22//! 1. Exact `(model_type, Some(hidden_size))` declarations beat wildcard
23//! `(model_type, None)` declarations (unchanged).
24//! 2. Within the winning tier, if exactly ONE target (by name) matches,
25//! it is selected (unchanged — covers every non-colliding model).
26//! 3. If SEVERAL differently-named targets match, the tie is broken by
27//! the checkpoint reference: each colliding target declares explicit
28//! `match_names` needles in its MODEL.toml (`[model] match_names`),
29//! and a candidate survives when any needle is a case-insensitive
30//! substring of any reference (HF id, `--model-name`, resolved model
31//! dir). Exactly one survivor wins — and the selection is explicit,
32//! because the needles are declared per target, not inferred.
33//! 4. Anything else — zero or multiple survivors — is
34//! [`TargetResolveError::Ambiguous`]. It never falls through to the
35//! wildcard tier and never picks by order.
36//! 5. `--kernel-target <name>` pins resolution to a named target,
37//! bypassing the tie-break — but the pinned target must still declare
38//! compatibility with the checkpoint's `(model_type, hidden_size)`,
39//! otherwise [`TargetResolveError::PinIncompatible`]. This is the
40//! escape hatch for checkpoints whose references carry no identity
41//! (e.g. `--model-from-path /model`).
42//!
43//! `build.rs` enforces at compile time that every set of differently-named
44//! targets sharing a `(model_type, hidden_size)` declaration carries
45//! explicit `match_names`, so rule 3 can never reach a colliding target
46//! with nothing declared.
47
48use crate::{ModelTypeMatch, TargetPtxSet};
49
50/// The resolution-relevant slice of one compiled target. Borrowed views so
51/// tests can drive the rules with synthetic declarations and production
52/// wraps `TargetPtxSet`s without copying module blobs.
53pub struct ResolveCandidate<'a> {
54 /// Kernel-target directory name (`KernelTarget::model`), e.g.
55 /// `"qwen3.8-27b"`. Multi-quant builds repeat a name with different
56 /// quants; same-name candidates are never ambiguous with each other
57 /// (the downstream quant-compat gate arbitrates quant).
58 pub name: &'a str,
59 /// `[[model_types]]` declarations from MODEL.toml.
60 pub type_matches: &'a [ModelTypeMatch],
61 /// `[model] match_names` needles from MODEL.toml. Empty when the
62 /// target never collides (build.rs enforces presence on collision).
63 pub match_names: &'a [&'a str],
64}
65
66/// Why resolution could not choose a target. Every variant is a hard error
67/// at the call site — resolution must never fall back to iteration order.
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub enum TargetResolveError {
70 /// More than one differently-named target claims the checkpoint's
71 /// `(model_type, hidden_size)` and the reference tie-break did not
72 /// leave exactly one.
73 Ambiguous {
74 model_type: String,
75 hidden_size: usize,
76 /// `"exact"` or `"wildcard"` — which declaration tier collided.
77 tier: &'static str,
78 /// Distinct target names in the colliding tier, with their needles.
79 candidates: Vec<(String, Vec<String>)>,
80 /// The subset whose needles matched a reference (empty = none did).
81 matched: Vec<String>,
82 /// The checkpoint references that were searched.
83 model_refs: Vec<String>,
84 },
85 /// `--kernel-target` named a target this binary did not compile.
86 PinNotFound { pin: String, available: Vec<String> },
87 /// `--kernel-target` named a compiled target that does not declare
88 /// support for the checkpoint's `(model_type, hidden_size)` — serving
89 /// would run another architecture's kernels.
90 PinIncompatible {
91 pin: String,
92 model_type: String,
93 hidden_size: usize,
94 declared: Vec<String>,
95 },
96}
97
98impl std::fmt::Display for TargetResolveError {
99 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
100 match self {
101 Self::Ambiguous {
102 model_type,
103 hidden_size,
104 tier,
105 candidates,
106 matched,
107 model_refs,
108 } => {
109 let cands = candidates
110 .iter()
111 .map(|(n, needles)| format!("{n} (match_names: {needles:?})"))
112 .collect::<Vec<_>>()
113 .join(", ");
114 let outcome = if matched.is_empty() {
115 "no checkpoint reference names any of them".to_string()
116 } else {
117 format!("the references match several of them: {matched:?}")
118 };
119 write!(
120 f,
121 "AMBIGUOUS kernel target: {} compiled targets declare {tier} support for \
122 (model_type '{model_type}', hidden_size {hidden_size}) — [{cands}] — and \
123 {outcome} (references searched: {model_refs:?}). Refusing to pick by build \
124 order. Fix: serve with a model id/path that contains exactly one target's \
125 match_names needle, pin explicitly with --kernel-target <name>, or build \
126 single-target with ATLAS_TARGET_MODEL=<name>.",
127 candidates.len(),
128 )
129 }
130 Self::PinNotFound { pin, available } => write!(
131 f,
132 "--kernel-target '{pin}' does not name a compiled kernel target \
133 (available: {available:?})"
134 ),
135 Self::PinIncompatible {
136 pin,
137 model_type,
138 hidden_size,
139 declared,
140 } => write!(
141 f,
142 "--kernel-target '{pin}' is compiled but declares no support for this \
143 checkpoint's (model_type '{model_type}', hidden_size {hidden_size}) — it \
144 declares {declared:?}. Serving another architecture's kernels would be \
145 garbage; refusing."
146 ),
147 }
148 }
149}
150
151impl std::error::Error for TargetResolveError {}
152
153/// Case-insensitive substring: does any declared needle appear in any
154/// checkpoint reference? References are HF ids, `--model-name` values, or
155/// resolved model directories — all of which normally embed the model name.
156fn needles_hit(match_names: &[&str], refs_lower: &[String]) -> bool {
157 match_names.iter().any(|needle| {
158 let n = needle.to_lowercase();
159 !n.is_empty() && refs_lower.iter().any(|r| r.contains(&n))
160 })
161}
162
163/// Distinct candidate names in first-seen order for a set of indices.
164fn distinct_names<'a>(candidates: &[ResolveCandidate<'a>], idxs: &[usize]) -> Vec<&'a str> {
165 let mut names: Vec<&str> = Vec::new();
166 for &i in idxs {
167 if !names.contains(&candidates[i].name) {
168 names.push(candidates[i].name);
169 }
170 }
171 names
172}
173
174/// Resolve which candidate serves `(model_type, hidden_size)` for a
175/// checkpoint identified by `model_refs`. Returns the index of the winning
176/// candidate, `Ok(None)` when nothing declares the pair at all, and
177/// [`TargetResolveError::Ambiguous`] when a collision cannot be broken to
178/// exactly one target name.
179pub fn resolve_target(
180 candidates: &[ResolveCandidate<'_>],
181 model_type: &str,
182 hidden_size: usize,
183 model_refs: &[&str],
184) -> Result<Option<usize>, TargetResolveError> {
185 let refs_lower: Vec<String> = model_refs.iter().map(|r| r.to_lowercase()).collect();
186
187 let tiers: [(&'static str, Option<usize>); 2] =
188 [("exact", Some(hidden_size)), ("wildcard", None)];
189 for (tier, want_hidden) in tiers {
190 let idxs: Vec<usize> = candidates
191 .iter()
192 .enumerate()
193 .filter(|(_, c)| {
194 c.type_matches
195 .iter()
196 .any(|m| m.model_type == model_type && m.hidden_size == want_hidden)
197 })
198 .map(|(i, _)| i)
199 .collect();
200 if idxs.is_empty() {
201 continue;
202 }
203 let names = distinct_names(candidates, &idxs);
204 if names.len() == 1 {
205 // Single target name (possibly several quant variants — the
206 // quant-compat gate downstream arbitrates those, as before).
207 return Ok(Some(idxs[0]));
208 }
209 // Collision: break the tie on declared match_names vs references.
210 let matched: Vec<&str> = names
211 .iter()
212 .copied()
213 .filter(|n| {
214 idxs.iter().any(|&i| {
215 candidates[i].name == *n && needles_hit(candidates[i].match_names, &refs_lower)
216 })
217 })
218 .collect();
219 if let [winner] = matched.as_slice() {
220 let idx = idxs
221 .iter()
222 .copied()
223 .find(|&i| candidates[i].name == *winner)
224 .expect("winner name came from idxs");
225 return Ok(Some(idx));
226 }
227 // Zero or several survivors: hard error. Deliberately does NOT
228 // fall through to the wildcard tier — a checkpoint that exact-
229 // matches colliding targets must be disambiguated, not quietly
230 // downgraded to a wildcard target.
231 return Err(TargetResolveError::Ambiguous {
232 model_type: model_type.to_string(),
233 hidden_size,
234 tier,
235 candidates: names
236 .iter()
237 .map(|n| {
238 let needles = idxs
239 .iter()
240 .filter(|&&i| candidates[i].name == *n)
241 .flat_map(|&i| candidates[i].match_names.iter().map(|s| s.to_string()))
242 .collect();
243 (n.to_string(), needles)
244 })
245 .collect(),
246 matched: matched.iter().map(|n| n.to_string()).collect(),
247 model_refs: model_refs.iter().map(|r| r.to_string()).collect(),
248 });
249 }
250 Ok(None)
251}
252
253/// Resolve a `--kernel-target` pin: the named target wins unconditionally
254/// over the tie-break, but must exist and must declare the checkpoint's
255/// `(model_type, hidden_size)` (exact or wildcard).
256pub fn resolve_pinned(
257 candidates: &[ResolveCandidate<'_>],
258 pin: &str,
259 model_type: &str,
260 hidden_size: usize,
261) -> Result<usize, TargetResolveError> {
262 let pinned: Vec<usize> = candidates
263 .iter()
264 .enumerate()
265 .filter(|(_, c)| c.name.eq_ignore_ascii_case(pin))
266 .map(|(i, _)| i)
267 .collect();
268 if pinned.is_empty() {
269 let all: Vec<usize> = (0..candidates.len()).collect();
270 return Err(TargetResolveError::PinNotFound {
271 pin: pin.to_string(),
272 available: distinct_names(candidates, &all)
273 .into_iter()
274 .map(String::from)
275 .collect(),
276 });
277 }
278 let compatible = pinned.iter().copied().find(|&i| {
279 candidates[i].type_matches.iter().any(|m| {
280 m.model_type == model_type
281 && (m.hidden_size.is_none() || m.hidden_size == Some(hidden_size))
282 })
283 });
284 compatible.ok_or_else(|| TargetResolveError::PinIncompatible {
285 pin: pin.to_string(),
286 model_type: model_type.to_string(),
287 hidden_size,
288 declared: pinned
289 .iter()
290 .flat_map(|&i| candidates[i].type_matches.iter())
291 .map(|m| format!("({}, {:?})", m.model_type, m.hidden_size))
292 .collect(),
293 })
294}
295
296/// Find the PTX module set matching a checkpoint.
297///
298/// Matching rules (full statement + rationale in this module):
299/// 1. Exact match on `(model_type, Some(hidden_size))` beats wildcard
300/// `(model_type, None)`.
301/// 2. When several differently-named targets declare the same pair (the
302/// configs of e.g. Qwen3.6-27B and Qwen3.8-27B are bit-identical), the
303/// tie is broken by matching each target's declared `match_names`
304/// needles against `model_refs` (HF id, `--model-name`, resolved model
305/// dir) — and a tie that does not break to exactly one target is
306/// `Err(TargetResolveError::Ambiguous)`, never a build-order pick.
307/// 3. `pinned_target` (`--kernel-target`) bypasses the tie-break but must
308/// name a compiled target that declares the `(model_type, hidden_size)`.
309/// 4. `Ok(None)` if no compiled target declares the pair at all.
310pub fn ptx_for_config(
311 model_type: &str,
312 hidden_size: usize,
313 model_refs: &[&str],
314 pinned_target: Option<&str>,
315) -> Result<Option<TargetPtxSet>, TargetResolveError> {
316 let targets = crate::all_ptx_sets();
317 let candidates: Vec<ResolveCandidate<'_>> = targets
318 .iter()
319 .map(|t| ResolveCandidate {
320 name: t.target.model,
321 type_matches: &t.model_type_matches,
322 match_names: t.match_names,
323 })
324 .collect();
325 let idx = match pinned_target {
326 Some(pin) => Some(resolve_pinned(&candidates, pin, model_type, hidden_size)?),
327 None => resolve_target(&candidates, model_type, hidden_size, model_refs)?,
328 };
329 drop(candidates);
330 Ok(idx.and_then(|i| targets.into_iter().nth(i)))
331}
332
333/// The compiled target with exactly this `(model, quant)` identity.
334///
335/// For consumers that already KNOW the resolved target (the dashboard's
336/// kernel table re-reads the target `serve` selected and published) —
337/// an exact lookup cannot re-introduce the ambiguity `ptx_for_config`
338/// just resolved.
339pub fn ptx_for_exact_target(model: &str, quant: &str) -> Option<TargetPtxSet> {
340 crate::all_ptx_sets()
341 .into_iter()
342 .find(|t| t.target.model == model && t.target.quant == quant)
343}
344
345#[cfg(test)]
346#[path = "resolve_tests.rs"]
347mod tests;