From 2e4878c5e644237d53a039a050e0d006bbffb680 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Mon, 19 Oct 2026 10:09:54 +0000 Subject: [PATCH] [infra]: leave FlashAttention out of the GB10 image FA2 gives no measured speedup on GB10, and with it installed FastH3 V2 on a Spark produces different audio and frames than a native Spark install (which has no FA2). Without it the image matches the native install bit for bit at the same speed. Add INSTALL_FLASH_ATTN (default 1), set it to 0 for the CUDA 13 arm64 leg, and document the first-run FlashInfer compile (persist the cache, cap MAX_JOBS). --- .github/workflows/infra-build-image.yml | 9 ++++++ docker/Dockerfile | 22 ++++++++++---- docs/contributing/developer_env/docker.md | 30 ++++++++++++++----- .../tests/modal/test_kernel_build_cache.py | 12 ++++++-- 4 files changed, 57 insertions(+), 16 deletions(-) diff --git a/.github/workflows/infra-build-image.yml b/.github/workflows/infra-build-image.yml index d7940f7b6..0f7299dfc 100644 --- a/.github/workflows/infra-build-image.yml +++ b/.github/workflows/infra-build-image.yml @@ -80,6 +80,9 @@ jobs: torch_cuda_arch_list: amd64: '9.0a' arm64: '9.0a' + install_flash_attn: + amd64: '1' + arm64: '1' install_fa4_cute: amd64: '1' arm64: '1' @@ -93,6 +96,11 @@ jobs: arm64: '12.1a' # FA4 cute does not run on GB10 (sm_12x), and its cutlass-dsl pin # breaks flashinfer there; GB10 FP4 attention uses fastvideo-kernel. + # FA2 gives no measured GB10 speedup and changes the output relative + # to a native Spark install, so the GB10 image matches that install. + install_flash_attn: + amd64: '1' + arm64: '0' install_fa4_cute: amd64: '1' arm64: '0' @@ -115,6 +123,7 @@ jobs: CUDA_VERSION=${{ matrix.cuda.version }} UV_TORCH_BACKEND=${{ matrix.cuda.torch_backend }} TORCH_CUDA_ARCH_LIST=${{ matrix.cuda.torch_cuda_arch_list[matrix.architecture.name] }} + INSTALL_FLASH_ATTN=${{ matrix.cuda.install_flash_attn[matrix.architecture.name] }} INSTALL_FA4_CUTE=${{ matrix.cuda.install_fa4_cute[matrix.architecture.name] }} CMAKE_BUILD_PARALLEL_LEVEL=${{ matrix.cuda.cmake_build_parallel_level }} FLASH_ATTN_WHEEL_TAG=${{ matrix.cuda.fa_tag }} diff --git a/docker/Dockerfile b/docker/Dockerfile index 2e85e5dbe..70e32f9a8 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -23,7 +23,8 @@ # # DGX Spark / GB10 (aarch64, Blackwell sm_121a): the matching arm64 # # prebuilt flash-attn wheel is selected automatically from TARGETARCH. # docker build --platform linux/arm64 -f docker/Dockerfile \ -# --build-arg TORCH_CUDA_ARCH_LIST=12.1a --build-arg INSTALL_FA4_CUTE=0 \ +# --build-arg TORCH_CUDA_ARCH_LIST=12.1a \ +# --build-arg INSTALL_FLASH_ATTN=0 --build-arg INSTALL_FA4_CUTE=0 \ # -t fastvideo-dev:spark . # # The ARG defaults below are the source of truth for the default image. @@ -61,6 +62,9 @@ ARG CMAKE_BUILD_PARALLEL_LEVEL=1 # cuXXXtorchY.Y tag must match UV_TORCH_BACKEND (below): cu130torch2.12 for the # default 13.0 base, cu126torch2.12 for 12.6 (both backends are on torch 2.12). ARG FLASH_ATTN_VERSION=2.8.3 +# INSTALL_FLASH_ATTN=0 leaves FA2 out (the DGX Spark image): on GB10 it gives no +# measured speedup and changes the output relative to a native Spark install. +ARG INSTALL_FLASH_ATTN=1 ARG FLASH_ATTN_WHEEL_TAG=cu130torch2.12 ARG FLASH_ATTN_WHEEL_RELEASE=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.17 ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22 @@ -178,7 +182,9 @@ RUN --mount=type=cache,target=/opt/uv/cache \ exit 1 \ ;; \ esac && \ - if [ -n "${FLASH_ATTN_WHEEL_TAG}" ]; then \ + if [ "${INSTALL_FLASH_ATTN}" != "1" ]; then \ + echo "Skipping flash-attn (INSTALL_FLASH_ATTN=${INSTALL_FLASH_ATTN})"; \ + elif [ -n "${FLASH_ATTN_WHEEL_TAG}" ]; then \ echo "Installing prebuilt flash-attn wheel (${FLASH_ATTN_WHEEL_TAG}, ${PYTAG}, ${FLASH_ATTN_WHEEL_PLATFORM})" && \ uv pip install "${FLASH_ATTN_TARGET_RELEASE}/flash_attn-${FLASH_ATTN_VERSION}+${FLASH_ATTN_WHEEL_TAG}-${PYTAG}-${PYTAG}-${FLASH_ATTN_WHEEL_PLATFORM}.whl"; \ else \ @@ -197,17 +203,23 @@ RUN --mount=type=cache,target=/opt/uv/cache \ # an FA2-less image. The pinned stack is validated on ARM64 GB200 (sm_100) as well # as x86; FA4 remains opt-in through FASTVIDEO_FA4=1 so lanes with FA2 baselines # keep their existing numerics. -# With INSTALL_FA4_CUTE=0 the wheel's stale cute is still removed, so +# With INSTALL_FA4_CUTE=0 any wheel's stale cute is still removed, so # `import flash_attn.cute` fails cleanly and FastVideo's guarded import falls back. RUN --mount=type=cache,target=/opt/uv/cache \ source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ python -c "import glob, shutil; [shutil.rmtree(d, ignore_errors=True) for d in glob.glob('/opt/venv/lib/python*/site-packages/flash_attn/cute')]" && \ - if [ "${INSTALL_FA4_CUTE}" = "1" ]; then \ + if [ "${INSTALL_FA4_CUTE}" = "1" ] && [ "${INSTALL_FLASH_ATTN}" != "1" ]; then \ + echo "INSTALL_FA4_CUTE=1 overlays the FA2 install; it requires INSTALL_FLASH_ATTN=1" >&2; \ + exit 1; \ + elif [ "${INSTALL_FA4_CUTE}" = "1" ]; then \ uv pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@${FA4_CUTE_REF}#subdirectory=flash_attn/cute" && \ python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'), 'FA2 was clobbered by the cute overlay'; import flash_attn.cute; print('FA2 + FA4 cute OK')"; \ else \ - python -c "import flash_attn, flashinfer; assert hasattr(flash_attn, 'flash_attn_func'); print('FA2 + flashinfer', flashinfer.__version__, 'OK; FA4 cute skipped')"; \ + python -c "import flashinfer; print('flashinfer', flashinfer.__version__, 'OK; FA4 cute skipped')" && \ + if [ "${INSTALL_FLASH_ATTN}" = "1" ]; then \ + python -c "import flash_attn; assert hasattr(flash_attn, 'flash_attn_func'); print('FA2 OK')"; \ + fi; \ fi COPY . . diff --git a/docs/contributing/developer_env/docker.md b/docs/contributing/developer_env/docker.md index 6208a3be7..484528fc5 100644 --- a/docs/contributing/developer_env/docker.md +++ b/docs/contributing/developer_env/docker.md @@ -37,18 +37,32 @@ native ARM64 host, build the CUDA 13 DGX Spark image with: ```bash docker build --platform linux/arm64 -f docker/Dockerfile \ - --build-arg TORCH_CUDA_ARCH_LIST=12.1a --build-arg INSTALL_FA4_CUTE=0 \ + --build-arg TORCH_CUDA_ARCH_LIST=12.1a \ + --build-arg INSTALL_FLASH_ATTN=0 --build-arg INSTALL_FA4_CUTE=0 \ -t fastvideo-dev:spark . ``` -This uses the prebuilt Linux ARM64 FlashAttention wheel. FastVideo's in-tree -CUDA kernel is still compiled for `sm_121a` as part of the image build, including -the FP4 attention kernels that `FASTVIDEO_H3_VSA_FP4=1` uses. +FastVideo's in-tree CUDA kernel is compiled for `sm_121a` as part of the image +build, including the FP4 attention kernels that `FASTVIDEO_H3_VSA_FP4=1` uses. +These settings match the published DGX Spark image, which in turn matches a +native Spark install: -`INSTALL_FA4_CUTE=0` (set for the published DGX Spark image) leaves -out the FA4 `flash_attn.cute` overlay: FA4 supports SM 9/10/11 only, and its pinned -`nvidia-cutlass-dsl` would break FlashInfer, which the NVFP4 paths need. FA2 is -still installed. +- **No FlashAttention (`INSTALL_FLASH_ATTN=0`).** It gives no measured speedup on + GB10, and with it installed the output no longer matches a native install + bit for bit. +- **No FA4 cute overlay (`INSTALL_FA4_CUTE=0`).** FA4 supports SM 9/10/11 only, + and its pinned `nvidia-cutlass-dsl` breaks FlashInfer, which the NVFP4 paths need. + +The first NVFP4 generation in a fresh container compiles FlashInfer's GB10 FP4 +kernels (about 3.5 minutes). Persist the cache and cap the compile jobs; at full +parallelism the compile can be killed for lack of memory once a model is loaded: + +```bash +docker run --gpus all --ipc=host -e MAX_JOBS=4 \ + -v ~/.cache/flashinfer:/root/.cache/flashinfer \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + -it ghcr.io/hao-ai-lab/fastvideo/fastvideo-dev:py3.12-cuda13.0.0-latest +``` ## Using the container diff --git a/fastvideo/tests/modal/test_kernel_build_cache.py b/fastvideo/tests/modal/test_kernel_build_cache.py index 92d91700c..0838b73d3 100644 --- a/fastvideo/tests/modal/test_kernel_build_cache.py +++ b/fastvideo/tests/modal/test_kernel_build_cache.py @@ -178,7 +178,8 @@ def test_ci_runner_image_targets_arm64_sm100_with_opencv_runtime() -> None: def test_gb10_image_builds_sm121a_without_fa4_cute() -> None: """The CUDA 13 arm64 leg is the DGX Spark image: sm_121a keeps fastvideo-kernel's FP4 - attention, and skipping the FA4 cute overlay keeps flashinfer importable.""" + attention, skipping the FA4 cute overlay keeps flashinfer importable, and leaving FA2 + out keeps the output identical to a native Spark install.""" workflow = yaml.load( (REPO_ROOT / ".github/workflows/infra-build-image.yml").read_text(encoding="utf-8"), Loader=yaml.BaseLoader, @@ -188,12 +189,17 @@ def test_gb10_image_builds_sm121a_without_fa4_cute() -> None: assert cuda["13.0.0"]["torch_cuda_arch_list"]["arm64"] == "12.1a" assert cuda["13.0.0"]["install_fa4_cute"] == {"amd64": "1", "arm64": "0"} + assert cuda["13.0.0"]["install_flash_attn"] == {"amd64": "1", "arm64": "0"} + assert cuda["12.6.3"]["install_flash_attn"] == {"amd64": "1", "arm64": "1"} assert cuda["12.6.3"]["install_fa4_cute"] == {"amd64": "1", "arm64": "1"} - assert ("INSTALL_FA4_CUTE=${{ matrix.cuda.install_fa4_cute[matrix.architecture.name] }}" - in job["with"]["build_args"]) + for arg in ("INSTALL_FLASH_ATTN", "INSTALL_FA4_CUTE"): + key = arg.lower() + assert f"{arg}=${{{{ matrix.cuda.{key}[matrix.architecture.name] }}}}" in job["with"]["build_args"] dockerfile = (REPO_ROOT / "docker/Dockerfile").read_text(encoding="utf-8") assert "ARG INSTALL_FA4_CUTE=1" in dockerfile + assert "ARG INSTALL_FLASH_ATTN=1" in dockerfile + assert 'if [ "${INSTALL_FLASH_ATTN}" != "1" ]; then' in dockerfile assert 'if [ "${INSTALL_FA4_CUTE}" = "1" ]; then' in dockerfile