rusty_stack/installers/components/
fastvideo.rs

1//! FastVideo installer — native Rust build for ROCm gfx11 support.
2//!
3//! Clones scooter-lacroix/FastVideo feature/rocm-gfx11-support branch,
4//! builds fastvideo-kernel with ROCm support using cmake, then pip installs.
5//!
6//! Instead of delegating to `./build.sh --rocm` (which calls `uv pip install`
7//! without `--system` and fails outside a venv), this module replicates the
8//! build steps individually with proper env var injection and pip prefix
9//! handling consistent with the rest of the rusty-stack installer ecosystem.
10//!
11//! # Validation Assertion
12//!
13//! - **VAL-INSTALL-050**: FastVideo builds from fork with ROCm gfx11 support
14
15use std::path::PathBuf;
16
17const FASTVIDEO_REPO: &str = "https://github.com/scooter-lacroix/FastVideo.git";
18const FASTVIDEO_BRANCH: &str = "feature/rocm-gfx11-support";
19const FASTVIDEO_BUILD_DIR: &str = "/tmp/FastVideo_ROCm_build";
20
21/// A constructed shell command.
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct ShellCommand {
24    /// The program to run.
25    pub program: String,
26    /// Arguments to pass.
27    pub args: Vec<String>,
28    /// Environment variables to set.
29    pub env: Vec<(String, String)>,
30    /// Working directory for the command.
31    pub working_dir: Option<PathBuf>,
32}
33
34/// Configuration for the FastVideo installer.
35#[derive(Debug, Clone)]
36pub struct FastVideoConfig {
37    /// GPU architecture (e.g., "gfx1100").
38    pub gpu_arch: String,
39    /// Python binary to use.
40    pub python_bin: String,
41}
42
43impl Default for FastVideoConfig {
44    fn default() -> Self {
45        Self {
46            gpu_arch: "gfx1100".to_string(),
47            python_bin: "python3".to_string(),
48        }
49    }
50}
51
52/// FastVideo installer.
53#[derive(Debug, Clone)]
54pub struct FastVideoInstaller {
55    repo_url: String,
56    branch: String,
57    build_dir: String,
58    config: FastVideoConfig,
59}
60
61impl Default for FastVideoInstaller {
62    fn default() -> Self {
63        Self::new(FastVideoConfig::default())
64    }
65}
66
67impl FastVideoInstaller {
68    pub fn new(config: FastVideoConfig) -> Self {
69        Self {
70            repo_url: FASTVIDEO_REPO.to_string(),
71            branch: FASTVIDEO_BRANCH.to_string(),
72            build_dir: FASTVIDEO_BUILD_DIR.to_string(),
73            config,
74        }
75    }
76
77    /// Build the sequence of commands to install FastVideo with ROCm support.
78    pub fn build_commands(&self) -> Vec<ShellCommand> {
79        vec![
80            self.mkdir_build_dir(),
81            self.git_clone(),
82            self.git_checkout(),
83            self.git_submodule_init(),
84            self.patch_cmake(),
85            self.install_build_deps(),
86            self.pip_install_kernel(),
87            self.cleanup(),
88        ]
89    }
90
91    pub fn mkdir_build_dir(&self) -> ShellCommand {
92        ShellCommand {
93            program: "mkdir".to_string(),
94            args: vec!["-p".to_string(), self.build_dir.clone()],
95            env: vec![],
96            working_dir: None,
97        }
98    }
99
100    pub fn git_clone(&self) -> ShellCommand {
101        ShellCommand {
102            program: "git".to_string(),
103            args: vec!["clone".to_string(), self.repo_url.clone(), ".".to_string()],
104            env: vec![],
105            working_dir: None,
106        }
107    }
108
109    pub fn git_checkout(&self) -> ShellCommand {
110        ShellCommand {
111            program: "git".to_string(),
112            args: vec!["checkout".to_string(), self.branch.clone()],
113            env: vec![],
114            working_dir: None,
115        }
116    }
117
118    /// Initialize git submodules (cutlass, ThunderKittens).
119    pub fn git_submodule_init(&self) -> ShellCommand {
120        ShellCommand {
121            program: "git".to_string(),
122            args: vec![
123                "submodule".to_string(),
124                "update".to_string(),
125                "--init".to_string(),
126                "--recursive".to_string(),
127            ],
128            env: vec![],
129            working_dir: None,
130        }
131    }
132
133    /// Patch source files to remove flash_attn_rocm.cpp from the build.
134    ///
135    /// The `feature/rocm-gfx11-support` branch includes `flash_attn_rocm.cpp`
136    /// which depends on AMD's composable_kernel MHA device implementations.
137    /// The CK submodule at the required commit (`5ff2d64`) is incompatible with
138    /// the current HIP compiler — it fails with missing CUDA builtins
139    /// (`threadIdx`/`blockDim`), missing standard library includes, and
140    /// template deduction errors in `device_grouped_mha_bwd_*_v2.hpp`.
141    ///
142    /// Three patches are applied:
143    /// 1. Remove `flash_attn_rocm.cpp` from CMakeLists.txt source list
144    /// 2. Remove `#include "attention/flash_attn_rocm.cpp"` and the forward
145    ///    declarations for `mha_varlen_fwd`/`mha_varlen_bwd` from
146    ///    `common_extension.cpp` (lines 4-23)
147    /// 3. Remove pybind11 bindings for `flash_attn_varlen_fwd`/`_bwd`
148    ///    from `common_extension.cpp`
149    ///
150    /// The rest of fastvideo-kernel (gemm, norm, quant) compiles correctly.
151    pub fn patch_cmake(&self) -> ShellCommand {
152        ShellCommand {
153            program: "bash".to_string(),
154            args: vec![
155                "-c".to_string(),
156                "sed -i '/flash_attn_rocm\\.cpp/d' CMakeLists.txt && \
157                 sed -i '4,23d' csrc/common_extension.cpp && \
158                 sed -i '/Flash Attention ROCm$/d;/flash_attn_varlen_fwd/d;/flash_attn_varlen_bwd/d' csrc/common_extension.cpp".to_string(),
159            ],
160            env: vec![],
161            working_dir: None,
162        }
163    }
164
165    /// Install build dependencies (scikit-build-core, cmake, ninja).
166    ///
167    /// Uses `python -m pip install` with `--break-system-packages` when needed,
168    /// matching the pattern used by Triton, AITER, and other source-build installers.
169    pub fn install_build_deps(&self) -> ShellCommand {
170        let use_break =
171            std::env::var("VIRTUAL_ENV").is_err() && std::env::var("CONDA_PREFIX").is_err();
172        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
173        if use_break {
174            args.push("--break-system-packages".to_string());
175        }
176        args.extend([
177            "--upgrade".to_string(),
178            "--no-cache-dir".to_string(),
179            "scikit-build-core".to_string(),
180            "cmake".to_string(),
181            "ninja".to_string(),
182        ]);
183
184        ShellCommand {
185            program: self.config.python_bin.clone(),
186            args,
187            env: vec![],
188            working_dir: None,
189        }
190    }
191
192    /// Build and install fastvideo-kernel with ROCm support.
193    ///
194    /// Sets CMAKE_ARGS for ROCm:
195    /// - `-DCMAKE_HIP_ARCHITECTURES=<gpu_arch>`
196    /// - `-DFASTVIDEO_KERNEL_BUILD_TK=OFF` (ThunderKittens not supported on ROCm)
197    /// - `-DGPU_BACKEND=ROCM`
198    ///
199    /// Uses `pip install --no-build-isolation` matching the upstream build.sh.
200    pub fn pip_install_kernel(&self) -> ShellCommand {
201        let use_break =
202            std::env::var("VIRTUAL_ENV").is_err() && std::env::var("CONDA_PREFIX").is_err();
203        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
204        if use_break {
205            args.push("--break-system-packages".to_string());
206        }
207        args.extend([
208            "-v".to_string(),
209            "--no-build-isolation".to_string(),
210            ".".to_string(),
211        ]);
212
213        let cmake_args = format!(
214            "-DCMAKE_HIP_ARCHITECTURES={} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DGPU_BACKEND=ROCM",
215            self.config.gpu_arch
216        );
217
218        let nproc = std::thread::available_parallelism()
219            .map(|n| n.get().saturating_sub(1))
220            .unwrap_or(4);
221
222        let env = vec![
223            ("CMAKE_ARGS".to_string(), cmake_args),
224            ("GPU_ARCHS".to_string(), self.config.gpu_arch.clone()),
225            ("MAX_JOBS".to_string(), nproc.to_string()),
226            ("GPU_BACKEND".to_string(), "ROCM".to_string()),
227        ];
228
229        ShellCommand {
230            program: self.config.python_bin.clone(),
231            args,
232            env,
233            working_dir: None,
234        }
235    }
236
237    pub fn cleanup(&self) -> ShellCommand {
238        ShellCommand {
239            program: "rm".to_string(),
240            args: vec!["-rf".to_string(), self.build_dir.clone()],
241            env: vec![],
242            working_dir: None,
243        }
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn test_build_commands_count() {
253        let inst = FastVideoInstaller::new(FastVideoConfig::default());
254        let cmds = inst.build_commands();
255        assert_eq!(cmds.len(), 8, "Should produce 8 commands");
256    }
257
258    #[test]
259    fn test_git_clone_uses_fork() {
260        let inst = FastVideoInstaller::new(FastVideoConfig::default());
261        let cmds = inst.build_commands();
262        let clone_cmd = &cmds[1];
263        assert!(clone_cmd.args.contains(&"clone".into()));
264        assert!(
265            clone_cmd.args.contains(&FASTVIDEO_REPO.into()),
266            "Should clone from scooter-lacroix fork"
267        );
268    }
269
270    #[test]
271    fn test_submodule_init_present() {
272        let inst = FastVideoInstaller::new(FastVideoConfig::default());
273        let cmds = inst.build_commands();
274        let submod_cmd = &cmds[3];
275        assert!(submod_cmd.args.contains(&"submodule".into()));
276        assert!(submod_cmd.args.contains(&"--recursive".into()));
277    }
278
279    #[test]
280    fn test_pip_install_has_rocm_cmake_args() {
281        let inst = FastVideoInstaller::new(FastVideoConfig {
282            gpu_arch: "gfx1100".to_string(),
283            python_bin: "python3".to_string(),
284        });
285        let cmd = inst.pip_install_kernel();
286        assert!(cmd.args.contains(&"--no-build-isolation".into()));
287        let cmake_env = cmd.env.iter().find(|(k, _)| k == "CMAKE_ARGS").unwrap();
288        assert!(cmake_env.1.contains("CMAKE_HIP_ARCHITECTURES=gfx1100"));
289        assert!(cmake_env.1.contains("GPU_BACKEND=ROCM"));
290        assert!(cmake_env.1.contains("FASTVIDEO_KERNEL_BUILD_TK=OFF"));
291    }
292
293    #[test]
294    fn test_pip_install_sets_gpu_archs() {
295        let inst = FastVideoInstaller::new(FastVideoConfig {
296            gpu_arch: "gfx1200".to_string(),
297            python_bin: "python3".to_string(),
298        });
299        let cmd = inst.pip_install_kernel();
300        let gpu_env = cmd.env.iter().find(|(k, _)| k == "GPU_ARCHS").unwrap();
301        assert_eq!(gpu_env.1, "gfx1200");
302    }
303
304    #[test]
305    fn test_cleanup_removes_build_dir() {
306        let inst = FastVideoInstaller::new(FastVideoConfig::default());
307        let cmds = inst.build_commands();
308        let cleanup = &cmds[7];
309        assert!(cleanup.args.contains(&"-rf".into()));
310        assert!(cleanup.args.contains(&FASTVIDEO_BUILD_DIR.into()));
311    }
312
313    #[test]
314    fn test_config_propagates_python_bin() {
315        let inst = FastVideoInstaller::new(FastVideoConfig {
316            gpu_arch: "gfx1100".to_string(),
317            python_bin: "python3.12".to_string(),
318        });
319        let cmd = inst.pip_install_kernel();
320        assert_eq!(cmd.program, "python3.12");
321    }
322}
rusty_stack/installers/components/fastvideo.rs

rusty_stack/installers/components/
fastvideo.rs
