metrale_core/
target.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Kernel target descriptor: the `(arch, model, quant)` key of a compiled kernel set.
4//!
5//! Owner: metrale-core.
6//! Invariants: none beyond the types.
7
8/// 2026-09-25: The `(arch, model, quant)` key of one compiled kernel set; each
9/// `metrale_kernels::TargetPtxSet` carries one as its `target`.
10#[derive(Debug, Clone, PartialEq, Eq, Hash)]
11pub struct KernelTarget {
12    /// 2026-09-25: Architecture identifier. For CUDA it is the base SM
13    /// (`sm_121`, `sm_90`): the build strips a HARDWARE.toml feature suffix
14    /// (`sm_90a` -> `sm_90`). Other backends' names are kept verbatim.
15    pub arch: &'static str,
16    /// 2026-09-25: Model directory name under `kernels/<hardware>/` (e.g. `qwen3-next-80b-a3b`).
17    pub model: &'static str,
18    /// 2026-09-25: Quantization directory name under the model (e.g. `nvfp4`, `fp8`, `bf16`).
19    pub quant: &'static str,
20}
21
22impl KernelTarget {
23    pub const GB10_QWEN3_NVFP4: Self = Self {
24        arch: "sm_121",
25        model: "qwen3-next-80b-a3b",
26        quant: "nvfp4",
27    };
28
29    pub const GB10_QWEN35_NVFP4: Self = Self {
30        arch: "sm_121",
31        model: "qwen3.5-35b-a3b",
32        quant: "nvfp4",
33    };
34
35    pub const GB10_QWEN35_122B_NVFP4: Self = Self {
36        arch: "sm_121",
37        model: "qwen3.5-122b-a10b",
38        quant: "nvfp4",
39    };
40
41    pub fn model_contains(&self, substring: &str) -> bool {
42        self.model.contains(substring)
43    }
44}
45
46impl std::fmt::Display for KernelTarget {
47    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
48        write!(f, "({}, {}, {})", self.arch, self.model, self.quant)
49    }
50}
51
52#[cfg(test)]
53mod tests {
54    use super::*;
55
56    #[test]
57    fn default_target_display() {
58        let t = KernelTarget::GB10_QWEN3_NVFP4;
59        assert_eq!(t.to_string(), "(sm_121, qwen3-next-80b-a3b, nvfp4)");
60    }
61
62    #[test]
63    fn target_equality_observes_each_dispatch_dimension() {
64        let a = KernelTarget::GB10_QWEN3_NVFP4;
65        assert_eq!(a, KernelTarget::GB10_QWEN3_NVFP4);
66        assert_ne!(
67            a,
68            KernelTarget {
69                arch: "sm_100a",
70                ..a
71            }
72        );
73        assert_ne!(
74            a,
75            KernelTarget {
76                model: "llama-70b",
77                ..a
78            }
79        );
80        assert_ne!(a, KernelTarget { quant: "fp8", ..a });
81    }
82}