From f496e4b838d1b7d5eacd74afe49184d56efe523f Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Tue, 10 Nov 2026 02:07:14 +0000 Subject: [PATCH] [misc]: extend the comment sweep to every file --- .../scripts/dreamverse-deploy.sh | 19 --- .buildkite/scripts/pr_test.sh | 2 - .github/mergify.yml | 8 - .pre-commit-config.yaml | 6 +- apps/dreamverse/dreamverse/gpu_pool.py | 2 +- .../dreamverse/session/controller.py | 1 - .../scripts/install_native_ffmpeg.sh | 5 - apps/dreamverse/scripts/launch/launch_demo.sh | 2 - apps/dreamverse/web/src/app/globals.css | 5 - apps/dreamverse/web/src/app/page.tsx | 1 - .../web/src/lib/media/avPipeline.test.ts | 5 - .../web/src/lib/media/avPipeline.ts | 1 - apps/fastvideo_studio/e2e/datasets.spec.ts | 1 - apps/fastvideo_studio/job_runner.py | 4 +- apps/fastvideo_studio/src/app/globals.css | 2 - .../datasets/DatasetSidebar.test.tsx | 3 - .../components/jobs/CreateJobModal.test.tsx | 1 - .../src/components/jobs/CreateJobModal.tsx | 1 - .../jobs/JobDetailsSidebar.test.tsx | 1 - .../src/components/jobs/JobQueue.tsx | 2 - comfyui/web/js/FastVideo.core.js | 18 +-- docs/cookbook/recipes/fastwan21.yaml | 3 - docs/cookbook/recipes/wan21-i2v.yaml | 3 - docs/generate_examples.py | 38 +---- examples/distill/MatrixGame2.0/distill_dmd.sh | 11 +- .../distill/MatrixGame2.0/distill_dmd.slurm | 1 - .../SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh | 4 +- examples/distill/SFWan2.2-A14B/distill_dmd.sh | 11 -- .../distill_dmd_VSA_t2v_1.3B.slurm | 3 - .../distill_dmd_VSA_t2v_14B.slurm | 3 - .../distill_dmd_t2v_1.3B.slurm | 3 - .../Data-free/distill_dmd_t2v_5B.sh | 3 - .../Data-free/distill_dmd_t2v_5B_VSA.sh | 3 - .../crush_smol/distill_dmd_VSA_t2v_5B.sh | 3 - .../crush_smol/distill_dmd_VSA_t2v_5B_lora.sh | 3 - .../basic/basic_matrixgame2_streaming.py | 1 - examples/inference/basic/basic_wan2_2.py | 1 - examples/inference/eval/all_metrics_demo.py | 2 +- .../gradio/serving/ray_serve_backend.py | 4 +- .../wan/self_forcing_causal_t2v_vsa.yaml | 1 - examples/train/configs/example.yaml | 12 -- .../wan/causal_t2v.yaml | 2 +- .../knowledge_distillation/wan/t2v.yaml | 2 +- examples/train/run.sh | 3 - examples/train/run_slurm.sh | 2 - .../scenario/qad_wan2_1_mixkit/run_stage1.sh | 1 - .../worldmodel/solaris/finetune_i2v.yaml | 3 - .../zelda/self_forcing_causal_i2v.yaml | 1 - .../streaming_long_tuning_causal_i2v.yaml | 1 - .../causal_ode_init/finetune_ode_init.sh | 1 - .../finetune_ode_init_16k_vidprom.slurm | 1 - .../finetune/MatrixGame2.0/finetune_df.sh | 4 - .../finetune/MatrixGame2.0/finetune_i2v.sh | 3 - .../preprocess_matrixgame_data_i2v.sh | 1 - .../crush_smol/finetune_i2v.sh | 4 - .../crush_smol/finetune_i2v.slurm | 3 - .../Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm | 3 - .../Wan-Syn-Data/T2V-1.3B-VSA.slurm | 3 - .../Wan-Syn-Data/T2V-1.3B-VSA_4_gpu_debug.sh | 13 +- .../Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm | 3 - .../finetune/cosmos2_5/finetune_t2w.sh | 1 - .../finetune/cosmos2_5/finetune_t2w_lora.sh | 1 - .../run_train_vggsound_ddp_from_scratch.sh | 2 - .../run_train_vggsound_from_scratch.sh | 2 - .../crush_smol/finetune_i2v.sh | 4 - .../crush_smol/finetune_i2v.slurm | 3 - .../crush_smol/finetune_i2v_lora.sh | 4 - .../wan_t2v_1.3B/crush_smol/finetune_t2v.sh | 2 - .../crush_smol/finetune_t2v.slurm | 3 - .../crush_smol/finetune_t2v_lora.sh | 2 - .../crush_smol/profile_finetune_t2v.sh | 2 - .../attn_qat_infer/blackwell/api.cu | 19 --- .../attn_qat_infer/blackwell/block_info.h | 2 - .../blackwell/epilogue_tma_ws.h | 4 - .../attn_qat_infer/blackwell/kernel_ws.h | 5 +- .../attn_qat_infer/blackwell/launch.h | 1 - .../blackwell/mainloop_tma_ws.h | 1 - .../attn_qat_infer/blackwell/tile_scheduler.h | 1 - .../attn_qat_infer/blackwell/utils.h | 2 - .../quantization/bench/bench_quant_v.py | 1 - .../quantization/bench/bench_utils.py | 4 - .../quantization/fp4_quantization_4d.cu | 11 -- fastvideo-kernel/build.sh | 2 - .../csrc/attention/block_sparse_h100.cu | 33 +--- .../csrc/attention/st_attn_h100.cu | 5 - .../csrc/comm/ulysses_all_to_all.cu | 1 - fastvideo-kernel/csrc/common_extension.cpp | 1 - .../csrc/turbodiffusion/gemm/kernel.hpp | 6 - .../csrc/turbodiffusion/norm/layernorm.hpp | 4 - .../csrc/turbodiffusion/norm/rmsnorm.hpp | 3 - .../csrc/turbodiffusion/quant/quant.hpp | 1 - .../triton_kernels/attn_qat_train.py | 27 ---- .../block_sparse_attn_triton.py | 10 -- .../triton_kernels/nvfp4_utils.py | 2 - .../triton_kernels/st_attn_triton.py | 8 - .../fastvideo_kernel/turbodiffusion_ops.py | 44 ------ .../python/fastvideo_kernel/vmoba.py | 20 +-- fastvideo-kernel/tests/test_attn_qat_train.py | 143 ------------------ fastvideo-kernel/tests/test_sta.py | 2 - fastvideo-kernel/tests/test_turbodiffusion.py | 12 -- .../tests/test_vmoba_correctness.py | 3 - fastvideo-kernel/tests/test_vsa.py | 6 - fastvideo-kernel/tests/test_vsa_forward.py | 3 +- fastvideo/dataset/preprocessing_datasets.py | 2 +- .../device_communicators/pynccl_wrapper.py | 8 - fastvideo/entrypoints/video_generator.py | 2 +- fastvideo/layers/fp4linear.py | 1 - fastvideo/mlx_runtime/refine.py | 1 - fastvideo/models/dits/cosmos2_5.py | 4 +- fastvideo/models/dits/hyworld/pose.py | 2 +- .../models/dits/hyworld/retrieval_context.py | 2 +- fastvideo/models/dits/hyworld/trajectory.py | 4 - .../models/dits/lingbotworld/cam_utils.py | 2 +- .../models/dits/lingbotworld2/cam_utils.py | 2 +- .../models/dits/matrixgame2/action_module.py | 2 +- .../scheduling_flow_match_euler_discrete.py | 3 - .../scheduling_flow_unipc_multistep.py | 1 - .../schedulers/scheduling_unipc_multistep.py | 2 - fastvideo/models/vaes/hyworldvae.py | 1 - .../basic/ltx2/stages/ltx2_text_encoding.py | 1 - ...ame2_preprocess_pipeline_ode_trajectory.py | 8 +- .../preprocess/preprocess_pipeline_base.py | 6 +- .../preprocess_pipeline_ode_trajectory.py | 6 +- .../preprocess/preprocess_pipeline_text.py | 4 +- .../pipelines/preprocess/v1_preprocess.py | 4 +- fastvideo/pipelines/stages/denoising.py | 4 +- fastvideo/pipelines/stages/encoding.py | 2 +- .../pipelines/stages/hyworld_denoising.py | 2 +- fastvideo/pipelines/stages/image_encoding.py | 2 +- .../pipelines/stages/matrixgame3_denoising.py | 2 +- fastvideo/pipelines/stages/sr_denoising.py | 4 +- fastvideo/tests/encoders/test_hyt5_encoder.py | 1 - .../entrypoints/streaming/test_server.py | 1 - .../tests/hooks/test_layerwise_offload.py | 1 - .../lora/test_lora_inference_similarity.py | 2 +- .../tests/transformers/test_cosmos2_5.py | 2 +- .../vaes/test_minimax_h3_nvfp4_decoder.py | 2 +- fastvideo/training/distillation_pipeline.py | 2 +- fastvideo/training/training_pipeline.py | 4 +- fastvideo/training/training_utils.py | 6 +- fastvideo/worker/ray_distributed_executor.py | 1 - .../convert_gen3c_to_fastvideo.py | 3 +- .../checkpoint_conversion/create_hf_repo.py | 2 +- .../checkpoint_conversion/wan_to_diffusers.py | 2 +- scripts/finetune/finetune_v1.sh | 3 - scripts/preprocess/launch_slurm_text.sh | 2 - scripts/preprocess/syn_text.slurm | 5 - .../preprocess/v1_preprocess_wan_data_i2v.sh | 1 - .../preprocess/v1_preprocess_wan_data_t2v.sh | 1 - tests/local_tests/gen3c/test_gen3c.py | 11 -- .../helpers/magi_human_upstream.py | 1 - .../ltx2_5/test_ltx2_5_conversion.py | 1 - .../ltx2_5/test_ltx2_5_registry.py | 2 - .../local_tests/ltx2_5/test_ltx2_5_sampler.py | 5 - tests/test_bsa.py | 2 - 155 files changed, 65 insertions(+), 752 deletions(-) diff --git a/.agents/skills/dreamverse-deploy/scripts/dreamverse-deploy.sh b/.agents/skills/dreamverse-deploy/scripts/dreamverse-deploy.sh index e16189b8f..0ced8fe1b 100755 --- a/.agents/skills/dreamverse-deploy/scripts/dreamverse-deploy.sh +++ b/.agents/skills/dreamverse-deploy/scripts/dreamverse-deploy.sh @@ -77,9 +77,6 @@ if [[ "${1:-}" == "--stop" ]]; then exit 0 fi -# --------------------------------------------------------------------------- -# Args -# --------------------------------------------------------------------------- usage() { cat <&2; exit 3; } @@ -234,9 +228,6 @@ echo " python: ${CONDA_ENV_PYTHON}" mkdir -p "${LOG_DIR}" -# --------------------------------------------------------------------------- -# Teardown anything on target ports -# --------------------------------------------------------------------------- echo "[1/8] killing any existing deploy on ports ${BACKEND_PORT}/${FRONTEND_PORT} and GPU ${GPU}..." @@ -281,9 +272,6 @@ done gpu_mem="$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | sed -n "$((GPU + 1))p" || echo 0)" echo " ports cleared; GPU${GPU} at ${gpu_mem} MiB" -# --------------------------------------------------------------------------- -# Launch backend -# --------------------------------------------------------------------------- echo "[2/8] launching backend on GPU ${GPU} port ${BACKEND_PORT} (warmup=${WARMUP} torch_compile=${TORCH_COMPILE} nvenc=${NVENC})..." @@ -367,9 +355,6 @@ fi echo "[4/8] backend /readyz OK" -# --------------------------------------------------------------------------- -# Launch frontend -# --------------------------------------------------------------------------- echo "[5/8] launching frontend on port ${FRONTEND_PORT}..." @@ -403,7 +388,6 @@ fi echo " frontend pid=${frontend_pid} log=${frontend_log}" -# Poll FE root echo "[6/8] polling http://127.0.0.1:${FRONTEND_PORT}/ ..." fe_ready=0 for i in $(seq 1 30); do @@ -428,9 +412,6 @@ fi echo "[7/8] frontend / OK" -# --------------------------------------------------------------------------- -# Print summary -# --------------------------------------------------------------------------- cwd="$(readlink "/proc/${backend_pid}/cwd" 2>/dev/null || echo unknown)" gpu_mem_now="$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | sed -n "$((GPU + 1))p" || echo 0)" diff --git a/.buildkite/scripts/pr_test.sh b/.buildkite/scripts/pr_test.sh index 3ea9fc7c4..0dffddc3c 100755 --- a/.buildkite/scripts/pr_test.sh +++ b/.buildkite/scripts/pr_test.sh @@ -43,7 +43,6 @@ if ! python3 -m modal --version &> /dev/null; then # --break-system-packages preserves prior `pip install --user` semantics on PEP 668 agents. uv pip install --system --break-system-packages modal - # Verify installation if ! python3 -m modal --version &> /dev/null; then log "Error: Failed to install modal. Please install it manually." exit 1 @@ -52,7 +51,6 @@ fi log "modal version: $(python3 -m modal --version)" -# Set up Modal authentication using Buildkite secrets log "Setting up Modal authentication from Buildkite secrets..." MODAL_TOKEN_ID=$(buildkite-agent secret get modal_token_id) MODAL_TOKEN_SECRET=$(buildkite-agent secret get modal_token_secret) diff --git a/.github/mergify.yml b/.github/mergify.yml index 253cc5878..668eb225f 100644 --- a/.github/mergify.yml +++ b/.github/mergify.yml @@ -33,9 +33,7 @@ merge_protections: pull_request_rules: - # ============================================================ # Type labels (from PR title prefix) - # ============================================================ - name: "label type: feat" conditions: @@ -117,9 +115,7 @@ pull_request_rules: label: add: ["type: skill"] - # ============================================================ # Scope labels (from changed files) - # ============================================================ - name: "label scope: training" conditions: @@ -227,9 +223,7 @@ pull_request_rules: label: add: ["scope: model"] - # ============================================================ # Merge conflict detection - # ============================================================ - name: label conflicting PRs conditions: @@ -249,9 +243,7 @@ pull_request_rules: label: remove: [needs-rebase] - # ============================================================ # Auto-merge - # ============================================================ - name: auto-merge when ready and all checks pass conditions: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index c0cc0d650..62067eb3f 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -24,7 +24,7 @@ repos: require_serial: true args: [--in-place, --verbose] language_version: python3.12 - additional_dependencies: [toml] # TODO: Remove when yapf is upgraded + additional_dependencies: [toml] - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.11.12 hooks: @@ -36,10 +36,6 @@ repos: - id: codespell additional_dependencies: ['tomli'] args: ['--toml', 'pyproject.toml'] -# - repo: https://github.com/PyCQA/isort -# rev: 6.0.1 -# hooks: -# - id: isort - repo: https://github.com/jackdewinter/pymarkdown rev: v0.9.30 hooks: diff --git a/apps/dreamverse/dreamverse/gpu_pool.py b/apps/dreamverse/dreamverse/gpu_pool.py index 2b206c0c6..920f66243 100644 --- a/apps/dreamverse/dreamverse/gpu_pool.py +++ b/apps/dreamverse/dreamverse/gpu_pool.py @@ -307,7 +307,7 @@ def gpu_worker_process( message=str(e), )) - return True # Continue loop + return True if first_cmd is not None: if first_cmd.type == CommandType.SHUTDOWN: diff --git a/apps/dreamverse/dreamverse/session/controller.py b/apps/dreamverse/dreamverse/session/controller.py index 04eaad324..089b4306f 100644 --- a/apps/dreamverse/dreamverse/session/controller.py +++ b/apps/dreamverse/dreamverse/session/controller.py @@ -1704,7 +1704,6 @@ class SessionController: and slot.shared_stream_buffer is not None): start = chunk_offset end = start + chunk_length - # Copy out of shared buffer for websocket send. chunk = bytes(slot.shared_stream_buffer[start:end]) else: chunk = chunk_bytes or b"" diff --git a/apps/dreamverse/scripts/install_native_ffmpeg.sh b/apps/dreamverse/scripts/install_native_ffmpeg.sh index c617eca90..6fed59315 100755 --- a/apps/dreamverse/scripts/install_native_ffmpeg.sh +++ b/apps/dreamverse/scripts/install_native_ffmpeg.sh @@ -49,7 +49,6 @@ # toolchain. set -euo pipefail -# ─── Defaults (override via env) ────────────────────────────────────────── INSTALL_PREFIX="${INSTALL_PREFIX:-$HOME/opt/ffmpeg-native}" SOURCE_DIR="${SOURCE_DIR:-$HOME/src/ffmpeg-native}" X264_REF="${X264_REF:-stable}" @@ -137,7 +136,6 @@ export CC CXX [[ -n "${AS:-}" ]] && export AS echo "[install_native_ffmpeg] toolchain: CC=$CC CXX=$CXX AS=${AS:-} (uname -m=$ARCH, source: $toolchain_source)" -# ─── Step 0: probe required tools ───────────────────────────────────────── required=("$CC" "$CXX" make pkg-config git) [[ "$(uname -m)" == "x86_64" ]] && required+=(nasm) missing=() @@ -155,7 +153,6 @@ if (( ${#missing[@]} > 0 )); then exit 1 fi -# ─── Step 0.5: destructive-path guards ──────────────────────────────────── guard_path() { local name="$1" value="$2" case "$value" in @@ -172,7 +169,6 @@ guard_path() { guard_path INSTALL_PREFIX "$INSTALL_PREFIX" guard_path SOURCE_DIR "$SOURCE_DIR" -# ─── Step 1: clean ──────────────────────────────────────────────────────── echo "[install_native_ffmpeg] cleaning prior install + sources" rm -rf "$INSTALL_PREFIX" "$SOURCE_DIR/x264" "$SOURCE_DIR/ffmpeg" "$SOURCE_DIR/nv-codec-headers" mkdir -p "$SOURCE_DIR" "$INSTALL_PREFIX/lib" @@ -271,7 +267,6 @@ git clone --depth 1 --branch "$FFMPEG_REF" \ make install ) -# ─── Step 4: sanity check ────────────────────────────────────────────────── ffmpeg_bin="$INSTALL_PREFIX/bin/ffmpeg" echo "[install_native_ffmpeg] verifying $ffmpeg_bin" "$ffmpeg_bin" -hide_banner -buildconf | grep -i -E 'libx264|lto' diff --git a/apps/dreamverse/scripts/launch/launch_demo.sh b/apps/dreamverse/scripts/launch/launch_demo.sh index 504328812..b61323c83 100755 --- a/apps/dreamverse/scripts/launch/launch_demo.sh +++ b/apps/dreamverse/scripts/launch/launch_demo.sh @@ -96,7 +96,6 @@ wait_for() { return 1 } -# --- backend --- echo "[launch-demo] backend logs: ${BE_LOG}" ( "${BE_SCRIPT}" "${BE_PORT_ARGS[@]}" >"${BE_LOG}" 2>&1 ) & BE_PID=$! @@ -117,7 +116,6 @@ if [[ "${HEALTH_PATH}" != "${READY_PATH}" ]]; then fi fi -# --- frontend --- if [[ "${NO_FRONTEND:-0}" != "1" ]]; then echo "[launch-demo] frontend logs: ${FE_LOG}" ( "${SCRIPT_DIR}/launch_frontend.sh" >"${FE_LOG}" 2>&1 ) & diff --git a/apps/dreamverse/web/src/app/globals.css b/apps/dreamverse/web/src/app/globals.css index 4c2641bfd..88cd0f5af 100644 --- a/apps/dreamverse/web/src/app/globals.css +++ b/apps/dreamverse/web/src/app/globals.css @@ -110,7 +110,6 @@ --radius-xl: calc(var(--radius) + 4px); } -/* ——— Resets ——— */ *, *::before, @@ -194,7 +193,6 @@ summary::-webkit-details-marker { color: #ffffff; } -/* ——— Base layer ——— */ @layer base { * { @@ -205,7 +203,6 @@ summary::-webkit-details-marker { } } -/* ——— Utilities ——— */ @layer utilities { .scrollbar-hidden { @@ -269,7 +266,6 @@ summary::-webkit-details-marker { animation: stt-pulse 1.4s ease-in-out infinite; } -/* ——— Theme transition ——— */ html.theme-transition, html.theme-transition *, @@ -284,7 +280,6 @@ html.theme-transition *::after { stroke 300ms ease !important; } -/* —————————————— CUSTOM TAILWIND —————————————— */ @layer components { .debug { diff --git a/apps/dreamverse/web/src/app/page.tsx b/apps/dreamverse/web/src/app/page.tsx index 50a15167d..a5151f226 100644 --- a/apps/dreamverse/web/src/app/page.tsx +++ b/apps/dreamverse/web/src/app/page.tsx @@ -2853,7 +2853,6 @@ export default function Page() { onSelectCurrent={() => { const clips = completedClips as any[]; if (sessionStarted && !sessionExpired) { - // Return to live playback streamStore.patch({ activeClipId: "", activePlaybackStartTime: 0 }); } else if (clips.length > 0) { selectClip(clips[clips.length - 1].id, 0); diff --git a/apps/dreamverse/web/src/lib/media/avPipeline.test.ts b/apps/dreamverse/web/src/lib/media/avPipeline.test.ts index ba7eda4ba..03d1bc367 100644 --- a/apps/dreamverse/web/src/lib/media/avPipeline.test.ts +++ b/apps/dreamverse/web/src/lib/media/avPipeline.test.ts @@ -219,15 +219,12 @@ describe('createAvPipeline', () => { updating = false; addEventListener(): void { - // noop for test } removeEventListener(): void { - // noop for test } appendBuffer(): void { - // noop for test } } @@ -335,11 +332,9 @@ describe('createAvPipeline', () => { appendBuffer = vi.fn(); addEventListener(): void { - // noop for test } removeEventListener(): void { - // noop for test } } diff --git a/apps/dreamverse/web/src/lib/media/avPipeline.ts b/apps/dreamverse/web/src/lib/media/avPipeline.ts index 0c29a6511..2f8b8ee47 100644 --- a/apps/dreamverse/web/src/lib/media/avPipeline.ts +++ b/apps/dreamverse/web/src/lib/media/avPipeline.ts @@ -270,7 +270,6 @@ export function createAvPipeline({ return; } stallRecoveryTimerId = setInterval(() => { - // Stop polling once playback is running normally. if (initialPlaybackStarted && !playbackResumePending) { clearStallRecovery(); return; diff --git a/apps/fastvideo_studio/e2e/datasets.spec.ts b/apps/fastvideo_studio/e2e/datasets.spec.ts index fff508943..d3fb668fc 100644 --- a/apps/fastvideo_studio/e2e/datasets.spec.ts +++ b/apps/fastvideo_studio/e2e/datasets.spec.ts @@ -10,7 +10,6 @@ test.describe('datasets', () => { skipWithoutMock(); test('lists the seeded datasets', async ({ page, request }) => { - // Read the seeded names from the mock so the assertion tracks the fixture. const res = await request.get(`${API_BASE}/datasets`); const datasets = (await res.json()) as Array<{ name: string }>; expect(datasets.length).toBeGreaterThan(0); diff --git a/apps/fastvideo_studio/job_runner.py b/apps/fastvideo_studio/job_runner.py index 2190b8462..aa94f392a 100644 --- a/apps/fastvideo_studio/job_runner.py +++ b/apps/fastvideo_studio/job_runner.py @@ -711,8 +711,8 @@ class JobRunner: job.finished_at = None job._stop_event.clear() job._log_buf = JobLogBuffer() # fresh log buffer - job.log_file_handler = None # Reset file handler - job._process = None # Reset subprocess handle + job.log_file_handler = None + job._process = None # Wrap _run_job in an additional safety layer to catch any exceptions # that might escape (though they shouldn't with our comprehensive handling) diff --git a/apps/fastvideo_studio/src/app/globals.css b/apps/fastvideo_studio/src/app/globals.css index 5a49bcc3c..9425dd3bd 100644 --- a/apps/fastvideo_studio/src/app/globals.css +++ b/apps/fastvideo_studio/src/app/globals.css @@ -120,7 +120,6 @@ --radius-xl: calc(var(--radius) + 4px); } -/* ——— Resets ——— */ html, body { @@ -211,7 +210,6 @@ html.theme-transition *::after { stroke 300ms ease !important; } -/* ——— Base layer ——— */ @layer base { button, diff --git a/apps/fastvideo_studio/src/components/datasets/DatasetSidebar.test.tsx b/apps/fastvideo_studio/src/components/datasets/DatasetSidebar.test.tsx index 33e5049cd..1e299479c 100644 --- a/apps/fastvideo_studio/src/components/datasets/DatasetSidebar.test.tsx +++ b/apps/fastvideo_studio/src/components/datasets/DatasetSidebar.test.tsx @@ -54,15 +54,12 @@ describe('DatasetSidebar', () => { it('lists dataset files after loading', async () => { render( {}} />); - // The first file's caption is rendered once loading resolves. expect(await screen.findByDisplayValue('cap a')).toBeInTheDocument(); expect(mockedApi.getDatasetFiles).toHaveBeenCalledWith('ds-1'); - // One caption editor per returned file. const captionFields = screen.getAllByPlaceholderText('Caption'); expect(captionFields).toHaveLength(2); - // Media URLs are requested per visible file. expect(mockedApi.getDatasetMediaUrl).toHaveBeenCalledWith('ds-1', 'a.mp4'); expect(mockedApi.getDatasetMediaUrl).toHaveBeenCalledWith('ds-1', 'b.mp4'); }); diff --git a/apps/fastvideo_studio/src/components/jobs/CreateJobModal.test.tsx b/apps/fastvideo_studio/src/components/jobs/CreateJobModal.test.tsx index a2ae5517d..63e60de90 100644 --- a/apps/fastvideo_studio/src/components/jobs/CreateJobModal.test.tsx +++ b/apps/fastvideo_studio/src/components/jobs/CreateJobModal.test.tsx @@ -23,7 +23,6 @@ const MODELS = [ ]; beforeEach(() => { - // Reset the shared options store to a known baseline for test isolation. defaultOptionsStore.set({ options: DEFAULT_OPTIONS }); vi.mocked(getModels).mockResolvedValue(MODELS); vi.mocked(getDatasets).mockResolvedValue([]); diff --git a/apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx b/apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx index cc96ebc29..ecd018348 100644 --- a/apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx +++ b/apps/fastvideo_studio/src/components/jobs/CreateJobModal.tsx @@ -302,7 +302,6 @@ export default function CreateJobModal({ } }, [isOpen, workloadType, inferenceWorkload, options, editingJobId]); - // Load the models available for this workload. React.useEffect(() => { if (!isOpen) return; // Ignore a superseded response so a slow fetch for a previous workload diff --git a/apps/fastvideo_studio/src/components/jobs/JobDetailsSidebar.test.tsx b/apps/fastvideo_studio/src/components/jobs/JobDetailsSidebar.test.tsx index 75ece9ecd..61a926a5f 100644 --- a/apps/fastvideo_studio/src/components/jobs/JobDetailsSidebar.test.tsx +++ b/apps/fastvideo_studio/src/components/jobs/JobDetailsSidebar.test.tsx @@ -81,7 +81,6 @@ describe('JobDetailsSidebar', () => { />, ); - // Flush the immediate poll fired on mount. await act(async () => { await vi.advanceTimersByTimeAsync(0); }); diff --git a/apps/fastvideo_studio/src/components/jobs/JobQueue.tsx b/apps/fastvideo_studio/src/components/jobs/JobQueue.tsx index bee02ca44..58c0013e5 100644 --- a/apps/fastvideo_studio/src/components/jobs/JobQueue.tsx +++ b/apps/fastvideo_studio/src/components/jobs/JobQueue.tsx @@ -101,12 +101,10 @@ export default function JobQueue({ jobType, jobTypesForList }: JobQueueProps) { fetchJobsRef.current = fetchJobs; }, [fetchJobs]); - // Initial fetch + refetch whenever an external refresh is triggered. React.useEffect(() => { fetchJobs(); }, [fetchJobs, nonce]); - // Poll every second while any job is running/pending; stop otherwise. const hasActive = jobs.some( (j) => j.status === 'running' || j.status === 'pending', ); diff --git a/comfyui/web/js/FastVideo.core.js b/comfyui/web/js/FastVideo.core.js index 53afb22ad..44ddccc1a 100644 --- a/comfyui/web/js/FastVideo.core.js +++ b/comfyui/web/js/FastVideo.core.js @@ -68,7 +68,6 @@ function drawAutoAnnotated(ctx, node, widget_width, y, H) { ctx.fillStyle = litegraph_base.WIDGET_BGCOLOR; ctx.fill(); - // Reset fill style for the teeth if (isAuto) { ctx.fillStyle = litegraph_base.WIDGET_TEXT_COLOR; } else { @@ -200,7 +199,6 @@ function mouseAutoAnnotated(event, [x, y], node) { return true; } - // Handle increment/decrement buttons if not in auto mode if (delta !== 0 && !this.isAuto) { if (this.config[0] === "FVAUTOCOMBO") { const options = this.options.values || []; @@ -241,7 +239,6 @@ function mouseAutoAnnotated(event, [x, y], node) { v += increment; - // Apply min/max constraints if (this.options.min != null) { v = Math.max(this.options.min, v); } @@ -272,7 +269,6 @@ function mouseAutoAnnotated(event, [x, y], node) { if (this.config[0] === "FVAUTOCOMBO") { const options = this.options.values || []; - // Create menu items const menuItems = options.map(opt => { const value = typeof opt === 'object' ? opt.value : opt; const label = typeof opt === 'object' ? opt.label : opt.toString(); @@ -323,7 +319,6 @@ function mouseAutoAnnotated(event, [x, y], node) { const d_callback = (v) => { this.value = this.parseValue?.(v) ?? Number(v); - // Apply min/max constraints if (this.options.min != null) { this.value = Math.max(this.options.min, this.value); } @@ -370,7 +365,6 @@ function makeAutoAnnotated(widget, inputData) { value: widget.value }; - // Add AUTO properties to the widget Object.assign(widget, { type: "BOOLEAN", draw: drawAutoAnnotated, @@ -381,7 +375,7 @@ function makeAutoAnnotated(widget, inputData) { cachedValue: widget.value, config: inputData, options: Object.assign({}, inputData[1], widget.options), - original: original, // Store original properties for reference + original: original, // Disable other potential mouse handlers with no-op functions onClick: function () { @@ -429,7 +423,6 @@ function makeAutoAnnotated(widget, inputData) { return v; }, serializeValue: function () { - // Return special value for AUTO mode return this.isAuto ? -99999 : this.value; }, deserializeValue: function (data) { @@ -473,28 +466,22 @@ app.registerExtension({ async beforeRegisterNodeDef(nodeType, nodeData, app) { if (nodeData?.name == "VideoGenerator" || nodeData?.name === "InferenceArgs" || nodeData?.name === "VAEConfig" || nodeData?.name === "TextEncoderConfig" || nodeData?.name === "DITConfig") { - // Add serialization support chainCallback(nodeType.prototype, "onSerialize", function (info) { if (!this.widgets) { return; } - // Ensure widgets_values exists if (!info.widgets_values) { info.widgets_values = {}; } - // Store AUTO widget states in a separate property if (!info.auto_widget_states) { info.auto_widget_states = {}; } - // Handle AUTO widgets specially for (const w of this.widgets) { if (w.type === "BOOLEAN" && w.isAuto !== undefined) { - // Store the serialized value (for Python node) info.widgets_values[w.name] = w.serializeValue(); - // Store the full state (for UI restoration) info.auto_widget_states[w.name] = { isAuto: w.isAuto, value: w.value, @@ -504,7 +491,6 @@ app.registerExtension({ } }); - // Add deserialization support chainCallback(nodeType.prototype, "onConfigure", function (info) { if (!this.widgets) { return; @@ -522,7 +508,6 @@ app.registerExtension({ } } - // Then, restore full state if available if (info.auto_widget_states) { for (const w of this.widgets) { if (w.type === "BOOLEAN" && w.isAuto !== undefined && w.name in info.auto_widget_states) { @@ -544,7 +529,6 @@ app.registerExtension({ // Override addInput to handle AUTO widgets chainCallback(nodeType.prototype, "onNodeCreated", function () { - // Convert any existing widgets to AUTO widgets if needed let new_widgets = []; const intWidgetNames = ["sp_size", "tp_size", "height", "width", "num_frames", "num_inference_steps", "flow_shift", "seed", "fps", "scale_factor", "tile_sample_min_height", "tile_sample_min_width", "tile_sample_min_num_frames", "tile_sample_stride_height", "tile_sample_stride_width", diff --git a/docs/cookbook/recipes/fastwan21.yaml b/docs/cookbook/recipes/fastwan21.yaml index 1189ef5b4..f9500d388 100644 --- a/docs/cookbook/recipes/fastwan21.yaml +++ b/docs/cookbook/recipes/fastwan21.yaml @@ -1,6 +1,3 @@ -# TODO(cookbook-review): Design sample only. Before merge, double-check the -# selected options, schema types/ranges, and baseline values for each deployment. -# TODO: Recheck field paths and choices when the typed API in PR #1913 is finalized. title: FastWan2.1 1.3B summary: Distilled text-to-video serving with sparse attention. diff --git a/docs/cookbook/recipes/wan21-i2v.yaml b/docs/cookbook/recipes/wan21-i2v.yaml index 72df0800a..4f5ba30c6 100644 --- a/docs/cookbook/recipes/wan21-i2v.yaml +++ b/docs/cookbook/recipes/wan21-i2v.yaml @@ -1,6 +1,3 @@ -# TODO(cookbook-review): Design sample only. Before merge, double-check the -# selected options, schema types/ranges, and baseline values for each deployment. -# TODO: Recheck field paths and choices when the typed API in PR #1913 is finalized. title: Wan2.1 I2V 14B summary: Image-conditioned video generation from a source frame. diff --git a/docs/generate_examples.py b/docs/generate_examples.py index 8203bead7..ceb45cd7a 100644 --- a/docs/generate_examples.py +++ b/docs/generate_examples.py @@ -14,7 +14,7 @@ ROOT_DIR = Path(__file__).parent.parent.resolve() ROOT_DIR_RELATIVE = '../..' EXAMPLE_DIR = ROOT_DIR / "examples" EXAMPLE_DOC_DIR = ROOT_DIR / "docs/getting_started/examples" -GITHUB_REPO = "hao-ai-lab/FastVideo" # Update this to your repo +GITHUB_REPO = "hao-ai-lab/FastVideo" COOKBOOK_DATA = ROOT_DIR / "docs/assets/cookbook-recipes.json" COOKBOOK_CLIENTS = { "python": ("video.py", "python -m pip install openai==3.6.0"), @@ -364,11 +364,8 @@ class Index: content = f"# {self.title}\n\n{self.description}\n\n" if self.caption: content += f"## {self.caption}\n\n" - # Generate a simple list of documentation links for doc in self.documents: - # Convert document path to proper link doc_link = doc.replace("\\", "/") - # Get just the filename for the link text doc_title = fix_case(Path(doc).stem.replace("_", " ").title()) content += f"- [{doc_title}]({doc_link}.md)\n" content += "\n" @@ -459,12 +456,10 @@ class Example: return fix_case(self.path.stem.replace("_", " ").title()) def generate(self) -> str: - # Create GitHub link to source github_path = str(self.path.relative_to(ROOT_DIR)).replace("\\", "/") github_url = f"https://github.com/{GITHUB_REPO}/blob/main/{github_path}" content = f"**Source:** [{github_path}]({github_url})\n\n" - # Add title for code files if self.main_file.suffix != ".md": content += f"# {self.title}\n\n" @@ -484,19 +479,16 @@ class Example: return content content += "## Additional Files\n\n" - # Define binary/non-text file extensions to skip binary_extensions = { '.mp4', '.avi', '.mov', '.mkv', '.gif', '.jpg', '.jpeg', '.png', '.webp', '.bmp', '.pdf', '.zip', '.tar', '.gz', '.mp3', '.wav' } for file in sorted(self.other_files): - # Skip binary files if file.suffix.lower() in binary_extensions: continue file_rel_path = file.relative_to(self.path) - # Use collapsible admonitions, normalized by the site preparation content += f"??? note \"{file_rel_path}\"\n\n" try: @@ -517,7 +509,6 @@ class Example: content += " ```\n" content += "\n" except UnicodeDecodeError: - # Skip files that can't be decoded as UTF-8 continue return content @@ -580,7 +571,6 @@ def create_category_indices() -> dict[str, Index]: ), } - # Ensure all category doc directories exist for index in category_indices.values(): if not index.path.parent.exists(): index.path.parent.mkdir(parents=True) @@ -604,7 +594,6 @@ def find_examples(category_indices: dict[str, Index], generate_main_index: bool) dir_name = category_dir_mapping.get(category, category) category_dir = EXAMPLE_DIR / dir_name - # Skip if directory doesn't exist if not category_dir.exists(): continue @@ -622,7 +611,6 @@ def find_examples(category_indices: dict[str, Index], generate_main_index: bool) examples.append(Example(path)) # Find examples in subdirectories for path in EXAMPLE_DIR.glob("*/*.md"): - # Skip categorised examples if path.parent.name in category_indices: continue examples.append(Example(path.parent)) @@ -643,7 +631,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic if example.category not in ["training", "distillation"]: continue - # Use mapped directory name if available dir_name = category_dir_mapping.get(example.category, example.category) category_dir = EXAMPLE_DIR / dir_name relative_path = example.path.relative_to(category_dir) @@ -656,7 +643,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic model = path_parts[1] # e.g., "wan_i2v_14b_480p" dataset = path_parts[2] # e.g., "crush_smol" - # Initialize nested structure if example.category not in nested_structures: nested_structures[example.category] = {} if method not in nested_structures[example.category]: @@ -664,7 +650,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic if model not in nested_structures[example.category][method]: nested_structures[example.category][method][model] = {} - # Store the nested structure nested_structures[example.category][method][model][dataset] = NestedStructure(category=example.category, method=method, model=model, @@ -677,7 +662,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic dataset = path_parts[1] # e.g., "Wan-Syn-Data-480P" method = "DMD" # Default method for distillation - # Initialize nested structure if example.category not in nested_structures: nested_structures[example.category] = {} if method not in nested_structures[example.category]: @@ -685,7 +669,6 @@ def create_nested_structures(examples: list[Example]) -> dict[str, dict[str, dic if model not in nested_structures[example.category][method]: nested_structures[example.category][method][model] = {} - # Store the nested structure nested_structures[example.category][method][model][dataset] = NestedStructure(category=example.category, method=method, model=model, @@ -700,7 +683,7 @@ def generate_flat_examples(examples: list[Example], category_indices: dict[str, """Generate documentation for flat structure examples (inference, etc.).""" for example in examples: if example.category in ["training", "distillation"]: - continue # Skip nested structure examples + continue # Determine which index to use for this example if example.category is not None and example.category in category_indices: @@ -711,7 +694,6 @@ def generate_flat_examples(examples: list[Example], category_indices: dict[str, else: continue - # Generate the example documentation doc_path = index.path.parent / f"{example.path.stem}.md" with open(doc_path, "w+", encoding="utf-8") as f: f.write(example.generate()) @@ -729,7 +711,6 @@ def generate_nested_examples(nested_structures: dict[str, dict[str, dict[str, di category_base_dir = category_index.path.parent for method, models in nested_structures[category_name].items(): - # Create method-level index method_index = Index(path=category_base_dir / f"{method}.md", title=fix_case(method), description=f"Examples using {method}.", @@ -737,35 +718,28 @@ def generate_nested_examples(nested_structures: dict[str, dict[str, dict[str, di maxdepth=2) for model, datasets in models.items(): - # Generate dataset examples using the Example class for dataset, nested_struct in datasets.items(): doc_path = category_base_dir / f"{nested_struct.filename}.md" with open(doc_path, "w+", encoding="utf-8") as f: f.write(nested_struct.example.generate()) - # Create model-level index model_index = Index(path=category_base_dir / f"{model}.md", title=fix_case(model.replace('_', ' ')), description=f"Examples for the {model} model.", caption=f"{fix_case(model.replace('_', ' '))} Datasets", maxdepth=1) - # Add dataset indices to model index for dataset, nested_struct in datasets.items(): model_index.documents.append(nested_struct.filename) - # Write model index with open(model_index.path, "w+", encoding="utf-8") as f: f.write(model_index.generate()) - # Add model to method index method_index.documents.append(model) - # Write method index with open(method_index.path, "w+", encoding="utf-8") as f: f.write(method_index.generate()) - # Add method to main category index category_index.documents.append(method) @@ -777,7 +751,6 @@ def generate_examples(generate_main_index: bool = False) -> None: generate_main_index (bool): Whether to generate the main examples index. If False, only category-specific indices will be generated. """ - # Create category indices category_indices = create_category_indices() # Create the main examples index only if requested @@ -796,29 +769,22 @@ def generate_examples(generate_main_index: bool = False) -> None: # Find all examples examples = find_examples(category_indices, generate_main_index) - # Create nested structures for training and distillation nested_structures = create_nested_structures(examples) - # Generate flat structure examples (inference, etc.) generate_flat_examples(examples, category_indices, examples_index, generate_main_index) - # Generate nested structure examples (training, distillation) generate_nested_examples(nested_structures, category_indices) - # Generate the index files for categories for category_index in category_indices.values(): if category_index.documents: - # Add to main index if it exists if generate_main_index and examples_index: main_index_dir = examples_index.path.parent rel_path = os.path.relpath(category_index.path, start=main_index_dir) examples_index.documents.insert(0, str(rel_path).replace("\\", "/").replace(".md", "")) - # Write the category index file with open(category_index.path, "w+", encoding="utf-8") as f: f.write(category_index.generate()) - # Write the main index file if requested if generate_main_index and examples_index: with open(examples_index.path, "w+", encoding="utf-8") as f: f.write(examples_index.generate()) diff --git a/examples/distill/MatrixGame2.0/distill_dmd.sh b/examples/distill/MatrixGame2.0/distill_dmd.sh index 708146c05..0493a7da0 100644 --- a/examples/distill/MatrixGame2.0/distill_dmd.sh +++ b/examples/distill/MatrixGame2.0/distill_dmd.sh @@ -7,7 +7,6 @@ export TOKENIZERS_PARALLELISM=false RUN_NAME=$(date +"%m%d_%H%M") echo "RUN_NAME: $RUN_NAME" -# Model paths for Self-Forcing DMD distillation. GENERATOR_MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers" REAL_SCORE_MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers" # Teacher FAKE_SCORE_MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers" # Critic @@ -16,7 +15,6 @@ DATA_DIR="data/matrixgame2" VALIDATION_DATASET_FILE="examples/distill/MatrixGame2.0/validation.json" NUM_GPUS=1 -# Training arguments training_args=( --tracker_project_name "matrixgame2_sf" --output_dir "checkpoints/matrixgame2_sf_${RUN_NAME}" @@ -32,12 +30,8 @@ training_args=( --simulate_generator_forward --num_frames 81 --num_frame_per_block 3 # Frame generation block size for self-forcing - # --enable_gradient_masking - # --gradient_mask_last_n_frames 21 - # --init_weights_from_safetensors "path/to/generator_ema.safetensors" ) -# Parallel arguments parallel_args=( --num_gpus $NUM_GPUS --sp_size 1 @@ -47,7 +41,7 @@ parallel_args=( ) model_args=( - --model_path $GENERATOR_MODEL_PATH # TODO: check if you can remove this in this script + --model_path $GENERATOR_MODEL_PATH --pretrained_model_name_or_path $GENERATOR_MODEL_PATH --real_score_model_path $REAL_SCORE_MODEL_PATH --fake_score_model_path $FAKE_SCORE_MODEL_PATH @@ -58,7 +52,6 @@ dataset_args=( --dataloader_num_workers 4 ) -# Validation arguments validation_args=( --log_validation --log_visualization @@ -69,7 +62,6 @@ validation_args=( --validation_guidance_scale "6.0" ) -# Optimizer arguments optimizer_args=( --learning_rate 3e-6 --mixed_precision "bf16" @@ -80,7 +72,6 @@ optimizer_args=( --max_grad_norm 1.0 ) -# Miscellaneous arguments miscellaneous_args=( --inference_mode False --checkpoints_total_limit 3 diff --git a/examples/distill/MatrixGame2.0/distill_dmd.slurm b/examples/distill/MatrixGame2.0/distill_dmd.slurm index 8a7a47b31..7a360c0d2 100644 --- a/examples/distill/MatrixGame2.0/distill_dmd.slurm +++ b/examples/distill/MatrixGame2.0/distill_dmd.slurm @@ -77,7 +77,6 @@ training_args=( --simulate_generator_forward --num_frames 81 --num_frame_per_block 3 - # --init_weights_from_safetensors "path/to/generator_ema.safetensors" ) parallel_args=( diff --git a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh index 1c72cab57..d9f86f226 100644 --- a/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh +++ b/examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh @@ -33,8 +33,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model DATA_DIR=your_data_dir VALIDATION_DATASET_FILE=your_validation_data_dir -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] training_args=( --tracker_project_name SFwan_t2v_distill_self_forcing_dmd @@ -64,7 +62,7 @@ parallel_args=( ) model_args=( - --model_path $GENERATOR_MODEL_PATH # TODO: check if you can remove this in this script + --model_path $GENERATOR_MODEL_PATH --pretrained_model_name_or_path $GENERATOR_MODEL_PATH --real_score_model_path $REAL_SCORE_MODEL_PATH --fake_score_model_path $FAKE_SCORE_MODEL_PATH diff --git a/examples/distill/SFWan2.2-A14B/distill_dmd.sh b/examples/distill/SFWan2.2-A14B/distill_dmd.sh index 4750cb1a3..d048fa77b 100644 --- a/examples/distill/SFWan2.2-A14B/distill_dmd.sh +++ b/examples/distill/SFWan2.2-A14B/distill_dmd.sh @@ -23,7 +23,6 @@ nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) export MASTER_ADDR=${nodes[0]} export TOKENIZERS_PARALLELISM=false export WANDB_API_KEY=YOUR_WANDB_API_KEY -# export WANDB_API_KEY='your_wandb_api_key_here' export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN @@ -31,23 +30,13 @@ export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # Configs NUM_GPUS=8 -# Model paths for Self-Forcing DMD distillation with Wan2.2: -# GENERATOR_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Updated to Wan2.2 -# REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Teacher model -# FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-T2V-A14B-Diffusers" # Critic model GENERATOR_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Updated to Wan2.2 REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Teacher model FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model -# DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/" DATA_DIR="/mnt/weka/home/hao.zhang/matthew/FastVideo/data/test-text-preprocessing" -# DATA_DIR=data/crush-smol_processed_t2v/combined_parquet_dataset -# DATA_DIR="/mnt/sharefs/users/hao.zhang/Vchitect-2M/Wan-Syn-upload/latents_i2v/train/" -# VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json" VALIDATION_DATASET_FILE="/mnt/weka/home/hao.zhang/wl/FastVideo/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/validation_64.json" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] training_args=( --tracker_project_name SFwan2.2_t2v_distill_self_forcing_dmd # Updated for Wan2.2 diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm index 9fbc795c5..091bd1296 100644 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm +++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_1.3B.slurm @@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -44,8 +43,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR=your_data_dir VALIDATION_DATASET_FILE=your_validation_dataset_file OUTPUT_DIR="checkpoints/wan_t2v_finetune" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm index ceefedf11..c7fbaabfa 100644 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm +++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_VSA_t2v_14B.slurm @@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -44,8 +43,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" DATA_DIR=your_data_dir VALIDATION_DATASET_FILE=your_validation_dataset_file OUTPUT_DIR="checkpoints/wan_t2v_finetune" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm index 7f1cbcc4c..b2433dd9f 100644 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm +++ b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.slurm @@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -44,8 +43,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR=your_data_dir VALIDATION_DATASET_FILE=your_validation_dataset_file OUTPUT_DIR="checkpoints/wan_t2v_finetune" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh index 0e320eb8b..bb65f4e54 100644 --- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh +++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B.sh @@ -32,7 +32,6 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN export WANDB_API_KEY=your_wandb_api_key -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -45,8 +44,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" DATA_DIR=your_data_dir VALIDATION_DIR=your_validation_path #(example:validation_64.json) OUTPUT_DIR="checkpoints/wan_t2v_finetune" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh index 96f8d02f5..d5cf1d204 100644 --- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh +++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/Data-free/distill_dmd_t2v_5B_VSA.sh @@ -32,7 +32,6 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN export WANDB_API_KEY=your_wandb_api_key -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -44,8 +43,6 @@ REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" DATA_DIR=your_data_dir VALIDATION_DIR=your_validation_path #(example:validation_64.json) -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh index 4ca0404d9..a6dd5d4e3 100644 --- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh +++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh @@ -9,7 +9,6 @@ export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA # Configs NUM_GPUS=1 @@ -19,8 +18,6 @@ FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json" OUTPUT_DIR="checkpoints/wan_t2v_finetune" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh index 39746249c..05cfd2584 100644 --- a/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh +++ b/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B_lora.sh @@ -9,7 +9,6 @@ export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA # Configs NUM_GPUS=1 @@ -18,8 +17,6 @@ REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.2-TI2V-5B-Diffusers" DATA_DIR="data/crush-smol_processed_ti2v/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/validation.json" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/inference/basic/basic_matrixgame2_streaming.py b/examples/inference/basic/basic_matrixgame2_streaming.py index d3e008988..5710975a2 100644 --- a/examples/inference/basic/basic_matrixgame2_streaming.py +++ b/examples/inference/basic/basic_matrixgame2_streaming.py @@ -63,7 +63,6 @@ async def main(): dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, - # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, ) diff --git a/examples/inference/basic/basic_wan2_2.py b/examples/inference/basic/basic_wan2_2.py index 8e115de9b..ecacdcfd2 100644 --- a/examples/inference/basic/basic_wan2_2.py +++ b/examples/inference/basic/basic_wan2_2.py @@ -17,7 +17,6 @@ def main(): dit_cpu_offload=True, # DiT need to be offloaded for MoE vae_cpu_offload=False, text_encoder_cpu_offload=True, - # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer pin_cpu_memory=True, ) diff --git a/examples/inference/eval/all_metrics_demo.py b/examples/inference/eval/all_metrics_demo.py index 77cf66f1b..f563332a3 100644 --- a/examples/inference/eval/all_metrics_demo.py +++ b/examples/inference/eval/all_metrics_demo.py @@ -118,7 +118,7 @@ def score_all_metrics(video_path: str) -> None: for name in ev.metric_names: r = results[0].get(name) if r is None: - continue # set metric, will show in corpus section + continue if r.score is None: print(f" {name}: SKIPPED — {r.details.get('skipped', 'no score')}") else: diff --git a/examples/inference/gradio/serving/ray_serve_backend.py b/examples/inference/gradio/serving/ray_serve_backend.py index 66a8b1c15..6c56d5cb0 100644 --- a/examples/inference/gradio/serving/ray_serve_backend.py +++ b/examples/inference/gradio/serving/ray_serve_backend.py @@ -187,8 +187,8 @@ class BaseModelDeployment: num_gpus=1, use_fsdp_inference=True, text_encoder_cpu_offload=config["text_encoder_cpu_offload"], - dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], # TODO: hardocde for I2V - dit_precision="fp32", # TODO: hardocde for I2V + dmd_denoising_steps=[1000, 850, 700, 550, 350, 275, 200, 125], + dit_precision="fp32", dit_cpu_offload=config["dit_cpu_offload"], vae_cpu_offload=config["vae_cpu_offload"], VSA_sparsity=config["VSA_sparsity"], diff --git a/examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v_vsa.yaml b/examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v_vsa.yaml index 5d3947f2c..2cc4dd4b0 100644 --- a/examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v_vsa.yaml +++ b/examples/train/configs/distribution_matching/wan/self_forcing_causal_t2v_vsa.yaml @@ -35,7 +35,6 @@ method: enable_gradient_in_rollout: true start_gradient_frame: 0 - # Critic optimizer fake_score_learning_rate: 8.0e-6 fake_score_betas: [0.0, 0.999] fake_score_lr_scheduler: constant diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml index d79cd317b..cdf40721d 100644 --- a/examples/train/configs/example.yaml +++ b/examples/train/configs/example.yaml @@ -1,4 +1,3 @@ -# ============================================================================== # Full configuration reference for fastvideo.train # # Legend: @@ -8,15 +7,12 @@ # Keys depend on the _target_ class constructor / method_config. # [RESOLVED] — parsed by PipelineConfig.from_kwargs(); auto-populated from # the model's config files. Only scalar overrides are useful. -# ============================================================================== -# ------------------------------------------------------------------------------ # models: [FREE] # # Each role is instantiated via _target_(*, training_config=..., **kwargs). # Keys here are constructor kwargs of the _target_ class (e.g. WanModel). # You can define any role name (student, teacher, critic, etc.). -# ------------------------------------------------------------------------------ models: student: _target_: fastvideo.train.models.wan.WanModel # required @@ -39,13 +35,11 @@ models: trainable: true disable_custom_init_weights: true -# ------------------------------------------------------------------------------ # method: [FREE] # # Instantiated via _target_(*, cfg=RunConfig, role_models=...). # All keys besides _target_ are available in self.method_config (a plain dict). # Keys depend entirely on the method class. -# ------------------------------------------------------------------------------ method: _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method # required @@ -71,11 +65,9 @@ method: # --- FineTuneMethod keys (if using finetune instead) --- # _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod -# ------------------------------------------------------------------------------ # training: [TYPED] -> TrainingConfig # # Every field below has a typed default. Unknown keys are ignored. -# ------------------------------------------------------------------------------ training: # --- training.distributed [TYPED] -> DistributedConfig --- @@ -166,13 +158,11 @@ training: dit_precision: fp32 # default: "fp32" (master weight precision) # model_path: ... # default: "" (auto-derived from models.student.init_from) -# ------------------------------------------------------------------------------ # callbacks: [FREE] # # Each callback is instantiated via _target_(*, **kwargs). # The callback name (e.g. "grad_clip") is arbitrary — only _target_ matters. # training_config is injected automatically (not from YAML). -# ------------------------------------------------------------------------------ callbacks: # --- GradNormClipCallback --- @@ -199,13 +189,11 @@ callbacks: output_dir: null # default: null (falls back to training.checkpoint.output_dir) sampling_timesteps: null # default: null (explicit timestep list for streaming) -# ------------------------------------------------------------------------------ # pipeline: [RESOLVED] -> PipelineConfig # # Parsed by PipelineConfig.from_kwargs(). Most fields are auto-populated from # the model's config files (vae_config, dit_config, text_encoder_configs, etc.). # Only scalar overrides are typically needed here. -# ------------------------------------------------------------------------------ pipeline: flow_shift: 3 # default: null (model-specific) # LTX-2 linear QAT (SM100+ with FlashInfer); distinct from the role-local diff --git a/examples/train/configs/knowledge_distillation/wan/causal_t2v.yaml b/examples/train/configs/knowledge_distillation/wan/causal_t2v.yaml index 0dc9940eb..9786ef73f 100644 --- a/examples/train/configs/knowledge_distillation/wan/causal_t2v.yaml +++ b/examples/train/configs/knowledge_distillation/wan/causal_t2v.yaml @@ -24,7 +24,7 @@ models: _target_: fastvideo.train.models.wan.WanCausalModel init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers trainable: true - teacher: # remove once cache is complete + teacher: _target_: fastvideo.train.models.wan.WanModel init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers trainable: false diff --git a/examples/train/configs/knowledge_distillation/wan/t2v.yaml b/examples/train/configs/knowledge_distillation/wan/t2v.yaml index 079434f39..9959ae32c 100644 --- a/examples/train/configs/knowledge_distillation/wan/t2v.yaml +++ b/examples/train/configs/knowledge_distillation/wan/t2v.yaml @@ -21,7 +21,7 @@ models: _target_: fastvideo.train.models.wan.WanModel init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers trainable: true - teacher: # remove once cache is complete + teacher: _target_: fastvideo.train.models.wan.WanModel init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers trainable: false diff --git a/examples/train/run.sh b/examples/train/run.sh index d3485ed2d..2291f5984 100755 --- a/examples/train/run.sh +++ b/examples/train/run.sh @@ -27,7 +27,6 @@ if [[ "${CONFIG}" != /* ]]; then CONFIG="$(pwd)/${CONFIG}" fi -# ── GPU / node settings ────────────────────────────────────────── NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | wc -l)}" NUM_GPUS="${NUM_GPUS:-1}" NNODES="${NNODES:-1}" @@ -38,11 +37,9 @@ export TOKENIZERS_PARALLELISM=false export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" cd "${REPO_ROOT}" -# ── W&B ────────────────────────────────────────────────────────── export WANDB_API_KEY="${WANDB_API_KEY:-}" export WANDB_MODE="${WANDB_MODE:-online}" -# ── Log file ───────────────────────────────────────────────────── CONFIG_NAME="$(basename "${CONFIG}" .yaml)" TIMESTAMP="$(date +%Y%m%d_%H%M%S)" LOG_DIR="${LOG_DIR:-examples/train}" diff --git a/examples/train/run_slurm.sh b/examples/train/run_slurm.sh index 153ba02b7..e92b400d8 100755 --- a/examples/train/run_slurm.sh +++ b/examples/train/run_slurm.sh @@ -49,7 +49,6 @@ JOB_NAME="${JOB_NAME:-${CONFIG_NAME}}" OUTPUT_DIR="${OUTPUT_DIR:-logs/slurm}" mkdir -p "${OUTPUT_DIR}" -# ── Build sbatch args ───────────────────────────────────────────── # Forward the exported W&B API key through Slurm's environment so the generated # job script and its xtrace output do not contain the credential value. SBATCH_ARGS=( @@ -70,7 +69,6 @@ SBATCH_ARGS=( if [[ -n "${EXCLUDE}" ]]; then SBATCH_ARGS+=(--exclude="${EXCLUDE}") fi -# Apply hardware and wall-clock requirements selected by an experiment launcher. if [[ -n "${SBATCH_CONSTRAINT}" ]]; then SBATCH_ARGS+=(--constraint="${SBATCH_CONSTRAINT}") fi diff --git a/examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh b/examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh index 84819bc51..cf1c9677e 100755 --- a/examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh +++ b/examples/train/scenario/qad_wan2_1_mixkit/run_stage1.sh @@ -1,5 +1,4 @@ #!/usr/bin/env bash -# Run the modular Attn-QAT finetune recipe. set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" diff --git a/examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml b/examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml index 29ed5f810..78b03551a 100644 --- a/examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml +++ b/examples/train/scenario/worldmodel/solaris/finetune_i2v.yaml @@ -45,9 +45,6 @@ training: checkpoints_total_limit: 3 resume_from_checkpoint: latest - # tracker: - # project_name: fastvideo - # run_name: matrixgame2_i2v_finetune tracker: trackers: [none] project_name: "" diff --git a/examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml b/examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml index 1ad114eb2..638d5ea89 100644 --- a/examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml +++ b/examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml @@ -33,7 +33,6 @@ method: enable_gradient_in_rollout: true start_gradient_frame: 0 - # Critic optimizer fake_score_learning_rate: 3.0e-7 fake_score_betas: [0.9, 0.95] fake_score_lr_scheduler: constant diff --git a/examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml b/examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml index 4b3424794..a6f0b3663 100644 --- a/examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml +++ b/examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml @@ -6,7 +6,6 @@ models: _target_: fastvideo.train.models.matrixgame2.matrixgame2_causal.MatrixGame2CausalModel init_from: mignonjia/mg_sf_distilled_zelda_1k_steps trainable: true - # transformer_override_safetensor: outputs/matrixgame_dmd/checkpoints/mg_zelda_sf_m2/checkpoint-1000_weight_only/ema/generator_ema.safetensors teacher: _target_: fastvideo.train.models.matrixgame2.matrixgame2.MatrixGame2Model init_from: mignonjia/mg_bidirectional_zelda diff --git a/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh b/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh index a0d611dce..42429a6fd 100644 --- a/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh +++ b/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh @@ -8,7 +8,6 @@ MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers" DATA_DIR="data/crush-smol_processed_t2v_1_3b_ode_init/" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=1 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init_16k_vidprom.slurm b/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init_16k_vidprom.slurm index b56941be6..74c7dc02d 100644 --- a/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init_16k_vidprom.slurm +++ b/examples/training/consistency_finetune/causal_ode_init/finetune_ode_init_16k_vidprom.slurm @@ -29,7 +29,6 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_API_KEY=your-wandb-api-key -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" diff --git a/examples/training/finetune/MatrixGame2.0/finetune_df.sh b/examples/training/finetune/MatrixGame2.0/finetune_df.sh index 7f810d773..9d7f086ce 100644 --- a/examples/training/finetune/MatrixGame2.0/finetune_df.sh +++ b/examples/training/finetune/MatrixGame2.0/finetune_df.sh @@ -3,14 +3,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers" DATA_DIR="footsies-dataset/preprocessed/combined_parquet_dataset" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=8 -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( @@ -80,7 +77,6 @@ miscellaneous_args=( --ema_start_step 0 ) -# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t torchrun \ --nnodes 1 \ --nproc_per_node $NUM_GPUS \ diff --git a/examples/training/finetune/MatrixGame2.0/finetune_i2v.sh b/examples/training/finetune/MatrixGame2.0/finetune_i2v.sh index 07473407e..046f5517a 100644 --- a/examples/training/finetune/MatrixGame2.0/finetune_i2v.sh +++ b/examples/training/finetune/MatrixGame2.0/finetune_i2v.sh @@ -3,14 +3,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="FastVideo/Matrix-Game-2.0-Base-Diffusers" DATA_DIR="footsies-dataset/preprocessed/combined_parquet_dataset" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=8 -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh b/examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh index 575c87a2f..59f5e34ca 100644 --- a/examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh +++ b/examples/training/finetune/MatrixGame2.0/preprocess_matrixgame_data_i2v.sh @@ -5,7 +5,6 @@ MODEL_PATH="Matrix-Game-2.0-Base-Diffusers" DATA_MERGE_PATH="footsies-dataset/merge.txt" OUTPUT_DIR="footsies-dataset/preprocessed/" -# export CUDA_VISIBLE_DEVICES=0 export MASTER_ADDR=localhost export MASTER_PORT=29500 export RANK=0 diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh index 8695a516d..1d53118fc 100644 --- a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.sh @@ -3,14 +3,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" NUM_GPUS=8 -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( @@ -80,7 +77,6 @@ miscellaneous_args=( --enable_gradient_checkpointing_type "full" ) -# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t torchrun \ --nnodes 1 \ --nproc_per_node $NUM_GPUS \ diff --git a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm index 71907d5b9..ef72f0a22 100644 --- a/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm +++ b/examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/finetune_i2v.slurm @@ -30,7 +30,6 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -41,8 +40,6 @@ NUM_GPUS=8 MODEL_PATH="weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers" DATA_DIR="data/crush-smol_processed_i2v_1_3b_inp/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/training/finetune/Wan2.1-Fun-1.3B-InP/crush_smol/validation.json" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t diff --git a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm index f284a6076..a7e0b2e90 100644 --- a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm +++ b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/I2V-14B-VSA.slurm @@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -41,8 +40,6 @@ NUM_GPUS=8 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-720P-Diffusers" DATA_DIR=your_data_dir VALIDATION_DATASET_FILE=your_validation_dataset_file -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Core training arguments core_training_args=( diff --git a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm index 3a173f4f9..53c70e7f6 100644 --- a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm +++ b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA.slurm @@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -41,8 +40,6 @@ NUM_GPUS=8 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR=your_data_dir VALIDATION_DATASET_FILE=your_validation_dataset_file -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Core training arguments core_training_args=( diff --git a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA_4_gpu_debug.sh b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA_4_gpu_debug.sh index 453195761..96e875f54 100755 --- a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA_4_gpu_debug.sh +++ b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-1.3B-VSA_4_gpu_debug.sh @@ -1,6 +1,5 @@ -# Basic info export WANDB_MODE="online" export NCCL_P2P_DISABLE=1 export TORCH_NCCL_ENABLE_MONITORING=0 @@ -11,15 +10,13 @@ export NODE_RANK=0 export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN -export WANDB_API_KEY="your_wandb_api_key_here" # TODO: Replace with your actual key or load from a secure location +export WANDB_API_KEY="your_wandb_api_key_here" -# Configs NUM_GPUS=4 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR=data/Wan-Syn_77x448x832_600k VALIDATION_DATASET_FILE=examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json -# Core training arguments core_training_args=( --tracker_project_name wan_t2v_VSA --output_dir "checkpoints/wan_t2v_finetune_VSA" @@ -31,14 +28,12 @@ core_training_args=( --enable_gradient_checkpointing_type "full" # if OOM enable this ) -# Validation generation shape arguments (used during training validation) validation_generation_shape_args=( --num_height 448 --num_width 832 --num_frames 77 ) -# Parallel arguments parallel_args=( --num_gpus $NUM_GPUS --sp_size 1 @@ -47,19 +42,16 @@ parallel_args=( --hsdp_shard_dim 1 ) -# Model arguments model_args=( --model_path "$MODEL_PATH" --pretrained_model_name_or_path "$MODEL_PATH" ) -# Dataset arguments dataset_args=( --data_path "$DATA_DIR" --dataloader_num_workers 4 ) -# Validation arguments validation_args=( --log_validation --validation_dataset_file "$VALIDATION_DATASET_FILE" @@ -68,7 +60,6 @@ validation_args=( --validation_guidance_scale "5.0" ) -# Optimizer arguments optimizer_args=( --learning_rate 1e-6 --mixed_precision "bf16" @@ -78,7 +69,6 @@ optimizer_args=( --max_grad_norm 1.0 ) -# Miscellaneous arguments miscellaneous_args=( --inference_mode False --checkpoints_total_limit 3 @@ -89,7 +79,6 @@ miscellaneous_args=( --seed 1000 ) -# VSA arguments vsa_args=( --VSA_decay_rate 0.03 --VSA_decay_interval_steps 50 diff --git a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm index 0449b54f3..0241c2c38 100644 --- a/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm +++ b/examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/T2V-14B-VSA.slurm @@ -31,7 +31,6 @@ export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export FASTVIDEO_ATTENTION_BACKEND=VIDEO_SPARSE_ATTN -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -41,8 +40,6 @@ NUM_GPUS=8 MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" DATA_DIR=your_data_dir VALIDATION_DATASET_FILE=your_validation_dataset_file -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Core training arguments core_training_args=( diff --git a/examples/training/finetune/cosmos2_5/finetune_t2w.sh b/examples/training/finetune/cosmos2_5/finetune_t2w.sh index c8ecc7f3f..7e9c5bcf9 100755 --- a/examples/training/finetune/cosmos2_5/finetune_t2w.sh +++ b/examples/training/finetune/cosmos2_5/finetune_t2w.sh @@ -16,7 +16,6 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="KyleShao/Cosmos-Predict2.5-2B-Diffusers" DATA_DIR="data/cosmos2_5_processed_t2w/combined_parquet_dataset/" diff --git a/examples/training/finetune/cosmos2_5/finetune_t2w_lora.sh b/examples/training/finetune/cosmos2_5/finetune_t2w_lora.sh index e7f7cd071..5d47c341f 100755 --- a/examples/training/finetune/cosmos2_5/finetune_t2w_lora.sh +++ b/examples/training/finetune/cosmos2_5/finetune_t2w_lora.sh @@ -10,7 +10,6 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="KyleShao/Cosmos-Predict2.5-2B-Diffusers" DATA_DIR="data/cosmos2_5_processed_t2w/combined_parquet_dataset/" diff --git a/examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh b/examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh index 6f5590c7c..ea72da5c1 100755 --- a/examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh +++ b/examples/training/finetune/mmaudio/run_train_vggsound_ddp_from_scratch.sh @@ -7,9 +7,7 @@ set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" -# --------------------------------------------------------------------------- # Edit the DDP run here. Do not launch while another job owns the same GPUs. -# --------------------------------------------------------------------------- FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/train" VALIDATION_FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/val" EMPTY_STRING_FEATURES="${REPO_ROOT}/official_weights/mmaudio/raw/ext_weights/empty_string.pth" diff --git a/examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh b/examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh index d4b88e4ab..0dff04d99 100755 --- a/examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh +++ b/examples/training/finetune/mmaudio/run_train_vggsound_from_scratch.sh @@ -6,9 +6,7 @@ set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" -# --------------------------------------------------------------------------- # Edit the training run here. No environment variables are required. -# --------------------------------------------------------------------------- VARIANT="small_44k" FEATURE_DIR="/mnt/lustre/vlm-kai/datasets/VGGSound/mmaudio_features_torio/train" diff --git a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh index df9e3b68e..a6b6089f2 100644 --- a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh +++ b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh @@ -3,14 +3,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=8 -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( @@ -80,7 +77,6 @@ miscellaneous_args=( --enable_gradient_checkpointing_type "full" ) -# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t torchrun \ --nnodes 1 \ --nproc_per_node $NUM_GPUS \ diff --git a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm index 024ff49cd..2840a6655 100644 --- a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm +++ b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.slurm @@ -30,7 +30,6 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -41,8 +40,6 @@ NUM_GPUS=8 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/training/finetune/wan_i2v_14b_480p/crush_smol/validation.json" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t diff --git a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh index da9de3c85..3a2044c21 100644 --- a/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh +++ b/examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v_lora.sh @@ -2,14 +2,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" DATA_DIR="data/crush-smol_processed_i2v/combined_parquet_dataset/" VALIDATION_DIR="data/crush-smol_processed_i2v/validation_parquet_dataset/" NUM_GPUS=4 -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( @@ -79,7 +76,6 @@ miscellaneous_args=( --ema_start_step 0 ) -# If you do not have 32 GPUs and to fit in memory, you can: 1. increase sp_size. 2. reduce num_latent_t torchrun \ --nnodes 1 \ --nproc_per_node $NUM_GPUS \ diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh index 9407c213b..07c690d97 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh @@ -3,13 +3,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=4 -# export CUDA_VISIBLE_DEVICES=4,5 # Training arguments diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm index 03dbbbd99..ff648c09f 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.slurm @@ -30,7 +30,6 @@ export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID export TOKENIZERS_PARALLELISM=false export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA echo "MASTER_ADDR: $MASTER_ADDR" echo "NODE_RANK: $NODE_RANK" @@ -40,8 +39,6 @@ NUM_GPUS=8 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1_3b/crush_smol/validation.json" -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Training arguments training_args=( diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh index 26c12e0cc..5e512d647 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh @@ -2,13 +2,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=1 -# export CUDA_VISIBLE_DEVICES=4,5 # Training arguments diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh index 34908c56c..23ac35bff 100644 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh +++ b/examples/training/finetune/wan_t2v_1.3B/crush_smol/profile_finetune_t2v.sh @@ -3,13 +3,11 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" VALIDATION_DATASET_FILE="$(dirname "$0")/validation.json" NUM_GPUS=4 -# export CUDA_VISIBLE_DEVICES=4,5 mkdir -p ../profiler_traces/wan_t2v_finetune/ diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/api.cu b/fastvideo-kernel/attn_qat_infer/blackwell/api.cu index c7a3624f0..24d5b92db 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/api.cu +++ b/fastvideo-kernel/attn_qat_infer/blackwell/api.cu @@ -68,9 +68,7 @@ void set_params_fprop(Flash_fwd_params ¶ms, bool single_level_p_quant=false, bool seqlenq_ngroups_swapped=false) { - // Reset the parameters params = {}; - // Set the pointers and strides. params.q_ptr = q.data_ptr(); params.k_ptr = k.data_ptr(); params.v_ptr = v.data_ptr(); @@ -119,13 +117,11 @@ void set_params_fprop(Flash_fwd_params ¶ms, params.cu_seqlens_k = static_cast(cu_seqlens_k_d); params.seqused_k = static_cast(seqused_k); - // P = softmax(QK^T) params.p_ptr = p_d; // Softmax sum params.softmax_lse_ptr = softmax_lse_d; - // Set the dimensions. params.b = b; params.h = h; params.h_k = h_k; @@ -140,7 +136,6 @@ void set_params_fprop(Flash_fwd_params ¶ms, params.head_divmod = cutlass::FastDivmod(int(h)); - // Set the different scale values. params.scale_softmax = softmax_scale; params.scale_softmax_log2 = softmax_scale * M_LOG2E; __half scale_softmax_log2_half = __float2half(params.scale_softmax_log2); @@ -149,10 +144,6 @@ void set_params_fprop(Flash_fwd_params ¶ms, // Set this to probability of keeping an element to simplify things. params.p_dropout = 1.f - p_dropout; - // Convert p from float to int so we don't have to convert the random uint to float to compare. - // [Minor] We want to round down since when we do the comparison we use <= instead of < - // params.p_dropout_in_uint = uint32_t(std::floor(params.p_dropout * 4294967295.0)); - // params.p_dropout_in_uint16_t = uint16_t(std::floor(params.p_dropout * 65535.0)); params.p_dropout_in_uint8_t = uint8_t(std::floor(params.p_dropout * 255.0)); params.rp_dropout = 1.f / params.p_dropout; params.scale_softmax_rp_dropout = params.rp_dropout * params.scale_softmax; @@ -276,18 +267,12 @@ mha_fwd_impl(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head CHECK_SHAPE(q, batch_size, num_heads, seqlen_q, head_size_og); CHECK_SHAPE(k, batch_size, num_heads_k, seqlen_k, head_size_og); CHECK_SHAPE(v, batch_size, num_heads_k, unpacked_head_size, seqlen_k/2); - // CHECK_SHAPE(delta_s, batch_size, num_heads, seqlen_q / 128, seqlen_k); - // CHECK_SHAPE(sfq, batch_size, seqlen_q, num_heads, unpacked_head_size); - // CHECK_SHAPE(sfk, batch_size, seqlen_k, num_heads_k, unpacked_head_size); - // CHECK_SHAPE(sfv, batch_size, unpacked_head_size, num_heads_k, seqlen_k); TORCH_CHECK(unpacked_head_size % 8 == 0, "head_size must be a multiple of 8"); auto dtype = is_bf16 ? at::ScalarType::BFloat16 : at::ScalarType::Half; at::Tensor out = torch::empty({batch_size, num_heads, seqlen_q, unpacked_head_size}, opts.dtype(dtype)); auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; }; - // const int head_size = round_multiple(head_size_og, 8); - // const int head_size_rounded = round_multiple(head_size, 32); const int seqlen_q_rounded = round_multiple(seqlen_q, flash::BLOCK_M); const int seqlen_k_rounded = round_multiple(seqlen_k, flash::BLOCK_N); @@ -348,10 +333,6 @@ mha_fwd_impl(at::Tensor &q, // batch_size x seqlen_q x num_heads x (head // if (out_.has_value()) { out_.value().copy_(out); } // } - // return {out, q_padded, k_padded, v_padded, out_padded, softmax_lse, p}; - // cudaDeviceSynchronize(); - // auto err = cudaGetLastError(); - // printf("%s\n", cudaGetErrorString(err)); return {out, softmax_lse}; } diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/block_info.h b/fastvideo-kernel/attn_qat_infer/blackwell/block_info.h index bd0d45488..3c09a80dc 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/block_info.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/block_info.h @@ -20,7 +20,6 @@ namespace flash { -//////////////////////////////////////////////////////////////////////////////////////////////////// template struct BlockInfo { @@ -55,6 +54,5 @@ struct BlockInfo { const int actual_seqlen_k; }; -//////////////////////////////////////////////////////////////////////////////////////////////////// } // namespace flash diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/epilogue_tma_ws.h b/fastvideo-kernel/attn_qat_infer/blackwell/epilogue_tma_ws.h index 85e82a053..0d11333d4 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/epilogue_tma_ws.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/epilogue_tma_ws.h @@ -119,7 +119,6 @@ struct CollectiveEpilogueFwd{ auto smem_thr_copy_O = smem_tiled_copy_O.get_thread_slice(thread_idx); constexpr int numel = decltype(size(tOrO))::value; cutlass::NumericArrayConverter convert_op; - // HACK: this requires tensor to be "contiguous" auto frag = convert_op(*reinterpret_cast *>(tOrO.data())); auto tOrO_out = make_tensor(make_rmem_ptr(&frag), tOrO.layout()); Tensor taccOrO = smem_thr_copy_O.retile_S(tOrO_out); // ((Atom,AtomNum), MMA_M, MMA_N) @@ -183,7 +182,6 @@ struct CollectiveEpilogueFwd{ tma_store_wait<0>(); } - // Write 0 to output and -inf to LSE CUTLASS_DEVICE void store_zero( Params const& epilogue_params, @@ -202,9 +200,7 @@ struct CollectiveEpilogueFwd{ Tensor tOgO = gmem_thr_copy_O.partition_D(gO); Tensor tOrO = make_fragment_like(tOgO); clear(tOrO); - // Construct identity layout for sO Tensor cO = cute::make_identity_tensor(select<0, 2>(TileShape_MNK{})); // (BLK_M,BLK_K) -> (blk_m,blk_k) - // Repeat the partitioning with identity layouts Tensor tOcO = gmem_thr_copy_O.partition_D(cO); Tensor tOpO = make_tensor(make_shape(size<2>(tOgO))); #pragma unroll diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h b/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h index d237499c3..6979b3e75 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/kernel_ws.h @@ -97,7 +97,6 @@ __global__ void __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp, CollectiveEpilogue::prefetch_tma_descriptors(epilogue_params); } - // Obtain warp index PipelineParams pipeline_params_v; pipeline_params_v.transaction_bytes = CollectiveMainloop::TmaTransactionBytesV; @@ -142,7 +141,7 @@ __global__ void __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp, cutlass::arch::warpgroup_reg_dealloc<24>(); TileScheduler scheduler; - if (producer_warp_role == ProducerWarpRole::Mainloop) { // Load Q, K, V + if (producer_warp_role == ProducerWarpRole::Mainloop) { PipelineStateQ smem_pipe_write_q = cutlass::make_producer_start_state(); PipelineState smem_pipe_write_k = cutlass::make_producer_start_state(); PipelineState smem_pipe_write_v = cutlass::make_producer_start_state(); @@ -179,8 +178,6 @@ __global__ void __launch_bounds__(Ktraits::kNWarps * cutlass::NumThreadsPerWarp, for (auto work_tile_info = scheduler.get_initial_work(); work_tile_info.is_valid(scheduler_params); work_tile_info = scheduler.get_next_work(scheduler_params, work_tile_info)) { // Attention output (GEMM-II) accumulator. Tensor tOrO = partition_fragment_C(tiled_mma_pv, select<0, 2>(TileShape_MNK{})); - // flash::Softmax<2 * (2 * kBlockM / NumMmaThreads)> softmax; - // Pass single_level_p_quant flag to control P quantization mode flash::SoftmaxFused<2 * (2 * kBlockM / NumMmaThreads)> softmax_fused(params.single_level_p_quant); auto block_coord = work_tile_info.get_block_coord(scheduler_params); auto [m_block, bidh, bidb] = block_coord; diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/launch.h b/fastvideo-kernel/attn_qat_infer/blackwell/launch.h index edc167afc..86f31fda2 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/launch.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/launch.h @@ -79,7 +79,6 @@ void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { num_blocks_m = cutlass::ceil_div(num_blocks_m, size<0>(ClusterShape{})) * size<0>(ClusterShape{}); typename Scheduler::Arguments scheduler_args = {num_blocks_m, params.h, params.b}; typename Scheduler::Params scheduler_params = Scheduler::to_underlying_arguments(scheduler_args); - // Get the ptr to kernel function. void *kernel; kernel = (void *)flash::compute_attn_ws; int smem_size = sizeof(typename Kernel_traits::SharedStorage); diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h b/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h index 51d9ff23a..b3bdcf83e 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/mainloop_tma_ws.h @@ -683,7 +683,6 @@ struct CollectiveMainloopFwd { Tensor tOrSFP = make_tensor(LayoutSFP{}); Tensor tOrSFP_flt = filter_zeros(tOrSFP); Tensor tSrDS = make_tensor(make_shape(_8{}, _4{}), make_stride(_1{}, _8{})); - // copy qk and sf from smem to rmem auto smem_tiled_copy_Q = make_tiled_copy_A(SmemCopyAtomQ{}, tiled_mma_qk); auto smem_thr_copy_Q = smem_tiled_copy_Q.get_thread_slice(thread_idx); Tensor tSsQ = smem_thr_copy_Q.partition_S(as_position_independent_swizzle_tensor(sQ)); diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/tile_scheduler.h b/fastvideo-kernel/attn_qat_infer/blackwell/tile_scheduler.h index a559bef07..38c3dd7a2 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/tile_scheduler.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/tile_scheduler.h @@ -65,7 +65,6 @@ public: // MSVC requires protecting use of CUDA-specific nonstandard syntax, // like blockIdx and gridDim, with __CUDA_ARCH__. #if defined(__CUDA_ARCH__) - // current_work_linear_idx_ = blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y; current_work_linear_idx_ = blockIdx.x; #else CUTLASS_ASSERT(false && "This line should never be reached"); diff --git a/fastvideo-kernel/attn_qat_infer/blackwell/utils.h b/fastvideo-kernel/attn_qat_infer/blackwell/utils.h index 793270196..f4e6b4cfd 100644 --- a/fastvideo-kernel/attn_qat_infer/blackwell/utils.h +++ b/fastvideo-kernel/attn_qat_infer/blackwell/utils.h @@ -136,7 +136,6 @@ __forceinline__ __device__ __half2 half_exp(__half2 x) { return out; } -// Apply the exp to all the elements. template __forceinline__ __device__ void max_scale_exp2_sum(Tensor &tensor, Tensor &max, Tensor &sum, const float scale) { static_assert(Layout0::rank == 2, "Only support 2D Tensor"); static_assert(Layout1::rank == 1, "Only support 1D Tensor"); CUTE_STATIC_ASSERT_V(size<0>(max) == size<0>(tensor)); @@ -164,7 +163,6 @@ __forceinline__ __device__ void max_scale_exp2_sum(Tensor &ten } } -// Apply the exp to all the elements. template __forceinline__ __device__ void scale_apply_exp2(Tensor &tensor, Tensor const &max, const float scale) { static_assert(Layout0::rank == 2, "Only support 2D Tensor"); diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py index a17885ec2..e3d5e05cc 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_quant_v.py @@ -79,7 +79,6 @@ if n % 128 != 0: else: q_padded = q -# use torch transpose + scaled_fp4_quant to get the ground truth q_padded = q_padded.transpose(2, 3).reshape(b, h, n_padded, d).contiguous() o_gt = torch.empty((b, h, n_padded, d // 2), dtype=torch.uint8, device='cuda') o_s_gt = torch.empty((b, h, n_padded, d // 16), dtype=torch.float8_e4m3fn, device='cuda') diff --git a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py index e757efcad..48350a912 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py +++ b/fastvideo-kernel/attn_qat_infer/quantization/bench/bench_utils.py @@ -127,11 +127,9 @@ def bench_kineto(fn, if not using_nsys: profiler.step() - # Return 1 if using Nsight Systems if using_nsys: return 1 - # Parse the profiling table assert isinstance(kernel_names, str) or isinstance(kernel_names, tuple) is_tupled = isinstance(kernel_names, tuple) prof_lines = profiler.key_averages().table(sort_by='cuda_time_total', max_name_column_width=100).split('\n') @@ -140,11 +138,9 @@ def bench_kineto(fn, for name in kernel_names: assert sum([name in line for line in prof_lines]) == 1, f'Errors of the kernel {name} in the profiling table' - # Save chrome traces if trace_path is not None: profiler.export_chrome_trace(trace_path) - # Return average kernel times units = {'ms': 1e3, 'us': 1e6} kernel_times = [] for name in kernel_names: diff --git a/fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu b/fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu index bf8a9bad5..f13dbf9e4 100644 --- a/fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu +++ b/fastvideo-kernel/attn_qat_infer/quantization/fp4_quantization_4d.cu @@ -125,7 +125,6 @@ using Type = __nv_bfloat162; }; - // Define a 32 bytes packed data type. template struct PackedVec { typename TypeConverter::Type elts[8]; @@ -152,7 +151,6 @@ constexpr uint32_t NUM_THREADS_PER_TOKEN = head_dim / CVT_FP4_ELTS_PER_THREAD; - // load input const int token_id = token_block_id * BLOCK_SIZE + threadIdx.x / NUM_THREADS_PER_TOKEN; int load_token_id; @@ -183,7 +181,6 @@ (threadIdx.x % NUM_THREADS_PER_TOKEN) * CVT_FP4_ELTS_PER_THREAD)[0]; // feature dim } - // calculate max of every consecutive 16 elements auto localMax = __habs2(in_vec.elts[0]); #pragma unroll for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { // local max @@ -204,7 +201,6 @@ float SFValueInv = (SFValue == 0.0f) ? 0.0f : 1.0f / SFValue; - // convert input to float2 and apply scale float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; #pragma unroll @@ -218,7 +214,6 @@ fp2Vals[i].y = fp2Vals[i].y * SFValueInv; } - // convert to e2m1 uint32_t e2m1Vals[CVT_FP4_ELTS_PER_THREAD / 8]; #pragma unroll for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 8; i++) { @@ -229,7 +224,6 @@ // is not a multiple of BLOCK_SIZE (out-of-bounds write fix). if (token_id >= num_tokens) return; - // save if constexpr (CVT_FP4_ELTS_PER_THREAD == 8) { reinterpret_cast(output + batch_id * stride_bz_output + @@ -284,7 +278,6 @@ constexpr uint32_t NUM_THREADS_PER_TOKEN = head_dim / CVT_FP4_ELTS_PER_THREAD; constexpr uint32_t NUM_THREADS_PER_SEQ = BLOCK_SIZE / CVT_FP4_ELTS_PER_THREAD; - // load input const int token_id = token_block_id * BLOCK_SIZE + threadIdx.x / NUM_THREADS_PER_TOKEN; PackedVec in_vec; @@ -312,7 +305,6 @@ in_vec.elts[i].y = shared_input[(threadIdx.x / NUM_THREADS_PER_SEQ) + ((threadIdx.x % NUM_THREADS_PER_SEQ) * CVT_FP4_ELTS_PER_THREAD + 2 * i + 1) * head_dim]; } - // calculate max of every consecutive 16 elements auto localMax = __habs2(in_vec.elts[0]); #pragma unroll for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { // local max @@ -333,7 +325,6 @@ float SFValueInv = (SFValue == 0.0f) ? 0.0f : 1.0f / SFValue; - // convert input to float2 and apply scale float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; #pragma unroll @@ -347,7 +338,6 @@ fp2Vals[i].y = fp2Vals[i].y * SFValueInv; } - // convert to e2m1 uint32_t e2m1Vals[CVT_FP4_ELTS_PER_THREAD / 8]; #pragma unroll for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 8; i++) { @@ -360,7 +350,6 @@ (threadIdx.x % NUM_THREADS_PER_SEQ) * CVT_FP4_ELTS_PER_THREAD; if (write_token_id >= num_tokens) return; - // save if constexpr (CVT_FP4_ELTS_PER_THREAD == 8) { reinterpret_cast(output + batch_id * stride_bz_output + diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index a0d6a6797..5d8511f4a 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -224,8 +224,6 @@ export CMAKE_ARGS echo "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST:-}" echo "CMAKE_ARGS: ${CMAKE_ARGS:-}" echo "GPU_BACKEND: ${GPU_BACKEND:-}" -# Build and install -# Use -v for verbose output if [ -n "${WHEEL_DIR}" ]; then mkdir -p "${WHEEL_DIR}" uv build --wheel -v --no-build-isolation --out-dir "${WHEEL_DIR}" . diff --git a/fastvideo-kernel/csrc/attention/block_sparse_h100.cu b/fastvideo-kernel/csrc/attention/block_sparse_h100.cu index e561e5ac7..1596e87b9 100644 --- a/fastvideo-kernel/csrc/attention/block_sparse_h100.cu +++ b/fastvideo-kernel/csrc/attention/block_sparse_h100.cu @@ -65,7 +65,7 @@ template struct fwd_globals { template __global__ __launch_bounds__(128, 4) -void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block size of 64 +void fwd_attend_ker(const __grid_constant__ fwd_globals g) { #if FASTVIDEO_TK_HOPPER extern __shared__ int __shm[]; tma_swizzle_allocator al((int*)&__shm[0]); @@ -138,7 +138,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block siz // wait k wait(k_smem_arrived, kv_idx % 2); - // compute QK^T warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]); copy(max_vec_last_scaled, max_vec); @@ -147,7 +146,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block siz warpgroup::mma_async_wait(); - // load K if (threadIdx.x == 0) { tma::expect_bytes(k_smem_arrived, sizeof(k_tile)); coord k_tile_idx = {blockIdx.z, kv_head_idx, kv_block_index, 0}; @@ -180,11 +178,9 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block siz // wait v wait(v_smem_arrived, kv_idx % 2); - // compute SV warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]); warpgroup::mma_async_wait(); - // load V if (threadIdx.x == 0) { tma::expect_bytes(v_smem_arrived, sizeof(v_tile)); coord v_tile_idx = {blockIdx.z, kv_head_idx, kv_block_index, 0}; @@ -198,7 +194,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block siz // wait k wait(k_smem_arrived, kv_idx % 2); - // compute QK^T warpgroup::mm_ABt(att_block, q_smem[0], k_smem[0]); copy(max_vec_last_scaled, max_vec); @@ -234,7 +229,6 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { // use block siz // wait v wait(v_smem_arrived, kv_idx % 2); - // compute SV warpgroup::mma_AB(o_reg, att_block_mma, v_smem[0]); warpgroup::mma_async_wait(); } @@ -550,7 +544,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { if constexpr (D == 64) { mul(ds_block_t, ds_block_t, 0.125f); } else { mul(ds_block_t, ds_block_t, 0.08838834764f); } - // load vec if (threadIdx.x == 0) { coord vec_idx = {blockIdx.z, blockIdx.y, 0, load_q_block_index}; tma::expect_bytes(vec_b[0], sizeof(l_smem[0]) + sizeof(d_smem[0])); @@ -564,7 +557,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { warpgroup::store(ds_smem_t[0], ds_block_t); warpgroup::mma_async_wait(); - // load og if (threadIdx.x == 0) { coord tile_idx = {blockIdx.z, blockIdx.y, load_q_block_index, 0}; tma::expect_bytes(o_b[0], sizeof(og_smem[0])); @@ -575,7 +567,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { warpgroup::mma_commit_group(); warpgroup::mma_async_wait(); - // load q if (threadIdx.x == 0) { coord q_tile_idx = {blockIdx.z, blockIdx.y, load_q_block_index, 0}; tma::expect_bytes(q_b[0], sizeof(q_smem[0])); @@ -587,11 +578,9 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK warpgroup::mma_commit_group(); warpgroup::mma_async_wait(); - // store qg to shared memory warpgroup::store(qg_smem, qg_reg); __syncthreads(); - // store and add dQ to global memory if (threadIdx.x / 32 == 0) { coord tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0}; tma::store_add_async(g.qg, qg_smem, tile_idx); @@ -643,11 +632,9 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { warpgroup::mm_AtB(qg_reg, ds_smem_t[0], k_smem[0]); //delat dQ = dSK warpgroup::mma_commit_group(); warpgroup::mma_async_wait(); - // store qg to shared memory warpgroup::store(qg_smem, qg_reg); __syncthreads(); - // store and add dQ to global memory if (threadIdx.x / 32 == 0) { coord tile_idx = {blockIdx.z, blockIdx.y, store_qg_block_index, 0}; tma::store_add_async(g.qg, qg_smem, tile_idx); @@ -655,7 +642,6 @@ void bwd_attend_ker(const __grid_constant__ bwd_globals g) { } } - // store kq and vq // ! the following two line seems unnecessary. // tma::store_async_wait(); // ensure qg is finished @@ -768,7 +754,6 @@ block_sparse_attention_forward( float* l_ptr = reinterpret_cast(l_vec.data_ptr()); float* d_l = reinterpret_cast(l_ptr); - //cudadevicesynchronize(); const c10::cuda::OptionalCUDAGuard device_guard(q.device()); const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); @@ -820,7 +805,6 @@ block_sparse_attention_forward( fwd_attend_ker<64><<>>(g); CHECK_CUDA_ERROR(cudaGetLastError()); - // cudaStreamSynchronize(stream); } if (head_dim == 128) { @@ -871,11 +855,9 @@ block_sparse_attention_forward( fwd_attend_ker<128><<>>(g); CHECK_CUDA_ERROR(cudaGetLastError()); - // cudaStreamSynchronize(stream); } return {o, l_vec}; - //cudadevicesynchronize(); } std::vector @@ -996,11 +978,9 @@ block_sparse_attention_backward(torch::Tensor q, constexpr int mem_size = kittens::MAX_SHARED_MEMORY; int threads = PREP_NUM_WARPS * kittens::WARP_THREADS; - //cudadevicesynchronize(); const c10::cuda::OptionalCUDAGuard device_guard(q.device()); const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - // cudaStreamSynchronize(stream); // TORCH_CHECK(seq_len % (4*kittens::TILE_DIM*4) == 0, "sequence length must be divisible by 256"); dim3 grid_bwd(q_seq_len/(PREP_NUM_WARPS*kittens::TILE_ROW_DIM*4), qo_heads, batch); @@ -1084,7 +1064,6 @@ block_sparse_attention_backward(torch::Tensor q, dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch); threads = 128; - //cudadevicesynchronize(); { cudaFuncSetAttribute( @@ -1101,12 +1080,6 @@ block_sparse_attention_backward(torch::Tensor q, bwd_attend_ker<64><<>>(bwd_global); } - // CHECK_CUDA_ERROR(cudaGetLastError()); - // cudaStreamSynchronize(stream); - //cudadevicesynchronize(); - // const auto kernel_end = std::chrono::high_resolution_clock::now(); - // std::cout << "Kernel Time: " << std::chrono::duration_cast(kernel_end - start).count() << "us" << std::endl; - // std::cout << "---" << std::endl; } if (head_dim == 128) { @@ -1188,7 +1161,6 @@ block_sparse_attention_backward(torch::Tensor q, dim3 grid_bwd_2(kv_seq_len/BLOCK_N, qo_heads, batch); threads = 128; - //cudadevicesynchronize(); { cudaFuncSetAttribute( @@ -1200,10 +1172,7 @@ block_sparse_attention_backward(torch::Tensor q, bwd_attend_ker<128><<>>(bwd_global); } - // cudaStreamSynchronize(stream); - //cudadevicesynchronize(); } return {qg, kg, vg}; - //cudadevicesynchronize(); } \ No newline at end of file diff --git a/fastvideo-kernel/csrc/attention/st_attn_h100.cu b/fastvideo-kernel/csrc/attention/st_attn_h100.cu index 5aab72acf..dd5d0eef6 100644 --- a/fastvideo-kernel/csrc/attention/st_attn_h100.cu +++ b/fastvideo-kernel/csrc/attention/st_attn_h100.cu @@ -310,9 +310,7 @@ void fwd_attend_ker(const __grid_constant__ fwd_globals g) { else { mul(max_vec_last_scaled, max_vec_last_scaled, 1.44269504089f*0.08838834764f); } warpgroup::mma_async_wait(); - // apply non-pad mask int offset = g.text_L - (kv_idx - (kv_iters + 1)) * K::kv_height; - // printf("k_idx_start: %d, k_idx_end: %d, text_end: %d, offset: %d\n", k_idx_start, k_idx_end, text_end, offset); right_fill(att_block, att_block, offset, base_types::constants::neg_infty()); @@ -437,7 +435,6 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, float* l_ptr = reinterpret_cast(l_vec.data_ptr()); float* d_l = reinterpret_cast(l_ptr); - //cudadevicesynchronize(); const c10::cuda::OptionalCUDAGuard device_guard(q.device()); const cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); @@ -578,10 +575,8 @@ sta_forward(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, } CHECK_CUDA_ERROR(cudaGetLastError()); - // cudaStreamSynchronize(stream); } return o; - //cudadevicesynchronize(); } diff --git a/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu b/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu index 724537e51..c1bb6268a 100644 --- a/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu +++ b/fastvideo-kernel/csrc/comm/ulysses_all_to_all.cu @@ -267,7 +267,6 @@ void ulysses_a2a_tuned(int64_t handle, torch::Tensor inp, torch::Tensor out, int #undef LAUNCH_ULYSSES_A2A TORCH_CHECK(cudaGetLastError() == cudaSuccess, "ulysses_a2a kernel launch failed"); - // Copy this rank's completed result out of the window. auto status = cudaMemcpyAsync(out.data_ptr(), ctx->buf, out_bytes, cudaMemcpyDeviceToDevice, stream); TORCH_CHECK(status == cudaSuccess, "ulysses_a2a copy-out failed: ", cudaGetErrorString(status)); diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index 1d0462808..12c80e577 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -2,7 +2,6 @@ #include #include -// Forward declarations #ifdef TK_COMPILE_ST_ATTN extern torch::Tensor sta_forward( torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor o, diff --git a/fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp b/fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp index 7a0eef8dd..c982cdb68 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp +++ b/fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp @@ -295,7 +295,6 @@ struct GemmKernel { SmemLayoutB{} ); - // register TiledMma tiled_mma; auto thr_mma = tiled_mma.get_slice(t_idx); auto tCrA = thr_mma.partition_fragment_A(gA(_, _, 0)); @@ -314,7 +313,6 @@ struct GemmKernel { clear(tDrD); - // global to shared copy G2STiledCopyA g2s_tiled_copy_a; auto g2s_thr_copy_a = g2s_tiled_copy_a.get_slice(t_idx); auto tAgA = g2s_thr_copy_a.partition_S(gA); @@ -371,7 +369,6 @@ struct GemmKernel { int32_t fallbackA_load = 0; int32_t fallbackA_mma = 0; - // copy first Stage - 1 tile CUTLASS_PRAGMA_UNROLL for (int i = 0, _i = min(Stage - 1, ntile); i < _i; ++i) { if (g2s_b_tile < ntile) { @@ -410,8 +407,6 @@ struct GemmKernel { // a s2r increase anyway s2r_a_smem = (s2r_a_smem + 1) % Stage; - // get next s2r b tile int64_t - // end of a block // dequant first dequant( @@ -425,7 +420,6 @@ struct GemmKernel { mma_block_a = mma_block_b; if (mma_block_a < size<1>(AS)) scale_a = AS(blk_m_coord, mma_block_a); - // load next stage if (g2s_b_tile < ntile) { g2s_tile_in_block = (g2s_tile_in_block + 1) % NumTilePerBlock; copy_AB(g2s_tiled_copy_a, tAgA, tAsA, tAcA, g2s_a_tile, g2s_a_smem, m_limit); diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp index 733e27550..4e396a80e 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp @@ -61,7 +61,6 @@ public: int tidx = threadIdx.x; float x[NumElementPerThread]; - // load Loader loader; loader.load(params.Iptr, x, params.m, params.n, blk_m, 0, tidx); @@ -86,7 +85,6 @@ public: } __syncthreads(); - // var reduction float v = sqrtf(_reduce_square(x, shared_data) / params.n + params.eps); CUTLASS_PRAGMA_UNROLL @@ -94,7 +92,6 @@ public: x[i] /= v; if constexpr (Affine) { - // load weight Loader weight_loader; float w[NumElementPerThread]; weight_loader.load(params.Wptr, w, 1, params.n, 0, 0, tidx); @@ -111,7 +108,6 @@ public: } } - // save y { Saver saver; if constexpr (std::is_same_v) { diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp index 55796c92d..8fe7fa148 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp @@ -56,14 +56,12 @@ public: int tidx = threadIdx.x; float x[NumElementPerThread]; - // load Loader loader; loader.load(params.Iptr, x, params.m, params.n, blk_m, 0, tidx); // rms reduction float rms = sqrtf(_reduce_square(x, shared_data) / params.n + params.eps); - // load weight Loader weight_loader; float w[NumElementPerThread]; loader.load(params.Wptr, w, 1, params.n, 0, 0, tidx); @@ -73,7 +71,6 @@ public: for (int i = 0; i < NumElementPerThread; ++i) x[i] = w[i] * x[i] / rms ; - // save y OutputDtype *output_reg = (OutputDtype*)x; if constexpr (!std::is_same_v) { output_reg = (OutputDtype*)w; diff --git a/fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp index 6c85b5035..e7e30d8bf 100644 --- a/fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp +++ b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp @@ -109,7 +109,6 @@ public: float float_reg[NumElementPerThread]; - // load float32 data Loader loader; loader.load(params.Iptr, float_reg, params.m, params.n, blk_m, blk_n, tidx); quantization( diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py index 7279708e4..661b06dfe 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/attn_qat_train.py @@ -9,7 +9,6 @@ import triton.language as tl from triton.tools.tensor_descriptor import TensorDescriptor from .quant_utils import fake_quantize_q, fake_quantize_kv, fake_quantize -# from attn_qat_infer.api import triton_group_mean def is_cuda(): @@ -141,7 +140,6 @@ def _attn_fwd_inner(acc, elif STAGE == 2: lo, hi = start_m * BLOCK_M, (start_m + 1) * BLOCK_M lo = tl.multiple_of(lo, BLOCK_M) - # causal = False else: lo, hi = 0, N_CTX offsetk_y = offset_y + lo # offset from the start of the current batch-head combination @@ -149,7 +147,6 @@ def _attn_fwd_inner(acc, offsetv_y = offset_y * HEAD_DIM + lo else: offsetv_y = offset_y + lo - # loop over k, v and update accumulator for start_n in tl.range(lo, hi, BLOCK_N, warp_specialize=warp_specialize): start_n = tl.multiple_of(start_n, BLOCK_N) kv_valid = (start_n + offs_n) < N_CTX @@ -191,7 +188,6 @@ def _attn_fwd_inner(acc, acc = acc * alpha[:, None] if IS_QAT: high_prec_acc = high_prec_acc * alpha[:, None] - # prepare p and v for the dot if dtype == tl.float8e5: v = desc_v.load([0, offsetv_y]).T else: @@ -298,7 +294,6 @@ configs = [ for s in NUM_STAGES_OPTIONS for w in [4, 8] ] if "PYTEST_VERSION" in os.environ: - # Use a single config in testing for reproducibility configs = [ triton.Config(dict(BLOCK_M=128, BLOCK_N=64), num_stages=2, num_warps=4, pre_hook=_host_descriptor_pre_hook), ] @@ -329,8 +324,6 @@ def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape): return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape) -# @triton.autotune(configs=list(filter(keep, configs)), key=["N_CTX_Q", "N_CTX_KV", "HEAD_DIM", "FP8_OUTPUT", "warp_specialize"], -# prune_configs_by={'early_config_prune': prune_invalid_configs}, cache_results=True) @triton.jit def _attn_fwd( sm_scale, @@ -398,11 +391,9 @@ def _attn_fwd( offset_y_q = off_z * (N_CTX_Q * H) + off_h * N_CTX_Q # offset for query tensor offset_y_kv = off_z * (N_CTX_KV * H) + off_h * N_CTX_KV # offset for key/value tensors qo_offset_y = offset_y_q + start_m * BLOCK_M - # initialize offsets offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = tl.arange(0, BLOCK_N) # for kv blocks q_valid = offs_m < N_CTX_Q - # initialize pointer to m and l m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) @@ -411,10 +402,8 @@ def _attn_fwd( else: # dummy value high_prec_acc = acc - # load scales qk_scale = sm_scale qk_scale *= 1.44269504 # 1/log(2) - # load q: it will stay in SRAM throughout q = desc_q.load([qo_offset_y, 0]) # load from start of q block and start of the entire head dimension q = tl.where(q_valid[:, None], q, 0.0) # stage 1: off-band @@ -434,7 +423,6 @@ def _attn_fwd( BLOCK_N, 4 - STAGE, offs_m, offs_n, N_CTX_KV, warp_specialize, IS_HOPPER, IS_QAT, fake_quant_P, two_level_quant_P, use_global_sf_P, JOIN_QAT_PV) - # stage 2: on-band if STAGE & 2: acc, high_prec_acc, l_i, m_i = _attn_fwd_inner(acc, high_prec_acc, l_i, m_i, q, q_valid, desc_k, desc_v, offset_y_kv, dtype, start_m, qk_scale, BLOCK_M, HEAD_DIM, @@ -505,12 +493,10 @@ def _attn_bwd_preprocess(O, DO, Delta, Z, H, N_CTX, BLOCK_M: tl.constexpr, HEAD_ off_hz = tl.program_id(1) off_n = tl.arange(0, HEAD_DIM) valid = off_m < N_CTX - # load o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None]) do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :], mask=valid[:, None]).to(tl.float32) delta = tl.sum(o * do, axis=1) - # write-back tl.store(Delta + off_hz * N_CTX + off_m, delta, mask=valid) @@ -573,7 +559,6 @@ def _attn_bwd_dkdv( qk = tl.where(mask, qk, -1.0e6) p = tl.math.exp2(qk - m[:, None]) do = tl.load(do_ptrs, mask=q_valid[:, None]) - # Compute dV. p_quant = p if IS_QAT and fake_quant_P: p_quant, _ = fake_quantize(src_tensor=p, @@ -586,7 +571,6 @@ def _attn_bwd_dkdv( dv += tl.dot(tl.trans(p_quant.to(tl.bfloat16)), do) # D (= delta) is pre-divided by ds_scale. Di = tl.load(D + offs_m, mask=q_valid) - # Compute dP and dS. dp = tl.dot(do, tl.trans(v)) ds = p * (dp - Di[:, None]) ds = ds.to(tl.bfloat16) @@ -659,7 +643,6 @@ def _attn_bwd_dq( mask = (offs_m[:, None] >= offs_n[None, :]) qk = tl.where(mask, qk, -1.0e6) p = tl.math.exp2(qk - m) - # Compute dP and dS. dp = tl.dot(do, tl.trans(v)) ds = p * (dp - Di[:, None]) # Compute dQ. @@ -1037,7 +1020,6 @@ def _attn_bwd( if SMOOTH_Q: Q_MEAN += adj - # load scales offs_k = tl.arange(0, HEAD_DIM) start_n = pid * BLOCK_N1 @@ -1050,7 +1032,6 @@ def _attn_bwd( dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32) dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32) - # load K and V: they stay in SRAM throughout the inner loop. k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=kv_valid[:, None]) v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d, mask=kv_valid[:, None]) @@ -1101,7 +1082,6 @@ def _attn_bwd( start_m = 0 num_steps = (N_CTX + BLOCK_M1 - 1) // BLOCK_M1 - # Compute dK and dV for non-masked blocks. dk, dv = _attn_bwd_dkdv(dk, dv, Q, @@ -1133,7 +1113,6 @@ def _attn_bwd( dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d tl.store(dv_ptrs, dv, mask=kv_valid[:, None]) - # Write back dK. dk *= sm_scale dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d tl.store(dk_ptrs, dk, mask=kv_valid[:, None]) @@ -1185,14 +1164,12 @@ def _attn_bwd( warp_specialize=warp_specialize, ) end_n -= num_steps * MASK_BLOCK_N2 - # stage 2: process KV blocks from 0 to end_n (before diagonal), ceiling division to include remainder num_steps = (end_n + BLOCK_N2 - 1) // BLOCK_N2 start_n = 0 else: # For non-causal, process all KV blocks from 0 to N_CTX (ceiling division to include remainder) num_steps = (N_CTX + BLOCK_N2 - 1) // BLOCK_N2 start_n = 0 - # stage 2 dq = _attn_bwd_dq( dq, q, @@ -1299,7 +1276,6 @@ class _attention(torch.autograd.Function): if IS_QAT: high_prec_o = torch.empty_like(q) else: - # Initialize to a dummy value high_prec_o = o stage = 3 if causal else 1 extra_kern_args = {} @@ -1442,7 +1418,6 @@ class _attention(torch.autograd.Function): use_global_sf=use_global_sf_QKV, ) - # Apply pre-hook to set block shapes on tensor descriptors _host_descriptor_pre_hook({ "BLOCK_M": fwd_block_m, "BLOCK_N": fwd_block_n, @@ -1647,7 +1622,6 @@ class _attention(torch.autograd.Function): num_stages=dkdv_launch[1], ) elif N_CTX_Q == N_CTX_KV: - # Use existing kernel for self-attention (same sequence lengths) grid = ((N_CTX_KV + BLOCK_N1 - 1) // BLOCK_N1, 1, BATCH * N_HEAD) _attn_bwd[grid](q, arg_k, @@ -1684,7 +1658,6 @@ class _attention(torch.autograd.Function): num_warps=NUM_WARPS, num_stages=NUM_STAGES) else: - # Use separate kernels for cross-attention (different sequence lengths) grid_dq = ((N_CTX_Q + BLOCK_M2 - 1) // BLOCK_M2, 1, BATCH * N_HEAD) _attn_bwd_dq_cross[grid_dq]( q, diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py index 22d0b13e5..30b317271 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/block_sparse_attn_triton.py @@ -183,11 +183,9 @@ def _attn_bwd_preprocess( off_m = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) off_hz = tl.program_id(1) off_n = tl.arange(0, HEAD_DIM) - # load o = tl.load(O + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]) do = tl.load(DO + off_hz * HEAD_DIM * N_CTX + off_m[:, None] * HEAD_DIM + off_n[None, :]).to(tl.float32) delta = tl.sum(o * do, axis=1) - # write-back tl.store(Delta + off_hz * N_CTX + off_m, delta) @@ -255,13 +253,11 @@ def _attn_bwd_dkdv( pT = tl.where(mask[:, None], pT, 0.0) do = tl.load(do_ptrs + block_sparse_offset * stride_tok) - # Compute dV. ppT = pT ppT = ppT.to(tl.bfloat16) dv += tl.dot(ppT, do) # D (= delta) is pre-divided by ds_scale. Di = tl.load(D + offs_m) - # Compute dP and dS. dpT = tl.dot(v, tl.trans(do)).to(tl.float32) dsT = pT * (dpT - Di[None, :]) dsT = dsT.to(tl.bfloat16) @@ -333,7 +329,6 @@ def _attn_bwd_dq( offs_in_block = half * step_n + tl.arange(0, BLOCK_N2) mask = offs_in_block < block_size p = tl.where(mask[None, :], p, 0.0) - # Compute dP and dS. dp = tl.dot(do, vT).to(tl.float32) ds = p * (dp - Di[:, None]) ds = ds.to(tl.bfloat16) @@ -392,7 +387,6 @@ def _attn_bwd( M += off_chz D += off_chz - # load scales offs_k = tl.arange(0, HEAD_DIM) start_n = pid * BLOCK_N1 @@ -403,7 +397,6 @@ def _attn_bwd( dv = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32) dk = tl.zeros([BLOCK_N1, HEAD_DIM], dtype=tl.float32) - # load K and V: they stay in SRAM throughout the inner loop. k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d) v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d) @@ -438,7 +431,6 @@ def _attn_bwd( dv_ptrs = DV + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d tl.store(dv_ptrs, dv) - # Write back dK. dk *= sm_scale dk_ptrs = DK + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d tl.store(dk_ptrs, dk) @@ -481,7 +473,6 @@ def _attn_bwd( end_n, num_steps # ) - # Write back dQ. dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d dq *= sm_scale tl.store(dq_ptrs, dq) @@ -557,7 +548,6 @@ def _attn_bwd_dkdv_kernel( start_n = kv_blk * BLOCK_N1 offs_n = start_n + tl.arange(0, BLOCK_N1) - # load K and V: they stay in SRAM throughout the inner loop. k = tl.load(K + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d) v = tl.load(V + offs_n[:, None] * stride_tok + offs_k[None, :] * stride_d) diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py index 7abc9f79a..4bd4a3bb9 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py @@ -165,7 +165,6 @@ def _compute_dequant( if dst_dtype == tl.float16: dst_scale = dst_scale.to(tl.float16) - # Now upcast the tensor. intermediate_dtype: tl.constexpr = tl.bfloat16 if dst_dtype == tl.float32 else dst_dtype if cuda_capability_geq(10, 0): assert is_fp4 @@ -215,7 +214,6 @@ def _compute_dequant( # 2) x is subnormal (x == 0bs001 where s is the sign): Map to +-0.5 in the dst type x0 = tl.where(em0 == 0x01, dst_0p5 | (x0 & 0x8000), x0) x1 = tl.where(em1 == 0x10, dst_0p5 | (x1 & 0x8000), x1) - # 3) x is zero, do nothing dst_tensor = tl.interleave(x0, x1).to(intermediate_dtype, bitcast=True) dst_tensor = dst_tensor.to(dst_dtype) diff --git a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py index 4ea18f2d2..46cefc58d 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/st_attn_triton.py @@ -78,12 +78,10 @@ def _attn_fwd_loop( exp_scores = tl.math.exp2(scores - new_m[:, None]) current_l = tl.sum(exp_scores, axis=1) - # Update L <- L * exp(M - M') + L1, M <- M' alpha = tl.math.exp2(m - new_m) l = l * alpha + current_l m = new_m - # Update O <- O * exp(M - M') + P @ V acc = (acc * alpha[:, None] + tl.dot(exp_scores.to(v.type.element_ty), v)) return m, l, acc @@ -286,8 +284,6 @@ def sliding_tile_attention_triton(q: torch.Tensor, total_tile_size = tile_t * tile_h * tile_w - # BLOCK_Q=128 - # BLOCK_KV=128 BLOCK_DIM = head_dim output = torch.empty_like(q) @@ -329,8 +325,6 @@ def sliding_tile_attention_triton(q: torch.Tensor, scale=1.0 / (head_dim**0.5), has_text=has_text, text_q=False, - # BLOCK_Q=BLOCK_Q, - # BLOCK_KV=BLOCK_KV, BLOCK_DIM=BLOCK_DIM, ) @@ -363,8 +357,6 @@ def sliding_tile_attention_triton(q: torch.Tensor, scale=1.0 / (head_dim**0.5), has_text=has_text, text_q=True, - # BLOCK_Q=BLOCK_Q, - # BLOCK_KV=BLOCK_KV, BLOCK_DIM=BLOCK_DIM, ) diff --git a/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py b/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py index 092e2f55f..aecdb158f 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py +++ b/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py @@ -129,21 +129,17 @@ def _rms_norm_fwd_fused( x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32) - # Compute variance _var = x * x var = tl.sum(_var, axis=1) / N rstd = 1 / tl.sqrt(var + eps) - # Write mean / rstd tl.store(Rstd + rows, rstd) rstd = tl.reshape(rstd, (BLOCK_M, 1)) - # Normalize and apply linear transformation w = tl.load(W + cols) x_hat = x * rstd y = x_hat * w - # Write output y = y.to(Y.type.element_ty) tl.store(y_ptr, y, mask=mask[None, :]) @@ -181,7 +177,6 @@ def rmsnorm(x, w, eps): else: BLOCK_M = 1 - # Call the triton kernel _rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M), )]( # x, y, @@ -229,20 +224,16 @@ def _layer_norm_param_fwd_fused( x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32) - # Compute mean and Variance mean = tl.sum(x, axis=1, keep_dims=True) / N - # Compute variance _var = (x - mean) * (x - mean) var = tl.sum(_var, axis=1, keep_dims=True) / N rstd = 1 / tl.sqrt(var + eps) - # Write mean / rstd _mean = tl.reshape(mean, (BLOCK_M)) _rstd = tl.reshape(rstd, (BLOCK_M)) tl.store(Mean + rows, _mean) tl.store(Rstd + rows, _rstd) - # Normalize and apply linear transformation x_hat = (x - mean) * rstd w = tl.load(W + cols) @@ -250,7 +241,6 @@ def _layer_norm_param_fwd_fused( x_hat = x_hat * w + b - # Write output x_hat = x_hat.to(Y.type.element_ty) tl.store(y_ptr, x_hat, mask=mask[None, :]) @@ -298,9 +288,6 @@ def layernorm_param(x, w, b, eps): return y, mean, rstd -######################################################## -# Elementwise_affine=False -######################################################## @triton.jit @@ -327,23 +314,18 @@ def _layer_norm_noparam_fwd_fused( x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32) - # Compute mean and Variance mean = tl.sum(x, axis=1, keep_dims=True) / N - # Compute variance _var = (x - mean) * (x - mean) var = tl.sum(_var, axis=1, keep_dims=True) / N rstd = 1 / tl.sqrt(var + eps) - # Write mean / rstd _mean = tl.reshape(mean, (BLOCK_M)) _rstd = tl.reshape(rstd, (BLOCK_M)) tl.store(Mean + rows, _mean) tl.store(Rstd + rows, _rstd) - # Normalize and apply linear transformation x_hat = (x - mean) * rstd - # Write output x_hat = x_hat.to(Y.type.element_ty) tl.store(y_ptr, x_hat, mask=mask[None, :]) @@ -440,7 +422,6 @@ class Int8LinearFunction(torch.autograd.Function): grad_bias = None if ctx.needs_input_grad[0]: - # grad_input = grad_output @ W w_float = dequantize_weight(w_q, w_s, grad_output.dtype) grad_input = torch.matmul(grad_output, w_float) @@ -489,34 +470,15 @@ class RMSNormFunction(torch.autograd.Function): grad_output_flat = grad_output dx = dx - # Standard RMSNorm backward - # dy = grad_output_flat - # x_hat = x * rstd - # w * dy - # c1 = mean(dy * w * x) * rstd^2 - # dx = (dy * w - c1 * x) * rstd - # More precise: - # y = x * rstd * w - # dy = grad_output - # dw = sum(dy * x * rstd, dim=0) grad_w = None if ctx.needs_input_grad[1]: # w - # grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0) - # More accurate gradient for weight when considering rstd was computed from x - # Actually, for Affine part (w), it is just sum(dL/dy * x_hat). - # y = x_hat * w - # dL/dw = sum(dL/dy * x_hat) x_hat = x_flat * rstd grad_w = (grad_output_flat * x_hat).sum(dim=0) grad_input = None if ctx.needs_input_grad[0]: # x - # x_hat = x * rstd - # dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat) - # but for RMSNorm: - # dx = rstd * (w * dy - (x * rstd^2) * mean(w * dy * x)) -> check formula # Using PyTorch autograd for reference logic: # y = x / sqrt(mean(x^2) + eps) * w @@ -530,8 +492,6 @@ class RMSNormFunction(torch.autograd.Function): # Let's compute: dy_w = grad_output_flat * w - # term2 = (dy_w * x_flat).sum(dim=-1, keepdim=True) / N # mean(dy*w*x) - # grad_input = rstd * (dy_w - x_flat * (rstd ** 2) * term2) # Re-derivation: # x_hat = x * rstd @@ -586,10 +546,6 @@ class LayerNormFunction(torch.autograd.Function): else: w_eff = w - # dx calculation - # x_hat = (x - mean) * rstd - # y = x_hat * w + b - # dL/dx_hat = dL/dy * w dy = grad_output_flat dx_hat = dy * w_eff diff --git a/fastvideo-kernel/python/fastvideo_kernel/vmoba.py b/fastvideo-kernel/python/fastvideo_kernel/vmoba.py index 4149a59f0..c84ae7037 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/vmoba.py +++ b/fastvideo-kernel/python/fastvideo_kernel/vmoba.py @@ -5,7 +5,7 @@ import random import time import torch try: - from flash_attn import flash_attn_varlen_func # Use the new flash attention function + from flash_attn import flash_attn_varlen_func from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward except ImportError: @@ -46,7 +46,6 @@ def calc_chunks(cu_seqlen, moba_chunk_size): chunk_to_batch[cu_num_chunk[1:-1]] = 1 chunk_to_batch = chunk_to_batch.cumsum(dim=0, dtype=torch.int32) - # Do not filter out any chunk filtered_chunk_indices = torch.arange(num_chunk, device=cu_seqlen.device, dtype=torch.int32) num_filtered_chunk = num_chunk @@ -66,7 +65,7 @@ def _select_threshold_query_head(gate: torch.Tensor, valid_gate_mask: torch.Tens eps = 1e-6 # LSE‐style normalization per (across chunks) - gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min row_min = gate_min_val.amin(dim=0) # (H, S) @@ -127,7 +126,7 @@ def _select_threshold_block(gate: torch.Tensor, valid_gate_mask: torch.Tensor, g eps = 1e-6 # LSE‐style normalization per block (across heads and queries) - gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min block_max = gate_masked.amax(dim=(1, 2), keepdim=True) # (C, 1, 1) @@ -193,7 +192,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, eps = 1e-6 # LSE‐style normalization globally across all valid entries - gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) # Use -inf for max + gate_masked = torch.where(valid_gate_mask, gate, -torch.inf) gate_min_val = torch.where(valid_gate_mask, gate, torch.inf) # Use +inf for min overall_max = gate_masked.max() # scalar @@ -251,7 +250,7 @@ def _select_threshold_overall(gate: torch.Tensor, valid_gate_mask: torch.Tensor, # 8) create the mask in the original flat shape others_mask_flat = torch.zeros_like(others_flat, dtype=torch.bool) # (C*H*S,) - if selected_other_indices.numel() > 0: # Check if any 'other' indices were selected + if selected_other_indices.numel() > 0: others_mask_flat[selected_other_indices] = True others_mask = others_mask_flat.view(C, H, S) # (C, H, S) @@ -349,8 +348,6 @@ class MixedAttention(torch.autograd.Function): ctx.moba_chunk_size = moba_chunk_size ctx.softmax_scale = softmax_scale = q.shape[-1]**(-0.5) - # Non-causal self-attention branch - # return out, softmax_lse, S_dmask, rng_state self_attn_out_sh, self_attn_lse_hs, _, _ = _flash_attn_varlen_forward( q=q, k=k, @@ -532,12 +529,10 @@ def moba_attn_varlen( kv = torch.stack((k, v), dim=1) seqlen, num_head, head_dim = q.shape - # Compute chunk boundaries. cu_chunk, filtered_chunk_indices, num_filtered_chunk, chunk_to_batch = calc_chunks(cu_seqlens, moba_chunk_size) self_attn_cu_seqlen = cu_chunk - # Update top-k selection to include the self chunk. moba_topk = min(moba_topk, num_filtered_chunk) # --- Build filtered KV from chunks --- @@ -553,16 +548,13 @@ def moba_attn_varlen( gathered = kv[indices.view(-1)].view(num_filtered_chunk, max_chunk_len, *kv.shape[1:]) gathered = gathered * valid_mask.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1).type_as(gathered) - # Compute key_gate_weight over valid tokens. key_values = gathered[:, :, 0].float() # [num_filtered_chunk, max_chunk_len, num_head, head_dim] valid_mask_exp = valid_mask.unsqueeze(-1).unsqueeze(-1) key_sum = (key_values * valid_mask_exp).sum(dim=1) divisor = valid_mask.sum(dim=1).unsqueeze(-1).unsqueeze(-1) key_gate_weight = key_sum / divisor # [num_filtered_chunk, num_head, head_dim] - # Compute gate logits between key_gate_weight and queries. q_float = q.float() - # gate = torch.einsum("nhd,shd->nhs", key_gate_weight, q_float) # [num_filtered_chunk, num_head, seqlen] gate = torch.bmm(key_gate_weight.permute(1, 0, 2), q_float.permute(1, 0, 2).transpose(1, 2)).permute(1, 0, 2) # Amplify the diagonal (self chunk) contributions. @@ -650,7 +642,6 @@ def moba_attn_varlen( moba_q_sh_indices = (moba_q_indices % seqlen) * num_head + (moba_q_indices // seqlen) moba_q = rearrange(q, "s h d -> (h s) d").index_select(0, moba_q_indices).unsqueeze(1) - # Build cumulative sequence lengths for the selected queries. moba_seqlen_q = gate_mask.sum(dim=-1).flatten() q_zero_mask = moba_seqlen_q == 0 valid_expert_mask = ~q_zero_mask @@ -814,7 +805,6 @@ def test_attn_varlen_moba_speed(batch, simsum_threshold=0.25, threshold_type='query_head'): """Speed test comparing flash_attn vs moba_attention""" - # Get data q, k, v, cu_seqlen, max_seqlen = generate_data(batch, seqlen, head, head_dim, dtype) print( f"batch:{batch} head:{head} seqlen:{seqlen} chunk:{moba_chunk_size} topk:{moba_topk} select_mode: {select_mode} simsum_threshold:{simsum_threshold}" diff --git a/fastvideo-kernel/tests/test_attn_qat_train.py b/fastvideo-kernel/tests/test_attn_qat_train.py index e221b50aa..6d9949c8d 100644 --- a/fastvideo-kernel/tests/test_attn_qat_train.py +++ b/fastvideo-kernel/tests/test_attn_qat_train.py @@ -53,15 +53,12 @@ def attn_qat_train_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=Non if sm_scale is None: sm_scale = 1.0 / sqrt(q_BLHD.shape[-1]) - # Convert from BLHD to BHLD format q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() - # Call attention with BHLD format o_BHLD = attention(q_BHLD, k_BHLD, v_BHLD, is_causal, sm_scale) - # Convert back from BHLD to BLHD format return o_BHLD.permute(0, 2, 1, 3).contiguous() @@ -84,7 +81,6 @@ def fused_attn_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=None, w if sm_scale is None: sm_scale = 1.0 / sqrt(q_BLHD.shape[-1]) - # Convert from BLHD to BHLD format q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -93,7 +89,6 @@ def fused_attn_wrapper(q_BLHD, k_BLHD, v_BLHD, is_causal=False, sm_scale=None, w # Note: fused_attention expects inputs in BHLD format and only supports same sequence lengths o_BHLD = fused_attention(q_BHLD, k_BHLD, v_BHLD, is_causal, sm_scale, warp_specialize) - # Convert back from BHLD to BLHD format return o_BHLD.permute(0, 2, 1, 3).contiguous() @@ -112,7 +107,6 @@ def cosine_similarity(tensor1, tensor2): t1_flat = tensor1.flatten().float() t2_flat = tensor2.flatten().float() - # Compute cosine similarity: (A · B) / (||A|| * ||B||) dot_product = torch.dot(t1_flat, t2_flat) norm1 = torch.norm(t1_flat) norm2 = torch.norm(t2_flat) @@ -139,7 +133,6 @@ def naive_attention(q, k, v, causal, sm_scale): Returns: Output tensor of shape (Z, H, N_CTX_Q, HEAD_DIM) """ - # Compute attention scores: QK^T scores = torch.matmul(q, k.transpose(-2, -1)) * sm_scale # Apply causal mask if needed @@ -156,10 +149,8 @@ def naive_attention(q, k, v, causal, sm_scale): # In cross-attention, typically no causal masking is applied pass - # Apply softmax attn_weights = torch.softmax(scores, dim=-1) - # Apply attention weights to values out = torch.matmul(attn_weights, v) return out @@ -175,7 +166,6 @@ def test_qat_attention_forward(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - # Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) @@ -183,18 +173,15 @@ def test_qat_attention_forward(): # Test with QAT (using wrapper that does permute/contiguous) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - # Check that outputs have correct shape assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) - # Check that outputs are finite assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() @@ -206,9 +193,6 @@ def test_qat_attention_forward(): f" QAT vs Naive Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM} - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # QAT output should be reasonably close to naive (within tolerance due to quantization) - # Using a reasonable tolerance for float16 and quantization effects - # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close, got max_diff={max_diff.item():.6f}" print("✓ Forward pass test passed.") @@ -223,12 +207,10 @@ def test_qat_attention_backward(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format for QAT q_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) @@ -250,7 +232,6 @@ def test_qat_attention_backward(): # Backward pass for naive out_naive_BLHD.backward(dout) - # Check that gradients are computed assert q_qat_BLHD.grad is not None assert k_qat_BLHD.grad is not None assert v_qat_BLHD.grad is not None @@ -258,17 +239,14 @@ def test_qat_attention_backward(): assert k_naive_BHLD.grad is not None assert v_naive_BHLD.grad is not None - # Check that gradients have correct shape assert q_qat_BLHD.grad.shape == q_qat_BLHD.shape assert k_qat_BLHD.grad.shape == k_qat_BLHD.shape assert v_qat_BLHD.grad.shape == v_qat_BLHD.shape - # Check that gradients are finite assert torch.isfinite(q_qat_BLHD.grad).all() assert torch.isfinite(k_qat_BLHD.grad).all() assert torch.isfinite(v_qat_BLHD.grad).all() - # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() @@ -292,10 +270,6 @@ def test_qat_attention_backward(): f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" ) - # Gradients should be reasonably close (within tolerance due to quantization) - # assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}" - # assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}" - # assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}" print("✓ Backward pass test passed.") @@ -313,7 +287,6 @@ def test_qat_attention_different_shapes(): causal = True for Z, H, N_CTX, HEAD_DIM in test_configs: - # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) @@ -321,7 +294,6 @@ def test_qat_attention_different_shapes(): sm_scale = 1.0 / sqrt(HEAD_DIM) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -341,8 +313,6 @@ def test_qat_attention_different_shapes(): f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs should be reasonably close - # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}" print(f"✓ Shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})") @@ -356,14 +326,12 @@ def test_qat_attention_non_causal(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -383,8 +351,6 @@ def test_qat_attention_non_causal(): f" QAT vs Naive (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs should be reasonably close - # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}" print("✓ Non-causal test passed.") @@ -398,14 +364,12 @@ def test_qat_attention_causal(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -425,8 +389,6 @@ def test_qat_attention_causal(): f" QAT vs Naive (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs should be reasonably close - # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}" print("✓ Causal test passed.") @@ -440,12 +402,10 @@ def test_qat_attention_causal_backward(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format for QAT q_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_qat_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) @@ -457,7 +417,6 @@ def test_qat_attention_causal_backward(): out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - # Create dummy gradient dout = torch.randn_like(out_qat_BLHD).contiguous() # Backward pass for QAT @@ -466,12 +425,10 @@ def test_qat_attention_causal_backward(): # Backward pass for naive out_naive_BHLD.backward(dout) - # Check that gradients are computed assert q_qat_BLHD.grad is not None assert k_qat_BLHD.grad is not None assert v_qat_BLHD.grad is not None - # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() @@ -513,7 +470,6 @@ def test_qat_attention_different_seq_lengths(): for Z, H, N_CTX_Q, N_CTX_KV, HEAD_DIM, causal in test_configs: sm_scale = 1.0 / sqrt(HEAD_DIM) - # Create input tensors in BLHD format with different sequence lengths q_BLHD = torch.randn((Z, N_CTX_Q, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE) @@ -521,7 +477,6 @@ def test_qat_attention_different_seq_lengths(): # Test with QAT (using wrapper that does permute/contiguous) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -532,7 +487,6 @@ def test_qat_attention_different_seq_lengths(): assert out_qat_BLHD.shape == (Z, N_CTX_Q, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX_Q, H, HEAD_DIM) - # Check that outputs are finite assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() @@ -544,8 +498,6 @@ def test_qat_attention_different_seq_lengths(): f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}, causal={causal}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs should be reasonably close - # assert max_diff < 1.0, f"QAT and naive outputs should be reasonably close for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}), got max_diff={max_diff.item():.6f}" print(f"✓ Different sequence lengths test passed for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV})") @@ -564,12 +516,10 @@ def test_qat_attention_different_seq_lengths_backward(): for Z, H, N_CTX_Q, N_CTX_KV, HEAD_DIM, causal in test_configs: sm_scale = 1.0 / sqrt(HEAD_DIM) - # Create input tensors in BLHD format with different sequence lengths q_qat_BLHD = torch.randn((Z, N_CTX_Q, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_qat_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_qat_BLHD = torch.randn((Z, N_CTX_KV, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_qat_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) @@ -581,7 +531,6 @@ def test_qat_attention_different_seq_lengths_backward(): out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - # Create dummy gradient dout = torch.randn_like(out_qat_BLHD).contiguous() # Backward pass for QAT @@ -590,7 +539,6 @@ def test_qat_attention_different_seq_lengths_backward(): # Backward pass for naive out_naive_BHLD.backward(dout) - # Check that gradients are computed assert q_qat_BLHD.grad is not None assert k_qat_BLHD.grad is not None assert v_qat_BLHD.grad is not None @@ -598,17 +546,14 @@ def test_qat_attention_different_seq_lengths_backward(): assert k_naive_BHLD.grad is not None assert v_naive_BHLD.grad is not None - # Check that gradients have correct shape assert q_qat_BLHD.grad.shape == q_qat_BLHD.shape assert k_qat_BLHD.grad.shape == k_qat_BLHD.shape assert v_qat_BLHD.grad.shape == v_qat_BLHD.shape - # Check that gradients are finite assert torch.isfinite(q_qat_BLHD.grad).all() assert torch.isfinite(k_qat_BLHD.grad).all() assert torch.isfinite(v_qat_BLHD.grad).all() - # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() @@ -632,10 +577,6 @@ def test_qat_attention_different_seq_lengths_backward(): f" (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV}) - dV: Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" ) - # Gradients should be reasonably close (within tolerance due to quantization) - # assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}" - # assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}" - # assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}" print(f"✓ Cross-attention backward test passed for (N_CTX_Q={N_CTX_Q}, N_CTX_KV={N_CTX_KV})") @@ -650,7 +591,6 @@ def test_fused_attention_forward(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - # Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) @@ -658,18 +598,15 @@ def test_fused_attention_forward(): # Test with fused attention (using wrapper that does permute/contiguous) out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() out_naive_BHLD = naive_attention(q_BHLD.clone(), k_BHLD.clone(), v_BHLD.clone(), causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - # Check that outputs have correct shape assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) - # Check that outputs are finite assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() @@ -681,8 +618,6 @@ def test_fused_attention_forward(): f" Fused vs Naive - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Fused attention should be reasonably close to naive (within tolerance for float16) - # assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close, got max_diff={max_diff.item():.6f}" print("✓ Fused attention forward pass test passed.") @@ -697,12 +632,10 @@ def test_fused_attention_backward(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format for fused attention q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) @@ -714,7 +647,6 @@ def test_fused_attention_backward(): out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - # Create dummy gradient (same for both, in BLHD format) dout = torch.randn_like(out_fused_BLHD).contiguous() # Backward pass for fused attention @@ -723,7 +655,6 @@ def test_fused_attention_backward(): # Backward pass for naive out_naive_BLHD.backward(dout) - # Check that gradients are computed assert q_fused_BLHD.grad is not None assert k_fused_BLHD.grad is not None assert v_fused_BLHD.grad is not None @@ -731,17 +662,14 @@ def test_fused_attention_backward(): assert k_naive_BHLD.grad is not None assert v_naive_BHLD.grad is not None - # Check that gradients have correct shape assert q_fused_BLHD.grad.shape == q_fused_BLHD.shape assert k_fused_BLHD.grad.shape == k_fused_BLHD.shape assert v_fused_BLHD.grad.shape == v_fused_BLHD.shape - # Check that gradients are finite assert torch.isfinite(q_fused_BLHD.grad).all() assert torch.isfinite(k_fused_BLHD.grad).all() assert torch.isfinite(v_fused_BLHD.grad).all() - # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() @@ -765,10 +693,6 @@ def test_fused_attention_backward(): f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" ) - # Gradients should be reasonably close (within tolerance for float16) - # assert dq_diff.max() < 1e-1, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}" - # assert dk_diff.max() < 1e-1, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}" - # assert dv_diff.max() < 1e-1, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}" print("✓ Fused attention backward pass test passed.") @@ -787,7 +711,6 @@ def test_fused_attention_different_shapes(): causal = True for Z, H, N_CTX, HEAD_DIM in test_configs: - # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) @@ -795,7 +718,6 @@ def test_fused_attention_different_shapes(): sm_scale = 1.0 / sqrt(HEAD_DIM) out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -815,8 +737,6 @@ def test_fused_attention_different_shapes(): f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs should be reasonably close - # assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}" print(f"✓ Fused attention shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})") @@ -830,14 +750,12 @@ def test_fused_attention_non_causal(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -857,8 +775,6 @@ def test_fused_attention_non_causal(): f" Fused vs Naive (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs should be reasonably close - # assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}" print("✓ Fused attention non-causal test passed.") @@ -872,14 +788,12 @@ def test_fused_attention_causal(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) out_fused_BLHD = fused_attn_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -899,8 +813,6 @@ def test_fused_attention_causal(): f" Fused vs Naive (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs should be reasonably close - # assert max_diff < 1e-1, f"Fused and naive outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}" print("✓ Fused attention causal test passed.") @@ -914,12 +826,10 @@ def test_fused_attention_causal_backward(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format for fused attention q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE, requires_grad=True) - # Create input tensors for naive (same values, convert to BHLD) q_naive_BHLD = q_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) k_naive_BHLD = k_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) v_naive_BHLD = v_fused_BLHD.clone().permute(0, 2, 1, 3).contiguous().detach().requires_grad_(True) @@ -931,7 +841,6 @@ def test_fused_attention_causal_backward(): out_naive_BHLD = naive_attention(q_naive_BHLD, k_naive_BHLD, v_naive_BHLD, causal, sm_scale) out_naive_BHLD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() - # Create dummy gradient dout = torch.randn_like(out_fused_BLHD).contiguous() # Backward pass for fused attention @@ -940,12 +849,10 @@ def test_fused_attention_causal_backward(): # Backward pass for naive out_naive_BHLD.backward(dout) - # Check that gradients are computed assert q_fused_BLHD.grad is not None assert k_fused_BLHD.grad is not None assert v_fused_BLHD.grad is not None - # Convert naive gradients to BLHD format for comparison q_naive_grad_BLHD = q_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() k_naive_grad_BLHD = k_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() v_naive_grad_BLHD = v_naive_BHLD.grad.permute(0, 2, 1, 3).contiguous() @@ -996,14 +903,11 @@ def test_fused_vs_qat_attention_forward(): v_qat_BLHD = v_BLHD.clone().to(torch.bfloat16) out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - # Convert QAT output back to float16 for comparison out_qat_BLHD = out_qat_BLHD.to(torch.bfloat16) - # Check that outputs have correct shape assert out_fused_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) - # Check that outputs are finite assert torch.isfinite(out_fused_BLHD).all() assert torch.isfinite(out_qat_BLHD).all() @@ -1015,8 +919,6 @@ def test_fused_vs_qat_attention_forward(): f" Fused vs QAT - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # They may differ due to quantization in QAT and different implementations - # assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close, got max_diff={max_diff.item():.6f}" print("✓ Fused vs QAT forward pass test passed.") @@ -1030,12 +932,10 @@ def test_fused_vs_qat_attention_backward(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format for fused attention (float16) q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) - # Create input tensors for QAT (float16, same values) q_qat_BLHD = q_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) k_qat_BLHD = k_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) v_qat_BLHD = v_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) @@ -1046,7 +946,6 @@ def test_fused_vs_qat_attention_backward(): # Forward pass with QAT out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - # Create dummy gradient (same for both, in BLHD format) dout = torch.randn_like(out_fused_BLHD).contiguous() dout_qat = dout.to(torch.bfloat16) @@ -1056,7 +955,6 @@ def test_fused_vs_qat_attention_backward(): # Backward pass for QAT out_qat_BLHD.backward(dout_qat) - # Check that gradients are computed assert q_fused_BLHD.grad is not None assert k_fused_BLHD.grad is not None assert v_fused_BLHD.grad is not None @@ -1064,17 +962,14 @@ def test_fused_vs_qat_attention_backward(): assert k_qat_BLHD.grad is not None assert v_qat_BLHD.grad is not None - # Check that gradients have correct shape assert q_fused_BLHD.grad.shape == q_fused_BLHD.shape assert k_fused_BLHD.grad.shape == k_fused_BLHD.shape assert v_fused_BLHD.grad.shape == v_fused_BLHD.shape - # Check that gradients are finite assert torch.isfinite(q_fused_BLHD.grad).all() assert torch.isfinite(k_fused_BLHD.grad).all() assert torch.isfinite(v_fused_BLHD.grad).all() - # Convert QAT gradients to float16 for comparison q_qat_grad_f16 = q_qat_BLHD.grad.to(torch.bfloat16) k_qat_grad_f16 = k_qat_BLHD.grad.to(torch.bfloat16) v_qat_grad_f16 = v_qat_BLHD.grad.to(torch.bfloat16) @@ -1098,10 +993,6 @@ def test_fused_vs_qat_attention_backward(): f" dV - Max diff: {dv_diff.max().item():.6f}, Mean diff: {dv_diff.mean().item():.6f}, Cosine sim: {dv_cos_sim:.6f}" ) - # Gradients may differ due to quantization in QAT and different implementations - # assert dq_diff.max() < 2.0, f"dQ gradients should be reasonably close, got max_diff={dq_diff.max().item():.6f}" - # assert dk_diff.max() < 2.0, f"dK gradients should be reasonably close, got max_diff={dk_diff.max().item():.6f}" - # assert dv_diff.max() < 2.0, f"dV gradients should be reasonably close, got max_diff={dv_diff.max().item():.6f}" print("✓ Fused vs QAT backward pass test passed.") @@ -1120,7 +1011,6 @@ def test_fused_vs_qat_attention_different_shapes(): for Z, H, N_CTX, HEAD_DIM in test_configs: sm_scale = 1.0 / sqrt(HEAD_DIM) - # Create input tensors in BLHD format (float16) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) @@ -1148,8 +1038,6 @@ def test_fused_vs_qat_attention_different_shapes(): f" (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs may differ due to quantization in QAT - # assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for shape (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}), got max_diff={max_diff.item():.6f}" print(f"✓ Fused vs QAT shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})") @@ -1162,7 +1050,6 @@ def test_fused_vs_qat_attention_non_causal(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = False - # Create input tensors in BLHD format (float16) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) @@ -1190,8 +1077,6 @@ def test_fused_vs_qat_attention_non_causal(): f" Fused vs QAT (non-causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs may differ due to quantization in QAT - # assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for non-causal, got max_diff={max_diff.item():.6f}" print("✓ Fused vs QAT non-causal test passed.") @@ -1204,7 +1089,6 @@ def test_fused_vs_qat_attention_causal(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format (float16) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE) @@ -1232,8 +1116,6 @@ def test_fused_vs_qat_attention_causal(): f" Fused vs QAT (causal) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}" ) - # Outputs may differ due to quantization in QAT - # assert max_diff < 1.0, f"Fused and QAT outputs should be reasonably close for causal, got max_diff={max_diff.item():.6f}" print("✓ Fused vs QAT causal test passed.") @@ -1246,12 +1128,10 @@ def test_fused_vs_qat_attention_causal_backward(): sm_scale = 1.0 / sqrt(HEAD_DIM) causal = True - # Create input tensors in BLHD format for fused attention (float16) q_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) k_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) v_fused_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=torch.bfloat16, device=DEVICE, requires_grad=True) - # Create input tensors for QAT (float16, same values) q_qat_BLHD = q_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) k_qat_BLHD = k_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) v_qat_BLHD = v_fused_BLHD.clone().to(torch.bfloat16).detach().requires_grad_(True) @@ -1262,7 +1142,6 @@ def test_fused_vs_qat_attention_causal_backward(): # Forward pass with QAT out_qat_BLHD = attn_qat_train_wrapper(q_qat_BLHD, k_qat_BLHD, v_qat_BLHD, causal, sm_scale) - # Create dummy gradient dout = torch.randn_like(out_fused_BLHD).contiguous() dout_qat = dout.to(torch.bfloat16) @@ -1272,7 +1151,6 @@ def test_fused_vs_qat_attention_causal_backward(): # Backward pass for QAT out_qat_BLHD.backward(dout_qat) - # Check that gradients are computed assert q_fused_BLHD.grad is not None assert k_fused_BLHD.grad is not None assert v_fused_BLHD.grad is not None @@ -1280,7 +1158,6 @@ def test_fused_vs_qat_attention_causal_backward(): assert k_qat_BLHD.grad is not None assert v_qat_BLHD.grad is not None - # Convert QAT gradients to float16 for comparison q_qat_grad_f16 = q_qat_BLHD.grad.to(torch.bfloat16) k_qat_grad_f16 = k_qat_BLHD.grad.to(torch.bfloat16) v_qat_grad_f16 = v_qat_BLHD.grad.to(torch.bfloat16) @@ -1319,7 +1196,6 @@ def test_qat_attention_wan_shape_forward(): print(f" Testing WAN shape: Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM}") - # Create input tensors in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM) q_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) k_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) v_BLHD = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) @@ -1327,7 +1203,6 @@ def test_qat_attention_wan_shape_forward(): # Test with QAT (using wrapper that does permute/contiguous) out_qat_BLHD = attn_qat_train_wrapper(q_BLHD.clone(), k_BLHD.clone(), v_BLHD.clone(), causal, sm_scale) - # Convert to BHLD for naive attention comparison q_BHLD = q_BLHD.permute(0, 2, 1, 3).contiguous() k_BHLD = k_BLHD.permute(0, 2, 1, 3).contiguous() v_BHLD = v_BLHD.permute(0, 2, 1, 3).contiguous() @@ -1338,11 +1213,9 @@ def test_qat_attention_wan_shape_forward(): print(f" Expected shape: {(Z, N_CTX, H, HEAD_DIM)}") print(f" out_qat_BLHD shape: {out_qat_BLHD.shape}") - # Check that outputs have correct shape assert out_qat_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) assert out_naive_BLHD.shape == (Z, N_CTX, H, HEAD_DIM) - # Check that outputs are finite assert torch.isfinite(out_qat_BLHD).all() assert torch.isfinite(out_naive_BLHD).all() @@ -1393,7 +1266,6 @@ def test_qat_attention_wan_shape_backward(): # ----------------------------- out_naive_BHLD = naive_attention(q_BHLD, k_BHLD, v_BHLD, causal, sm_scale) - # Convert naive output back to BLHD for comparing outputs out_naive_BLHD = out_naive_BHLD.permute(0, 2, 1, 3).contiguous() # ----------------------------- @@ -1597,23 +1469,8 @@ if __name__ == "__main__": print("Running Fused attention tests...") print() - # test_fused_attention_forward() - # test_fused_attention_backward() - # test_fused_attention_different_shapes() - # test_fused_attention_non_causal() - # test_fused_attention_causal() - # test_fused_attention_causal_backward() - # print() - # print("Running Fused vs QAT attention comparison tests...") - # print() - # test_fused_vs_qat_attention_forward() - # test_fused_vs_qat_attention_backward() - # test_fused_vs_qat_attention_different_shapes() - # test_fused_vs_qat_attention_non_causal() - # test_fused_vs_qat_attention_causal() - # test_fused_vs_qat_attention_causal_backward() print() print("All tests passed! ✓") diff --git a/fastvideo-kernel/tests/test_sta.py b/fastvideo-kernel/tests/test_sta.py index 02b0bd4ef..6ac019c21 100644 --- a/fastvideo-kernel/tests/test_sta.py +++ b/fastvideo-kernel/tests/test_sta.py @@ -2,7 +2,6 @@ import torch from .support_flex_sta import get_sliding_tile_attention_mask from fastvideo_kernel import sliding_tile_attention from torch.nn.attention.flex_attention import flex_attention -# from flash_attn_interface import flash_attn_func from tqdm import tqdm flex_attention = torch.compile(flex_attention, dynamic=False) @@ -83,7 +82,6 @@ def test_sliding_tile_attention(): mean = 1e-1 std = 10 - # Run correctness check directly results = check_correctness(b, h, n, d, causal, mean, std, error_mode='output') assert results['TK vs FLEX'][ 'avg_diff'] < 3e-6, f"Average difference: {results['TK vs FLEX']['avg_diff']} is too large" diff --git a/fastvideo-kernel/tests/test_turbodiffusion.py b/fastvideo-kernel/tests/test_turbodiffusion.py index 02ae43bdf..676489a34 100644 --- a/fastvideo-kernel/tests/test_turbodiffusion.py +++ b/fastvideo-kernel/tests/test_turbodiffusion.py @@ -3,8 +3,6 @@ import pytest import sys import os -# Ensure local package is imported -# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../python"))) try: from fastvideo_kernel._C import fastvideo_kernel_ops @@ -65,20 +63,13 @@ class TestTurboDiffusion: M, N, K = 32, 64, 128 x = torch.randn(M, K, dtype=dtype, device="cuda") - # Create weights - # For simplicity in testing, let's create random int8 weights and scales w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda") - # Scale shape: The Int8Linear class uses: - # row_blocks = cdiv(out_features, b=128) - # col_blocks = cdiv(in_features, b=128) - # scale shape: (row_blocks, col_blocks) row_blocks = (N + 127) // 128 col_blocks = (K + 127) // 128 w_s = torch.randn(row_blocks, col_blocks, dtype=torch.float32, device="cuda").abs() - # Run int8_linear output = turbodiffusion_ops.int8_linear(x, w_q, w_s) assert output.shape == (M, N) @@ -92,7 +83,6 @@ class TestTurboDiffusion: M, N, K = 32, 64, 128 x = torch.randn(M, K, dtype=dtype, device="cuda", requires_grad=True) - # Weights (frozen) w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda") row_blocks = (N + 127) // 128 col_blocks = (K + 127) // 128 @@ -100,7 +90,6 @@ class TestTurboDiffusion: bias = torch.randn(N, dtype=dtype, device="cuda", requires_grad=True) - # Use Int8LinearFunction output = turbodiffusion_ops.Int8LinearFunction.apply(x, w_q, w_s, bias) loss = output.sum() loss.backward() @@ -109,7 +98,6 @@ class TestTurboDiffusion: assert bias.grad is not None assert x.grad.shape == (M, K) - # Check correctness against dequantized weight w_float = turbodiffusion_ops.dequantize_weight(w_q, w_s, dtype) # Reference diff --git a/fastvideo-kernel/tests/test_vmoba_correctness.py b/fastvideo-kernel/tests/test_vmoba_correctness.py index 2f31bafd3..2dea54490 100644 --- a/fastvideo-kernel/tests/test_vmoba_correctness.py +++ b/fastvideo-kernel/tests/test_vmoba_correctness.py @@ -14,7 +14,6 @@ def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, dev random.seed(42) torch.cuda.manual_seed_all(42) - # Generate sequence lengths for each item in the batch if batch_size > 1: # Ensure sequence lengths are reasonably distributed avg_seqlen = total_seqlen // batch_size @@ -26,7 +25,6 @@ def generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype, dev seqlens.append(avg_seqlen) current_sum = sum(seqlens) seqlens[-1] -= (current_sum - total_seqlen) - # Ensure all lengths are positive seqlens = [max(1, s) for s in seqlens] # Final adjustment to match total_seqlen seqlens[-1] += total_seqlen - sum(seqlens) @@ -64,7 +62,6 @@ def test_moba_attn_varlen_forward(batch_size, total_seqlen, num_heads, head_dim, q, k, v, cu_seqlens, max_seqlen = generate_test_data(batch_size, total_seqlen, num_heads, head_dim, dtype) - # Ensure chunk size is not larger than the smallest sequence length min_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).min().item() if moba_chunk_size > min_seqlen: pytest.skip("moba_chunk_size is larger than the minimum sequence length in the batch") diff --git a/fastvideo-kernel/tests/test_vsa.py b/fastvideo-kernel/tests/test_vsa.py index 85bb85e43..cdc83f2a1 100644 --- a/fastvideo-kernel/tests/test_vsa.py +++ b/fastvideo-kernel/tests/test_vsa.py @@ -6,8 +6,6 @@ import numpy as np from tqdm import tqdm from .utils import generate_block_sparse_mask_for_function, create_full_mask_from_block_mask -# Use installed package -# from fastvideo_kernel import video_sparse_attn as block_sparse_attn BLOCK_M = 64 BLOCK_N = 64 @@ -118,10 +116,7 @@ def check_correctness(h, d, num_blocks, k, num_iterations=20, error_mode='all'): V = generate_tensor((1, h, S, d), torch.bfloat16, device) dO = generate_tensor((1, h, S, d), torch.bfloat16, device) - # print(Q.shape, K.shape, V.shape, dO.shape) - # dO_padded = torch.zeros_like(dO_padded) - # dO_padded[:, :, non_pad_index, :] = dO pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO) bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), variable_block_sizes, @@ -190,7 +185,6 @@ def check_correctness_qkdiff(h, d, num_q_blocks, num_kv_blocks, k, num_iteration V = generate_tensor((1, h, S_kv, d), torch.bfloat16, device) dO = generate_tensor((1, h, S_q, d), torch.bfloat16, device) - # print(Q.shape, K.shape, V.shape, dO.shape) pt_o, pt_qg, pt_kg, pt_vg = pytorch_test(Q, K, V, full_mask, dO) bs_o, bs_qg, bs_kg, bs_vg = block_sparse_kernel_test(Q, K, V, block_mask.unsqueeze(0), kv_variable_block_sizes, diff --git a/fastvideo-kernel/tests/test_vsa_forward.py b/fastvideo-kernel/tests/test_vsa_forward.py index b7406d4d1..7921f3cd1 100644 --- a/fastvideo-kernel/tests/test_vsa_forward.py +++ b/fastvideo-kernel/tests/test_vsa_forward.py @@ -9,7 +9,7 @@ from .utils import ( generate_block_sparse_mask_for_function, create_full_mask_from_block_mask, ) -from .test_vsa import BLOCK_M # Import from local test_vsa +from .test_vsa import BLOCK_M from . import test_vsa as ref @@ -66,7 +66,6 @@ def block_sparse_forward_test( except RuntimeError as e: pytest.skip(str(e)) - # Remove padding on the query side out = out_padded[:, :, q_non_pad_index, :] return out diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index e0123bf37..c212a3f5c 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -569,7 +569,7 @@ class TextDataset(torch.utils.data.IterableDataset, torch.distributed.checkpoint with open(self.data_merge_path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() - if line: # Skip empty lines + if line: prompts.append(line) logger.info(f"Loaded {len(prompts)} text prompts from {self.data_merge_path}") diff --git a/fastvideo/distributed/device_communicators/pynccl_wrapper.py b/fastvideo/distributed/device_communicators/pynccl_wrapper.py index fdf4e6944..30d74a9cc 100644 --- a/fastvideo/distributed/device_communicators/pynccl_wrapper.py +++ b/fastvideo/distributed/device_communicators/pynccl_wrapper.py @@ -127,7 +127,6 @@ class Function: class NCCLLibrary: exported_functions = [ - # const char* ncclGetErrorString(ncclResult_t result) Function("ncclGetErrorString", ctypes.c_char_p, [ncclResult_t]), # ncclResult_t ncclGetVersion(int *version); Function("ncclGetVersion", ncclResult_t, [ctypes.POINTER(ctypes.c_int)]), @@ -166,9 +165,6 @@ class NCCLLibrary: Function("ncclReduceScatter", ncclResult_t, [buffer_type, buffer_type, ctypes.c_size_t, ncclDataType_t, ncclRedOp_t, ncclComm_t, cudaStream_t]), - # ncclResult_t ncclSend( - # const void* sendbuff, size_t count, ncclDataType_t datatype, - # int dest, ncclComm_t comm, cudaStream_t stream); Function("ncclSend", ncclResult_t, [buffer_type, ctypes.c_size_t, ncclDataType_t, ctypes.c_int, ncclComm_t, cudaStream_t]), @@ -193,12 +189,8 @@ class NCCLLibrary: Function("ncclCommDestroy", ncclResult_t, [ncclComm_t]), ] - # class attribute to store the mapping from the path to the library - # to avoid loading the same library multiple times path_to_library_cache: dict[str, Any] = {} - # class attribute to store the mapping from library path - # to the corresponding dictionary path_to_dict_mapping: dict[str, dict[str, Any]] = {} def __init__(self, so_file: str | None = None): diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 8be4cd670..6656285a7 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -321,7 +321,7 @@ class VideoGenerator: return cls( fastvideo_args=fastvideo_args, executor_class=executor_class, - log_stats=False, # TODO: implement + log_stats=False, log_queue=log_queue, ) diff --git a/fastvideo/layers/fp4linear.py b/fastvideo/layers/fp4linear.py index 0a1841a75..d5fd6712a 100644 --- a/fastvideo/layers/fp4linear.py +++ b/fastvideo/layers/fp4linear.py @@ -79,7 +79,6 @@ class _LinearFWD4BWD16Fn(torch.autograd.Function): if bias_cast is not None: out2d.add_(bias_cast) - # save tensors for backward (keep original dtypes) ctx.save_for_backward(x2d, weight, bias) ctx.k = k ctx.n = n diff --git a/fastvideo/mlx_runtime/refine.py b/fastvideo/mlx_runtime/refine.py index 81c627ed7..5b80babe3 100644 --- a/fastvideo/mlx_runtime/refine.py +++ b/fastvideo/mlx_runtime/refine.py @@ -291,7 +291,6 @@ def _upsample_latents_numpy( x1 = np.minimum(x0 + 1, w - 1) wy = (ys - y0.astype(np.float32))[:, None] wx = (xs - x0.astype(np.float32))[None, :] - # Gather the four corners: shape (N, out_h, out_w). Ia = src[:, y0[:, None], x0[None, :]] Ib = src[:, y0[:, None], x1[None, :]] Ic = src[:, y1[:, None], x0[None, :]] diff --git a/fastvideo/models/dits/cosmos2_5.py b/fastvideo/models/dits/cosmos2_5.py index c722eeed5..b787fa774 100644 --- a/fastvideo/models/dits/cosmos2_5.py +++ b/fastvideo/models/dits/cosmos2_5.py @@ -96,7 +96,7 @@ class Cosmos25TimestepEmbedding(nn.Module): if self.use_adaln_lora: adaln_lora = emb # (B, T, 3D) - emb_standard = sample # Use input as standard embedding + emb_standard = sample else: emb_standard = emb adaln_lora = None @@ -736,7 +736,7 @@ class Cosmos25Transformer3DModel(BaseDiT): patch_embed_in_channels = config.in_channels # Base VAE channels (16) patch_embed_in_channels += 1 # Always add 1 for condition_mask if config.concat_padding_mask: - patch_embed_in_channels += 1 # Add 1 for padding_mask + patch_embed_in_channels += 1 # Total: 16 + 1 + 1 = 18 (with concat_padding_mask=True) self.patch_embed = Cosmos25PatchEmbed(patch_embed_in_channels, inner_dim, config.patch_size) diff --git a/fastvideo/models/dits/hyworld/pose.py b/fastvideo/models/dits/hyworld/pose.py index 5edb7b3b8..325712175 100644 --- a/fastvideo/models/dits/hyworld/pose.py +++ b/fastvideo/models/dits/hyworld/pose.py @@ -235,7 +235,7 @@ def pose_to_input( if move_norms > move_norm_valid: # threshold for movement move_norm_dirs = move_dirs / move_norms angles_rad = np.arccos(move_norm_dirs.clip(-1.0, 1.0)) - trans_angles_deg = angles_rad * (180.0 / np.pi) # convert to degrees + trans_angles_deg = angles_rad * (180.0 / np.pi) else: trans_angles_deg = np.zeros(3) diff --git a/fastvideo/models/dits/hyworld/retrieval_context.py b/fastvideo/models/dits/hyworld/retrieval_context.py index 8947e4b5c..f126957c4 100644 --- a/fastvideo/models/dits/hyworld/retrieval_context.py +++ b/fastvideo/models/dits/hyworld/retrieval_context.py @@ -252,7 +252,7 @@ def select_aligned_memory_frames( historical_clip_indices = list(range(4, current_frame_idx - temporal_context_size, 4)) - memory_frames_indices = [0, 1, 2, 3] # add the first chunk as context + memory_frames_indices = [0, 1, 2, 3] memory_frames = memory_frames - temporal_context_size for hist_idx in historical_clip_indices: diff --git a/fastvideo/models/dits/hyworld/trajectory.py b/fastvideo/models/dits/hyworld/trajectory.py index 5d7699b0b..a023fcc1c 100644 --- a/fastvideo/models/dits/hyworld/trajectory.py +++ b/fastvideo/models/dits/hyworld/trajectory.py @@ -35,24 +35,20 @@ def generate_camera_trajectory_local(motions): poses.append(T.copy()) for move in motions: - # Rotate (Left or Right) if "yaw" in move: R = rot_y(move["yaw"]) T[:3, :3] = T[:3, :3] @ R - # Rotate (Up or Down) if "pitch" in move: R = rot_x(move["pitch"]) T[:3, :3] = T[:3, :3] @ R - # Translation (Z-direction of the camera's local coordinate system) forward = move.get("forward", 0.0) if forward != 0: local_t = np.array([0, 0, forward]) world_t = T[:3, :3] @ local_t T[:3, 3] += world_t - # Translation (Z-direction of the camera's local coordinate system) right = move.get("right", 0.0) if right != 0: local_t = np.array([right, 0, 0]) diff --git a/fastvideo/models/dits/lingbotworld/cam_utils.py b/fastvideo/models/dits/lingbotworld/cam_utils.py index 5384fbaff..5ac8a77e4 100644 --- a/fastvideo/models/dits/lingbotworld/cam_utils.py +++ b/fastvideo/models/dits/lingbotworld/cam_utils.py @@ -179,7 +179,7 @@ def prepare_camera_embedding( height_final=height, width_final=width, ) - Ks = Ks[0] # use first frame + Ks = Ks[0] len_c2ws = len(c2ws) num_latent_frames = (len_c2ws - 1) // 4 + 1 diff --git a/fastvideo/models/dits/lingbotworld2/cam_utils.py b/fastvideo/models/dits/lingbotworld2/cam_utils.py index b0e9c030f..71fb21338 100644 --- a/fastvideo/models/dits/lingbotworld2/cam_utils.py +++ b/fastvideo/models/dits/lingbotworld2/cam_utils.py @@ -182,7 +182,7 @@ def prepare_camera_embedding( height_final=height, width_final=width, ) - Ks = Ks[0] # use first frame + Ks = Ks[0] len_c2ws = len(c2ws) num_latent_frames = (len_c2ws - 1) // 4 + 1 diff --git a/fastvideo/models/dits/matrixgame2/action_module.py b/fastvideo/models/dits/matrixgame2/action_module.py index c74a2039f..5e3eadea6 100644 --- a/fastvideo/models/dits/matrixgame2/action_module.py +++ b/fastvideo/models/dits/matrixgame2/action_module.py @@ -418,7 +418,7 @@ class ActionModule(nn.Module): ] group_mouse = torch.stack(group_mouse, dim=1) - actual_num_frames = group_mouse.shape[1] # Use actual stacked frame count + actual_num_frames = group_mouse.shape[1] S = th * tw group_mouse = group_mouse.unsqueeze(-1).expand(B, actual_num_frames, pad_t, C, S) diff --git a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py index 2a559dfce..598a9345c 100644 --- a/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -529,7 +529,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler def _convert_to_karras(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: """Constructs the noise schedule of Karras et al. (2022).""" - # Hack to make sure that other schedulers which copy this function don't break sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None @@ -548,7 +547,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler def _convert_to_exponential(self, in_sigmas: torch.Tensor, num_inference_steps: int) -> torch.Tensor: """Constructs an exponential noise schedule.""" - # Hack to make sure that other schedulers which copy this function don't break sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None @@ -567,7 +565,6 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler beta: float = 0.6) -> torch.Tensor: """From "Beta Sampling is All You Need" [arXiv:2407.12173] (Lee et. al, 2024)""" - # Hack to make sure that other schedulers which copy this function don't break sigma_min = self.config.sigma_min if hasattr(self.config, "sigma_min") else None sigma_max = self.config.sigma_max if hasattr(self.config, "sigma_max") else None diff --git a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py index feeaa20b3..187dd4671 100644 --- a/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py @@ -243,7 +243,6 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if self.solver_p: self.solver_p.set_timesteps(self.num_inference_steps, device=device) - # add an index counter for schedulers that allow duplicated timesteps self._step_index = None self._begin_index = None self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication diff --git a/fastvideo/models/schedulers/scheduling_unipc_multistep.py b/fastvideo/models/schedulers/scheduling_unipc_multistep.py index dff0eb33a..df19f2631 100644 --- a/fastvideo/models/schedulers/scheduling_unipc_multistep.py +++ b/fastvideo/models/schedulers/scheduling_unipc_multistep.py @@ -425,7 +425,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): if self.solver_p: self.solver_p.set_timesteps(self.num_inference_steps, device=device) - # add an index counter for schedulers that allow duplicated timesteps self._step_index = None self._begin_index = None self.sigmas = self.sigmas.to("cpu") # to avoid too much CPU/GPU communication @@ -1008,7 +1007,6 @@ class UniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): noise: torch.Tensor, timesteps: torch.IntTensor, ) -> torch.Tensor: - # Make sure sigmas and timesteps have the same device and dtype as original_samples sigmas = self.sigmas.to(device=original_samples.device, dtype=original_samples.dtype) if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): # mps does not support float64 diff --git a/fastvideo/models/vaes/hyworldvae.py b/fastvideo/models/vaes/hyworldvae.py index c0626ad1f..9cb45779f 100644 --- a/fastvideo/models/vaes/hyworldvae.py +++ b/fastvideo/models/vaes/hyworldvae.py @@ -865,7 +865,6 @@ class AutoencoderKLHYWorld(nn.Module, ParallelTiledVAE): self.clear_cache() - # Process in chunks: first frame alone, then groups of 4 frames iter_ = 1 + (num_frame - 1) // 4 for i in range(iter_): self._enc_conv_idx = [0] diff --git a/fastvideo/pipelines/basic/ltx2/stages/ltx2_text_encoding.py b/fastvideo/pipelines/basic/ltx2/stages/ltx2_text_encoding.py index fb6caa640..52f718cc1 100644 --- a/fastvideo/pipelines/basic/ltx2/stages/ltx2_text_encoding.py +++ b/fastvideo/pipelines/basic/ltx2/stages/ltx2_text_encoding.py @@ -158,7 +158,6 @@ class LTX2TextEncodingStage(TextEncodingStage): # LTX2 audio negative embeddings has_audio_neg = int(d["_has_audio_neg"].item()) == 1 if has_audio_neg: - # Use same count as audio embeds num_audio = int(d["_num_audio_embeds"].item()) audio_neg = [d[f"audio_neg_embed_{i}"] for i in range(num_audio)] batch.extra["ltx2_audio_negative_embeds"] = audio_neg diff --git a/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline_ode_trajectory.py index 4bcd6935d..983902651 100644 --- a/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline_ode_trajectory.py +++ b/fastvideo/pipelines/preprocess/matrixgame2/matrixgame2_preprocess_pipeline_ode_trajectory.py @@ -155,7 +155,7 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): # Create mask_cond: ones for first frame, zeros for rest # Shape: (B, 16, latent_frames, latent_height, latent_width) mask_cond = torch.ones_like(latent_condition) - mask_cond[:, :, 1:] = 0 # Set all frames except first to 0 + mask_cond[:, :, 1:] = 0 # Create cond_concat: first 4 channels of mask + all 16 channels of img_cond # Shape: (B, 20, latent_frames, latent_height, latent_width) cond_concat = torch.cat([mask_cond[:, :4], latent_condition], dim=1) @@ -193,7 +193,7 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): # Filter out invalid samples (those with all zeros) valid_indices = [] for i, pixel_values in enumerate(data["pixel_values"]): - if not torch.all(pixel_values == 0): # Check if all values are zero + if not torch.all(pixel_values == 0): valid_indices.append(i) self.num_processed_samples += len(valid_indices) @@ -362,8 +362,8 @@ class PreprocessPipeline_MatrixGame2_ODE_Trajectory(BasePreprocessPipeline): unit="batch", disable=self.local_rank != 0) - self.video_data: dict[str, Any] = {} # Store video metadata and paths - self.latent_data: dict[str, Any] = {} # Store latent tensors + self.video_data: dict[str, Any] = {} + self.latent_data: dict[str, Any] = {} self.preprocess_action_and_trajectory(fastvideo_args, args) diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py index 9420f858b..8ff4cf563 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_base.py @@ -43,8 +43,8 @@ class BasePreprocessPipeline(ComposedPipelineBase): if not self.post_init_called: self.post_init() - self.video_data: dict[str, Any] = {} # Store video metadata and paths - self.latent_data: dict[str, Any] = {} # Store latent tensors + self.video_data: dict[str, Any] = {} + self.latent_data: dict[str, Any] = {} self.preprocess_video_and_text(fastvideo_args, args) def get_extra_features(self, valid_data: dict[str, Any], fastvideo_args: FastVideoArgs) -> dict[str, Any]: @@ -256,7 +256,7 @@ class BasePreprocessPipeline(ComposedPipelineBase): # Filter out invalid samples (those with all zeros) valid_indices = [] for i, pixel_values in enumerate(data["pixel_values"]): - if not torch.all(pixel_values == 0): # Check if all values are zero + if not torch.all(pixel_values == 0): valid_indices.append(i) num_processed_samples += len(valid_indices) diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py index 3c53aea97..5aabcb06c 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_ode_trajectory.py @@ -87,7 +87,7 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): # Filter out samples without text valid_indices = [] for i, text in enumerate(data["text"]): - if text and text.strip(): # Check if text is not empty + if text and text.strip(): valid_indices.append(i) self.num_processed_samples += len(valid_indices) @@ -265,8 +265,8 @@ class PreprocessPipeline_ODE_Trajectory(BasePreprocessPipeline): unit="batch", disable=self.local_rank != 0) - self.video_data: dict[str, Any] = {} # Store video metadata and paths - self.latent_data: dict[str, Any] = {} # Store latent tensors + self.video_data: dict[str, Any] = {} + self.latent_data: dict[str, Any] = {} self.preprocess_text_and_trajectory(fastvideo_args, args) diff --git a/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py b/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py index edcf70e10..09546b317 100644 --- a/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py +++ b/fastvideo/pipelines/preprocess/preprocess_pipeline_text.py @@ -61,7 +61,7 @@ class PreprocessPipeline_Text(BasePreprocessPipeline): # Filter out samples without text valid_indices = [] for i, text in enumerate(data["text"]): - if text and text.strip(): # Check if text is not empty + if text and text.strip(): valid_indices.append(i) self.num_processed_samples += len(valid_indices) @@ -160,7 +160,7 @@ class PreprocessPipeline_Text(BasePreprocessPipeline): unit="batch", disable=self.local_rank != 0) - self.text_data: dict[str, Any] = {} # Store text metadata and paths + self.text_data: dict[str, Any] = {} self.preprocess_text_only(fastvideo_args, args) diff --git a/fastvideo/pipelines/preprocess/v1_preprocess.py b/fastvideo/pipelines/preprocess/v1_preprocess.py index 2c8bc6f79..db5b132e0 100644 --- a/fastvideo/pipelines/preprocess/v1_preprocess.py +++ b/fastvideo/pipelines/preprocess/v1_preprocess.py @@ -105,8 +105,8 @@ if __name__ == "__main__": parser.add_argument("--max_height", type=int, default=480) parser.add_argument("--max_width", type=int, default=848) parser.add_argument("--video_length_tolerance_range", type=int, default=2.0) - parser.add_argument("--group_frame", action="store_true") # TODO - parser.add_argument("--group_resolution", action="store_true") # TODO + parser.add_argument("--group_frame", action="store_true") + parser.add_argument("--group_resolution", action="store_true") parser.add_argument("--flow_shift", type=float, default=None) parser.add_argument( "--preprocess_task", diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 04ca2ea09..4fd0e8455 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -70,11 +70,11 @@ class DenoisingStage(PipelineStage): attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads self.attn_backend = get_attn_backend( head_size=attn_head_size, - dtype=torch.float16, # TODO(will): hack + dtype=torch.float16, supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.BSA_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, AttentionBackendEnum.SAGE_ATTN_THREE, - AttentionBackendEnum.ATTN_QAT_INFER), # hack + AttentionBackendEnum.ATTN_QAT_INFER), # Build metadata for the backend this transformer actually resolved # instead of re-deriving it from the environment. The two agreed # only when the request arrived via the env var: a request passed as diff --git a/fastvideo/pipelines/stages/encoding.py b/fastvideo/pipelines/stages/encoding.py index 4901b7b54..0a79dcab5 100644 --- a/fastvideo/pipelines/stages/encoding.py +++ b/fastvideo/pipelines/stages/encoding.py @@ -11,7 +11,7 @@ from fastvideo.logger import init_logger from fastvideo.models.vaes.common import ParallelTiledVAE from fastvideo.pipelines.pipeline_batch_info import ForwardBatch from fastvideo.pipelines.stages.base import PipelineStage -from fastvideo.pipelines.stages.validators import V # Import validators +from fastvideo.pipelines.stages.validators import V from fastvideo.pipelines.stages.validators import VerificationResult from fastvideo.utils import PRECISION_TO_TYPE diff --git a/fastvideo/pipelines/stages/hyworld_denoising.py b/fastvideo/pipelines/stages/hyworld_denoising.py index 0a192dbbd..553e3f320 100644 --- a/fastvideo/pipelines/stages/hyworld_denoising.py +++ b/fastvideo/pipelines/stages/hyworld_denoising.py @@ -170,7 +170,7 @@ class HYWorldDenoisingStage(DenoisingStage): for chunk_start_idx in range( current_frame_idx, current_frame_idx + chunk_latent_frames, - 4, # Process every 4 frames + 4, ): selected_history_frame_id = select_aligned_memory_frames( viewmats[0].cpu().detach().numpy(), diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index b0061a6af..7c55bd5af 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -945,7 +945,7 @@ class MatrixGame2ImageVAEEncodingStage(ImageVAEEncodingStage): # Create mask_cond: ones for first frame, zeros for rest # Shape: (B, 16, latent_frames, latent_height, latent_width) mask_cond = torch.ones_like(img_cond) - mask_cond[:, :, 1:] = 0 # Set all frames except first to 0 + mask_cond[:, :, 1:] = 0 # Create cond_concat: first 4 channels of mask + all 16 channels of img_cond # Shape: (B, 20, latent_frames, latent_height, latent_width) diff --git a/fastvideo/pipelines/stages/matrixgame3_denoising.py b/fastvideo/pipelines/stages/matrixgame3_denoising.py index 6c5b1fc19..9962c1759 100644 --- a/fastvideo/pipelines/stages/matrixgame3_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame3_denoising.py @@ -99,7 +99,7 @@ class MatrixGame3DenoisingStage(DenoisingStage): target_h = latent_h * spatial_ratio target_w = latent_w * spatial_ratio num_iterations = self._infer_num_iterations(batch) - clip_frame = 56 # hardcode for now + clip_frame = 56 first_clip_frame = clip_frame + 1 past_frame = 16 diff --git a/fastvideo/pipelines/stages/sr_denoising.py b/fastvideo/pipelines/stages/sr_denoising.py index 0f32be765..0c05fac0c 100644 --- a/fastvideo/pipelines/stages/sr_denoising.py +++ b/fastvideo/pipelines/stages/sr_denoising.py @@ -61,10 +61,10 @@ class SRDenoisingStage(PipelineStage): attn_head_size = self.transformer.hidden_size // self.transformer.num_attention_heads self.attn_backend = get_attn_backend( head_size=attn_head_size, - dtype=torch.float16, # TODO(will): hack + dtype=torch.float16, supported_attention_backends=(AttentionBackendEnum.VIDEO_SPARSE_ATTN, AttentionBackendEnum.VMOBA_ATTN, AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA, - AttentionBackendEnum.SAGE_ATTN_THREE), # hack + AttentionBackendEnum.SAGE_ATTN_THREE), # See DenoisingStage: use this transformer's recorded decision # rather than the environment. requested=component_attention_backend(self.transformer), diff --git a/fastvideo/tests/encoders/test_hyt5_encoder.py b/fastvideo/tests/encoders/test_hyt5_encoder.py index e5510863d..d6c05cd90 100644 --- a/fastvideo/tests/encoders/test_hyt5_encoder.py +++ b/fastvideo/tests/encoders/test_hyt5_encoder.py @@ -112,7 +112,6 @@ def test_t5_encoder(t5_model_paths_and_config): logger.info("Maximum difference in last hidden states: %s", max_diff_hidden.item()) logger.info("Mean difference in last hidden states: %s", mean_diff_hidden.item()) logger.info("Max memory allocated: %s GB", torch.cuda.max_memory_allocated() / 1024**3) - # Check if outputs are similar (allowing for small numerical differences) assert mean_diff_hidden < 1e-4, \ f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" assert max_diff_hidden < 1e-4, \ diff --git a/fastvideo/tests/entrypoints/streaming/test_server.py b/fastvideo/tests/entrypoints/streaming/test_server.py index c98e32cd5..d982db821 100644 --- a/fastvideo/tests/entrypoints/streaming/test_server.py +++ b/fastvideo/tests/entrypoints/streaming/test_server.py @@ -218,7 +218,6 @@ class TestSegmentFlow: assert step["type"] == "step_complete" media_init = ws.receive_json() assert media_init["type"] == "media_init" - # Then one or more binary frames until media_segment_complete. saw_binary = False while True: msg = ws.receive() diff --git a/fastvideo/tests/hooks/test_layerwise_offload.py b/fastvideo/tests/hooks/test_layerwise_offload.py index 86a9b3af4..34859072f 100644 --- a/fastvideo/tests/hooks/test_layerwise_offload.py +++ b/fastvideo/tests/hooks/test_layerwise_offload.py @@ -98,7 +98,6 @@ def test_layerwise_offload_bf16(): with torch.no_grad(): offloaded_output = model(input_tensor.clone()) - # Check output correctness (looser tolerance for bf16) assert torch.allclose(reference_output, offloaded_output, rtol=1e-2, atol=1e-3), "Output with offloading should match reference output for bf16" diff --git a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py index 5dfaeb615..f3e893be6 100644 --- a/fastvideo/tests/inference/lora/test_lora_inference_similarity.py +++ b/fastvideo/tests/inference/lora/test_lora_inference_similarity.py @@ -191,7 +191,7 @@ def test_lora_inference_similarity(ATTENTION_BACKEND, model_id, env_overrides): # Try match with unsanitized prefix (for legacy reference videos) # Remove .mp4 and compare the base names after sanitization - base_filename = filename[:-4] # Remove .mp4 + base_filename = filename[:-4] if _sanitize_filename_component(base_filename) == output_video_name: reference_video_name = filename break diff --git a/fastvideo/tests/transformers/test_cosmos2_5.py b/fastvideo/tests/transformers/test_cosmos2_5.py index 9ff950b7b..bf0af6db9 100644 --- a/fastvideo/tests/transformers/test_cosmos2_5.py +++ b/fastvideo/tests/transformers/test_cosmos2_5.py @@ -327,7 +327,7 @@ def test_cosmos25_transformer(): forward_batch=forward_batch, ): # FastVideo expects padding_mask in [B, 1, H, W] format - padding_mask_fv = padding_mask.unsqueeze(1) # Add channel dimension for FastVideo + padding_mask_fv = padding_mask.unsqueeze(1) # FastVideo supports both [B] and [B, T] formats - use [B, T] to match official model # This ensures each frame gets its own timestep embedding (even if values are the same) timestep_fv = timestep # Already in [B, T] format diff --git a/fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py b/fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py index 99b41ec89..af3e4f340 100644 --- a/fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py +++ b/fastvideo/tests/vaes/test_minimax_h3_nvfp4_decoder.py @@ -497,7 +497,7 @@ def _load_h3_vae(vae_dir, cpu: bool = True, decoder: str | None = None, pipeline model_path=str(vae_dir), pipeline_config=pipeline_config or PipelineConfig(vae_config=MiniMaxH3VideoVAEConfig(), vae_precision="fp32"), h3_vae_nvfp4_decoder=decoder) - args.vae_cpu_offload = cpu # load on CPU + args.vae_cpu_offload = cpu return VAELoader().load(str(vae_dir), args) diff --git a/fastvideo/training/distillation_pipeline.py b/fastvideo/training/distillation_pipeline.py index 77b61868c..00ec47931 100644 --- a/fastvideo/training/distillation_pipeline.py +++ b/fastvideo/training/distillation_pipeline.py @@ -1117,7 +1117,7 @@ class DistillationPipeline(TrainingPipeline): if self.rank_in_sp_group == 0: if self.global_rank == 0: # Global rank 0 collects results from all sp_group leaders - all_videos = step_videos # Start with own results + all_videos = step_videos all_captions = step_captions all_audios = step_audios all_audio_sample_rates = step_audio_sample_rates diff --git a/fastvideo/training/training_pipeline.py b/fastvideo/training/training_pipeline.py index 9a139cce6..f3865810c 100644 --- a/fastvideo/training/training_pipeline.py +++ b/fastvideo/training/training_pipeline.py @@ -457,9 +457,7 @@ class TrainingPipeline(LoRAPipeline, ABC): # old sharding code, need to shard latents and noise but not input # Shard latents across sp groups training_batch.latents = training_batch.latents[:, :, :self.training_args.num_latent_t] - # shard noisy_model_input to match training_batch.noisy_model_input = training_batch.noisy_model_input[:, :, :self.training_args.num_latent_t] - # shard noise to match latents training_batch.noise = training_batch.noise[:, :, :self.training_args.num_latent_t] training_batch = self._build_attention_metadata(training_batch) @@ -765,7 +763,7 @@ class TrainingPipeline(LoRAPipeline, ABC): # results to global rank 0 if self.rank_in_sp_group == 0 and self.global_rank == 0: # Global rank 0 collects results from all sp_group leaders - all_videos = step_videos # Start with own results + all_videos = step_videos all_captions = step_captions all_ref_videos = step_ref_videos all_audios = step_audio diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index 169a204e9..bb75500b2 100644 --- a/fastvideo/training/training_utils.py +++ b/fastvideo/training/training_utils.py @@ -166,7 +166,7 @@ def save_checkpoint(transformer, config_dict = transformer.hf_config if "dtype" in config_dict: - del config_dict["dtype"] # TODO + del config_dict["dtype"] config_path = os.path.join(transformer_save_dir, "config.json") with open(config_path, "w") as f: json.dump(config_dict, f, indent=4) @@ -425,7 +425,7 @@ def save_distillation_checkpoint( config_dict = generator_transformer.hf_config if "dtype" in config_dict: - del config_dict["dtype"] # TODO + del config_dict["dtype"] config_path = os.path.join(inference_save_dir, "config.json") with open(config_path, "w") as f: json.dump(config_dict, f, indent=4) @@ -454,7 +454,7 @@ def save_distillation_checkpoint( config_dict_2 = generator_transformer_2.hf_config if "dtype" in config_dict_2: - del config_dict_2["dtype"] # TODO + del config_dict_2["dtype"] config_path_2 = os.path.join(inference_save_dir_2, "config.json") with open(config_path_2, "w") as f: json.dump(config_dict_2, f, indent=4) diff --git a/fastvideo/worker/ray_distributed_executor.py b/fastvideo/worker/ray_distributed_executor.py index 60e3f068f..e53b47268 100644 --- a/fastvideo/worker/ray_distributed_executor.py +++ b/fastvideo/worker/ray_distributed_executor.py @@ -170,7 +170,6 @@ class RayDistributedExecutor(Executor): # The remaining workers are the actual ray actors. self.workers: list[RayWorkerWrapper] = [] - # use the first N bundles that have GPU resources. bundle_indices: list[int] = [] for bundle_id, bundle in enumerate(placement_group.bundle_specs): if bundle.get(current_platform.ray_device_key, 0): diff --git a/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py b/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py index c2b786e77..db0da72de 100644 --- a/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py +++ b/scripts/checkpoint_conversion/convert_gen3c_to_fastvideo.py @@ -142,7 +142,7 @@ def load_checkpoint(path: Path, mmap: bool = True) -> dict[str, torch.Tensor]: if mmap: # Check if PyTorch version supports mmap (2.1+) try: - version_str = torch.__version__.split("+")[0] # Remove +cu118 suffix etc. + version_str = torch.__version__.split("+")[0] major, minor = version_str.split(".")[:2] torch_version = (int(major), int(minor)) if torch_version >= (2, 1): @@ -280,7 +280,6 @@ def convert_weights( if verbose: print(f" {key} -> {new_key}") - # Free memory by deleting processed tensor from source if memory_efficient: del state_dict[key] if i % 100 == 0: diff --git a/scripts/checkpoint_conversion/create_hf_repo.py b/scripts/checkpoint_conversion/create_hf_repo.py index 25e4eda85..f83e303ac 100644 --- a/scripts/checkpoint_conversion/create_hf_repo.py +++ b/scripts/checkpoint_conversion/create_hf_repo.py @@ -81,7 +81,7 @@ def main(): # Remove existing weights in target component dir to avoid mixing shards/files existing_weights = glob.glob(os.path.join(target_component_dir, "*.safetensors")) + \ glob.glob(os.path.join(target_component_dir, "*.bin")) + \ - glob.glob(os.path.join(target_component_dir, "*.index.json")) # Remove index if present + glob.glob(os.path.join(target_component_dir, "*.index.json")) for ew in existing_weights: print(f"Removing existing weight file: {ew}") diff --git a/scripts/checkpoint_conversion/wan_to_diffusers.py b/scripts/checkpoint_conversion/wan_to_diffusers.py index 08866cc73..72aa28fc0 100644 --- a/scripts/checkpoint_conversion/wan_to_diffusers.py +++ b/scripts/checkpoint_conversion/wan_to_diffusers.py @@ -81,7 +81,7 @@ for k, v in state_dict.items(): for pattern, replacement in _self_forcing_to_diffusers_param_names_mapping.items(): if re.match(pattern, k): new_key = re.sub(pattern, replacement, k) - break # Stop at the first match + break else: raise ValueError(f"No match found for {k}") new_state_dict[new_key] = v diff --git a/scripts/finetune/finetune_v1.sh b/scripts/finetune/finetune_v1.sh index dad52c5b9..4f2cb64d8 100644 --- a/scripts/finetune/finetune_v1.sh +++ b/scripts/finetune/finetune_v1.sh @@ -1,13 +1,10 @@ export WANDB_BASE_URL="https://api.wandb.ai" export WANDB_MODE=online export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA DATA_DIR=[your data dir] VALIDATION_DATASET_FILE=[your validation dataset file] NUM_GPUS=4 -# export CUDA_VISIBLE_DEVICES=4,5 -# IP=[MASTER NODE IP] # Make sure that num_latent_t is a multiple of sp_size torchrun --nnodes 1 --nproc_per_node $NUM_GPUS \ diff --git a/scripts/preprocess/launch_slurm_text.sh b/scripts/preprocess/launch_slurm_text.sh index f6a33ecec..027ebbed7 100644 --- a/scripts/preprocess/launch_slurm_text.sh +++ b/scripts/preprocess/launch_slurm_text.sh @@ -1,12 +1,10 @@ #!/bin/bash -# Create output directory if it doesn't exist mkdir -p preprocess_output_text # Launch 8 jobs, one for each node # Each node processes 8 consecutive files (64 total files / 8 nodes = 8 files per node) for node_id in {0..7}; do - # Calculate the starting file number for this node start_file=$((node_id * 8 + 1)) echo "Launching text-only node $node_id with files v2m_${start_file}.txt to v2m_$((start_file + 7)).txt" diff --git a/scripts/preprocess/syn_text.slurm b/scripts/preprocess/syn_text.slurm index 009e330e4..6a7c3cd62 100644 --- a/scripts/preprocess/syn_text.slurm +++ b/scripts/preprocess/syn_text.slurm @@ -29,10 +29,8 @@ NODE_ID=${2:-0} # Node identifier (0-7) num_gpus=1 export MODEL_BASE=Wan-AI/Wan2.1-T2V-1.3B-Diffusers -# Start port number - we'll increment for each job base_port=$((29603 + NODE_ID * 100)) # Different port range per node -# Create an array of CUDA device IDs gpu_ids=(0 1 2 3 4 5 6 7) GPU_NUM=1 @@ -43,12 +41,9 @@ echo "START_FILE: $START_FILE" echo "Base port for this node: $base_port" echo "Processing TEXT-ONLY data" -# Run 8 parallel preprocessing jobs on this node for i in {1..8}; do - # Calculate port for this job port=$((base_port + i)) - # Get GPU ID using modulo to cycle through available GPUs gpu=${gpu_ids[((i-1))]} # Calculate which file this GPU should process diff --git a/scripts/preprocess/v1_preprocess_wan_data_i2v.sh b/scripts/preprocess/v1_preprocess_wan_data_i2v.sh index 600e1fd1b..7aba7d006 100644 --- a/scripts/preprocess/v1_preprocess_wan_data_i2v.sh +++ b/scripts/preprocess/v1_preprocess_wan_data_i2v.sh @@ -1,4 +1,3 @@ -# export WANDB_MODE="offline" GPU_NUM=1 # 2,4,8 MODEL_PATH="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers" MODEL_TYPE="wan" diff --git a/scripts/preprocess/v1_preprocess_wan_data_t2v.sh b/scripts/preprocess/v1_preprocess_wan_data_t2v.sh index b335d5990..e32ea99fc 100644 --- a/scripts/preprocess/v1_preprocess_wan_data_t2v.sh +++ b/scripts/preprocess/v1_preprocess_wan_data_t2v.sh @@ -1,4 +1,3 @@ -# export WANDB_MODE="offline" GPU_NUM=1 # 2,4,8 MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" MODEL_TYPE="wan" diff --git a/tests/local_tests/gen3c/test_gen3c.py b/tests/local_tests/gen3c/test_gen3c.py index 678752fbe..4b46c07be 100644 --- a/tests/local_tests/gen3c/test_gen3c.py +++ b/tests/local_tests/gen3c/test_gen3c.py @@ -45,7 +45,6 @@ def create_test_config( config.arch_config.max_size = max_size config.arch_config.frame_buffer_max = frame_buffer_max - # Ensure buffer channels are recalculated config.arch_config.__post_init__() return config @@ -57,7 +56,6 @@ def test_gen3c_model_instantiation(): model = Gen3CTransformer3DModel(config=config, hf_config={}) - # Verify model attributes assert model.hidden_size == config.arch_config.num_attention_heads * config.arch_config.attention_head_dim assert model.in_channels == config.arch_config.in_channels assert model.out_channels == config.arch_config.out_channels @@ -79,7 +77,6 @@ def test_gen3c_forward_pass_shapes(): if not torch.cuda.is_available(): pytest.skip("CUDA required for GEN3C forward pass test") - # Check if distributed is initialized try: from fastvideo.distributed.parallel_state import get_sp_group get_sp_group() @@ -93,7 +90,6 @@ def test_gen3c_forward_pass_shapes(): model = Gen3CTransformer3DModel(config=config, hf_config={}).to(device=device, dtype=dtype) model.eval() - # Create test inputs batch_size = 1 num_frames = 8 # Divisible by patch_t=1 height = 16 # Divisible by patch_h=2 @@ -160,7 +156,6 @@ def test_gen3c_forward_pass_shapes(): padding_mask=padding_mask, ) - # Verify output shape matches input latent shape assert output.shape == hidden_states.shape, ( f"Output shape {output.shape} does not match input shape {hidden_states.shape}") @@ -179,7 +174,6 @@ def test_gen3c_forward_pass_no_conditioning(): if not torch.cuda.is_available(): pytest.skip("CUDA required for GEN3C forward pass test") - # Check if distributed is initialized try: from fastvideo.distributed.parallel_state import get_sp_group get_sp_group() @@ -307,7 +301,6 @@ def test_gen3c_weight_loading(): total_params = sum(p.numel() for p in model.parameters()) print(f"[GEN3C TEST] Total parameters: {total_params:,}") - # Verify model is usable model.eval() assert model is not None @@ -332,7 +325,6 @@ def test_gen3c_transformer_parity(): device = torch.device("cuda:0") dtype = torch.bfloat16 - # Load FastVideo model config = Gen3CVideoConfig() args = FastVideoArgs( model_path=str(fastvideo_path), @@ -346,7 +338,6 @@ def test_gen3c_transformer_parity(): fastvideo_model = loader.load(str(fastvideo_path), args).to(device=device, dtype=dtype) fastvideo_model.eval() - # Load official GEN3C model try: from cosmos_predict1.diffusion.networks.general_dit_video_conditioned import VideoExtendGeneralDIT @@ -383,7 +374,6 @@ def test_gen3c_transformer_parity(): except Exception as e: pytest.skip(f"Failed to load official GEN3C: {e}") - # Create test inputs batch_size = 1 num_frames = 8 height = 16 @@ -443,7 +433,6 @@ def test_gen3c_transformer_parity(): condition_video_pose=condition_video_pose, ) - # Prepare inputs for official model x_official = torch.cat([ hidden_states, condition_video_input_mask, diff --git a/tests/local_tests/helpers/magi_human_upstream.py b/tests/local_tests/helpers/magi_human_upstream.py index 724e734e1..f47ad1e1c 100644 --- a/tests/local_tests/helpers/magi_human_upstream.py +++ b/tests/local_tests/helpers/magi_human_upstream.py @@ -316,7 +316,6 @@ def load_upstream_dit(base_shard_dir, device=None, dtype=None, local_attn_layers model = DiTModel(model_config=model_config) - # Load all base shards into a single state dict. base_shard_dir = Path(base_shard_dir) shard_paths = sorted(base_shard_dir.glob("*.safetensors")) state = {} diff --git a/tests/local_tests/ltx2_5/test_ltx2_5_conversion.py b/tests/local_tests/ltx2_5/test_ltx2_5_conversion.py index 3d0ba8a76..15abc5b30 100644 --- a/tests/local_tests/ltx2_5/test_ltx2_5_conversion.py +++ b/tests/local_tests/ltx2_5/test_ltx2_5_conversion.py @@ -276,7 +276,6 @@ def test_split_model_index_enables_distilled_refine_without_lora() -> None: def test_split_ltx2_5_conversion_preserves_enabled_transformer_gates(tmp_path: Path) -> None: """Transformer config filtering retains True-valued 2.5 architecture gates.""" - # Create a fresh fixture with all 2.5 gates enabled transformer_config = { "transformer": { "num_attention_heads": 2, diff --git a/tests/local_tests/ltx2_5/test_ltx2_5_registry.py b/tests/local_tests/ltx2_5/test_ltx2_5_registry.py index 9fa437175..7d1957bc2 100644 --- a/tests/local_tests/ltx2_5/test_ltx2_5_registry.py +++ b/tests/local_tests/ltx2_5/test_ltx2_5_registry.py @@ -44,7 +44,6 @@ def test_ltx2_5_metadata_variant_selects_preset(tmp_path) -> None: import json from pathlib import Path - # Create a local directory with LTX2Pipeline _class_name and distilled variant metadata distilled_dir = tmp_path / "ltx2_5_distilled_local" distilled_dir.mkdir() model_index = { @@ -57,7 +56,6 @@ def test_ltx2_5_metadata_variant_selects_preset(tmp_path) -> None: assert preset == "ltx2_5_distilled_two_stage" assert family == "ltx2" - # Create a local directory with LTX2Pipeline _class_name and dev variant metadata dev_dir = tmp_path / "ltx2_5_dev_local" dev_dir.mkdir() model_index = { diff --git a/tests/local_tests/ltx2_5/test_ltx2_5_sampler.py b/tests/local_tests/ltx2_5/test_ltx2_5_sampler.py index 62d453d96..4aee42c83 100644 --- a/tests/local_tests/ltx2_5/test_ltx2_5_sampler.py +++ b/tests/local_tests/ltx2_5/test_ltx2_5_sampler.py @@ -86,7 +86,6 @@ def test_ancestral_step_captures_noise_from_stage_forward(monkeypatch) -> None: mock_step, ) - # Create a minimal stage with a mock transformer mock_transformer = MagicMock() mock_transformer.return_value = (torch.zeros(1, 12, 4), torch.zeros(1, 18, 4)) stage = LTX2DenoisingStage( @@ -95,7 +94,6 @@ def test_ancestral_step_captures_noise_from_stage_forward(monkeypatch) -> None: num_inference_steps_override=1, ) - # Create a batch with video and audio latents matching expected shapes batch = ForwardBatch( latents=torch.randn(1, 4, 3, 2, 2), # video: (B, C, T, H, W) extra={ @@ -113,10 +111,8 @@ def test_ancestral_step_captures_noise_from_stage_forward(monkeypatch) -> None: ltx2_use_ancestral_sampler=True, ) - # Run the stage forward to generate ancestral noise stage.forward(batch, args) - # Verify we captured noise tensors (one per ancestral step) assert len(captured_noises) == 1 noise = captured_noises[0] @@ -125,7 +121,6 @@ def test_ancestral_step_captures_noise_from_stage_forward(monkeypatch) -> None: # Combined flattened noise should have 48 + 24 = 72 elements assert noise.numel() == 48 + 24 - # Verify reproducibility with the seeded offset expected_seed = 17 + ANCESTRAL_NOISE_SEED_OFFSET generator = torch.Generator(device="cpu").manual_seed(expected_seed) video_noise_expected = torch.randn((1, 4, 3, 2, 2), generator=generator) diff --git a/tests/test_bsa.py b/tests/test_bsa.py index b74956eac..89055fa01 100644 --- a/tests/test_bsa.py +++ b/tests/test_bsa.py @@ -393,7 +393,6 @@ class TestReconstructPruned: full_output = _reconstruct_pruned(sparse_output, keep_indices, block_size) - # Check that kept positions match exactly for n in range(N): for i, pos in enumerate(keep_indices[0, 0, n].tolist()): assert torch.allclose( @@ -426,7 +425,6 @@ class TestReconstructPruned: assert torch.allclose(full_output, sparse_output, atol=1e-6) -# --------------------------------------------------------------------------- # End-to-end BSAAttentionImpl.forward tests