Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 root
|
23c285bfb2 | ||
|
root
|
ceec06dc31 |
@@ -90,13 +90,6 @@ steps:
|
||||
limit: 2
|
||||
agents:
|
||||
queue: "default"
|
||||
- label: ":microscope: Unit Tests"
|
||||
key: "unit-ci"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "unit_test_ci"
|
||||
command: "/opt/fastvideo-ci-runner/run-unit"
|
||||
timeout_in_minutes: 90
|
||||
agents:
|
||||
queue: "ci-runner"
|
||||
- label: ":microscope: DreamVerse App Tests"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "dreamverse_app"
|
||||
command: "timeout 90m .buildkite/scripts/pr_test.sh"
|
||||
|
||||
@@ -1,22 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
exec pytest \
|
||||
./fastvideo/tests/api/ \
|
||||
./fastvideo/tests/contract/ \
|
||||
./fastvideo/tests/dataset/ \
|
||||
./fastvideo/tests/workflow/ \
|
||||
./fastvideo/tests/entrypoints/ \
|
||||
./fastvideo/tests/train/ \
|
||||
./fastvideo/tests/stages/ \
|
||||
./fastvideo/tests/ops/ \
|
||||
./fastvideo/tests/worker/ \
|
||||
./fastvideo/tests/training/test_trackers.py \
|
||||
./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py \
|
||||
./fastvideo/tests/modal/test_kernel_build_cache.py \
|
||||
./fastvideo/tests/modal/test_pr_test.py \
|
||||
./fastvideo/tests/modal/test_ssim_test.py \
|
||||
--ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py \
|
||||
--ignore=./fastvideo/tests/train/models \
|
||||
--ignore=./fastvideo/tests/train/methods \
|
||||
-vs
|
||||
@@ -190,7 +190,6 @@ jobs:
|
||||
if: ${{ !inputs.push_by_digest }}
|
||||
run: |
|
||||
echo "✅ Python ${{ inputs.python_version }} image successfully built and pushed to ${{ steps.image.outputs.name }}:${{ inputs.tag_suffix }}-sha-${GITHUB_SHA::7}"
|
||||
echo "Digest: ${{ steps.build-push.outputs.digest }}"
|
||||
echo "To run tests with this image, manually trigger the 'Run Tests' workflow."
|
||||
|
||||
- name: Digest success message
|
||||
|
||||
@@ -129,7 +129,7 @@ jobs:
|
||||
set -euo pipefail
|
||||
TEST_NAME=$(echo "$COMMENT" | grep -oP '(?<=/test\s)\S+' | head -1 || true)
|
||||
|
||||
VALID="encoder vae transformer kernel unit unit-ci dreamverse ssim golden-gate training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit"
|
||||
VALID="encoder vae transformer kernel unit dreamverse ssim golden-gate training lora-inference lora-training lora-extraction distillation self-forcing vsa vmoba performance api train-framework eval full fastcheck pre-commit"
|
||||
if [ -z "$TEST_NAME" ] || ! echo "$VALID" | grep -qw "$TEST_NAME"; then
|
||||
echo "Unknown test: '$TEST_NAME'. Valid: $VALID"
|
||||
exit 1
|
||||
@@ -137,8 +137,7 @@ jobs:
|
||||
|
||||
declare -A MAP=(
|
||||
[encoder]=encoder [vae]=vae [transformer]=transformer
|
||||
[kernel]=kernel_tests [unit]=unit_test [unit-ci]=unit_test_ci
|
||||
[dreamverse]=dreamverse_app
|
||||
[kernel]=kernel_tests [unit]=unit_test [dreamverse]=dreamverse_app
|
||||
[ssim]=ssim [golden-gate]=golden_gate [training]=training
|
||||
[lora-inference]=inference_lora [lora-training]=training_lora
|
||||
[lora-extraction]=lora_extraction
|
||||
|
||||
@@ -13,11 +13,6 @@ on:
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
build_ci_runner_image:
|
||||
description: 'Build the ARM64 CUDA 13 CI runner image (sm_100)'
|
||||
required: false
|
||||
default: false
|
||||
type: boolean
|
||||
# Auto-rebuild the CUDA images when a repository-controlled image input
|
||||
# changes on main. This includes the trusted SM89 kernel artifact's source,
|
||||
# metadata/key helper, ABI dependency metadata, and build orchestration.
|
||||
@@ -203,28 +198,6 @@ jobs:
|
||||
docker buildx imagetools create "${TAG_ARGS[@]}" "${IMAGE_REFS[@]}"
|
||||
docker buildx imagetools inspect "${TAGS[0]}"
|
||||
|
||||
# The CI runner is ARM64 like DGX Spark, but its GB200 GPUs are sm_100 rather
|
||||
# than sm_121. Publish a single-architecture variant so Slurm CI can reuse
|
||||
# the exact prebuilt kernel instead of compiling it in every job.
|
||||
build-ci-runner-image:
|
||||
if: ${{ (github.event_name == 'push' && github.repository == 'hao-ai-lab/FastVideo') || github.event.inputs.build_ci_runner_image == 'true' }}
|
||||
uses: ./.github/workflows/_template-build-image.yml
|
||||
with:
|
||||
python_version: '3.12'
|
||||
dockerfile_path: docker/Dockerfile
|
||||
tag_suffix: py3.12-cuda13.0.0-sm100
|
||||
runner: ubuntu-24.04-arm
|
||||
architecture: arm64
|
||||
build_args: |
|
||||
PYTHON_VERSION=3.12
|
||||
CUDA_VERSION=13.0.0
|
||||
UV_TORCH_BACKEND=cu130
|
||||
TORCH_CUDA_ARCH_LIST=10.0
|
||||
CMAKE_BUILD_PARALLEL_LEVEL=1
|
||||
FLASH_ATTN_WHEEL_TAG=cu130torch2.12
|
||||
FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.9.22
|
||||
secrets: inherit
|
||||
|
||||
# Dreamverse matrix: {backend, UI} x {12.6.3, 13.0.0}, Python 3.12. Torch backend
|
||||
# matches the base CUDA (cu126 / cu130). Keep these images amd64-only until the
|
||||
# required FA4 dependency stack is available and validated on arm64.
|
||||
|
||||
+1
-5
@@ -68,7 +68,7 @@ ARG FLASH_ATTN_WHEEL_RELEASE_ARM64=https://github.com/mjun0812/flash-attention-p
|
||||
# cutlass-4.4 `cute.core.ThrMma` API, which crashes on the cutlass-dsl 4.5 that
|
||||
# flashinfer/quack pull in. After the wheel install we overlay this cutlass-4.5-safe
|
||||
# upstream cute (flash-attn-4) so the image runs FA4 instead of the FA2 fallback.
|
||||
ARG FA4_CUTE_REF=14c377950125c70b7a9dabf9c561fca53715ac7d
|
||||
ARG FA4_CUTE_REF=82d6441eec5d4dfec120153db2c0145ae855a083
|
||||
|
||||
# Provided automatically by BuildKit/buildx (e.g. "amd64" / "arm64") and used to
|
||||
# select the prebuilt flash-attn wheel. Empty under a plain `docker build` without
|
||||
@@ -89,9 +89,6 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
zsh \
|
||||
vim \
|
||||
curl \
|
||||
ffmpeg \
|
||||
libgl1 \
|
||||
libglib2.0-0 \
|
||||
gcc-11 \
|
||||
g++-11 \
|
||||
clang-11 \
|
||||
@@ -141,7 +138,6 @@ RUN --mount=type=cache,target=/opt/uv/cache \
|
||||
source /opt/venv/bin/activate && \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install --excludes docker/uv-excludes ".[dev]" && \
|
||||
python -c "import cv2; print('OpenCV', cv2.__version__)" && \
|
||||
PYTAG=cp$(echo "${PYTHON_VERSION}" | tr -d .) && \
|
||||
case "${TARGETARCH:-amd64}" in \
|
||||
amd64) \
|
||||
|
||||
@@ -9,33 +9,82 @@ FastVideo exposes a process-wide torch profiler that you can enable via environm
|
||||
|
||||
```bash
|
||||
FASTVIDEO_TORCH_PROFILER_DIR=/mnt/traces/fastvideo \
|
||||
FASTVIDEO_TORCH_PROFILE_REGIONS="profiler_region_model_loading,profiler_region_training_step"
|
||||
FASTVIDEO_TORCH_PROFILE_REGIONS="model_loading,training_train_one_step" \
|
||||
bash examples/train/run.sh /path/to/config.yaml
|
||||
```
|
||||
|
||||
All profiled regions must be registered in `fastvideo.profiler`; the current list includes:
|
||||
The `profiler_region_` prefix is optional. All profiled regions must be
|
||||
registered in `fastvideo.profiler`; the current list includes:
|
||||
|
||||
- `profiler_region_model_loading` — pipeline/module loading
|
||||
- `profiler_region_inference_pre_denoising`
|
||||
- `profiler_region_inference_denoising`
|
||||
- `profiler_region_inference_post_denoising`
|
||||
- `profiler_region_training_checkpoint_saving`
|
||||
- `profiler_region_training_dit`
|
||||
- `profiler_region_training_train` — the complete training run
|
||||
- `profiler_region_training_train_one_step` — one complete optimizer step
|
||||
- `profiler_region_training_dataloader` — fetch the next batch in the trainer process
|
||||
- `profiler_region_training_forward` — method forward and loss computation
|
||||
- `profiler_region_training_backward` — backward pass
|
||||
- `profiler_region_training_optimizer` — gradient clipping, optimizer/scheduler step, and zeroing gradients
|
||||
- `profiler_region_training_callbacks` — end-of-step callbacks such as EMA updates
|
||||
- `profiler_region_training_validation`
|
||||
- `profiler_region_training_epoch`
|
||||
- `profiler_region_training_step`
|
||||
- `profiler_region_training_backward`
|
||||
- `profiler_region_training_optimizer`
|
||||
- `profiler_region_training_save_checkpoint`
|
||||
- `profiler_region_distillation_teacher_forward`
|
||||
- `profiler_region_distillation_student_forward`
|
||||
- `profiler_region_distillation_loss`
|
||||
- `profiler_region_distillation_update`
|
||||
- `profiler_region_dmd2_student_rollout` — one DMD2 student rollout, including simulated prefix steps
|
||||
- `profiler_region_dmd2_generator_loss` — teacher/critic scoring for the generator loss
|
||||
- `profiler_region_dmd2_critic_loss` — critic flow-matching loss, including its student rollout
|
||||
|
||||
### Profiling modular training
|
||||
|
||||
The YAML-driven trainer under `fastvideo/train/` initializes and flushes the
|
||||
profiler automatically. For example, this captures six DMD2 steps: one warmup
|
||||
step, four steady-state critic updates, and the configured 1-in-5 generator
|
||||
update, without changing the method's update cadence:
|
||||
|
||||
```bash
|
||||
TRACE_DIR="$(pwd)/profiler_traces/wan_dmd2"
|
||||
mkdir -p "$TRACE_DIR"
|
||||
|
||||
NUM_GPUS=4 \
|
||||
WANDB_MODE=offline \
|
||||
FASTVIDEO_TORCH_PROFILER_DIR="$TRACE_DIR" \
|
||||
FASTVIDEO_TORCH_PROFILE_REGIONS="training_train,training_train_one_step,training_dataloader,training_forward,training_backward,training_optimizer,training_callbacks,dmd2_student_rollout,dmd2_generator_loss,dmd2_critic_loss" \
|
||||
bash examples/train/run.sh \
|
||||
examples/train/configs/distribution_matching/wan/dmd2_t2v.yaml \
|
||||
--training.loop.max_train_steps 6 \
|
||||
--training.checkpoint.training_state_checkpointing_steps 0 \
|
||||
--callbacks.validation.every_steps 0
|
||||
```
|
||||
|
||||
The two overrides disable validation and checkpoint writes so their I/O does
|
||||
not contaminate training-step measurements. Remove them when profiling those
|
||||
regions. Methods that manage optimization internally emit the enclosing
|
||||
`training_train_one_step` region but do not emit the generic
|
||||
dataloader/forward/backward/optimizer child regions because those boundaries
|
||||
belong to the method. With `dataloader_num_workers > 0`, the
|
||||
`training_dataloader` region measures the training process waiting for and
|
||||
receiving the next batch; work performed inside dataloader worker subprocesses
|
||||
does not appear in that process's torch-profiler trace.
|
||||
|
||||
The modular trainer also logs `dataloader_time_sec` on every ordinary training
|
||||
step (including DMD2), so routine runs can monitor rank 0's
|
||||
`next(dataloader)` wait without enabling the torch profiler. It is
|
||||
intentionally not reduced across ranks; cross-rank synchronization would
|
||||
perturb the hot path being measured.
|
||||
|
||||
While profiling is enabled, FastVideo records additional annotations:
|
||||
|
||||
- `fastvideo.region::<name>` spans are emitted when entering a region.
|
||||
- `fastvideo.profiler.enable_collection` / `fastvideo.profiler.disable_collection` events mark when torch profiler collection is toggled on or off.
|
||||
|
||||
Only one profiler instance is created per process; subsequent pipelines reuse the same controller. If you set `FASTVIDEO_TORCH_PROFILE_REGIONS` incorrectly (e.g. misspelled name), FastVideo logs a warning and ignores that entry.
|
||||
Only one profiler controller is created per process; subsequent pipelines
|
||||
reuse it. Each outermost enabled region invocation produces one complete
|
||||
CPU/CUDA trace segment. Enabled regions nested inside it are annotations in
|
||||
that segment. Include an enclosing region such as `training_train_one_step`
|
||||
when selecting its child regions so they share one trace and profiler startup
|
||||
does not perturb each phase independently. If you set
|
||||
`FASTVIDEO_TORCH_PROFILE_REGIONS` incorrectly (e.g. misspelled name), FastVideo
|
||||
logs a warning and ignores that entry.
|
||||
|
||||
Additional knobs:
|
||||
|
||||
@@ -44,13 +93,18 @@ Additional knobs:
|
||||
- `FASTVIDEO_TORCH_PROFILER_WITH_STACK`
|
||||
- `FASTVIDEO_TORCH_PROFILER_WITH_FLOPS`
|
||||
|
||||
Traces can be visualized using <https://ui.perfetto.dev/>.
|
||||
Traces can be visualized using <https://ui.perfetto.dev/>. Each rank also
|
||||
writes `summary_rank<N>_segment<MMMM>_<region>.txt` and the corresponding JSON
|
||||
beside every trace segment for a quick operator-level view without loading the
|
||||
full timeline. The JSON contains CPU and device totals; input-shape grouping is
|
||||
only enabled with `FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES=1` because it makes
|
||||
summary generation materially more expensive on large traces.
|
||||
|
||||
### Best Practices
|
||||
|
||||
- Keep the profiled step count small; traces can be large and slow down job shutdown while the profiler flushes data.
|
||||
- After profiling, clean up trace directories to avoid filling disk storage.
|
||||
- When adding new regions, register them in `fastvideo.profiler` and wrap the corresponding code block with `with self.profiler_controller.region("your_region"):` or the `@profile_region` decorator.
|
||||
- When adding new regions, register them in `fastvideo.profiler` and wrap the corresponding code block with `profiler_region("your_region")` or the `@profile_region` decorator.
|
||||
|
||||
## Related: Activation Trace Mode
|
||||
|
||||
|
||||
@@ -119,14 +119,12 @@ surfaces:
|
||||
vae_precision: "Precision override pending dedicated typed component precision design."
|
||||
vae_decode_precision: "Decode-only precision override pending dedicated typed component precision design."
|
||||
image_encoder_precision: "Precision override pending dedicated typed component precision design."
|
||||
image_encoder_precisions: "Precision overrides pending dedicated typed component precision design."
|
||||
text_encoder_precisions: "Precision override pending dedicated typed component precision design."
|
||||
internal_only:
|
||||
dit_config: "Legacy internal component config object."
|
||||
upsampler_config: "Legacy internal component config object."
|
||||
vae_config: "Legacy internal component config object."
|
||||
image_encoder_config: "Legacy internal component config object."
|
||||
image_encoder_configs: "Legacy internal component config objects."
|
||||
text_encoder_configs: "Legacy internal component config object."
|
||||
preprocess_text_funcs: "Internal text preprocessing hooks."
|
||||
postprocess_text_funcs: "Internal text postprocessing hooks."
|
||||
@@ -363,30 +361,6 @@ surfaces:
|
||||
sources: [fastvideo.configs.pipelines.matrixgame2.MatrixGame2I2V480PConfig]
|
||||
num_frames_per_block:
|
||||
sources: [fastvideo.configs.pipelines.matrixgame2.MatrixGame2I2V480PConfig]
|
||||
duration_s:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
spectrogram_frame_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
latent_downsample_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
clip_frame_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_frame_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_segment_size:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_segment_stride:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_downsample_rate:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
clip_image_size:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_image_size:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
clip_batch_size_multiplier:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
sync_batch_size_multiplier:
|
||||
sources: [fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig]
|
||||
audio_channels:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
|
||||
@@ -401,7 +375,6 @@ surfaces:
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
|
||||
max_audio_duration_s:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
|
||||
sample_size:
|
||||
@@ -410,7 +383,6 @@ surfaces:
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
|
||||
sampling_rate:
|
||||
sources:
|
||||
- fastvideo.configs.pipelines.mmaudio.MMAudioV2AConfig
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioT2AConfig
|
||||
- fastvideo.configs.pipelines.stable_audio.StableAudioOpenSmallConfig
|
||||
audio_txt_guidance_scale:
|
||||
|
||||
@@ -431,7 +431,9 @@ callbacks:
|
||||
```
|
||||
|
||||
The EMA callback owns its own state and checkpoints independently — EMA weights
|
||||
are saved and restored automatically on resume.
|
||||
are saved and restored automatically on resume. It advances only after the
|
||||
training method reports that the student optimizer updated; alternating methods
|
||||
such as DMD2 therefore skip EMA work on critic-only steps.
|
||||
|
||||
### ValidationCallback
|
||||
|
||||
|
||||
+11
-40
@@ -17,7 +17,7 @@ Runtime-JIT kernels (no build step, ship in every wheel/image):
|
||||
| Kernels | Where | Used when |
|
||||
|---|---|---|
|
||||
| Triton: STA, VSA block-sparse, SLA, fused compress+topk, FP4 QAT training, quant/norm utils | `python/fastvideo_kernel/triton_kernels/` | automatic fallback when the matching C++ op is absent (`ops.py`, `turbodiffusion_ops.py`) |
|
||||
| FA4 CuTe-DSL block-sparse forward/backward (VSA-128/256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
|
||||
| FA4 CuTe-DSL block-sparse (VSA-256 fastpath on `sm_100`) | `block_sparse_attn_cute_fwd.py` | optional `flash_attn.cute` dependency, see below |
|
||||
| VMoBA `moba_attn_varlen` | `vmoba.py` | wraps flash-attn varlen |
|
||||
|
||||
## What gets built where, and when
|
||||
@@ -64,49 +64,28 @@ cd fastvideo-kernel
|
||||
./build.sh --rocm
|
||||
```
|
||||
|
||||
### Optional: FA4 CuTe block-sparse backend (VSA-128/256 fastpath)
|
||||
### Optional: FA4 CuTe block-sparse backend (VSA-256 fastpath)
|
||||
|
||||
The VSA-128/256 fastpaths (tile volume 128 or 256, on NVIDIA Blackwell / sm_100) route to the
|
||||
The VSA-256 fastpath (tile volume 256, on NVIDIA Blackwell / sm_100) routes to the
|
||||
FlashAttention-4 CuTe-DSL block-sparse kernel exposed as `flash_attn.cute`. This is
|
||||
an **optional** dependency: it is imported lazily, and `video_sparse_attn`
|
||||
transparently falls back to the Triton backend when it is absent (so the package is
|
||||
fully usable without it).
|
||||
|
||||
The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`
|
||||
and the public/private forward-backward bridges in `flash_attn.cute.interface`) are provided upstream by
|
||||
The symbols the fastpath needs (`flash_attn.cute.block_sparsity.BlockSparseTensorsTorch`,
|
||||
`flash_attn.cute.interface._flash_attn_fwd`) are provided upstream by
|
||||
[Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention). Pin to
|
||||
commit `14c377950125c70b7a9dabf9c561fca53715ac7d`, the revision FastVideo pins as
|
||||
commit `940cd9680f3315f2f06b43ab5bea2c2cf2d96806`, the revision FastVideo pins as
|
||||
the `flash-attn-4` source in the repo-root `pyproject.toml`; other revisions may
|
||||
have incompatible block-sparse forward/backward interfaces.
|
||||
|
||||
Install it under its distribution name so its own runtime stack resolves with it.
|
||||
Do **not** pre-install `nvidia-cutlass-dsl` by hand: this revision pins
|
||||
`nvidia-cutlass-dsl==4.6.0.dev0` exactly, and a hand-installed 4.5.x floor either
|
||||
gets silently upgraded or, if something else holds it back, leaves the CuTe
|
||||
kernels broken.
|
||||
have an incompatible `_flash_attn_fwd` signature.
|
||||
|
||||
```bash
|
||||
pip install torchvision
|
||||
pip install "flash-attn-4 @ git+https://github.com/Dao-AILab/flash-attention.git@14c377950125c70b7a9dabf9c561fca53715ac7d#subdirectory=flash_attn/cute"
|
||||
pip install "nvidia-cutlass-dsl>=4.5.0" torchvision
|
||||
pip install "git+https://github.com/Dao-AILab/flash-attention.git@940cd9680f3315f2f06b43ab5bea2c2cf2d96806#subdirectory=flash_attn/cute"
|
||||
```
|
||||
|
||||
That resolves `nvidia-cutlass-dsl` to 4.6.0.dev0 and `quack-kernels` to 0.5.3, a
|
||||
combination this revision works with. A mismatched CuTe DSL only surfaces when the
|
||||
kernel JIT-compiles, so the error points at CuTe internals rather than at the
|
||||
install:
|
||||
|
||||
| Error on first VSA-128/256 CuTe call | Cause |
|
||||
|---|---|
|
||||
| `TypeError: fmax() missing 1 required positional argument: 'b'` | `nvidia-cutlass-dsl` 4.5.x |
|
||||
| `AttributeError: module 'cutlass.cute.core' has no attribute 'ThrMma'` | `quack-kernels` older than 0.5.1 |
|
||||
| `ImportError: cannot import name 'alloc_reserved_mbarrier'` | `quack-kernels` 0.6.2 or newer |
|
||||
|
||||
An environment whose `flash_attn.cute` came from a prebuilt flash-attn wheel rather
|
||||
than from this pin hits the first row; that is what the overlay step in
|
||||
`docker/Dockerfile` works around.
|
||||
|
||||
The CuTe kernels JIT-compile on first use. Forward and backward are verified on
|
||||
Blackwell (sm_100) against `tests/test_vsa128_*.py` and `tests/test_vsa256_*.py`.
|
||||
The CuTe kernel JIT-compiles on first use. Verified on Blackwell (sm_100) against
|
||||
`tests/test_vsa256_forward*.py`.
|
||||
|
||||
## Usage
|
||||
|
||||
@@ -163,14 +142,6 @@ After building/installing `fastvideo-kernel`, run:
|
||||
```bash
|
||||
cd fastvideo-kernel
|
||||
python benchmarks/bench_vsa.py --batch_size 1 --num_heads 16 --head_dim 128 --q_seq_lens 49152 --topk 64
|
||||
|
||||
# VSA-256 FA4 CuTe forward/backward on Blackwell
|
||||
python benchmarks/bench_vsa.py --block_size 256 --use_cute \
|
||||
--batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 20
|
||||
|
||||
# VSA-128 FA4 CuTe forward/backward on Blackwell
|
||||
python benchmarks/bench_vsa.py --block_size 128 --use_cute \
|
||||
--batch_size 1 --num_heads 12 --head_dim 128 --q_seq_lens 39936 --topk 40
|
||||
```
|
||||
|
||||
### TurboDiffusion Kernels
|
||||
|
||||
@@ -2,9 +2,8 @@
|
||||
"""
|
||||
Benchmark VSA *wrapper* performance (forward + backward) and report TFLOPs.
|
||||
|
||||
This script benchmarks the autograd-enabled wrappers:
|
||||
- 64-token TK/Triton: fastvideo_kernel.block_sparse_attn.block_sparse_attn
|
||||
- 128/256-token Triton/CuTe: fastvideo_kernel.block_sparse_attn_256
|
||||
This script benchmarks the autograd-enabled wrapper:
|
||||
- fastvideo_kernel.block_sparse_attn.block_sparse_attn
|
||||
|
||||
So measured time includes wrapper overhead (map->index conversion, dispatch) plus kernel time.
|
||||
"""
|
||||
@@ -24,6 +23,9 @@ try:
|
||||
except Exception as e: # pragma: no cover
|
||||
raise ImportError("This benchmark requires triton (for triton.testing.do_bench).") from e
|
||||
|
||||
BLOCK_M = 64
|
||||
BLOCK_N = 64
|
||||
|
||||
|
||||
def set_seed(seed: int = 42) -> None:
|
||||
random.seed(seed)
|
||||
@@ -39,11 +41,7 @@ def parse_arguments() -> argparse.Namespace:
|
||||
p.add_argument("--num_heads", type=int, default=12)
|
||||
p.add_argument("--head_dim", type=int, default=128, choices=[64, 128])
|
||||
p.add_argument("--topk", type=int, default=None, help="KV blocks per Q block (default: ~90%% sparsity)")
|
||||
p.add_argument("--q_seq_lens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[49152],
|
||||
help="Q sequence lengths (must be divisible by --block_size)")
|
||||
p.add_argument("--q_seq_lens", type=int, nargs="+", default=[49152], help="Q sequence lengths (must be /64)")
|
||||
p.add_argument("--kv_seq_lens",
|
||||
type=int,
|
||||
nargs="+",
|
||||
@@ -53,13 +51,9 @@ def parse_arguments() -> argparse.Namespace:
|
||||
p.add_argument("--rep", type=int, default=20)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp16"])
|
||||
p.add_argument("--block_size", type=int, default=64, choices=[64, 128, 256])
|
||||
p.add_argument("--force_triton",
|
||||
action="store_true",
|
||||
help="Force wrapper to use Triton path (if supported by shapes).")
|
||||
p.add_argument("--use_cute",
|
||||
action="store_true",
|
||||
help="Use the optional FA4 CuTe forward/backward path (requires --block_size 128 or 256).")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
@@ -90,38 +84,18 @@ def bench_ms(fn: Callable[[], object], warmup: int, rep: int) -> float:
|
||||
return do_bench(fn, warmup=warmup, rep=rep, quantiles=None)
|
||||
|
||||
|
||||
def _configure_backend(args: argparse.Namespace) -> None:
|
||||
if args.use_cute and args.block_size not in (128, 256):
|
||||
raise ValueError("--use_cute requires --block_size 128 or 256")
|
||||
if args.use_cute and args.force_triton:
|
||||
raise ValueError("--use_cute and --force_triton are mutually exclusive")
|
||||
|
||||
if args.force_triton:
|
||||
os.environ.pop("FASTVIDEO_VSA_CUTEDSL", None)
|
||||
os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
|
||||
elif args.use_cute:
|
||||
os.environ.pop("FASTVIDEO_VSA_TRITON", None)
|
||||
os.environ.pop("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", None)
|
||||
os.environ["FASTVIDEO_VSA_CUTEDSL"] = "1"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_arguments()
|
||||
set_seed(args.seed)
|
||||
_configure_backend(args)
|
||||
|
||||
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
|
||||
|
||||
if args.force_triton:
|
||||
os.environ["FASTVIDEO_KERNEL_VSA_FORCE_TRITON"] = "1"
|
||||
|
||||
from fastvideo_kernel.block_sparse_attn import block_sparse_attn
|
||||
from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128, block_sparse_attn_256
|
||||
|
||||
bs, h, d = args.batch_size, args.num_heads, args.head_dim
|
||||
block_size = args.block_size
|
||||
attention = {
|
||||
64: block_sparse_attn,
|
||||
128: block_sparse_attn_128,
|
||||
256: block_sparse_attn_256,
|
||||
}[block_size]
|
||||
kv_seq_lens = args.kv_seq_lens
|
||||
if kv_seq_lens is None:
|
||||
kv_seq_lens = args.q_seq_lens
|
||||
@@ -131,22 +105,20 @@ def main() -> None:
|
||||
print("VSA Block-Sparse Attention Benchmark (WRAPPER)")
|
||||
print(f"device: {torch.cuda.get_device_name(0)}")
|
||||
print(f"batch={bs}, heads={h}, head_dim={d}, dtype={args.dtype}")
|
||||
print(f"block_size={block_size}")
|
||||
print(f"BLOCK_M={BLOCK_M}, BLOCK_N={BLOCK_N}")
|
||||
print("NOTE: timings include wrapper overhead (map->index + dispatch).")
|
||||
if args.use_cute:
|
||||
print("dispatch: FA4 CuTe")
|
||||
elif args.force_triton:
|
||||
if args.force_triton:
|
||||
print("dispatch: forced Triton (FASTVIDEO_KERNEL_VSA_FORCE_TRITON=1)")
|
||||
else:
|
||||
print("dispatch: SM90 if available, else Triton")
|
||||
|
||||
for q_len, kv_len in zip(args.q_seq_lens, kv_seq_lens):
|
||||
if q_len % block_size != 0 or kv_len % block_size != 0:
|
||||
print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by {block_size}")
|
||||
if q_len % BLOCK_M != 0 or kv_len % BLOCK_N != 0:
|
||||
print(f"[skip] q_len={q_len}, kv_len={kv_len} must be divisible by 64")
|
||||
continue
|
||||
|
||||
num_q_blocks = q_len // block_size
|
||||
num_kv_blocks = kv_len // block_size
|
||||
num_q_blocks = q_len // BLOCK_M
|
||||
num_kv_blocks = kv_len // BLOCK_N
|
||||
topk = args.topk if args.topk is not None else max(1, num_kv_blocks // 10)
|
||||
topk = min(topk, num_kv_blocks)
|
||||
|
||||
@@ -157,11 +129,11 @@ def main() -> None:
|
||||
q, k, v = create_qkv(bs, h, q_len, kv_len, d, dtype)
|
||||
block_map = make_block_map(bs, h, num_q_blocks, num_kv_blocks, topk)
|
||||
|
||||
# Variable block sizes: default full logical blocks.
|
||||
variable_block_sizes = torch.full((num_kv_blocks, ), block_size, dtype=torch.int32, device="cuda")
|
||||
# Variable block sizes: default full blocks (64 tokens per KV block)
|
||||
variable_block_sizes = torch.full((num_kv_blocks, ), BLOCK_N, dtype=torch.int32, device="cuda")
|
||||
|
||||
def _fwd():
|
||||
return attention(q, k, v, block_map, variable_block_sizes)
|
||||
return block_sparse_attn(q, k, v, block_map, variable_block_sizes)
|
||||
|
||||
fwd_ms = bench_ms(_fwd, warmup=args.warmup, rep=args.rep)
|
||||
|
||||
@@ -170,7 +142,7 @@ def main() -> None:
|
||||
q_ = q.detach().requires_grad_(True)
|
||||
k_ = k.detach().requires_grad_(True)
|
||||
v_ = v.detach().requires_grad_(True)
|
||||
o_, _aux_ = attention(q_, k_, v_, block_map, variable_block_sizes)
|
||||
o_, _aux_ = block_sparse_attn(q_, k_, v_, block_map, variable_block_sizes)
|
||||
og = torch.randn_like(o_)
|
||||
loss = (o_ * og).sum()
|
||||
|
||||
@@ -184,7 +156,7 @@ def main() -> None:
|
||||
rep=max(5, args.rep // 2),
|
||||
)
|
||||
|
||||
flops = flops_sparse_attention(bs, h, d, q_len, topk, block_size)
|
||||
flops = flops_sparse_attention(bs, h, d, q_len, topk, BLOCK_N)
|
||||
fwd_tflops = flops / fwd_ms * 1e-12 * 1e3
|
||||
# Rough backward multiplier (attention backward typically ~2-3x forward)
|
||||
bwd_tflops = (2.5 * flops) / bwd_ms * 1e-12 * 1e3
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
"""VSA-128/256 block-sparse attention wrappers.
|
||||
"""VSA-256 block-sparse attention wrapper.
|
||||
|
||||
The default 256-block path is Triton: it expands the logical 256-block map
|
||||
to the existing 64-block Triton kernel via a dense 4x4 expansion per logical
|
||||
@@ -7,8 +7,8 @@ edge ("route A"), and requires no optional dependencies.
|
||||
The FA4 CuTe block-sparse fastpath (intended for Blackwell sm_100+) is
|
||||
*opt-in* via ``FASTVIDEO_VSA_CUTEDSL=1``. It routes to
|
||||
:mod:`fastvideo_kernel.block_sparse_attn_cute_fwd`, which natively operates
|
||||
on 128-token Q/KV blocks (the 256 wrapper expands its logical KV map and
|
||||
sizes into that physical representation). The CuTe kernel
|
||||
on 128-token KV blocks (this wrapper expands the logical 256-block map /
|
||||
sizes into that physical 128-block representation). The CuTe kernel
|
||||
(``flash_attn.cute`` with block-sparsity) is an optional dependency,
|
||||
imported lazily only when this fastpath is selected.
|
||||
|
||||
@@ -35,7 +35,7 @@ _KV_BLOCK_TRITON = 64 # Existing Triton path uses 64-token KV blocks.
|
||||
|
||||
|
||||
def _resolve_backend() -> str:
|
||||
"""Pick the backend for the 128/256-block VSA paths.
|
||||
"""Pick the backend for the 256-block VSA path.
|
||||
|
||||
Default is Triton (no optional deps). The FA4 CuTe fastpath is opt-in
|
||||
via ``FASTVIDEO_VSA_CUTEDSL=1`` and requires the optional FA4 CuTe
|
||||
@@ -49,26 +49,6 @@ def _resolve_backend() -> str:
|
||||
return "triton"
|
||||
|
||||
|
||||
def _expand_mask_and_sizes_128_to_64(
|
||||
logical_mask_128: torch.Tensor,
|
||||
logical_kv_sizes_128: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Expand a [B, H, Qb128, KVb128] map to 64-token Triton tiles."""
|
||||
expanded_mask = logical_mask_128.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3)
|
||||
sizes_i32 = logical_kv_sizes_128.to(torch.int32)
|
||||
offsets = torch.tensor(
|
||||
[0, _KV_BLOCK_TRITON],
|
||||
dtype=torch.int32,
|
||||
device=sizes_i32.device,
|
||||
)
|
||||
expanded_sizes = torch.clamp(
|
||||
sizes_i32[:, None] - offsets[None, :],
|
||||
min=0,
|
||||
max=_KV_BLOCK_TRITON,
|
||||
).reshape(-1)
|
||||
return expanded_mask, expanded_sizes
|
||||
|
||||
|
||||
def _expand_mask_and_sizes_256_to_128(
|
||||
logical_mask_256: torch.Tensor,
|
||||
logical_kv_sizes_256: torch.Tensor,
|
||||
@@ -132,63 +112,6 @@ def _triton_via_route_a(
|
||||
return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
|
||||
|
||||
|
||||
def _triton_via_route_a_128(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
logical_mask_128: torch.Tensor,
|
||||
logical_kv_sizes_128: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
from .triton_kernels.index import map_to_index as triton_map_to_index
|
||||
|
||||
mask_64, sizes_64 = _expand_mask_and_sizes_128_to_64(logical_mask_128, logical_kv_sizes_128)
|
||||
q2k_idx, q2k_num = triton_map_to_index(mask_64.to(torch.bool))
|
||||
return block_sparse_attn_triton(q, k, v, q2k_idx, q2k_num, sizes_64)
|
||||
|
||||
|
||||
def block_sparse_attn_128(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
logical_block_map_128: torch.Tensor,
|
||||
logical_variable_block_sizes_128: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""VSA-128 sparse-branch entrypoint for [B, H, S, D] inputs."""
|
||||
if logical_block_map_128.dim() == 3:
|
||||
logical_block_map_128 = logical_block_map_128.unsqueeze(0)
|
||||
|
||||
if _resolve_backend() == "triton":
|
||||
return _triton_via_route_a_128(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
|
||||
|
||||
from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd
|
||||
return block_sparse_attn_cute_fwd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
|
||||
|
||||
|
||||
def block_sparse_attn_128_bshd(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
logical_block_map_128: torch.Tensor,
|
||||
logical_variable_block_sizes_128: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""VSA-128 sparse-branch entrypoint for [B, S, H, D] inputs."""
|
||||
if logical_block_map_128.dim() == 3:
|
||||
logical_block_map_128 = logical_block_map_128.unsqueeze(0)
|
||||
|
||||
if _resolve_backend() == "triton":
|
||||
out_bhsd, aux = _triton_via_route_a_128(
|
||||
q.transpose(1, 2).contiguous(),
|
||||
k.transpose(1, 2).contiguous(),
|
||||
v.transpose(1, 2).contiguous(),
|
||||
logical_block_map_128,
|
||||
logical_variable_block_sizes_128,
|
||||
)
|
||||
return out_bhsd.transpose(1, 2).contiguous(), aux
|
||||
|
||||
from .block_sparse_attn_cute_fwd import block_sparse_attn_cute_fwd_bshd
|
||||
return block_sparse_attn_cute_fwd_bshd(q, k, v, logical_block_map_128, logical_variable_block_sizes_128)
|
||||
|
||||
|
||||
def block_sparse_attn_256(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
|
||||
@@ -1,18 +1,18 @@
|
||||
"""FA4 CuTe-DSL block-sparse attention adapter.
|
||||
"""CuTe-DSL block-sparse attention forward kernel.
|
||||
|
||||
This module adapts VSA's ``(block_map, variable_block_sizes)`` inputs into
|
||||
FA4's forward and backward ``BlockSparseTensorsTorch`` representations.
|
||||
FA4's public ``flash_attn_func`` owns the forward/backward autograd bridge.
|
||||
Thin wrapper around `flash_attn.cute.interface._flash_attn_fwd` that adapts
|
||||
VSA's `(block_map, variable_block_sizes)` inputs into FA4's
|
||||
`BlockSparseTensorsTorch` representation and the per-KV-block validity mask.
|
||||
|
||||
Both [B, H, S, D] (BHSD) and [B, S, H, D] (BSHD) entrypoints are provided.
|
||||
The BSHD variant is preferred from VSA-128/256 callers to avoid layout
|
||||
The BSHD variant is preferred from VSA-256 callers to avoid layout
|
||||
round-trips on the hot path.
|
||||
|
||||
The FA4 CuTe block-sparse kernel (``flash_attn.cute`` with
|
||||
``block_sparsity``) is an *optional* dependency: it is imported lazily and
|
||||
only exercised when the VSA-128/256 CuTe fastpath is explicitly selected
|
||||
(``FASTVIDEO_VSA_CUTEDSL=1``). The default path is Triton and does not require
|
||||
it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``.
|
||||
only exercised when the VSA-256 CuTe fastpath is explicitly selected
|
||||
(``FASTVIDEO_VSA_CUTEDSL=1``). The default VSA-256 path is Triton and does
|
||||
not require it. Also needs ``nvidia-cutlass-dsl`` and ``quack-kernels``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -22,14 +22,13 @@ from typing import Tuple
|
||||
|
||||
import torch
|
||||
|
||||
_FA4_IMPORT_HINT = ("VSA-128/256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
|
||||
_FA4_IMPORT_HINT = ("VSA-256 CuTe fastpath requires a FlashAttention-4 CuTe build that "
|
||||
"provides `flash_attn.cute` with block-sparsity support (plus "
|
||||
"`nvidia-cutlass-dsl` and `quack-kernels`). This is an optional "
|
||||
"dependency; the default path is Triton. Install the FA4 CuTe "
|
||||
"dependency; the default VSA-256 path is Triton. Install the FA4 CuTe "
|
||||
"build and set FASTVIDEO_VSA_CUTEDSL=1 to enable the CuTe fastpath.")
|
||||
|
||||
|
||||
@functools.lru_cache(maxsize=1)
|
||||
def _load_fa4_cute():
|
||||
"""Lazily import the optional FA4 CuTe block-sparse symbols.
|
||||
|
||||
@@ -39,39 +38,14 @@ def _load_fa4_cute():
|
||||
"""
|
||||
try:
|
||||
from flash_attn.cute.block_sparsity import BlockSparseTensorsTorch
|
||||
from flash_attn.cute.interface import (
|
||||
_flash_attn_bwd,
|
||||
_flash_attn_fwd,
|
||||
flash_attn_func,
|
||||
)
|
||||
from flash_attn.cute.interface import _flash_attn_fwd
|
||||
except ImportError as exc: # pragma: no cover - optional dependency
|
||||
raise ImportError(_FA4_IMPORT_HINT) from exc
|
||||
return BlockSparseTensorsTorch, flash_attn_func, _flash_attn_fwd, _flash_attn_bwd
|
||||
return BlockSparseTensorsTorch, _flash_attn_fwd
|
||||
|
||||
|
||||
# FA4's physical Q tile size; KV block size comes from the VSA caller.
|
||||
_FA4_Q_BLOCK_SIZE = 128
|
||||
|
||||
|
||||
class _SingleQStageLength(int):
|
||||
"""Keep the real length while selecting FA4's one-stage Q128 path.
|
||||
|
||||
On sm_100 FA4 derives ``q_stage`` from ``max_seqlen_q > tile_m``. Its
|
||||
kernel supports one 128-token Q stage, but the fixed-length public wrapper
|
||||
does not expose that choice. VSA-128 must select it explicitly; otherwise
|
||||
adjacent logical Q blocks are merged into a 256-token sparse block.
|
||||
"""
|
||||
|
||||
def __mul__(self, other):
|
||||
return type(self)(int(self) * int(other))
|
||||
|
||||
def __rmul__(self, other):
|
||||
return type(self)(int(other) * int(self))
|
||||
|
||||
def __gt__(self, other):
|
||||
if int(other) == _FA4_Q_BLOCK_SIZE:
|
||||
return False
|
||||
return int(self) > int(other)
|
||||
# Q-side tile size; kv_block_size comes from the caller's VSA logical KV block.
|
||||
_M_BLOCK_SIZE_DEFAULT = 128
|
||||
|
||||
|
||||
def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
@@ -90,12 +64,12 @@ def _map_to_index(block_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
return triton_map_to_index(block_map)
|
||||
|
||||
|
||||
def _choose_q_sparse_block_size(q_len: int, q_tile_size: int = _FA4_Q_BLOCK_SIZE) -> int:
|
||||
# FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > q_tile_size.
|
||||
def _choose_q_sparse_block_size(q_len: int, m_block_size: int = _M_BLOCK_SIZE_DEFAULT) -> int:
|
||||
# FA4 supports a doubled Q sparsity granularity on sm_100+ when q_len > m_block_size.
|
||||
major, _ = torch.cuda.get_device_capability()
|
||||
if major >= 10 and q_len > q_tile_size:
|
||||
return 2 * q_tile_size
|
||||
return q_tile_size
|
||||
if major >= 10 and q_len > m_block_size:
|
||||
return 2 * m_block_size
|
||||
return m_block_size
|
||||
|
||||
|
||||
def _aggregate_q_block_map(
|
||||
@@ -160,35 +134,23 @@ def _build_vbs_mask_mod(kv_block_size: int):
|
||||
return _vbs_mask_mod
|
||||
|
||||
|
||||
def _build_sparse_tensors(
|
||||
def _cute_forward(
|
||||
q_bshd: torch.Tensor,
|
||||
k_bshd: torch.Tensor,
|
||||
v_bshd: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
*,
|
||||
q_len: int,
|
||||
q_block_size: int,
|
||||
kv_block_size: int,
|
||||
need_backward: bool,
|
||||
force_q_sparse_block_size: int | None = None,
|
||||
) -> Tuple[object, object | None]:
|
||||
"""Build the Q-owned forward and KV-owned backward sparse metadata.
|
||||
|
||||
``need_backward`` is False on inference-only calls: the backward metadata
|
||||
is a pair of dense ``[B, H, kv_blocks, q_blocks]`` int32 index tensors that
|
||||
FA4 keeps alive on its autograd ctx until backward runs, so building it
|
||||
when nothing requires grad is pure overhead (~80 MiB per call at Wan-14B
|
||||
720p shape).
|
||||
"""
|
||||
BlockSparseTensorsTorch, _, _, _ = _load_fa4_cute()
|
||||
if force_q_sparse_block_size is None:
|
||||
q_sparse_candidate = _choose_q_sparse_block_size(q_len)
|
||||
q_sparse_block_size = max(
|
||||
q_block_size,
|
||||
((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size,
|
||||
)
|
||||
else:
|
||||
q_sparse_block_size = force_q_sparse_block_size
|
||||
if q_sparse_block_size < q_block_size or q_sparse_block_size % q_block_size != 0:
|
||||
raise ValueError("force_q_sparse_block_size must be a positive multiple of q_block_size")
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Internal: FA4 CuTe BSA fwd with BSHD inputs."""
|
||||
BlockSparseTensorsTorch, _flash_attn_fwd = _load_fa4_cute()
|
||||
q_sparse_candidate = _choose_q_sparse_block_size(q_bshd.shape[1])
|
||||
q_sparse_block_size = max(
|
||||
q_block_size,
|
||||
((q_sparse_candidate + q_block_size - 1) // q_block_size) * q_block_size,
|
||||
)
|
||||
sparse_map = _aggregate_q_block_map(
|
||||
block_map,
|
||||
q_sparse_block_size=q_sparse_block_size,
|
||||
@@ -196,166 +158,35 @@ def _build_sparse_tensors(
|
||||
)
|
||||
kv_full = (variable_block_sizes == kv_block_size).view(1, 1, 1, -1)
|
||||
kv_partial = ((variable_block_sizes > 0) & (variable_block_sizes < kv_block_size)).view(1, 1, 1, -1)
|
||||
full_map = sparse_map & kv_full
|
||||
mask_map = sparse_map & kv_partial
|
||||
|
||||
def from_maps(full_map: torch.Tensor, mask_map: torch.Tensor) -> object:
|
||||
full_block_idx, full_block_cnt = _map_to_index(full_map.contiguous())
|
||||
mask_block_idx, mask_block_cnt = _map_to_index(mask_map.contiguous())
|
||||
return BlockSparseTensorsTorch(
|
||||
full_block_cnt=full_block_cnt.to(torch.int32).contiguous(),
|
||||
full_block_idx=full_block_idx.to(torch.int32).contiguous(),
|
||||
mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(),
|
||||
mask_block_idx=mask_block_idx.to(torch.int32).contiguous(),
|
||||
block_size=(q_sparse_block_size, kv_block_size),
|
||||
)
|
||||
full_block_idx, full_block_cnt = _map_to_index(full_map)
|
||||
mask_block_idx, mask_block_cnt = _map_to_index(mask_map)
|
||||
|
||||
forward_sparse_tensors = from_maps(
|
||||
sparse_map & kv_full,
|
||||
sparse_map & kv_partial,
|
||||
sparse_tensors = BlockSparseTensorsTorch(
|
||||
full_block_cnt=full_block_cnt.to(torch.int32).contiguous(),
|
||||
full_block_idx=full_block_idx.to(torch.int32).contiguous(),
|
||||
mask_block_cnt=mask_block_cnt.to(torch.int32).contiguous(),
|
||||
mask_block_idx=mask_block_idx.to(torch.int32).contiguous(),
|
||||
block_size=(q_sparse_block_size, kv_block_size),
|
||||
)
|
||||
|
||||
if not need_backward:
|
||||
return forward_sparse_tensors, None
|
||||
|
||||
# FA4 backward is KV-owned: for each physical KV tile, list the sparse
|
||||
# query tiles that selected it. Full and partial KV tiles stay separate
|
||||
# so the token-level validity mask only runs for padded tiles.
|
||||
backward_sparse_tensors = from_maps(
|
||||
(sparse_map & kv_full).transpose(2, 3),
|
||||
(sparse_map & kv_partial).transpose(2, 3),
|
||||
)
|
||||
return forward_sparse_tensors, backward_sparse_tensors
|
||||
|
||||
|
||||
def _cute_attention_q128_forward(
|
||||
q_bshd: torch.Tensor,
|
||||
k_bshd: torch.Tensor,
|
||||
v_bshd: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
*,
|
||||
need_backward: bool,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, object | None]:
|
||||
"""Run FA4 with one physical Q stage per logical VSA-128 block."""
|
||||
_, _, flash_attn_fwd, _ = _load_fa4_cute()
|
||||
forward_sparse_tensors, backward_sparse_tensors = _build_sparse_tensors(
|
||||
block_map,
|
||||
variable_block_sizes,
|
||||
q_len=q_bshd.shape[1],
|
||||
q_block_size=_FA4_Q_BLOCK_SIZE,
|
||||
kv_block_size=_FA4_Q_BLOCK_SIZE,
|
||||
need_backward=need_backward,
|
||||
force_q_sparse_block_size=_FA4_Q_BLOCK_SIZE,
|
||||
)
|
||||
out, lse = flash_attn_fwd(
|
||||
# _flash_attn_fwd returns (out, lse, p, row_max); keep the first two.
|
||||
out, lse = _flash_attn_fwd(
|
||||
q_bshd,
|
||||
k_bshd,
|
||||
v_bshd,
|
||||
tile_mn=(_FA4_Q_BLOCK_SIZE, _FA4_Q_BLOCK_SIZE),
|
||||
max_seqlen_q=_SingleQStageLength(q_bshd.shape[1]),
|
||||
mask_mod=_build_vbs_mask_mod(_FA4_Q_BLOCK_SIZE),
|
||||
block_sparse_tensors=forward_sparse_tensors,
|
||||
tile_mn=(_M_BLOCK_SIZE_DEFAULT, kv_block_size),
|
||||
mask_mod=_build_vbs_mask_mod(kv_block_size),
|
||||
block_sparse_tensors=sparse_tensors,
|
||||
aux_tensors=[variable_block_sizes],
|
||||
causal=False,
|
||||
return_lse=True,
|
||||
)[:2]
|
||||
return out, lse, backward_sparse_tensors
|
||||
|
||||
|
||||
class _CuteAttentionQ128(torch.autograd.Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes):
|
||||
out, lse, backward_sparse_tensors = _cute_attention_q128_forward(
|
||||
q_bshd,
|
||||
k_bshd,
|
||||
v_bshd,
|
||||
block_map,
|
||||
variable_block_sizes,
|
||||
need_backward=True,
|
||||
)
|
||||
ctx.save_for_backward(q_bshd, k_bshd, v_bshd, out, lse, variable_block_sizes)
|
||||
ctx.backward_sparse_tensors = backward_sparse_tensors
|
||||
ctx.mark_non_differentiable(lse)
|
||||
ctx.set_materialize_grads(False)
|
||||
return out, lse
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_out, grad_lse):
|
||||
del grad_lse
|
||||
q_bshd, k_bshd, v_bshd, out, lse, variable_block_sizes = ctx.saved_tensors
|
||||
if grad_out is None:
|
||||
grad_out = torch.zeros_like(out)
|
||||
_, _, _, flash_attn_bwd = _load_fa4_cute()
|
||||
dq, dk, dv = flash_attn_bwd(
|
||||
q_bshd,
|
||||
k_bshd,
|
||||
v_bshd,
|
||||
out,
|
||||
grad_out.contiguous(),
|
||||
lse,
|
||||
softmax_scale=q_bshd.shape[-1]**-0.5,
|
||||
mask_mod=_build_vbs_mask_mod(_FA4_Q_BLOCK_SIZE),
|
||||
aux_tensors=[variable_block_sizes],
|
||||
block_sparse_tensors=ctx.backward_sparse_tensors,
|
||||
)
|
||||
return dq, dk, dv, None, None
|
||||
|
||||
|
||||
def _cute_attention_q128(
|
||||
q_bshd: torch.Tensor,
|
||||
k_bshd: torch.Tensor,
|
||||
v_bshd: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q_bshd, k_bshd, v_bshd))
|
||||
if need_backward:
|
||||
return _CuteAttentionQ128.apply(q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes)
|
||||
out, lse, _ = _cute_attention_q128_forward(
|
||||
q_bshd,
|
||||
k_bshd,
|
||||
v_bshd,
|
||||
block_map,
|
||||
variable_block_sizes,
|
||||
need_backward=False,
|
||||
)
|
||||
return out, lse
|
||||
|
||||
|
||||
def _cute_attention(
|
||||
q_bshd: torch.Tensor,
|
||||
k_bshd: torch.Tensor,
|
||||
v_bshd: torch.Tensor,
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Run FA4's autograd-enabled block-sparse attention with BSHD inputs."""
|
||||
_, flash_attn_func, _, _ = _load_fa4_cute()
|
||||
q_block_size = q_bshd.shape[1] // block_map.shape[2]
|
||||
kv_block_size = k_bshd.shape[1] // block_map.shape[3]
|
||||
if q_block_size == kv_block_size == _FA4_Q_BLOCK_SIZE:
|
||||
return _cute_attention_q128(q_bshd, k_bshd, v_bshd, block_map, variable_block_sizes)
|
||||
need_backward = torch.is_grad_enabled() and any(t.requires_grad for t in (q_bshd, k_bshd, v_bshd))
|
||||
forward_sparse_tensors, backward_sparse_tensors = _build_sparse_tensors(
|
||||
block_map,
|
||||
variable_block_sizes,
|
||||
q_len=q_bshd.shape[1],
|
||||
q_block_size=q_block_size,
|
||||
kv_block_size=kv_block_size,
|
||||
need_backward=need_backward,
|
||||
)
|
||||
return flash_attn_func(
|
||||
q_bshd,
|
||||
k_bshd,
|
||||
v_bshd,
|
||||
mask_mod=_build_vbs_mask_mod(kv_block_size),
|
||||
aux_tensors=[variable_block_sizes],
|
||||
block_sparse_tensors=forward_sparse_tensors,
|
||||
block_sparse_tensors_bwd=backward_sparse_tensors,
|
||||
return_lse=True,
|
||||
)
|
||||
|
||||
|
||||
def block_sparse_attn_cute_fwd(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
@@ -363,25 +194,34 @@ def block_sparse_attn_cute_fwd(
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Autograd-enabled CuTe block-sparse attention for [B, H, S, D]."""
|
||||
"""CuTe forward-only block-sparse attention with [B, H, S, D] inputs."""
|
||||
if block_map.dim() == 3:
|
||||
block_map = block_map.unsqueeze(0)
|
||||
q_block_size = q.shape[2] // block_map.shape[2]
|
||||
kv_block_size = k.shape[2] // block_map.shape[3]
|
||||
|
||||
q_bshd = q.transpose(1, 2).contiguous()
|
||||
k_bshd = k.transpose(1, 2).contiguous()
|
||||
v_bshd = v.transpose(1, 2).contiguous()
|
||||
out_bshd, lse = _cute_attention(
|
||||
out_bshd, lse_bshd = _cute_forward(
|
||||
q_bshd,
|
||||
k_bshd,
|
||||
v_bshd,
|
||||
block_map,
|
||||
variable_block_sizes,
|
||||
q_block_size=q_block_size,
|
||||
kv_block_size=kv_block_size,
|
||||
)
|
||||
out = out_bshd.transpose(1, 2).contiguous()
|
||||
# FA4 already returns lse as [B, H, S], matching the Triton path's aux
|
||||
# contract, so it needs no transpose. Detach before any further op: the
|
||||
# value is informational and callers never backprop through it.
|
||||
return out, lse.detach()
|
||||
if lse_bshd is None:
|
||||
lse = torch.empty(
|
||||
(q.shape[0], q.shape[1], q.shape[2]),
|
||||
dtype=torch.float32,
|
||||
device=q.device,
|
||||
)
|
||||
else:
|
||||
lse = lse_bshd.transpose(1, 2).contiguous()
|
||||
return out, lse
|
||||
|
||||
|
||||
def block_sparse_attn_cute_fwd_bshd(
|
||||
@@ -391,16 +231,27 @@ def block_sparse_attn_cute_fwd_bshd(
|
||||
block_map: torch.Tensor,
|
||||
variable_block_sizes: torch.Tensor,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Autograd-enabled CuTe block-sparse attention for [B, S, H, D]."""
|
||||
"""CuTe forward-only block-sparse attention with [B, S, H, D] inputs."""
|
||||
if block_map.dim() == 3:
|
||||
block_map = block_map.unsqueeze(0)
|
||||
q_block_size = q.shape[1] // block_map.shape[2]
|
||||
kv_block_size = k.shape[1] // block_map.shape[3]
|
||||
|
||||
out, lse = _cute_attention(
|
||||
out, lse_bshd = _cute_forward(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
block_map,
|
||||
variable_block_sizes,
|
||||
q_block_size=q_block_size,
|
||||
kv_block_size=kv_block_size,
|
||||
)
|
||||
# lse is [B, H, S] regardless of the q/k/v layout; see above.
|
||||
return out, lse.detach()
|
||||
if lse_bshd is None:
|
||||
lse = torch.empty(
|
||||
(q.shape[0], q.shape[2], q.shape[1]),
|
||||
dtype=torch.float32,
|
||||
device=q.device,
|
||||
)
|
||||
else:
|
||||
lse = lse_bshd.transpose(1, 2).contiguous()
|
||||
return out, lse
|
||||
|
||||
@@ -2,8 +2,6 @@ import math
|
||||
import torch
|
||||
from .block_sparse_attn import block_sparse_attn
|
||||
from .block_sparse_attn_256 import (
|
||||
block_sparse_attn_128,
|
||||
block_sparse_attn_128_bshd,
|
||||
block_sparse_attn_256,
|
||||
block_sparse_attn_256_bshd,
|
||||
)
|
||||
@@ -76,13 +74,12 @@ def video_sparse_attn(
|
||||
|
||||
Dispatches the sparse branch by ``block_elements = prod(block_size)``:
|
||||
- 64 -> existing TK/Triton path (see ``block_sparse_attn_from_indices``).
|
||||
- 128 -> Triton fallback or CuTe FA4 block-sparse attention.
|
||||
- 256 -> CuTe FA4 block-sparse attention (see ``block_sparse_attn_256``).
|
||||
|
||||
Backend overrides:
|
||||
- ``FASTVIDEO_VSA_TRITON=1`` forces Triton in either path.
|
||||
- ``FASTVIDEO_VSA_TK=1`` prefers the sm_90 TK kernel in the 64-block path.
|
||||
- ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 128/256-block paths.
|
||||
- ``FASTVIDEO_VSA_CUTEDSL=1`` prefers CuTe in the 256-block path.
|
||||
"""
|
||||
if isinstance(block_size, int):
|
||||
block_size = (block_size, block_size, block_size)
|
||||
@@ -122,9 +119,8 @@ def video_sparse_attn(
|
||||
# Sparse branch (fused Triton topk mask)
|
||||
mask = fused_topk_mask(scores, topk)
|
||||
|
||||
if block_elements in (128, 256):
|
||||
attention = block_sparse_attn_128 if block_elements == 128 else block_sparse_attn_256
|
||||
out_s = attention(q, k, v, mask, variable_block_sizes)[0]
|
||||
if block_elements == 256:
|
||||
out_s = block_sparse_attn_256(q, k, v, mask, variable_block_sizes)[0]
|
||||
else:
|
||||
out_s = block_sparse_attn(q, k, v, mask, variable_block_sizes)[0]
|
||||
|
||||
@@ -146,14 +142,14 @@ def video_sparse_attn_bshd(
|
||||
"""VSA entrypoint for [B, S, H, D] tensors.
|
||||
|
||||
Avoids the BHSD<->BSHD round-trip that ``video_sparse_attn`` performs on
|
||||
the CuTe 128/256-block paths; the 64-block path still expects BHSD and is not
|
||||
the CuTe 256-block path; the 64-block path still expects BHSD and is not
|
||||
supported here.
|
||||
"""
|
||||
if isinstance(block_size, int):
|
||||
block_size = (block_size, block_size, block_size)
|
||||
block_elements = block_size[0] * block_size[1] * block_size[2]
|
||||
if block_elements not in (128, 256):
|
||||
raise ValueError("video_sparse_attn_bshd is only defined for block_elements=128 or 256 "
|
||||
if block_elements != 256:
|
||||
raise ValueError("video_sparse_attn_bshd is only defined for block_elements=256 "
|
||||
f"(got {block_elements}); use video_sparse_attn for the 64-block path.")
|
||||
|
||||
batch, q_seq_len, heads, dim = q.shape
|
||||
@@ -175,15 +171,19 @@ def video_sparse_attn_bshd(
|
||||
raise ValueError(f"q_variable_block_sizes must have length q_num_blocks={q_num_blocks}, "
|
||||
f"got {q_variable_block_sizes.numel()}")
|
||||
|
||||
# Compression branch (BSHD-native: match fused_block_mean's semantics).
|
||||
# Padding values are expected to be zero; gradients are broadcast across
|
||||
# the full padded block, just like the BHSD fused common path.
|
||||
# Compression branch (BSHD-native: mean over the 256-token axis).
|
||||
token_idx = torch.arange(block_elements, device=q.device, dtype=torch.int32)
|
||||
q_token_valid = (token_idx.view(1, -1) < q_variable_block_sizes.view(-1,
|
||||
1)).view(1, q_num_blocks, block_elements, 1, 1)
|
||||
kv_token_valid = (token_idx.view(1, -1) < variable_block_sizes.view(-1,
|
||||
1)).view(1, kv_num_blocks, block_elements, 1, 1)
|
||||
|
||||
q_c = q.view(batch, q_num_blocks, block_elements, heads, dim)
|
||||
k_c = k.view(batch, kv_num_blocks, block_elements, heads, dim)
|
||||
v_c = v.view(batch, kv_num_blocks, block_elements, heads, dim)
|
||||
q_c = (q_c.float().sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
|
||||
k_c = (k_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
|
||||
v_c = (v_c.float().sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
|
||||
q_c = ((q_c.float() * q_token_valid).sum(dim=2) / q_variable_block_sizes.view(1, -1, 1, 1)).to(q.dtype)
|
||||
k_c = ((k_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(k.dtype)
|
||||
v_c = ((v_c.float() * kv_token_valid).sum(dim=2) / variable_block_sizes.view(1, -1, 1, 1)).to(v.dtype)
|
||||
q_ch = q_c.permute(0, 2, 1, 3).contiguous()
|
||||
k_ch = k_c.permute(0, 2, 1, 3).contiguous()
|
||||
v_ch = v_c.permute(0, 2, 1, 3).contiguous()
|
||||
@@ -195,15 +195,13 @@ def video_sparse_attn_bshd(
|
||||
|
||||
# Sparse branch (fused Triton topk mask + CuTe BSHD).
|
||||
mask = fused_topk_mask(scores, topk)
|
||||
attention = block_sparse_attn_128_bshd if block_elements == 128 else block_sparse_attn_256_bshd
|
||||
out_s, _ = attention(q, k, v, mask, variable_block_sizes)
|
||||
out_s, _ = block_sparse_attn_256_bshd(q, k, v, mask, variable_block_sizes)
|
||||
|
||||
# Out-of-place: ``out_s`` is the tensor FA4's autograd node saved for its
|
||||
# backward, so mutating it in place invalidates the graph.
|
||||
out_view = out_s.view(batch, q_num_blocks, block_elements, heads, dim)
|
||||
out = out_s
|
||||
out_view = out.view(batch, q_num_blocks, block_elements, heads, dim)
|
||||
if compress_attn_weight is not None:
|
||||
gate_view = compress_attn_weight.view(batch, q_num_blocks, block_elements, heads, dim)
|
||||
out = out_view + out_c_blk.unsqueeze(2) * gate_view
|
||||
out_view.add_(out_c_blk.unsqueeze(2) * gate_view)
|
||||
else:
|
||||
out = out_view + out_c_blk.unsqueeze(2)
|
||||
return out.view(batch, q_seq_len, heads, dim)
|
||||
out_view.add_(out_c_blk.unsqueeze(2))
|
||||
return out
|
||||
|
||||
@@ -1,146 +0,0 @@
|
||||
"""VSA-128 CuTe/Triton forward and backward parity on Blackwell."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo_kernel import video_sparse_attn, video_sparse_attn_bshd
|
||||
from fastvideo_kernel.block_sparse_attn_256 import block_sparse_attn_128
|
||||
|
||||
from .test_vsa256_triton import _metrics, _torch_vsa256_reference
|
||||
|
||||
_BLOCK = 128
|
||||
_BLOCK_SIZE_3D = (2, 8, 8)
|
||||
|
||||
|
||||
def _select_backend(monkeypatch, backend: str) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA is required")
|
||||
if backend == "cute":
|
||||
pytest.importorskip(
|
||||
"flash_attn.cute.block_sparsity",
|
||||
reason="optional FA4 CuTe build (flash_attn.cute) not installed",
|
||||
)
|
||||
monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
|
||||
monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False)
|
||||
monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
|
||||
monkeypatch.delenv("FASTVIDEO_VSA_CUTEDSL", raising=False)
|
||||
|
||||
|
||||
def _dense_sparse_reference(q, k, v, block_map, variable_block_sizes):
|
||||
token_mask = block_map.repeat_interleave(_BLOCK, dim=2).repeat_interleave(_BLOCK, dim=3)
|
||||
kv_valid = torch.arange(_BLOCK, device=k.device) < variable_block_sizes[:, None]
|
||||
token_mask = token_mask & kv_valid.reshape(1, 1, 1, -1)
|
||||
logits = torch.matmul(q.float(), k.float().transpose(-2, -1)) / math.sqrt(q.shape[-1])
|
||||
probabilities = torch.softmax(logits.masked_fill(~token_mask, float("-inf")), dim=-1)
|
||||
return torch.matmul(probabilities, v.float()).to(q.dtype)
|
||||
|
||||
|
||||
def _check(tag: str, expected: torch.Tensor, actual: torch.Tensor, avg_tol: float, rel_tol: float) -> None:
|
||||
assert torch.isfinite(actual).all().item(), f"{tag}: non-finite values"
|
||||
avg_abs, max_rel = _metrics(expected, actual)
|
||||
print(f" {tag}: avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
|
||||
assert avg_abs < avg_tol
|
||||
assert max_rel < rel_tol
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.parametrize("backend", ["cute", "triton"])
|
||||
def test_vsa128_explicit_routes_forward_backward(backend: str, monkeypatch) -> None:
|
||||
"""Adjacent Q128 blocks must keep independent routes instead of merging."""
|
||||
_select_backend(monkeypatch, backend)
|
||||
torch.manual_seed(53)
|
||||
shape = (1, 1, 3 * _BLOCK, 128)
|
||||
base = [torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)]
|
||||
grad_output = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
|
||||
variable_block_sizes = torch.tensor([128, 91, 37], device="cuda", dtype=torch.int32)
|
||||
block_map = torch.eye(3, device="cuda", dtype=torch.bool).view(1, 1, 3, 3)
|
||||
|
||||
actual_inputs = [tensor.detach().clone().requires_grad_() for tensor in base]
|
||||
actual, _ = block_sparse_attn_128(*actual_inputs, block_map, variable_block_sizes)
|
||||
(actual * grad_output).sum().backward()
|
||||
|
||||
reference_inputs = [tensor.detach().clone().requires_grad_() for tensor in base]
|
||||
expected = _dense_sparse_reference(*reference_inputs, block_map, variable_block_sizes)
|
||||
(expected * grad_output).sum().backward()
|
||||
|
||||
print(f"[vsa128-explicit-{backend}]")
|
||||
_check("out", expected, actual, 1e-3, 0.2)
|
||||
for name, reference, candidate in zip(("dq", "dk", "dv"), reference_inputs, actual_inputs, strict=True):
|
||||
_check(name, reference.grad, candidate.grad, 2e-2, 0.5)
|
||||
|
||||
|
||||
def _zero_kv_tail(x: torch.Tensor, variable_block_sizes: torch.Tensor) -> torch.Tensor:
|
||||
valid = torch.arange(_BLOCK, device=x.device) < variable_block_sizes[:, None]
|
||||
valid = valid.view(1, 1, -1, _BLOCK, 1).expand_as(x.view(1, x.shape[1], -1, _BLOCK, x.shape[-1]))
|
||||
return x * valid.reshape_as(x).to(x.dtype)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.parametrize("backend", ["cute", "triton"])
|
||||
@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
|
||||
def test_vsa128_wrapper_forward_backward(backend: str, layout: str, monkeypatch) -> None:
|
||||
_select_backend(monkeypatch, backend)
|
||||
torch.manual_seed(59)
|
||||
batch, heads, dim = 1, 2, 128
|
||||
q_blocks, kv_blocks, topk = 3, 4, 2
|
||||
q_shape = (batch, heads, q_blocks * _BLOCK, dim)
|
||||
kv_shape = (batch, heads, kv_blocks * _BLOCK, dim)
|
||||
q_base = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16)
|
||||
kv_sizes = torch.tensor([128, 91, 37, 128], device="cuda", dtype=torch.int32)
|
||||
q_sizes = torch.full((q_blocks, ), _BLOCK, device="cuda", dtype=torch.int32)
|
||||
k_base = _zero_kv_tail(torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16), kv_sizes)
|
||||
v_base = _zero_kv_tail(torch.randn(kv_shape, device="cuda", dtype=torch.bfloat16), kv_sizes)
|
||||
gate_base = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
grad_output = torch.randn(q_shape, device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
if layout == "bhsd":
|
||||
actual_inputs = [tensor.detach().clone().requires_grad_() for tensor in (q_base, k_base, v_base)]
|
||||
actual_gate = gate_base.detach().clone().requires_grad_()
|
||||
actual = video_sparse_attn(
|
||||
*actual_inputs,
|
||||
kv_sizes,
|
||||
q_sizes,
|
||||
topk,
|
||||
block_size=_BLOCK_SIZE_3D,
|
||||
compress_attn_weight=actual_gate,
|
||||
)
|
||||
actual_grads = actual_inputs
|
||||
else:
|
||||
bshd_inputs = [tensor.transpose(1, 2).contiguous().detach().requires_grad_()
|
||||
for tensor in (q_base, k_base, v_base)]
|
||||
bshd_gate = gate_base.transpose(1, 2).contiguous().detach().requires_grad_()
|
||||
actual = video_sparse_attn_bshd(
|
||||
*bshd_inputs,
|
||||
kv_sizes,
|
||||
q_sizes,
|
||||
topk,
|
||||
block_size=_BLOCK_SIZE_3D,
|
||||
compress_attn_weight=bshd_gate,
|
||||
).transpose(1, 2)
|
||||
actual_grads = bshd_inputs
|
||||
(actual * grad_output).sum().backward()
|
||||
|
||||
reference_inputs = [tensor.detach().clone().requires_grad_() for tensor in (q_base, k_base, v_base)]
|
||||
reference_gate = gate_base.detach().clone().requires_grad_()
|
||||
expected = _torch_vsa256_reference(
|
||||
*reference_inputs,
|
||||
q_sizes,
|
||||
kv_sizes,
|
||||
topk,
|
||||
compress_attn_weight=reference_gate,
|
||||
)
|
||||
(expected * grad_output).sum().backward()
|
||||
|
||||
print(f"[vsa128-wrapper-{backend}-{layout}]")
|
||||
_check("out", expected, actual, 1e-3, 0.2)
|
||||
for name, reference, candidate in zip(("dq", "dk", "dv"), reference_inputs, actual_grads, strict=True):
|
||||
candidate_grad = candidate.grad if layout == "bhsd" else candidate.grad.transpose(1, 2)
|
||||
_check(name, reference.grad, candidate_grad, 2e-2, 0.5)
|
||||
actual_gate_grad = actual_gate.grad if layout == "bhsd" else bshd_gate.grad.transpose(1, 2)
|
||||
_check("dgate", reference_gate.grad, actual_gate_grad, 1e-3, 0.2)
|
||||
@@ -1,224 +0,0 @@
|
||||
"""VSA-256 FA4 CuTe forward/backward parity for BHSD and BSHD APIs.
|
||||
|
||||
Covers the shapes the CuTe backward actually sees in production: the gated
|
||||
compression branch (`compress_attn_weight`), partially filled Q tiles,
|
||||
and q_len != kv_len. Also pins the inference fast path, which must skip the
|
||||
KV-owned backward metadata without changing the forward result.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo_kernel import video_sparse_attn, video_sparse_attn_bshd
|
||||
|
||||
from .test_vsa256_triton import _metrics, _torch_vsa256_reference
|
||||
|
||||
_BLOCK = 256
|
||||
_BLOCK_SIZE_3D = (4, 8, 8) # prod == 256
|
||||
|
||||
# Measured on GB200 (sm_100) with bf16 inputs: grads land around 1e-4 avg_abs
|
||||
# and <=0.11 max_rel across every case below, so these leave ~10x headroom
|
||||
# without being loose enough to hide a real regression.
|
||||
_OUT_TOL = (1e-3, 0.2)
|
||||
_GRAD_TOL = (1e-3, 0.25)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _require_cute_backend(monkeypatch):
|
||||
pytest.importorskip(
|
||||
"flash_attn.cute.block_sparsity",
|
||||
reason="optional FA4 CuTe build (flash_attn.cute) not installed",
|
||||
)
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA is required")
|
||||
monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
|
||||
monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False)
|
||||
monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False)
|
||||
|
||||
|
||||
def _zero_pad_tail(x: torch.Tensor, var: torch.Tensor) -> torch.Tensor:
|
||||
"""Zero the padded tail of every 256-token tile of a [B, H, S, D] tensor.
|
||||
|
||||
VSA callers scatter into a zeroed tile buffer, so padded slots are zero;
|
||||
both the kernel and the reference rely on that.
|
||||
"""
|
||||
bsz, heads, _, dim = x.shape
|
||||
blocks = var.numel()
|
||||
token_idx = torch.arange(_BLOCK, device=x.device, dtype=torch.int32)
|
||||
valid = (token_idx.view(1, -1) < var.view(-1, 1)).view(1, 1, blocks, _BLOCK, 1)
|
||||
valid = valid.expand(bsz, heads, blocks, _BLOCK, dim).reshape_as(x)
|
||||
return x * valid.to(x.dtype)
|
||||
|
||||
|
||||
def _make_inputs(
|
||||
q_blocks: int,
|
||||
kv_blocks: int,
|
||||
kv_var: torch.Tensor,
|
||||
q_var: torch.Tensor,
|
||||
heads: int = 2,
|
||||
dim: int = 128,
|
||||
seed: int = 42,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
torch.manual_seed(seed)
|
||||
device = torch.device("cuda")
|
||||
dtype = torch.bfloat16
|
||||
sq, skv = q_blocks * _BLOCK, kv_blocks * _BLOCK
|
||||
q = torch.randn(1, heads, sq, dim, device=device, dtype=dtype)
|
||||
k = torch.randn(1, heads, skv, dim, device=device, dtype=dtype)
|
||||
v = torch.randn(1, heads, skv, dim, device=device, dtype=dtype)
|
||||
grad_out = torch.randn_like(q)
|
||||
return _zero_pad_tail(q, q_var), _zero_pad_tail(k, kv_var), _zero_pad_tail(v, kv_var), grad_out
|
||||
|
||||
|
||||
def _check(tag: str, ref: torch.Tensor, got: torch.Tensor, tol: Tuple[float, float]) -> None:
|
||||
assert torch.isfinite(got).all().item(), f"{tag}: non-finite values"
|
||||
avg_abs, max_rel = _metrics(ref, got)
|
||||
print(f" {tag}: avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
|
||||
assert avg_abs < tol[0], f"{tag}: avg_abs {avg_abs:.3e} >= {tol[0]:.3e}"
|
||||
assert max_rel < tol[1], f"{tag}: max_rel {max_rel:.3e} >= {tol[1]:.3e}"
|
||||
|
||||
|
||||
def _run_bhsd(q, k, v, kv_var, q_var, topk, gate=None):
|
||||
qg, kg, vg = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
|
||||
out = video_sparse_attn(qg, kg, vg, kv_var, q_var, topk, block_size=_BLOCK_SIZE_3D, compress_attn_weight=gate)
|
||||
return out, (qg, kg, vg)
|
||||
|
||||
|
||||
def _run_bshd(q, k, v, kv_var, q_var, topk, gate=None):
|
||||
qg, kg, vg = (t.transpose(1, 2).contiguous().requires_grad_(True) for t in (q, k, v))
|
||||
gate_bshd = None if gate is None else gate.transpose(1, 2).contiguous()
|
||||
out = video_sparse_attn_bshd(qg,
|
||||
kg,
|
||||
vg,
|
||||
kv_var,
|
||||
q_var,
|
||||
topk,
|
||||
block_size=_BLOCK_SIZE_3D,
|
||||
compress_attn_weight=gate_bshd)
|
||||
return out.transpose(1, 2), (qg, kg, vg)
|
||||
|
||||
|
||||
def _reference(q, k, v, q_var, kv_var, topk, gate=None):
|
||||
qr, kr, vr = (t.detach().clone().requires_grad_(True) for t in (q, k, v))
|
||||
out = _torch_vsa256_reference(qr, kr, vr, q_var, kv_var, topk, compress_attn_weight=gate)
|
||||
return out, (qr, kr, vr)
|
||||
|
||||
|
||||
def _compare(tag, layout, q, k, v, kv_var, q_var, topk, grad_out, gate=None):
|
||||
runner = _run_bhsd if layout == "bhsd" else _run_bshd
|
||||
out, (qg, kg, vg) = runner(q, k, v, kv_var, q_var, topk, gate=gate)
|
||||
(out * grad_out).sum().backward()
|
||||
grads = [g.grad if g.grad.dim() == 4 and layout == "bhsd" else g.grad for g in (qg, kg, vg)]
|
||||
if layout == "bshd":
|
||||
grads = [g.transpose(1, 2) for g in grads]
|
||||
|
||||
out_ref, refs = _reference(q, k, v, q_var, kv_var, topk, gate=gate)
|
||||
(out_ref * grad_out).sum().backward()
|
||||
|
||||
print(f"[{tag}-{layout}]")
|
||||
_check("out", out_ref, out, _OUT_TOL)
|
||||
for name, ref, got in zip(("dq", "dk", "dv"), refs, grads):
|
||||
_check(name, ref.grad, got, _GRAD_TOL)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
|
||||
def test_vsa256_cute_forward_backward_vs_torch_ref(layout: str) -> None:
|
||||
kv_var = torch.tensor([256, 173, 79, 256], dtype=torch.int32, device="cuda")
|
||||
q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
|
||||
q, k, v, grad_out = _make_inputs(3, 4, kv_var, q_var)
|
||||
_compare("vsa256-cute", layout, q, k, v, kv_var, q_var, 2, grad_out)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
|
||||
def test_vsa256_cute_backward_with_compress_gate(layout: str) -> None:
|
||||
"""The gated compression branch is what Wan and MiniMax-H3 actually run.
|
||||
|
||||
It is also the branch that composes the sparse output with the compression
|
||||
output, so it is the one that breaks if that composition mutates FA4's
|
||||
saved output in place.
|
||||
"""
|
||||
kv_var = torch.tensor([256, 200, 256, 91], dtype=torch.int32, device="cuda")
|
||||
q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
|
||||
q, k, v, grad_out = _make_inputs(3, 4, kv_var, q_var, seed=7)
|
||||
gate = torch.randn_like(q) * 0.1
|
||||
_compare("vsa256-cute-gated", layout, q, k, v, kv_var, q_var, 2, grad_out, gate=gate)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
|
||||
def test_vsa256_cute_backward_partial_q_blocks(layout: str) -> None:
|
||||
"""Q tiles that are not full: only the compression divisor depends on it,
|
||||
but it is the one axis the existing coverage held constant."""
|
||||
kv_var = torch.tensor([256, 128, 256], dtype=torch.int32, device="cuda")
|
||||
q_var = torch.tensor([256, 61, 199], dtype=torch.int32, device="cuda")
|
||||
q, k, v, grad_out = _make_inputs(3, 3, kv_var, q_var, seed=11)
|
||||
_compare("vsa256-cute-partial-q", layout, q, k, v, kv_var, q_var, 2, grad_out)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.parametrize("layout", ["bhsd", "bshd"])
|
||||
def test_vsa256_cute_backward_cross_q_kv(layout: str) -> None:
|
||||
"""q_len != kv_len: forward has coverage, backward did not."""
|
||||
kv_var = torch.tensor([256, 143, 256, 256, 88], dtype=torch.int32, device="cuda")
|
||||
q_var = torch.full((2, ), _BLOCK, dtype=torch.int32, device="cuda")
|
||||
q, k, v, grad_out = _make_inputs(2, 5, kv_var, q_var, seed=13)
|
||||
_compare("vsa256-cute-cross", layout, q, k, v, kv_var, q_var, 3, grad_out)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
def test_vsa256_cute_inference_matches_training_forward() -> None:
|
||||
"""The KV-owned backward metadata is only built when something requires
|
||||
grad. Skipping it must not perturb the forward result."""
|
||||
kv_var = torch.tensor([256, 173, 79, 256], dtype=torch.int32, device="cuda")
|
||||
q_var = torch.full((3, ), _BLOCK, dtype=torch.int32, device="cuda")
|
||||
q, k, v, _ = _make_inputs(3, 4, kv_var, q_var, seed=5)
|
||||
|
||||
with torch.no_grad():
|
||||
out_infer = video_sparse_attn_bshd(
|
||||
q.transpose(1, 2).contiguous(),
|
||||
k.transpose(1, 2).contiguous(),
|
||||
v.transpose(1, 2).contiguous(),
|
||||
kv_var,
|
||||
q_var,
|
||||
2,
|
||||
block_size=_BLOCK_SIZE_3D,
|
||||
compress_attn_weight=None,
|
||||
)
|
||||
|
||||
out_train, _ = _run_bshd(q, k, v, kv_var, q_var, 2)
|
||||
torch.testing.assert_close(out_infer, out_train.transpose(1, 2).detach(), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
def test_vsa256_cute_lse_is_bhs() -> None:
|
||||
"""The aux return is [B, H, S] on both entrypoints, matching the Triton
|
||||
path's contract."""
|
||||
from fastvideo_kernel.block_sparse_attn_256 import (block_sparse_attn_256, block_sparse_attn_256_bshd)
|
||||
|
||||
device = torch.device("cuda")
|
||||
heads, dim, q_blocks, kv_blocks = 2, 128, 3, 4
|
||||
sq, skv = q_blocks * _BLOCK, kv_blocks * _BLOCK
|
||||
q = torch.randn(1, heads, sq, dim, device=device, dtype=torch.bfloat16)
|
||||
k = torch.randn(1, heads, skv, dim, device=device, dtype=torch.bfloat16)
|
||||
v = torch.randn(1, heads, skv, dim, device=device, dtype=torch.bfloat16)
|
||||
vbs = torch.full((kv_blocks, ), _BLOCK, dtype=torch.int32, device=device)
|
||||
mask = torch.zeros(1, heads, q_blocks, kv_blocks, dtype=torch.bool, device=device)
|
||||
mask[..., :2] = True
|
||||
|
||||
_, lse_bhsd = block_sparse_attn_256(q, k, v, mask, vbs)
|
||||
assert lse_bhsd.shape == (1, heads, sq), lse_bhsd.shape
|
||||
|
||||
_, lse_bshd = block_sparse_attn_256_bshd(
|
||||
q.transpose(1, 2).contiguous(),
|
||||
k.transpose(1, 2).contiguous(),
|
||||
v.transpose(1, 2).contiguous(),
|
||||
mask,
|
||||
vbs,
|
||||
)
|
||||
assert lse_bshd.shape == (1, heads, sq), lse_bshd.shape
|
||||
@@ -22,7 +22,6 @@ def _torch_vsa256_reference(
|
||||
q_var: torch.Tensor,
|
||||
kv_var: torch.Tensor,
|
||||
topk_logical: int,
|
||||
compress_attn_weight: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
bsz, heads, _sq, dim = q.shape
|
||||
q_blocks = q_var.numel()
|
||||
@@ -56,8 +55,6 @@ def _torch_vsa256_reference(
|
||||
logits = logits.masked_fill(~token_mask, float("-inf"))
|
||||
prob = torch.softmax(logits, dim=-1)
|
||||
out_s = torch.matmul(prob, vf).to(q.dtype)
|
||||
if compress_attn_weight is not None:
|
||||
return out_c * compress_attn_weight + out_s
|
||||
return out_c + out_s
|
||||
|
||||
|
||||
|
||||
@@ -23,20 +23,6 @@ from fastvideo_kernel.block_sparse_attn import (
|
||||
from fastvideo_kernel.block_sparse_attn_varlen import block_sparse_attn_varlen
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _seed_rng():
|
||||
"""Pin the RNG so these cases do not depend on what ran before them.
|
||||
|
||||
Every tensor and every variable block size here comes from the global
|
||||
torch RNG, and the gradient checks use a tight max_rel threshold. Without
|
||||
a seed the inputs shift whenever an earlier test file draws a different
|
||||
number of randoms, which surfaces as an unrelated-looking failure in
|
||||
whichever case happens to land on unlucky data.
|
||||
"""
|
||||
torch.manual_seed(0)
|
||||
torch.cuda.manual_seed_all(0)
|
||||
|
||||
|
||||
def _reference_per_sequence(
|
||||
q_list,
|
||||
k_list,
|
||||
|
||||
@@ -319,13 +319,8 @@ class MiniMaxH3VSAImpl(AttentionImpl):
|
||||
v_pooled = _pool_tiles(value, attn_metadata.variable_block_sizes)
|
||||
out_c = torch.matmul(torch.softmax(scores, dim=-1), v_pooled) # [B, H, n_tiles, D]
|
||||
out_c = out_c.permute(0, 2, 1, 3).to(out.dtype) # [B, n_tiles, H, D]
|
||||
batch, seq_len, heads, dim = out.shape
|
||||
batch, _, heads, dim = out.shape
|
||||
n_tiles = attn_metadata.variable_block_sizes.numel()
|
||||
# Out-of-place: on the CuTe backend ``out`` is the tensor FA4's
|
||||
# autograd node saved for its backward, so an in-place add here
|
||||
# bumps its version counter and backward dies with "one of the
|
||||
# variables needed for gradient computation has been modified".
|
||||
out_tiled = out.view(batch, n_tiles, _TILE_ELEMS, heads, dim)
|
||||
gate_tiled = gate_compress.view(batch, n_tiles, _TILE_ELEMS, heads, dim)
|
||||
out = (out_tiled + out_c.unsqueeze(2) * gate_tiled).view(batch, seq_len, heads, dim)
|
||||
out.view(batch, n_tiles, _TILE_ELEMS, heads,
|
||||
dim).addcmul_(out_c.unsqueeze(2), gate_compress.view(batch, n_tiles, _TILE_ELEMS, heads, dim))
|
||||
return out
|
||||
|
||||
+1
-6
@@ -250,12 +250,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
|
||||
# not profile flops.
|
||||
"FASTVIDEO_TORCH_PROFILER_WITH_FLOPS":
|
||||
lambda: bool(os.getenv("FASTVIDEO_TORCH_PROFILER_WITH_FLOPS", "0") != "0"),
|
||||
# Wait steps per profiling cycle (torch.profiler.schedule wait parameter)
|
||||
# Defaults to 2 if not set.
|
||||
# Warmup steps per profiling cycle (torch.profiler.schedule warmup parameter)
|
||||
# Defaults to 1 if not set.
|
||||
# Active steps per profiling cycle (torch.profiler.schedule active parameter)
|
||||
# Defaults to 2 if not set.
|
||||
# Comma-separated names of registered profiler regions to capture.
|
||||
"FASTVIDEO_TORCH_PROFILE_REGIONS":
|
||||
lambda: os.getenv("FASTVIDEO_TORCH_PROFILE_REGIONS", ""),
|
||||
|
||||
|
||||
+185
-84
@@ -2,18 +2,17 @@
|
||||
"""Utilities for managing the PyTorch profiler within FastVideo.
|
||||
|
||||
The profiler is shared across the process; this module adds a light-weight
|
||||
controller that gates collection based on named *regions*. Regions may be
|
||||
enabled through dedicated environment variables (e.g.
|
||||
``FASTVIDEO_TORCH_PROFILE_MODEL_LOADING=1``) or via the consolidated
|
||||
``FASTVIDEO_TORCH_PROFILE_REGIONS`` comma-separated list. Short names work
|
||||
controller that gates collection based on named *regions*. Regions are enabled
|
||||
through the ``FASTVIDEO_TORCH_PROFILE_REGIONS`` comma-separated list. Short names work
|
||||
(``FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_train`` resolves the
|
||||
``profiler_region_`` prefix automatically).
|
||||
|
||||
Typical usage from client code::
|
||||
|
||||
controller = TorchProfilerController(profiler, activities)
|
||||
with controller.region("training_dit"):
|
||||
controller = get_or_create_profiler("/tmp/fastvideo-traces")
|
||||
with controller.region("training_train"):
|
||||
run_training_step()
|
||||
controller.stop()
|
||||
|
||||
To introduce a new region, register it via :func:`register_profiler_region`
|
||||
and wrap the corresponding code in :meth:`TorchProfilerController.region`.
|
||||
@@ -22,11 +21,11 @@ and wrap the corresponding code in :meth:`TorchProfilerController.region`.
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import os
|
||||
from collections.abc import Callable, Iterable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from collections.abc import Callable
|
||||
import functools
|
||||
from collections.abc import Iterable
|
||||
|
||||
import torch
|
||||
|
||||
@@ -142,6 +141,26 @@ register_profiler_region(
|
||||
name="profiler_region_training_train_one_step",
|
||||
description="Single optimizer step including forward/backward passes.",
|
||||
)
|
||||
register_profiler_region(
|
||||
name="profiler_region_training_forward",
|
||||
description="Training method forward pass and loss computation.",
|
||||
)
|
||||
register_profiler_region(
|
||||
name="profiler_region_training_dataloader",
|
||||
description="Fetch the next training batch in the trainer process.",
|
||||
)
|
||||
register_profiler_region(
|
||||
name="profiler_region_training_backward",
|
||||
description="Training backward pass.",
|
||||
)
|
||||
register_profiler_region(
|
||||
name="profiler_region_training_optimizer",
|
||||
description="Gradient clipping, optimizer/scheduler steps, and zero_grad.",
|
||||
)
|
||||
register_profiler_region(
|
||||
name="profiler_region_training_callbacks",
|
||||
description="End-of-step training callbacks such as EMA updates.",
|
||||
)
|
||||
register_profiler_region(
|
||||
name="profiler_region_training_train",
|
||||
description="High-level step orchestration in the training loop.",
|
||||
@@ -165,6 +184,21 @@ register_profiler_region(
|
||||
description="Parameter updates specific to distillation workflows.",
|
||||
)
|
||||
|
||||
# DMD2 method regions. These sit inside ``training_forward`` and make the
|
||||
# method's multi-model forward path distinguishable in a single trace.
|
||||
register_profiler_region(
|
||||
name="profiler_region_dmd2_student_rollout",
|
||||
description="DMD2 student rollout, including its simulated prefix steps.",
|
||||
)
|
||||
register_profiler_region(
|
||||
name="profiler_region_dmd2_generator_loss",
|
||||
description="DMD2 generator loss, including teacher and critic scoring.",
|
||||
)
|
||||
register_profiler_region(
|
||||
name="profiler_region_dmd2_critic_loss",
|
||||
description="DMD2 critic flow-matching loss, including its student rollout.",
|
||||
)
|
||||
|
||||
|
||||
def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController:
|
||||
"""Create or reuse the process-wide torch profiler controller."""
|
||||
@@ -189,25 +223,28 @@ def get_or_create_profiler(trace_dir: str | None) -> TorchProfilerController:
|
||||
)
|
||||
logger.info("FASTVIDEO_TORCH_PROFILE_REGIONS=%s", envs.FASTVIDEO_TORCH_PROFILE_REGIONS)
|
||||
|
||||
profiler = torch.profiler.profile(
|
||||
activities=_DEFAULT_ACTIVITIES,
|
||||
record_shapes=envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES,
|
||||
profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
|
||||
with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
|
||||
with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
|
||||
# No schedule: nothing in the codebase calls profiler.step(), so a
|
||||
# wait/warmup schedule never advances and the profiler records nothing.
|
||||
# Region toggling gates collection; the single trace exports at stop().
|
||||
on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir, use_gzip=True),
|
||||
def profiler_factory() -> Any:
|
||||
return torch.profiler.profile(
|
||||
activities=_DEFAULT_ACTIVITIES,
|
||||
record_shapes=envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES,
|
||||
profile_memory=envs.FASTVIDEO_TORCH_PROFILER_WITH_PROFILE_MEMORY,
|
||||
with_stack=envs.FASTVIDEO_TORCH_PROFILER_WITH_STACK,
|
||||
with_flops=envs.FASTVIDEO_TORCH_PROFILER_WITH_FLOPS,
|
||||
on_trace_ready=torch.profiler.tensorboard_trace_handler(trace_dir, use_gzip=True),
|
||||
)
|
||||
|
||||
controller = TorchProfilerController(
|
||||
None,
|
||||
_DEFAULT_ACTIVITIES,
|
||||
profiler_factory=profiler_factory,
|
||||
trace_dir=trace_dir,
|
||||
)
|
||||
controller = TorchProfilerController(profiler, _DEFAULT_ACTIVITIES)
|
||||
controller._trace_dir = trace_dir
|
||||
controller.start()
|
||||
# The trace only exports at stop(); inference paths have no shutdown hook
|
||||
# that calls it, so register one. stop() is idempotent.
|
||||
# Region exit normally exports each trace segment. Keep an atexit hook for
|
||||
# exceptions or process shutdown while a region is still active.
|
||||
import atexit
|
||||
atexit.register(controller.stop)
|
||||
logger.info("Torch profiler started")
|
||||
logger.info("Torch profiler armed; collection starts at the first enabled region")
|
||||
return controller
|
||||
|
||||
|
||||
@@ -260,7 +297,14 @@ class TorchProfilerConfig:
|
||||
|
||||
|
||||
class TorchProfilerController:
|
||||
"""Helper that toggles torch profiler collection for named regions.
|
||||
"""Create complete torch-profiler trace segments for named regions.
|
||||
|
||||
PyTorch's dynamic CUDA collection toggle can fail to re-enable CUPTI on
|
||||
some supported stacks. In that failure mode it emits CPU operators while
|
||||
silently dropping every CUDA kernel. This controller therefore starts a
|
||||
fresh profiler at each outermost enabled region and stops it when that
|
||||
region exits. Nested enabled regions become annotations in the same
|
||||
CPU/CUDA trace segment.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -273,12 +317,17 @@ class TorchProfilerController:
|
||||
config:
|
||||
Optional :class:`TorchProfilerConfig`. If omitted, :meth:`from_env`
|
||||
constructs one during initialization.
|
||||
profiler_factory:
|
||||
Factory for fresh profiler instances. Required to profile more than
|
||||
one outermost region invocation.
|
||||
trace_dir:
|
||||
Directory for per-segment summaries.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Enabling an existing region from the command line::
|
||||
|
||||
FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_dit \
|
||||
FASTVIDEO_TORCH_PROFILE_REGIONS=model_loading,training_train \
|
||||
python fastvideo/training/wan_training_pipeline.py ...
|
||||
|
||||
Wrapping a code block in a registered region::
|
||||
@@ -299,24 +348,29 @@ class TorchProfilerController:
|
||||
activities: Iterable[torch.profiler.ProfilerActivity],
|
||||
config: TorchProfilerConfig | None = None,
|
||||
disabled: bool = False,
|
||||
profiler_factory: Callable[[], Any] | None = None,
|
||||
trace_dir: str | None = None,
|
||||
) -> None:
|
||||
activities_tuple = tuple(activities)
|
||||
existing = get_global_controller()
|
||||
if existing is not None and not disabled:
|
||||
raise RuntimeError("TorchProfilerController already initialized globally. Use get_global_controller().")
|
||||
self._profiler = profiler
|
||||
self._profiler_factory = profiler_factory
|
||||
self._activities = activities_tuple
|
||||
self._trace_dir = trace_dir
|
||||
self._segment_index = 0
|
||||
self._segment_region: str | None = None
|
||||
self._active_region_depth = 0
|
||||
self._collection_enabled = False
|
||||
if disabled:
|
||||
self._profiler = None
|
||||
self._configured = False
|
||||
self._armed = False
|
||||
return
|
||||
|
||||
self._profiler = profiler
|
||||
self._activities = activities_tuple
|
||||
self._config = config or TorchProfilerConfig.from_env()
|
||||
# torch.profiler collects from start(); reflect that so the initial
|
||||
# _set_collection(False) in start() actually toggles it off instead of
|
||||
# short-circuiting (which captured everything before the first region).
|
||||
self._collection_enabled = True
|
||||
self._active_region_depth = 0
|
||||
self._trace_dir: str | None = None
|
||||
self._configured = True
|
||||
self._armed = False
|
||||
logger.info("PROFILER: TorchProfilerController initialized with config: %s", self._config)
|
||||
set_global_controller(self)
|
||||
|
||||
@@ -324,29 +378,61 @@ class TorchProfilerController:
|
||||
def is_enabled(self) -> bool:
|
||||
"""Return ``True`` when the underlying profiler is collecting."""
|
||||
|
||||
if self._profiler is None:
|
||||
return False
|
||||
return self._collection_enabled
|
||||
|
||||
def is_region_enabled(self, region: str) -> bool:
|
||||
"""Return ``True`` if ``region`` should be collected."""
|
||||
|
||||
if self._profiler is None:
|
||||
if not self.has_profiler:
|
||||
return False
|
||||
resolved = resolve_profiler_region(region)
|
||||
if resolved is None:
|
||||
return False
|
||||
return self._config.regions.get(resolved.name, False)
|
||||
|
||||
def _set_collection(self, enabled: bool) -> None:
|
||||
if self._profiler is None:
|
||||
def _new_profiler(self) -> Any:
|
||||
if self._profiler is not None:
|
||||
profiler = self._profiler
|
||||
self._profiler = None
|
||||
return profiler
|
||||
if self._profiler_factory is None:
|
||||
raise RuntimeError("Torch profiler cannot start another trace segment without a profiler_factory")
|
||||
return self._profiler_factory()
|
||||
|
||||
def _start_segment(self, region: str) -> None:
|
||||
if self._collection_enabled:
|
||||
return
|
||||
if self._collection_enabled == enabled:
|
||||
self._profiler = self._new_profiler()
|
||||
logger.info(
|
||||
"PROFILER: Starting segment %d for region %s",
|
||||
self._segment_index,
|
||||
region,
|
||||
)
|
||||
self._profiler.start()
|
||||
self._segment_region = region
|
||||
self._collection_enabled = True
|
||||
|
||||
def _finish_segment(self) -> None:
|
||||
if self._profiler is None or not self._collection_enabled:
|
||||
return
|
||||
event = ("fastvideo.profiler.enable_collection" if enabled else "fastvideo.profiler.disable_collection")
|
||||
with torch.profiler.record_function(event):
|
||||
self._profiler.toggle_collection_dynamic(enabled, self._activities)
|
||||
self._collection_enabled = enabled
|
||||
profiler = self._profiler
|
||||
segment_index = self._segment_index
|
||||
segment_region = self._segment_region or "unknown"
|
||||
logger.info(
|
||||
"PROFILER: Stopping segment %d for region %s",
|
||||
segment_index,
|
||||
segment_region,
|
||||
)
|
||||
profiler.stop()
|
||||
self._write_summary(
|
||||
profiler,
|
||||
segment_index=segment_index,
|
||||
segment_region=segment_region,
|
||||
)
|
||||
self._profiler = None
|
||||
self._collection_enabled = False
|
||||
self._segment_region = None
|
||||
self._segment_index += 1
|
||||
|
||||
_warned_unregistered: set[str] = set()
|
||||
|
||||
@@ -354,7 +440,7 @@ class TorchProfilerController:
|
||||
def region(self, region: str):
|
||||
"""Context manager that enables profiling for ``region`` if configured."""
|
||||
|
||||
if self._profiler is None:
|
||||
if not self.has_profiler:
|
||||
yield
|
||||
return
|
||||
|
||||
@@ -371,64 +457,70 @@ class TorchProfilerController:
|
||||
yield
|
||||
return
|
||||
|
||||
# NVTX range so the same region names are visible in nsys timelines
|
||||
if self._active_region_depth == 0:
|
||||
self._start_segment(region)
|
||||
|
||||
# NVTX range so the same region names are visible in nsys timelines.
|
||||
# Push after profiler startup so Kineto also records the annotation.
|
||||
nvtx = torch.cuda.is_available()
|
||||
if nvtx:
|
||||
torch.cuda.nvtx.range_push(f"fastvideo.region::{region}")
|
||||
self._active_region_depth += 1
|
||||
if self._active_region_depth == 1:
|
||||
logger.info("PROFILER: Setting collection to True (depth=%s) for region %s", self._active_region_depth,
|
||||
region)
|
||||
self._set_collection(True)
|
||||
try:
|
||||
# record_function opens after collection is enabled so the region
|
||||
# marker itself lands in the trace.
|
||||
with torch.profiler.record_function(f"fastvideo.region::{region}"):
|
||||
yield
|
||||
finally:
|
||||
self._active_region_depth -= 1
|
||||
logger.info("PROFILER: Decreasing active region depth to %s", self._active_region_depth)
|
||||
if self._active_region_depth == 0:
|
||||
logger.info("PROFILER: Setting collection to False upon exiting region %s", region)
|
||||
self._set_collection(False)
|
||||
if nvtx:
|
||||
torch.cuda.nvtx.range_pop()
|
||||
try:
|
||||
if nvtx:
|
||||
torch.cuda.nvtx.range_pop()
|
||||
finally:
|
||||
# Close NVTX before stopping Kineto so both profilers see a
|
||||
# balanced outermost range in the exported segment.
|
||||
if self._active_region_depth == 0:
|
||||
self._finish_segment()
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the profiler and pause collection until a region is entered."""
|
||||
"""Arm the controller; collection begins at an enabled region."""
|
||||
|
||||
logger.info("PROFILER: Starting profiler...")
|
||||
if self._profiler is None:
|
||||
if not self._configured:
|
||||
return
|
||||
self._profiler.start()
|
||||
logger.info("PROFILER: Profiler started")
|
||||
# Profiler starts with collection disabled by default.
|
||||
logger.info("PROFILER: Setting collection to False")
|
||||
self._set_collection(False)
|
||||
logger.info("PROFILER: Profiler started with collection disabled")
|
||||
self._armed = True
|
||||
logger.info("PROFILER: Controller armed")
|
||||
|
||||
def _write_summary(self) -> None:
|
||||
def _write_summary(
|
||||
self,
|
||||
profiler: Any,
|
||||
*,
|
||||
segment_index: int,
|
||||
segment_region: str,
|
||||
) -> None:
|
||||
"""Compact per-rank op summary next to the trace: a key_averages table
|
||||
and a JSON with input shapes, so operator-split analysis does not
|
||||
require parsing multi-GB chrome traces."""
|
||||
if self._profiler is None or not self._trace_dir:
|
||||
if not self._trace_dir:
|
||||
return
|
||||
try:
|
||||
import json as _json
|
||||
import os as _os
|
||||
rank = _os.environ.get("RANK", "0")
|
||||
averages = self._profiler.key_averages(group_by_input_shape=True)
|
||||
stem = _os.path.join(self._trace_dir, f"summary_rank{rank}")
|
||||
with open(f"{stem}.txt", "w") as fh:
|
||||
rank = os.environ.get("RANK", "0")
|
||||
short_region = segment_region.removeprefix("profiler_region_")
|
||||
stem = os.path.join(
|
||||
self._trace_dir,
|
||||
f"summary_rank{rank}_segment{segment_index:04d}_{short_region}",
|
||||
)
|
||||
averages = profiler.key_averages(group_by_input_shape=envs.FASTVIDEO_TORCH_PROFILER_RECORD_SHAPES, )
|
||||
with open(f"{stem}.txt", "w", encoding="utf-8") as fh:
|
||||
fh.write(averages.table(sort_by="self_device_time_total", row_limit=60))
|
||||
rows = [{
|
||||
"name": e.key,
|
||||
"shapes": str(e.input_shapes),
|
||||
"self_cpu_us": e.self_cpu_time_total,
|
||||
"cpu_us": e.cpu_time_total,
|
||||
"self_device_us": e.self_device_time_total,
|
||||
"device_us": e.device_time_total,
|
||||
"count": e.count,
|
||||
} for e in averages]
|
||||
with open(f"{stem}.json", "w") as fh:
|
||||
with open(f"{stem}.json", "w", encoding="utf-8") as fh:
|
||||
_json.dump(rows, fh)
|
||||
if rank == "0":
|
||||
logger.info("PROFILER: summary written to %s.txt", stem)
|
||||
@@ -436,24 +528,25 @@ class TorchProfilerController:
|
||||
logger.exception("PROFILER: summary generation failed")
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the profiler after disabling collection and clearing state."""
|
||||
"""Flush any active segment and disable this controller."""
|
||||
|
||||
if self._profiler is None:
|
||||
if not self._configured:
|
||||
return
|
||||
|
||||
logger.info("PROFILER: Stopping profiler...")
|
||||
self._profiler.stop()
|
||||
self._write_summary()
|
||||
self._profiler = None # makes stop() idempotent (atexit may re-enter)
|
||||
self._finish_segment()
|
||||
self._profiler = None
|
||||
self._configured = False
|
||||
self._armed = False
|
||||
logger.info("PROFILER: Profiler stopped")
|
||||
self._active_region_depth = 0
|
||||
set_global_controller(None)
|
||||
|
||||
@property
|
||||
def has_profiler(self) -> bool:
|
||||
"""Return ``True`` when a profiler instance is available."""
|
||||
"""Return ``True`` when this controller is configured and armed."""
|
||||
|
||||
return self._profiler is not None
|
||||
return self._configured and self._armed
|
||||
|
||||
@property
|
||||
def activities(self) -> tuple[torch.profiler.ProfilerActivity, ...]:
|
||||
@@ -476,13 +569,21 @@ def profiler_region(region: str):
|
||||
|
||||
|
||||
def profile_region(region: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
|
||||
"""Wrap a bound method so it runs inside a profiler region if available."""
|
||||
"""Wrap a bound method so it runs inside a profiler region if available.
|
||||
|
||||
Prefer a controller attached to the instance, then fall back to the
|
||||
process-wide controller. The fallback lets lightweight owners such as the
|
||||
modular trainer and its callbacks add regions without threading profiler
|
||||
plumbing through their public constructors.
|
||||
"""
|
||||
|
||||
def decorator(fn: Callable[..., Any]) -> Callable[..., Any]:
|
||||
|
||||
@functools.wraps(fn)
|
||||
def wrapped(self, *args, **kwargs):
|
||||
controller = getattr(self, "profiler_controller", None)
|
||||
if controller is None:
|
||||
controller = get_global_controller()
|
||||
if controller is None or not controller.has_profiler:
|
||||
return fn(self, *args, **kwargs)
|
||||
with controller.region(region):
|
||||
|
||||
@@ -4,11 +4,11 @@ the real platform resolver.
|
||||
|
||||
``is_attn_qat_infer_available()`` used to test only whether the kernel
|
||||
extension imports. CUDA 13 wheel builds can carry the sm_120/sm_121
|
||||
extension on any host (e.g. H100 sm_90, GB200 sm_100). Without the
|
||||
capability gate, that successful import selects the consumer-Blackwell
|
||||
backend and defers failure until the first unsupported kernel call.
|
||||
Explicit ATTN_QAT_INFER requests must instead fail closed during
|
||||
resolution.
|
||||
extension on any host (e.g. H100 sm_90, GB200 sm_100): the import
|
||||
succeeds, ``CudaPlatformBase.get_attn_backend_cls`` selects the
|
||||
consumer-Blackwell backend, and the first kernel call fails with an
|
||||
unsupported-capability error -- instead of the FlashAttention fallback
|
||||
the QAD README documents for non-sm_120 GPUs.
|
||||
|
||||
These tests drive the REAL resolver (``fastvideo.platforms.cuda``) and
|
||||
the REAL availability function; only the two physical facts are faked --
|
||||
@@ -17,7 +17,7 @@ active" (``torch.cuda``). The stage-guard test
|
||||
(fastvideo/tests/stages/test_kandinsky5_attention_backend_guard.py)
|
||||
injects an already-resolved backend and by design cannot see this bug.
|
||||
|
||||
CPU-only: the ATTN_QAT_INFER branch and its error path never require a
|
||||
CPU-only: the ATTN_QAT_INFER branch and its fallback never require a
|
||||
physical GPU to *resolve* (only to run).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
@@ -35,6 +35,13 @@ from fastvideo.platforms.cuda import NonNvmlCudaPlatform
|
||||
from fastvideo.platforms.interface import AttentionBackendEnum
|
||||
|
||||
ATTN_QAT_INFER_CLS = "fastvideo.attention.backends.attn_qat_infer.AttnQatInferBackend"
|
||||
# What the resolver's fallthrough legitimately returns when ATTN_QAT_INFER
|
||||
# is unavailable: FlashAttention, or SDPA when flash_attn isn't installed
|
||||
# in the running environment (e.g. CPU-only CI).
|
||||
FALLBACK_CLASSES = {
|
||||
"fastvideo.attention.backends.flash_attn.FlashAttentionBackend",
|
||||
"fastvideo.attention.backends.sdpa.SDPABackend",
|
||||
}
|
||||
|
||||
|
||||
def _fake_gpu(monkeypatch, *, capability: tuple[int, int], extension_imports: bool, fa4_imports: bool = False) -> None:
|
||||
@@ -59,24 +66,22 @@ def _resolve() -> str:
|
||||
)
|
||||
|
||||
|
||||
def test_sm90_host_with_bundled_extension_fails_closed(monkeypatch):
|
||||
def test_sm90_host_with_bundled_extension_falls_back(monkeypatch):
|
||||
"""The reviewed failure: H100 + CUDA 13 wheel that bundles the sm_120
|
||||
extension. Import succeeds; explicit selection must still fail closed."""
|
||||
extension. Import succeeds; selection must still fall back."""
|
||||
_fake_gpu(monkeypatch, capability=(9, 0), extension_imports=True)
|
||||
|
||||
assert not is_attn_qat_infer_available()
|
||||
with pytest.raises(ImportError, match="ATTN_QAT_INFER selected but"):
|
||||
_resolve()
|
||||
assert _resolve() in FALLBACK_CLASSES
|
||||
|
||||
|
||||
def test_sm100_host_with_bundled_extension_fails_closed(monkeypatch):
|
||||
def test_sm100_host_with_bundled_extension_falls_back(monkeypatch):
|
||||
"""sm_100 with only the (unrunnable) bundled sm_12x extension and no
|
||||
FP4 FA4 kernel still fails closed -- the original reviewed failure class."""
|
||||
FP4 FA4 kernel still falls back -- the original reviewed failure class."""
|
||||
_fake_gpu(monkeypatch, capability=(10, 0), extension_imports=True, fa4_imports=False)
|
||||
|
||||
assert not is_attn_qat_infer_available()
|
||||
with pytest.raises(ImportError, match="ATTN_QAT_INFER selected but"):
|
||||
_resolve()
|
||||
assert _resolve() in FALLBACK_CLASSES
|
||||
|
||||
|
||||
@pytest.mark.parametrize("capability", [(10, 0), (10, 3)])
|
||||
@@ -97,12 +102,11 @@ def test_consumer_blackwell_with_extension_selects_backend(monkeypatch, capabili
|
||||
assert _resolve() == ATTN_QAT_INFER_CLS
|
||||
|
||||
|
||||
def test_consumer_blackwell_without_extension_fails_closed(monkeypatch):
|
||||
def test_consumer_blackwell_without_extension_falls_back(monkeypatch):
|
||||
_fake_gpu(monkeypatch, capability=(12, 0), extension_imports=False)
|
||||
|
||||
assert not is_attn_qat_infer_available()
|
||||
with pytest.raises(ImportError, match="ATTN_QAT_INFER selected but"):
|
||||
_resolve()
|
||||
assert _resolve() in FALLBACK_CLASSES
|
||||
|
||||
|
||||
def test_no_cuda_reports_unavailable(monkeypatch):
|
||||
|
||||
@@ -1,110 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""GPU backward checks for the VSA-H3 backend.
|
||||
|
||||
The CuTe backend returns FA4's own output tensor, which FA4's autograd node
|
||||
saved for its backward. Composing the compression branch onto it in place
|
||||
therefore poisons the graph, and the failure only appears once the VSA-256
|
||||
CuTe path has a backward at all. These tests pin the composition.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo.attention.backends.video_sparse_attn_h3 import (MiniMaxH3VSAImpl, MiniMaxH3VSAMetadataBuilder)
|
||||
|
||||
_SPEC = dict(raw_latent_shape=(16, 16, 24), patch_size=(1, 2, 2), prefix_segments=(64, 32, 16))
|
||||
_HEADS = 2
|
||||
_DIM = 128
|
||||
|
||||
|
||||
def _build_meta(device, sparsity=0.5):
|
||||
return MiniMaxH3VSAMetadataBuilder().build(
|
||||
current_timestep=0,
|
||||
raw_latent_shape=_SPEC["raw_latent_shape"],
|
||||
patch_size=_SPEC["patch_size"],
|
||||
VSA_sparsity=sparsity,
|
||||
prefix_segments=_SPEC["prefix_segments"],
|
||||
device=device,
|
||||
)
|
||||
|
||||
|
||||
def _select_backend(monkeypatch, backend):
|
||||
if backend == "cute":
|
||||
pytest.importorskip(
|
||||
"flash_attn.cute.block_sparsity",
|
||||
reason="optional FA4 CuTe build (flash_attn.cute) not installed",
|
||||
)
|
||||
monkeypatch.setenv("FASTVIDEO_VSA_CUTEDSL", "1")
|
||||
monkeypatch.delenv("FASTVIDEO_VSA_TRITON", raising=False)
|
||||
monkeypatch.delenv("FASTVIDEO_KERNEL_VSA_FORCE_TRITON", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("FASTVIDEO_VSA_TRITON", "1")
|
||||
monkeypatch.delenv("FASTVIDEO_VSA_CUTEDSL", raising=False)
|
||||
|
||||
|
||||
def _forward_backward(impl, meta, gate_compress, device):
|
||||
seq = meta.total_seq_length
|
||||
torch.manual_seed(0)
|
||||
q, k, v = (torch.randn(1, seq, _HEADS, _DIM, device=device, dtype=torch.bfloat16, requires_grad=True)
|
||||
for _ in range(3))
|
||||
tq, tk, tv = (impl.tile(t, meta).clone() for t in (q, k, v))
|
||||
|
||||
gate = None
|
||||
if gate_compress:
|
||||
gate = torch.randn(1, tq.shape[1], _HEADS, _DIM, device=device, dtype=torch.bfloat16) * 0.1
|
||||
|
||||
out = impl.forward(tq, tk, tv, gate, meta)
|
||||
out = impl.postprocess_output(out, meta)
|
||||
out.float().pow(2).sum().backward()
|
||||
return out, (q, k, v)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["triton", "cute"])
|
||||
@pytest.mark.parametrize("gate_compress", [False, True])
|
||||
def test_h3_vsa_backward_runs(monkeypatch, backend: str, gate_compress: bool) -> None:
|
||||
"""Regression: with the CuTe backend and a non-zero gate this used to die
|
||||
with "one of the variables needed for gradient computation has been
|
||||
modified by an inplace operation ... output 0 of FlashAttnFuncBackward".
|
||||
"""
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA is required")
|
||||
_select_backend(monkeypatch, backend)
|
||||
|
||||
device = torch.device("cuda")
|
||||
meta = _build_meta(device)
|
||||
impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
|
||||
|
||||
out, leaves = _forward_backward(impl, meta, gate_compress, device)
|
||||
|
||||
assert torch.isfinite(out).all().item()
|
||||
for name, leaf in zip(("q", "k", "v"), leaves):
|
||||
assert leaf.grad is not None, f"{name} received no gradient"
|
||||
assert torch.isfinite(leaf.grad).all().item(), f"{name}.grad has non-finite values"
|
||||
assert leaf.grad.abs().sum().item() > 0, f"{name}.grad is all zero"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gate_compress", [False, True])
|
||||
def test_h3_vsa_backward_cute_matches_triton(monkeypatch, gate_compress: bool) -> None:
|
||||
"""CuTe and Triton take different routes to the same math; their gradients
|
||||
should agree to bf16 tolerance."""
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA is required")
|
||||
|
||||
device = torch.device("cuda")
|
||||
impl = MiniMaxH3VSAImpl(num_heads=_HEADS, head_size=_DIM, causal=False, softmax_scale=_DIM**-0.5)
|
||||
|
||||
grads = {}
|
||||
for backend in ("triton", "cute"):
|
||||
with monkeypatch.context() as m:
|
||||
_select_backend(m, backend)
|
||||
meta = _build_meta(device)
|
||||
_, leaves = _forward_backward(impl, meta, gate_compress, device)
|
||||
grads[backend] = [leaf.grad.detach().float() for leaf in leaves]
|
||||
|
||||
for name, ref, got in zip(("dq", "dk", "dv"), grads["triton"], grads["cute"]):
|
||||
diff = (ref - got).abs()
|
||||
avg_abs = diff.mean().item()
|
||||
max_rel = (diff.max() / (ref.abs().mean() + 1e-6)).item()
|
||||
print(f"[h3-vsa gate={gate_compress}] {name}: avg_abs={avg_abs:.6e}, max_rel={max_rel:.6e}")
|
||||
assert avg_abs < 1e-2, f"{name}: avg_abs {avg_abs:.3e}"
|
||||
assert max_rel < 0.5, f"{name}: max_rel {max_rel:.3e}"
|
||||
@@ -23,7 +23,6 @@ CI_SOURCES = [
|
||||
TESTS_ROOT / "modal" / "ssim_test.py",
|
||||
*sorted((REPO_ROOT / ".buildkite").rglob("*.yml")),
|
||||
*sorted((REPO_ROOT / ".buildkite").rglob("*.sh")),
|
||||
*sorted((REPO_ROOT / ".github/workflows").glob("ci-*.yml")),
|
||||
]
|
||||
|
||||
# Directories that intentionally have no CI lane today. Every entry needs a
|
||||
@@ -69,22 +68,6 @@ def test_local_tests_stays_out_of_ci():
|
||||
"reference or move the tests into a fastvideo/tests/ lane.")
|
||||
|
||||
|
||||
def test_unit_ci_routes_to_trusted_static_driver():
|
||||
slash_commands = (REPO_ROOT / ".github/workflows/ci-slash-commands.yml").read_text()
|
||||
pipeline = (REPO_ROOT / ".buildkite/pipeline.yml").read_text()
|
||||
|
||||
valid_line = next(line for line in slash_commands.splitlines() if line.strip().startswith("VALID="))
|
||||
assert "unit-ci" in valid_line
|
||||
assert "[unit-ci]=unit_test_ci" in slash_commands
|
||||
assert ''' - label: ":microscope: Unit Tests"
|
||||
key: "unit-ci"
|
||||
if: build.env("TEST_SCOPE") == "direct" && build.env("TEST_TYPE") == "unit_test_ci"
|
||||
command: "/opt/fastvideo-ci-runner/run-unit"
|
||||
timeout_in_minutes: 90
|
||||
agents:
|
||||
queue: "ci-runner"''' in pipeline
|
||||
|
||||
|
||||
def test_allowlist_entries_are_still_real_directories():
|
||||
# A stale allowlist hides regressions; entries must track reality.
|
||||
missing = [name for name in ALLOWLIST if name != "modal" and not (TESTS_ROOT / name).is_dir()]
|
||||
|
||||
@@ -27,13 +27,13 @@ def _function_strings(path: Path, function_name: str) -> str:
|
||||
|
||||
|
||||
def test_generic_l40s_launcher_defaults_fa4_off():
|
||||
source = ast.unparse(ast.parse(LAUNCH_L40S_JOB.read_text(encoding="utf-8")))
|
||||
assert "'FASTVIDEO_FA4': os.environ.get('FASTVIDEO_FA4', '0')" in source
|
||||
source = LAUNCH_L40S_JOB.read_text(encoding="utf-8")
|
||||
assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "0")' in source
|
||||
|
||||
|
||||
def test_ssim_launcher_keeps_fa4_enabled_by_default():
|
||||
source = ast.unparse(ast.parse(SSIM_TEST.read_text(encoding="utf-8")))
|
||||
assert "'FASTVIDEO_FA4': os.environ.get('FASTVIDEO_FA4', '1')" in source
|
||||
source = SSIM_TEST.read_text(encoding="utf-8")
|
||||
assert '"FASTVIDEO_FA4": os.environ.get("FASTVIDEO_FA4", "1")' in source
|
||||
|
||||
|
||||
def test_performance_identity_env_reaches_modal_runtime():
|
||||
|
||||
@@ -16,6 +16,9 @@ import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
# Five-window child: ops before any region, inside a region, between regions,
|
||||
# inside a second (short-named) region, after the last region. Exits without
|
||||
# calling stop() — export must happen via the atexit hook. Each window is
|
||||
@@ -61,19 +64,21 @@ def _run_child(tmp_path):
|
||||
return trace_dir, proc.stdout + proc.stderr
|
||||
|
||||
|
||||
def _trace_event_names(trace_dir):
|
||||
def _trace_events(trace_dir):
|
||||
traces = glob.glob(os.path.join(trace_dir, "**", "*.json*"), recursive=True)
|
||||
traces = [t for t in traces if "summary" not in os.path.basename(t)]
|
||||
assert traces, f"no trace exported to {trace_dir} (atexit hook missing?)"
|
||||
opener = gzip.open if traces[0].endswith(".gz") else open
|
||||
with opener(traces[0], "rt") as fh:
|
||||
events = json.load(fh).get("traceEvents", [])
|
||||
return {e.get("name", "") for e in events}
|
||||
events = []
|
||||
for trace in traces:
|
||||
opener = gzip.open if trace.endswith(".gz") else open
|
||||
with opener(trace, "rt") as fh:
|
||||
events.extend(json.load(fh).get("traceEvents", []))
|
||||
return events
|
||||
|
||||
|
||||
def test_regions_gate_collection_and_atexit_exports(tmp_path):
|
||||
trace_dir, output = _run_child(tmp_path)
|
||||
names = _trace_event_names(trace_dir)
|
||||
names = {event.get("name", "") for event in _trace_events(trace_dir)}
|
||||
|
||||
assert "win_region1" in names, "op inside an enabled region was not captured"
|
||||
assert "win_region2" in names, "short region name did not resolve/capture"
|
||||
@@ -87,8 +92,61 @@ def test_regions_gate_collection_and_atexit_exports(tmp_path):
|
||||
assert output.count("is not registered") == 1
|
||||
|
||||
# per-rank op summary written next to the trace
|
||||
summaries = glob.glob(os.path.join(trace_dir, "summary_rank0.*"))
|
||||
assert sorted(os.path.splitext(s)[1] for s in summaries) == [".json", ".txt"]
|
||||
summaries = glob.glob(os.path.join(trace_dir, "summary_rank0_segment*.*"))
|
||||
assert sorted(os.path.splitext(s)[1] for s in summaries) == [
|
||||
".json",
|
||||
".json",
|
||||
".txt",
|
||||
".txt",
|
||||
]
|
||||
with open(next(s for s in summaries if s.endswith(".json")), encoding="utf-8") as fh:
|
||||
summary_rows = json.load(fh)
|
||||
assert summary_rows
|
||||
assert {
|
||||
"name",
|
||||
"shapes",
|
||||
"self_cpu_us",
|
||||
"cpu_us",
|
||||
"self_device_us",
|
||||
"device_us",
|
||||
"count",
|
||||
} <= summary_rows[0].keys()
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
|
||||
def test_cuda_region_exports_kernel_events(tmp_path):
|
||||
trace_dir = str(tmp_path / "cuda_traces")
|
||||
child = r"""
|
||||
import torch
|
||||
from fastvideo.profiler import get_or_create_profiler, profiler_region
|
||||
|
||||
controller = get_or_create_profiler({trace_dir!r})
|
||||
x = torch.randn(1024, 1024, device="cuda")
|
||||
y = torch.randn(1024, 1024, device="cuda")
|
||||
torch.cuda.synchronize()
|
||||
with profiler_region("training_forward"):
|
||||
torch.mm(x, y)
|
||||
torch.cuda.synchronize()
|
||||
controller.stop()
|
||||
""".format(trace_dir=trace_dir)
|
||||
env = os.environ.copy()
|
||||
env["FASTVIDEO_TORCH_PROFILER_DIR"] = trace_dir
|
||||
env["FASTVIDEO_TORCH_PROFILE_REGIONS"] = "training_forward"
|
||||
proc = subprocess.run(
|
||||
[sys.executable, "-c", child],
|
||||
env=env,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=300,
|
||||
)
|
||||
assert proc.returncode == 0, proc.stderr
|
||||
|
||||
events = _trace_events(trace_dir)
|
||||
categories = {event.get("cat", "") for event in events}
|
||||
names = {event.get("name", "") for event in events}
|
||||
assert "kernel" in categories
|
||||
assert "cuda_runtime" in categories
|
||||
assert "fastvideo.region::training_forward" in names
|
||||
|
||||
|
||||
def test_noop_without_profiler_dir(tmp_path):
|
||||
|
||||
@@ -22,7 +22,7 @@ import zipfile
|
||||
from email.parser import Parser
|
||||
from pathlib import Path
|
||||
|
||||
CACHE_SCHEMA_VERSION = 4
|
||||
CACHE_SCHEMA_VERSION = 3
|
||||
DEFAULT_PREBUILT_INFO_PATH = "/opt/fastvideo-kernel-prebuilt"
|
||||
KERNEL_RELATIVE_DIR = "fastvideo-kernel"
|
||||
DEFAULT_BUILD_INFO_OUTPUT = "/opt/fastvideo-kernel-prebuilt/default/metadata.json"
|
||||
@@ -166,7 +166,6 @@ def _torch_metadata() -> dict[str, object]:
|
||||
return {
|
||||
"torch_version": str(torch.__version__),
|
||||
"torch_cuda_version": str(torch.version.cuda),
|
||||
"torch_git_version": str(getattr(torch.version, "git_version", "")),
|
||||
"torch_file": str(getattr(torch, "__file__", "")),
|
||||
"torch_config": str(torch.__config__.show()),
|
||||
"cxx11_abi": cxx11_abi,
|
||||
@@ -175,7 +174,6 @@ def _torch_metadata() -> dict[str, object]:
|
||||
return {
|
||||
"torch_version": f"<unavailable: {error}>",
|
||||
"torch_cuda_version": "<unavailable>",
|
||||
"torch_git_version": "<unavailable>",
|
||||
"torch_file": "<unavailable>",
|
||||
"torch_config": "<unavailable>",
|
||||
"cxx11_abi": "<unavailable>",
|
||||
@@ -225,11 +223,6 @@ def _compiler_libc_metadata() -> dict[str, object]:
|
||||
def _build_metadata(repo_root: Path) -> dict[str, object]:
|
||||
explicit_arch = os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip()
|
||||
resolved_arch = explicit_arch or _detect_arch_from_torch()
|
||||
torch_metadata = _torch_metadata()
|
||||
torch_cache_metadata = {
|
||||
name: torch_metadata[name]
|
||||
for name in ("torch_version", "torch_cuda_version", "torch_git_version", "cxx11_abi")
|
||||
}
|
||||
cache_key_build = {
|
||||
"gpu_backend": os.environ.get("GPU_BACKEND", "CUDA"),
|
||||
"resolved_torch_cuda_arch_list": resolved_arch,
|
||||
@@ -249,7 +242,7 @@ def _build_metadata(repo_root: Path) -> dict[str, object]:
|
||||
"platform": sysconfig.get_platform(),
|
||||
"machine": platform.machine(),
|
||||
},
|
||||
"torch": torch_cache_metadata,
|
||||
"torch": _torch_metadata(),
|
||||
"cuda": {
|
||||
"cuda_home": os.environ.get("CUDA_HOME", ""),
|
||||
"nvcc": _selected_command_metadata("CUDACXX", "nvcc"),
|
||||
@@ -259,7 +252,6 @@ def _build_metadata(repo_root: Path) -> dict[str, object]:
|
||||
}
|
||||
metadata = {
|
||||
**cache_key_metadata,
|
||||
"torch": torch_metadata,
|
||||
"build": {
|
||||
**cache_key_build,
|
||||
"torch_cuda_arch_list": explicit_arch,
|
||||
|
||||
@@ -349,7 +349,15 @@ def run_self_forcing_tests():
|
||||
|
||||
@app.function(gpu="L40S:1", image=image, timeout=900, secrets=[ci_env_secret])
|
||||
def run_unit_test():
|
||||
run_test("bash .buildkite/scripts/unit_test.sh")
|
||||
run_test("pytest ./fastvideo/tests/api/ ./fastvideo/tests/contract/ ./fastvideo/tests/dataset/ "
|
||||
"./fastvideo/tests/workflow/ ./fastvideo/tests/entrypoints/ ./fastvideo/tests/train/ "
|
||||
"./fastvideo/tests/stages/ ./fastvideo/tests/ops/ ./fastvideo/tests/worker/ "
|
||||
"./fastvideo/tests/training/test_trackers.py "
|
||||
"./fastvideo/tests/attention/test_sdpa_metadata_mask_contract.py "
|
||||
"./fastvideo/tests/modal/test_kernel_build_cache.py ./fastvideo/tests/modal/test_pr_test.py "
|
||||
"./fastvideo/tests/modal/test_ssim_test.py "
|
||||
"--ignore=./fastvideo/tests/entrypoints/test_openai_api_integration.py "
|
||||
"--ignore=./fastvideo/tests/train/models --ignore=./fastvideo/tests/train/methods -vs")
|
||||
|
||||
|
||||
# TODO: David: GPU only used to resolve import time requirement (not needed for this test). Maybe make those imports lazy?
|
||||
|
||||
@@ -106,7 +106,6 @@ def _patch_stable_metadata(monkeypatch) -> None:
|
||||
lambda: {
|
||||
"torch_version": "2.9.0",
|
||||
"torch_cuda_version": "12.8",
|
||||
"torch_git_version": "stable-build-commit",
|
||||
"torch_file": "/opt/venv/lib/python3.12/site-packages/torch/__init__.py",
|
||||
"torch_config": "USE_CUDA=ON",
|
||||
"cxx11_abi": True,
|
||||
@@ -156,27 +155,6 @@ def test_kernel_only_main_change_republishes_trusted_l40s_artifact() -> None:
|
||||
assert '--output "${l40s_wheel_dir}/metadata.json"' in dockerfile
|
||||
|
||||
|
||||
def test_ci_runner_image_targets_arm64_sm100_with_opencv_runtime() -> None:
|
||||
workflow = yaml.load(
|
||||
(REPO_ROOT / ".github/workflows/infra-build-image.yml").read_text(encoding="utf-8"),
|
||||
Loader=yaml.BaseLoader,
|
||||
)
|
||||
job = workflow["jobs"]["build-ci-runner-image"]
|
||||
|
||||
assert job["with"]["architecture"] == "arm64"
|
||||
assert job["with"]["runner"] == "ubuntu-24.04-arm"
|
||||
assert job["with"]["tag_suffix"] == "py3.12-cuda13.0.0-sm100"
|
||||
assert "CUDA_VERSION=13.0.0" in job["with"]["build_args"]
|
||||
assert "UV_TORCH_BACKEND=cu130" in job["with"]["build_args"]
|
||||
assert "TORCH_CUDA_ARCH_LIST=10.0" in job["with"]["build_args"]
|
||||
|
||||
dockerfile = (REPO_ROOT / "docker/Dockerfile").read_text(encoding="utf-8")
|
||||
assert " ffmpeg \\\n" in dockerfile
|
||||
assert " libgl1 \\\n" in dockerfile
|
||||
assert " libglib2.0-0 \\\n" in dockerfile
|
||||
assert 'python -c "import cv2; print(\'OpenCV\', cv2.__version__)"' in dockerfile
|
||||
|
||||
|
||||
def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, tmp_path) -> None:
|
||||
_patch_stable_metadata(monkeypatch)
|
||||
monkeypatch.setattr(kernel_build_cache, "_detect_arch_from_torch", lambda: "9.0a")
|
||||
@@ -196,25 +174,6 @@ def test_cache_key_uses_resolved_arch_not_raw_env(monkeypatch, tmp_path) -> None
|
||||
assert detected_l40s["cache_key"] != detected_hopper["cache_key"]
|
||||
|
||||
|
||||
def test_cache_key_ignores_runtime_only_torch_config(monkeypatch, tmp_path) -> None:
|
||||
_patch_stable_metadata(monkeypatch)
|
||||
build_host = kernel_build_cache._build_metadata(tmp_path)
|
||||
torch_metadata = kernel_build_cache._torch_metadata()
|
||||
monkeypatch.setattr(
|
||||
kernel_build_cache,
|
||||
"_torch_metadata",
|
||||
lambda: {
|
||||
**torch_metadata,
|
||||
"torch_config": torch_metadata["torch_config"] + "\nCUDA Runtime 13.0\nCuDNN 92.0",
|
||||
},
|
||||
)
|
||||
|
||||
gpu_host = kernel_build_cache._build_metadata(tmp_path)
|
||||
|
||||
assert gpu_host["torch"]["torch_config"] != build_host["torch"]["torch_config"]
|
||||
assert gpu_host["cache_key"] == build_host["cache_key"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("environment_name", ["CFLAGS", "CXXFLAGS", "LDFLAGS"])
|
||||
def test_cache_key_changes_with_build_flags(monkeypatch, tmp_path, environment_name) -> None:
|
||||
_patch_stable_metadata(monkeypatch)
|
||||
@@ -254,7 +213,7 @@ def test_cache_key_changes_with_compiler_or_torch_abi(monkeypatch, tmp_path) ->
|
||||
monkeypatch.setattr(kernel_build_cache, "_torch_metadata", lambda: {**torch_metadata, "cxx11_abi": False})
|
||||
torch_abi_changed = kernel_build_cache._build_metadata(tmp_path)
|
||||
|
||||
assert baseline["schema_version"] == 4
|
||||
assert baseline["schema_version"] == 3
|
||||
assert baseline["cache_key"] != compiler_changed["cache_key"]
|
||||
assert baseline["cache_key"] != torch_abi_changed["cache_key"]
|
||||
|
||||
|
||||
@@ -316,18 +316,17 @@ def test_run_test_command_uses_nonshared_kernel_install_before_tests(monkeypatch
|
||||
assert not hasattr(module, "kernel_cache_vol")
|
||||
|
||||
|
||||
def test_run_unit_test_uses_shared_command(monkeypatch):
|
||||
def test_run_unit_test_collects_modal_cache_runner_tests(monkeypatch):
|
||||
module = _load_pr_test_module(monkeypatch)
|
||||
commands = []
|
||||
monkeypatch.setattr(module, "run_test", commands.append)
|
||||
|
||||
module.run_unit_test()
|
||||
|
||||
assert commands == ["bash .buildkite/scripts/unit_test.sh"]
|
||||
unit_command = (Path(__file__).resolve().parents[3] / ".buildkite/scripts/unit_test.sh").read_text()
|
||||
assert len(commands) == 1
|
||||
for test_path in (
|
||||
"./fastvideo/tests/modal/test_kernel_build_cache.py",
|
||||
"./fastvideo/tests/modal/test_pr_test.py",
|
||||
"./fastvideo/tests/modal/test_ssim_test.py",
|
||||
):
|
||||
assert test_path in unit_command
|
||||
assert test_path in commands[0]
|
||||
|
||||
@@ -42,9 +42,15 @@ class _Method:
|
||||
self,
|
||||
transformer: torch.nn.Module | None,
|
||||
tracker: Any | None = None,
|
||||
updated_iterations: set[int] | None = None,
|
||||
) -> None:
|
||||
self.student = _Student(transformer)
|
||||
self.tracker = tracker
|
||||
self.updated_iterations = updated_iterations
|
||||
|
||||
def did_update_role(self, role: str, iteration: int) -> bool:
|
||||
assert role == "student"
|
||||
return (self.updated_iterations is None or iteration in self.updated_iterations)
|
||||
|
||||
|
||||
def _tiny_transformer(*, fill: float = 0.0) -> torch.nn.Module:
|
||||
@@ -155,6 +161,44 @@ class TestOnTrainingStepEnd:
|
||||
|
||||
assert any(payload.get("ema/decay") == 0.99 and step == 0 for payload, step in tracker.entries)
|
||||
|
||||
def test_only_updates_after_student_optimizer_step(self) -> None:
|
||||
transformer = _tiny_transformer(fill=1.0)
|
||||
tracker = _RecordingTracker()
|
||||
method = _Method(
|
||||
transformer,
|
||||
tracker=tracker,
|
||||
updated_iterations={5},
|
||||
)
|
||||
cb = EMACallback(decay=0.9, start_iter=0)
|
||||
cb.on_train_start(method, iteration=0)
|
||||
|
||||
with torch.no_grad():
|
||||
transformer.weight.fill_(7.0)
|
||||
cb.on_training_step_end(method, loss_dict={}, iteration=4)
|
||||
assert not cb._ema_started
|
||||
assert torch.allclose(
|
||||
cb.student_ema.shadow["weight"],
|
||||
torch.full((2, 4), 1.0),
|
||||
)
|
||||
assert tracker.entries == []
|
||||
|
||||
cb.on_training_step_end(method, loss_dict={}, iteration=5)
|
||||
assert cb._ema_started
|
||||
assert torch.allclose(
|
||||
cb.student_ema.shadow["weight"],
|
||||
torch.full((2, 4), 7.0),
|
||||
)
|
||||
assert tracker.entries == [({"ema/decay": 0.9}, 5)]
|
||||
|
||||
with torch.no_grad():
|
||||
transformer.weight.fill_(11.0)
|
||||
cb.on_training_step_end(method, loss_dict={}, iteration=6)
|
||||
assert torch.allclose(
|
||||
cb.student_ema.shadow["weight"],
|
||||
torch.full((2, 4), 7.0),
|
||||
)
|
||||
assert tracker.entries == [({"ema/decay": 0.9}, 5)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# C. ema_context
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
|
||||
|
||||
|
||||
class _RecordingStudent:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.predict_calls: list[torch.Tensor] = []
|
||||
self.add_noise_calls: list[torch.Tensor] = []
|
||||
|
||||
def predict_x0(
|
||||
self,
|
||||
noisy_latents: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
batch: Any,
|
||||
**kwargs: Any,
|
||||
) -> torch.Tensor:
|
||||
del batch, kwargs
|
||||
self.predict_calls.append(timestep.detach().clone())
|
||||
return noisy_latents + timestep.to(noisy_latents.dtype)
|
||||
|
||||
def add_noise(
|
||||
self,
|
||||
clean_latents: torch.Tensor,
|
||||
noise: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
self.add_noise_calls.append(timestep.detach().clone())
|
||||
return clean_latents + noise
|
||||
|
||||
|
||||
def _rollout_method(seed: int = 1234) -> tuple[DMD2Method, _RecordingStudent]:
|
||||
method = object.__new__(DMD2Method)
|
||||
torch.nn.Module.__init__(method)
|
||||
student = _RecordingStudent()
|
||||
method.student = student
|
||||
method._rollout_mode = "simulate"
|
||||
method._cfg_uncond = None
|
||||
method._denoising_step_list = torch.tensor([1000, 750, 500, 250])
|
||||
method.cuda_generator = torch.Generator(device="cpu").manual_seed(seed)
|
||||
return method, student
|
||||
|
||||
|
||||
def _legacy_full_rollout_reference(
|
||||
*,
|
||||
seed: int,
|
||||
target_idx: int,
|
||||
shape: tuple[int, ...],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Reproduce the pre-optimization simulate rollout and RNG state."""
|
||||
step_list = torch.tensor([1000, 750, 500, 250])
|
||||
generator = torch.Generator(device="cpu").manual_seed(seed)
|
||||
current = torch.randn(shape, generator=generator)
|
||||
initial = current.clone()
|
||||
noise_latents: list[torch.Tensor] = []
|
||||
|
||||
for step_idx in range(len(step_list) - 1):
|
||||
pred_clean = current + step_list[step_idx].to(current.dtype)
|
||||
noise = torch.randn(shape, generator=generator)
|
||||
current = pred_clean + noise
|
||||
noise_latents.append(current.clone())
|
||||
|
||||
noisy_input: torch.Tensor
|
||||
if target_idx == 0:
|
||||
noisy_input = initial
|
||||
else:
|
||||
noisy_input = noise_latents[target_idx - 1]
|
||||
output = noisy_input + step_list[target_idx].to(noisy_input.dtype)
|
||||
return output, generator.get_state()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target_idx", range(4))
|
||||
def test_simulate_rollout_only_runs_required_prefix_forwards(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
target_idx: int,
|
||||
) -> None:
|
||||
method, student = _rollout_method()
|
||||
batch = SimpleNamespace(
|
||||
latents=torch.zeros((1, 2)),
|
||||
dmd_latent_vis_dict={},
|
||||
)
|
||||
|
||||
def _fixed_target(*args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
del args, kwargs
|
||||
return torch.tensor([target_idx], dtype=torch.long)
|
||||
|
||||
monkeypatch.setattr(torch, "randint", _fixed_target)
|
||||
method._student_rollout(batch, with_grad=True)
|
||||
|
||||
# One prediction per required prefix step, plus the differentiable target
|
||||
# prediction. Prefix noising only happens for the required prefix steps.
|
||||
assert len(student.predict_calls) == target_idx + 1
|
||||
assert len(student.add_noise_calls) == target_idx
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target_idx", range(4))
|
||||
def test_simulate_rollout_preserves_method_generator_progress(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
target_idx: int,
|
||||
) -> None:
|
||||
seed = 4321
|
||||
method, _ = _rollout_method(seed)
|
||||
shape = (1, 2)
|
||||
batch = SimpleNamespace(
|
||||
latents=torch.zeros(shape),
|
||||
dmd_latent_vis_dict={},
|
||||
)
|
||||
|
||||
def _fixed_target(*args: Any, **kwargs: Any) -> torch.Tensor:
|
||||
del args, kwargs
|
||||
return torch.tensor([target_idx], dtype=torch.long)
|
||||
|
||||
monkeypatch.setattr(torch, "randint", _fixed_target)
|
||||
output = method._student_rollout(batch, with_grad=False)
|
||||
|
||||
# The previous implementation drew the initial latent and one noise tensor
|
||||
# for each of the three possible prefix transitions. Keep consuming those
|
||||
# draws so subsequent DMD2 randomness remains aligned across the change.
|
||||
reference_output, reference_state = _legacy_full_rollout_reference(
|
||||
seed=seed,
|
||||
target_idx=target_idx,
|
||||
shape=shape,
|
||||
)
|
||||
assert torch.equal(output, reference_output)
|
||||
assert torch.equal(
|
||||
method.cuda_generator.get_state(),
|
||||
reference_state,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target_idx", range(4))
|
||||
def test_simulate_rollout_uses_global_max_prefix_without_changing_local_output(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
target_idx: int,
|
||||
) -> None:
|
||||
seed = 9876
|
||||
method, student = _rollout_method(seed)
|
||||
shape = (1, 2)
|
||||
batch = SimpleNamespace(
|
||||
latents=torch.zeros(shape),
|
||||
dmd_latent_vis_dict={},
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
torch,
|
||||
"randint",
|
||||
lambda *args, **kwargs: torch.tensor([target_idx], dtype=torch.long),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
method,
|
||||
"_max_rollout_target_idx_across_ranks",
|
||||
lambda sampled_idx: 3,
|
||||
)
|
||||
|
||||
output = method._student_rollout(batch, with_grad=False)
|
||||
reference_output, reference_state = _legacy_full_rollout_reference(
|
||||
seed=seed,
|
||||
target_idx=target_idx,
|
||||
shape=shape,
|
||||
)
|
||||
|
||||
# Every rank participates in the globally required three prefix forwards,
|
||||
# then evaluates its own target. The local result and method-owned RNG
|
||||
# sequence still match the legacy full rollout exactly.
|
||||
assert len(student.predict_calls) == 4
|
||||
assert len(student.add_noise_calls) == 3
|
||||
assert torch.equal(output, reference_output)
|
||||
assert torch.equal(method.cuda_generator.get_state(), reference_state)
|
||||
|
||||
|
||||
def test_dmd2_role_update_cadence_follows_selected_optimizers() -> None:
|
||||
method = object.__new__(DMD2Method)
|
||||
torch.nn.Module.__init__(method)
|
||||
method.method_config = {"generator_update_interval": 5}
|
||||
method._student_optimizer = object()
|
||||
method._critic_optimizer = object()
|
||||
|
||||
assert not method.did_update_role("student", iteration=4)
|
||||
assert method.did_update_role("student", iteration=5)
|
||||
assert method.did_update_role("critic", iteration=4)
|
||||
assert not method.did_update_role("teacher", iteration=5)
|
||||
@@ -0,0 +1,61 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""CPU-only tests for modular training profiler lifecycle ownership."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fastvideo.train.entrypoint.train import run_training_from_config
|
||||
|
||||
|
||||
class _RecordingProfiler:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.events: list[tuple[str, str] | tuple[str]] = []
|
||||
|
||||
@contextmanager
|
||||
def region(self, name: str):
|
||||
self.events.append(("enter", name))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.events.append(("exit", name))
|
||||
|
||||
def stop(self) -> None:
|
||||
self.events.append(("stop", ))
|
||||
|
||||
|
||||
def test_dry_run_profiles_model_build_and_flushes(monkeypatch) -> None:
|
||||
profiler = _RecordingProfiler()
|
||||
training = SimpleNamespace(
|
||||
distributed=SimpleNamespace(tp_size=1, sp_size=1),
|
||||
vsa_sparsity=0.0,
|
||||
model_path="model",
|
||||
)
|
||||
cfg = SimpleNamespace(training=training)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.train.utils.config.load_run_config",
|
||||
lambda *args, **kwargs: cfg,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.distributed.maybe_init_distributed_environment_and_model_parallel",
|
||||
lambda *args, **kwargs: None,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.train.utils.builder.build_from_config",
|
||||
lambda loaded_cfg: (loaded_cfg.training, object(), object(), 0),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.train.entrypoint.train.get_or_create_profiler",
|
||||
lambda trace_dir: profiler,
|
||||
)
|
||||
|
||||
run_training_from_config("unused.yaml", dry_run=True)
|
||||
|
||||
assert profiler.events == [
|
||||
("enter", "profiler_region_model_loading"),
|
||||
("exit", "profiler_region_model_loading"),
|
||||
("stop", ),
|
||||
]
|
||||
@@ -0,0 +1,137 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""CPU-only tests for modular Trainer profiler boundaries."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from fastvideo.profiler import list_profiler_regions
|
||||
from fastvideo.train.trainer import Trainer
|
||||
from fastvideo.train.utils.training_config import TrainingConfig
|
||||
|
||||
|
||||
class _RecordingProfiler:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.events: list[tuple[str, str]] = []
|
||||
|
||||
@property
|
||||
def has_profiler(self) -> bool:
|
||||
return True
|
||||
|
||||
@contextmanager
|
||||
def region(self, name: str):
|
||||
self.events.append(("enter", name))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
self.events.append(("exit", name))
|
||||
|
||||
|
||||
class _DummyTracker:
|
||||
|
||||
def log(self, metrics: dict[str, float], step: int) -> None:
|
||||
del metrics, step
|
||||
|
||||
def finish(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class _DummyMethod:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.weight = torch.nn.Parameter(torch.tensor(1.0))
|
||||
|
||||
def set_tracker(self, tracker: Any) -> None:
|
||||
del tracker
|
||||
|
||||
def on_train_start(self) -> None:
|
||||
pass
|
||||
|
||||
def manages_optimization(self) -> bool:
|
||||
return False
|
||||
|
||||
def single_train_step(
|
||||
self,
|
||||
batch: dict[str, Any],
|
||||
iteration: int,
|
||||
) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, float]]:
|
||||
del batch, iteration
|
||||
return {"total_loss": self.weight.square()}, {}, {}
|
||||
|
||||
def backward(
|
||||
self,
|
||||
loss_map: dict[str, torch.Tensor],
|
||||
outputs: dict[str, Any],
|
||||
*,
|
||||
grad_accum_rounds: int,
|
||||
) -> None:
|
||||
del outputs
|
||||
(loss_map["total_loss"] / grad_accum_rounds).backward()
|
||||
|
||||
def optimizers_schedulers_step(self, iteration: int) -> None:
|
||||
del iteration
|
||||
|
||||
def optimizers_zero_grad(self, iteration: int) -> None:
|
||||
del iteration
|
||||
self.weight.grad = None
|
||||
|
||||
|
||||
def test_modular_trainer_emits_nested_step_regions(monkeypatch) -> None:
|
||||
group = SimpleNamespace(rank=0, local_rank=0, rank_in_group=0, world_size=1)
|
||||
profiler = _RecordingProfiler()
|
||||
|
||||
monkeypatch.setattr("fastvideo.train.trainer.get_world_group", lambda: group)
|
||||
monkeypatch.setattr("fastvideo.train.trainer.get_sp_group", lambda: group)
|
||||
monkeypatch.setattr(
|
||||
"fastvideo.train.trainer.build_tracker",
|
||||
lambda *args, **kwargs: _DummyTracker(),
|
||||
)
|
||||
monkeypatch.setattr("fastvideo.profiler._GLOBAL_CONTROLLER", profiler)
|
||||
|
||||
trainer = Trainer(TrainingConfig())
|
||||
trainer.run(
|
||||
_DummyMethod(),
|
||||
dataloader=[{}],
|
||||
max_steps=1,
|
||||
)
|
||||
|
||||
assert profiler.events == [
|
||||
("enter", "profiler_region_training_train"),
|
||||
("enter", "profiler_region_training_train_one_step"),
|
||||
("enter", "profiler_region_training_dataloader"),
|
||||
("exit", "profiler_region_training_dataloader"),
|
||||
("enter", "profiler_region_training_forward"),
|
||||
("exit", "profiler_region_training_forward"),
|
||||
("enter", "profiler_region_training_backward"),
|
||||
("exit", "profiler_region_training_backward"),
|
||||
("enter", "profiler_region_training_optimizer"),
|
||||
("exit", "profiler_region_training_optimizer"),
|
||||
("enter", "profiler_region_training_callbacks"),
|
||||
("exit", "profiler_region_training_callbacks"),
|
||||
("exit", "profiler_region_training_train_one_step"),
|
||||
("exit", "profiler_region_training_train"),
|
||||
]
|
||||
|
||||
|
||||
def test_modular_training_regions_are_registered() -> None:
|
||||
names = {region.name for region in list_profiler_regions()}
|
||||
assert {
|
||||
"profiler_region_model_loading",
|
||||
"profiler_region_training_train",
|
||||
"profiler_region_training_train_one_step",
|
||||
"profiler_region_training_dataloader",
|
||||
"profiler_region_training_forward",
|
||||
"profiler_region_training_backward",
|
||||
"profiler_region_training_optimizer",
|
||||
"profiler_region_training_callbacks",
|
||||
"profiler_region_training_save_checkpoint",
|
||||
"profiler_region_training_validation",
|
||||
"profiler_region_dmd2_student_rollout",
|
||||
"profiler_region_dmd2_generator_loss",
|
||||
"profiler_region_dmd2_critic_loss",
|
||||
} <= names
|
||||
@@ -136,4 +136,5 @@ def test_trainer_runs_validation_callback_during_training(monkeypatch, ) -> None
|
||||
assert method.zero_grad_steps == [0, 1, 2, 3]
|
||||
assert method.optimizer_steps == [1, 2, 3]
|
||||
assert [step for _, step in tracker.logs] == [1, 2, 3]
|
||||
assert all("dataloader_time_sec" in metrics for metrics, _ in tracker.logs)
|
||||
assert tracker.finished is True
|
||||
|
||||
@@ -94,6 +94,8 @@ class EMACallback(Callback):
|
||||
|
||||
if iteration < self._start_iter:
|
||||
return
|
||||
if not method.did_update_role("student", iteration):
|
||||
return
|
||||
if not self._ema_started:
|
||||
logger.info(
|
||||
"Starting EMA updates at iteration %d "
|
||||
|
||||
@@ -31,6 +31,7 @@ from fastvideo.distributed import (
|
||||
get_world_group,
|
||||
)
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.profiler import profile_region
|
||||
from fastvideo.pipelines import ForwardBatch
|
||||
from fastvideo.train.callbacks.callback import Callback
|
||||
from fastvideo.train.utils.instantiate import resolve_target
|
||||
@@ -285,6 +286,7 @@ class ValidationCallback(Callback):
|
||||
# Core validation logic
|
||||
# ----------------------------------------------------------
|
||||
|
||||
@profile_region("profiler_region_training_validation")
|
||||
def _run_validation(
|
||||
self,
|
||||
method: TrainingMethod,
|
||||
|
||||
@@ -24,7 +24,9 @@ from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
import fastvideo.envs as envs
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.profiler import get_or_create_profiler
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -74,47 +76,55 @@ def run_training_from_config(
|
||||
tc.distributed.sp_size,
|
||||
)
|
||||
|
||||
_, method, dataloader, start_step = build_from_config(cfg)
|
||||
profiler_controller = get_or_create_profiler(envs.FASTVIDEO_TORCH_PROFILER_DIR, )
|
||||
try:
|
||||
with profiler_controller.region("profiler_region_model_loading"):
|
||||
_, method, dataloader, start_step = build_from_config(cfg)
|
||||
|
||||
if dry_run:
|
||||
logger.info("Dry-run: config parsed and "
|
||||
"build_from_config succeeded.")
|
||||
return
|
||||
if dry_run:
|
||||
logger.info("Dry-run: config parsed and "
|
||||
"build_from_config succeeded.")
|
||||
return
|
||||
|
||||
trainer = Trainer(
|
||||
tc,
|
||||
config=cfg.resolved_config(),
|
||||
callback_configs=cfg.callbacks,
|
||||
)
|
||||
trainer = Trainer(
|
||||
tc,
|
||||
config=cfg.resolved_config(),
|
||||
callback_configs=cfg.callbacks,
|
||||
)
|
||||
|
||||
# Attach the exact YAML used for this run to the
|
||||
# tracker (e.g., W&B Files).
|
||||
trainer.tracker.log_file(
|
||||
os.path.abspath(os.path.expanduser(config_path)),
|
||||
name="run.yaml",
|
||||
)
|
||||
# Attach the exact YAML used for this run to the
|
||||
# tracker (e.g., W&B Files).
|
||||
trainer.tracker.log_file(
|
||||
os.path.abspath(os.path.expanduser(config_path)),
|
||||
name="run.yaml",
|
||||
)
|
||||
|
||||
ckpt_config = CheckpointConfig(
|
||||
save_steps=int(tc.checkpoint.training_state_checkpointing_steps or 0),
|
||||
keep_last=int(tc.checkpoint.checkpoints_total_limit or 0),
|
||||
)
|
||||
ckpt_config = CheckpointConfig(
|
||||
save_steps=int(tc.checkpoint.training_state_checkpointing_steps or 0),
|
||||
keep_last=int(tc.checkpoint.checkpoints_total_limit or 0),
|
||||
)
|
||||
|
||||
checkpoint_manager = CheckpointManager(
|
||||
method=method,
|
||||
dataloader=dataloader,
|
||||
output_dir=tc.checkpoint.output_dir,
|
||||
config=ckpt_config,
|
||||
callbacks=trainer.callbacks,
|
||||
raw_config=cfg.raw,
|
||||
)
|
||||
checkpoint_manager = CheckpointManager(
|
||||
method=method,
|
||||
dataloader=dataloader,
|
||||
output_dir=tc.checkpoint.output_dir,
|
||||
config=ckpt_config,
|
||||
callbacks=trainer.callbacks,
|
||||
raw_config=cfg.raw,
|
||||
)
|
||||
|
||||
trainer.run(
|
||||
method,
|
||||
dataloader=dataloader,
|
||||
max_steps=tc.loop.max_train_steps,
|
||||
start_step=start_step,
|
||||
checkpoint_manager=checkpoint_manager,
|
||||
)
|
||||
trainer.run(
|
||||
method,
|
||||
dataloader=dataloader,
|
||||
max_steps=tc.loop.max_train_steps,
|
||||
start_step=start_step,
|
||||
checkpoint_manager=checkpoint_manager,
|
||||
)
|
||||
finally:
|
||||
# torch.profiler exports its trace at stop(). Keep atexit as a fallback
|
||||
# for abrupt exits, but flush here so normal training returns only after
|
||||
# all per-rank traces and summaries are complete.
|
||||
profiler_controller.stop()
|
||||
|
||||
|
||||
def main(
|
||||
|
||||
@@ -208,6 +208,23 @@ class TrainingMethod(torch.nn.Module, ABC):
|
||||
"""
|
||||
return False
|
||||
|
||||
def did_update_role(
|
||||
self,
|
||||
role: str,
|
||||
iteration: int,
|
||||
) -> bool:
|
||||
"""Whether ``role`` completed an optimizer update this iteration.
|
||||
|
||||
Callbacks run after optimization and use this hook to follow the
|
||||
method's actual update cadence. The default implementation matches
|
||||
the role's optimizer against the optimizers selected for this
|
||||
iteration; methods that step optimizers internally can override it.
|
||||
"""
|
||||
role_optimizer = self._optimizer_dict.get(role)
|
||||
if role_optimizer is None:
|
||||
return False
|
||||
return any(optimizer is role_optimizer for optimizer in self.get_optimizers(iteration))
|
||||
|
||||
def managed_train_step(
|
||||
self,
|
||||
data_stream: Any,
|
||||
|
||||
@@ -6,8 +6,10 @@ from __future__ import annotations
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
|
||||
from fastvideo.profiler import profile_region
|
||||
from fastvideo.train.methods.base import TrainingMethod, LogScalar
|
||||
from fastvideo.train.models.base import ModelBase
|
||||
from fastvideo.train.utils.optimizer import (
|
||||
@@ -431,6 +433,23 @@ class DMD2Method(TrainingMethod):
|
||||
)
|
||||
return step_list[index]
|
||||
|
||||
@staticmethod
|
||||
def _max_rollout_target_idx_across_ranks(target_timestep_idx: torch.Tensor, ) -> int:
|
||||
"""Return the largest sampled rollout index across all data ranks.
|
||||
|
||||
Each rank intentionally samples its own DMD2 timestep, but FSDP ranks
|
||||
must execute the same number of transformer forwards so their
|
||||
collectives stay ordered. The largest local target is therefore the
|
||||
minimum prefix length every rank must execute this iteration.
|
||||
"""
|
||||
max_target_timestep_idx = target_timestep_idx.detach().clone()
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
dist.all_reduce(
|
||||
max_target_timestep_idx,
|
||||
op=dist.ReduceOp.MAX,
|
||||
)
|
||||
return int(max_target_timestep_idx.item())
|
||||
|
||||
def _parse_score_timestep_bounds(self) -> tuple[int, int]:
|
||||
"""Resolve the score-model timestep window used by legacy DMD.
|
||||
|
||||
@@ -476,6 +495,7 @@ class DMD2Method(TrainingMethod):
|
||||
self._score_max_timestep,
|
||||
)
|
||||
|
||||
@profile_region("profiler_region_dmd2_student_rollout")
|
||||
def _student_rollout(
|
||||
self,
|
||||
batch: Any,
|
||||
@@ -516,6 +536,7 @@ class DMD2Method(TrainingMethod):
|
||||
generator=self.cuda_generator,
|
||||
)
|
||||
target_timestep_idx_int = int(target_timestep_idx.item())
|
||||
synchronized_target_idx = self._max_rollout_target_idx_across_ranks(target_timestep_idx, )
|
||||
target_timestep = step_list[target_timestep_idx]
|
||||
|
||||
current_noise_latents = torch.randn(
|
||||
@@ -524,56 +545,62 @@ class DMD2Method(TrainingMethod):
|
||||
dtype=dtype,
|
||||
generator=self.cuda_generator,
|
||||
)
|
||||
current_noise_latents_copy = (current_noise_latents.clone())
|
||||
|
||||
max_target_idx = len(step_list) - 1
|
||||
noise_latents: list[torch.Tensor] = []
|
||||
noise_latent_index = target_timestep_idx_int - 1
|
||||
noisy_input = current_noise_latents
|
||||
|
||||
if max_target_idx > 0:
|
||||
with torch.no_grad():
|
||||
for step_idx in range(max_target_idx):
|
||||
current_timestep = step_list[step_idx]
|
||||
current_timestep_tensor = (current_timestep * torch.ones(
|
||||
1,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
))
|
||||
# FSDP ranks must run the same number of forwards. Ranks
|
||||
# whose local target is earlier keep advancing a throwaway
|
||||
# trajectory until the largest target sampled globally.
|
||||
needs_prefix_step = step_idx < synchronized_target_idx
|
||||
pred_clean: torch.Tensor | None = None
|
||||
noise_dtype = dtype
|
||||
if needs_prefix_step:
|
||||
current_timestep = step_list[step_idx]
|
||||
current_timestep_tensor = (current_timestep * torch.ones(
|
||||
1,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
))
|
||||
|
||||
pred_clean = self.student.predict_x0(
|
||||
current_noise_latents,
|
||||
current_timestep_tensor,
|
||||
batch,
|
||||
conditional=True,
|
||||
cfg_uncond=self._cfg_uncond,
|
||||
attn_kind="vsa",
|
||||
)
|
||||
pred_clean = self.student.predict_x0(
|
||||
current_noise_latents,
|
||||
current_timestep_tensor,
|
||||
batch,
|
||||
conditional=True,
|
||||
cfg_uncond=self._cfg_uncond,
|
||||
attn_kind="vsa",
|
||||
)
|
||||
noise_dtype = pred_clean.dtype
|
||||
|
||||
next_timestep = step_list[step_idx + 1]
|
||||
next_timestep_tensor = (next_timestep * torch.ones(
|
||||
1,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
))
|
||||
# Preserve the method-owned generator sequence even when
|
||||
# the corresponding prefix forward is unnecessary. This
|
||||
# keeps all later DMD2 random draws aligned with the old
|
||||
# full-rollout implementation.
|
||||
noise = torch.randn(
|
||||
latents.shape,
|
||||
device=device,
|
||||
dtype=pred_clean.dtype,
|
||||
dtype=noise_dtype,
|
||||
generator=self.cuda_generator,
|
||||
)
|
||||
current_noise_latents = (self.student.add_noise(
|
||||
pred_clean,
|
||||
noise,
|
||||
next_timestep_tensor,
|
||||
))
|
||||
noise_latents.append(current_noise_latents.clone())
|
||||
|
||||
if noise_latent_index >= 0:
|
||||
if noise_latent_index >= len(noise_latents):
|
||||
raise RuntimeError("noise_latent_index is out of bounds")
|
||||
noisy_input = noise_latents[noise_latent_index]
|
||||
else:
|
||||
noisy_input = current_noise_latents_copy
|
||||
if needs_prefix_step:
|
||||
assert pred_clean is not None
|
||||
next_timestep = step_list[step_idx + 1]
|
||||
next_timestep_tensor = (next_timestep * torch.ones(
|
||||
1,
|
||||
device=device,
|
||||
dtype=torch.long,
|
||||
))
|
||||
current_noise_latents = (self.student.add_noise(
|
||||
pred_clean,
|
||||
noise,
|
||||
next_timestep_tensor,
|
||||
))
|
||||
if step_idx + 1 == target_timestep_idx_int:
|
||||
noisy_input = current_noise_latents
|
||||
|
||||
if with_grad:
|
||||
pred_x0 = self.student.predict_x0(
|
||||
@@ -598,6 +625,7 @@ class DMD2Method(TrainingMethod):
|
||||
batch.dmd_latent_vis_dict["generator_timestep"] = target_timestep.float().detach()
|
||||
return pred_x0
|
||||
|
||||
@profile_region("profiler_region_dmd2_critic_loss")
|
||||
def _critic_flow_matching_loss(
|
||||
self,
|
||||
batch: Any,
|
||||
@@ -638,6 +666,7 @@ class DMD2Method(TrainingMethod):
|
||||
outputs,
|
||||
)
|
||||
|
||||
@profile_region("profiler_region_dmd2_generator_loss")
|
||||
def _dmd_loss(
|
||||
self,
|
||||
generator_pred_x0: torch.Tensor,
|
||||
|
||||
+113
-77
@@ -11,6 +11,7 @@ import torch
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from fastvideo.distributed import get_sp_group, get_world_group
|
||||
from fastvideo.profiler import profile_region, profiler_region
|
||||
from fastvideo.train.callbacks.callback import CallbackDict
|
||||
from fastvideo.train.methods.base import LogScalar, TrainingMethod
|
||||
from fastvideo.train.utils.tracking import build_tracker
|
||||
@@ -98,6 +99,113 @@ class Trainer:
|
||||
if self.global_rank == 0 and validation_metrics:
|
||||
self.tracker.log(validation_metrics, iteration)
|
||||
|
||||
@profile_region("profiler_region_training_train_one_step")
|
||||
def _run_train_step(
|
||||
self,
|
||||
method: TrainingMethod,
|
||||
*,
|
||||
data_stream: Iterator[dict[str, Any]],
|
||||
step: int,
|
||||
grad_accum: int,
|
||||
method_manages_optimization: bool,
|
||||
) -> None:
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# Accumulate on GPU during grad-accum; materialise to CPU once per
|
||||
# step right before logging.
|
||||
loss_sums: dict[str, float | torch.Tensor] = {}
|
||||
metric_sums: dict[str, float | torch.Tensor] = {}
|
||||
dataloader_time_sec = 0.0
|
||||
if method_manages_optimization:
|
||||
# Managed methods own their forward/backward/optimizer boundaries,
|
||||
# so the enclosing training_train_one_step region is the truthful
|
||||
# granularity available to the trainer.
|
||||
loss_map, outputs, step_metrics = method.managed_train_step(
|
||||
data_stream,
|
||||
step,
|
||||
)
|
||||
for k, v in loss_map.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
loss_sums[k] = v.detach()
|
||||
for k, v in step_metrics.items():
|
||||
if k in loss_sums:
|
||||
raise ValueError(f"Metric key {k!r} collides "
|
||||
"with loss key. Use a "
|
||||
"different name (e.g. prefix "
|
||||
"with 'train/').")
|
||||
metric_sums[k] = _coerce_log_scalar(
|
||||
v,
|
||||
where=("method.managed_train_step()"
|
||||
f".metrics[{k!r}]"),
|
||||
)
|
||||
else:
|
||||
for _ in range(grad_accum):
|
||||
dataloader_t0 = time.perf_counter()
|
||||
with profiler_region("profiler_region_training_dataloader"):
|
||||
batch = next(data_stream)
|
||||
dataloader_time_sec += (time.perf_counter() - dataloader_t0)
|
||||
with profiler_region("profiler_region_training_forward"):
|
||||
loss_map, outputs, step_metrics = (method.single_train_step(
|
||||
batch,
|
||||
step,
|
||||
))
|
||||
|
||||
with profiler_region("profiler_region_training_backward"):
|
||||
method.backward(
|
||||
loss_map,
|
||||
outputs,
|
||||
grad_accum_rounds=grad_accum,
|
||||
)
|
||||
|
||||
for k, v in loss_map.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
prev = loss_sums.get(k, 0.0)
|
||||
loss_sums[k] = prev + v.detach()
|
||||
for k, v in step_metrics.items():
|
||||
if k in loss_sums:
|
||||
raise ValueError(f"Metric key {k!r} collides "
|
||||
"with loss key. Use a "
|
||||
"different name (e.g. prefix "
|
||||
"with 'train/').")
|
||||
prev = metric_sums.get(k, 0.0)
|
||||
metric_sums[k] = (prev + _coerce_log_scalar(
|
||||
v,
|
||||
where=("method.single_train_step()"
|
||||
f".metrics[{k!r}]"),
|
||||
))
|
||||
|
||||
if not method_manages_optimization:
|
||||
with profiler_region("profiler_region_training_optimizer"):
|
||||
self.callbacks.on_before_optimizer_step(
|
||||
method,
|
||||
iteration=step,
|
||||
)
|
||||
method.optimizers_schedulers_step(step)
|
||||
method.optimizers_zero_grad(step)
|
||||
|
||||
# Single CPU sync point: materialise GPU tensors to float right before
|
||||
# logging.
|
||||
divisor = 1 if method_manages_optimization else grad_accum
|
||||
metrics = {k: float(v) / divisor for k, v in loss_sums.items()}
|
||||
metrics.update({k: float(v) / divisor for k, v in metric_sums.items()})
|
||||
metrics["step_time_sec"] = (time.perf_counter() - t0)
|
||||
if not method_manages_optimization:
|
||||
# This is the local training process's wait for next(data_stream).
|
||||
# Track it without a cross-rank reduction to avoid adding a
|
||||
# synchronization to every training step.
|
||||
metrics["dataloader_time_sec"] = dataloader_time_sec
|
||||
metrics["vsa_sparsity"] = float(self.training_config.vsa_sparsity)
|
||||
if self.global_rank == 0 and metrics:
|
||||
self.tracker.log(metrics, step)
|
||||
|
||||
with profiler_region("profiler_region_training_callbacks"):
|
||||
self.callbacks.on_training_step_end(
|
||||
method,
|
||||
metrics,
|
||||
iteration=step,
|
||||
)
|
||||
|
||||
@profile_region("profiler_region_training_train")
|
||||
def run(
|
||||
self,
|
||||
method: TrainingMethod,
|
||||
@@ -150,84 +258,12 @@ class Trainer:
|
||||
# Allow method-specific optimization flow (e.g. DiffusionNFT).
|
||||
method_manages_optimization = bool(method.manages_optimization())
|
||||
for step in progress:
|
||||
t0 = time.perf_counter()
|
||||
|
||||
# Accumulate on GPU during grad-accum; materialise
|
||||
# to CPU once per step right before logging.
|
||||
loss_sums: dict[str, float | torch.Tensor] = {}
|
||||
metric_sums: dict[str, float | torch.Tensor] = {}
|
||||
if method_manages_optimization:
|
||||
loss_map, outputs, step_metrics = method.managed_train_step(
|
||||
data_stream,
|
||||
step,
|
||||
)
|
||||
for k, v in loss_map.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
loss_sums[k] = v.detach()
|
||||
for k, v in step_metrics.items():
|
||||
if k in loss_sums:
|
||||
raise ValueError(f"Metric key {k!r} collides "
|
||||
"with loss key. Use a "
|
||||
"different name (e.g. prefix "
|
||||
"with 'train/').")
|
||||
metric_sums[k] = _coerce_log_scalar(
|
||||
v,
|
||||
where=("method.managed_train_step()"
|
||||
f".metrics[{k!r}]"),
|
||||
)
|
||||
else:
|
||||
for accum_iter in range(grad_accum):
|
||||
batch = next(data_stream)
|
||||
loss_map, outputs, step_metrics = (method.single_train_step(
|
||||
batch,
|
||||
step,
|
||||
))
|
||||
|
||||
method.backward(
|
||||
loss_map,
|
||||
outputs,
|
||||
grad_accum_rounds=grad_accum,
|
||||
)
|
||||
|
||||
for k, v in loss_map.items():
|
||||
if isinstance(v, torch.Tensor):
|
||||
prev = loss_sums.get(k, 0.0)
|
||||
loss_sums[k] = prev + v.detach()
|
||||
for k, v in step_metrics.items():
|
||||
if k in loss_sums:
|
||||
raise ValueError(f"Metric key {k!r} collides "
|
||||
"with loss key. Use a "
|
||||
"different name (e.g. prefix "
|
||||
"with 'train/').")
|
||||
prev = metric_sums.get(k, 0.0)
|
||||
metric_sums[k] = (prev + _coerce_log_scalar(
|
||||
v,
|
||||
where=("method.single_train_step()"
|
||||
f".metrics[{k!r}]"),
|
||||
))
|
||||
|
||||
if not method_manages_optimization:
|
||||
self.callbacks.on_before_optimizer_step(
|
||||
method,
|
||||
iteration=step,
|
||||
)
|
||||
method.optimizers_schedulers_step(step)
|
||||
method.optimizers_zero_grad(step)
|
||||
|
||||
# Single CPU sync point: materialise GPU tensors
|
||||
# to float right before logging.
|
||||
divisor = 1 if method_manages_optimization else grad_accum
|
||||
metrics = {k: float(v) / divisor for k, v in loss_sums.items()}
|
||||
metrics.update({k: float(v) / divisor for k, v in metric_sums.items()})
|
||||
metrics["step_time_sec"] = (time.perf_counter() - t0)
|
||||
metrics["vsa_sparsity"] = float(tc.vsa_sparsity)
|
||||
if self.global_rank == 0 and metrics:
|
||||
self.tracker.log(metrics, step)
|
||||
|
||||
self.callbacks.on_training_step_end(
|
||||
self._run_train_step(
|
||||
method,
|
||||
metrics,
|
||||
iteration=step,
|
||||
data_stream=data_stream,
|
||||
step=step,
|
||||
grad_accum=grad_accum,
|
||||
method_manages_optimization=method_manages_optimization,
|
||||
)
|
||||
|
||||
if checkpoint_manager is not None:
|
||||
|
||||
@@ -23,6 +23,7 @@ from torch.distributed.checkpoint.state_dict import (
|
||||
from torch.distributed.checkpoint.stateful import Stateful
|
||||
|
||||
from fastvideo.logger import init_logger
|
||||
from fastvideo.profiler import profile_region
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -246,6 +247,7 @@ class CheckpointManager:
|
||||
return
|
||||
self.save(step)
|
||||
|
||||
@profile_region("profiler_region_training_save_checkpoint")
|
||||
def save(self, step: int) -> None:
|
||||
checkpoint_dir = self._checkpoint_dir(step)
|
||||
dcp_dir = self._dcp_dir(step)
|
||||
|
||||
+3
-5
@@ -122,11 +122,9 @@ fastvideo-kernel = [
|
||||
# uv reads UV_TORCH_BACKEND / --torch-backend (not a pyproject setting). The
|
||||
# supported backends (cu126, cu130) both provide torch 2.12.0.
|
||||
imagebind = { git = "https://github.com/facebookresearch/ImageBind.git", rev = "53680b02d7e37b19b124fa37bae4b6c98c38f5be" }
|
||||
# FA4 cute. This revision pins nvidia-cutlass-dsl==4.6.0.dev0; holding cutlass-dsl
|
||||
# back at 4.5.x breaks the CuTe kernels at JIT time, not at import
|
||||
# (see fastvideo-kernel/README.md).
|
||||
# torch.compile support comes from FastVideo's own custom_op wrappers.
|
||||
flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "14c377950125c70b7a9dabf9c561fca53715ac7d", subdirectory = "flash_attn/cute" }
|
||||
# FA4 cute, pinned to a cutlass-4.5-compatible revision. torch.compile support
|
||||
# comes from FastVideo's own custom_op wrappers.
|
||||
flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = "82d6441eec5d4dfec120153db2c0145ae855a083", subdirectory = "flash_attn/cute" }
|
||||
|
||||
[project.optional-dependencies]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user