Skip to main content

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 --no-deps .` matching the upstream
200    /// build.sh, with `--no-deps` added to enforce the No-CUDA hard-prime tenet
201    /// (prevents pip from resolving/overriding the ROCm torch). Runtime deps are
202    /// installed explicitly and filtered via the unified nvidia_blocklist.
203    pub fn pip_install_kernel(&self) -> ShellCommand {
204        let use_break =
205            std::env::var("VIRTUAL_ENV").is_err() && std::env::var("CONDA_PREFIX").is_err();
206        let mut args = vec!["-m".to_string(), "pip".to_string(), "install".to_string()];
207        if use_break {
208            args.push("--break-system-packages".to_string());
209        }
210        args.extend([
211            "-v".to_string(),
212            "--no-build-isolation".to_string(),
213            "--no-deps".to_string(),
214            ".".to_string(),
215        ]);
216
217        let cmake_args = format!(
218            "-DCMAKE_HIP_ARCHITECTURES={} -DFASTVIDEO_KERNEL_BUILD_TK=OFF -DGPU_BACKEND=ROCM",
219            self.config.gpu_arch
220        );
221
222        let nproc = std::thread::available_parallelism()
223            .map(|n| n.get().saturating_sub(1))
224            .unwrap_or(4);
225
226        let env = vec![
227            ("CMAKE_ARGS".to_string(), cmake_args),
228            ("GPU_ARCHS".to_string(), self.config.gpu_arch.clone()),
229            ("MAX_JOBS".to_string(), nproc.to_string()),
230            ("GPU_BACKEND".to_string(), "ROCM".to_string()),
231        ];
232
233        ShellCommand {
234            program: self.config.python_bin.clone(),
235            args,
236            env,
237            working_dir: None,
238        }
239    }
240
241    pub fn cleanup(&self) -> ShellCommand {
242        ShellCommand {
243            program: "rm".to_string(),
244            args: vec!["-rf".to_string(), self.build_dir.clone()],
245            env: vec![],
246            working_dir: None,
247        }
248    }
249}
250
251#[cfg(test)]
252mod tests {
253    use super::*;
254
255    #[test]
256    fn test_build_commands_count() {
257        let inst = FastVideoInstaller::new(FastVideoConfig::default());
258        let cmds = inst.build_commands();
259        assert_eq!(cmds.len(), 8, "Should produce 8 commands");
260    }
261
262    #[test]
263    fn test_git_clone_uses_fork() {
264        let inst = FastVideoInstaller::new(FastVideoConfig::default());
265        let cmds = inst.build_commands();
266        let clone_cmd = &cmds[1];
267        assert!(clone_cmd.args.contains(&"clone".into()));
268        assert!(
269            clone_cmd.args.contains(&FASTVIDEO_REPO.into()),
270            "Should clone from scooter-lacroix fork"
271        );
272    }
273
274    #[test]
275    fn test_submodule_init_present() {
276        let inst = FastVideoInstaller::new(FastVideoConfig::default());
277        let cmds = inst.build_commands();
278        let submod_cmd = &cmds[3];
279        assert!(submod_cmd.args.contains(&"submodule".into()));
280        assert!(submod_cmd.args.contains(&"--recursive".into()));
281    }
282
283    #[test]
284    fn test_pip_install_has_rocm_cmake_args() {
285        let inst = FastVideoInstaller::new(FastVideoConfig {
286            gpu_arch: "gfx1100".to_string(),
287            python_bin: "python3".to_string(),
288        });
289        let cmd = inst.pip_install_kernel();
290        assert!(cmd.args.contains(&"--no-build-isolation".into()));
291        let cmake_env = cmd.env.iter().find(|(k, _)| k == "CMAKE_ARGS").unwrap();
292        assert!(cmake_env.1.contains("CMAKE_HIP_ARCHITECTURES=gfx1100"));
293        assert!(cmake_env.1.contains("GPU_BACKEND=ROCM"));
294        assert!(cmake_env.1.contains("FASTVIDEO_KERNEL_BUILD_TK=OFF"));
295    }
296
297    #[test]
298    fn test_pip_install_sets_gpu_archs() {
299        let inst = FastVideoInstaller::new(FastVideoConfig {
300            gpu_arch: "gfx1200".to_string(),
301            python_bin: "python3".to_string(),
302        });
303        let cmd = inst.pip_install_kernel();
304        let gpu_env = cmd.env.iter().find(|(k, _)| k == "GPU_ARCHS").unwrap();
305        assert_eq!(gpu_env.1, "gfx1200");
306    }
307
308    #[test]
309    fn test_cleanup_removes_build_dir() {
310        let inst = FastVideoInstaller::new(FastVideoConfig::default());
311        let cmds = inst.build_commands();
312        let cleanup = &cmds[7];
313        assert!(cleanup.args.contains(&"-rf".into()));
314        assert!(cleanup.args.contains(&FASTVIDEO_BUILD_DIR.into()));
315    }
316
317    #[test]
318    fn test_config_propagates_python_bin() {
319        let inst = FastVideoInstaller::new(FastVideoConfig {
320            gpu_arch: "gfx1100".to_string(),
321            python_bin: "python3.12".to_string(),
322        });
323        let cmd = inst.pip_install_kernel();
324        assert_eq!(cmd.program, "python3.12");
325    }
326}