rusty_stack/installers/components/
flash_attention_ck.rs

1//! Flash Attention CK installer — ports `scripts/install_flash_attention_ck.sh`.
2//!
3//! Constructs correct cmake commands with HIP compiler flags and ROCm path,
4//! make commands with job count, and pip install from build output.
5//!
6//! Flash Attention CK depends on PyTorch (ROCm-enabled) and ROCm.
7//!
8//! # Validation Assertions
9//!
10//! - **VAL-INSTALL-008**: Flash Attention correct cmake command
11//! - **VAL-INSTALL-009**: Flash Attention correct make command
12//! - **VAL-INSTALL-014**: Source builds parse build output for errors
13//! - **VAL-INSTALL-044**: Flash Attention declares dependency on PyTorch and ROCm
14
15use crate::installers::common::RocmEnv;
16use std::fmt;
17use std::path::PathBuf;
18
19// ===========================================================================
20// Types
21// ===========================================================================
22
23/// Installation method.
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub enum InstallMethod {
26    /// Install globally.
27    Global,
28    /// Install in a virtual environment.
29    Venv,
30    /// Try global, fallback to venv.
31    Auto,
32}
33
34impl fmt::Display for InstallMethod {
35    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
36        match self {
37            InstallMethod::Global => write!(f, "global"),
38            InstallMethod::Venv => write!(f, "venv"),
39            InstallMethod::Auto => write!(f, "auto"),
40        }
41    }
42}
43
44/// GPU architecture specification for building.
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub enum GpuArch {
47    /// Detected or specified GPU architecture (e.g., "gfx1100").
48    Known(String),
49    /// Default fallback architecture.
50    Default,
51}
52
53impl GpuArch {
54    /// Get the architecture string.
55    pub fn as_str(&self) -> &str {
56        match self {
57            GpuArch::Known(s) => s,
58            GpuArch::Default => "gfx1100",
59        }
60    }
61}
62
63/// Configuration for the Flash Attention CK installer.
64#[derive(Debug, Clone)]
65pub struct FlashAttentionConfig {
66    /// Installation method.
67    pub method: InstallMethod,
68    /// Python binary to use.
69    pub python_bin: String,
70    /// Whether to force reinstall.
71    pub force_reinstall: bool,
72    /// Whether to run in dry-run mode.
73    pub dry_run: bool,
74    /// Installation directory (defaults to `$HOME/ml_stack/flash_attn_amd`).
75    pub install_dir: Option<PathBuf>,
76    /// GPU architecture override.
77    pub gpu_arch: GpuArch,
78    /// Number of parallel build jobs (defaults to nproc).
79    pub build_jobs: Option<usize>,
80}
81
82impl Default for FlashAttentionConfig {
83    fn default() -> Self {
84        Self {
85            method: InstallMethod::Auto,
86            python_bin: "python3".to_string(),
87            force_reinstall: false,
88            dry_run: false,
89            install_dir: None,
90            gpu_arch: GpuArch::Default,
91            build_jobs: None,
92        }
93    }
94}
95
96impl FlashAttentionConfig {
97    /// Get the installation directory path.
98    pub fn install_dir(&self) -> PathBuf {
99        self.install_dir
100            .clone()
101            .unwrap_or_else(|| dirs_home().join("ml_stack").join("flash_attn_amd"))
102    }
103
104    /// Get the number of build jobs (defaults to num_cpus).
105    pub fn build_jobs(&self) -> usize {
106        self.build_jobs.unwrap_or_else(num_cpus)
107    }
108}
109
110/// A constructed shell command.
111#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct ShellCommand {
113    /// The program to run.
114    pub program: String,
115    /// Arguments to pass.
116    pub args: Vec<String>,
117    /// Environment variables to set.
118    pub env: Vec<(String, String)>,
119    /// Working directory for the command.
120    pub working_dir: Option<PathBuf>,
121}
122
123impl ShellCommand {
124    /// Format as a shell command string.
125    pub fn to_command_string(&self) -> String {
126        let env_prefix = self
127            .env
128            .iter()
129            .map(|(k, v)| format!("{k}={v}"))
130            .collect::<Vec<_>>()
131            .join(" ");
132        let cmd = if self.args.is_empty() {
133            self.program.clone()
134        } else {
135            format!("{} {}", self.program, self.args.join(" "))
136        };
137        if env_prefix.is_empty() {
138            cmd
139        } else {
140            format!("{env_prefix} {cmd}")
141        }
142    }
143}
144
145/// The Flash Attention CK installer.
146pub struct FlashAttentionInstaller {
147    config: FlashAttentionConfig,
148}
149
150impl FlashAttentionInstaller {
151    /// Create a new Flash Attention installer with the given config.
152    pub fn new(config: FlashAttentionConfig) -> Self {
153        Self { config }
154    }
155
156    /// Create with default config.
157    pub fn with_defaults() -> Self {
158        Self::new(FlashAttentionConfig::default())
159    }
160
161    // -----------------------------------------------------------------------
162    // Dependencies (VAL-INSTALL-044)
163    // -----------------------------------------------------------------------
164
165    /// Get the list of required dependencies.
166    ///
167    /// Flash Attention CK depends on:
168    /// - PyTorch (ROCm-enabled)
169    /// - ROCm
170    pub fn dependencies(&self) -> &[&str] {
171        &["pytorch", "rocm"]
172    }
173
174    // -----------------------------------------------------------------------
175    // Git clone command
176    // -----------------------------------------------------------------------
177
178    /// Construct the git clone command for Flash Attention.
179    ///
180    /// The original script clones from:
181    /// `https://github.com/ROCmSoftwarePlatform/flash-attention.git`
182    /// and checks out the `main` branch.
183    pub fn build_git_clone_command(&self) -> ShellCommand {
184        let install_dir = self.config.install_dir();
185        ShellCommand {
186            program: "git".to_string(),
187            args: vec![
188                "clone".to_string(),
189                "https://github.com/ROCmSoftwarePlatform/flash-attention.git".to_string(),
190                install_dir.to_string_lossy().to_string(),
191            ],
192            env: vec![],
193            working_dir: None,
194        }
195    }
196
197    /// Construct the git checkout command.
198    pub fn build_git_checkout_command(&self) -> ShellCommand {
199        ShellCommand {
200            program: "git".to_string(),
201            args: vec!["checkout".to_string(), "main".to_string()],
202            env: vec![],
203            working_dir: Some(self.config.install_dir()),
204        }
205    }
206
207    // -----------------------------------------------------------------------
208    // CMake command (VAL-INSTALL-008)
209    // -----------------------------------------------------------------------
210
211    /// Construct the cmake configure command.
212    ///
213    /// The original script configures cmake with:
214    /// - CMAKE_PREFIX_PATH: torch cmake prefix + ROCm path
215    /// - CMAKE_BUILD_TYPE=Release
216    /// - GPU_TARGETS=gfx1100 (or detected)
217    /// - CMAKE_CXX_FLAGS="-Wno-error"
218    /// - ROCM_PATH=/opt/rocm
219    ///
220    /// Working directory: `<install_dir>/build`
221    pub fn build_cmake_command(
222        &self,
223        rocm_env: &RocmEnv,
224        torch_cmake_prefix: &str,
225    ) -> ShellCommand {
226        let rocm_path = rocm_env
227            .path()
228            .map(|p| p.to_string_lossy().to_string())
229            .unwrap_or_else(|| "/opt/rocm".to_string());
230
231        let prefix_path = format!("{};{};/opt/rocm", torch_cmake_prefix, rocm_path);
232
233        let mut args = vec![
234            "..".to_string(),
235            format!("-DCMAKE_PREFIX_PATH={}", prefix_path),
236            "-DCMAKE_BUILD_TYPE=Release".to_string(),
237            format!("-DGPU_TARGETS={}", self.config.gpu_arch.as_str()),
238            "-DCMAKE_CXX_FLAGS=-Wno-error".to_string(),
239        ];
240
241        args.push(format!("-DROCM_PATH={}", rocm_path));
242
243        ShellCommand {
244            program: "cmake".to_string(),
245            args,
246            env: vec![],
247            working_dir: Some(self.config.install_dir().join("build")),
248        }
249    }
250
251    // -----------------------------------------------------------------------
252    // Make / build command (VAL-INSTALL-009)
253    // -----------------------------------------------------------------------
254
255    /// Construct the cmake --build command.
256    ///
257    /// The original script uses:
258    /// `cmake --build . --config Release -j $(nproc)`
259    ///
260    /// Working directory: `<install_dir>/build`
261    pub fn build_make_command(&self) -> ShellCommand {
262        let jobs = self.config.build_jobs();
263        ShellCommand {
264            program: "cmake".to_string(),
265            args: vec![
266                "--build".to_string(),
267                ".".to_string(),
268                "--config".to_string(),
269                "Release".to_string(),
270                "-j".to_string(),
271                jobs.to_string(),
272            ],
273            env: vec![],
274            working_dir: Some(self.config.install_dir().join("build")),
275        }
276    }
277
278    // -----------------------------------------------------------------------
279    // Setup.py install command
280    // -----------------------------------------------------------------------
281
282    /// Construct the Python setup.py install command.
283    ///
284    /// The original script runs:
285    /// `python3 setup_flash_attn_amd.py install`
286    pub fn build_setup_install_command(&self) -> ShellCommand {
287        let args = vec!["setup_flash_attn_amd.py".to_string(), "install".to_string()];
288
289        ShellCommand {
290            program: self.config.python_bin.clone(),
291            args,
292            env: vec![],
293            working_dir: Some(self.config.install_dir()),
294        }
295    }
296
297    // -----------------------------------------------------------------------
298    // Build error detection (VAL-INSTALL-014)
299    // -----------------------------------------------------------------------
300
301    /// Check build output for errors.
302    ///
303    /// Detects common cmake and make failure patterns in the output.
304    /// Returns an error message if a failure is detected.
305    pub fn check_build_output(
306        &self,
307        stdout: &str,
308        stderr: &str,
309        exit_code: i32,
310    ) -> Result<(), String> {
311        if exit_code != 0 {
312            // Check for specific error patterns
313            let error_patterns = [
314                "CMake Error",
315                "error:",
316                "Error:",
317                "FAILED:",
318                "fatal error:",
319                "undefined reference",
320                "No such file or directory",
321                "Permission denied",
322            ];
323
324            let combined = format!("{}\n{}", stdout, stderr);
325            let mut detected_errors: Vec<&str> = Vec::new();
326
327            for pattern in &error_patterns {
328                if combined.contains(pattern) {
329                    detected_errors.push(*pattern);
330                }
331            }
332
333            if detected_errors.is_empty() {
334                return Err(format!(
335                    "Build command failed with exit code {} (no specific error pattern detected)",
336                    exit_code
337                ));
338            }
339
340            return Err(format!(
341                "Build failed with exit code {}. Detected errors: {}",
342                exit_code,
343                detected_errors.join(", ")
344            ));
345        }
346
347        Ok(())
348    }
349}
350
351/// Get the home directory.
352fn dirs_home() -> PathBuf {
353    std::env::var("HOME")
354        .map(PathBuf::from)
355        .unwrap_or_else(|_| PathBuf::from("/root"))
356}
357
358/// Get the number of CPUs.
359fn num_cpus() -> usize {
360    std::thread::available_parallelism()
361        .map(|n| n.get())
362        .unwrap_or(4)
363}
364
365// ===========================================================================
366// Tests
367// ===========================================================================
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372    use std::path::PathBuf;
373
374    // --- VAL-INSTALL-008: Flash Attention correct cmake command ---
375
376    #[test]
377    fn test_cmake_command_has_correct_flags() {
378        let installer = FlashAttentionInstaller::new(FlashAttentionConfig {
379            gpu_arch: GpuArch::Known("gfx1100".to_string()),
380            ..Default::default()
381        });
382        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
383        let cmd = installer.build_cmake_command(
384            &rocm_env,
385            "/usr/lib/python3/dist-packages/torch/share/cmake",
386        );
387
388        assert_eq!(cmd.program, "cmake");
389        assert!(cmd.args.contains(&"..".to_string()));
390        assert!(cmd.args.iter().any(|a| a.contains("CMAKE_PREFIX_PATH")));
391        assert!(cmd
392            .args
393            .iter()
394            .any(|a| a.contains("CMAKE_BUILD_TYPE=Release")));
395        assert!(cmd.args.iter().any(|a| a.contains("GPU_TARGETS=gfx1100")));
396        assert!(cmd
397            .args
398            .iter()
399            .any(|a| a.contains("CMAKE_CXX_FLAGS=-Wno-error")));
400        assert!(cmd.args.iter().any(|a| a.contains("ROCM_PATH=/opt/rocm")));
401    }
402
403    #[test]
404    fn test_cmake_command_includes_torch_prefix() {
405        let installer = FlashAttentionInstaller::with_defaults();
406        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
407        let cmd = installer.build_cmake_command(&rocm_env, "/torch/cmake/prefix");
408        let prefix_arg = cmd
409            .args
410            .iter()
411            .find(|a| a.starts_with("-DCMAKE_PREFIX_PATH="));
412        assert!(prefix_arg.is_some());
413        let prefix = prefix_arg.unwrap();
414        assert!(prefix.contains("/torch/cmake/prefix"));
415        assert!(prefix.contains("/opt/rocm"));
416    }
417
418    #[test]
419    fn test_cmake_command_working_dir() {
420        let installer = FlashAttentionInstaller::new(FlashAttentionConfig {
421            install_dir: Some(PathBuf::from("/home/user/ml_stack/flash_attn_amd")),
422            ..Default::default()
423        });
424        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
425        let cmd = installer.build_cmake_command(&rocm_env, "/torch/cmake");
426        assert_eq!(
427            cmd.working_dir,
428            Some(PathBuf::from("/home/user/ml_stack/flash_attn_amd/build"))
429        );
430    }
431
432    #[test]
433    fn test_cmake_command_custom_gpu_arch() {
434        let installer = FlashAttentionInstaller::new(FlashAttentionConfig {
435            gpu_arch: GpuArch::Known("gfx1030".to_string()),
436            ..Default::default()
437        });
438        let rocm_env = RocmEnv::from_known(Some(PathBuf::from("/opt/rocm")), "7.2.0".to_string());
439        let cmd = installer.build_cmake_command(&rocm_env, "/torch/cmake");
440        assert!(cmd.args.iter().any(|a| a.contains("GPU_TARGETS=gfx1030")));
441    }
442
443    // --- VAL-INSTALL-009: Flash Attention correct make command ---
444
445    #[test]
446    fn test_make_command_correct_flags() {
447        let installer = FlashAttentionInstaller::with_defaults();
448        let cmd = installer.build_make_command();
449
450        assert_eq!(cmd.program, "cmake");
451        assert!(cmd.args.contains(&"--build".to_string()));
452        assert!(cmd.args.contains(&".".to_string()));
453        assert!(cmd.args.contains(&"--config".to_string()));
454        assert!(cmd.args.contains(&"Release".to_string()));
455        assert!(cmd.args.contains(&"-j".to_string()));
456    }
457
458    #[test]
459    fn test_make_command_job_count() {
460        let installer = FlashAttentionInstaller::new(FlashAttentionConfig {
461            build_jobs: Some(8),
462            ..Default::default()
463        });
464        let cmd = installer.build_make_command();
465        assert!(cmd.args.contains(&"8".to_string()));
466    }
467
468    #[test]
469    fn test_make_command_working_dir() {
470        let installer = FlashAttentionInstaller::new(FlashAttentionConfig {
471            install_dir: Some(PathBuf::from("/home/user/ml_stack/flash_attn_amd")),
472            ..Default::default()
473        });
474        let cmd = installer.build_make_command();
475        assert_eq!(
476            cmd.working_dir,
477            Some(PathBuf::from("/home/user/ml_stack/flash_attn_amd/build"))
478        );
479    }
480
481    // --- Git clone command ---
482
483    #[test]
484    fn test_git_clone_command() {
485        let installer = FlashAttentionInstaller::with_defaults();
486        let cmd = installer.build_git_clone_command();
487        assert_eq!(cmd.program, "git");
488        assert!(cmd.args.contains(&"clone".to_string()));
489        assert!(cmd
490            .args
491            .iter()
492            .any(|a| a.contains("ROCmSoftwarePlatform/flash-attention")));
493    }
494
495    #[test]
496    fn test_git_checkout_command() {
497        let installer = FlashAttentionInstaller::with_defaults();
498        let cmd = installer.build_git_checkout_command();
499        assert_eq!(cmd.program, "git");
500        assert!(cmd.args.contains(&"checkout".to_string()));
501        assert!(cmd.args.contains(&"main".to_string()));
502    }
503
504    // --- Setup.py install command ---
505
506    #[test]
507    fn test_setup_install_command() {
508        let installer = FlashAttentionInstaller::with_defaults();
509        let cmd = installer.build_setup_install_command();
510        assert_eq!(cmd.program, "python3");
511        assert!(cmd.args.contains(&"setup_flash_attn_amd.py".to_string()));
512        assert!(cmd.args.contains(&"install".to_string()));
513    }
514
515    // --- VAL-INSTALL-044: Flash Attention declares dependency on PyTorch and ROCm ---
516
517    #[test]
518    fn test_dependencies() {
519        let installer = FlashAttentionInstaller::with_defaults();
520        let deps = installer.dependencies();
521        assert!(deps.contains(&"pytorch"));
522        assert!(deps.contains(&"rocm"));
523    }
524
525    // --- VAL-INSTALL-014: Build error detection ---
526
527    #[test]
528    fn test_build_error_detection_cmake_error() {
529        let installer = FlashAttentionInstaller::with_defaults();
530        let result =
531            installer.check_build_output("some output", "CMake Error at CMakeLists.txt:10", 1);
532        assert!(result.is_err());
533        assert!(result.unwrap_err().contains("CMake Error"));
534    }
535
536    #[test]
537    fn test_build_error_detection_make_error() {
538        let installer = FlashAttentionInstaller::with_defaults();
539        let result = installer.check_build_output(
540            "compiling...",
541            "error: 'undefined_var' was not declared",
542            2,
543        );
544        assert!(result.is_err());
545        assert!(result.unwrap_err().contains("error:"));
546    }
547
548    #[test]
549    fn test_build_error_detection_success() {
550        let installer = FlashAttentionInstaller::with_defaults();
551        let result = installer.check_build_output("Build finished successfully", "", 0);
552        assert!(result.is_ok());
553    }
554
555    #[test]
556    fn test_build_error_detection_nonzero_no_pattern() {
557        let installer = FlashAttentionInstaller::with_defaults();
558        let result = installer.check_build_output("clean output", "clean output", 1);
559        assert!(result.is_err());
560        assert!(result.unwrap_err().contains("exit code 1"));
561    }
562
563    // --- Config defaults ---
564
565    #[test]
566    fn test_config_defaults() {
567        let config = FlashAttentionConfig::default();
568        assert_eq!(config.method, InstallMethod::Auto);
569        assert_eq!(config.python_bin, "python3");
570        assert!(!config.force_reinstall);
571        assert!(!config.dry_run);
572        assert!(config.install_dir.is_none());
573        assert_eq!(config.gpu_arch, GpuArch::Default);
574        assert!(config.build_jobs.is_none());
575    }
576
577    #[test]
578    fn test_gpu_arch_default() {
579        assert_eq!(GpuArch::Default.as_str(), "gfx1100");
580        assert_eq!(GpuArch::Known("gfx1030".to_string()).as_str(), "gfx1030");
581    }
582
583    #[test]
584    fn test_command_string_format() {
585        let cmd = ShellCommand {
586            program: "cmake".to_string(),
587            args: vec!["--build".to_string(), ".".to_string()],
588            env: vec![("ROCM_PATH".to_string(), "/opt/rocm".to_string())],
589            working_dir: None,
590        };
591        let s = cmd.to_command_string();
592        assert!(s.contains("ROCM_PATH=/opt/rocm"));
593        assert!(s.contains("cmake --build ."));
594    }
595}
rusty_stack/installers/components/flash_attention_ck.rs

rusty_stack/installers/components/
flash_attention_ck.rs
