1use std::collections::BTreeSet;
4
5use ferrum_types::NativeOperatorBackend;
6use thiserror::Error;
7
8use crate::ResolvedNativeOperatorArtifactSet;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
11pub enum CudaNativeBuildUnit {
12 Marlin,
13 VllmMarlin,
14 VllmMoeMarlin,
15 VllmPagedAttentionV2,
16}
17
18impl CudaNativeBuildUnit {
19 pub const ALL: [Self; 4] = [
20 Self::Marlin,
21 Self::VllmMarlin,
22 Self::VllmMoeMarlin,
23 Self::VllmPagedAttentionV2,
24 ];
25
26 pub const fn as_str(self) -> &'static str {
27 match self {
28 Self::Marlin => "marlin",
29 Self::VllmMarlin => "vllm_marlin",
30 Self::VllmMoeMarlin => "vllm_moe_marlin",
31 Self::VllmPagedAttentionV2 => "vllm_paged_attention_v2",
32 }
33 }
34
35 pub const fn artifact_operator(self) -> &'static str {
36 match self {
37 Self::Marlin => "ferrum.cuda.marlin",
38 Self::VllmMarlin => "ferrum.cuda.vllm_marlin",
39 Self::VllmMoeMarlin => "ferrum.cuda.vllm_moe_marlin",
40 Self::VllmPagedAttentionV2 => "ferrum.cuda.vllm_paged_attention_v2",
41 }
42 }
43
44 pub fn from_artifact_operator(operator: &str) -> Option<Self> {
45 Self::ALL
46 .into_iter()
47 .find(|unit| unit.artifact_operator() == operator)
48 }
49
50 pub const fn required_exports(self) -> &'static [&'static str] {
51 match self {
52 Self::Marlin => &["marlin_cuda", "marlin_cuda_moe"],
53 Self::VllmMarlin => &[
54 "ferrum_block_fp8_group128_repack",
55 "ferrum_block_fp8_group128_scales",
56 "ferrum_marlin_mm",
57 "ferrum_marlin_mm_f16_u4b8",
58 "ferrum_vllm_gptq_marlin_repack",
59 ],
60 Self::VllmMoeMarlin => &[
61 "ferrum_vllm_marlin_moe_clear_profile_config",
62 "ferrum_vllm_marlin_moe_f16",
63 "ferrum_vllm_marlin_moe_fp8_f16",
64 "ferrum_vllm_marlin_moe_mxfp4_bf16",
65 "ferrum_vllm_marlin_moe_set_profile_config",
66 ],
67 Self::VllmPagedAttentionV2 => &[
68 "ferrum_vllm_paged_attention_v1_f16_h128_b16",
69 "ferrum_vllm_paged_attention_v1_f16_h256_b16",
70 "ferrum_vllm_paged_attention_v2_f16_h128_b16",
71 "ferrum_vllm_paged_attention_v2_f16_h256_b16",
72 "ferrum_vnext_vllm_paged_attention_v1_f16_h128_b16_addressed",
73 "ferrum_vnext_vllm_paged_attention_v1_f16_h256_b16_addressed",
74 "ferrum_vnext_vllm_paged_attention_v2_f16_h128_b16_addressed",
75 "ferrum_vnext_vllm_paged_attention_v2_f16_h256_b16_addressed",
76 ],
77 }
78 }
79}
80
81#[derive(Debug, Clone, PartialEq, Eq)]
82pub struct ResolvedCudaNativeBuildCoverage {
83 units: BTreeSet<CudaNativeBuildUnit>,
84}
85
86impl ResolvedCudaNativeBuildCoverage {
87 pub fn resolve(
88 artifacts: &ResolvedNativeOperatorArtifactSet,
89 required: impl IntoIterator<Item = CudaNativeBuildUnit>,
90 ) -> Result<Self, CudaNativeBuildCoverageError> {
91 let views = artifacts
92 .artifacts
93 .iter()
94 .map(|artifact| ArtifactView {
95 operator: artifact.resolved.manifest.operator.as_str(),
96 backend: artifact.resolved.manifest.backend,
97 exports: artifact.resolved.manifest.exports.as_slice(),
98 })
99 .collect::<Vec<_>>();
100 resolve_views(&views, required)
101 }
102
103 pub fn contains(&self, unit: CudaNativeBuildUnit) -> bool {
104 self.units.contains(&unit)
105 }
106
107 pub fn iter(&self) -> impl Iterator<Item = CudaNativeBuildUnit> + '_ {
108 self.units.iter().copied()
109 }
110}
111
112#[derive(Debug, Error, PartialEq, Eq)]
113pub enum CudaNativeBuildCoverageError {
114 #[error(
115 "native artifact set does not provide CUDA build unit {unit} (required operator={operator})"
116 )]
117 MissingArtifact {
118 unit: &'static str,
119 operator: &'static str,
120 },
121 #[error("native artifact for CUDA build unit {unit} has backend {actual:?}, expected cuda")]
122 WrongBackend {
123 unit: &'static str,
124 actual: NativeOperatorBackend,
125 },
126 #[error("native artifact for CUDA build unit {unit} is missing required export {export}")]
127 MissingExport {
128 unit: &'static str,
129 export: &'static str,
130 },
131}
132
133struct ArtifactView<'a> {
134 operator: &'a str,
135 backend: NativeOperatorBackend,
136 exports: &'a [String],
137}
138
139fn resolve_views(
140 artifacts: &[ArtifactView<'_>],
141 required: impl IntoIterator<Item = CudaNativeBuildUnit>,
142) -> Result<ResolvedCudaNativeBuildCoverage, CudaNativeBuildCoverageError> {
143 let required = required.into_iter().collect::<BTreeSet<_>>();
144 for unit in &required {
145 let artifact = artifacts
146 .iter()
147 .find(|artifact| artifact.operator == unit.artifact_operator())
148 .ok_or(CudaNativeBuildCoverageError::MissingArtifact {
149 unit: unit.as_str(),
150 operator: unit.artifact_operator(),
151 })?;
152 if artifact.backend != NativeOperatorBackend::Cuda {
153 return Err(CudaNativeBuildCoverageError::WrongBackend {
154 unit: unit.as_str(),
155 actual: artifact.backend,
156 });
157 }
158 for required_export in unit.required_exports() {
159 if !artifact
160 .exports
161 .iter()
162 .any(|export| export == required_export)
163 {
164 return Err(CudaNativeBuildCoverageError::MissingExport {
165 unit: unit.as_str(),
166 export: required_export,
167 });
168 }
169 }
170 }
171 Ok(ResolvedCudaNativeBuildCoverage { units: required })
172}
173
174#[cfg(test)]
175mod tests {
176 use super::*;
177
178 fn exports(unit: CudaNativeBuildUnit) -> Vec<String> {
179 unit.required_exports()
180 .iter()
181 .map(|value| (*value).to_string())
182 .collect()
183 }
184
185 #[test]
186 fn resolves_exact_required_build_units() {
187 let operator_exports = CudaNativeBuildUnit::ALL.map(exports);
188 let artifacts = CudaNativeBuildUnit::ALL
189 .iter()
190 .enumerate()
191 .map(|(index, unit)| ArtifactView {
192 operator: unit.artifact_operator(),
193 backend: NativeOperatorBackend::Cuda,
194 exports: &operator_exports[index],
195 })
196 .collect::<Vec<_>>();
197
198 let coverage =
199 resolve_views(&artifacts, CudaNativeBuildUnit::ALL).expect("complete artifact set");
200
201 assert!(CudaNativeBuildUnit::ALL
202 .into_iter()
203 .all(|unit| coverage.contains(unit)));
204 }
205
206 #[test]
207 fn rejects_missing_artifact_before_source_build_can_fallback() {
208 let error = resolve_views(&[], [CudaNativeBuildUnit::VllmMoeMarlin]).unwrap_err();
209
210 assert_eq!(
211 error,
212 CudaNativeBuildCoverageError::MissingArtifact {
213 unit: "vllm_moe_marlin",
214 operator: "ferrum.cuda.vllm_moe_marlin",
215 }
216 );
217 }
218
219 #[test]
220 fn rejects_artifact_that_does_not_export_the_linked_rust_abi() {
221 let exports = vec!["ferrum_vllm_marlin_moe_f16".to_string()];
222 let artifacts = [ArtifactView {
223 operator: CudaNativeBuildUnit::VllmMoeMarlin.artifact_operator(),
224 backend: NativeOperatorBackend::Cuda,
225 exports: &exports,
226 }];
227
228 let error = resolve_views(&artifacts, [CudaNativeBuildUnit::VllmMoeMarlin]).unwrap_err();
229
230 assert_eq!(
231 error,
232 CudaNativeBuildCoverageError::MissingExport {
233 unit: "vllm_moe_marlin",
234 export: "ferrum_vllm_marlin_moe_clear_profile_config",
235 }
236 );
237 }
238
239 #[test]
240 fn rejects_moe_artifact_without_fp8_entrypoint_before_link() {
241 let exports = CudaNativeBuildUnit::VllmMoeMarlin
242 .required_exports()
243 .iter()
244 .copied()
245 .filter(|export| *export != "ferrum_vllm_marlin_moe_fp8_f16")
246 .map(str::to_owned)
247 .collect::<Vec<_>>();
248 let artifacts = [ArtifactView {
249 operator: CudaNativeBuildUnit::VllmMoeMarlin.artifact_operator(),
250 backend: NativeOperatorBackend::Cuda,
251 exports: &exports,
252 }];
253
254 let error = resolve_views(&artifacts, [CudaNativeBuildUnit::VllmMoeMarlin]).unwrap_err();
255
256 assert_eq!(
257 error,
258 CudaNativeBuildCoverageError::MissingExport {
259 unit: "vllm_moe_marlin",
260 export: "ferrum_vllm_marlin_moe_fp8_f16",
261 }
262 );
263 }
264}