atlas_core/
target.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Kernel target descriptor for (Hardware, Model_quantization) tuples.
4//!
5//! Every set of Atlas kernels is hyperoptimized for a specific target.
6//! This module provides the `KernelTarget` type that serves as the
7//! indexing key for PTX module sets, benchmarks, and kernel dispatch.
8
9/// Describes a (Hardware, Model, Quantization) target for kernel selection.
10///
11/// Each unique `KernelTarget` maps to a distinct set of PTX modules
12/// that have been hyperoptimized for that specific combination.
13#[derive(Debug, Clone, PartialEq, Eq, Hash)]
14pub struct KernelTarget {
15    /// SM architecture identifier (e.g., "sm_121", "sm_100a").
16    pub arch: &'static str,
17    /// Model identifier (e.g., "qwen3-next-80b-a3b").
18    pub model: &'static str,
19    /// Quantization scheme (e.g., "nvfp4", "fp8", "bf16").
20    pub quant: &'static str,
21}
22
23impl KernelTarget {
24    /// GB10 + Qwen3-Next-80B-A3B + NVFP4.
25    pub const GB10_QWEN3_NVFP4: Self = Self {
26        arch: "sm_121",
27        model: "qwen3-next-80b-a3b",
28        quant: "nvfp4",
29    };
30
31    /// GB10 + Qwen3.5-35B-A3B + NVFP4.
32    pub const GB10_QWEN35_NVFP4: Self = Self {
33        arch: "sm_121",
34        model: "qwen3.5-35b-a3b",
35        quant: "nvfp4",
36    };
37
38    /// GB10 + Qwen3.5-122B-A10B + NVFP4.
39    pub const GB10_QWEN35_122B_NVFP4: Self = Self {
40        arch: "sm_121",
41        model: "qwen3.5-122b-a10b",
42        quant: "nvfp4",
43    };
44
45    /// Check if this target's model name contains the given substring.
46    pub fn model_contains(&self, substring: &str) -> bool {
47        self.model.contains(substring)
48    }
49}
50
51impl std::fmt::Display for KernelTarget {
52    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53        write!(f, "({}, {}, {})", self.arch, self.model, self.quant)
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn default_target_display() {
63        let t = KernelTarget::GB10_QWEN3_NVFP4;
64        assert_eq!(t.to_string(), "(sm_121, qwen3-next-80b-a3b, nvfp4)");
65    }
66
67    #[test]
68    fn target_equality_observes_each_dispatch_dimension() {
69        let a = KernelTarget::GB10_QWEN3_NVFP4;
70        assert_eq!(a, KernelTarget::GB10_QWEN3_NVFP4);
71        assert_ne!(
72            a,
73            KernelTarget {
74                arch: "sm_100a",
75                ..a
76            }
77        );
78        assert_ne!(
79            a,
80            KernelTarget {
81                model: "llama-70b",
82                ..a
83            }
84        );
85        assert_ne!(a, KernelTarget { quant: "fp8", ..a });
86    }
87}
atlas_core/target.rs

atlas_core/
target.rs
