1#[derive(Debug, Clone, PartialEq, Eq, Hash)]
11pub struct KernelTarget {
12 pub arch: &'static str,
16 pub model: &'static str,
18 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}