nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[infra]: leave FlashAttention out of the GB10 image

Browse Source 
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).
This commit is contained in:
SolitaryThinker
2026-10-19 11:15:23 +00:00
parent fba50356b1
commit 2e4878c5e6
4 changed files with 57 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/infra-build-image.yml
+9
View File
@@ -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 }}
docker/Dockerfile
+17 -5
View File
@@ -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 . .
docs/contributing/developer_env/docker.md
+22 -8
View File
@@ -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
fastvideo/tests/modal/test_kernel_build_cache.py
+9 -3
View File
@@ -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