From ad099ca18b607e13479e24d52f5af64b1791a29b Mon Sep 17 00:00:00 2001 From: RandNMR73 Date: Wed, 8 Apr 2026 09:10:11 +0000 Subject: [PATCH] clean --- .codex | 0 .gitignore | 2 + download.sh | 1 - .../Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.sh | 124 --- .../distill_dmd_t2v_1.3B_linear_fp4.sh | 120 --- .../distill_dmd_t2v_1.3B_qat.sh | 122 --- .../distill_dmd_t2v_1.3B_sage3.sh | 137 --- .../distill_dmd_t2v_1.3B_vidprom.sh | 122 --- .../distill_dmd_t2v_1.3B_vidprom_e2e_fp4.sh | 124 --- ...distill_dmd_t2v_1.3B_vidprom_linear_fp4.sh | 123 --- .../distill_dmd_t2v_1.3B_vidprom_qat.sh | 123 --- .../crush_smol/finetune_t2v_fp4_e2e.sh | 98 -- .../crush_smol/finetune_t2v_fp4_linear.sh | 97 -- .../crush_smol/finetune_t2v_nvfp4_5090.sh | 94 -- .../crush_smol/finetune_t2v_qat_16_M2.sh | 124 --- .../crush_smol/finetune_t2v_qat_32.sh | 116 --- .../crush_smol/finetune_t2v_qat_32_M2.sh | 124 --- .../crush_smol/finetune_t2v_qat_5090.sh | 97 -- .../crush_smol/finetune_t2v_qat_8.sh | 119 --- .../crush_smol/finetune_t2v_qat_8_M2.sh | 123 --- .../crush_smol/finetune_t2v_qat_8_NVL72.sh | 121 --- ...t2v_qat_16.sh => finetune_t2v_qat_attn.sh} | 0 .../crush_smol/finetune_t2v_sage3_5090.sh | 94 -- .../finetune/wan_t2v_14B/finetune_t2v_B200.sh | 115 --- .../finetune_t2v_B200_qat_debug.sh | 125 --- .../finetune/wan_t2v_14B/finetune_t2v_qat.sh | 93 -- .../wan_t2v_14B/finetune_t2v_qat_B200.sh | 116 --- ..._32_14B_M2.sh => finetune_t2v_qat_attn.sh} | 0 .../wan_t2v_14B/finetune_t2v_sage3.sh | 113 --- .../modified_sageattn}/__init__.py | 0 .../modified_sageattn}/api.py | 0 .../modified_sageattn}/blackwell/__init__.py | 0 .../modified_sageattn}/blackwell/api.cu | 0 .../blackwell/block_config.h | 0 .../modified_sageattn}/blackwell/block_info.h | 0 .../blackwell/blockscaled_layout.h | 0 .../blackwell/cute_extension.h | 0 .../blackwell/epilogue_tma_ws.h | 0 .../blackwell/kernel_traits.h | 0 .../modified_sageattn}/blackwell/kernel_ws.h | 0 .../modified_sageattn}/blackwell/launch.h | 0 .../blackwell/mainloop_tma_ws.h | 0 .../blackwell/named_barrier.h | 0 .../modified_sageattn}/blackwell/params.h | 0 .../blackwell/softmax_fused.h | 0 .../blackwell/static_switch.h | 0 .../blackwell/tile_scheduler.h | 0 .../modified_sageattn}/blackwell/utils.h | 0 .../quantization/__init__.py | 0 .../quantization/bench/bench_quant_k.py | 0 .../quantization/bench/bench_quant_q.py | 0 .../quantization/bench/bench_quant_v.py | 0 .../quantization/bench/bench_utils.py | 0 .../quantization/cuda_utils.h | 0 .../quantization/fp4_quantization_4d.cu | 0 .../triton_kernels/nvfp4_utils.py | 0 .../triton_kernels/qat_attn.py | 6 +- .../triton_kernels/quant_utils.py | 0 .../tests/test_fake_quant.py | 5 +- .../tests/test_qat_attn.py | 6 +- .../benchmark_blockscaled_fp4_attn.py | 5 +- .../benchmark_scripts/benchmark_combined.py | 5 +- .../benchmark_scripts/benchmark_flashattn2.py | 3 +- .../benchmark_scripts/benchmark_sageattn3.py | 5 +- ...{setup.py => modified_sageattn_3_setup.py} | 13 +- fastvideo/layers/quantization/fp8_config.py | 109 --- fastvideo/models/dits/hunyuanvideo15.py | 766 +++++++++++++++ fastvideo/models/schedulers/scheduling_rcm.py | 326 +++++++ fastvideo/pipelines/stages/conditioning.py | 6 +- .../pipelines/stages/matrixgame_denoising.py | 739 +++++++++++++++ fastvideo/platforms/cuda.py | 2 +- fused_attention.py | 763 --------------- test_fake_quant_attn.py | 416 --------- test_fake_quant_basic.py | 883 ------------------ test_fake_quantize_kernels.py | 452 --------- test_get_fake_quant.py | 336 ------- test_quantization_precision.py | 461 --------- test_scale_debug.py | 92 -- test_triton_fp4_vs_cutlass.py | 276 ------ triton_fp4.py | 320 ------- triton_original.py | 302 ------ 81 files changed, 1864 insertions(+), 7200 deletions(-) create mode 100644 .codex delete mode 100644 download.sh delete mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.sh delete mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_linear_fp4.sh delete mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_qat.sh delete mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh delete mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom.sh delete mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_e2e_fp4.sh delete mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_linear_fp4.sh delete mode 100644 examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_qat.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_e2e.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_linear.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4_5090.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_16_M2.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_32.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_32_M2.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_5090.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8_M2.sh delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8_NVL72.sh rename examples/training/finetune/wan_t2v_1.3B/crush_smol/{finetune_t2v_qat_16.sh => finetune_t2v_qat_attn.sh} (100%) delete mode 100644 examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh delete mode 100644 examples/training/finetune/wan_t2v_14B/finetune_t2v_B200.sh delete mode 100644 examples/training/finetune/wan_t2v_14B/finetune_t2v_B200_qat_debug.sh delete mode 100644 examples/training/finetune/wan_t2v_14B/finetune_t2v_qat.sh delete mode 100644 examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_B200.sh rename examples/training/finetune/wan_t2v_14B/{finetune_t2v_qat_32_14B_M2.sh => finetune_t2v_qat_attn.sh} (100%) delete mode 100644 examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/__init__.py (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/api.py (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/__init__.py (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/api.cu (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/block_config.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/block_info.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/blockscaled_layout.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/cute_extension.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/epilogue_tma_ws.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/kernel_traits.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/kernel_ws.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/launch.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/mainloop_tma_ws.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/named_barrier.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/params.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/softmax_fused.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/static_switch.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/tile_scheduler.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/blackwell/utils.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/quantization/__init__.py (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/quantization/bench/bench_quant_k.py (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/quantization/bench/bench_quant_q.py (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/quantization/bench/bench_quant_v.py (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/quantization/bench/bench_utils.py (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/quantization/cuda_utils.h (100%) rename {fastvideo/attention/backends/sageattn => fastvideo-kernel/modified_sageattn}/quantization/fp4_quantization_4d.cu (100%) rename nvfp4_utils.py => fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py (100%) rename qat_attn.py => fastvideo-kernel/python/fastvideo_kernel/triton_kernels/qat_attn.py (99%) rename quant_utils.py => fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py (100%) rename test_fake_quant.py => fastvideo-kernel/tests/test_fake_quant.py (99%) rename test_qat_attn.py => fastvideo-kernel/tests/test_qat_attn.py (99%) rename fastvideo/attention/backends/{setup.py => modified_sageattn_3_setup.py} (94%) delete mode 100644 fastvideo/layers/quantization/fp8_config.py delete mode 100644 fused_attention.py delete mode 100644 test_fake_quant_attn.py delete mode 100644 test_fake_quant_basic.py delete mode 100644 test_fake_quantize_kernels.py delete mode 100644 test_get_fake_quant.py delete mode 100644 test_quantization_precision.py delete mode 100644 test_scale_debug.py delete mode 100644 test_triton_fp4_vs_cutlass.py delete mode 100644 triton_fp4.py delete mode 100644 triton_original.py diff --git a/.codex b/.codex new file mode 100644 index 000000000..e69de29bb diff --git a/.gitignore b/.gitignore index 222c3c3a8..bdd4d92ee 100644 --- a/.gitignore +++ b/.gitignore @@ -92,3 +92,5 @@ preprocess_output_text/ .sisyphus/ openspec/ fastvideo/tests/ssim/reference_videos/** +data2/ +FastVideo/ \ No newline at end of file diff --git a/download.sh b/download.sh deleted file mode 100644 index cebb6c13a..000000000 --- a/download.sh +++ /dev/null @@ -1 +0,0 @@ -python scripts/huggingface/download_hf.py --repo_id "FastVideo/Wan-Syn_77x448x832_600k" --local_dir "data2/Wan-Syn_77x448x832_600k" --repo_type "dataset" \ No newline at end of file diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.sh deleted file mode 100644 index b3ccd78d6..000000000 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B.sh +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################ -# Single node, 6 GPUs -############################################ -NUM_GPUS=1 -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -# export CUDA_VISIBLE_DEVICES=0,1,2,3 - -export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_MODE="online" -export WANDB_BASE_URL="https://api.wandb.ai" -export TOKENIZERS_PARALLELISM=false -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -export TRITON_CACHE_DIR=/tmp/triton_cache - -echo "NUM_GPUS: $NUM_GPUS" -echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" - -############################################ -# Configs -############################################ -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -OUTPUT_DIR="checkpoints/base_distill" - -# Training arguments -training_args=( - --tracker_project_name wan_t2v_distill_dmd_qat - --output_dir "$OUTPUT_DIR" - --max_train_steps 4000 - --train_batch_size 64 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 12 - --num_height 480 - --num_width 832 - --num_frames 45 - --enable_gradient_checkpointing_type "full" -) - -# Parallel arguments (adjusted to 6 GPUs) -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# Model arguments -model_args=( - --model_path "$MODEL_PATH" - --pretrained_model_name_or_path "$MODEL_PATH" - --real_score_model_path "$REAL_SCORE_MODEL_PATH" - --fake_score_model_path "$FAKE_SCORE_MODEL_PATH" -) - -# Dataset arguments -dataset_args=( - --data_path "$DATA_DIR" - --dataloader_num_workers 1 -) - -# Validation arguments -validation_args=( - --log_validation - --validation_dataset_file "$VALIDATION_DATASET_FILE" - --validation_steps 100 - --validation_sampling_steps "3" - --validation_guidance_scale "6.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 2e-6 - --mixed_precision "bf16" - --training_state_checkpointing_steps 2000 - --weight_only_checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.0 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 8 - --seed 1000 -) - -# DMD arguments -dmd_args=( - --dmd_denoising_steps '1000,757,522' - --min_timestep_ratio 0.02 - --max_timestep_ratio 0.98 - --generator_update_interval 5 - --real_score_guidance_scale 3.5 -) - -############################################ -# Launch (no rendezvous flags needed) -############################################ -set -x -torchrun \ - --nproc_per_node="$NUM_GPUS" \ - fastvideo/training/wan_distillation_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" \ - "${dmd_args[@]}" diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_linear_fp4.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_linear_fp4.sh deleted file mode 100644 index 4ba2c5571..000000000 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_linear_fp4.sh +++ /dev/null @@ -1,120 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################ -# Single node, 6 GPUs -############################################ -NUM_GPUS=6 -export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5 - - -export WANDB_MODE="online" -export WANDB_BASE_URL="https://api.wandb.ai" -export TOKENIZERS_PARALLELISM=false -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -export TRITON_CACHE_DIR=/tmp/triton_cache - -echo "NUM_GPUS: $NUM_GPUS" -echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" - -############################################ -# Configs -############################################ -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k/train -VALIDATION_DATASET_FILE=data/Wan-Syn_77x448x832_600k/validation_6.json -OUTPUT_DIR="checkpoints/linear_fp4_distill" - -# Training arguments -training_args=( - --tracker_project_name wan_t2v_distill_dmd - --output_dir "$OUTPUT_DIR" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 21 - --num_height 480 - --num_width 832 - --num_frames 81 - --enable_gradient_checkpointing_type "full" - --generator_4bit_linear True -) - -# Parallel arguments (adjusted to 6 GPUs) -parallel_args=( - --num_gpus 6 - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim 6 -) - -# 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" - --validation_steps 100 - --validation_sampling_steps "3" - --validation_guidance_scale "6.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 2e-6 - --mixed_precision "bf16" - --training_state_checkpointing_steps 2000 - --weight_only_checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.0 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 8 - --seed 1000 -) - -# DMD arguments -dmd_args=( - --dmd_denoising_steps '1000,757,522' - --min_timestep_ratio 0.02 - --max_timestep_ratio 0.98 - --generator_update_interval 5 - --real_score_guidance_scale 3.5 -) - -############################################ -# Launch (no rendezvous flags needed) -############################################ -set -x -torchrun \ - --nproc_per_node="$NUM_GPUS" \ - fastvideo/training/wan_distillation_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" \ - "${dmd_args[@]}" diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_qat.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_qat.sh deleted file mode 100644 index ce3f13d5a..000000000 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_qat.sh +++ /dev/null @@ -1,122 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################ -# Single node, 6 GPUs -############################################ -NUM_GPUS=1 -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -export WANDB_MODE="online" -export WANDB_BASE_URL="https://api.wandb.ai" -export TOKENIZERS_PARALLELISM=false -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -export TRITON_CACHE_DIR=/tmp/triton_cache - -echo "NUM_GPUS: $NUM_GPUS" -echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" - -############################################ -# Configs -############################################ -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -OUTPUT_DIR="checkpoints/wan_1.3B_t2v_distill_dmd_qat" -REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - -# Training arguments -training_args=( - --tracker_project_name wan_t2v_distill_dmd_qat - --output_dir $OUTPUT_DIR - --max_train_steps 4000 - --train_batch_size 64 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 21 - --num_height 480 - --num_width 832 - --num_frames 81 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments (adjusted to 6 GPUs) -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# Model arguments -model_args=( - --model_path $MODEL_PATH - --pretrained_model_name_or_path $MODEL_PATH - --real_score_model_path $REAL_SCORE_MODEL_PATH - --fake_score_model_path $FAKE_SCORE_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" - --validation_steps 200 - --validation_sampling_steps "3" - --validation_guidance_scale "6.0" # not used for dmd inference -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 2e-6 - --mixed_precision "bf16" - --training_state_checkpointing_steps 500 - --weight_only_checkpointing_steps 500 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.0 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 8 - --seed 1000 -) - -# DMD arguments -dmd_args=( - --dmd_denoising_steps '1000,757,522' - --min_timestep_ratio 0.02 - --max_timestep_ratio 0.98 - --generator_update_interval 5 - --real_score_guidance_scale 3.5 -) - -############################################ -# Launch (no rendezvous flags needed) -############################################ -set -x -torchrun \ - --nproc_per_node="$NUM_GPUS" \ - fastvideo/training/wan_distillation_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" \ - "${dmd_args[@]}" diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh deleted file mode 100644 index 528c1219c..000000000 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_sage3.sh +++ /dev/null @@ -1,137 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################ -# Single node, 6 GPUs -############################################ -<<<<<<< HEAD -NUM_GPUS=4 -======= -NUM_GPUS=1 ->>>>>>> 98ce47c7 (checkpoint (qat attn in progress)) -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -# export CUDA_VISIBLE_DEVICES=0,1,2,3 - -export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_MODE="online" -export WANDB_BASE_URL="https://api.wandb.ai" -export TOKENIZERS_PARALLELISM=false -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -export TRITON_CACHE_DIR=/tmp/triton_cache - -echo "NUM_GPUS: $NUM_GPUS" -echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" - -############################################ -# Configs -############################################ -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -<<<<<<< HEAD -DATA_DIR=data/Wan-Syn_77x448x832_600k -======= -DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/" ->>>>>>> 98ce47c7 (checkpoint (qat attn in progress)) -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -OUTPUT_DIR="checkpoints/sage3_distill" -REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - -# Training arguments -training_args=( - --tracker_project_name wan_t2v_distill_dmd_qat - --output_dir "$OUTPUT_DIR" - --max_train_steps 4000 - --train_batch_size 16 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 12 - --num_height 480 - --num_width 832 - --num_frames 45 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments (adjusted to 6 GPUs) -parallel_args=( - --num_gpus $NUM_GPUS -<<<<<<< HEAD - --sp_size 4 -======= - --sp_size 1 ->>>>>>> 98ce47c7 (checkpoint (qat attn in progress)) - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# Model arguments -model_args=( - --model_path "$MODEL_PATH" - --pretrained_model_name_or_path "$MODEL_PATH" - --real_score_model_path "$REAL_SCORE_MODEL_PATH" - --fake_score_model_path "$FAKE_SCORE_MODEL_PATH" -) - -# Dataset arguments -dataset_args=( - --data_path "$DATA_DIR" - --dataloader_num_workers 1 -) - -# Validation arguments -validation_args=( - --log_validation - --validation_dataset_file "$VALIDATION_DATASET_FILE" - --validation_steps 100 - --validation_sampling_steps "3" - --validation_guidance_scale "6.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 2e-6 - --mixed_precision "bf16" - --training_state_checkpointing_steps 2000 - --weight_only_checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.0 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 8 - --seed 1000 -) - -# DMD arguments -dmd_args=( - --dmd_denoising_steps '1000,757,522' - --min_timestep_ratio 0.02 - --max_timestep_ratio 0.98 - --generator_update_interval 5 - --real_score_guidance_scale 3.5 -) - -############################################ -# Launch (no rendezvous flags needed) -############################################ -set -x -torchrun \ - --nproc_per_node="$NUM_GPUS" \ - fastvideo/training/wan_distillation_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" \ - "${dmd_args[@]}" diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom.sh deleted file mode 100644 index bc258a031..000000000 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom.sh +++ /dev/null @@ -1,122 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################ -# Single node, 6 GPUs -############################################ -NUM_GPUS=1 -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -export WANDB_MODE="online" -export WANDB_BASE_URL="https://api.wandb.ai" -export TOKENIZERS_PARALLELISM=false -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -export TRITON_CACHE_DIR=/tmp/triton_cache - -echo "NUM_GPUS: $NUM_GPUS" -echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" - -############################################ -# Configs -############################################ -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/vidprom_16k_umt5_text_embed -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -OUTPUT_DIR="checkpoints/wan_1.3B_t2v_distill_dmd_vidprom" -REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - -# Training arguments -training_args=( - --tracker_project_name wan_t2v_distill_dmd_qat - --output_dir $OUTPUT_DIR - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 21 - --num_height 480 - --num_width 832 - --num_frames 81 - --enable_gradient_checkpointing_type "full" -) - -# Parallel arguments (adjusted to 6 GPUs) -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# Model arguments -model_args=( - --model_path $MODEL_PATH - --pretrained_model_name_or_path $MODEL_PATH - --real_score_model_path $REAL_SCORE_MODEL_PATH - --fake_score_model_path $FAKE_SCORE_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" - --validation_steps 200 - --validation_sampling_steps "3" - --validation_guidance_scale "6.0" # not used for dmd inference -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 2e-6 - --mixed_precision "bf16" - --training_state_checkpointing_steps 500 - --weight_only_checkpointing_steps 500 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.0 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 8 - --seed 1000 -) - -# DMD arguments -dmd_args=( - --dmd_denoising_steps '1000,757,522' - --min_timestep_ratio 0.02 - --max_timestep_ratio 0.98 - --generator_update_interval 5 - --real_score_guidance_scale 3.5 - --simulate_generator_forward -) - -############################################ -# Launch (no rendezvous flags needed) -############################################ -set -x -torchrun \ - --nproc_per_node="$NUM_GPUS" \ - fastvideo/training/wan_distillation_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" \ - "${dmd_args[@]}" \ No newline at end of file diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_e2e_fp4.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_e2e_fp4.sh deleted file mode 100644 index fb30044b2..000000000 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_e2e_fp4.sh +++ /dev/null @@ -1,124 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################ -# Single node, 6 GPUs -############################################ -NUM_GPUS=1 -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -export WANDB_MODE="online" -export WANDB_BASE_URL="https://api.wandb.ai" -export TOKENIZERS_PARALLELISM=false -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -export TRITON_CACHE_DIR=/tmp/triton_cache - -echo "NUM_GPUS: $NUM_GPUS" -echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" - -############################################ -# Configs -############################################ -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/vidprom_16k_umt5_text_embed -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -OUTPUT_DIR="checkpoints/wan_1.3B_t2v_distill_dmd_e2e_fp4_vidprom" -REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - -# Training arguments -training_args=( - --tracker_project_name wan_t2v_distill_dmd_qat - --output_dir $OUTPUT_DIR - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 21 - --num_height 480 - --num_width 832 - --num_frames 81 - --enable_gradient_checkpointing_type "full" - --generator_4bit_linear True - --generator_4bit_attn True -) - -# Parallel arguments (adjusted to 6 GPUs) -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# Model arguments -model_args=( - --model_path $MODEL_PATH - --pretrained_model_name_or_path $MODEL_PATH - --real_score_model_path $REAL_SCORE_MODEL_PATH - --fake_score_model_path $FAKE_SCORE_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" - --validation_steps 200 - --validation_sampling_steps "3" - --validation_guidance_scale "6.0" # not used for dmd inference -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 2e-6 - --mixed_precision "bf16" - --training_state_checkpointing_steps 500 - --weight_only_checkpointing_steps 500 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.0 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 8 - --seed 1000 -) - -# DMD arguments -dmd_args=( - --dmd_denoising_steps '1000,757,522' - --min_timestep_ratio 0.02 - --max_timestep_ratio 0.98 - --generator_update_interval 5 - --real_score_guidance_scale 3.5 - --simulate_generator_forward -) - -############################################ -# Launch (no rendezvous flags needed) -############################################ -set -x -torchrun \ - --nproc_per_node="$NUM_GPUS" \ - fastvideo/training/wan_distillation_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" \ - "${dmd_args[@]}" \ No newline at end of file diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_linear_fp4.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_linear_fp4.sh deleted file mode 100644 index c3aa922bb..000000000 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_linear_fp4.sh +++ /dev/null @@ -1,123 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################ -# Single node, 6 GPUs -############################################ -NUM_GPUS=1 -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -export WANDB_MODE="online" -export WANDB_BASE_URL="https://api.wandb.ai" -export TOKENIZERS_PARALLELISM=false -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -export TRITON_CACHE_DIR=/tmp/triton_cache - -echo "NUM_GPUS: $NUM_GPUS" -echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" - -############################################ -# Configs -############################################ -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/vidprom_16k_umt5_text_embed -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -OUTPUT_DIR="checkpoints/wan_1.3B_t2v_distill_dmd_linear_fp4_vidprom" -REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - -# Training arguments -training_args=( - --tracker_project_name wan_t2v_distill_dmd_qat - --output_dir $OUTPUT_DIR - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 21 - --num_height 480 - --num_width 832 - --num_frames 81 - --enable_gradient_checkpointing_type "full" - --generator_4bit_linear True -) - -# Parallel arguments (adjusted to 6 GPUs) -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# Model arguments -model_args=( - --model_path $MODEL_PATH - --pretrained_model_name_or_path $MODEL_PATH - --real_score_model_path $REAL_SCORE_MODEL_PATH - --fake_score_model_path $FAKE_SCORE_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" - --validation_steps 200 - --validation_sampling_steps "3" - --validation_guidance_scale "6.0" # not used for dmd inference -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 2e-6 - --mixed_precision "bf16" - --training_state_checkpointing_steps 500 - --weight_only_checkpointing_steps 500 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.0 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 8 - --seed 1000 -) - -# DMD arguments -dmd_args=( - --dmd_denoising_steps '1000,757,522' - --min_timestep_ratio 0.02 - --max_timestep_ratio 0.98 - --generator_update_interval 5 - --real_score_guidance_scale 3.5 - --simulate_generator_forward -) - -############################################ -# Launch (no rendezvous flags needed) -############################################ -set -x -torchrun \ - --nproc_per_node="$NUM_GPUS" \ - fastvideo/training/wan_distillation_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" \ - "${dmd_args[@]}" \ No newline at end of file diff --git a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_qat.sh b/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_qat.sh deleted file mode 100644 index 88e84cde0..000000000 --- a/examples/distill/Wan2.1-T2V/Wan-Syn-Data-480P/distill_dmd_t2v_1.3B_vidprom_qat.sh +++ /dev/null @@ -1,123 +0,0 @@ -#!/usr/bin/env bash -set -euo pipefail - -############################################ -# Single node, 6 GPUs -############################################ -NUM_GPUS=1 -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -export WANDB_MODE="online" -export WANDB_BASE_URL="https://api.wandb.ai" -export TOKENIZERS_PARALLELISM=false -export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN # or TORCH_SDPA -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -export TRITON_CACHE_DIR=/tmp/triton_cache - -echo "NUM_GPUS: $NUM_GPUS" -echo "CUDA_VISIBLE_DEVICES: $CUDA_VISIBLE_DEVICES" - -############################################ -# Configs -############################################ -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/vidprom_16k_umt5_text_embed -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -OUTPUT_DIR="checkpoints/wan_1.3B_t2v_distill_dmd_qat_vidprom" -REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" - -# Training arguments -training_args=( - --tracker_project_name wan_t2v_distill_dmd_qat - --output_dir $OUTPUT_DIR - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 21 - --num_height 480 - --num_width 832 - --num_frames 81 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments (adjusted to 6 GPUs) -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# Model arguments -model_args=( - --model_path $MODEL_PATH - --pretrained_model_name_or_path $MODEL_PATH - --real_score_model_path $REAL_SCORE_MODEL_PATH - --fake_score_model_path $FAKE_SCORE_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" - --validation_steps 200 - --validation_sampling_steps "3" - --validation_guidance_scale "6.0" # not used for dmd inference -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 2e-6 - --mixed_precision "bf16" - --training_state_checkpointing_steps 500 - --weight_only_checkpointing_steps 500 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.0 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 8 - --seed 1000 -) - -# DMD arguments -dmd_args=( - --dmd_denoising_steps '1000,757,522' - --min_timestep_ratio 0.02 - --max_timestep_ratio 0.98 - --generator_update_interval 5 - --real_score_guidance_scale 3.5 - --simulate_generator_forward -) - -############################################ -# Launch (no rendezvous flags needed) -############################################ -set -x -torchrun \ - --nproc_per_node="$NUM_GPUS" \ - fastvideo/training/wan_distillation_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" \ - "${dmd_args[@]}" \ No newline at end of file diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_e2e.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_e2e.sh deleted file mode 100644 index 9bc207f98..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_e2e.sh +++ /dev/null @@ -1,98 +0,0 @@ -#!/bin/bash - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=1 -# export CUDA_VISIBLE_DEVICES=4,5 - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_fp4_e2e_77" - --max_train_steps 5000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 448 - --num_width 832 - --num_frames 77 - --enable_gradient_checkpointing_type "full" - --generator_4bit_linear True - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size $NUM_GPUS - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# 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 - --validation_steps 100 - --validation_sampling_steps "50" - --validation_guidance_scale "6.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 500 - --training_state_checkpointing_steps 500 - --weight_decay 1e-4 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --multi_phased_distill_schedule "4000-1" - --not_apply_cfg_solver - --dit_precision "fp32" - --num_euler_timesteps 50 - --ema_start_step 0 - --enable_gradient_checkpointing_type "full" - # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500" -) - -torchrun \ - --nnodes 1 \ - --nproc_per_node $NUM_GPUS \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_linear.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_linear.sh deleted file mode 100644 index c86825555..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_fp4_linear.sh +++ /dev/null @@ -1,97 +0,0 @@ -#!/bin/bash - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=1 -# export CUDA_VISIBLE_DEVICES=4,5 - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_fp4_linear_77" - --max_train_steps 5000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 448 - --num_width 832 - --num_frames 77 - --enable_gradient_checkpointing_type "full" - --generator_4bit_linear True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size $NUM_GPUS - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# 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 - --validation_steps 100 - --validation_sampling_steps "50" - --validation_guidance_scale "6.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 500 - --training_state_checkpointing_steps 500 - --weight_decay 1e-4 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --multi_phased_distill_schedule "4000-1" - --not_apply_cfg_solver - --dit_precision "fp32" - --num_euler_timesteps 50 - --ema_start_step 0 - --enable_gradient_checkpointing_type "full" - # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500" -) - -torchrun \ - --nnodes 1 \ - --nproc_per_node $NUM_GPUS \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4_5090.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4_5090.sh deleted file mode 100644 index 56c7da9fd..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_nvfp4_5090.sh +++ /dev/null @@ -1,94 +0,0 @@ -#!/bin/bash - -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/Wan-Syn_77x448x832_600k/train/" -VALIDATION_DATASET_FILE="data/Wan-Syn_77x448x832_600k/validation_6.json" -NUM_GPUS=6 -# export CUDA_VISIBLE_DEVICES=4,5 - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune" - --output_dir "checkpoints/wan_t2v_finetune" - --max_train_steps 5000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 6 - --num_height 480 - --num_width 832 - --num_frames 21 - --enable_gradient_checkpointing_type "full" - --generator_4bit_linear True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# 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 1 -) - -# Validation arguments -validation_args=( - --log_validation - --validation_dataset_file $VALIDATION_DATASET_FILE - --validation_steps 50 - --validation_sampling_steps "50" - --validation_guidance_scale "6.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --checkpointing_steps 2000 - --weight_decay 1e-4 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --multi_phased_distill_schedule "4000-1" - --not_apply_cfg_solver - --dit_precision "fp32" - --num_euler_timesteps 50 - --ema_start_step 0 - --enable_gradient_checkpointing_type "full" - # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500" -) - -torchrun \ - --nnodes 1 \ - --nproc_per_node $NUM_GPUS \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_16_M2.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_16_M2.sh deleted file mode 100644 index ffc7220e3..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_16_M2.sh +++ /dev/null @@ -1,124 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_1.3B_finetune_qat_16 -#SBATCH --partition=main -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=1 -#SBATCH --gres=gpu:8 -#SBATCH --cpus-per-task=128 -#SBATCH --mem=1440G -#SBATCH --output=logs/wan_t2v_1.3B_finetune_qat_16.out -#SBATCH --error=logs/wan_t2v_1.3B_finetune_qat_16.err -#SBATCH --exclusive - -source ~/conda/miniconda/bin/activate -conda activate matthew-fv - -# Basic Info -export WANDB_MODE="online" -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -# different cache dir for different processes -export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} -export MASTER_PORT=29500 -export NODE_RANK=$SLURM_PROCID -nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) -export MASTER_ADDR=${nodes[0]} -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" - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=/mnt/weka/home/hao.zhang/wl/sharefs/Vchitect-2M/Wan-Syn-upload/latents_i2v/train -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS_PER_NODE=8 -TOTAL_GPUS=$((NUM_GPUS_PER_NODE * SLURM_JOB_NUM_NODES)) -# export CUDA_VISIBLE_DEVICES=4,5 - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_qat_16" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 448 - --num_width 832 - --num_frames 77 - --enable_gradient_checkpointing_type "full" # if OOM enable this - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $TOTAL_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim $TOTAL_GPUS - --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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 200 - --training_state_checkpointing_steps 200 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 1 - --seed 1000 -) - -srun torchrun \ ---nnodes $SLURM_JOB_NUM_NODES \ ---nproc_per_node $NUM_GPUS_PER_NODE \ ---node_rank $SLURM_PROCID \ ---rdzv_backend=c10d \ ---rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_32.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_32.sh deleted file mode 100644 index f894354a1..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_32.sh +++ /dev/null @@ -1,116 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_1.3B_finetune_qat_32 -#SBATCH --partition=all -#SBATCH --nodes=8 -#SBATCH --gres=gpu:4 -#SBATCH --ntasks-per-node=1 -#SBATCH --output=logs/wan_t2v_1.3B_finetune_qat_32.out -#SBATCH --error=logs/wan_t2v_1.3B_finetune_qat_32.err - -source .venv/bin/activate - - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=32 -# export CUDA_VISIBLE_DEVICES=4,5 - -set -euo pipefail - -# ---- torchrun rendezvous (multi-node) ---- -# Launch ONE torchrun per node (via srun) and let torchrun spawn 4 workers per node. -MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)" -MASTER_PORT="${MASTER_PORT:-29500}" -export MASTER_ADDR MASTER_PORT - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_qat_32" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 448 - --num_width 832 - --num_frames 77 - --enable_gradient_checkpointing_type "full" # if OOM enable this - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim $NUM_GPUS - --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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 200 - --training_state_checkpointing_steps 200 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 1 - --seed 1000 -) - -srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ - torchrun \ - --nnodes "$SLURM_NNODES" \ - --nproc_per_node 4 \ - --rdzv_backend c10d \ - --rdzv_endpoint "${MASTER_ADDR}:${MASTER_PORT}" \ - --rdzv_id "$SLURM_JOB_ID" \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_32_M2.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_32_M2.sh deleted file mode 100644 index b16aaa17d..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_32_M2.sh +++ /dev/null @@ -1,124 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_1.3B_finetune_qat_32 -#SBATCH --partition=main -#SBATCH --nodes=4 -#SBATCH --ntasks-per-node=1 -#SBATCH --gres=gpu:8 -#SBATCH --cpus-per-task=128 -#SBATCH --mem=1440G -#SBATCH --output=logs/wan_t2v_1.3B_finetune_qat_32.out -#SBATCH --error=logs/wan_t2v_1.3B_finetune_qat_32.err -#SBATCH --exclusive - -source ~/conda/miniconda/bin/activate -conda activate matthew-fv - -# Basic Info -export WANDB_MODE="online" -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -# different cache dir for different processes -export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} -export MASTER_PORT=29500 -export NODE_RANK=$SLURM_PROCID -nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) -export MASTER_ADDR=${nodes[0]} -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" - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=/mnt/weka/home/hao.zhang/wl/sharefs/Vchitect-2M/Wan-Syn-upload/latents_i2v/train -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS_PER_NODE=8 -TOTAL_GPUS=$((NUM_GPUS_PER_NODE * SLURM_JOB_NUM_NODES)) -# export CUDA_VISIBLE_DEVICES=4,5 - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_qat_32" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 448 - --num_width 832 - --num_frames 77 - --enable_gradient_checkpointing_type "full" # if OOM enable this - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $TOTAL_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim $TOTAL_GPUS - --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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 200 - --training_state_checkpointing_steps 200 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 1 - --seed 1000 -) - -srun torchrun \ ---nnodes $SLURM_JOB_NUM_NODES \ ---nproc_per_node $NUM_GPUS_PER_NODE \ ---node_rank $SLURM_PROCID \ ---rdzv_backend=c10d \ ---rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_5090.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_5090.sh deleted file mode 100644 index 3ff8498bb..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_5090.sh +++ /dev/null @@ -1,97 +0,0 @@ -#!/bin/bash - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=4 -# export CUDA_VISIBLE_DEVICES=4,5 - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_qat" - --max_train_steps 5000 - --train_batch_size 8 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 12 - --num_height 480 - --num_width 832 - --num_frames 45 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size $NUM_GPUS - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# 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 - --validation_steps 100 - --validation_sampling_steps "50" - --validation_guidance_scale "6.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 500 - --training_state_checkpointing_steps 500 - --weight_decay 1e-4 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --multi_phased_distill_schedule "4000-1" - --not_apply_cfg_solver - --dit_precision "fp32" - --num_euler_timesteps 50 - --ema_start_step 0 - --enable_gradient_checkpointing_type "full" - # --resume_from_checkpoint "checkpoints/wan_t2v_finetune/checkpoint-2500" -) - -torchrun \ - --nnodes 1 \ - --nproc_per_node $NUM_GPUS \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8.sh deleted file mode 100644 index 21008bade..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8.sh +++ /dev/null @@ -1,119 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_1.3B_finetune_qat_8 -#SBATCH --partition=all -#SBATCH --nodes=2 -#SBATCH --gres=gpu:4 -#SBATCH --ntasks-per-node=1 -#SBATCH --output=logs/wan_t2v_1.3B_finetune_qat_8.out -#SBATCH --error=logs/wan_t2v_1.3B_finetune_qat_8.err - -source .venv/bin/activate - - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# Use node-local Triton cache to avoid stale file handle errors on shared filesystems -export TRITON_CACHE_DIR="/tmp/triton_cache_${SLURM_JOB_ID}_${SLURM_NODEID}" - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=8 -# export CUDA_VISIBLE_DEVICES=4,5 - -set -euo pipefail - -# ---- torchrun rendezvous (multi-node) ---- -# Launch ONE torchrun per node (via srun) and let torchrun spawn 4 workers per node. -MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)" -MASTER_PORT="${MASTER_PORT:-29500}" -export MASTER_ADDR MASTER_PORT - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_qat_8" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 448 - --num_width 832 - --num_frames 77 - --enable_gradient_checkpointing_type "full" # if OOM enable this - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim $NUM_GPUS - --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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 200 - --training_state_checkpointing_steps 200 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 1 - --seed 1000 -) - -srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ - torchrun \ - --nnodes "$SLURM_NNODES" \ - --nproc_per_node 4 \ - --rdzv_backend c10d \ - --rdzv_endpoint "${MASTER_ADDR}:${MASTER_PORT}" \ - --rdzv_id "$SLURM_JOB_ID" \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8_M2.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8_M2.sh deleted file mode 100644 index 13dd80f3a..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8_M2.sh +++ /dev/null @@ -1,123 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_1.3B_finetune_qat_8 -#SBATCH --partition=main -#SBATCH --nodes=1 -#SBATCH --ntasks-per-node=1 -#SBATCH --gres=gpu:8 -#SBATCH --cpus-per-task=128 -#SBATCH --mem=1440G -#SBATCH --output=logs/wan_t2v_1.3B_finetune_qat_8.out -#SBATCH --error=logs/wan_t2v_1.3B_finetune_qat_8.err -#SBATCH --exclusive - -source ~/conda/miniconda/bin/activate -conda activate matthew-fv - -# Basic Info -export WANDB_MODE="online" -export NCCL_P2P_DISABLE=1 -export TORCH_NCCL_ENABLE_MONITORING=0 -# different cache dir for different processes -export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID} -export MASTER_PORT=29500 -export NODE_RANK=$SLURM_PROCID -nodes=( $(scontrol show hostnames $SLURM_JOB_NODELIST) ) -export MASTER_ADDR=${nodes[0]} -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" - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=/mnt/weka/home/hao.zhang/wl/sharefs/Vchitect-2M/Wan-Syn-upload/latents_i2v/train -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=8 -# export CUDA_VISIBLE_DEVICES=4,5 - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_qat_8" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 448 - --num_width 832 - --num_frames 77 - --enable_gradient_checkpointing_type "full" # if OOM enable this - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim $NUM_GPUS - --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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 200 - --training_state_checkpointing_steps 200 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 1 - --seed 1000 -) - -srun torchrun \ ---nnodes $SLURM_JOB_NUM_NODES \ ---nproc_per_node $NUM_GPUS \ ---node_rank $SLURM_PROCID \ ---rdzv_backend=c10d \ ---rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8_NVL72.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8_NVL72.sh deleted file mode 100644 index 7eb2db8bd..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_8_NVL72.sh +++ /dev/null @@ -1,121 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_1.3B_finetune_qat_8_NVL72 -#SBATCH --nodes=4 -#SBATCH --gres=gpu:4 -#SBATCH --ntasks-per-node=1 -#SBATCH --output=logs/wan_t2v_1.3B_finetune_qat_8_NVL72.out -#SBATCH --error=logs/wan_t2v_1.3B_finetune_qat_8_NVL72.err - -source .venv/bin/activate - - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -# export CUDA_VISIBLE_DEVICES=4,5 - -set -euo pipefail - -# ---- torchrun rendezvous (multi-node) ---- -# 1. Get the hostname of the first node (Master) -nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) ) -nodes_array=($nodes) -head_node=${nodes_array[0]} -MASTER_ADDR=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address) -MASTER_PORT=29500 - -# 2. Get the node count automatically -NNODES=$SLURM_NNODES -GPUS_PER_NODE=$SLURM_GPUS_ON_NODE -NUM_GPUS=$((NNODES * GPUS_PER_NODE)) - -echo "MASTER_ADDR=$MASTER_ADDR MASTER_PORT=$MASTER_PORT NNODES=$NNODES" - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_14B" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 768 - --num_width 1280 - --num_frames 77 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 4 - --tp_size 1 - --hsdp_replicate_dim 4 - --hsdp_shard_dim 4 -) - -# 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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 1000 - --training_state_checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 5 - --seed 1000 -) - -srun torchrun \ - --nnodes $NNODES \ - --nproc_per_node $GPUS_PER_NODE \ - --node_rank $SLURM_PROCID \ - --rdzv_backend c10d \ - --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_16.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh similarity index 100% rename from examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_16.sh rename to examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_qat_attn.sh diff --git a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh b/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh deleted file mode 100644 index bcfd77cd4..000000000 --- a/examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_sage3_5090.sh +++ /dev/null @@ -1,94 +0,0 @@ -#!/bin/bash - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=4 -# export CUDA_VISIBLE_DEVICES=4,5 - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_sage3" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 16 - --num_latent_t 20 - --num_height 448 - --num_width 832 - --num_frames 77 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 4 - --tp_size 1 - --hsdp_replicate_dim $NUM_GPUS - --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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 200 - --training_state_checkpointing_steps 200 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 1 - --seed 1000 -) - -torchrun \ - --nnodes 1 \ - --nproc_per_node $NUM_GPUS \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200.sh deleted file mode 100644 index 01d612c9b..000000000 --- a/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200.sh +++ /dev/null @@ -1,115 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_14B_finetune_B200 -#SBATCH --partition=all -#SBATCH --nodes=16 -#SBATCH --gres=gpu:4 -#SBATCH --ntasks-per-node=1 -#SBATCH --output=logs/wan_t2v_14B_finetune_B200.out -#SBATCH --error=logs/wan_t2v_14B_finetune_B200.err - -source .venv/bin/activate - - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" -DATA_DIR=data/Wan-Syn_77x768x1280_250k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_14B/validation.json" -NUM_GPUS=64 -# export CUDA_VISIBLE_DEVICES=4,5 - -set -euo pipefail - -# ---- torchrun rendezvous (multi-node) ---- -# Launch ONE torchrun per node (via srun) and let torchrun spawn 4 workers per node. -MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)" -MASTER_PORT="${MASTER_PORT:-29500}" -export MASTER_ADDR MASTER_PORT - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_14B" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 768 - --num_width 1280 - --num_frames 77 - --enable_gradient_checkpointing_type "full" -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# 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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 1000 - --training_state_checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 5 - --seed 1000 -) - -srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ - torchrun \ - --nnodes "$SLURM_NNODES" \ - --nproc_per_node 4 \ - --rdzv_backend c10d \ - --rdzv_endpoint "${MASTER_ADDR}:${MASTER_PORT}" \ - --rdzv_id "$SLURM_JOB_ID" \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200_qat_debug.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200_qat_debug.sh deleted file mode 100644 index 2876fabb6..000000000 --- a/examples/training/finetune/wan_t2v_14B/finetune_t2v_B200_qat_debug.sh +++ /dev/null @@ -1,125 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_14B_finetune_B200 -#SBATCH --nodes=1 -#SBATCH --gres=gpu:4 -#SBATCH --ntasks-per-node=1 -#SBATCH --output=logs/wan_t2v_14B_finetune_B200.out -#SBATCH --error=logs/wan_t2v_14B_finetune_B200.err - -source .venv/bin/activate - -export NCCL_P2P_DISABLE=1 - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# Use node-local Triton cache to avoid stale file handle errors on shared filesystems -export TRITON_CACHE_DIR="/tmp/triton_cache_${SLURM_JOB_ID}_${SLURM_NODEID}" - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -# export CUDA_VISIBLE_DEVICES=4,5 - -set -euo pipefail - -# ---- torchrun rendezvous (multi-node) ---- -# 1. Get the hostname of the first node (Master) -nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) ) -nodes_array=($nodes) -head_node=${nodes_array[0]} -MASTER_ADDR=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address) -MASTER_PORT=29500 - -# 2. Get the node count automatically -NNODES=$SLURM_NNODES -GPUS_PER_NODE=$SLURM_GPUS_ON_NODE -NUM_GPUS=$((NNODES * GPUS_PER_NODE)) - -echo "MASTER_ADDR=$MASTER_ADDR MASTER_PORT=$MASTER_PORT NNODES=$NNODES" - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_1.3B_debug" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 768 - --num_width 1280 - --num_frames 77 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim 4 -) - -# 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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 1000 - --training_state_checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 5 - --seed 1000 -) - -srun torchrun \ - --nnodes $NNODES \ - --nproc_per_node $GPUS_PER_NODE \ - --node_rank $SLURM_PROCID \ - --rdzv_backend c10d \ - --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat.sh deleted file mode 100644 index 7d1082a42..000000000 --- a/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat.sh +++ /dev/null @@ -1,93 +0,0 @@ -#!/bin/bash - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=4 -# export CUDA_VISIBLE_DEVICES=4,5 - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_qat" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 21 - --num_height 720 - --num_width 1280 - --num_frames 81 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# 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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-5 - --mixed_precision "bf16" - --checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 5 - --seed 1000 -) - -torchrun \ - --nnodes 1 \ - --nproc_per_node $NUM_GPUS \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_B200.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_B200.sh deleted file mode 100644 index 99a693d0a..000000000 --- a/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_B200.sh +++ /dev/null @@ -1,116 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_14B_finetune_B200 -#SBATCH --partition=all -#SBATCH --nodes=16 -#SBATCH --gres=gpu:4 -#SBATCH --ntasks-per-node=1 -#SBATCH --output=logs/wan_t2v_14B_finetune_B200.out -#SBATCH --error=logs/wan_t2v_14B_finetune_B200.err - -source .venv/bin/activate - - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" -DATA_DIR=data/Wan-Syn_77x768x1280_250k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_14B/validation.json" -NUM_GPUS=64 -# export CUDA_VISIBLE_DEVICES=4,5 - -set -euo pipefail - -# ---- torchrun rendezvous (multi-node) ---- -# Launch ONE torchrun per node (via srun) and let torchrun spawn 4 workers per node. -MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)" -MASTER_PORT="${MASTER_PORT:-29500}" -export MASTER_ADDR MASTER_PORT - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_14B" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 20 - --num_height 768 - --num_width 1280 - --num_frames 77 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 1 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# 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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-6 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 1000 - --training_state_checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 5 - --seed 1000 -) - -srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ - torchrun \ - --nnodes "$SLURM_NNODES" \ - --nproc_per_node 4 \ - --rdzv_backend c10d \ - --rdzv_endpoint "${MASTER_ADDR}:${MASTER_PORT}" \ - --rdzv_id "$SLURM_JOB_ID" \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_32_14B_M2.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh similarity index 100% rename from examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_32_14B_M2.sh rename to examples/training/finetune/wan_t2v_14B/finetune_t2v_qat_attn.sh diff --git a/examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh b/examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh deleted file mode 100644 index 8849c7246..000000000 --- a/examples/training/finetune/wan_t2v_14B/finetune_t2v_sage3.sh +++ /dev/null @@ -1,113 +0,0 @@ -#!/bin/bash -#SBATCH --job-name=wan_t2v_14B_finetune_sage3 -#SBATCH --partition=all -#SBATCH --nodes=2 -#SBATCH --gres=gpu:4 -#SBATCH --ntasks-per-node=4 -#SBATCH --output=logs/wan_t2v_14B_finetune_sage3_%j.out -#SBATCH --error=logs/wan_t2v_14B_finetune_sage3_%j.err - -export WANDB_BASE_URL="https://api.wandb.ai" -export WANDB_MODE=online -export TOKENIZERS_PARALLELISM=false -# export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA - -# export TRITON_PRINT_AUTOTUNING=1 # to print the best config -export WANDB_API_KEY=2f25ad37933894dbf0966c838c0b8494987f9f2f -MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" -DATA_DIR=data/Wan-Syn_77x448x832_600k -VALIDATION_DATASET_FILE="examples/training/finetune/wan_t2v_1.3B/crush_smol/validation.json" -NUM_GPUS=8 -# export CUDA_VISIBLE_DEVICES=4,5 - -set -euo pipefail - -# ---- torchrun rendezvous (multi-node) ---- -# Launch ONE torchrun per node (via srun) and let torchrun spawn 4 workers per node. -MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)" -MASTER_PORT="${MASTER_PORT:-29500}" -export MASTER_ADDR MASTER_PORT - - -# Training arguments -training_args=( - --tracker_project_name "wan_t2v_finetune_qat" - --output_dir "checkpoints/wan_t2v_finetune_sage3_14B" - --max_train_steps 4000 - --train_batch_size 1 - --train_sp_batch_size 1 - --gradient_accumulation_steps 1 - --num_latent_t 32 - --num_height 720 - --num_width 1280 - --num_frames 125 - --enable_gradient_checkpointing_type "full" - --generator_4bit_attn True -) - -# Parallel arguments -parallel_args=( - --num_gpus $NUM_GPUS - --sp_size 8 - --tp_size 1 - --hsdp_replicate_dim 1 - --hsdp_shard_dim $NUM_GPUS -) - -# 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 - --validation_steps 200 - --validation_sampling_steps "50" - --validation_guidance_scale "5.0" -) - -# Optimizer arguments -optimizer_args=( - --learning_rate 1e-5 - --mixed_precision "bf16" - --weight_only_checkpointing_steps 1000 - --training_state_checkpointing_steps 1000 - --weight_decay 0.01 - --max_grad_norm 1.0 -) - -# Miscellaneous arguments -miscellaneous_args=( - --inference_mode False - --checkpoints_total_limit 3 - --training_cfg_rate 0.1 - --dit_precision "fp32" - --ema_start_step 0 - --flow_shift 5 - --seed 1000 -) - -srun --nodes="$SLURM_NNODES" --ntasks="$SLURM_NNODES" --ntasks-per-node=1 \ - torchrun \ - --nnodes "$SLURM_NNODES" \ - --nproc_per_node 4 \ - --rdzv_backend c10d \ - --rdzv_endpoint "${MASTER_ADDR}:${MASTER_PORT}" \ - --node_rank "$SLURM_NODEID" \ - fastvideo/training/wan_training_pipeline.py \ - "${parallel_args[@]}" \ - "${model_args[@]}" \ - "${dataset_args[@]}" \ - "${training_args[@]}" \ - "${optimizer_args[@]}" \ - "${validation_args[@]}" \ - "${miscellaneous_args[@]}" diff --git a/fastvideo/attention/backends/sageattn/__init__.py b/fastvideo-kernel/modified_sageattn/__init__.py similarity index 100% rename from fastvideo/attention/backends/sageattn/__init__.py rename to fastvideo-kernel/modified_sageattn/__init__.py diff --git a/fastvideo/attention/backends/sageattn/api.py b/fastvideo-kernel/modified_sageattn/api.py similarity index 100% rename from fastvideo/attention/backends/sageattn/api.py rename to fastvideo-kernel/modified_sageattn/api.py diff --git a/fastvideo/attention/backends/sageattn/blackwell/__init__.py b/fastvideo-kernel/modified_sageattn/blackwell/__init__.py similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/__init__.py rename to fastvideo-kernel/modified_sageattn/blackwell/__init__.py diff --git a/fastvideo/attention/backends/sageattn/blackwell/api.cu b/fastvideo-kernel/modified_sageattn/blackwell/api.cu similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/api.cu rename to fastvideo-kernel/modified_sageattn/blackwell/api.cu diff --git a/fastvideo/attention/backends/sageattn/blackwell/block_config.h b/fastvideo-kernel/modified_sageattn/blackwell/block_config.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/block_config.h rename to fastvideo-kernel/modified_sageattn/blackwell/block_config.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/block_info.h b/fastvideo-kernel/modified_sageattn/blackwell/block_info.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/block_info.h rename to fastvideo-kernel/modified_sageattn/blackwell/block_info.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/blockscaled_layout.h b/fastvideo-kernel/modified_sageattn/blackwell/blockscaled_layout.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/blockscaled_layout.h rename to fastvideo-kernel/modified_sageattn/blackwell/blockscaled_layout.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/cute_extension.h b/fastvideo-kernel/modified_sageattn/blackwell/cute_extension.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/cute_extension.h rename to fastvideo-kernel/modified_sageattn/blackwell/cute_extension.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/epilogue_tma_ws.h b/fastvideo-kernel/modified_sageattn/blackwell/epilogue_tma_ws.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/epilogue_tma_ws.h rename to fastvideo-kernel/modified_sageattn/blackwell/epilogue_tma_ws.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/kernel_traits.h b/fastvideo-kernel/modified_sageattn/blackwell/kernel_traits.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/kernel_traits.h rename to fastvideo-kernel/modified_sageattn/blackwell/kernel_traits.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/kernel_ws.h b/fastvideo-kernel/modified_sageattn/blackwell/kernel_ws.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/kernel_ws.h rename to fastvideo-kernel/modified_sageattn/blackwell/kernel_ws.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/launch.h b/fastvideo-kernel/modified_sageattn/blackwell/launch.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/launch.h rename to fastvideo-kernel/modified_sageattn/blackwell/launch.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/mainloop_tma_ws.h b/fastvideo-kernel/modified_sageattn/blackwell/mainloop_tma_ws.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/mainloop_tma_ws.h rename to fastvideo-kernel/modified_sageattn/blackwell/mainloop_tma_ws.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/named_barrier.h b/fastvideo-kernel/modified_sageattn/blackwell/named_barrier.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/named_barrier.h rename to fastvideo-kernel/modified_sageattn/blackwell/named_barrier.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/params.h b/fastvideo-kernel/modified_sageattn/blackwell/params.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/params.h rename to fastvideo-kernel/modified_sageattn/blackwell/params.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/softmax_fused.h b/fastvideo-kernel/modified_sageattn/blackwell/softmax_fused.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/softmax_fused.h rename to fastvideo-kernel/modified_sageattn/blackwell/softmax_fused.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/static_switch.h b/fastvideo-kernel/modified_sageattn/blackwell/static_switch.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/static_switch.h rename to fastvideo-kernel/modified_sageattn/blackwell/static_switch.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/tile_scheduler.h b/fastvideo-kernel/modified_sageattn/blackwell/tile_scheduler.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/tile_scheduler.h rename to fastvideo-kernel/modified_sageattn/blackwell/tile_scheduler.h diff --git a/fastvideo/attention/backends/sageattn/blackwell/utils.h b/fastvideo-kernel/modified_sageattn/blackwell/utils.h similarity index 100% rename from fastvideo/attention/backends/sageattn/blackwell/utils.h rename to fastvideo-kernel/modified_sageattn/blackwell/utils.h diff --git a/fastvideo/attention/backends/sageattn/quantization/__init__.py b/fastvideo-kernel/modified_sageattn/quantization/__init__.py similarity index 100% rename from fastvideo/attention/backends/sageattn/quantization/__init__.py rename to fastvideo-kernel/modified_sageattn/quantization/__init__.py diff --git a/fastvideo/attention/backends/sageattn/quantization/bench/bench_quant_k.py b/fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_k.py similarity index 100% rename from fastvideo/attention/backends/sageattn/quantization/bench/bench_quant_k.py rename to fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_k.py diff --git a/fastvideo/attention/backends/sageattn/quantization/bench/bench_quant_q.py b/fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_q.py similarity index 100% rename from fastvideo/attention/backends/sageattn/quantization/bench/bench_quant_q.py rename to fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_q.py diff --git a/fastvideo/attention/backends/sageattn/quantization/bench/bench_quant_v.py b/fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_v.py similarity index 100% rename from fastvideo/attention/backends/sageattn/quantization/bench/bench_quant_v.py rename to fastvideo-kernel/modified_sageattn/quantization/bench/bench_quant_v.py diff --git a/fastvideo/attention/backends/sageattn/quantization/bench/bench_utils.py b/fastvideo-kernel/modified_sageattn/quantization/bench/bench_utils.py similarity index 100% rename from fastvideo/attention/backends/sageattn/quantization/bench/bench_utils.py rename to fastvideo-kernel/modified_sageattn/quantization/bench/bench_utils.py diff --git a/fastvideo/attention/backends/sageattn/quantization/cuda_utils.h b/fastvideo-kernel/modified_sageattn/quantization/cuda_utils.h similarity index 100% rename from fastvideo/attention/backends/sageattn/quantization/cuda_utils.h rename to fastvideo-kernel/modified_sageattn/quantization/cuda_utils.h diff --git a/fastvideo/attention/backends/sageattn/quantization/fp4_quantization_4d.cu b/fastvideo-kernel/modified_sageattn/quantization/fp4_quantization_4d.cu similarity index 100% rename from fastvideo/attention/backends/sageattn/quantization/fp4_quantization_4d.cu rename to fastvideo-kernel/modified_sageattn/quantization/fp4_quantization_4d.cu diff --git a/nvfp4_utils.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py similarity index 100% rename from nvfp4_utils.py rename to fastvideo-kernel/python/fastvideo_kernel/triton_kernels/nvfp4_utils.py diff --git a/qat_attn.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/qat_attn.py similarity index 99% rename from qat_attn.py rename to fastvideo-kernel/python/fastvideo_kernel/triton_kernels/qat_attn.py index 0f792e502..159f586c3 100644 --- a/qat_attn.py +++ b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/qat_attn.py @@ -8,8 +8,8 @@ import triton 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 fastvideo.attention.backends.sageattn.api import triton_group_mean +from .quant_utils import fake_quantize_q, fake_quantize_kv, fake_quantize +# from modified_sageattn.api import triton_group_mean def is_cuda(): @@ -1116,4 +1116,4 @@ class _attention(torch.autograd.Function): return dq, dk, dv, None, None, None, None, None, None, None, None, None, None, None, None -attention = _attention.apply \ No newline at end of file +attention = _attention.apply diff --git a/quant_utils.py b/fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py similarity index 100% rename from quant_utils.py rename to fastvideo-kernel/python/fastvideo_kernel/triton_kernels/quant_utils.py diff --git a/test_fake_quant.py b/fastvideo-kernel/tests/test_fake_quant.py similarity index 99% rename from test_fake_quant.py rename to fastvideo-kernel/tests/test_fake_quant.py index f544ee08d..00be67f8b 100644 --- a/test_fake_quant.py +++ b/fastvideo-kernel/tests/test_fake_quant.py @@ -10,7 +10,10 @@ import torch import triton import triton.language as tl from flashinfer import SfLayout, nvfp4_quantize, e2m1_and_ufp8sf_scale_to_float -from nvfp4_utils import _compute_quant_and_scale, _compute_dequant +from fastvideo_kernel.triton_kernels.nvfp4_utils import ( + _compute_quant_and_scale, + _compute_dequant, +) from typing import Optional # MXFP_BLOCK_SIZE is 16 - use Python int for runtime checks diff --git a/test_qat_attn.py b/fastvideo-kernel/tests/test_qat_attn.py similarity index 99% rename from test_qat_attn.py rename to fastvideo-kernel/tests/test_qat_attn.py index d36ad3a85..07a6e51fc 100644 --- a/test_qat_attn.py +++ b/fastvideo-kernel/tests/test_qat_attn.py @@ -5,8 +5,8 @@ Tests forward and backward passes with and without QAT enabled. """ import torch -from qat_attn import _attention -from fused_attention import attention as fused_attention +from fastvideo_kernel.triton_kernels.qat_attn import _attention +from fastvideo_kernel.triton_kernels.fused_attention import attention as fused_attention from math import sqrt attention = _attention.apply @@ -1500,4 +1500,4 @@ if __name__ == "__main__": # test_fused_vs_qat_attention_causal_backward() print() - print("All tests passed! ✓") \ No newline at end of file + print("All tests passed! ✓") diff --git a/fastvideo/attention/backends/benchmark_scripts/benchmark_blockscaled_fp4_attn.py b/fastvideo/attention/backends/benchmark_scripts/benchmark_blockscaled_fp4_attn.py index 2f49d1f26..fe487de45 100644 --- a/fastvideo/attention/backends/benchmark_scripts/benchmark_blockscaled_fp4_attn.py +++ b/fastvideo/attention/backends/benchmark_scripts/benchmark_blockscaled_fp4_attn.py @@ -3,14 +3,14 @@ import sys import os import traceback -from fastvideo.attention.backends.sageattn.api import ( +from modified_sageattn.api import ( blockscaled_fp4_attn, preprocess_qkv, scale_and_quant_fp4, scale_and_quant_fp4_permute, scale_and_quant_fp4_transpose, ) -from fastvideo.attention.backends.sageattn.quantization.bench.bench_utils import bench +from modified_sageattn.quantization.bench.bench_utils import bench def calculate_attention_flops(batch_size, num_heads, seq_len_q, seq_len_k, head_dim, is_causal=False): @@ -281,4 +281,3 @@ if __name__ == "__main__": print(" excluding quantization overhead.") sys.stdout.flush() run_benchmark_suite() - diff --git a/fastvideo/attention/backends/benchmark_scripts/benchmark_combined.py b/fastvideo/attention/backends/benchmark_scripts/benchmark_combined.py index c0c60d743..e7af8d6d4 100644 --- a/fastvideo/attention/backends/benchmark_scripts/benchmark_combined.py +++ b/fastvideo/attention/backends/benchmark_scripts/benchmark_combined.py @@ -9,10 +9,10 @@ import numpy as np from typing import Dict, List, Tuple, Optional from flash_attn import flash_attn_func -from fastvideo.attention.backends.sageattn.quantization.bench.bench_utils import bench +from modified_sageattn.quantization.bench.bench_utils import bench # Import SageAttn components for direct control -from fastvideo.attention.backends.sageattn.api import ( +from modified_sageattn.api import ( preprocess_qkv, scale_and_quant_fp4, scale_and_quant_fp4_permute, @@ -377,4 +377,3 @@ if __name__ == "__main__": seq_lens=args.seq_lens, output_file=args.output ) - diff --git a/fastvideo/attention/backends/benchmark_scripts/benchmark_flashattn2.py b/fastvideo/attention/backends/benchmark_scripts/benchmark_flashattn2.py index a3a42e06e..fd79eabcc 100644 --- a/fastvideo/attention/backends/benchmark_scripts/benchmark_flashattn2.py +++ b/fastvideo/attention/backends/benchmark_scripts/benchmark_flashattn2.py @@ -4,7 +4,7 @@ import os import traceback from flash_attn import flash_attn_func -from fastvideo.attention.backends.sageattn.quantization.bench.bench_utils import bench +from modified_sageattn.quantization.bench.bench_utils import bench def calculate_attention_flops(batch_size, num_heads, seq_len_q, seq_len_k, head_dim, is_causal=False): @@ -228,4 +228,3 @@ if __name__ == "__main__": print("Example: python benchmark_flashattn2.py --batch-size 1 --num-heads 16 --seq-len 4096 --head-dim 128") sys.stdout.flush() run_benchmark_suite() - diff --git a/fastvideo/attention/backends/benchmark_scripts/benchmark_sageattn3.py b/fastvideo/attention/backends/benchmark_scripts/benchmark_sageattn3.py index 776c6e695..1649da084 100644 --- a/fastvideo/attention/backends/benchmark_scripts/benchmark_sageattn3.py +++ b/fastvideo/attention/backends/benchmark_scripts/benchmark_sageattn3.py @@ -3,8 +3,8 @@ import sys import os import traceback -from fastvideo.attention.backends.sageattn.api import sageattn_blackwell -from fastvideo.attention.backends.sageattn.quantization.bench.bench_utils import bench +from modified_sageattn.api import sageattn_blackwell +from modified_sageattn.quantization.bench.bench_utils import bench def calculate_attention_flops(batch_size, num_heads, seq_len_q, seq_len_k, head_dim, is_causal=False): @@ -247,4 +247,3 @@ if __name__ == "__main__": print("Example: python benchmark_sageattn3.py --batch-size 1 --num-heads 16 --seq-len 4096 --head-dim 128") sys.stdout.flush() run_benchmark_suite() - diff --git a/fastvideo/attention/backends/setup.py b/fastvideo/attention/backends/modified_sageattn_3_setup.py similarity index 94% rename from fastvideo/attention/backends/setup.py rename to fastvideo/attention/backends/modified_sageattn_3_setup.py index 8a88e0427..15d1ee5e4 100644 --- a/fastvideo/attention/backends/setup.py +++ b/fastvideo/attention/backends/modified_sageattn_3_setup.py @@ -11,7 +11,7 @@ from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtensio this_dir = os.path.dirname(os.path.abspath(__file__)) -PACKAGE_NAME = "sageattn3" +PACKAGE_NAME = "modified_sageattn" # FORCE_BUILD: Force a fresh build locally, instead of attempting to find prebuilt wheels # SKIP_CUDA_BUILD: Intended to allow CI to use a simple `python setup.py sdist` run to copy over raw files, without any cuda compilation @@ -68,7 +68,10 @@ if not SKIP_CUDA_BUILD: # https://github.com/pytorch/pytorch/blob/8472c24e3b5b60150096486616d98b7bea01500b/torch/utils/cpp_extension.py#L920 if FORCE_CXX11_ABI: torch._C._GLIBCXX_USE_CXX11_ABI = True - repo_dir = Path(this_dir) + repo_dir = Path(this_dir).resolve() + modified_sageattn_dir = ( + repo_dir.parents[2] / "fastvideo-kernel" / "modified_sageattn" + ) cutlass_dir = repo_dir / "csrc" / "cutlass" (repo_dir / "csrc").mkdir(parents=True, exist_ok=True) if not cutlass_dir.exists(): @@ -100,7 +103,7 @@ if not SKIP_CUDA_BUILD: "-DDQINRMEM", ] include_dirs = [ - repo_dir / "sageattn", + modified_sageattn_dir, cutlass_dir / "include", cutlass_dir / "tools" / "util" / "include", ] @@ -108,7 +111,7 @@ if not SKIP_CUDA_BUILD: ext_modules.append( CUDAExtension( name="fp4attn_cuda", - sources=["sageattn/blackwell/api.cu"], + sources=[str(modified_sageattn_dir / "blackwell/api.cu")], extra_compile_args={ "cxx": ["-O3", "-std=c++17"], "nvcc": append_nvcc_threads( @@ -123,7 +126,7 @@ if not SKIP_CUDA_BUILD: ext_modules.append( CUDAExtension( name="fp4quant_cuda", - sources=["sageattn/quantization/fp4_quantization_4d.cu"], + sources=[str(modified_sageattn_dir / "quantization/fp4_quantization_4d.cu")], extra_compile_args={ "cxx": ["-O3", "-std=c++17"], "nvcc": append_nvcc_threads( diff --git a/fastvideo/layers/quantization/fp8_config.py b/fastvideo/layers/quantization/fp8_config.py deleted file mode 100644 index 25e938f59..000000000 --- a/fastvideo/layers/quantization/fp8_config.py +++ /dev/null @@ -1,109 +0,0 @@ -# # SPDX-License-Identifier: Apache-2.0 -# from fastvideo.layers.quantization.base_config import QuantizationConfig, QuantizeMethodBase -# import torch -# from torch.nn.parameter import Parameter - -# from typing import Any, Tuple -# import deep_gemm -# from deep_gemm import ceil_div, get_mn_major_tma_aligned_tensor -# from deep_gemm.utils import per_token_cast_to_fp8, per_block_cast_to_fp8 -# from fastvideo.models.utils import set_weight_attrs - -# class FP8QuantizeMethod(QuantizeMethodBase): -# def __init__(self): -# super().__init__() -# self.weight_fp8 = None -# self.weight_scale = None - -# def create_weights(self, layer: torch.nn.Module, -# input_size_per_partition: int, -# output_partition_sizes: list[int], input_size: int, -# output_size: int, params_dtype: torch.dtype, -# **extra_weight_attrs): -# """Create weights for a linear layer. Note the corrected signature to match LinearMethodBase.""" -# weight = Parameter(torch.empty( -# sum(output_partition_sizes), -# input_size_per_partition, -# dtype=params_dtype, -# ), -# requires_grad=False) -# set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0}) -# layer.register_parameter("weight", weight) -# set_weight_attrs(weight, extra_weight_attrs) - -# @torch.compile -# def apply(self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor: -# """Apply FP8 quantized computation.""" -# out_dim = layer.weight.shape[0] -# # Need contiguous tensors for collectives. -# assert x.dtype == torch.bfloat16 or x.dtype == torch.float16 or x.dtype == torch.float32, f"only allow bf16/fp16/fp32 inputs to fp8 linear, got {x.dtype}" -# x_fp8, x_scale = per_token_cast_to_fp8(x.view(-1, x.shape[-1]), use_ue8m0=False) -# x_scale = get_mn_major_tma_aligned_tensor(x_scale) -# weight_fp8 = layer._fp8_weight -# weight_scale = layer._fp8_weight_scale - -# original_shape = x.shape -# out = torch.zeros((x_fp8.shape[0], out_dim), device=x.device, dtype=x.dtype) -# deep_gemm.fp8_gemm_nt( -# (x_fp8, x_scale), -# (weight_fp8, weight_scale), -# out, -# disable_ue8m0_cast=False -# ) - -# if bias is not None: -# if bias.device != out.device or bias.dtype != out.dtype: -# bias = bias.to(device=out.device, dtype=out.dtype) -# out = out + bias - -# if len(original_shape) == 3: -# out = out.view(original_shape[0], original_shape[1], out_dim) - -# return out - - -# class FP8Config(QuantizationConfig): -# def __init__(self): -# super().__init__() - -# def get_name(self): -# return "fp8" - -# def get_supported_act_dtypes(self): -# return [torch.bfloat16, torch.float16, torch.float32] - -# @classmethod -# def get_min_capability(cls): -# return 90 - -# @staticmethod -# def get_config_filenames(): -# return [] - -# @classmethod -# def from_config(cls, config: dict[str, Any]) -> "FP8Config": -# return cls() - -# def get_quant_method(self, layer: torch.nn.Module, prefix: str): -# from fastvideo.layers.linear import LinearBase -# fp8_layers = ["ffn.fc_in", "ffn.fc_out", "to_q", "to_k", "to_v", "to_out"] -# if isinstance(layer, LinearBase) and any(layer_name in prefix for layer_name in fp8_layers): -# return FP8QuantizeMethod() -# return None - -# @torch.compile -# def convert_model_to_fp8(model: torch.nn.Module): -# from torch.distributed.tensor import DTensor # type: ignore -# for mod in model.modules(): -# qm = getattr(mod, "quant_method", None) -# if isinstance(qm, FP8QuantizeMethod): -# weight = getattr(mod, "weight", None) -# if weight is None: -# continue -# if isinstance(weight, DTensor): # type: ignore -# weight_local = weight.to_local() -# else: -# weight_local = weight -# fp8_w, fp8_s = per_block_cast_to_fp8(weight_local, use_ue8m0=False) -# mod.register_buffer("_fp8_weight", fp8_w, persistent=False) -# mod.register_buffer("_fp8_weight_scale", fp8_s, persistent=False) \ No newline at end of file diff --git a/fastvideo/models/dits/hunyuanvideo15.py b/fastvideo/models/dits/hunyuanvideo15.py index e69de29bb..b8e5562eb 100644 --- a/fastvideo/models/dits/hunyuanvideo15.py +++ b/fastvideo/models/dits/hunyuanvideo15.py @@ -0,0 +1,766 @@ +# Copyright 2025 The Hunyuan Team and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Any, Dict, Optional, List + +import torch +import torch.nn as nn + +from fastvideo.attention import DistributedAttention, LocalAttention +from fastvideo.forward_context import set_forward_context +from fastvideo.distributed.communication_op import ( + sequence_model_parallel_all_gather_with_unpad, + sequence_model_parallel_shard) +from fastvideo.configs.models.dits import HunyuanVideo15Config +from fastvideo.layers.layernorm import (LayerNormScaleShift, ScaleResidual, + ScaleResidualLayerNormScaleShift) +from fastvideo.layers.linear import ReplicatedLinear +# TODO(will-PY-refactor): RMSNorm .... +from fastvideo.layers.mlp import MLP +from fastvideo.layers.rotary_embedding import get_rotary_pos_embed +from fastvideo.layers.visual_embedding import (ModulateProjection, PatchEmbed, + TimestepEmbedder, unpatchify) +from fastvideo.models.dits.base import BaseDiT +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.logger import init_logger + +from fastvideo.distributed.parallel_state import get_sp_world_size + +logger = init_logger(__name__) + +class HunyuanRMSNorm(nn.Module): + + def __init__( + self, + dim: int, + elementwise_affine=True, + eps: float = 1e-6, + device=None, + dtype=None, + ): + """ + Initialize the RMSNorm normalization layer. + + Args: + dim (int): The dimension of the input tensor. + eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6. + + Attributes: + eps (float): A small value added to the denominator for numerical stability. + weight (nn.Parameter): Learnable scaling parameter. + + """ + factory_kwargs = {"device": device, "dtype": dtype} + super().__init__() + self.eps = eps + if elementwise_affine: + self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs)) + + def _norm(self, x) -> torch.Tensor: + """ + Apply the RMSNorm normalization to the input tensor. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The normalized tensor. + + """ + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + """ + Forward pass through the RMSNorm layer. + + Args: + x (torch.Tensor): The input tensor. + + Returns: + torch.Tensor: The output tensor after applying RMSNorm. + + """ + output = self._norm(x.float()).type_as(x) + if hasattr(self, "weight"): + output = output * self.weight + return output + + +class HunyuanVideo15TimeEmbedding(nn.Module): + r""" + Time embedding for HunyuanVideo 1.5. + + Supports standard timestep embedding and optional reference timestep embedding for MeanFlow-based super-resolution + models. + + Args: + embedding_dim (`int`): + The dimension of the output embedding. + """ + + def __init__(self, embedding_dim: int, use_meanflow: bool = False): + super().__init__() + + self.timestep_embedder = TimestepEmbedder(hidden_size=embedding_dim) + + self.use_meanflow = use_meanflow + self.time_proj_r = None + self.timestep_embedder_r = None + if use_meanflow: + self.timestep_embedder_r = TimestepEmbedder(hidden_size=embedding_dim) + + def forward( + self, + timestep: torch.Tensor, + timestep_r: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + timesteps_emb = self.timestep_embedder(timestep) + + if timestep_r is not None: + timesteps_emb_r = self.timestep_embedder_r(timestep_r) + timesteps_emb = timesteps_emb + timesteps_emb_r + + return timesteps_emb + + +class HunyuanVideo15ByT5TextProjection(nn.Module): + def __init__(self, in_features: int, hidden_size: int, out_features: int): + super().__init__() + self.norm = nn.LayerNorm(in_features) + self.linear_1 = nn.Linear(in_features, hidden_size) + self.linear_2 = nn.Linear(hidden_size, hidden_size) + self.linear_3 = nn.Linear(hidden_size, out_features) + self.act_fn = nn.GELU() + + def forward(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm(encoder_hidden_states) + hidden_states = self.linear_1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_2(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_3(hidden_states) + return hidden_states + + +class HunyuanVideo15ImageProjection(nn.Module): + def __init__(self, in_channels: int, hidden_size: int): + super().__init__() + self.norm_in = nn.LayerNorm(in_channels) + self.linear_1 = nn.Linear(in_channels, in_channels) + self.act_fn = nn.GELU() + self.linear_2 = nn.Linear(in_channels, hidden_size) + self.norm_out = nn.LayerNorm(hidden_size) + + def forward(self, image_embeds: torch.Tensor) -> torch.Tensor: + hidden_states = self.norm_in(image_embeds) + hidden_states = self.linear_1(hidden_states) + hidden_states = self.act_fn(hidden_states) + hidden_states = self.linear_2(hidden_states) + hidden_states = self.norm_out(hidden_states) + return hidden_states + + +class MMDoubleStreamBlock(nn.Module): + """ + A multimodal DiT block with separate modulation for text and image/video, + using distributed attention and linear layers. + """ + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + mlp_ratio: float, + dtype: torch.dtype | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = "", + ): + super().__init__() + + self.deterministic = False + self.num_attention_heads = num_attention_heads + head_dim = hidden_size // num_attention_heads + mlp_hidden_dim = int(hidden_size * mlp_ratio) + + # Image modulation components + self.img_mod = ModulateProjection( + hidden_size, + factor=6, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.img_mod", + ) + + # Fused operations for image stream + self.img_attn_norm = LayerNormScaleShift(hidden_size, + norm_type="layer", + elementwise_affine=False, + dtype=dtype) + self.img_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( + hidden_size, + norm_type="layer", + elementwise_affine=False, + dtype=dtype) + self.img_mlp_residual = ScaleResidual() + + # Image attention components + self.img_attn_qkv = ReplicatedLinear(hidden_size, + hidden_size * 3, + bias=True, + params_dtype=dtype, + prefix=f"{prefix}.img_attn_qkv") + + self.img_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) + self.img_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) + + self.img_attn_proj = ReplicatedLinear(hidden_size, + hidden_size, + bias=True, + params_dtype=dtype, + prefix=f"{prefix}.img_attn_proj") + + self.img_mlp = MLP(hidden_size, + mlp_hidden_dim, + bias=True, + dtype=dtype, + prefix=f"{prefix}.img_mlp") + + # Text modulation components + self.txt_mod = ModulateProjection( + hidden_size, + factor=6, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.txt_mod", + ) + + # Fused operations for text stream + self.txt_attn_norm = LayerNormScaleShift(hidden_size, + norm_type="layer", + elementwise_affine=False, + dtype=dtype) + self.txt_attn_residual_mlp_norm = ScaleResidualLayerNormScaleShift( + hidden_size, + norm_type="layer", + elementwise_affine=False, + dtype=dtype) + self.txt_mlp_residual = ScaleResidual() + + # Text attention components + self.txt_attn_qkv = ReplicatedLinear(hidden_size, + hidden_size * 3, + bias=True, + params_dtype=dtype) + + # QK norm layers for text + self.txt_attn_q_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) + self.txt_attn_k_norm = HunyuanRMSNorm(head_dim, eps=1e-6, dtype=dtype) + + self.txt_attn_proj = ReplicatedLinear(hidden_size, + hidden_size, + bias=True, + params_dtype=dtype) + + self.txt_mlp = MLP(hidden_size, mlp_hidden_dim, bias=True, dtype=dtype) + + # Distributed attention + self.attn = DistributedAttention( + num_heads=num_attention_heads, + head_size=head_dim, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn") + + def forward( + self, + img: torch.Tensor, + txt: torch.Tensor, + vec: torch.Tensor, + freqs_cis: tuple, + original_seq_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Process modulation vectors + img_mod_outputs = self.img_mod(vec) + ( + img_attn_shift, + img_attn_scale, + img_attn_gate, + img_mlp_shift, + img_mlp_scale, + img_mlp_gate, + ) = torch.chunk(img_mod_outputs, 6, dim=-1) + + txt_mod_outputs = self.txt_mod(vec) + ( + txt_attn_shift, + txt_attn_scale, + txt_attn_gate, + txt_mlp_shift, + txt_mlp_scale, + txt_mlp_gate, + ) = torch.chunk(txt_mod_outputs, 6, dim=-1) + + # Prepare image for attention using fused operation + img_attn_input = self.img_attn_norm(img, img_attn_shift, img_attn_scale) + # Get QKV for image + img_qkv, _ = self.img_attn_qkv(img_attn_input) + batch_size, image_seq_len = img_qkv.shape[0], img_qkv.shape[1] + + # Split QKV + img_qkv = img_qkv.view(batch_size, image_seq_len, 3, + self.num_attention_heads, -1) + img_q, img_k, img_v = img_qkv[:, :, 0], img_qkv[:, :, 1], img_qkv[:, :, + 2] + + # Apply QK-Norm if needed + img_q = self.img_attn_q_norm(img_q).to(img_v) + img_k = self.img_attn_k_norm(img_k).to(img_v) + + # Prepare text for attention using fused operation + txt_attn_input = self.txt_attn_norm(txt, txt_attn_shift, txt_attn_scale) + + # Get QKV for text + txt_qkv, _ = self.txt_attn_qkv(txt_attn_input) + batch_size, text_seq_len = txt_qkv.shape[0], txt_qkv.shape[1] + + # Split QKV + txt_qkv = txt_qkv.view(batch_size, text_seq_len, 3, + self.num_attention_heads, -1) + txt_q, txt_k, txt_v = txt_qkv[:, :, 0], txt_qkv[:, :, 1], txt_qkv[:, :, + 2] + # Apply QK-Norm if needed + txt_q = self.txt_attn_q_norm(txt_q).to(txt_q.dtype) + txt_k = self.txt_attn_k_norm(txt_k).to(txt_k.dtype) + + + img_attn, txt_attn = self.attn(img_q, img_k, img_v, original_seq_len, txt_q, txt_k, txt_v, freqs_cis=freqs_cis) + + img_attn_out, _ = self.img_attn_proj( + img_attn.view(batch_size, image_seq_len, -1)) + # Use fused operation for residual connection, normalization, and modulation + img_mlp_input, img_residual = self.img_attn_residual_mlp_norm( + img, img_attn_out, img_attn_gate, img_mlp_shift, img_mlp_scale) + + # Process image MLP + img_mlp_out = self.img_mlp(img_mlp_input) + img = self.img_mlp_residual(img_residual, img_mlp_out, img_mlp_gate) + + # Process text attention output + txt_attn_out, _ = self.txt_attn_proj( + txt_attn.reshape(batch_size, text_seq_len, -1)) + + # Use fused operation for residual connection, normalization, and modulation + txt_mlp_input, txt_residual = self.txt_attn_residual_mlp_norm( + txt, txt_attn_out, txt_attn_gate, txt_mlp_shift, txt_mlp_scale) + + # Process text MLP + txt_mlp_out = self.txt_mlp(txt_mlp_input) + txt = self.txt_mlp_residual(txt_residual, txt_mlp_out, txt_mlp_gate) + + return img, txt + + +class HunyuanVideo15Transformer3DModel(BaseDiT): + r""" + A Transformer model for video-like data used in [HunyuanVideo1.5](https://huggingface.co/tencent/HunyuanVideo1.5). + """ + + # shard single stream, double stream blocks, and refiner_blocks + _fsdp_shard_conditions = HunyuanVideo15Config()._fsdp_shard_conditions + _compile_conditions = HunyuanVideo15Config()._compile_conditions + _supported_attention_backends = HunyuanVideo15Config( + )._supported_attention_backends + param_names_mapping = HunyuanVideo15Config().param_names_mapping + reverse_param_names_mapping = HunyuanVideo15Config( + ).reverse_param_names_mapping + lora_param_names_mapping = HunyuanVideo15Config().lora_param_names_mapping + + def __init__( + self, + config: HunyuanVideo15Config, + hf_config: dict[str, Any], + ) -> None: + super().__init__(config=config, hf_config=hf_config) + + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_channels_latents = config.num_channels_latents + self.out_channels = config.out_channels or config.in_channels + self.patch_size = (config.patch_size_t, config.patch_size, config.patch_size) + + # 1. Latent and condition embedders + self.img_in = PatchEmbed(self.patch_size, + config.in_channels, + self.hidden_size, + prefix=f"{config.prefix}.img_in") + self.image_embedder = HunyuanVideo15ImageProjection(config.image_embed_dim, self.hidden_size) + + self.txt_in = SingleTokenRefiner(config.text_embed_dim, + self.hidden_size, + config.num_attention_heads, + depth=config.num_refiner_layers, + dtype=None, + prefix=f"{config.prefix}.txt_in") + + self.txt_in_2 = HunyuanVideo15ByT5TextProjection(config.text_embed_2_dim, 2048, self.hidden_size) + + self.time_in = HunyuanVideo15TimeEmbedding(self.hidden_size, use_meanflow=config.use_meanflow) + + self.cond_type_embed = nn.Embedding(3, self.hidden_size) + + # 3. Dual stream transformer blocks + + self.double_blocks = nn.ModuleList( + [ + MMDoubleStreamBlock( + hidden_size=self.hidden_size, + num_attention_heads=config.num_attention_heads, + mlp_ratio=config.mlp_ratio, + dtype=None, + supported_attention_backends=self._supported_attention_backends, + prefix=f"{config.prefix}.double_blocks.{i}" + ) + for i in range(config.num_layers) + ] + ) + + # 5. Output projection + self.final_layer = FinalLayer(self.hidden_size, + self.patch_size, + self.out_channels, + prefix=f"{config.prefix}.final_layer") + + self.gradient_checkpointing = False + + self.__post_init__() + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: List[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: List[torch.Tensor], + guidance: Optional[torch.Tensor] = None, + timestep_r: Optional[torch.LongTensor] = None, + attention_kwargs: Optional[Dict[str, Any]] = None, + ): + encoder_hidden_states_image = encoder_hidden_states_image[0] + encoder_hidden_states, encoder_hidden_states_2 = encoder_hidden_states + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.config.patch_size_t, self.config.patch_size, self.config.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + + # 1. RoPE + # Get rotary embeddings + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames, post_patch_height, post_patch_width), self.hidden_size, + self.num_attention_heads, self.config.rope_axes_dim, self.config.rope_theta) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + # 2. Conditional embeddings + temb = self.time_in(timestep, timestep_r=timestep_r) + + hidden_states = self.img_in(hidden_states) + hidden_states, original_seq_len = sequence_model_parallel_shard(hidden_states, dim=1) + + current_seq_len = hidden_states.shape[1] + sp_world_size = get_sp_world_size() + padded_seq_len = current_seq_len * sp_world_size + + + + encoder_hidden_states = self.txt_in(encoder_hidden_states, timestep) + + encoder_hidden_states_cond_emb = self.cond_type_embed( + torch.zeros_like(encoder_hidden_states[:, :, 0], dtype=torch.long) + ) + encoder_hidden_states = encoder_hidden_states + encoder_hidden_states_cond_emb + + # byt5 text embedding + encoder_hidden_states_2 = self.txt_in_2(encoder_hidden_states_2) + + encoder_hidden_states_2_cond_emb = self.cond_type_embed( + torch.ones_like(encoder_hidden_states_2[:, :, 0], dtype=torch.long) + ) + encoder_hidden_states_2 = encoder_hidden_states_2 + encoder_hidden_states_2_cond_emb + + # image embed + encoder_hidden_states_3 = self.image_embedder(encoder_hidden_states_image) + is_t2v = torch.all(encoder_hidden_states_image == 0) + encoder_hidden_states_3_cond_emb = self.cond_type_embed( + 2 + * torch.ones_like( + encoder_hidden_states_3[:, :, 0], + dtype=torch.long, + ) + ) + encoder_hidden_states_3 = encoder_hidden_states_3 + encoder_hidden_states_3_cond_emb + + encoder_hidden_states = torch.cat([encoder_hidden_states_2, encoder_hidden_states], dim=1) if is_t2v else torch.cat([encoder_hidden_states_3, encoder_hidden_states_2, encoder_hidden_states], dim=1) + + # 4. Transformer blocks + if torch.is_grad_enabled() and self.gradient_checkpointing: + for block in self.double_blocks: + hidden_states, encoder_hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + temb, + freqs_cis, + original_seq_len + ) + + else: + for block in self.double_blocks: + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + freqs_cis, + original_seq_len + ) + + # Final layer processing + if get_sp_world_size() > 1: + hidden_states = hidden_states.contiguous() + hidden_states = sequence_model_parallel_all_gather_with_unpad(hidden_states, original_seq_len, dim=1) + hidden_states = self.final_layer(hidden_states, temb) + # Unpatchify to get original shape + hidden_states = unpatchify(hidden_states, post_patch_num_frames, post_patch_height, post_patch_width, self.patch_size, self.out_channels) + + return hidden_states + +class SingleTokenRefiner(nn.Module): + """ + A token refiner that processes text embeddings with attention to improve + their representation for cross-attention with image features. + """ + + def __init__( + self, + in_channels, + hidden_size, + num_attention_heads, + depth=2, + qkv_bias=True, + dtype=None, + prefix: str = "", + ) -> None: + super().__init__() + + # Input projection + # self.input_embedder = ReplicatedLinear( + # in_channels, + # hidden_size, + # bias=True, + # params_dtype=dtype, + # prefix=f"{prefix}.input_embedder") + self.input_embedder = nn.Linear(in_channels, hidden_size, bias=True) + + # Timestep embedding + self.t_embedder = TimestepEmbedder(hidden_size, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.t_embedder") + + # Context embedding + self.c_embedder = MLP(in_channels, + hidden_size, + hidden_size, + act_type="silu", + dtype=dtype, + prefix=f"{prefix}.c_embedder") + + # Refiner blocks + self.refiner_blocks = nn.ModuleList([ + IndividualTokenRefinerBlock( + hidden_size, + num_attention_heads, + qkv_bias=qkv_bias, + dtype=dtype, + prefix=f"{prefix}.refiner_blocks.{i}", + ) for i in range(depth) + ]) + + def forward(self, x, t, mask=None): + # Get timestep embeddings + timestep_aware_representations = self.t_embedder(t) + + original_dtype = x.dtype + if mask is None: + context_aware_representations = x.mean(dim=1) + else: + mask_float = mask.float().unsqueeze(-1) # [B, L, 1] + context_aware_representations = (x * mask_float).sum(dim=1) / mask_float.sum(dim=1) + + context_aware_representations = self.c_embedder( + context_aware_representations) + c = timestep_aware_representations + context_aware_representations + # Project input + x = self.input_embedder(x) + # Process through refiner blocks + for block in self.refiner_blocks: + x = block(x, c, mask=mask) + return x + +class IndividualTokenRefinerBlock(nn.Module): + """ + A transformer block for refining individual tokens with self-attention. + """ + + def __init__( + self, + hidden_size, + num_attention_heads, + mlp_ratio=4.0, + qkv_bias=True, + dtype=None, + prefix: str = "", + ) -> None: + super().__init__() + self.num_attention_heads = num_attention_heads + mlp_hidden_dim = int(hidden_size * mlp_ratio) + + # Normalization and attention + self.norm1 = nn.LayerNorm(hidden_size, + eps=1e-6, + elementwise_affine=True, + dtype=dtype) + + self.self_attn_qkv = ReplicatedLinear(hidden_size, + hidden_size * 3, + bias=qkv_bias, + params_dtype=dtype, + prefix=f"{prefix}.self_attn_qkv") + + self.self_attn_proj = ReplicatedLinear( + hidden_size, + hidden_size, + bias=qkv_bias, + params_dtype=dtype, + prefix=f"{prefix}.self_attn_proj") + + # MLP + self.norm2 = nn.LayerNorm(hidden_size, + eps=1e-6, + elementwise_affine=True, + dtype=dtype) + self.mlp = MLP(hidden_size, + mlp_hidden_dim, + bias=True, + act_type="silu", + dtype=dtype, + prefix=f"{prefix}.mlp") + + # Modulation + self.adaLN_modulation = ModulateProjection( + hidden_size, + factor=2, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.adaLN_modulation") + + # Scaled dot product attention + self.attn = LocalAttention( + num_heads=num_attention_heads, + head_size=hidden_size // num_attention_heads, + # TODO: remove hardcode + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA), + ) + + def forward(self, x, c, mask=None): + + # Get modulation parameters + gate_msa, gate_mlp = self.adaLN_modulation(c).chunk(2, dim=-1) + # Self-attention + norm_x = self.norm1(x) + qkv, _ = self.self_attn_qkv(norm_x) + + batch_size, seq_len = qkv.shape[0], qkv.shape[1] + qkv = qkv.view(batch_size, seq_len, 3, self.num_attention_heads, -1) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + if mask is None: + attn_output = self.attn(q, k, v) + else: + from fastvideo.attention.backends.flash_attn import FlashAttnMetadataBuilder + attn_metadata = FlashAttnMetadataBuilder().build( + current_timestep=0, + attn_mask=mask, + ) + # Run distributed attention + with set_forward_context(current_timestep=0, attn_metadata=attn_metadata): + attn_output = self.attn(q, k, v) # [B, L, H, D] + attn_output = attn_output.reshape(batch_size, seq_len, + -1) + + # Project and apply residual connection with gating + attn_out, _ = self.self_attn_proj(attn_output) + x = x + attn_out * gate_msa.unsqueeze(1) + + # MLP + mlp_out = self.mlp(self.norm2(x)) + x = x + mlp_out * gate_mlp.unsqueeze(1) + + return x + +class FinalLayer(nn.Module): + """ + The final layer of DiT that projects features to pixel space. + """ + + def __init__(self, + hidden_size, + patch_size, + out_channels, + dtype=None, + prefix: str = "") -> None: + super().__init__() + + # Normalization + self.norm_final = nn.LayerNorm(hidden_size, + eps=1e-6, + elementwise_affine=False, + dtype=dtype) + + output_dim = patch_size[0] * patch_size[1] * patch_size[2] * out_channels + + self.linear = ReplicatedLinear(hidden_size, + output_dim, + bias=True, + params_dtype=dtype, + prefix=f"{prefix}.linear") + + # Modulation + self.adaLN_modulation = ModulateProjection( + hidden_size, + factor=2, + act_layer="silu", + dtype=dtype, + prefix=f"{prefix}.adaLN_modulation") + + def forward(self, x, c): + # What the heck HF? Why you change the scale and shift order here??? + scale, shift = self.adaLN_modulation(c).chunk(2, dim=-1) + x = self.norm_final(x) * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1) + x, _ = self.linear(x) + return x + +# Entry point for model registry +EntryClass = HunyuanVideo15Transformer3DModel diff --git a/fastvideo/models/schedulers/scheduling_rcm.py b/fastvideo/models/schedulers/scheduling_rcm.py index e69de29bb..71cc714b0 100644 --- a/fastvideo/models/schedulers/scheduling_rcm.py +++ b/fastvideo/models/schedulers/scheduling_rcm.py @@ -0,0 +1,326 @@ +# SPDX-License-Identifier: Apache-2.0 +# rCM (recurrent Consistency Model) Scheduler for TurboDiffusion support +# +# This scheduler implements the rCM sampling method from TurboDiffusion, +# enabling 1-4 step video generation with distilled checkpoints. +# +# Reference: +# TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times +# https://arxiv.org/pdf/2512.16093 + +import math +from dataclasses import dataclass +from typing import Any + +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput + +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.base import BaseScheduler + +logger = init_logger(__name__) + + +@dataclass +class RCMSchedulerOutput(BaseOutput): + """ + Output class for the RCM scheduler's `step` function output. + + Args: + prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, ...)`): + Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used + as next model input in the denoising loop. + """ + + prev_sample: torch.FloatTensor + + +class RCMScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): + """ + rCM (recurrent Consistency Model) scheduler for TurboDiffusion. + + This scheduler implements the rCM sampling method which enables 1-4 step + video generation using distilled checkpoints. It uses: + 1. TrigFlow → RectifiedFlow timestep conversion + 2. SDE sampling formula: x = (1 - t_next) * (x - t_cur * v_pred) + t_next * noise + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps used to train the model. + sigma_max (`float`, defaults to 80.0): + The initial sigma value for rCM sampling. Controls the noise level + at the start of sampling. + mid_timesteps (`list[float]`, *optional*): + Custom intermediate timesteps. If None, uses optimized defaults + [1.5, 1.4, 1.0] for better visual quality. + """ + + _compatibles: list[Any] = [] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + sigma_max: float = 80.0, + mid_timesteps: list[float] | None = None, + ): + # Default mid timesteps optimized for visual quality + if mid_timesteps is None: + mid_timesteps = [1.5, 1.4, 1.0] + self._mid_timesteps = mid_timesteps + + self.num_train_timesteps = num_train_timesteps + self.sigma_max = sigma_max + + # Initialize with default timesteps (will be set properly via set_timesteps) + self.timesteps = torch.tensor([1.0, 0.0], dtype=torch.float64) + self.sigmas = self.timesteps.clone() + + self._step_index: int | None = None + self._begin_index: int | None = None + + BaseScheduler.__init__(self) + + @property + def step_index(self) -> int | None: + """ + The index counter for current timestep. Increases by 1 after each scheduler step. + """ + return self._step_index + + @property + def begin_index(self) -> int | None: + """ + The index for the first timestep. + """ + return self._begin_index + + @property + def init_noise_sigma(self) -> float: + """ + Initial noise sigma for scaling latents. + + In rCM, initial noise is scaled by the first sigma value: + x_0 = noise * sigmas[0] + + This property is used by LatentPreparationStage to scale initial latents. + """ + return float(self.sigmas[0]) + + def set_begin_index(self, begin_index: int = 0) -> None: + """ + Sets the begin index for the scheduler. + + Args: + begin_index (`int`): + The begin index for the scheduler. + """ + self._begin_index = begin_index + + def set_shift(self, shift: float) -> None: + """ + rCM doesn't use shift parameter, but required by BaseScheduler. + """ + pass + + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device | None = None, + sigma_max: float | None = None, + ) -> None: + """ + Sets the discrete timesteps used for the rCM sampling process. + + The timesteps are computed using TrigFlow → RectifiedFlow conversion: + 1. Start with atan(sigma_max) and intermediate values + 2. Convert via: t = sin(t) / (cos(t) + sin(t)) + + Args: + num_inference_steps (`int`): + The number of diffusion steps (1-4 for rCM). + device (`str` or `torch.device`, *optional*): + The device to move timesteps to. + sigma_max (`float`, *optional*): + Override the initial sigma value. + """ + if num_inference_steps < 1 or num_inference_steps > 4: + logger.warning( + "rCM is optimized for 1-4 steps, got %d steps. " + "Performance may be suboptimal.", num_inference_steps + ) + + self.num_inference_steps = num_inference_steps + + if sigma_max is not None: + self.sigma_max = sigma_max + + # Build timestep schedule + mid_t = self._mid_timesteps[:num_inference_steps - 1] + + # TrigFlow timesteps: [atan(sigma_max), mid_t..., 0] + t_steps = torch.tensor( + [math.atan(self.sigma_max), *mid_t, 0], + dtype=torch.float64, + device=device, + ) + + # Convert TrigFlow → RectifiedFlow: t = sin(t) / (cos(t) + sin(t)) + t_steps = torch.sin(t_steps) / (torch.cos(t_steps) + torch.sin(t_steps)) + + # Store raw sigmas for use in step() formula + self.sigmas = t_steps.clone() + + # Scale timesteps by 1000 for model input (as per TurboDiffusion) + self.timesteps = t_steps * 1000 + + self._step_index = None + self._begin_index = None + + logger.debug("rCM timesteps (scaled): %s", self.timesteps.tolist()) + logger.debug("rCM sigmas (raw): %s", self.sigmas.tolist()) + + def _init_step_index(self, timestep: torch.FloatTensor | None = None) -> None: + """Initialize step index at the beginning of sampling.""" + if self._begin_index is None: + self._step_index = 0 + else: + self._step_index = self._begin_index + + def scale_model_input( + self, + sample: torch.Tensor, + timestep: int | None = None, + ) -> torch.Tensor: + """ + rCM doesn't scale model input, returns sample as-is. + """ + return sample + + def scale_noise( + self, + sample: torch.FloatTensor, + timestep: torch.FloatTensor | None = None, + noise: torch.FloatTensor | None = None, + ) -> torch.FloatTensor: + """ + Scale initial noise for rCM sampling. + + In rCM, initial noise is scaled by the first timestep (raw sigma): + x_0 = noise * t_steps[0] + + Args: + sample: Not used (for API compatibility) + timestep: Not used (for API compatibility) + noise: The noise tensor to scale + + Returns: + Scaled noise tensor ready for sampling + """ + if noise is None: + raise ValueError("noise must be provided for rCM scale_noise") + + # Use raw sigma (not scaled timestep) for initial noise scaling + t_initial = self.sigmas[0] + return noise.to(torch.float64) * t_initial + + def step( + self, + model_output: torch.FloatTensor, + timestep: int | torch.Tensor, + sample: torch.FloatTensor, + generator: torch.Generator | None = None, + return_dict: bool = True, + ) -> RCMSchedulerOutput | tuple[torch.FloatTensor, ...]: + """ + Predict the sample from the previous timestep using rCM update rule. + + The rCM update formula is: + x_{t+1} = (1 - t_next) * (x_t - t_cur * v_pred) + t_next * noise + + Args: + model_output (`torch.FloatTensor`): + The velocity prediction from the model (v_pred). + timestep (`int` or `torch.Tensor`): + The current timestep index (not the actual timestep value). + For rCM, this should be the index into self.timesteps. + sample (`torch.FloatTensor`): + Current sample x_t. + generator (`torch.Generator`, *optional*): + Random number generator for noise. + return_dict (`bool`): + Whether to return RCMSchedulerOutput or tuple. + + Returns: + `RCMSchedulerOutput` or `tuple`: + The denoised sample for the next step. + """ + if self._step_index is None: + self._init_step_index() + + assert self._step_index is not None + + # Get current and next sigma values (raw, unscaled) for rCM formula + # Note: self.timesteps is scaled by 1000 for model input, + # but we need raw values for the step formula + t_cur = self.sigmas[self._step_index] + # On the final step, t_next should be 0 (fully denoised) + if self._step_index + 1 < len(self.sigmas): + t_next = self.sigmas[self._step_index + 1] + else: + t_next = torch.tensor(0.0, device=sample.device, dtype=torch.float64) + + # Ensure we're working in float64 for precision + sample = sample.to(torch.float64) + model_output = model_output.to(torch.float64) + + # rCM update: x = (1 - t_next) * (x - t_cur * v_pred) + t_next * noise + x_denoised = sample - t_cur * model_output + + # Generate noise for SDE sampling + if isinstance(generator, list): + generator = generator[0] + noise = torch.randn( + sample.shape, + dtype=torch.float32, + device="cpu", + generator=generator, + ).to(sample.device).to(torch.float64) + + prev_sample = (1 - t_next) * x_denoised + t_next * noise + + # Increment step counter + self._step_index += 1 + + # Cast back to model output dtype + prev_sample = prev_sample.to(model_output.dtype) + + if not return_dict: + return (prev_sample,) + + return RCMSchedulerOutput(prev_sample=prev_sample) + + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + """ + Add noise to samples (forward diffusion process). + + Not typically used for rCM inference, but provided for API compatibility. + """ + raise NotImplementedError( + "add_noise is not implemented for RCMScheduler. " + "Use scale_noise for initializing the sampling process." + ) + + def __len__(self) -> int: + return self.config.num_train_timesteps + +# Entry point for model registry +EntryClass = RCMScheduler diff --git a/fastvideo/pipelines/stages/conditioning.py b/fastvideo/pipelines/stages/conditioning.py index bbc639ec8..80fb3b6cc 100644 --- a/fastvideo/pipelines/stages/conditioning.py +++ b/fastvideo/pipelines/stages/conditioning.py @@ -86,5 +86,9 @@ class ConditioningStage(PipelineStage): def verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify conditioning stage outputs.""" result = VerificationResult() - result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + if batch.prompt_embeds is None or not batch.prompt_embeds: + batch.do_classifier_free_guidance = False + return result + if batch.do_classifier_free_guidance or batch.prompt_embeds: + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) return result diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py index e69de29bb..7e1496857 100644 --- a/fastvideo/pipelines/stages/matrixgame_denoising.py +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -0,0 +1,739 @@ +from __future__ import annotations +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch # type: ignore + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.denoising import DenoisingStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +try: + from fastvideo.attention.backends.video_sparse_attn import (VideoSparseAttentionBackend) + vsa_available = True +except ImportError: + vsa_available = False + VideoSparseAttentionBackend = None # type: ignore + +logger = init_logger(__name__) + + +@dataclass +class BlockProcessingContext: + """Dataclass contains for block processing.""" + + batch: ForwardBatch + + block_idx: int + start_index: int + + kv_cache1: list[dict[Any, Any]] + kv_cache2: list[dict[Any, Any]] | None + kv_cache_mouse: list[dict[Any, Any]] | None + kv_cache_keyboard: list[dict[Any, Any]] | None + crossattn_cache: list[dict[Any, Any]] + + timesteps: torch.Tensor + block_sizes: list[int] + noise_pool: list[torch.Tensor] | None + + fastvideo_args: FastVideoArgs + target_dtype: torch.dtype + autocast_enabled: bool + boundary_timestep: float | None + high_noise_timesteps: torch.Tensor | None + context_noise: float + + image_kwargs: dict[str, Any] + pos_cond_kwargs: dict[str, Any] + + def get_kv_cache(self, timestep_val: float) -> list[dict[Any, Any]]: + if self.boundary_timestep is not None: + if timestep_val >= self.boundary_timestep: + return self.kv_cache1 + else: + assert self.kv_cache2 is not None, "kv_cache2 is not initialized" + return self.kv_cache2 + return self.kv_cache1 + + +class MatrixGameCausalDenoisingStage(DenoisingStage): + + def __init__(self, transformer, scheduler, pipeline=None, transformer_2=None, vae=None) -> None: + super().__init__(transformer, scheduler, pipeline, transformer_2, vae) + self.transformer = transformer + self.transformer_2 = transformer_2 + self.vae = vae + self.num_transformer_blocks = len(self.transformer.blocks) + + if hasattr(self.transformer, 'config') and hasattr(self.transformer.config, 'arch_config'): + self.num_frame_per_block = getattr(self.transformer.config.arch_config, 'num_frames_per_block', + getattr(self.transformer, 'num_frame_per_block', 1)) + self.sliding_window_num_frames = getattr(self.transformer.config.arch_config, 'sliding_window_num_frames', + 15) + else: + self.num_frame_per_block = getattr(self.transformer, 'num_frame_per_block', 1) + self.sliding_window_num_frames = 15 + + try: + self.local_attn_size = getattr(self.transformer, "local_attn_size", -1) + except Exception: + self.local_attn_size = -1 + + assert self.local_attn_size != -1, (f"local_attn_size must be set for Matrix-Game causal inference, " + f"got {self.local_attn_size}. Check MatrixGameWanVideoArchConfig.") + assert self.num_frame_per_block > 0, (f"num_frame_per_block must be positive, got {self.num_frame_per_block}") + + logger.info("MatrixGame causal inference initialized: " + "local_attn_size=%s, num_frame_per_block=%s", self.local_attn_size, self.num_frame_per_block) + + self.action_config = getattr(self.transformer, 'action_config', {}) + self.use_action_module = len(self.action_config) > 0 + + self._streaming_initialized: bool = False + self._streaming_ctx: BlockProcessingContext | None = None + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast + + latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] + patch_size = self.transformer.patch_size + patch_ratio = patch_size[-1] * patch_size[-2] + self.frame_seq_length = latent_seq_length // patch_ratio + + timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu() + if fastvideo_args.pipeline_config.warp_denoising_step: + scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) + timesteps = scheduler_timesteps[1000 - timesteps] + timesteps = timesteps.to(get_local_torch_device()) + + boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, 'boundary_ratio', None) + if boundary_ratio is not None: + boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps + high_noise_timesteps = timesteps[timesteps >= boundary_timestep] + else: + boundary_timestep = None + high_noise_timesteps = None + + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert torch.isnan(image_embeds[0]).sum() == 0 + image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds] + + # directly set the kwarg. + image_kwargs = {"encoder_hidden_states_image": image_embeds} + pos_cond_kwargs: dict[str, Any] = {} + + assert batch.latents is not None, "latents must be provided" + latents = batch.latents + b, c, t, h, w = latents.shape + prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + + kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) + kv_cache2 = None + if boundary_timestep is not None: + kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + kv_cache_mouse = None + kv_cache_keyboard = None + if self.use_action_module: + kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + crossattn_cache = self._initialize_crossattn_cache( + batch_size=latents.shape[0], + max_text_len=257, # 1 CLS + 256 patch tokens + dtype=target_dtype, + device=latents.device) + + if t % self.num_frame_per_block != 0: + raise ValueError("num_frames must be divisible by num_frame_per_block for causal denoising") + num_blocks = t // self.num_frame_per_block + block_sizes = [self.num_frame_per_block] * num_blocks + start_index = 0 + + if boundary_timestep is not None: + block_sizes[0] = 1 + + # NOTE: MatrixGame does NOT process the first frame separately. + # The first frame information is already encoded in batch.image_latent (cond_concat) + # and will be used by the model via channel concatenation: torch.cat([x, cond_concat], dim=1) + + ctx = BlockProcessingContext( + batch=batch, + block_idx=0, + start_index=0, + kv_cache1=kv_cache1, + kv_cache2=kv_cache2, + kv_cache_mouse=kv_cache_mouse, + kv_cache_keyboard=kv_cache_keyboard, + crossattn_cache=crossattn_cache, + timesteps=timesteps, + block_sizes=block_sizes, + noise_pool=None, + fastvideo_args=fastvideo_args, + target_dtype=target_dtype, + autocast_enabled=autocast_enabled, + boundary_timestep=boundary_timestep, + high_noise_timesteps=high_noise_timesteps, + context_noise=getattr(fastvideo_args.pipeline_config, "context_noise", 0), + image_kwargs=image_kwargs, + pos_cond_kwargs=pos_cond_kwargs, + ) + + context_noise = getattr(fastvideo_args.pipeline_config, "context_noise", 0) + + with self.progress_bar(total=len(block_sizes) * len(timesteps)) as progress_bar: + for block_idx, current_num_frames in enumerate(block_sizes): + ctx.block_idx = block_idx + ctx.start_index = start_index + current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] + + action_kwargs = self._prepare_action_kwargs(batch, start_index, current_num_frames) + + current_latents = self._process_single_block( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + timesteps=timesteps, + ctx=ctx, + action_kwargs=action_kwargs, + progress_bar=progress_bar, + ) + + latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents + + # Update KV caches with clean context + self._update_context_cache( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + ctx=ctx, + action_kwargs=action_kwargs, + context_noise=context_noise, + ) + + start_index += current_num_frames + + if boundary_timestep is not None: + num_frames_to_remove = self.num_frame_per_block - 1 + if num_frames_to_remove > 0: + latents = latents[:, :, :-num_frames_to_remove, :, :] + + batch.latents = latents + return batch + + def _prepare_action_kwargs(self, batch: ForwardBatch, start_index: int, num_frames: int) -> dict[str, Any]: + action_kwargs: dict[str, Any] = {} + if not self.use_action_module: + return action_kwargs + + vae_time_compression_ratio = 4 + end_frame_idx = 1 + vae_time_compression_ratio * (start_index + num_frames - 1) + if hasattr(batch, 'mouse_cond') and batch.mouse_cond is not None: + action_kwargs['mouse_cond'] = batch.mouse_cond[:, :end_frame_idx] + if hasattr(batch, 'keyboard_cond') and batch.keyboard_cond is not None: + action_kwargs['keyboard_cond'] = batch.keyboard_cond[:, :end_frame_idx] + + # CRITICAL: Pass num_frame_per_block to model - this should be num_frames (current block size) + action_kwargs['num_frame_per_block'] = num_frames + return action_kwargs + + def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device) -> list[dict]: + kv_cache = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = getattr(self.transformer, 'attention_head_dim', + self.transformer.hidden_size // num_attention_heads) + if self.local_attn_size != -1: + kv_cache_size = self.local_attn_size * self.frame_seq_length + else: + kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames + + for _ in range(self.num_transformer_blocks): + kv_cache.append({ + "k": + torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, kv_cache_size, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + + return kv_cache + + def _initialize_action_kv_cache(self, batch_size: int, dtype: torch.dtype, device: torch.device): + kv_cache_mouse = [] + kv_cache_keyboard = [] + + action_heads = self.action_config.get('heads_num', 16) + mouse_head_dim = self.action_config.get('mouse_hidden_dim', 1024) // action_heads + keyboard_head_dim = self.action_config.get('keyboard_hidden_dim', 1024) // action_heads + + kv_cache_size = self.local_attn_size if self.local_attn_size != -1 else 15 + + for _ in range(self.num_transformer_blocks): + kv_cache_keyboard.append({ + "k": + torch.zeros([batch_size, kv_cache_size, action_heads, keyboard_head_dim], dtype=dtype, device=device), + "v": + torch.zeros([batch_size, kv_cache_size, action_heads, keyboard_head_dim], dtype=dtype, device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + kv_cache_mouse.append({ + "k": + torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, action_heads, mouse_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size * self.frame_seq_length, kv_cache_size, action_heads, mouse_head_dim], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + + return kv_cache_mouse, kv_cache_keyboard + + def _initialize_crossattn_cache(self, batch_size: int, max_text_len: int, dtype: torch.dtype, + device: torch.device) -> list[dict]: + crossattn_cache = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = getattr(self.transformer, 'attention_head_dim', + self.transformer.hidden_size // num_attention_heads) + for _ in range(self.num_transformer_blocks): + crossattn_cache.append({ + "k": + torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "v": + torch.zeros([batch_size, max_text_len, num_attention_heads, attention_head_dim], + dtype=dtype, + device=device), + "is_init": + False, + }) + return crossattn_cache + + def _process_single_block( + self, + current_latents: torch.Tensor, + batch: ForwardBatch, + start_index: int, + current_num_frames: int, + timesteps: torch.Tensor, + ctx: BlockProcessingContext, + action_kwargs: dict[str, Any], + noise_generator: Callable[[tuple, torch.dtype, int], torch.Tensor] + | None = None, + progress_bar: Any | None = None, + ) -> torch.Tensor: + prompt_embeds = batch.prompt_embeds + noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) + + for i, t_cur in enumerate(timesteps): + if ctx.boundary_timestep is not None and t_cur < ctx.boundary_timestep: + current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer + else: + current_model = self.transformer + + noise_latents = noise_latents_btchw.clone() + latent_model_input = current_latents.to(ctx.target_dtype) + + independent_first_frame = getattr(self.transformer, 'independent_first_frame', False) + if batch.image_latent is not None and independent_first_frame and start_index == 0: + latent_model_input = torch.cat([latent_model_input, batch.image_latent.to(ctx.target_dtype)], dim=2) + + # t_expand needs to be [batch * frames] to match flattened pred_noise/noise_latents + t_expand = t_cur.repeat(latent_model_input.shape[0] * current_num_frames) + + # Build attention metadata if VSA is available + if vsa_available and self.attn_backend == VideoSparseAttentionBackend: + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls() + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls() + h, w = current_latents.shape[-2:] + attn_metadata = self.attn_metadata_builder.build( + current_timestep=i, + raw_latent_shape=(current_num_frames, h, w), + patch_size=ctx.fastvideo_args.pipeline_config.dit_config.patch_size, + VSA_sparsity=ctx.fastvideo_args.VSA_sparsity, + device=get_local_torch_device(), + ) + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + else: + attn_metadata = None + + with torch.autocast(device_type="cuda", + dtype=ctx.target_dtype, + enabled=ctx.autocast_enabled), \ + set_forward_context(current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch): + # Expand timestep to per-frame format [batch, num_frames] for causal model + t_expanded_noise = t_cur * torch.ones((latent_model_input.shape[0], current_num_frames), + device=latent_model_input.device, + dtype=torch.long) + + model_kwargs = { + "kv_cache": ctx.get_kv_cache(t_cur), + "crossattn_cache": ctx.crossattn_cache, + "current_start": start_index * self.frame_seq_length, + "start_frame": start_index, + } + + if self.use_action_module and current_model == self.transformer: + model_kwargs.update({ + "kv_cache_mouse": ctx.kv_cache_mouse, + "kv_cache_keyboard": ctx.kv_cache_keyboard, + }) + model_kwargs.update(action_kwargs) + + pred_noise_btchw = current_model( + latent_model_input, + prompt_embeds, + t_expanded_noise, + **ctx.image_kwargs, + **ctx.pos_cond_kwargs, + **model_kwargs, + ).permute(0, 2, 1, 3, 4) + + if ctx.boundary_timestep is not None and t_cur >= ctx.boundary_timestep: + pred_video_btchw = pred_noise_to_x_bound( + pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + boundary_timestep=torch.ones_like(t_expand) * ctx.boundary_timestep, + scheduler=self.scheduler).unflatten(0, pred_noise_btchw.shape[:2]) + else: + pred_video_btchw = pred_noise_to_pred_video(pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + scheduler=self.scheduler).unflatten( + 0, pred_noise_btchw.shape[:2]) + + if i < len(timesteps) - 1: + next_timestep = timesteps[i + 1] * torch.ones([1], dtype=torch.long, device=pred_video_btchw.device) + + # Use custom noise generator if provided (for streaming), else generate + if noise_generator is not None: + noise = noise_generator(pred_video_btchw.shape, pred_video_btchw.dtype, i) + else: + noise = torch.randn( + pred_video_btchw.shape, + dtype=pred_video_btchw.dtype, + generator=(batch.generator[0] if isinstance(batch.generator, list) else batch.generator)).to( + pred_video_btchw.device) + + noise_btchw = noise + if ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and i < len( + ctx.high_noise_timesteps) - 1: + noise_latents_btchw = self.scheduler.add_noise_high( + pred_video_btchw.flatten(0, 1), noise_btchw.flatten(0, 1), next_timestep, + torch.ones_like(next_timestep) * ctx.boundary_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) + elif ctx.boundary_timestep is not None and ctx.high_noise_timesteps is not None and i == len( + ctx.high_noise_timesteps) - 1: + noise_latents_btchw = pred_video_btchw + else: + noise_latents_btchw = self.scheduler.add_noise(pred_video_btchw.flatten(0, 1), + noise_btchw.flatten(0, 1), next_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) + current_latents = noise_latents_btchw.permute(0, 2, 1, 3, 4) + else: + current_latents = pred_video_btchw.permute(0, 2, 1, 3, 4) + + if progress_bar is not None: + progress_bar.update() + + return current_latents + + def _update_context_cache( + self, + current_latents: torch.Tensor, + batch: ForwardBatch, + start_index: int, + current_num_frames: int, + ctx: BlockProcessingContext, + action_kwargs: dict[str, Any], + context_noise: float, + ) -> None: + prompt_embeds = batch.prompt_embeds + latents_device = current_latents.device + + # Expand context timestep to per-frame format [batch, num_frames] for causal model + t_context = torch.ones([current_latents.shape[0], current_num_frames], device=latents_device, + dtype=torch.long) * int(context_noise) + context_bcthw = current_latents.to(ctx.target_dtype) + + with torch.autocast(device_type="cuda", + dtype=ctx.target_dtype, + enabled=ctx.autocast_enabled), \ + set_forward_context(current_timestep=0, + attn_metadata=None, + forward_batch=batch): + + context_model_kwargs = { + "kv_cache": ctx.kv_cache1, + "crossattn_cache": ctx.crossattn_cache, + "current_start": start_index * self.frame_seq_length, + "start_frame": start_index, + } + + if self.use_action_module: + context_model_kwargs.update({ + "kv_cache_mouse": ctx.kv_cache_mouse, + "kv_cache_keyboard": ctx.kv_cache_keyboard, + }) + context_model_kwargs.update(action_kwargs) + + if ctx.boundary_timestep is not None and self.transformer_2 is not None: + self.transformer_2( + context_bcthw, + prompt_embeds, + t_context, + kv_cache=ctx.kv_cache2, + crossattn_cache=ctx.crossattn_cache, + current_start=start_index * self.frame_seq_length, + start_frame=start_index, + **ctx.image_kwargs, + **ctx.pos_cond_kwargs, + ) + + self.transformer( + context_bcthw, + prompt_embeds, + t_context, + **ctx.image_kwargs, + **ctx.pos_cond_kwargs, + **context_model_kwargs, + ) + + def streaming_reset(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast + + latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] + patch_size = self.transformer.patch_size + patch_ratio = patch_size[-1] * patch_size[-2] + self.frame_seq_length = latent_seq_length // patch_ratio + + timesteps = torch.tensor(fastvideo_args.pipeline_config.dmd_denoising_steps, dtype=torch.long).cpu() + if fastvideo_args.pipeline_config.warp_denoising_step: + scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) + timesteps = scheduler_timesteps[1000 - timesteps] + timesteps = timesteps.to(get_local_torch_device()) + + boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, 'boundary_ratio', None) + if boundary_ratio is not None: + boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps + high_noise_timesteps = timesteps[timesteps >= boundary_timestep] + else: + boundary_timestep = None + high_noise_timesteps = None + + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert torch.isnan(image_embeds[0]).sum() == 0 + image_embeds = [image_embed.to(target_dtype) for image_embed in image_embeds] + + # directly set the kwarg. + image_kwargs = {"encoder_hidden_states_image": image_embeds} + pos_cond_kwargs: dict[str, Any] = {} + + assert batch.latents is not None, "latents must be provided" + latents = batch.latents + b, c, t, h, w = latents.shape + prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + + # Initialize caches + kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], dtype=target_dtype, device=latents.device) + kv_cache2 = None + if boundary_timestep is not None: + kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + kv_cache_mouse = None + kv_cache_keyboard = None + if self.use_action_module: + kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + crossattn_cache = self._initialize_crossattn_cache( + batch_size=latents.shape[0], + max_text_len=257, # 1 CLS + 256 patch tokens + dtype=target_dtype, + device=latents.device) + + # Calculate block sizes + if t % self.num_frame_per_block != 0: + raise ValueError("num_frames must be divisible by num_frame_per_block for causal denoising") + num_blocks = t // self.num_frame_per_block + block_sizes = [self.num_frame_per_block] * num_blocks + if boundary_timestep is not None: + block_sizes[0] = 1 + + # Pre-allocate noise pool + num_denoising_steps = len(timesteps) + noise_shape = (b, self.num_frame_per_block, c, h, w) + noise_pool = [ + torch.randn( + noise_shape, + dtype=target_dtype, + device=latents.device, + ) for _ in range(num_denoising_steps - 1) + ] + + # Create and store context + self._streaming_ctx = BlockProcessingContext( + batch=batch, + block_idx=0, + start_index=0, + kv_cache1=kv_cache1, + kv_cache2=kv_cache2, + kv_cache_mouse=kv_cache_mouse, + kv_cache_keyboard=kv_cache_keyboard, + crossattn_cache=crossattn_cache, + timesteps=timesteps, + block_sizes=block_sizes, + noise_pool=noise_pool, + fastvideo_args=fastvideo_args, + target_dtype=target_dtype, + autocast_enabled=autocast_enabled, + boundary_timestep=boundary_timestep, + high_noise_timesteps=high_noise_timesteps, + context_noise=getattr(fastvideo_args.pipeline_config, "context_noise", 0), + image_kwargs=image_kwargs, + pos_cond_kwargs=pos_cond_kwargs, + ) + + self._streaming_initialized = True + return batch + + def streaming_step(self, + keyboard_action: torch.Tensor | None = None, + mouse_action: torch.Tensor | None = None) -> ForwardBatch: + if not self._streaming_initialized or self._streaming_ctx is None: + raise RuntimeError("Streaming not initialized! Call streaming_reset first.") + + ctx = self._streaming_ctx + if ctx.block_idx >= len(ctx.block_sizes): + return ctx.batch + + batch = ctx.batch + latents = batch.latents + assert latents is not None, "latents must be set in batch" + + current_num_frames = ctx.block_sizes[ctx.block_idx] + start_index = ctx.start_index + + current_latents = latents[:, :, start_index:start_index + current_num_frames, :, :] + + # Update batch with new actions for this block + if keyboard_action is not None or mouse_action is not None: + vae_ratio = 4 + start_frame = 0 if start_index == 0 else 1 + vae_ratio * (start_index - 1) + + if keyboard_action is not None: + n = keyboard_action.shape[1] + batch.keyboard_cond[:, start_frame:start_frame + n] = keyboard_action.to(batch.keyboard_cond.device) + if mouse_action is not None: + n = mouse_action.shape[1] + batch.mouse_cond[:, start_frame:start_frame + n] = mouse_action.to(batch.mouse_cond.device) + + action_kwargs = self._prepare_action_kwargs(batch, start_index, current_num_frames) + + # Create noise generator that uses pre-allocated noise pool + def streaming_noise_generator(shape: tuple, dtype: torch.dtype, step_idx: int) -> torch.Tensor: + if ctx.noise_pool is not None and step_idx < len(ctx.noise_pool): + return ctx.noise_pool[step_idx][:, :shape[1], :, :, :].to(latents.device) + else: + # Fallback to dynamic allocation if pool not available + return torch.randn( + shape, + dtype=dtype, + generator=(batch.generator[0] if isinstance(batch.generator, list) else batch.generator)).to( + latents.device) + + current_latents = self._process_single_block( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + timesteps=ctx.timesteps, + ctx=ctx, + action_kwargs=action_kwargs, + noise_generator=streaming_noise_generator, + ) + + latents[:, :, start_index:start_index + current_num_frames, :, :] = current_latents + + # Update KV caches with clean context + self._update_context_cache( + current_latents=current_latents, + batch=batch, + start_index=start_index, + current_num_frames=current_num_frames, + ctx=ctx, + action_kwargs=action_kwargs, + context_noise=ctx.context_noise, + ) + + # Advance streaming state + ctx.start_index += current_num_frames + ctx.block_idx += 1 + + return batch + + def streaming_clear(self) -> None: + self._streaming_initialized = False + self._streaming_ctx = None + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check("image_embeds", batch.image_embeds, V.is_list) + result.add_check("image_latent", batch.image_latent, V.none_or_tensor_with_dims(5)) + result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) + result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) + result.add_check("eta", batch.eta, V.non_negative_float) + result.add_check("generator", batch.generator, V.generator_or_list_generators) + result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value) + result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds, + lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)) + return result diff --git a/fastvideo/platforms/cuda.py b/fastvideo/platforms/cuda.py index aa6c10ec2..1b048fcda 100644 --- a/fastvideo/platforms/cuda.py +++ b/fastvideo/platforms/cuda.py @@ -130,7 +130,7 @@ class CudaPlatformBase(Platform): logger.info("Sage Attention backend is not installed. Fall back to Flash Attention.") elif selected_backend == AttentionBackendEnum.SAGE_ATTN_THREE: try: - from sageattn3 import sageattn3_blackwell # noqa: F401 + from modified_sageattn import sageattn_blackwell # noqa: F401 from fastvideo.attention.backends.sage_attn3 import ( # noqa: F401 SageAttention3Backend) diff --git a/fused_attention.py b/fused_attention.py deleted file mode 100644 index e3d9f1cef..000000000 --- a/fused_attention.py +++ /dev/null @@ -1,763 +0,0 @@ -""" -Fused Attention -=============== - -This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao (https://tridao.me/publications/flash2/flash2.pdf) - -Credits: OpenAI kernel team - -Extra Credits: - -* Original flash attention paper (https://arxiv.org/abs/2205.14135) -* Rabe and Staats (https://arxiv.org/pdf/2112.05682v2.pdf) - -""" - -import pytest -import torch -import os - -import triton -import triton.language as tl -from triton.tools.tensor_descriptor import TensorDescriptor - -DEVICE = triton.runtime.driver.active.get_active_torch_device() - - -def is_hip(): - return triton.runtime.driver.active.get_current_target().backend == "hip" - - -def is_cuda(): - return triton.runtime.driver.active.get_current_target().backend == "cuda" - - -def supports_host_descriptor(): - return is_cuda() and torch.cuda.get_device_capability()[0] >= 9 - - -def is_blackwell(): - return is_cuda() and torch.cuda.get_device_capability()[0] == 10 - - -def is_hopper(): - return is_cuda() and torch.cuda.get_device_capability()[0] == 9 - - -@triton.jit -def _attn_fwd_inner(acc, l_i, m_i, q, # - desc_k, desc_v, # - offset_y, dtype: tl.constexpr, start_m, qk_scale, # - BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr, BLOCK_N: tl.constexpr, # - STAGE: tl.constexpr, offs_m: tl.constexpr, offs_n: tl.constexpr, # - N_CTX: tl.constexpr, warp_specialize: tl.constexpr, IS_HOPPER: tl.constexpr): - # range of values handled by this stage - if STAGE == 1: - lo, hi = 0, start_m * BLOCK_M - 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 - if dtype == tl.float8e5: - 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) - # -- compute qk ---- - k = desc_k.load([offsetk_y, 0]).T - qk = tl.dot(q, k) - if STAGE == 2: - mask = offs_m[:, None] >= (start_n + offs_n[None, :]) - qk = qk * qk_scale + tl.where(mask, 0, -1.0e6) - m_ij = tl.maximum(m_i, tl.max(qk, 1)) - qk -= m_ij[:, None] - else: - m_ij = tl.maximum(m_i, tl.max(qk, 1) * qk_scale) - qk = qk * qk_scale - m_ij[:, None] - p = tl.math.exp2(qk) - # -- compute correction factor - alpha = tl.math.exp2(m_i - m_ij) - l_ij = tl.sum(p, 1) - # -- update output accumulator -- - if not IS_HOPPER and warp_specialize and BLOCK_M == 128 and HEAD_DIM == 128: - BM: tl.constexpr = acc.shape[0] - BN: tl.constexpr = acc.shape[1] - acc0, acc1 = acc.reshape([BM, 2, BN // 2]).permute(0, 2, 1).split() - acc0 = acc0 * alpha[:, None] - acc1 = acc1 * alpha[:, None] - acc = tl.join(acc0, acc1).permute(0, 2, 1).reshape([BM, BN]) - else: - acc = acc * alpha[:, None] - # prepare p and v for the dot - if dtype == tl.float8e5: - v = desc_v.load([0, offsetv_y]).T - else: - v = desc_v.load([offsetv_y, 0]) - p = p.to(dtype) - # note that this non transposed v for FP8 is only supported on Blackwell - acc = tl.dot(p, v.to(dtype), acc) - # update m_i and l_i - # place this at the end of the loop to reduce register pressure - l_i = l_i * alpha + l_ij - m_i = m_ij - offsetk_y += BLOCK_N - offsetv_y += BLOCK_N - return acc, l_i, m_i - - -def _host_descriptor_pre_hook(nargs): - BLOCK_M = nargs["BLOCK_M"] - BLOCK_N = nargs["BLOCK_N"] - HEAD_DIM = nargs["HEAD_DIM"] - if not isinstance(nargs["desc_q"], TensorDescriptor): - return - nargs["desc_q"].block_shape = [BLOCK_M, HEAD_DIM] - if nargs["FP8_OUTPUT"]: - nargs["desc_v"].block_shape = [HEAD_DIM, BLOCK_N] - else: - nargs["desc_v"].block_shape = [BLOCK_N, HEAD_DIM] - nargs["desc_k"].block_shape = [BLOCK_N, HEAD_DIM] - nargs["desc_o"].block_shape = [BLOCK_M, HEAD_DIM] - - -if is_hip(): - NUM_STAGES_OPTIONS = [1] -elif supports_host_descriptor(): - NUM_STAGES_OPTIONS = [2, 3, 4] -else: - NUM_STAGES_OPTIONS = [2, 3, 4] - -configs = [ - triton.Config({'BLOCK_M': BM, 'BLOCK_N': BN}, num_stages=s, num_warps=w, pre_hook=_host_descriptor_pre_hook) \ - for BM in [64, 128]\ - for BN in [32, 64, 128]\ - 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), - ] - - -def keep(conf): - BLOCK_M = conf.kwargs["BLOCK_M"] - BLOCK_N = conf.kwargs["BLOCK_N"] - return not (is_cuda() and torch.cuda.get_device_capability()[0] == 9 and BLOCK_M * BLOCK_N < 128 * 128 - and conf.num_warps == 8) - - -def prune_invalid_configs(configs, named_args, **kwargs): - N_CTX = kwargs["N_CTX"] - - # Filter out configs where BLOCK_M > N_CTX - return [conf for conf in configs if conf.kwargs.get("BLOCK_M", 0) <= N_CTX] - - -@triton.jit -def _maybe_make_tensor_desc(desc_or_ptr, shape, strides, block_shape): - if isinstance(desc_or_ptr, tl.tensor_descriptor): - return desc_or_ptr - else: - return tl.make_tensor_descriptor(desc_or_ptr, shape, strides, block_shape) - - -@triton.autotune(configs=list(filter(keep, configs)), key=["N_CTX", "HEAD_DIM", "FP8_OUTPUT", "warp_specialize"], - prune_configs_by={'early_config_prune': prune_invalid_configs}) -@triton.jit -def _attn_fwd(sm_scale, M, # - Z, H, desc_q, desc_k, desc_v, desc_o, N_CTX, # - HEAD_DIM: tl.constexpr, # - BLOCK_M: tl.constexpr, # - BLOCK_N: tl.constexpr, # - FP8_OUTPUT: tl.constexpr, # - STAGE: tl.constexpr, # - warp_specialize: tl.constexpr, # - IS_HOPPER: tl.constexpr, # - ): - dtype = tl.float8e5 if FP8_OUTPUT else tl.float16 - tl.static_assert(BLOCK_N <= HEAD_DIM) - start_m = tl.program_id(0) - off_hz = tl.program_id(1) - off_z = off_hz // H - off_h = off_hz % H - - y_dim = Z * H * N_CTX - desc_q = _maybe_make_tensor_desc(desc_q, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], - block_shape=[BLOCK_M, HEAD_DIM]) - if FP8_OUTPUT: - desc_v = _maybe_make_tensor_desc(desc_v, shape=[HEAD_DIM, y_dim], strides=[N_CTX, 1], - block_shape=[HEAD_DIM, BLOCK_N]) - else: - desc_v = _maybe_make_tensor_desc(desc_v, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], - block_shape=[BLOCK_N, HEAD_DIM]) - desc_k = _maybe_make_tensor_desc(desc_k, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], - block_shape=[BLOCK_N, HEAD_DIM]) - desc_o = _maybe_make_tensor_desc(desc_o, shape=[y_dim, HEAD_DIM], strides=[HEAD_DIM, 1], - block_shape=[BLOCK_M, HEAD_DIM]) - - offset_y = off_z * (N_CTX * H) + off_h * N_CTX - qo_offset_y = offset_y + start_m * BLOCK_M - # initialize offsets - offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = tl.arange(0, BLOCK_N) - # 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) - # 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]) - # stage 1: off-band - # For causal = True, STAGE = 3 and _attn_fwd_inner gets 1 as its STAGE - # For causal = False, STAGE = 1, and _attn_fwd_inner gets 3 as its STAGE - if STAGE & 1: - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, # - desc_k, desc_v, # - offset_y, dtype, start_m, qk_scale, # - BLOCK_M, HEAD_DIM, BLOCK_N, # - 4 - STAGE, offs_m, offs_n, N_CTX, # - warp_specialize, IS_HOPPER) - # stage 2: on-band - if STAGE & 2: - acc, l_i, m_i = _attn_fwd_inner(acc, l_i, m_i, q, # - desc_k, desc_v, # - offset_y, dtype, start_m, qk_scale, # - BLOCK_M, HEAD_DIM, BLOCK_N, # - 2, offs_m, offs_n, N_CTX, # - warp_specialize, IS_HOPPER) - # epilogue - m_i += tl.math.log2(l_i) - acc = acc / l_i[:, None] - m_ptrs = M + off_hz * N_CTX + offs_m - tl.store(m_ptrs, m_i) - desc_o.store([qo_offset_y, 0], acc.to(dtype)) - - -@triton.jit -def _attn_bwd_preprocess(O, DO, # - Delta, # - Z, H, N_CTX, # - BLOCK_M: tl.constexpr, HEAD_DIM: tl.constexpr # - ): - 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) - - -# The main inner-loop logic for computing dK and dV. -@triton.jit -def _attn_bwd_dkdv(dk, dv, # - Q, k, v, sm_scale, # - DO, # - M, D, # - # shared by Q/K/V/DO. - stride_tok, stride_d, # - H, N_CTX, BLOCK_M1: tl.constexpr, # - BLOCK_N1: tl.constexpr, # - HEAD_DIM: tl.constexpr, # - # Filled in by the wrapper. - start_n, start_m, num_steps, # - MASK: tl.constexpr): - offs_m = start_m + tl.arange(0, BLOCK_M1) - offs_n = start_n + tl.arange(0, BLOCK_N1) - offs_k = tl.arange(0, HEAD_DIM) - qT_ptrs = Q + offs_m[None, :] * stride_tok + offs_k[:, None] * stride_d - do_ptrs = DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d - # BLOCK_N1 must be a multiple of BLOCK_M1, otherwise the code wouldn't work. - tl.static_assert(BLOCK_N1 % BLOCK_M1 == 0) - curr_m = start_m - step_m = BLOCK_M1 - for blk_idx in range(num_steps): - qT = tl.load(qT_ptrs) - # Load m before computing qk to reduce pipeline stall. - offs_m = curr_m + tl.arange(0, BLOCK_M1) - m = tl.load(M + offs_m) - qkT = tl.dot(k, qT) - pT = tl.math.exp2(qkT - m[None, :]) - # Autoregressive masking. - if MASK: - mask = (offs_m[None, :] >= offs_n[:, None]) - pT = tl.where(mask, pT, 0.0) - do = tl.load(do_ptrs) - # Compute dV. - ppT = pT - ppT = ppT.to(tl.float16) - dv += tl.dot(ppT, do.to(tl.float16)) - # 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.float16) - dk += tl.dot(dsT, tl.trans(qT.to(tl.float16))) - # Increment pointers. - curr_m += step_m - qT_ptrs += step_m * stride_tok - do_ptrs += step_m * stride_tok - return dk, dv - - -# the main inner-loop logic for computing dQ -@triton.jit -def _attn_bwd_dq(dq, q, K, V, # - do, m, D, - # shared by Q/K/V/DO. - stride_tok, stride_d, # - H, N_CTX, # - BLOCK_M2: tl.constexpr, # - BLOCK_N2: tl.constexpr, # - HEAD_DIM: tl.constexpr, - # Filled in by the wrapper. - start_m, start_n, num_steps, # - MASK: tl.constexpr): - offs_m = start_m + tl.arange(0, BLOCK_M2) - offs_n = start_n + tl.arange(0, BLOCK_N2) - offs_k = tl.arange(0, HEAD_DIM) - kT_ptrs = K + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d - vT_ptrs = V + offs_n[None, :] * stride_tok + offs_k[:, None] * stride_d - # D (= delta) is pre-divided by ds_scale. - Di = tl.load(D + offs_m) - # BLOCK_M2 must be a multiple of BLOCK_N2, otherwise the code wouldn't work. - tl.static_assert(BLOCK_M2 % BLOCK_N2 == 0) - curr_n = start_n - step_n = BLOCK_N2 - for blk_idx in range(num_steps): - kT = tl.load(kT_ptrs) - vT = tl.load(vT_ptrs) - qk = tl.dot(q, kT) - p = tl.math.exp2(qk - m) - # Autoregressive masking. - if MASK: - offs_n = curr_n + tl.arange(0, BLOCK_N2) - mask = (offs_m[:, None] >= offs_n[None, :]) - p = tl.where(mask, p, 0.0) - # Compute dP and dS. - dp = tl.dot(do, vT).to(tl.float32) - ds = p * (dp - Di[:, None]) - ds = ds.to(tl.float16) - # Compute dQ. - # NOTE: We need to de-scale dq in the end, because kT was pre-scaled. - dq += tl.dot(ds, tl.trans(kT.to(tl.float16))) - # Increment pointers. - curr_n += step_n - kT_ptrs += step_n * stride_tok - vT_ptrs += step_n * stride_tok - return dq - - -@triton.jit -def _attn_bwd(Q, K, V, sm_scale, # - DO, # - DQ, DK, DV, # - M, D, - # shared by Q/K/V/DO. - stride_z, stride_h, stride_tok, stride_d, # - H, N_CTX, # - BLOCK_M1: tl.constexpr, # - BLOCK_N1: tl.constexpr, # - BLOCK_M2: tl.constexpr, # - BLOCK_N2: tl.constexpr, # - BLK_SLICE_FACTOR: tl.constexpr, # - HEAD_DIM: tl.constexpr): - LN2: tl.constexpr = 0.6931471824645996 # = ln(2) - - bhid = tl.program_id(2) - off_chz = (bhid * N_CTX).to(tl.int64) - adj = (stride_h * (bhid % H) + stride_z * (bhid // H)).to(tl.int64) - pid = tl.program_id(0) - - # offset pointers for batch/head - Q += adj - K += adj - V += adj - DO += adj - DQ += adj - DK += adj - DV += adj - M += off_chz - D += off_chz - - # load scales - offs_k = tl.arange(0, HEAD_DIM) - - start_n = pid * BLOCK_N1 - start_m = start_n - - MASK_BLOCK_M1: tl.constexpr = BLOCK_M1 // BLK_SLICE_FACTOR - offs_n = start_n + tl.arange(0, BLOCK_N1) - - 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) - - num_steps = BLOCK_N1 // MASK_BLOCK_M1 - - dk, dv = _attn_bwd_dkdv(dk, dv, # - Q, k, v, sm_scale, # - DO, # - M, D, # - stride_tok, stride_d, # - H, N_CTX, # - MASK_BLOCK_M1, BLOCK_N1, HEAD_DIM, # - start_n, start_m, num_steps, # - MASK=True # - ) - - start_m += num_steps * MASK_BLOCK_M1 - num_steps = (N_CTX - start_m) // BLOCK_M1 - - # Compute dK and dV for non-masked blocks. - dk, dv = _attn_bwd_dkdv( # - dk, dv, # - Q, k, v, sm_scale, # - DO, # - M, D, # - stride_tok, stride_d, # - H, N_CTX, # - BLOCK_M1, BLOCK_N1, HEAD_DIM, # - start_n, start_m, num_steps, # - MASK=False # - ) - - 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) - - # THIS BLOCK DOES DQ: - start_m = pid * BLOCK_M2 - end_n = start_m + BLOCK_M2 - - MASK_BLOCK_N2: tl.constexpr = BLOCK_N2 // BLK_SLICE_FACTOR - offs_m = start_m + tl.arange(0, BLOCK_M2) - - q = tl.load(Q + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d) - dq = tl.zeros([BLOCK_M2, HEAD_DIM], dtype=tl.float32) - do = tl.load(DO + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d) - - m = tl.load(M + offs_m) - m = m[:, None] - - # Compute dQ for masked (diagonal) blocks. - # NOTE: This code scans each row of QK^T backward (from right to left, - # but inside each call to _attn_bwd_dq, from left to right), but that's - # not due to anything important. I just wanted to reuse the loop - # structure for dK & dV above as much as possible. - num_steps = BLOCK_M2 // MASK_BLOCK_N2 - dq = _attn_bwd_dq(dq, q, K, V, # - do, m, D, # - stride_tok, stride_d, # - H, N_CTX, # - BLOCK_M2, MASK_BLOCK_N2, HEAD_DIM, # - start_m, end_n - num_steps * MASK_BLOCK_N2, num_steps, # - MASK=True # - ) - end_n -= num_steps * MASK_BLOCK_N2 - # stage 2 - num_steps = end_n // BLOCK_N2 - dq = _attn_bwd_dq(dq, q, K, V, # - do, m, D, # - stride_tok, stride_d, # - H, N_CTX, # - BLOCK_M2, BLOCK_N2, HEAD_DIM, # - start_m, end_n - num_steps * BLOCK_N2, num_steps, # - MASK=False # - ) - # Write back dQ. - dq_ptrs = DQ + offs_m[:, None] * stride_tok + offs_k[None, :] * stride_d - dq *= LN2 - tl.store(dq_ptrs, dq) - - -class _attention(torch.autograd.Function): - - @staticmethod - def forward(ctx, q, k, v, causal, sm_scale, warp_specialize=True): - # shape constraints - HEAD_DIM_Q, HEAD_DIM_K = q.shape[-1], k.shape[-1] - # when v is in float8_e5m2 it is transposed. - HEAD_DIM_V = v.shape[-1] - assert HEAD_DIM_Q == HEAD_DIM_K and HEAD_DIM_K == HEAD_DIM_V - assert HEAD_DIM_K in {16, 32, 64, 128, 256} - o = torch.empty_like(q) - stage = 3 if causal else 1 - extra_kern_args = {} - # Tuning for AMD target - if is_hip(): - waves_per_eu = 3 if HEAD_DIM_K <= 64 else 2 - extra_kern_args = {"waves_per_eu": waves_per_eu, "allow_flush_denorm": True} - - M = torch.empty((q.shape[0], q.shape[1], q.shape[2]), device=q.device, dtype=torch.float32) - # Use device_descriptor for Hopper + warpspec. - if supports_host_descriptor() and not (is_hopper() and warp_specialize): - # Note that on Hopper we cannot perform a FP8 dot with a non-transposed second tensor - y_dim = q.shape[0] * q.shape[1] * q.shape[2] - - dummy_block = [1, 1] - desc_q = TensorDescriptor(q, shape=[y_dim, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block) - if q.dtype == torch.float8_e5m2: - desc_v = TensorDescriptor(v, shape=[HEAD_DIM_K, y_dim], strides=[q.shape[2], 1], - block_shape=dummy_block) - else: - desc_v = TensorDescriptor(v, shape=[y_dim, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], - block_shape=dummy_block) - desc_k = TensorDescriptor(k, shape=[y_dim, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block) - desc_o = TensorDescriptor(o, shape=[y_dim, HEAD_DIM_K], strides=[HEAD_DIM_K, 1], block_shape=dummy_block) - else: - desc_q = q - desc_v = v - desc_k = k - desc_o = o - - def alloc_fn(size: int, align: int, _): - return torch.empty(size, dtype=torch.int8, device="cuda") - - triton.set_allocator(alloc_fn) - - def grid(META): - return (triton.cdiv(q.shape[2], META["BLOCK_M"]), q.shape[0] * q.shape[1], 1) - - ctx.grid = grid - if is_blackwell() and warp_specialize: - if HEAD_DIM_K == 128 and q.dtype == torch.float16: - extra_kern_args["maxnreg"] = 168 - else: - extra_kern_args["maxnreg"] = 80 - _attn_fwd[grid]( - sm_scale, M, # - q.shape[0], q.shape[1], # - desc_q, desc_k, desc_v, desc_o, # - N_CTX=q.shape[2], # - HEAD_DIM=HEAD_DIM_K, # - FP8_OUTPUT=q.dtype == torch.float8_e5m2, # - STAGE=stage, # - warp_specialize=warp_specialize, # - IS_HOPPER=is_hopper(), # - **extra_kern_args) - - ctx.save_for_backward(q, k, v, o, M) - ctx.sm_scale = sm_scale - ctx.HEAD_DIM = HEAD_DIM_K - ctx.causal = causal - return o - - @staticmethod - def backward(ctx, do): - q, k, v, o, M = ctx.saved_tensors - do = do.contiguous() - assert do.is_contiguous() - assert q.stride() == k.stride() == v.stride() == o.stride() == do.stride() - dq = torch.empty_like(q) - dk = torch.empty_like(k) - dv = torch.empty_like(v) - BATCH, N_HEAD, N_CTX = q.shape[:3] - PRE_BLOCK = 128 - NUM_WARPS, NUM_STAGES = 4, 5 - BLOCK_M1, BLOCK_N1, BLOCK_M2, BLOCK_N2 = 32, 128, 128, 32 - BLK_SLICE_FACTOR = 2 - RCP_LN2 = 1.4426950408889634 # = 1.0 / ln(2) - arg_k = k - arg_k = arg_k * (ctx.sm_scale * RCP_LN2) - PRE_BLOCK = 128 - assert N_CTX % PRE_BLOCK == 0 - pre_grid = (N_CTX // PRE_BLOCK, BATCH * N_HEAD) - delta = torch.empty_like(M) - _attn_bwd_preprocess[pre_grid]( - o, do, # - delta, # - BATCH, N_HEAD, N_CTX, # - BLOCK_M=PRE_BLOCK, HEAD_DIM=ctx.HEAD_DIM # - ) - grid = (N_CTX // BLOCK_N1, 1, BATCH * N_HEAD) - _attn_bwd[grid]( - q, arg_k, v, ctx.sm_scale, do, dq, dk, dv, # - M, delta, # - q.stride(0), q.stride(1), q.stride(2), q.stride(3), # - N_HEAD, N_CTX, # - BLOCK_M1=BLOCK_M1, BLOCK_N1=BLOCK_N1, # - BLOCK_M2=BLOCK_M2, BLOCK_N2=BLOCK_N2, # - BLK_SLICE_FACTOR=BLK_SLICE_FACTOR, # - HEAD_DIM=ctx.HEAD_DIM, # - num_warps=NUM_WARPS, # - num_stages=NUM_STAGES # - ) - - return dq, dk, dv, None, None, None, None - - -attention = _attention.apply - -TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2') - - -@pytest.mark.parametrize("Z", [1, 4]) -@pytest.mark.parametrize("H", [2, 48]) -@pytest.mark.parametrize("N_CTX", [128, 1024, (2 if is_hip() else 4) * 1024]) -@pytest.mark.parametrize("HEAD_DIM", [64, 128]) -@pytest.mark.parametrize("causal", [True]) # FIXME: Non-causal tests do not pass at the moment. -@pytest.mark.parametrize("warp_specialize", [False, True] if is_blackwell() else [False]) -@pytest.mark.parametrize("mode", ["fwd", "bwd"]) -@pytest.mark.parametrize("provider", ["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else [])) -def test_op(Z, H, N_CTX, HEAD_DIM, causal, warp_specialize, mode, provider, dtype=torch.float16): - if mode == "fwd" and "fp16" in provider: - pytest.skip("Avoid running the forward computation twice.") - if mode == "bwd" and "fp8" in provider: - pytest.skip("Backward pass with FP8 is not supported.") - torch.manual_seed(20) - q = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) - k = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) - v = (torch.empty((Z, H, N_CTX, HEAD_DIM), dtype=dtype, device=DEVICE).normal_(mean=0.0, std=0.5).requires_grad_()) - sm_scale = 0.5 - # reference implementation - ref_dtype = dtype - if mode == "fwd" and "fp8" in provider: - ref_dtype = torch.float32 - q = q.to(ref_dtype) - k = k.to(ref_dtype) - v = v.to(ref_dtype) - M = torch.tril(torch.ones((N_CTX, N_CTX), device=DEVICE)) - p = torch.matmul(q, k.transpose(2, 3)) * sm_scale - if causal: - p[:, :, M == 0] = float("-inf") - p = torch.softmax(p.float(), dim=-1) - p = p.to(ref_dtype) - # p = torch.exp(p) - ref_out = torch.matmul(p, v).half() - if mode == "bwd": - dout = torch.randn_like(q) - ref_out.backward(dout) - ref_dv, v.grad = v.grad.clone(), None - ref_dk, k.grad = k.grad.clone(), None - ref_dq, q.grad = q.grad.clone(), None - # triton implementation - if mode == "fwd" and "fp8" in provider: - q = q.to(torch.float8_e5m2) - k = k.to(torch.float8_e5m2) - v = v.permute(0, 1, 3, 2).contiguous() - v = v.permute(0, 1, 3, 2) - v = v.to(torch.float8_e5m2) - tri_out = attention(q, k, v, causal, sm_scale, warp_specialize).half() - if mode == "fwd": - atol = 3 if "fp8" in provider else 1e-2 - torch.testing.assert_close(tri_out, ref_out, atol=atol, rtol=0) - return - tri_out.backward(dout) - tri_dv, v.grad = v.grad.clone(), None - tri_dk, k.grad = k.grad.clone(), None - tri_dq, q.grad = q.grad.clone(), None - # compare - torch.testing.assert_close(tri_out, ref_out, atol=1e-2, rtol=0) - rtol = 0.0 - # Relative tolerance workaround for known hardware limitation of CDNA2 GPU. - # For details see https://pytorch.org/docs/stable/notes/numerical_accuracy.html#reduced-precision-fp16-and-bf16-gemms-and-convolutions-on-amd-instinct-mi200-devices - if torch.version.hip is not None and triton.runtime.driver.active.get_current_target().arch == "gfx90a": - rtol = 1e-2 - torch.testing.assert_close(tri_dv, ref_dv, atol=1e-2, rtol=rtol) - torch.testing.assert_close(tri_dk, ref_dk, atol=1e-2, rtol=rtol) - torch.testing.assert_close(tri_dq, ref_dq, atol=1e-2, rtol=rtol) - - -try: - from flash_attn.flash_attn_interface import \ - flash_attn_qkvpacked_func as flash_attn_func - HAS_FLASH = True -except BaseException: - HAS_FLASH = False - -TORCH_HAS_FP8 = hasattr(torch, 'float8_e5m2') -BATCH, N_HEADS = 4, 32 -# vary seq length for fixed head and batch=4 -configs = [] -for HEAD_DIM in [64, 128]: - for mode in ["fwd", "bwd"]: - for causal in [True, False]: - # Enable warpspec for causal fwd on Hopper - enable_ws = mode == "fwd" and (is_blackwell() or (is_hopper() and not causal)) - for warp_specialize in [False, True] if enable_ws else [False]: - configs.append( - triton.testing.Benchmark( - x_names=["N_CTX"], - x_vals=[2**i for i in range(10, 15)], - line_arg="provider", - line_vals=["triton-fp16"] + (["triton-fp8"] if TORCH_HAS_FP8 else []) + - (["flash"] if HAS_FLASH else []), - line_names=["Triton [FP16]"] + (["Triton [FP8]"] if TORCH_HAS_FP8 else []) + - (["Flash-2"] if HAS_FLASH else []), - styles=[("red", "-"), ("blue", "-"), ("green", "-")], - ylabel="TFLOPS", - plot_name= - f"fused-attention-batch{BATCH}-head{N_HEADS}-d{HEAD_DIM}-{mode}-causal={causal}-warp_specialize={warp_specialize}", - args={ - "H": N_HEADS, - "BATCH": BATCH, - "HEAD_DIM": HEAD_DIM, - "mode": mode, - "causal": causal, - "warp_specialize": warp_specialize, - }, - )) - - -@triton.testing.perf_report(configs) -def bench_flash_attention(BATCH, H, N_CTX, HEAD_DIM, causal, warp_specialize, mode, provider, device=DEVICE): - assert mode in ["fwd", "bwd"] - dtype = torch.float16 - if "triton" in provider: - q = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True) - k = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True) - v = torch.randn((BATCH, H, N_CTX, HEAD_DIM), dtype=dtype, device=device, requires_grad=True) - if mode == "fwd" and "fp8" in provider: - q = q.to(torch.float8_e5m2) - k = k.to(torch.float8_e5m2) - v = v.permute(0, 1, 3, 2).contiguous() - v = v.permute(0, 1, 3, 2) - v = v.to(torch.float8_e5m2) - sm_scale = 1.3 - fn = lambda: attention(q, k, v, causal, sm_scale, warp_specialize) - if mode == "bwd": - o = fn() - do = torch.randn_like(o) - fn = lambda: o.backward(do, retain_graph=True) - ms = triton.testing.do_bench(fn) - - if provider == "flash": - qkv = torch.randn((BATCH, N_CTX, 3, H, HEAD_DIM), dtype=dtype, device=device, requires_grad=True) - fn = lambda: flash_attn_func(qkv, causal=causal) - if mode == "bwd": - o = fn() - do = torch.randn_like(o) - fn = lambda: o.backward(do, retain_graph=True) - ms = triton.testing.do_bench(fn) - flops_per_matmul = 2.0 * BATCH * H * N_CTX * N_CTX * HEAD_DIM - total_flops = 2 * flops_per_matmul - if causal: - total_flops *= 0.5 - if mode == "bwd": - total_flops *= 2.5 # 2.0(bwd) + 0.5(recompute) - return total_flops * 1e-12 / (ms * 1e-3) - - -if __name__ == "__main__": - # only works on post-Ampere GPUs right now - bench_flash_attention.run(save_path=".", print_data=True) diff --git a/test_fake_quant_attn.py b/test_fake_quant_attn.py deleted file mode 100644 index f7cc0931e..000000000 --- a/test_fake_quant_attn.py +++ /dev/null @@ -1,416 +0,0 @@ -#!/usr/bin/env python3 -""" -Test script for fake quantization of attention computation using get_fake_quant method. -Tests attention computation with fake-quantized Q, K, V tensors and compares with reference. -""" - -import torch -import torch.nn.functional as F -from math import sqrt -from test_quantization_precision import get_fake_quant -import sys -import os -from contextlib import contextmanager - - -# Global flag to control verbosity -VERBOSE = False - - -@contextmanager -def suppress_stdout(): - """Context manager to suppress stdout.""" - if not VERBOSE: - with open(os.devnull, 'w') as devnull: - old_stdout = sys.stdout - sys.stdout = devnull - try: - yield - finally: - sys.stdout = old_stdout - else: - yield - - -def scaled_dot_product_attention_reference(q, k, v, causal=False, scale=None): - """ - Reference implementation of scaled dot-product attention. - - Args: - q: Query tensor [..., seq_len_q, head_dim] - k: Key tensor [..., seq_len_kv, head_dim] - v: Value tensor [..., seq_len_kv, head_dim] - causal: Whether to apply causal masking - scale: Scale factor (default: 1/sqrt(head_dim)) - - Returns: - Output tensor [..., seq_len_q, head_dim] - """ - if scale is None: - scale = 1.0 / sqrt(q.shape[-1]) - - # Compute attention scores: Q @ K^T - scores = torch.matmul(q, k.transpose(-2, -1)) * scale - - # Apply causal mask if needed - if causal: - seq_len = scores.shape[-1] - causal_mask = torch.triu(torch.ones(seq_len, seq_len, device=scores.device, dtype=scores.dtype), diagonal=1) - causal_mask = causal_mask.masked_fill(causal_mask == 1, float('-inf')) - scores = scores + causal_mask - - # Apply softmax - attn_weights = F.softmax(scores, dim=-1) - - # Apply attention weights to values - output = torch.matmul(attn_weights, v) - - return output - - -def attention_with_fake_quant(q, k, v, causal=False, scale=None, quantize_qkv=True): - """ - Compute attention with fake quantization applied to Q, K, V tensors. - - Args: - q: Query tensor [..., seq_len_q, head_dim] - k: Key tensor [..., seq_len_kv, head_dim] - v: Value tensor [..., seq_len_kv, head_dim] - causal: Whether to apply causal masking - scale: Scale factor (default: 1/sqrt(head_dim)) - quantize_qkv: Whether to apply fake quantization to Q, K, V - - Returns: - Output tensor [..., seq_len_q, head_dim] - """ - if quantize_qkv: - # Apply fake quantization to Q, K, V - if VERBOSE: - print(" Applying fake quantization to Q, K, V...") - with suppress_stdout(): - q = get_fake_quant(q) - k = get_fake_quant(k) - v = get_fake_quant(v) - - # Use PyTorch's scaled_dot_product_attention - return F.scaled_dot_product_attention( - q, k, v, - is_causal=causal, - scale=scale - ) - - -def test_attention_fake_quant_basic(): - """Test basic attention computation with fake quantization.""" - print("\n" + "="*60) - print("Test 1: Basic Attention with Fake Quantization") - print("="*60) - - if not torch.cuda.is_available(): - print("CUDA not available, skipping test") - return - - device = "cuda" - dtype = torch.bfloat16 - - # Test parameters - batch_size = 2 - num_heads = 4 - seq_len = 128 - head_dim = 64 - - # Create input tensors - torch.manual_seed(42) - q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - v = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - - print(f"Input shapes: Q={q.shape}, K={k.shape}, V={v.shape}") - print(f"Input dtype: {dtype}") - print(f"Q stats: min={q.min().item():.4f}, max={q.max().item():.4f}, mean={q.mean().item():.4f}") - - # Reference attention (no quantization) - print("\nComputing reference attention (no quantization)...") - with torch.no_grad(): - ref_output = F.scaled_dot_product_attention(q, k, v, is_causal=False) - - # Attention with fake quantization - print("\nComputing attention with fake quantization...") - with torch.no_grad(): - quant_output = attention_with_fake_quant(q, k, v, causal=False, quantize_qkv=True) - - # Compare results - diff = torch.abs(ref_output - quant_output) - max_diff = diff.max().item() - mean_diff = diff.mean().item() - relative_diff = (diff / (torch.abs(ref_output) + 1e-8)).max().item() - - print(f"\nResults:") - print(f" Reference output stats: min={ref_output.min().item():.4f}, max={ref_output.max().item():.4f}, mean={ref_output.mean().item():.4f}") - print(f" Quantized output stats: min={quant_output.min().item():.4f}, max={quant_output.max().item():.4f}, mean={quant_output.mean().item():.4f}") - print(f" Max absolute difference: {max_diff:.6f}") - print(f" Mean absolute difference: {mean_diff:.6f}") - print(f" Max relative difference: {relative_diff:.6f}") - - # Check that outputs are finite - assert torch.isfinite(ref_output).all(), "Reference output contains NaN/Inf" - assert torch.isfinite(quant_output).all(), "Quantized output contains NaN/Inf" - - print("✓ Basic test passed!") - return max_diff, mean_diff, relative_diff - - -def test_attention_fake_quant_causal(): - """Test causal attention with fake quantization.""" - print("\n" + "="*60) - print("Test 2: Causal Attention with Fake Quantization") - print("="*60) - - if not torch.cuda.is_available(): - print("CUDA not available, skipping test") - return - - device = "cuda" - dtype = torch.bfloat16 - - batch_size = 1 - num_heads = 2 - seq_len = 64 - head_dim = 32 - - torch.manual_seed(42) - q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - v = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - - print(f"Input shapes: Q={q.shape}, K={k.shape}, V={v.shape}") - - # Reference causal attention - print("\nComputing reference causal attention...") - with torch.no_grad(): - ref_output = F.scaled_dot_product_attention(q, k, v, is_causal=True) - - # Causal attention with fake quantization - print("\nComputing causal attention with fake quantization...") - with torch.no_grad(): - quant_output = attention_with_fake_quant(q, k, v, causal=True, quantize_qkv=True) - - # Compare results - diff = torch.abs(ref_output - quant_output) - max_diff = diff.max().item() - mean_diff = diff.mean().item() - - print(f"\nResults:") - print(f" Max absolute difference: {max_diff:.6f}") - print(f" Mean absolute difference: {mean_diff:.6f}") - - assert torch.isfinite(ref_output).all() - assert torch.isfinite(quant_output).all() - - print("✓ Causal attention test passed!") - return max_diff, mean_diff - - -def test_attention_fake_quant_different_shapes(): - """Test attention with different tensor shapes.""" - print("\n" + "="*60) - print("Test 3: Different Shapes") - print("="*60) - - if not torch.cuda.is_available(): - print("CUDA not available, skipping test") - return - - device = "cuda" - dtype = torch.bfloat16 - - test_configs = [ - (1, 1, 32, 16, "small"), - (2, 4, 128, 64, "medium"), - (1, 8, 256, 128, "large_seq"), - (2, 4, 64, 256, "large_head_dim"), - ] - - results = [] - - for batch_size, num_heads, seq_len, head_dim, name in test_configs: - print(f"\nTesting {name} configuration: B={batch_size}, H={num_heads}, L={seq_len}, D={head_dim}") - - torch.manual_seed(42) - q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - v = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - - with torch.no_grad(): - ref_output = F.scaled_dot_product_attention(q, k, v, is_causal=False) - quant_output = attention_with_fake_quant(q, k, v, causal=False, quantize_qkv=True) - - diff = torch.abs(ref_output - quant_output) - max_diff = diff.max().item() - mean_diff = diff.mean().item() - - print(f" Max diff: {max_diff:.6f}, Mean diff: {mean_diff:.6f}") - - assert torch.isfinite(ref_output).all() - assert torch.isfinite(quant_output).all() - assert ref_output.shape == quant_output.shape - - results.append((name, max_diff, mean_diff)) - - print("\n✓ All shape tests passed!") - return results - - -def test_attention_fake_quant_gradient(): - """Test that gradients can flow through fake-quantized attention.""" - print("\n" + "="*60) - print("Test 4: Gradient Flow") - print("="*60) - - if not torch.cuda.is_available(): - print("CUDA not available, skipping test") - return - - device = "cuda" - dtype = torch.bfloat16 - - batch_size = 1 - num_heads = 2 - seq_len = 32 - head_dim = 32 - - torch.manual_seed(42) - q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype, requires_grad=True) - k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype, requires_grad=True) - v = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype, requires_grad=True) - - print(f"Input shapes: Q={q.shape}, K={k.shape}, V={v.shape}") - print("Note: get_fake_quant uses no_grad internally, so gradients won't flow through quantization") - - # Forward pass with fake quantization - # Note: get_fake_quant uses no_grad, so we'll test that it doesn't crash - with torch.no_grad(): - q_quant = get_fake_quant(q) - k_quant = get_fake_quant(k) - v_quant = get_fake_quant(v) - - # Now compute attention with quantized tensors (requires_grad=False) - q_quant.requires_grad_(True) - k_quant.requires_grad_(True) - v_quant.requires_grad_(True) - - output = F.scaled_dot_product_attention(q_quant, k_quant, v_quant, is_causal=False) - - # Backward pass - loss = output.sum() - loss.backward() - - # Check that gradients exist - assert q_quant.grad is not None - assert k_quant.grad is not None - assert v_quant.grad is not None - - print("✓ Gradient flow test passed!") - print(f" Gradient shapes: Q={q_quant.grad.shape}, K={k_quant.grad.shape}, V={v_quant.grad.shape}") - - -def test_attention_fake_quant_precision(): - """Test precision comparison between quantized and non-quantized attention.""" - print("\n" + "="*60) - print("Test 5: Precision Analysis") - print("="*60) - - if not torch.cuda.is_available(): - print("CUDA not available, skipping test") - return - - device = "cuda" - dtype = torch.bfloat16 - - batch_size = 2 - num_heads = 4 - seq_len = 128 - head_dim = 64 - - torch.manual_seed(42) - q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - v = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) - - print(f"Input shapes: Q={q.shape}, K={k.shape}, V={v.shape}") - - # Compute reference - with torch.no_grad(): - ref_output = F.scaled_dot_product_attention(q, k, v, is_causal=False) - - # Compute with quantization - with torch.no_grad(): - quant_output = attention_with_fake_quant(q, k, v, causal=False, quantize_qkv=True) - - # Detailed analysis - diff = torch.abs(ref_output - quant_output) - - print("\nPrecision Analysis:") - print(f" Output shape: {ref_output.shape}") - print(f" Total elements: {ref_output.numel()}") - print(f" Max absolute error: {diff.max().item():.6f}") - print(f" Mean absolute error: {diff.mean().item():.6f}") - print(f" Median absolute error: {diff.median().item():.6f}") - print(f" Std of errors: {diff.std().item():.6f}") - - # Relative error - relative_diff = diff / (torch.abs(ref_output) + 1e-8) - print(f" Max relative error: {relative_diff.max().item():.6f}") - print(f" Mean relative error: {relative_diff.mean().item():.6f}") - - # Error distribution (convert to float32 for quantile computation) - error_percentiles = torch.quantile(diff.float(), torch.tensor([0.5, 0.9, 0.95, 0.99], device=device, dtype=torch.float32)) - print(f" Error percentiles (50th, 90th, 95th, 99th): {error_percentiles.cpu().tolist()}") - - print("\n✓ Precision analysis completed!") - - -def main(): - """Run all tests.""" - global VERBOSE - - # Check for verbose flag - VERBOSE = "--verbose" in sys.argv or "-v" in sys.argv - - print("="*60) - print("Fake Quantization Attention Test Suite") - print("="*60) - - if not torch.cuda.is_available(): - print("ERROR: CUDA is not available. These tests require CUDA.") - sys.exit(1) - - print(f"CUDA Device: {torch.cuda.get_device_name()}") - print(f"PyTorch version: {torch.__version__}") - if VERBOSE: - print("Verbose mode: ON (will show quantization details)") - print() - - try: - # Run tests - test_attention_fake_quant_basic() - test_attention_fake_quant_causal() - test_attention_fake_quant_different_shapes() - test_attention_fake_quant_gradient() - test_attention_fake_quant_precision() - - print("\n" + "="*60) - print("All tests passed! ✓") - print("="*60) - - except Exception as e: - print(f"\n❌ Test failed with error: {e}") - import traceback - traceback.print_exc() - sys.exit(1) - - -if __name__ == "__main__": - main() - - diff --git a/test_fake_quant_basic.py b/test_fake_quant_basic.py deleted file mode 100644 index c29a45b3b..000000000 --- a/test_fake_quant_basic.py +++ /dev/null @@ -1,883 +0,0 @@ -#!/usr/bin/env python3 -""" -Basic tests for the fake_quantize function. -Tests core functionality with simple inputs and edge cases. -""" - -import torch -import triton -import triton.language as tl -from typing import Tuple -from qat_attn import fake_quantize -from nvfp4_utils import _downcast_to_mxfp, _upcast_from_mxfp, MXFP_BLOCK_SIZE, _compute_quant_and_scale, _compute_dequant - -# Import get_fake_quant from test_quantization_precision -try: - from test_quantization_precision import get_fake_quant -except ImportError: - # Fallback: define get_fake_quant here if import fails - from flashinfer import SfLayout, nvfp4_quantize, e2m1_and_ufp8sf_scale_to_float - def get_fake_quant(x: torch.Tensor): - orig_shape = x.shape - orig_dtype = x.dtype - device = x.device - x = x.view(-1, x.shape[-1]) - x_global_sf = (448 * 6) / x.float().abs().nan_to_num().max() - x_fp4, x_scale = nvfp4_quantize(x, x_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) - x_dequant = e2m1_and_ufp8sf_scale_to_float(x_fp4, x_scale, 1 / x_global_sf) - val = x_dequant.view(orig_shape).to(orig_dtype).to(device) - return val - -DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - -def cosine_similarity(tensor1, tensor2): - """ - Compute cosine similarity between two tensors. - - Args: - tensor1: First tensor - tensor2: Second tensor (same shape as tensor1) - - Returns: - Cosine similarity value (scalar) - """ - # Flatten tensors for computation - 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) - - # Avoid division by zero - if norm1 == 0 or norm2 == 0: - return 0.0 - - cos_sim = dot_product / (norm1 * norm2) - return cos_sim.item() - - -@triton.jit -def fake_quantize_kernel( - src_ptr, - dst_ptr, - high_prec_ptr, - BLOCK_SIZE_OUT_DIM: tl.constexpr, - BLOCK_SIZE_QUANT_DIM: tl.constexpr, - dst_dtype: tl.constexpr, - mx_tensor_dtype: tl.constexpr, - stride_src_outer, - stride_src_quant, - stride_dst_outer, - stride_dst_quant, - stride_high_prec_outer, - stride_high_prec_quant, - outer_dim, - quant_dim, -): - """ - Kernel wrapper to test fake_quantize function. - """ - outer_idx = tl.program_id(0) - quant_idx = tl.program_id(1) - - # Compute offsets - start_outer = outer_idx * BLOCK_SIZE_OUT_DIM - start_quant = quant_idx * BLOCK_SIZE_QUANT_DIM - - # Create offset arrays - offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None] - offs_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :] - - # Create masks for valid elements - mask_outer = (start_outer + offs_outer) < outer_dim - mask_quant = (start_quant + offs_quant) < quant_dim - full_mask = mask_outer & mask_quant - - # Load source tensor - src_offsets = (start_outer + offs_outer) * stride_src_outer + (start_quant + offs_quant) * stride_src_quant - src_tensor = tl.load(src_ptr + src_offsets, mask=full_mask, other=0.0) - - # Call fake_quantize - quantized_tensor, high_prec_tensor = fake_quantize( - src_tensor=src_tensor, - BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, - BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, - dst_dtype=dst_dtype, - mx_tensor_dtype=mx_tensor_dtype - ) - - # Store results - dst_offsets = (start_outer + offs_outer) * stride_dst_outer + (start_quant + offs_quant) * stride_dst_quant - high_prec_offsets = (start_outer + offs_outer) * stride_high_prec_outer + (start_quant + offs_quant) * stride_high_prec_quant - - tl.store(dst_ptr + dst_offsets, quantized_tensor, mask=full_mask) - tl.store(high_prec_ptr + high_prec_offsets, high_prec_tensor, mask=full_mask) - - -def call_fake_quantize( - x: torch.Tensor, - BLOCK_SIZE_OUT_DIM: int = 128, - BLOCK_SIZE_QUANT_DIM: int = 128, - dst_dtype: torch.dtype = None, - mx_tensor_dtype = tl.float8e4nv -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Call fake_quantize on a tensor. - - Args: - x: Input tensor (2D or can be reshaped to 2D) - BLOCK_SIZE_OUT_DIM: Block size for outer dimension - BLOCK_SIZE_QUANT_DIM: Block size for quantization dimension (must be multiple of 16) - dst_dtype: Output dtype (defaults to input dtype) - mx_tensor_dtype: Quantization dtype (tl.float8e4nv or tl.uint8 for fp4) - - Returns: - Tuple of (quantized_tensor, high_precision_tensor) - """ - assert x.is_cuda, "Input must be on CUDA" - assert BLOCK_SIZE_QUANT_DIM % 16 == 0, f"BLOCK_SIZE_QUANT_DIM must be multiple of 16, got {BLOCK_SIZE_QUANT_DIM}" - - orig_shape = x.shape - orig_dtype = x.dtype - - # Reshape to 2D - x_2d = x.view(-1, x.shape[-1]) - outer_dim, quant_dim = x_2d.shape - - if dst_dtype is None: - dst_dtype = orig_dtype - - # Map torch dtype to triton dtype - dtype_map = { - torch.float32: tl.float32, - torch.float16: tl.float16, - torch.bfloat16: tl.bfloat16, - } - triton_dst_dtype = dtype_map.get(dst_dtype, tl.float16) - - # Allocate outputs - quantized_output = torch.empty_like(x_2d, dtype=dst_dtype) - high_prec_output = torch.empty_like(x_2d, dtype=torch.float32) - - # Launch kernel - grid = ( - triton.cdiv(outer_dim, BLOCK_SIZE_OUT_DIM), - triton.cdiv(quant_dim, BLOCK_SIZE_QUANT_DIM), - ) - - fake_quantize_kernel[grid]( - src_ptr=x_2d, - dst_ptr=quantized_output, - high_prec_ptr=high_prec_output, - BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, - BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, - dst_dtype=triton_dst_dtype, - mx_tensor_dtype=mx_tensor_dtype, - stride_src_outer=x_2d.stride(0), - stride_src_quant=x_2d.stride(1), - stride_dst_outer=quantized_output.stride(0), - stride_dst_quant=quantized_output.stride(1), - stride_high_prec_outer=high_prec_output.stride(0), - stride_high_prec_quant=high_prec_output.stride(1), - outer_dim=outer_dim, - quant_dim=quant_dim, - ) - - return quantized_output.view(orig_shape).to(dst_dtype), high_prec_output.view(orig_shape) - - -def test_basic_functionality(): - """Test basic functionality of fake_quantize.""" - print("Test 1: Basic functionality") - - torch.manual_seed(42) - x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) - - quantized, high_prec = call_fake_quantize(x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - - # Check shapes - assert quantized.shape == x.shape, f"Shape mismatch: {quantized.shape} vs {x.shape}" - assert high_prec.shape == x.shape, f"Shape mismatch: {high_prec.shape} vs {x.shape}" - - # Check that quantization happened (values should be different) - diff = (quantized.float() - x.float()).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - assert max_diff > 0, "Quantization should change values" - - # Compute cosine similarity between quantized and original - cos_sim = cosine_similarity(quantized, x) - - # Check that high_prec matches original - high_prec_diff = (high_prec - x.float()).abs().max().item() - assert high_prec_diff < 1e-5, f"High precision tensor should match original, got diff: {high_prec_diff}" - - print(f" ✓ Shape check passed") - print(f" ✓ Quantization occurred (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") - print(f" ✓ High precision tensor matches original (diff: {high_prec_diff:.6e})") - - -def test_different_shapes(): - """Test fake_quantize with different tensor shapes.""" - print("\nTest 2: Different shapes") - - torch.manual_seed(42) - test_shapes = [ - (64, 64), - (128, 128), - (256, 128), - (128, 256), - (32, 64), - ] - - for shape in test_shapes: - x = torch.randn(shape, device=DEVICE, dtype=torch.float16) - - # Use block sizes that divide evenly - block_out = min(128, shape[0]) - block_quant = min(128, shape[1]) - - # Ensure block_quant is multiple of 16 - block_quant = (block_quant // 16) * 16 - if block_quant == 0: - block_quant = 16 - - quantized, high_prec = call_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=block_out, - BLOCK_SIZE_QUANT_DIM=block_quant - ) - - assert quantized.shape == x.shape, f"Shape mismatch for {shape}" - assert high_prec.shape == x.shape, f"Shape mismatch for {shape}" - - print(f" ✓ Shape {shape} passed") - - -def test_different_dtypes(): - """Test fake_quantize with different input dtypes.""" - print("\nTest 3: Different dtypes") - - torch.manual_seed(42) - shape = (128, 64) - dtypes = [torch.float16, torch.bfloat16, torch.float32] - - for dtype in dtypes: - x = torch.randn(shape, device=DEVICE, dtype=dtype) - - quantized, high_prec = call_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=128, - BLOCK_SIZE_QUANT_DIM=64, - dst_dtype=dtype - ) - - assert quantized.dtype == dtype, f"Dtype mismatch: {quantized.dtype} vs {dtype}" - assert quantized.shape == x.shape, f"Shape mismatch for dtype {dtype}" - - print(f" ✓ Dtype {dtype} passed") - - -def test_edge_cases(): - """Test fake_quantize with edge cases.""" - print("\nTest 4: Edge cases") - - # Test with zeros - print(" Testing zeros...") - x_zeros = torch.zeros(128, 64, device=DEVICE, dtype=torch.float16) - quantized, high_prec = call_fake_quantize(x_zeros, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - assert quantized.shape == x_zeros.shape - # Quantized zeros should still be close to zero - assert quantized.abs().max().item() < 1e-3, "Quantized zeros should be close to zero" - print(" ✓ Zeros passed") - - # Test with ones (FAIL) - print(" Testing ones...") - x_ones = torch.ones(128, 64, device=DEVICE, dtype=torch.float16) - print(f"x_ones: {x_ones}") - quantized, high_prec = call_fake_quantize(x_ones, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - print(f"quantized: {quantized}") - print(f"high_prec: {high_prec}") - assert quantized.shape == x_ones.shape - # Quantized ones should be close to 1.0 - diff = (quantized - x_ones).abs().max().item() - assert diff < 0.1, f"Quantized ones should be close to 1.0, got max diff: {diff}" - print(f" ✓ Ones passed (max diff: {diff:.6f})") - - # Test with small values - print(" Testing small values...") - x_small = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) * 0.001 - quantized, high_prec = call_fake_quantize(x_small, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - assert quantized.shape == x_small.shape - print(" ✓ Small values passed") - - # Test with large values - print(" Testing large values...") - x_large = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) * 100.0 - quantized, high_prec = call_fake_quantize(x_large, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - assert quantized.shape == x_large.shape - print(" ✓ Large values passed") - - -def test_fp4_quantization(): - """Test fake_quantize with FP4 quantization (uint8).""" - print("\nTest 5: FP4 quantization (uint8)") - - torch.manual_seed(42) - x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) - - quantized, high_prec = call_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=128, - BLOCK_SIZE_QUANT_DIM=64, - mx_tensor_dtype=tl.uint8 # FP4 uses uint8 - ) - - assert quantized.shape == x.shape - assert high_prec.shape == x.shape - - # Check that quantization happened - diff = (quantized.float() - x.float()).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - assert max_diff > 0, "Quantization should change values" - - # Compute cosine similarity between quantized and original - cos_sim = cosine_similarity(quantized, x) - - print(f" ✓ FP4 quantization passed (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") - - - -def test_finite_values(): - """Test that fake_quantize produces finite values.""" - print("\nTest 7: Finite values check") - - torch.manual_seed(42) - x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) - - quantized, high_prec = call_fake_quantize(x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - - # Check for NaN and Inf - assert not torch.isnan(quantized).any(), "Quantized tensor contains NaN" - assert not torch.isinf(quantized).any(), "Quantized tensor contains Inf" - assert not torch.isnan(high_prec).any(), "High precision tensor contains NaN" - assert not torch.isinf(high_prec).any(), "High precision tensor contains Inf" - - print(" ✓ All values are finite") - - -def test_different_block_sizes(): - """Test fake_quantize with different block sizes.""" - print("\nTest 8: Different block sizes") - - torch.manual_seed(42) - x = torch.randn(256, 128, device=DEVICE, dtype=torch.float16) - - block_configs = [ - (64, 64), - (128, 64), - (64, 128), - (128, 128), - ] - - for block_out, block_quant in block_configs: - quantized, high_prec = call_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=block_out, - BLOCK_SIZE_QUANT_DIM=block_quant - ) - - assert quantized.shape == x.shape - assert high_prec.shape == x.shape - - print(f" ✓ Block size ({block_out}, {block_quant}) passed") - - -# ============================================================================ -# MXFP Kernel Tests -# ============================================================================ - -@triton.jit -def mxfp_fake_quantize_kernel( - src_ptr, - dst_ptr, - high_prec_ptr, - BLOCK_SIZE_OUT_DIM: tl.constexpr, - BLOCK_SIZE_QUANT_DIM: tl.constexpr, - dst_dtype: tl.constexpr, - mx_tensor_dtype: tl.constexpr, - stride_src_outer, - stride_src_quant, - stride_dst_outer, - stride_dst_quant, - stride_high_prec_outer, - stride_high_prec_quant, - outer_dim, - quant_dim, -): - """ - Kernel wrapper to test MXFP fake_quantize using _downcast_to_mxfp and _upcast_from_mxfp. - """ - outer_idx = tl.program_id(0) - quant_idx = tl.program_id(1) - - # Determine if we are dealing with fp4 types - is_fp4: tl.constexpr = mx_tensor_dtype == tl.uint8 - K_DIVISOR: tl.constexpr = 2 if is_fp4 else 1 - BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // MXFP_BLOCK_SIZE - BLOCK_SIZE_QUANT_MX_TENSOR: tl.constexpr = BLOCK_SIZE_QUANT_DIM // K_DIVISOR - - # Allocate temporary storage for quantized tensor and scale - # These need to be stored in shared memory or we need to allocate them - # For simplicity, we'll compute quantize and dequantize in the same kernel - - # Compute offsets - start_outer = outer_idx * BLOCK_SIZE_OUT_DIM - start_quant = quant_idx * BLOCK_SIZE_QUANT_DIM - start_mx_scale_quant = quant_idx * BLOCK_SIZE_QUANT_MX_SCALE - start_mx_quant = quant_idx * BLOCK_SIZE_QUANT_MX_TENSOR - - # Create offset arrays - offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None] - offs_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :] - offs_mx_scale_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_SCALE)[None, :] - offs_mx_quant = tl.arange(0, BLOCK_SIZE_QUANT_MX_TENSOR)[None, :] - - # Create masks for valid elements - mask_outer = (start_outer + offs_outer) < outer_dim - mask_quant = (start_quant + offs_quant) < quant_dim - full_mask = mask_outer & mask_quant - - # Load source tensor - src_offsets = (start_outer + offs_outer) * stride_src_outer + (start_quant + offs_quant) * stride_src_quant - src_tensor = tl.load(src_ptr + src_offsets, mask=full_mask, other=0.0) - - # Store high precision tensor (original) - high_prec_offsets = (start_outer + offs_outer) * stride_high_prec_outer + (start_quant + offs_quant) * stride_high_prec_quant - tl.store(high_prec_ptr + high_prec_offsets, src_tensor.to(tl.float32), mask=full_mask) - - # Quantize using _compute_quant_and_scale (same as fake_quantize uses internally) - valid_mask = tl.full(src_tensor.shape, 1, tl.int1) - quantized_tensor, scale_tensor, s_dec = _compute_quant_and_scale( - src_tensor=src_tensor, - valid_src_mask=valid_mask, - mx_tensor_dtype=mx_tensor_dtype - ) - - # Dequantize using _compute_dequant (same as fake_quantize uses internally) - dequantized_tensor = _compute_dequant( - mx_tensor=quantized_tensor, - scale=scale_tensor, - s_dec=s_dec, - BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, - BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, - dst_dtype=dst_dtype - ) - - # Store result - dst_offsets = (start_outer + offs_outer) * stride_dst_outer + (start_quant + offs_quant) * stride_dst_quant - tl.store(dst_ptr + dst_offsets, dequantized_tensor, mask=full_mask) - - -def call_mxfp_fake_quantize( - x: torch.Tensor, - BLOCK_SIZE_OUT_DIM: int = 128, - BLOCK_SIZE_QUANT_DIM: int = 128, - dst_dtype: torch.dtype = None, - mx_tensor_dtype = tl.uint8 # Default to FP4 (uint8) -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Call MXFP fake_quantize on a tensor using _downcast_to_mxfp and _upcast_from_mxfp logic. - - Args: - x: Input tensor (2D or can be reshaped to 2D) - BLOCK_SIZE_OUT_DIM: Block size for outer dimension - BLOCK_SIZE_QUANT_DIM: Block size for quantization dimension (must be multiple of 16) - dst_dtype: Output dtype (defaults to input dtype) - mx_tensor_dtype: Quantization dtype (tl.uint8 for fp4 or tl.float8e4nv for fp8) - - Returns: - Tuple of (quantized_tensor, high_precision_tensor) - """ - assert x.is_cuda, "Input must be on CUDA" - assert BLOCK_SIZE_QUANT_DIM % 16 == 0, f"BLOCK_SIZE_QUANT_DIM must be multiple of 16, got {BLOCK_SIZE_QUANT_DIM}" - - orig_shape = x.shape - orig_dtype = x.dtype - - # Reshape to 2D - x_2d = x.view(-1, x.shape[-1]) - outer_dim, quant_dim = x_2d.shape - - if dst_dtype is None: - dst_dtype = orig_dtype - - # Map torch dtype to triton dtype - dtype_map = { - torch.float32: tl.float32, - torch.float16: tl.float16, - torch.bfloat16: tl.bfloat16, - } - triton_dst_dtype = dtype_map.get(dst_dtype, tl.float16) - - # Allocate outputs - quantized_output = torch.empty_like(x_2d, dtype=dst_dtype) - high_prec_output = torch.empty_like(x_2d, dtype=torch.float32) - - # Launch kernel - grid = ( - triton.cdiv(outer_dim, BLOCK_SIZE_OUT_DIM), - triton.cdiv(quant_dim, BLOCK_SIZE_QUANT_DIM), - ) - - mxfp_fake_quantize_kernel[grid]( - src_ptr=x_2d, - dst_ptr=quantized_output, - high_prec_ptr=high_prec_output, - BLOCK_SIZE_OUT_DIM=BLOCK_SIZE_OUT_DIM, - BLOCK_SIZE_QUANT_DIM=BLOCK_SIZE_QUANT_DIM, - dst_dtype=triton_dst_dtype, - mx_tensor_dtype=mx_tensor_dtype, - stride_src_outer=x_2d.stride(0), - stride_src_quant=x_2d.stride(1), - stride_dst_outer=quantized_output.stride(0), - stride_dst_quant=quantized_output.stride(1), - stride_high_prec_outer=high_prec_output.stride(0), - stride_high_prec_quant=high_prec_output.stride(1), - outer_dim=outer_dim, - quant_dim=quant_dim, - ) - - return quantized_output.view(orig_shape).to(dst_dtype), high_prec_output.view(orig_shape) - - -def test_mxfp_basic_functionality(): - """Test basic functionality of MXFP fake_quantize.""" - print("\n" + "=" * 60) - print("MXFP Kernel Tests") - print("=" * 60) - print("\nTest 1: Basic functionality (MXFP)") - - torch.manual_seed(42) - x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) - - quantized, high_prec = call_mxfp_fake_quantize(x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - - # Check shapes - assert quantized.shape == x.shape, f"Shape mismatch: {quantized.shape} vs {x.shape}" - assert high_prec.shape == x.shape, f"Shape mismatch: {high_prec.shape} vs {x.shape}" - - # Check that quantization happened (values should be different) - diff = (quantized.float() - x.float()).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - assert max_diff > 0, "Quantization should change values" - - # Compute cosine similarity between quantized and original - cos_sim = cosine_similarity(quantized, x) - - # Check that high_prec matches original - high_prec_diff = (high_prec - x.float()).abs().max().item() - assert high_prec_diff < 1e-5, f"High precision tensor should match original, got diff: {high_prec_diff}" - - print(f" ✓ Shape check passed") - print(f" ✓ Quantization occurred (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") - print(f" ✓ High precision tensor matches original (diff: {high_prec_diff:.6e})") - - -def test_mxfp_different_shapes(): - """Test MXFP fake_quantize with different tensor shapes.""" - print("\nTest 2: Different shapes (MXFP)") - - torch.manual_seed(42) - test_shapes = [ - (64, 64), - (128, 128), - (256, 128), - (128, 256), - (32, 64), - ] - - for shape in test_shapes: - x = torch.randn(shape, device=DEVICE, dtype=torch.float16) - - # Use block sizes that divide evenly - block_out = min(128, shape[0]) - block_quant = min(128, shape[1]) - - # Ensure block_quant is multiple of 16 - block_quant = (block_quant // 16) * 16 - if block_quant == 0: - block_quant = 16 - - quantized, high_prec = call_mxfp_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=block_out, - BLOCK_SIZE_QUANT_DIM=block_quant - ) - - assert quantized.shape == x.shape, f"Shape mismatch for {shape}" - assert high_prec.shape == x.shape, f"Shape mismatch for {shape}" - - print(f" ✓ Shape {shape} passed") - - -def test_mxfp_different_dtypes(): - """Test MXFP fake_quantize with different input dtypes.""" - print("\nTest 3: Different dtypes (MXFP)") - - torch.manual_seed(42) - shape = (128, 64) - dtypes = [torch.float16, torch.bfloat16, torch.float32] - - for dtype in dtypes: - x = torch.randn(shape, device=DEVICE, dtype=dtype) - - quantized, high_prec = call_mxfp_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=128, - BLOCK_SIZE_QUANT_DIM=64, - dst_dtype=dtype - ) - - assert quantized.dtype == dtype, f"Dtype mismatch: {quantized.dtype} vs {dtype}" - assert quantized.shape == x.shape, f"Shape mismatch for dtype {dtype}" - - print(f" ✓ Dtype {dtype} passed") - - -def test_mxfp_edge_cases(): - """Test MXFP fake_quantize with edge cases.""" - print("\nTest 4: Edge cases (MXFP)") - - # Test with zeros - print(" Testing zeros...") - x_zeros = torch.zeros(128, 64, device=DEVICE, dtype=torch.float16) - quantized, high_prec = call_mxfp_fake_quantize(x_zeros, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - assert quantized.shape == x_zeros.shape - # Quantized zeros should still be close to zero - assert quantized.abs().max().item() < 1e-3, "Quantized zeros should be close to zero" - print(" ✓ Zeros passed") - - # Test with ones - print(" Testing ones...") - x_ones = torch.ones(128, 64, device=DEVICE, dtype=torch.float16) - quantized, high_prec = call_mxfp_fake_quantize(x_ones, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - assert quantized.shape == x_ones.shape - # Quantized ones should be close to 1.0 - diff = (quantized - x_ones).abs().max().item() - assert diff < 0.1, f"Quantized ones should be close to 1.0, got max diff: {diff}" - print(f" ✓ Ones passed (max diff: {diff:.6f})") - - # Test with small values - print(" Testing small values...") - x_small = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) * 0.001 - quantized, high_prec = call_mxfp_fake_quantize(x_small, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - assert quantized.shape == x_small.shape - print(" ✓ Small values passed") - - # Test with large values - print(" Testing large values...") - x_large = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) * 100.0 - quantized, high_prec = call_mxfp_fake_quantize(x_large, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - assert quantized.shape == x_large.shape - print(" ✓ Large values passed") - - -def test_mxfp_fp4_quantization(): - """Test MXFP fake_quantize with FP4 quantization (uint8).""" - print("\nTest 5: FP4 quantization (MXFP, uint8)") - - torch.manual_seed(42) - x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) - - quantized, high_prec = call_mxfp_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=128, - BLOCK_SIZE_QUANT_DIM=64, - mx_tensor_dtype=tl.uint8 # FP4 uses uint8 - ) - - assert quantized.shape == x.shape - assert high_prec.shape == x.shape - - # Check that quantization happened - diff = (quantized.float() - x.float()).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - assert max_diff > 0, "Quantization should change values" - - # Compute cosine similarity between quantized and original - cos_sim = cosine_similarity(quantized, x) - - print(f" ✓ FP4 quantization passed (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") - - -def test_mxfp_fp8_quantization(): - """Test MXFP fake_quantize with FP8 quantization (float8e4nv).""" - print("\nTest 6: FP8 quantization (MXFP, float8e4nv)") - - torch.manual_seed(42) - x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) - - quantized, high_prec = call_mxfp_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=128, - BLOCK_SIZE_QUANT_DIM=64, - mx_tensor_dtype=tl.float8e4nv # FP8 uses float8e4nv - ) - - assert quantized.shape == x.shape - assert high_prec.shape == x.shape - - # Check that quantization happened - diff = (quantized.float() - x.float()).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - assert max_diff > 0, "Quantization should change values" - - # Compute cosine similarity between quantized and original - cos_sim = cosine_similarity(quantized, x) - - print(f" ✓ FP8 quantization passed (max diff: {max_diff:.6f}, mean diff: {mean_diff:.6f}, cosine sim: {cos_sim:.6f})") - - -def test_mxfp_finite_values(): - """Test that MXFP fake_quantize produces finite values.""" - print("\nTest 7: Finite values check (MXFP)") - - torch.manual_seed(42) - x = torch.randn(128, 64, device=DEVICE, dtype=torch.float16) - - quantized, high_prec = call_mxfp_fake_quantize(x, BLOCK_SIZE_OUT_DIM=128, BLOCK_SIZE_QUANT_DIM=64) - - # Check for NaN and Inf - assert not torch.isnan(quantized).any(), "Quantized tensor contains NaN" - assert not torch.isinf(quantized).any(), "Quantized tensor contains Inf" - assert not torch.isnan(high_prec).any(), "High precision tensor contains NaN" - assert not torch.isinf(high_prec).any(), "High precision tensor contains Inf" - - print(" ✓ All values are finite") - - -def test_mxfp_different_block_sizes(): - """Test MXFP fake_quantize with different block sizes.""" - print("\nTest 8: Different block sizes (MXFP)") - - torch.manual_seed(42) - x = torch.randn(256, 128, device=DEVICE, dtype=torch.float16) - - block_configs = [ - (64, 64), - (128, 64), - (64, 128), - (128, 128), - ] - - for block_out, block_quant in block_configs: - quantized, high_prec = call_mxfp_fake_quantize( - x, - BLOCK_SIZE_OUT_DIM=block_out, - BLOCK_SIZE_QUANT_DIM=block_quant - ) - - assert quantized.shape == x.shape - assert high_prec.shape == x.shape - - print(f" ✓ Block size ({block_out}, {block_quant}) passed") - - -def test_get_fake_quant_metrics(): - """Test get_fake_quant function with cosine similarity, mean_diff, and max_diff metrics.""" - print("\n" + "=" * 60) - print("Testing get_fake_quant with cosine similarity, mean_diff, and max_diff") - print("=" * 60) - - if not torch.cuda.is_available(): - print("CUDA not available, skipping test") - return - - device = torch.device("cuda") - torch.manual_seed(42) - - # Test with different shapes and dtypes - test_cases = [ - ((128, 64), torch.bfloat16, "2D bfloat16"), - ((256, 128), torch.bfloat16, "2D bfloat16 large"), - ((2, 8, 128), torch.bfloat16, "3D bfloat16"), - ((128, 64), torch.float16, "2D float16"), - ((1, 8, 256, 128), torch.bfloat16, "4D bfloat16"), - ] - - for shape, dtype, name in test_cases: - print(f"\nTest: {name} - Shape: {shape}") - x = torch.randn(*shape, device=device, dtype=dtype) - x_fq = get_fake_quant(x) - - # Check shapes match - assert x_fq.shape == x.shape, f"Shape mismatch: {x_fq.shape} vs {x.shape}" - assert x_fq.dtype == x.dtype, f"Dtype mismatch: {x_fq.dtype} vs {x.dtype}" - - # Compute metrics - diff = (x_fq.float() - x.float()).abs() - max_diff = diff.max().item() - mean_diff = diff.mean().item() - cos_sim = cosine_similarity(x_fq, x) - - print(f" Max diff: {max_diff:.6f}") - print(f" Mean diff: {mean_diff:.6f}") - print(f" Cosine similarity: {cos_sim:.6f}") - - # Basic sanity checks - assert max_diff < 10.0, f"Max difference too large: {max_diff}" - assert cos_sim > 0.5, f"Cosine similarity too low: {cos_sim}" - - print("\n✓ All metric tests passed!") - - -if __name__ == "__main__": - if not torch.cuda.is_available(): - print("CUDA not available. These tests require CUDA.") - exit(1) - - print("=" * 60) - print("Running basic tests for fake_quantize function") - print("=" * 60) - print(f"Device: {DEVICE}") - print() - - try: - test_basic_functionality() - test_different_shapes() - test_different_dtypes() - test_edge_cases() - test_fp4_quantization() - test_finite_values() - - # Test get_fake_quant with metrics - test_get_fake_quant_metrics() - - # Run MXFP kernel tests - # test_mxfp_basic_functionality() - # test_mxfp_different_shapes() - # test_mxfp_different_dtypes() - # test_mxfp_edge_cases() - # test_mxfp_fp4_quantization() - # test_mxfp_fp8_quantization() - # test_mxfp_finite_values() - # test_mxfp_different_block_sizes() - - print("\n" + "=" * 60) - print("All tests passed! ✓") - print("=" * 60) - except Exception as e: - print(f"\n❌ Test failed with error: {e}") - import traceback - traceback.print_exc() - exit(1) - - diff --git a/test_fake_quantize_kernels.py b/test_fake_quantize_kernels.py deleted file mode 100644 index 01f991e14..000000000 --- a/test_fake_quantize_kernels.py +++ /dev/null @@ -1,452 +0,0 @@ -#!/usr/bin/env python3 -""" -Precision test for fake_quantize_q and fake_quantize_kv kernels. - -Tests that the kernels correctly: -1. Quantize and dequantize Q, K, V tensors -2. Preserve approximate values within quantization error -3. Handle edge cases (zeros, small/large values, padding) -4. Work with different shapes and dtypes -""" - -import torch -import triton -import triton.language as tl -import numpy as np -from typing import Tuple, Dict, Any -import unittest - -from quant_utils import fake_quantize_q, fake_quantize_kv - - -def cosine_similarity(tensor1: torch.Tensor, tensor2: torch.Tensor) -> float: - """Compute cosine similarity between two tensors.""" - t1_flat = tensor1.flatten().float() - t2_flat = tensor2.flatten().float() - - dot_product = torch.dot(t1_flat, t2_flat) - norm1 = torch.norm(t1_flat) - norm2 = torch.norm(t2_flat) - - if norm1 == 0 or norm2 == 0: - return 0.0 - - return (dot_product / (norm1 * norm2)).item() - - -def compute_metrics(original: torch.Tensor, quantized: torch.Tensor, mask: torch.Tensor = None) -> Dict[str, float]: - """Compute precision metrics between original and quantized tensors.""" - if mask is not None: - # Only compute metrics on valid (non-padding) elements - original = original[mask] - quantized = quantized[mask] - - diff = (original.float() - quantized.float()).abs() - - metrics = { - 'max_diff': diff.max().item(), - 'mean_diff': diff.mean().item(), - 'std_diff': diff.std().item(), - 'cosine_similarity': cosine_similarity(original, quantized), - } - - # Relative error (avoid division by zero) - abs_original = original.float().abs() - rel_error = diff / (abs_original + 1e-8) - metrics['max_relative_error'] = rel_error.max().item() - metrics['mean_relative_error'] = rel_error.mean().item() - - return metrics - - -def test_fake_quantize_q( - q: torch.Tensor, - BLOCK_M: int = 128, - HEAD_DIM: int = 128, - verbose: bool = False -) -> Tuple[torch.Tensor, Dict[str, Any]]: - """ - Test fake_quantize_q kernel. - - Args: - q: Input Q tensor with shape (Z, H, N_CTX_Q, HEAD_DIM) - BLOCK_M: Block size for sequence dimension - HEAD_DIM: Head dimension (must be multiple of 16) - verbose: Whether to print detailed information - - Returns: - Tuple of (fake_q, metrics_dict) - """ - assert q.is_cuda, "Input must be on CUDA" - assert q.dim() == 4, f"Q must be 4D, got shape {q.shape}" - assert HEAD_DIM % 16 == 0, f"HEAD_DIM must be multiple of 16, got {HEAD_DIM}" - - Z, H, N_CTX_Q, HEAD_DIM_actual = q.shape - assert HEAD_DIM_actual == HEAD_DIM, f"HEAD_DIM mismatch: {HEAD_DIM_actual} != {HEAD_DIM}" - - # Create output tensor - fake_q = torch.empty_like(q) - - # Launch kernel - grid = (triton.cdiv(N_CTX_Q, BLOCK_M), Z * H, 1) - - fake_quantize_q[grid]( - q, fake_q, - q.stride(0), q.stride(1), # stride_z_q, stride_h_q - q.stride(2), q.stride(3), # stride_tok_q, stride_d_q - fake_q.stride(0), fake_q.stride(1), # fake_stride_z_q, fake_stride_h_q - fake_q.stride(2), fake_q.stride(3), # fake_stride_tok_q, fake_stride_d_q - H, N_CTX_Q, - BLOCK_M=BLOCK_M, HEAD_DIM=HEAD_DIM - ) - - # Compute metrics - metrics = compute_metrics(q, fake_q) - metrics['shape'] = list(q.shape) - metrics['dtype'] = str(q.dtype) - - if verbose: - print(f"fake_quantize_q test:") - print(f" Shape: {q.shape}") - print(f" Dtype: {q.dtype}") - print(f" Max diff: {metrics['max_diff']:.6f}") - print(f" Mean diff: {metrics['mean_diff']:.6f}") - print(f" Cosine similarity: {metrics['cosine_similarity']:.6f}") - print(f" Max relative error: {metrics['max_relative_error']:.6f}") - - return fake_q, metrics - - -def test_fake_quantize_kv( - k: torch.Tensor, - v: torch.Tensor, - BLOCK_N: int = 128, - HEAD_DIM: int = 128, - verbose: bool = False -) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]: - """ - Test fake_quantize_kv kernel. - - Args: - k: Input K tensor with shape (Z, H, N_CTX_KV, HEAD_DIM) - v: Input V tensor with shape (Z, H, N_CTX_KV, HEAD_DIM) - BLOCK_N: Block size for sequence dimension - HEAD_DIM: Head dimension (must be multiple of 16) - verbose: Whether to print detailed information - - Returns: - Tuple of (fake_k, fake_v, metrics_dict) - """ - assert k.is_cuda and v.is_cuda, "Inputs must be on CUDA" - assert k.shape == v.shape, f"K and V must have same shape: {k.shape} != {v.shape}" - assert k.dim() == 4, f"K must be 4D, got shape {k.shape}" - assert HEAD_DIM % 16 == 0, f"HEAD_DIM must be multiple of 16, got {HEAD_DIM}" - - Z, H, N_CTX_KV, HEAD_DIM_actual = k.shape - assert HEAD_DIM_actual == HEAD_DIM, f"HEAD_DIM mismatch: {HEAD_DIM_actual} != {HEAD_DIM}" - - # Create output tensors - fake_k = torch.empty_like(k) - fake_v = torch.empty_like(v) - - # Launch kernel - grid = (triton.cdiv(N_CTX_KV, BLOCK_N), Z * H, 1) - - fake_quantize_kv[grid]( - k, v, fake_k, fake_v, - k.stride(0), k.stride(1), # stride_z_kv, stride_h_kv - k.stride(2), k.stride(3), # stride_tok_kv, stride_d_kv - fake_k.stride(0), fake_k.stride(1), # fake_stride_z_kv, fake_stride_h_kv - fake_k.stride(2), fake_k.stride(3), # fake_stride_tok_kv, fake_stride_d_kv - H, N_CTX_KV, - BLOCK_N=BLOCK_N, HEAD_DIM=HEAD_DIM - ) - - # Compute metrics for both K and V - metrics_k = compute_metrics(k, fake_k) - metrics_v = compute_metrics(v, fake_v) - - metrics = { - 'k': metrics_k, - 'v': metrics_v, - 'shape': list(k.shape), - 'dtype': str(k.dtype), - } - - if verbose: - print(f"fake_quantize_kv test:") - print(f" Shape: {k.shape}") - print(f" Dtype: {k.dtype}") - print(f" K - Max diff: {metrics_k['max_diff']:.6f}, Cosine sim: {metrics_k['cosine_similarity']:.6f}") - print(f" V - Max diff: {metrics_v['max_diff']:.6f}, Cosine sim: {metrics_v['cosine_similarity']:.6f}") - - return fake_k, fake_v, metrics - - -class FakeQuantizeKernelsTest: - """Test suite for fake_quantize_q and fake_quantize_kv kernels.""" - - def __init__(self, device: str = "cuda", dtype: torch.dtype = torch.bfloat16): - self.device = device - self.dtype = dtype - self.HEAD_DIM = 128 - self.BLOCK_M = 128 - self.BLOCK_N = 128 - - def test_basic_q(self, verbose: bool = False) -> Dict[str, Any]: - """Test fake_quantize_q with basic random input.""" - print("\n" + "="*60) - print("Test: Basic Q quantization") - print("="*60) - - Z, H, N_CTX_Q = 2, 8, 256 - q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype) - - fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose) - - # Assertions - assert fake_q.shape == q.shape, f"Shape mismatch: {fake_q.shape} != {q.shape}" - assert fake_q.dtype == q.dtype, f"Dtype mismatch: {fake_q.dtype} != {q.dtype}" - assert metrics['cosine_similarity'] > 0.9, f"Cosine similarity too low: {metrics['cosine_similarity']}" - - print("✓ Basic Q test passed") - return metrics - - def test_basic_kv(self, verbose: bool = False) -> Dict[str, Any]: - """Test fake_quantize_kv with basic random inputs.""" - print("\n" + "="*60) - print("Test: Basic KV quantization") - print("="*60) - - Z, H, N_CTX_KV = 2, 8, 512 - k = torch.randn(Z, H, N_CTX_KV, self.HEAD_DIM, device=self.device, dtype=self.dtype) - v = torch.randn(Z, H, N_CTX_KV, self.HEAD_DIM, device=self.device, dtype=self.dtype) - - fake_k, fake_v, metrics = test_fake_quantize_kv( - k, v, BLOCK_N=self.BLOCK_N, HEAD_DIM=self.HEAD_DIM, verbose=verbose - ) - - # Assertions - assert fake_k.shape == k.shape, f"K shape mismatch: {fake_k.shape} != {k.shape}" - assert fake_v.shape == v.shape, f"V shape mismatch: {fake_v.shape} != {v.shape}" - assert metrics['k']['cosine_similarity'] > 0.9, f"K cosine similarity too low: {metrics['k']['cosine_similarity']}" - assert metrics['v']['cosine_similarity'] > 0.9, f"V cosine similarity too low: {metrics['v']['cosine_similarity']}" - - print("✓ Basic KV test passed") - return metrics - - def test_small_values(self, verbose: bool = False) -> Dict[str, Any]: - """Test with small input values.""" - print("\n" + "="*60) - print("Test: Small values") - print("="*60) - - Z, H, N_CTX_Q = 1, 4, 128 - q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype) * 0.01 - - fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose) - - print("✓ Small values test passed") - return metrics - - def test_large_values(self, verbose: bool = False) -> Dict[str, Any]: - """Test with large input values.""" - print("\n" + "="*60) - print("Test: Large values") - print("="*60) - - Z, H, N_CTX_Q = 1, 4, 128 - q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype) * 10.0 - - fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose) - - print("✓ Large values test passed") - return metrics - - def test_zeros(self, verbose: bool = False) -> Dict[str, Any]: - """Test with zero input.""" - print("\n" + "="*60) - print("Test: Zeros") - print("="*60) - - Z, H, N_CTX_Q = 1, 4, 128 - q = torch.zeros(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype) - - fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose) - - # Zeros should remain approximately zero after quantization - max_val = fake_q.abs().max().item() - assert max_val < 1e-3, f"Zeros should remain near zero, got max: {max_val}" - - print("✓ Zeros test passed") - return metrics - - def test_different_shapes(self, verbose: bool = False) -> Dict[str, Any]: - """Test with different tensor shapes.""" - print("\n" + "="*60) - print("Test: Different shapes") - print("="*60) - - test_shapes = [ - (1, 1, 64, self.HEAD_DIM), # Single batch, single head, small seq - (1, 8, 128, self.HEAD_DIM), # Single batch, multiple heads - (4, 16, 256, self.HEAD_DIM), # Multiple batches, many heads - (2, 8, 512, self.HEAD_DIM), # Medium size - ] - - all_metrics = [] - for shape in test_shapes: - Z, H, N_CTX_Q, HEAD_DIM = shape - q = torch.randn(*shape, device=self.device, dtype=self.dtype) - - fake_q, metrics = test_fake_quantize_q( - q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose - ) - - assert fake_q.shape == q.shape - assert metrics['cosine_similarity'] > 0.85 - all_metrics.append(metrics) - - if verbose: - print(f" Shape {shape}: cosine_sim={metrics['cosine_similarity']:.4f}") - - print("✓ Different shapes test passed") - return all_metrics - - def test_different_dtypes(self, verbose: bool = False) -> Dict[str, Any]: - """Test with different dtypes.""" - print("\n" + "="*60) - print("Test: Different dtypes") - print("="*60) - - Z, H, N_CTX_Q = 1, 4, 128 - dtypes = [torch.bfloat16, torch.float16, torch.float32] - - all_metrics = [] - for dtype in dtypes: - q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=dtype) - - fake_q, metrics = test_fake_quantize_q( - q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose - ) - - assert fake_q.dtype == q.dtype - assert metrics['cosine_similarity'] > 0.85 - all_metrics.append(metrics) - - if verbose: - print(f" {dtype}: cosine_sim={metrics['cosine_similarity']:.4f}") - - print("✓ Different dtypes test passed") - return all_metrics - - def test_non_divisible_sequence_length(self, verbose: bool = False) -> Dict[str, Any]: - """Test with sequence lengths that don't divide evenly by block size.""" - print("\n" + "="*60) - print("Test: Non-divisible sequence length") - print("="*60) - - Z, H = 1, 4 - # Use sequence length that doesn't divide evenly by BLOCK_M - N_CTX_Q = 150 # Not divisible by 128 - - q = torch.randn(Z, H, N_CTX_Q, self.HEAD_DIM, device=self.device, dtype=self.dtype) - - fake_q, metrics = test_fake_quantize_q(q, BLOCK_M=self.BLOCK_M, HEAD_DIM=self.HEAD_DIM, verbose=verbose) - - # Check that padding is handled correctly (last block should have zeros for padding) - assert fake_q.shape == q.shape - assert metrics['cosine_similarity'] > 0.85 - - print("✓ Non-divisible sequence length test passed") - return metrics - - def test_kv_different_shapes(self, verbose: bool = False) -> Dict[str, Any]: - """Test fake_quantize_kv with different shapes.""" - print("\n" + "="*60) - print("Test: KV different shapes") - print("="*60) - - test_shapes = [ - (1, 1, 64, self.HEAD_DIM), - (1, 8, 128, self.HEAD_DIM), - (2, 16, 256, self.HEAD_DIM), - ] - - all_metrics = [] - for shape in test_shapes: - Z, H, N_CTX_KV, HEAD_DIM = shape - k = torch.randn(*shape, device=self.device, dtype=self.dtype) - v = torch.randn(*shape, device=self.device, dtype=self.dtype) - - fake_k, fake_v, metrics = test_fake_quantize_kv( - k, v, BLOCK_N=self.BLOCK_N, HEAD_DIM=self.HEAD_DIM, verbose=verbose - ) - - assert fake_k.shape == k.shape - assert fake_v.shape == v.shape - assert metrics['k']['cosine_similarity'] > 0.85 - assert metrics['v']['cosine_similarity'] > 0.85 - all_metrics.append(metrics) - - print("✓ KV different shapes test passed") - return all_metrics - - def run_all_tests(self, verbose: bool = False) -> Dict[str, Any]: - """Run all tests.""" - print("\n" + "="*60) - print("Running Fake Quantize Kernels Precision Tests") - print("="*60) - - if not torch.cuda.is_available(): - print("CUDA not available, skipping tests") - return {} - - results = {} - - try: - results['basic_q'] = self.test_basic_q(verbose=verbose) - results['basic_kv'] = self.test_basic_kv(verbose=verbose) - results['small_values'] = self.test_small_values(verbose=verbose) - results['large_values'] = self.test_large_values(verbose=verbose) - results['zeros'] = self.test_zeros(verbose=verbose) - results['different_shapes'] = self.test_different_shapes(verbose=verbose) - results['different_dtypes'] = self.test_different_dtypes(verbose=verbose) - results['non_divisible_seq'] = self.test_non_divisible_sequence_length(verbose=verbose) - results['kv_different_shapes'] = self.test_kv_different_shapes(verbose=verbose) - - print("\n" + "="*60) - print("✓ All tests passed!") - print("="*60) - - except Exception as e: - print(f"\n✗ Test failed with error: {e}") - import traceback - traceback.print_exc() - raise - - return results - - -def main(): - """Main function to run the precision tests.""" - tester = FakeQuantizeKernelsTest(device="cuda", dtype=torch.bfloat16) - results = tester.run_all_tests(verbose=True) - - # Print summary - print("\n" + "="*60) - print("Test Summary") - print("="*60) - - if 'basic_q' in results: - print(f"Basic Q - Cosine similarity: {results['basic_q']['cosine_similarity']:.6f}") - print(f"Basic Q - Max diff: {results['basic_q']['max_diff']:.6f}") - - if 'basic_kv' in results: - print(f"Basic K - Cosine similarity: {results['basic_kv']['k']['cosine_similarity']:.6f}") - print(f"Basic V - Cosine similarity: {results['basic_kv']['v']['cosine_similarity']:.6f}") - - -if __name__ == "__main__": - main() - diff --git a/test_get_fake_quant.py b/test_get_fake_quant.py deleted file mode 100644 index a8d4fad6f..000000000 --- a/test_get_fake_quant.py +++ /dev/null @@ -1,336 +0,0 @@ -#!/usr/bin/env python3 -""" -Precision test to compare numerical differences for get_fake_quant Python function: -1. get_fake_quant (Python implementation using FlashInfer) -2. Tests various shapes, dtypes, and value ranges -3. Evaluates cosine similarity, max diff, and mean diff between input and output -""" - -import torch -from test_quantization_precision import get_fake_quant - -DEVICE = torch.device("cuda") - - -def cosine_similarity(tensor1, tensor2): - """ - Compute cosine similarity between two tensors. - - Args: - tensor1: First tensor - tensor2: Second tensor (same shape as tensor1) - - Returns: - Cosine similarity value (scalar) - - Returns 1.0 if both tensors are zero (identical zero vectors) - - Returns 0.0 if only one tensor is zero (orthogonal to non-zero vector) - """ - # Flatten tensors for computation - 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) - - # Handle zero vectors - if norm1 == 0 and norm2 == 0: - # Both are zero vectors - they are identical, so similarity is 1.0 - return 1.0 - elif norm1 == 0 or norm2 == 0: - # One is zero, one is not - they are orthogonal, so similarity is 0.0 - return 0.0 - - cos_sim = dot_product / (norm1 * norm2) - return cos_sim.item() - - -def test_get_fake_quant_basic(): - """Test basic functionality of get_fake_quant.""" - torch.manual_seed(42) - - # Test parameters - shape = (128, 128) - dtype = torch.bfloat16 - - # Create input tensor - x = torch.randn(shape, dtype=dtype, device=DEVICE) - - # Test get_fake_quant - x_fq = get_fake_quant(x) - - # Check that outputs have correct shape - assert x_fq.shape == x.shape - assert x_fq.dtype == x.dtype - - # Check that outputs are finite - assert torch.isfinite(x_fq).all() - - # Compare input and output - max_diff = (x_fq - x).abs().max() - mean_diff = (x_fq - x).abs().mean() - cos_sim = cosine_similarity(x_fq, x) - print(f" Input vs Output - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - print("✓ Basic test passed.") - - -def test_get_fake_quant_different_shapes(): - """Test get_fake_quant with different input shapes.""" - torch.manual_seed(42) - - test_configs = [ - (128, 64), # Small - (256, 128), # Medium - (512, 256), # Large - (128, 256), # Rectangular - (256, 128), # Rectangular (reversed) - ] - - dtype = torch.bfloat16 - - for outer_dim, quant_dim in test_configs: - # Create input tensor - x = torch.randn((outer_dim, quant_dim), dtype=dtype, device=DEVICE) - - # Test get_fake_quant - x_fq = get_fake_quant(x) - - assert x_fq.shape == x.shape - assert torch.isfinite(x_fq).all() - - # Compare input and output - max_diff = (x_fq - x).abs().max() - mean_diff = (x_fq - x).abs().mean() - cos_sim = cosine_similarity(x_fq, x) - print(f" Shape ({outer_dim}, {quant_dim}) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - print(f"✓ Shape test passed for ({outer_dim}, {quant_dim})") - - -def test_get_fake_quant_different_dtypes(): - """Test get_fake_quant with different input dtypes.""" - torch.manual_seed(42) - - shape = (128, 128) - dtypes = [torch.float32, torch.float16, torch.bfloat16] - - for dtype in dtypes: - # Create input tensor - x = torch.randn(shape, dtype=dtype, device=DEVICE) - - # Test get_fake_quant - x_fq = get_fake_quant(x) - - assert x_fq.shape == x.shape - assert x_fq.dtype == x.dtype - assert torch.isfinite(x_fq).all() - - # Compare input and output - max_diff = (x_fq - x).abs().max() - mean_diff = (x_fq - x).abs().mean() - cos_sim = cosine_similarity(x_fq, x) - print(f" Dtype {dtype} - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - print(f"✓ Dtype test passed for {dtype}") - - -def test_get_fake_quant_3d_4d_tensors(): - """Test get_fake_quant with 3D and 4D tensors (reshaped to 2D internally).""" - torch.manual_seed(42) - - dtype = torch.bfloat16 - - # Test 3D tensor (B, H, D) - print("\nTesting 3D tensor (B, H, D)") - x_3d = torch.randn((2, 8, 128), dtype=dtype, device=DEVICE) - x_fq_3d = get_fake_quant(x_3d) - - assert x_fq_3d.shape == x_3d.shape - assert torch.isfinite(x_fq_3d).all() - - max_diff = (x_fq_3d - x_3d).abs().max() - mean_diff = (x_fq_3d - x_3d).abs().mean() - cos_sim = cosine_similarity(x_fq_3d, x_3d) - print(f" 3D (2, 8, 128) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - # Test 4D tensor (B, H, L, D) - print("\nTesting 4D tensor (B, H, L, D)") - x_4d = torch.randn((1, 8, 256, 128), dtype=dtype, device=DEVICE) - x_fq_4d = get_fake_quant(x_4d) - - assert x_fq_4d.shape == x_4d.shape - assert torch.isfinite(x_fq_4d).all() - - max_diff = (x_fq_4d - x_4d).abs().max() - mean_diff = (x_fq_4d - x_4d).abs().mean() - cos_sim = cosine_similarity(x_fq_4d, x_4d) - print(f" 4D (1, 8, 256, 128) - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - print("✓ 3D/4D tensor test passed.") - - -def test_get_fake_quant_edge_cases(): - """Test edge cases like zeros, ones, extreme values.""" - torch.manual_seed(42) - - dtype = torch.bfloat16 - shape = (128, 128) - - test_cases = [ - ("zeros", torch.zeros(shape, dtype=dtype, device=DEVICE)), - ("ones", torch.ones(shape, dtype=dtype, device=DEVICE)), - ("negative_ones", -torch.ones(shape, dtype=dtype, device=DEVICE)), - ("very_small", torch.randn(shape, dtype=dtype, device=DEVICE) * 1e-6), - ("very_large", torch.randn(shape, dtype=dtype, device=DEVICE) * 1e6), - ] - - for name, x in test_cases: - x_fq = get_fake_quant(x) - - assert x_fq.shape == x.shape - assert torch.isfinite(x_fq).all() - - max_diff = (x_fq - x).abs().max() - mean_diff = (x_fq - x).abs().mean() - cos_sim = cosine_similarity(x_fq, x) - print(f" {name} - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - print("✓ Edge cases test passed.") - - -def test_get_fake_quant_attention_shapes(): - """Test get_fake_quant with attention-like shapes (similar to test_qat_attn.py).""" - torch.manual_seed(42) - - test_configs = [ - (2, 4, 128, 64), # Z, H, N_CTX, HEAD_DIM - Medium - (1, 8, 256, 128), # Large head dim - (1, 40, 9360, 128), # WAN shape - ] - - dtype = torch.bfloat16 - - for Z, H, N_CTX, HEAD_DIM in test_configs: - # Create input tensor in BLHD format (B, L, H, D) = (Z, N_CTX, H, HEAD_DIM) - x = torch.randn((Z, N_CTX, H, HEAD_DIM), dtype=dtype, device=DEVICE) - - # Test get_fake_quant - x_fq = get_fake_quant(x) - - assert x_fq.shape == x.shape - assert torch.isfinite(x_fq).all() - - # Compare input and output - max_diff = (x_fq - x).abs().max() - mean_diff = (x_fq - x).abs().mean() - cos_sim = cosine_similarity(x_fq, x) - print(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}") - - print(f"✓ Attention shape test passed for (Z={Z}, H={H}, N_CTX={N_CTX}, HEAD_DIM={HEAD_DIM})") - - -def test_get_fake_quant_non_divisible_blocks(): - """Test get_fake_quant with shapes that are not divisible by block sizes.""" - torch.manual_seed(42) - - dtype = torch.bfloat16 - - # Test with non-divisible dimensions - test_shapes = [ - (100, 100), # Not divisible by 128 - (150, 200), # Not divisible by 128 - (256, 100), # One dimension divisible, one not - ] - - for shape in test_shapes: - x = torch.randn(shape, dtype=dtype, device=DEVICE) - - # Test get_fake_quant - x_fq = get_fake_quant(x) - - assert x_fq.shape == x.shape - assert torch.isfinite(x_fq).all() - - max_diff = (x_fq - x).abs().max() - mean_diff = (x_fq - x).abs().mean() - cos_sim = cosine_similarity(x_fq, x) - print(f" Shape {shape} - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - print(f"✓ Non-divisible block test passed for {shape}") - - -def test_get_fake_quant_idempotency(): - """Test that quantizing twice produces similar results (idempotency check).""" - torch.manual_seed(42) - - dtype = torch.bfloat16 - shape = (128, 128) - - # Create input tensor - x = torch.randn(shape, dtype=dtype, device=DEVICE) - - # Quantize once - x_fq1 = get_fake_quant(x) - - # Quantize again - x_fq2 = get_fake_quant(x_fq1) - - # Compare the two quantizations - max_diff = (x_fq1 - x_fq2).abs().max() - mean_diff = (x_fq1 - x_fq2).abs().mean() - cos_sim = cosine_similarity(x_fq1, x_fq2) - print(f" First quantization vs Second quantization - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - print("✓ Idempotency test passed.") - - -def test_get_fake_quant_different_value_ranges(): - """Test get_fake_quant with different value ranges.""" - torch.manual_seed(42) - - dtype = torch.bfloat16 - shape = (128, 128) - - test_cases = [ - ("normal_range", lambda: torch.randn(shape, dtype=dtype, device=DEVICE)), - ("small_range", lambda: torch.randn(shape, dtype=dtype, device=DEVICE) * 0.01), - ("medium_range", lambda: torch.randn(shape, dtype=dtype, device=DEVICE) * 1.0), - ("large_range", lambda: torch.randn(shape, dtype=dtype, device=DEVICE) * 10.0), - ("very_large_range", lambda: torch.randn(shape, dtype=dtype, device=DEVICE) * 100.0), - ] - - for name, gen_func in test_cases: - x = gen_func() - x_fq = get_fake_quant(x) - - assert x_fq.shape == x.shape - assert torch.isfinite(x_fq).all() - - max_diff = (x_fq - x).abs().max() - mean_diff = (x_fq - x).abs().mean() - cos_sim = cosine_similarity(x_fq, x) - print(f" {name} - Max diff: {max_diff.item():.6f}, Mean diff: {mean_diff.item():.6f}, Cosine sim: {cos_sim:.6f}") - - print("✓ Different value ranges test passed.") - - -if __name__ == "__main__": - print("Running get_fake_quant precision tests...") - print(f"Device: {DEVICE}") - print() - - test_get_fake_quant_basic() - test_get_fake_quant_different_shapes() - test_get_fake_quant_different_dtypes() - test_get_fake_quant_3d_4d_tensors() - test_get_fake_quant_edge_cases() - test_get_fake_quant_attention_shapes() - test_get_fake_quant_non_divisible_blocks() - test_get_fake_quant_idempotency() - test_get_fake_quant_different_value_ranges() - - print() - print("All tests passed! ✓") - diff --git a/test_quantization_precision.py b/test_quantization_precision.py deleted file mode 100644 index c72a01ab3..000000000 --- a/test_quantization_precision.py +++ /dev/null @@ -1,461 +0,0 @@ -#!/usr/bin/env python3 -""" -Precision test to compare numerical differences between two quantization methods: -1. dequant(quant(X) @ quant(W)) -2. dequant(quant(X)) @ dequant(quant(W)) - -Uses FlashInfer API for 4-bit quantization producing bf16 matrices. -""" - -import torch -import numpy as np -from flashinfer import SfLayout, mm_fp4, nvfp4_quantize, e2m1_and_ufp8sf_scale_to_float -import unittest -from typing import Tuple, Dict, Any - -import torch -import triton -import triton.language as tl -import triton.profiler as proton -from triton.tools.tensor_descriptor import TensorDescriptor -from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor - -def get_fake_quant(x: torch.Tensor): - # TODO: dtype needs to be float32? - orig_shape = x.shape - orig_dtype = x.dtype - device = x.device - x = x.view(-1, x.shape[-1]) - x_global_sf = (448 * 6) / x.float().abs().nan_to_num().max() - x_fp4, x_scale = nvfp4_quantize(x, x_global_sf, sfLayout=SfLayout.layout_128x4, do_shuffle=False) - x_dequant = e2m1_and_ufp8sf_scale_to_float(x_fp4, x_scale, 1 / x_global_sf) # Note: this function does dequantization on the cpu - return x_dequant.view(orig_shape).to(orig_dtype).to(device) - - -def generate_fake_quant(a_f32: torch.Tensor, - b_f32: torch.Tensor, - vec_size: int = 16, - max_mxf4: float = 6.0, - eps: float = 1e-8): - """ - Fake-quantize fp32 A (M,K) and B (K,N) with nvfp4 (E2M1) using FP8 E4M3 (finite-only) block scales. - Returns dequantized fp32 tensors (A_qdq, B_qdq). Assumes: - - M % 128 == 0, N % 128 == 0, K % (vec_size * 4) == 0 - - A is (M, K), B is (K, N) - - Scales are uniform per 128x64 (A) and 64x128 (B) tiles (matching your kernel’s layout) - """ - assert a_f32.is_cuda and b_f32.is_cuda - M, K = a_f32.shape - Kb, N = b_f32.shape - assert Kb == K, "A and B inner dims must match (A:(M,K), B:(K,N))" - assert M % 128 == 0 and N % 128 == 0, "M,N must be multiples of 128" - assert K % (vec_size * 4) == 0, f"K must be multiple of {vec_size*4}" - - device = a_f32.device - - # ---- tile-scale helpers (produce FP8 E4M3 scales with shape [*, *, 32, 16]) ---- - def _tile_scales_A(a_fp32): - m_chunks = M // 128 - k_chunks = K // (vec_size * 4) # 64-wide tiles along K - x = a_fp32.view(m_chunks, 128, k_chunks, 4, vec_size) - tile_max = x.abs().amax(dim=(1, 3, 4)) # [m_chunks, k_chunks] - scale = torch.clamp(tile_max / max_mxf4, min=eps) # fp32 - s = scale[:, :, None, None].expand(m_chunks, k_chunks, 32, 16) - return s.to(torch.float8_e4m3fn) # [m_chunks, k_chunks, 32, 16] - - def _tile_scales_B(b_fp32): - k_chunks = K // (vec_size * 4) - n_chunks = N // 128 - x = b_fp32.view(k_chunks, 4, vec_size, n_chunks, 128) - tile_max = x.abs().amax(dim=(1, 2, 4)) # [k_chunks, n_chunks] - scale = torch.clamp(tile_max / max_mxf4, min=eps).t() # [n_chunks, k_chunks] - s = scale[:, :, None, None].expand(n_chunks, k_chunks, 32, 16) - return s.to(torch.float8_e4m3fn) # [n_chunks, k_chunks, 32, 16] - - # ---- expand tile scales to per-element fp32 ---- - def _expand_A_scale_full(a_scale_fp8): - m_chunks, k_chunks, _, _ = a_scale_fp8.shape - base = a_scale_fp8[..., 0, 0].to(a_f32.dtype) # [m_chunks, k_chunks] - tile = base[:, :, None, None].expand(m_chunks, k_chunks, 128, 64) - return tile.permute(0, 2, 1, 3).reshape(M, k_chunks * 64).contiguous() # [M,K] - - def _expand_B_scale_full(b_scale_fp8): - n_chunks, k_chunks, _, _ = b_scale_fp8.shape - base = b_scale_fp8[..., 0, 0].to(b_f32.dtype) # [n_chunks, k_chunks] - tile = base[:, :, None, None].expand(n_chunks, k_chunks, 64, 128) - return tile.permute(1, 2, 0, 3).reshape(k_chunks * 64, n_chunks * 128).contiguous() # [K,N] - - # ---- 1) compute FP8 E4M3 tile scales from fp32 ---- - a_scale_fp8 = _tile_scales_A(a_f32) # [M//128, K//64, 32, 16] - b_scale_fp8 = _tile_scales_B(b_f32) # [N//128, K//64, 32, 16] - - # ---- 2) expand scales to full fp32 for quant/dequant math ---- - a_scale_full = _expand_A_scale_full(a_scale_fp8) # [M, K] fp32 - b_scale_full = _expand_B_scale_full(b_scale_fp8) # [K, N] fp32 - - # ---- 3) quantize to nvfp4 codes (E2M1) ---- - a_codes = MXFP4Tensor(data=(a_f32 / a_scale_full)).data # [M, K] uint8 (low nibble) - b_codes = MXFP4Tensor(data=(b_f32 / b_scale_full)).data # [K, N] uint8 - - # ---- 4) dequantize back to fp32 (fake-quant output) ---- - a_qdq = MXFP4Tensor(size=(M, K), device=device); a_qdq.data = a_codes - b_qdq = MXFP4Tensor(size=(K, N), device=device); b_qdq.data = b_codes - - a_qdq = a_qdq.to(torch.float32).to(a_f32.dtype) * a_scale_full - b_qdq = b_qdq.to(torch.float32).to(b_f32.dtype) * b_scale_full - - return a_qdq, b_qdq - - -class QuantizationPrecisionTest: - """Test class for comparing quantization precision methods.""" - - def __init__(self, device: str = "cuda", dtype: torch.dtype = torch.bfloat16): - self.device = device - self.dtype = dtype - self.sf_layout = SfLayout.layout_128x4 - - def _compute_scale_factor(self, tensor: torch.Tensor) -> torch.Tensor: - """Compute scale factor for quantization using the same method as FastVideo.""" - maxabs = tensor.float().abs().nan_to_num().max() - maxabs = torch.maximum(maxabs, torch.tensor(1e-12, device=tensor.device, dtype=maxabs.dtype)) - return (448.0 * 6.0) / maxabs - - def quantize_matrix(self, matrix: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Quantize a matrix to 4-bit using FlashInfer. - - Args: - matrix: Input matrix to quantize - - Returns: - Tuple of (quantized_matrix, inverse_scale_factor) - """ - # scale_factor = self._compute_scale_factor(matrix) - scale_factor = torch.tensor(1.0, device=matrix.device, dtype=torch.float32) - quantized, inv_scale = nvfp4_quantize( - matrix, - scale_factor, - sfLayout=self.sf_layout, - do_shuffle=False - ) - return quantized, inv_scale, scale_factor - - def dequantize_matrix(self, shape: Tuple[int, int], quantized_matrix: torch.Tensor, inv_scale: torch.Tensor, global_scale_factor: torch.Tensor) -> torch.Tensor: - """ - Dequantize a 4-bit matrix back to bf16. - - Note: FlashInfer doesn't provide direct dequantization, so we use a workaround - by multiplying with an identity matrix to extract the dequantized values. - - Args: - quantized_matrix: 4-bit quantized matrix - inv_scale: Inverse scale factor from quantization - - Returns: - Dequantized matrix in bf16 - """ - # Create identity matrix for dequantization - if len(quantized_matrix.shape) == 2: - # 2D case - identity = torch.eye(shape[1], device=self.device, dtype=self.dtype) - identity_fp4, identity_inv_scale = nvfp4_quantize( - identity, - torch.tensor(1.0, device=self.device, dtype=torch.float32), - sfLayout=self.sf_layout, - do_shuffle=False - ) - dequantized = torch.empty(shape, device=self.device, dtype=self.dtype) - mm_fp4( - quantized_matrix, identity_fp4.T, inv_scale, identity_inv_scale.T, - 1 / global_scale_factor, - self.dtype, dequantized, - block_size=16, - use_8x4_sf_layout=False, - backend="cutlass" - ) - else: - # 3D case (batch dimension) - identity = torch.eye(quantized_matrix.shape[-1], device=self.device, dtype=self.dtype) - dequantized = torch.empty_like(quantized_matrix) - for i in range(quantized_matrix.shape[0]): - mm_fp4( - quantized_matrix[i], identity, inv_scale, inv_scale, - torch.tensor(1.0, device=self.device, dtype=torch.float32), - self.dtype, dequantized[i], - block_size=16, - use_8x4_sf_layout=False, - backend="cutlass" - ) - - return dequantized - - def real_quant_matmul(self, X: torch.Tensor, W: torch.Tensor) -> torch.Tensor: - """ - Method 1: dequant(quant(X) @ quant(W)) - Perform matrix multiplication in quantized space, then dequantize. - """ - # Quantize both matrices - X_quant, X_inv_scale, X_global_scale_factor = self.quantize_matrix(X) - W_quant, W_inv_scale, W_global_scale_factor = self.quantize_matrix(W) - - # Compute alpha scaling factor - alpha = 1.0 / (X_global_scale_factor * W_global_scale_factor) - - result = torch.empty((X.shape[0], W.shape[0]), device=self.device, dtype=self.dtype) - mm_fp4( - X_quant, W_quant.T, X_inv_scale, W_inv_scale.T, alpha, - self.dtype, result, - block_size=16, - use_8x4_sf_layout=False, - backend="cutlass" - ) - - - return result - - def fake_quant_matmul(self, X: torch.Tensor, W: torch.Tensor) -> torch.Tensor: - """ - Method 2: dequant(quant(X)) @ dequant(quant(W)) - Dequantize both matrices first, then perform matrix multiplication. - """ - # Quantize both matrices - - X_quant, X_inv_scale, X_global_scale_factor = self.quantize_matrix(X) - W_quant, W_inv_scale, W_global_scale_factor = self.quantize_matrix(W) - # Dequantize both matrices - X_dequant = self.dequantize_matrix(X.shape, X_quant, X_inv_scale, X_global_scale_factor) - W_dequant = self.dequantize_matrix(W.shape, W_quant, W_inv_scale, W_global_scale_factor) - - result = torch.matmul(X_dequant, W_dequant.T) - - - # X_fake_quant, W_fake_quant_T = generate_fake_quant(X, W.T.contiguous()) - # import pdb; pdb.set_trace() - # result = torch.matmul(X_fake_quant, W_fake_quant_T) - - return result - - def compare_methods(self, X: torch.Tensor, W: torch.Tensor) -> Dict[str, Any]: - """ - Compare the two quantization methods and return detailed analysis. - - Args: - X: Input matrix X - W: Weight matrix W - - Returns: - Dictionary containing comparison results - """ - # Ensure matrices are on correct device and dtype - X = X.to(device=self.device, dtype=self.dtype) - W = W.to(device=self.device, dtype=self.dtype) - - # Compute results using both methods - result1 = self.real_quant_matmul(X, W) - result2 = self.fake_quant_matmul(X, W) - import pdb; pdb.set_trace() - # Compute reference result (no quantization) - reference = torch.matmul(X, W.T) - # Calculate differences - diff_methods = torch.abs(result1 - result2) - diff1_ref = torch.abs(result1 - reference) - diff2_ref = torch.abs(result2 - reference) - - # Statistical analysis - analysis = { - 'method1_result': result1, - 'method2_result': result2, - 'reference_result': reference, - 'method_difference': { - 'max': torch.max(diff_methods).item(), - 'mean': torch.mean(diff_methods).item(), - 'std': torch.std(diff_methods).item(), - 'relative_max': torch.max(diff_methods / (torch.abs(reference) + 1e-8)).item(), - 'relative_mean': torch.mean(diff_methods / (torch.abs(reference) + 1e-8)).item(), - }, - 'method1_vs_reference': { - 'max': torch.max(diff1_ref).item(), - 'mean': torch.mean(diff1_ref).item(), - 'relative_max': torch.max(diff1_ref / (torch.abs(reference) + 1e-8)).item(), - }, - 'method2_vs_reference': { - 'max': torch.max(diff2_ref).item(), - 'mean': torch.mean(diff2_ref).item(), - 'relative_max': torch.max(diff2_ref / (torch.abs(reference) + 1e-8)).item(), - }, - 'input_shapes': { - 'X': list(X.shape), - 'W': list(W.shape), - } - } - - return analysis - - def run_comprehensive_test(self) -> Dict[str, Any]: - """Run comprehensive tests with different matrix sizes and distributions.""" - test_results = {} - - # Test cases with different sizes and distributions - test_cases = [ - # (X_shape, W_shape, distribution_name) - ((128, 128), (256, 128), "small_matrices"), - ((256, 512), (1024, 512), "medium_matrices"), - ((1024, 2048), (4096, 2048), "large_matrices"), - ] - - distributions = [ - ("normal", lambda shape: torch.randn(shape, device=self.device)), - ("uniform", lambda shape: torch.rand(shape, device=self.device) * 2 - 1), - ] - - for x_shape, w_shape, size_name in test_cases: - test_results[size_name] = {} - - for dist_name, dist_func in distributions: - # Generate test matrices - X = dist_func(x_shape).to(dtype=self.dtype) - W = dist_func(w_shape).to(dtype=self.dtype) - - # Run comparison - analysis = self.compare_methods(X, W) - test_results[size_name][dist_name] = analysis - - print(f"Test: {size_name} - {dist_name}") - print(f" Max difference between methods: {analysis['method_difference']['max']:.6f}") - print(f" Mean difference between methods: {analysis['method_difference']['mean']:.6f}") - - print(f" Max difference between method1 (real quant) and reference: {analysis['method1_vs_reference']['max']:.6f}") - print(f" Mean difference between method1 (real quant) and reference: {analysis['method1_vs_reference']['mean']:.6f}") - print(f" Max difference between method2 (fake quant) and reference: {analysis['method2_vs_reference']['max']:.6f}") - print(f" Mean difference between method2 (fake quant) and reference: {analysis['method2_vs_reference']['mean']:.6f}") - print() - - return test_results - - - -def main(): - """Main function to run the precision test.""" - print("Starting Quantization Precision Test") - print("=" * 50) - - # Check if CUDA is available - if not torch.cuda.is_available(): - print("CUDA not available, using CPU (may be slower)") - device = "cpu" - else: - device = "cuda" - print(f"Using CUDA device: {torch.cuda.get_device_name()}") - - # Create tester - tester = QuantizationPrecisionTest(device=device) - - # Run comprehensive test - print("\nRunning comprehensive tests...") - results = tester.run_comprehensive_test() - - -def test_get_fake_quant(): - """Test the get_fake_quant function with various inputs.""" - print("Testing get_fake_quant function...") - - if not torch.cuda.is_available(): - print("CUDA not available, skipping test") - return - - device = "cuda" - dtype = torch.bfloat16 - - # Test 1: Basic functionality - 2D tensor - print("\nTest 1: Basic 2D tensor") - x = torch.randn(128, 64, device=device, dtype=dtype) - x_fq = get_fake_quant(x) - assert x_fq.shape == x.shape, f"Shape mismatch: {x_fq.shape} vs {x.shape}" - assert x_fq.dtype == x.dtype, f"Dtype mismatch: {x_fq.dtype} vs {x.dtype}" - assert x_fq.device == x.device, f"Device mismatch: {x_fq.device} vs {x.device}" - max_diff = (x - x_fq).abs().max().item() - print(f" Max difference: {max_diff:.6f}") - assert max_diff < 1.0, f"Difference too large: {max_diff}" - - # Test 2: 3D tensor (B, H, D) - print("\nTest 2: 3D tensor [B, H, D]") - x = torch.randn(2, 8, 128, device=device, dtype=dtype) - x_fq = get_fake_quant(x) - assert x_fq.shape == x.shape, f"Shape mismatch: {x_fq.shape} vs {x.shape}" - assert x_fq.dtype == x.dtype, f"Dtype mismatch: {x_fq.dtype} vs {x.dtype}" - max_diff = (x - x_fq).abs().max().item() - print(f" Max difference: {max_diff:.6f}") - assert max_diff < 1.0, f"Difference too large: {max_diff}" - - # Test 3: 4D tensor (B, H, L, D) - print("\nTest 3: 4D tensor [B, H, L, D]") - x = torch.randn(1, 8, 256, 128, device=device, dtype=dtype) - x_fq = get_fake_quant(x) - assert x_fq.shape == x.shape, f"Shape mismatch: {x_fq.shape} vs {x.shape}" - assert x_fq.dtype == x.dtype, f"Dtype mismatch: {x_fq.dtype} vs {x.dtype}" - max_diff = (x - x_fq).abs().max().item() - print(f" Max difference: {max_diff:.6f}") - assert max_diff < 1.0, f"Difference too large: {max_diff}" - - # Test 4: Small values - print("\nTest 4: Small values") - x = torch.randn(128, 64, device=device, dtype=dtype) * 0.01 - x_fq = get_fake_quant(x) - assert x_fq.shape == x.shape - max_diff = (x - x_fq).abs().max().item() - print(f" Max difference: {max_diff:.6f}") - - # Test 5: Large values - print("\nTest 5: Large values") - x = torch.randn(128, 64, device=device, dtype=dtype) * 10.0 - x_fq = get_fake_quant(x) - assert x_fq.shape == x.shape - max_diff = (x - x_fq).abs().max().item() - print(f" Max difference: {max_diff:.6f}") - - # Test 6: Zeros - print("\nTest 6: Zeros") - x = torch.zeros(128, 64, device=device, dtype=dtype) - x_fq = get_fake_quant(x) - assert x_fq.shape == x.shape - max_diff = (x - x_fq).abs().max().item() - print(f" Max difference: {max_diff:.6f}") - assert max_diff < 1e-5, f"Zeros should remain zeros: {max_diff}" - - # Test 7: Different dtypes - print("\nTest 7: Float16 dtype") - x = torch.randn(128, 64, device=device, dtype=torch.float16) - x_fq = get_fake_quant(x) - assert x_fq.shape == x.shape - assert x_fq.dtype == x.dtype - max_diff = (x - x_fq).abs().max().item() - print(f" Max difference: {max_diff:.6f}") - - # Test 8: Idempotency check (quantizing twice should be similar) - print("\nTest 8: Idempotency") - x = torch.randn(128, 64, device=device, dtype=dtype) - x_fq1 = get_fake_quant(x) - x_fq2 = get_fake_quant(x_fq1) - # Second quantization should be close to first - max_diff = (x_fq1 - x_fq2).abs().max().item() - print(f" Max difference between two quantizations: {max_diff:.6f}") - - # Test 9: Requires grad - print("\nTest 9: Gradient flow") - x = torch.randn(128, 64, device=device, dtype=dtype, requires_grad=True) - x_fq = get_fake_quant(x) - # Note: get_fake_quant uses no_grad internally, so x_fq won't have grad - # But we can check it doesn't crash - assert x_fq.shape == x.shape - print(" Gradient test passed (no crash)") - - print("\n✓ All tests passed!") - - -if __name__ == "__main__": - import sys - if len(sys.argv) > 1 and sys.argv[1] == "test_fake_quant": - test_get_fake_quant() - else: - main() diff --git a/test_scale_debug.py b/test_scale_debug.py deleted file mode 100644 index 4785b7313..000000000 --- a/test_scale_debug.py +++ /dev/null @@ -1,92 +0,0 @@ -#!/usr/bin/env python3 -""" -Debug script to understand the scale computation issue. -""" - -import torch -import triton -import triton.language as tl -from nvfp4_utils import _compute_quant_and_scale, _compute_dequant - -DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") - -@triton.jit -def debug_scale_kernel( - src_ptr, - scale_ptr, - s_dec_ptr, - BLOCK_SIZE_OUT_DIM: tl.constexpr, - BLOCK_SIZE_QUANT_DIM: tl.constexpr, - stride_src_outer, - stride_src_quant, - outer_dim, - quant_dim, -): - """Debug kernel to print scale values.""" - outer_idx = tl.program_id(0) - quant_idx = tl.program_id(1) - - start_outer = outer_idx * BLOCK_SIZE_OUT_DIM - start_quant = quant_idx * BLOCK_SIZE_QUANT_DIM - - offs_outer = tl.arange(0, BLOCK_SIZE_OUT_DIM)[:, None] - offs_quant = tl.arange(0, BLOCK_SIZE_QUANT_DIM)[None, :] - - mask_outer = (start_outer + offs_outer) < outer_dim - mask_quant = (start_quant + offs_quant) < quant_dim - full_mask = mask_outer & mask_quant - - src_offsets = (start_outer + offs_outer) * stride_src_outer + (start_quant + offs_quant) * stride_src_quant - src_tensor = tl.load(src_ptr + src_offsets, mask=full_mask, other=0.0) - - valid_mask = tl.full(src_tensor.shape, 1, tl.int1) - quantized_tensor, scale_tensor, s_dec = _compute_quant_and_scale( - src_tensor=src_tensor, - valid_src_mask=valid_mask, - mx_tensor_dtype=tl.float8e4nv - ) - - # Store scale and s_dec for inspection - # scale_tensor has shape [BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE] where BLOCK_SIZE_QUANT_MX_SCALE = BLOCK_SIZE_QUANT_DIM // 16 - BLOCK_SIZE_QUANT_MX_SCALE: tl.constexpr = BLOCK_SIZE_QUANT_DIM // 16 - scale_offsets = (start_outer + offs_outer) * stride_src_outer + (start_quant // 16) * stride_src_quant - tl.store(scale_ptr + scale_offsets, scale_tensor.reshape([BLOCK_SIZE_OUT_DIM, BLOCK_SIZE_QUANT_MX_SCALE])[:, 0].to(tl.float32), mask=mask_outer) - tl.store(s_dec_ptr + outer_idx, s_dec) - - -if __name__ == "__main__": - if not torch.cuda.is_available(): - print("CUDA not available") - exit(1) - - # Test with all ones - x = torch.ones(128, 64, device=DEVICE, dtype=torch.float16) - x_2d = x.view(-1, x.shape[-1]) - outer_dim, quant_dim = x_2d.shape - - scale_output = torch.zeros(outer_dim, quant_dim // 16, device=DEVICE, dtype=torch.float32) - s_dec_output = torch.zeros(outer_dim // 128, device=DEVICE, dtype=torch.float32) - - grid = ( - triton.cdiv(outer_dim, 128), - triton.cdiv(quant_dim, 64), - ) - - debug_scale_kernel[grid]( - src_ptr=x_2d, - scale_ptr=scale_output, - s_dec_ptr=s_dec_output, - BLOCK_SIZE_OUT_DIM=128, - BLOCK_SIZE_QUANT_DIM=64, - stride_src_outer=x_2d.stride(0), - stride_src_quant=x_2d.stride(1), - outer_dim=outer_dim, - quant_dim=quant_dim, - ) - - print(f"Scale values (first block): {scale_output[0, :4]}") - print(f"s_dec value: {s_dec_output[0]}") - print(f"Expected s_dec: {1.0 / (6 * 448 / 1.0)} = {1.0 / 2688}") - print(f"Expected scale: {(1.0 / 6) * 2688} = {2688 / 6} = {448}") - - diff --git a/test_triton_fp4_vs_cutlass.py b/test_triton_fp4_vs_cutlass.py deleted file mode 100644 index c028425fb..000000000 --- a/test_triton_fp4_vs_cutlass.py +++ /dev/null @@ -1,276 +0,0 @@ -import torch -from triton.tools.tensor_descriptor import TensorDescriptor -from triton.tools.mxfp import MXFP4Tensor -from triton_fp4 import block_scaled_matmul -from flashinfer import SfLayout, mm_fp4, nvfp4_quantize - - - -def triton_nvfp4(A: torch.Tensor, - B: torch.Tensor, - out_dtype: torch.dtype = torch.float16): - """ - Block-scaled nvfp4 GEMM: (A @ B) with A:(M,K), B:(K,N) given in fp32/bf16. - Returns the GEMM output in `out_dtype`. - - Requirements (match your kernel): - - M % 128 == 0, N % 128 == 0, K % 64 == 0 (since vec_size=16 and 4 groups → 64) - - Device: CUDA, Blackwell (SM=10) - """ - assert A.is_cuda and B.is_cuda, "Inputs must be on CUDA" - assert A.dtype in (torch.float32, torch.bfloat16), "A must be fp32/bf16" - assert B.dtype in (torch.float32, torch.bfloat16), "B must be fp32/bf16" - M, K = A.shape - Kb, N = B.shape - assert Kb == K, "Inner dims must match" - assert M % 128 == 0 and N % 128 == 0 and K % 64 == 0, "Shape must respect tiling" - - # ---- kernel constants / config ---- - BLOCK_M = 128 - BLOCK_N = 256 - BLOCK_K = 256 - VEC_SIZE = 16 - ELEM_PER_BYTE_A = 2 - ELEM_PER_BYTE_B = 2 - MAX_MXF4 = 6.0 # nvfp4 (E2M1) max magnitude - EPS = 1e-8 - - device = A.device - - # ---- tile-scale helpers (FP8 E4M3, finite-only) ---- - # A tiles are 128 x 64 along (M,K) - def _tile_scales_A(a_fp32): - m_chunks = M // 128 - k_chunks = K // (VEC_SIZE * 4) # 64 wide - x = a_fp32.view(m_chunks, 128, k_chunks, 4, VEC_SIZE) - tile_max = x.abs().amax(dim=(1, 3, 4)) # [m_chunks, k_chunks] - scale = torch.clamp(tile_max / MAX_MXF4, min=EPS) # fp32 - s = scale[:, :, None, None].expand(m_chunks, k_chunks, 32, 16) - return s.to(torch.float8_e4m3fn) # [m_chunks, k_chunks, 32, 16] - - # B tiles are 64 x 128 over (K,N) - def _tile_scales_B(b_fp32): - k_chunks = K // (VEC_SIZE * 4) - n_chunks = N // 128 - x = b_fp32.view(k_chunks, 4, VEC_SIZE, n_chunks, 128) - tile_max = x.abs().amax(dim=(1, 2, 4)) # [k_chunks, n_chunks] - scale = torch.clamp(tile_max / MAX_MXF4, min=EPS).t() # [n_chunks, k_chunks] - s = scale[:, :, None, None].expand(n_chunks, k_chunks, 32, 16) - return s.to(torch.float8_e4m3fn) # [n_chunks, k_chunks, 32, 16] - - def _expand_A_scale_full(a_scale_fp8): - m_chunks, k_chunks, _, _ = a_scale_fp8.shape - base = a_scale_fp8[..., 0, 0].to(A.dtype) # [m_chunks, k_chunks] - tile = base[:, :, None, None].expand(m_chunks, k_chunks, 128, 64) - return tile.permute(0, 2, 1, 3).reshape(M, k_chunks * 64).contiguous() # [M,K] - - def _expand_B_scale_full(b_scale_fp8): - n_chunks, k_chunks, _, _ = b_scale_fp8.shape - base = b_scale_fp8[..., 0, 0].to(B.dtype) # [n_chunks, k_chunks] - tile = base[:, :, None, None].expand(n_chunks, k_chunks, 64, 128) - return tile.permute(1, 2, 0, 3).reshape(k_chunks * 64, n_chunks * 128).contiguous() # [K,N] - - # ---- 1) scales from fp32 ---- - a_scale_fp8 = _tile_scales_A(A) # [M//128, K//64, 32, 16] - b_scale_fp8 = _tile_scales_B(B) # [N//128, K//64, 32, 16] - a_scale_full = _expand_A_scale_full(a_scale_fp8) # [M,K] fp32 - b_scale_full = _expand_B_scale_full(b_scale_fp8) # [K,N] fp32 - - # ---- 2) quantize to nvfp4 codes (E2M1) ---- - a_codes = MXFP4Tensor(data=(A / a_scale_full)).data # [M,K] uint8 (low nibble used) - b_codes_KN = MXFP4Tensor(data=(B / b_scale_full)).data # [K,N] uint8 - - # Kernel expects B laid out as (N, K) packed along K (like the original example) - b_codes_NK = b_codes_KN.transpose(0, 1).contiguous() # [N,K] - - # ---- 3) pack to 2-per-byte along K and build descriptors ---- - a_fp4 = MXFP4Tensor(size=(M, K), device=device); a_fp4.data = a_codes - b_fp4_desc = MXFP4Tensor(size=(N, K), device=device); b_fp4_desc.data = b_codes_NK - - a_packed = a_fp4.to_packed_tensor(dim=1) # (M, K/2) - b_packed = b_fp4_desc.to_packed_tensor(dim=1) # (N, K/2) - - a_desc = TensorDescriptor.from_tensor(a_packed, [BLOCK_M, BLOCK_K // ELEM_PER_BYTE_A]) - b_desc = TensorDescriptor.from_tensor(b_packed, [BLOCK_N, BLOCK_K // ELEM_PER_BYTE_B]) - - # reps for scale swizzle - rep_m = BLOCK_M // 128 - rep_n = BLOCK_N // 128 - rep_k = BLOCK_K // VEC_SIZE // 4 - - # reshape FP8 scales into 5D TMA blocks [1, rep_*, rep_k, 2, 256] - a_scale_5d = a_scale_fp8.reshape(1, a_scale_fp8.shape[0], a_scale_fp8.shape[1], 2, 256) - b_scale_5d = b_scale_fp8.reshape(1, b_scale_fp8.shape[0], b_scale_fp8.shape[1], 2, 256) - a_scale_desc = TensorDescriptor.from_tensor(a_scale_5d, [1, rep_m, rep_k, 2, 256]) - b_scale_desc = TensorDescriptor.from_tensor(b_scale_5d, [1, rep_n, rep_k, 2, 256]) - - a_packed_cutlass, a_scale_fp8_cutlass = nvfp4_quantize( - A, torch.tensor(1.0, device=device, dtype=torch.float32), sfLayout=SfLayout.layout_128x4, do_shuffle=False - ) - import pdb; pdb.set_trace() - # ---- 4) launch GEMM ---- - configs = { - "BLOCK_SIZE_M": BLOCK_M, - "BLOCK_SIZE_N": BLOCK_N, - "BLOCK_SIZE_K": BLOCK_K, - "num_stages": 4, - "ELEM_PER_BYTE_A": ELEM_PER_BYTE_A, - "ELEM_PER_BYTE_B": ELEM_PER_BYTE_B, - "VEC_SIZE": VEC_SIZE, - } - # Uses your already-defined block_scaled_matmul(...) - C = block_scaled_matmul( - a_desc, a_scale_desc, - b_desc, b_scale_desc, - out_dtype, M, N, K, - rep_m, rep_n, rep_k, configs - ) - return C - - - -def cutlass_nvfp4(A: torch.Tensor, - B: torch.Tensor, - out_dtype: torch.dtype = torch.bfloat16, - sf_layout: SfLayout = SfLayout.layout_128x4, - block_size: int = 16, - do_shuffle: bool = False, - scale_factor_A: torch.Tensor | float | None = None, - scale_factor_B: torch.Tensor | float | None = None) -> torch.Tensor: - """ - FlashInfer/CUTLASS nvfp4 GEMM: returns (A @ B) using 4-bit E2M1 operands with per-tile scaling. - - Args: - A: (M, K) tensor, fp32 or bf16, CUDA. - B: (K, N) tensor, fp32 or bf16, CUDA. - out_dtype: torch.float32, torch.float16, or torch.bfloat16 (recommended). - sf_layout: FlashInfer scale-factor layout (default SfLayout.layout_128x4). - block_size: internal tiling, default 16. - do_shuffle: pass-through to nvfp4_quantize. - scale_factor_A, scale_factor_B: optional global scale factors (float or 0-dim tensor). - If None, defaults to 1.0 (FlashInfer still applies its internal per-tile scales). - - Returns: - C: (M, N) tensor in out_dtype - """ - assert A.is_cuda and B.is_cuda, "Inputs must be CUDA tensors" - assert A.dim() == 2 and B.dim() == 2, "A and B must be 2D" - M, K = A.shape - Kb, N = B.shape - assert K == Kb, f"Incompatible shapes: A is (M={M},K={K}), B is (K={Kb},N={N})" - - # Cast inputs to supported accumulation input types for FlashInfer quantizer. - # (fp32/bf16 are fine; we keep user dtype if already bf16 to save bandwidth.) - A_in = A if A.dtype in (torch.bfloat16, torch.float32) else A.float() - B_in = B if B.dtype in (torch.bfloat16, torch.float32) else B.float() - B_in = B_in.T.contiguous() - - # Choose global scaling (FlashInfer also uses per-tile scales internally). - if scale_factor_A is None: - scale_factor_A = torch.tensor(1.0, device=A.device, dtype=torch.float32) - elif not torch.is_tensor(scale_factor_A): - scale_factor_A = torch.tensor(float(scale_factor_A), device=A.device, dtype=torch.float32) - - if scale_factor_B is None: - scale_factor_B = torch.tensor(1.0, device=B.device, dtype=torch.float32) - elif not torch.is_tensor(scale_factor_B): - scale_factor_B = torch.tensor(float(scale_factor_B), device=B.device, dtype=torch.float32) - - # Quantize A and B to nvfp4; inv scales are FlashInfer’s per-tile inverse scales. - # Returned quantized matrices are FP4-packed formats understood by mm_fp4. - A_q, A_inv_sf = nvfp4_quantize( - A_in, scale_factor_A, sfLayout=sf_layout, do_shuffle=do_shuffle - ) - B_q, B_inv_sf = nvfp4_quantize( - B_in, scale_factor_B, sfLayout=sf_layout, do_shuffle=do_shuffle - ) - - # Alpha rescales the product of the two global scale factors. - alpha = 1.0 / (scale_factor_A * scale_factor_B) - - # Pick FlashInfer flag for the scale-factor layout - use_8x4_sf_layout = (sf_layout == SfLayout.layout_8x4) - - # Allocate output and call CUTLASS-backed FP4 GEMM. - C = torch.empty((M, N), device=A.device, dtype=out_dtype) - mm_fp4( - A_q, B_q.T, # note the transpose on B_q (expecting KxN → NxK in kernel) - A_inv_sf, B_inv_sf.T, # matching transposition of per-tile inverse scales - alpha, # global alpha - out_dtype, C, # output dtype & buffer - block_size=block_size, - use_8x4_sf_layout=use_8x4_sf_layout, - backend="cutlass" - ) - return C - - -def main(): - import argparse, math - torch.manual_seed(0) - - parser = argparse.ArgumentParser() - parser.add_argument("--M", type=int, default=256, help="Rows of A / C (multiple of 128)") - parser.add_argument("--N", type=int, default=256, help="Cols of B / C (multiple of 128)") - parser.add_argument("--K", type=int, default=512, help="Inner dim (multiple of 64)") - parser.add_argument("--dtype", type=str, choices=["fp32", "bf16"], default="bf16") - parser.add_argument("--trials", type=int, default=1) - parser.add_argument("--verbose", action="store_true") - args = parser.parse_args() - - if not torch.cuda.is_available(): - raise RuntimeError("CUDA required") - - M, N, K = args.M, args.N, args.K - assert M % 128 == 0 and N % 128 == 0 and K % 64 == 0, "Shape must satisfy: M%128==0, N%128==0, K%64==0." - - dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 - device = "cuda" - - def stats(name, X): - X = X.float() - return f"{name}: max={X.abs().max().item():.6e} mean={X.abs().mean().item():.6e} std={X.std().item():.6e}" - - for t in range(args.trials): - # Inputs - A = torch.randn((M, K), device=device, dtype=dtype) - B = torch.randn((K, N), device=device, dtype=dtype) - - # GEMMs - C_triton = triton_nvfp4(A, B, out_dtype=torch.bfloat16 if dtype is torch.bfloat16 else torch.float32) - C_cutlass = cutlass_nvfp4(A, B, out_dtype=torch.bfloat16 if dtype is torch.bfloat16 else torch.float32) - # Bring to fp32 for error analysis - Ct = C_triton.float() - Cc = C_cutlass.float() - Cref = (A.float() @ B.float()) - # Triton vs Cutlass - diff_tc = (Ct - Cc) - denom_tc = Cc.abs().clamp_min(1e-8) - rel_tc = (diff_tc.abs() / denom_tc) - - # Each vs reference (optional sanity) - diff_t_ref = (Ct - Cref) - diff_c_ref = (Cc - Cref) - rel_t_ref = diff_t_ref.abs() / Cref.abs().clamp_min(1e-8) - rel_c_ref = diff_c_ref.abs() / Cref.abs().clamp_min(1e-8) - - - print("\nTriton vs Cutlass:") - print(f" abs diff -> max={diff_tc.abs().max().item():.6e} mean={diff_tc.abs().mean().item():.6e}") - print(f" rel diff -> max={rel_tc.max().item():.6e} mean={rel_tc.mean().item():.6e}") - - print("\nTriton vs Reference:") - print(f" abs diff -> max={diff_t_ref.abs().max().item():.6e} mean={diff_t_ref.abs().mean().item():.6e}") - print(f" rel diff -> max={rel_t_ref.max().item():.6e} mean={rel_t_ref.mean().item():.6e}") - - print("\nCutlass vs Reference:") - print(f" abs diff -> max={diff_c_ref.abs().max().item():.6e} mean={diff_c_ref.abs().mean().item():.6e}") - print(f" rel diff -> max={rel_c_ref.max().item():.6e} mean={rel_c_ref.mean().item():.6e}") - - if args.verbose: - # a tiny sample print to eyeball values - i, j = 0, 0 - print(f"\nSample C[0,0]: triton={Ct[i,j].item():.6e}, cutlass={Cc[i,j].item():.6e}, ref={Cref[i,j].item():.6e}") - -if __name__ == "__main__": - main() \ No newline at end of file diff --git a/triton_fp4.py b/triton_fp4.py deleted file mode 100644 index 782838132..000000000 --- a/triton_fp4.py +++ /dev/null @@ -1,320 +0,0 @@ -import argparse - -import torch -import triton -import triton.language as tl -import triton.profiler as proton -from triton.tools.tensor_descriptor import TensorDescriptor -from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor -from test_quantization_precision import generate_fake_quant - -VEC_SIZE = 16 -MAX_MXF4 = 6.0 # max magnitude for E2M1 (nvfp4) with M∈{0,1}, E∈{0..3} -EPS = 1e-8 - - - -def is_cuda(): - return triton.runtime.driver.active.get_current_target().backend == "cuda" - - -def supports_block_scaling(): - return is_cuda() and torch.cuda.get_device_capability()[0] == 10 - - -def _matmul_launch_metadata(grid, kernel, args): - ret = {} - M, N, K = args["M"], args["N"], args["K"] - kernel_name = kernel.name - if "ELEM_PER_BYTE_A" and "ELEM_PER_BYTE_B" and "VEC_SIZE" in args: - kernel_name += "_nvfp4" - ret["name"] = f"{kernel_name} [M={M}, N={N}, K={K}]" - ret["flops"] = 2.0 * M * N * K - return ret - - -@triton.jit(launch_metadata=_matmul_launch_metadata) -def block_scaled_matmul_kernel( # - a_desc, # - a_scale_desc, # - b_desc, # - b_scale_desc, # - c_desc, # - M: tl.constexpr, # - N: tl.constexpr, # - K: tl.constexpr, # - output_type: tl.constexpr, # - ELEM_PER_BYTE_A: tl.constexpr, # - ELEM_PER_BYTE_B: tl.constexpr, # - VEC_SIZE: tl.constexpr, # - BLOCK_M: tl.constexpr, # - BLOCK_N: tl.constexpr, # - BLOCK_K: tl.constexpr, # - rep_m: tl.constexpr, # - rep_n: tl.constexpr, # - rep_k: tl.constexpr, # - NUM_STAGES: tl.constexpr, # -): # - if output_type == 0: - output_dtype = tl.float32 - elif output_type == 1: - output_dtype = tl.float16 - elif output_type == 2: - output_dtype = tl.bfloat16 - elif output_type == 3: - output_dtype = tl.float8e4nv - - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_M) - pid_m = pid % num_pid_m - pid_n = pid // num_pid_m - offs_am = pid_m * BLOCK_M - offs_bn = pid_n * BLOCK_N - offs_k_a = 0 - offs_k_b = 0 - offs_scale_m = pid_m * rep_m - offs_scale_n = pid_n * rep_n - offs_scale_k = 0 - - MIXED_PREC: tl.constexpr = ELEM_PER_BYTE_A == 1 and ELEM_PER_BYTE_B == 2 - - accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k in tl.range(0, tl.cdiv(K, BLOCK_K), num_stages=NUM_STAGES): - a = a_desc.load([offs_am, offs_k_a]) - b = b_desc.load([offs_bn, offs_k_b]) - scale_a = a_scale_desc.load([0, offs_scale_m, offs_scale_k, 0, 0]) - scale_b = b_scale_desc.load([0, offs_scale_n, offs_scale_k, 0, 0]) - - scale_a = scale_a.reshape(rep_m, rep_k, 32, 4, 4).trans(0, 3, 2, 1, 4).reshape(BLOCK_M, BLOCK_K // VEC_SIZE) - scale_b = scale_b.reshape(rep_n, rep_k, 32, 4, 4).trans(0, 3, 2, 1, 4).reshape(BLOCK_N, BLOCK_K // VEC_SIZE) - - if MIXED_PREC: - accumulator = tl.dot_scaled(a, scale_a, "e4m3", b.T, scale_b, "e2m1", accumulator) - elif ELEM_PER_BYTE_A == 2 and ELEM_PER_BYTE_B == 2: - accumulator = tl.dot_scaled(a, scale_a, "e2m1", b.T, scale_b, "e2m1", accumulator) - else: - accumulator = tl.dot_scaled(a, scale_a, "e4m3", b.T, scale_b, "e4m3", accumulator) - - offs_k_a += BLOCK_K // ELEM_PER_BYTE_A - offs_k_b += BLOCK_K // ELEM_PER_BYTE_B - offs_scale_k += rep_k - - c_desc.store([offs_am, offs_bn], accumulator.to(output_dtype)) - - -def block_scaled_matmul(a_desc, a_scale_desc, b_desc, b_scale_desc, dtype_dst, M, N, K, rep_m, rep_n, rep_k, configs): - output = torch.empty((M, N), dtype=dtype_dst, device="cuda") - if dtype_dst == torch.float32: - dtype_dst = 0 - elif dtype_dst == torch.float16: - dtype_dst = 1 - elif dtype_dst == torch.bfloat16: - dtype_dst = 2 - elif dtype_dst == torch.float8_e4m3fn: - dtype_dst = 3 - else: - raise ValueError(f"Unsupported dtype: {dtype_dst}") - - BLOCK_M = configs["BLOCK_SIZE_M"] - BLOCK_N = configs["BLOCK_SIZE_N"] - c_desc = TensorDescriptor.from_tensor(output, [BLOCK_M, BLOCK_N]) - - grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1) - block_scaled_matmul_kernel[grid]( - a_desc, - a_scale_desc, - b_desc, - b_scale_desc, - c_desc, - M, - N, - K, - dtype_dst, - configs["ELEM_PER_BYTE_A"], - configs["ELEM_PER_BYTE_B"], - configs["VEC_SIZE"], - configs["BLOCK_SIZE_M"], - configs["BLOCK_SIZE_N"], - configs["BLOCK_SIZE_K"], - rep_m, - rep_n, - rep_k, - configs["num_stages"], - ) - return output -def initialize_block_scaled(M, N, K, block_scale_type="nvfp4", compute_reference=False, input_output_dtype=torch.bfloat16): - assert block_scale_type == "nvfp4", f"Only nvfp4 supported here (got {block_scale_type})" - - # --- constants --- - BLOCK_M = 128 - BLOCK_N = 256 - BLOCK_K = 256 - VEC_SIZE = 16 - ELEM_PER_BYTE_A = 2 - ELEM_PER_BYTE_B = 2 - device = "cuda" - MAX_MXF4 = 6.0 - EPS = 1e-8 - - assert M % 128 == 0 - assert N % 128 == 0 - assert K % (VEC_SIZE * 4) == 0 - - # --- helpers: tile scales (FP8 E4M3) --- - def _tile_scales_A(a_fp32): - m_chunks = M // 128 - k_chunks = K // (VEC_SIZE * 4) # 64-wide tiles along K - x = a_fp32.view(m_chunks, 128, k_chunks, 4, VEC_SIZE) - tile_max = x.abs().amax(dim=(1, 3, 4)) # [m_chunks, k_chunks] - scale = torch.clamp(tile_max / MAX_MXF4, min=EPS) # fp32 - s = scale[:, :, None, None].expand(m_chunks, k_chunks, 32, 16) - return s.to(torch.float8_e4m3fn) # [m_chunks, k_chunks, 32, 16] - - def _tile_scales_B(b_fp32): # b_fp32 is (K, N) - k_chunks = K // (VEC_SIZE * 4) - n_chunks = N // 128 - x = b_fp32.view(k_chunks, 4, VEC_SIZE, n_chunks, 128) - tile_max = x.abs().amax(dim=(1, 2, 4)) # [k_chunks, n_chunks] - scale = torch.clamp(tile_max / MAX_MXF4, min=EPS).t() # [n_chunks, k_chunks] - s = scale[:, :, None, None].expand(n_chunks, k_chunks, 32, 16) - return s.to(torch.float8_e4m3fn) # [n_chunks, k_chunks, 32, 16] - - # --- expand tile scales to per-element fp32 (for quant/dequant math) --- - def _expand_A_scale_full(a_scale_fp8): - m_chunks, k_chunks, _, _ = a_scale_fp8.shape - base = a_scale_fp8[..., 0, 0].to(input_output_dtype) # [m_chunks, k_chunks] - tile = base[:, :, None, None].expand(m_chunks, k_chunks, 128, 64) - return tile.permute(0, 2, 1, 3).reshape(M, k_chunks * 64).contiguous() # [M, K] - - def _expand_B_scale_full(b_scale_fp8): - n_chunks, k_chunks, _, _ = b_scale_fp8.shape - base = b_scale_fp8[..., 0, 0].to(input_output_dtype) # [n_chunks, k_chunks] - tile = base[:, :, None, None].expand(n_chunks, k_chunks, 64, 128) - return tile.permute(1, 2, 0, 3).reshape(k_chunks * 64, n_chunks * 128).contiguous() # [K, N] - - # --- 1) fp32 sources (pre-quant) --- - torch.manual_seed(0) - a_f32_pre = torch.randn((M, K), device=device, dtype=input_output_dtype) - b_f32_pre = torch.randn((K, N), device=device, dtype=input_output_dtype) - - # --- 2) tile scales (FP8 E4M3, finite-only) --- - a_scale_fp8 = _tile_scales_A(a_f32_pre) # [M//128, K//64, 32, 16] - b_scale_fp8 = _tile_scales_B(b_f32_pre) # [N//128, K//64, 32, 16] - - # --- 3) expand scales to full for quantization --- - a_scale_full = _expand_A_scale_full(a_scale_fp8) # [M, K] fp32 - b_scale_full = _expand_B_scale_full(b_scale_fp8) # [K, N] fp32 - - # --- 4) quantize to nvfp4 codes (E2M1) --- - a_codes = MXFP4Tensor(data=(a_f32_pre / a_scale_full)).data # [M, K] uint8 (nibble used) - b_codes_KN = MXFP4Tensor(data=(b_f32_pre / b_scale_full)).data # [K, N] uint8 - - # IMPORTANT: the kernel expects B laid out as (N, K_packed) (like original) - # so create a view for descriptors with shape (N, K) and pack along dim=1. - b_codes_NK = b_codes_KN.transpose(0, 1).contiguous() # [N, K] uint8 - - # --- 5) pack to 2-per-byte along K and build descriptors --- - a_fp4 = MXFP4Tensor(size=(M, K), device=device); a_fp4.data = a_codes - b_fp4_for_desc = MXFP4Tensor(size=(N, K), device=device); b_fp4_for_desc.data = b_codes_NK - - a_packed = a_fp4.to_packed_tensor(dim=1) # (M, K/2) - b_packed = b_fp4_for_desc.to_packed_tensor(dim=1) # (N, K/2) <-- matches kernel indexing - - a_desc = TensorDescriptor.from_tensor(a_packed, [BLOCK_M, BLOCK_K // ELEM_PER_BYTE_A]) - b_desc = TensorDescriptor.from_tensor(b_packed, [BLOCK_N, BLOCK_K // ELEM_PER_BYTE_B]) - - rep_m = BLOCK_M // 128 - rep_n = BLOCK_N // 128 - rep_k = BLOCK_K // VEC_SIZE // 4 - - # reshape FP8 scales to 5D TMA blocks [1, rep_*, rep_k, 2, 256] - a_scale_5d = a_scale_fp8.reshape(1, a_scale_fp8.shape[0], a_scale_fp8.shape[1], 2, 256) - b_scale_5d = b_scale_fp8.reshape(1, b_scale_fp8.shape[0], b_scale_fp8.shape[1], 2, 256) - a_scale_desc = TensorDescriptor.from_tensor(a_scale_5d, [1, rep_m, rep_k, 2, 256]) - b_scale_desc = TensorDescriptor.from_tensor(b_scale_5d, [1, rep_n, rep_k, 2, 256]) - - # --- 6) dequantize back to fp32 for the reference (post-quant tensors) --- - a_ref_fp4 = MXFP4Tensor(size=(M, K), device=device); a_ref_fp4.data = a_codes - b_ref_fp4 = MXFP4Tensor(size=(K, N), device=device); b_ref_fp4.data = b_codes_KN - - a_ref_f32, b_ref_f32 = generate_fake_quant(a_f32_pre, b_f32_pre) - - reference = a_ref_f32 @ b_ref_f32 if compute_reference else None - - configs = { - "BLOCK_SIZE_M": BLOCK_M, - "BLOCK_SIZE_N": BLOCK_N, - "BLOCK_SIZE_K": BLOCK_K, - "num_stages": 4, - "ELEM_PER_BYTE_A": ELEM_PER_BYTE_A, - "ELEM_PER_BYTE_B": ELEM_PER_BYTE_B, - "VEC_SIZE": VEC_SIZE, - } - # RETURN DESCRIPTORS (not raw tensors!) - return a_desc, a_scale_desc, b_desc, b_scale_desc, rep_m, rep_n, rep_k, configs, reference - - -def validate_block_scaled(M, N, K, block_scale_type="nvfp4", input_output_dtype=torch.bfloat16): - a_desc, a_scale, b_desc, b_scale, rep_m, rep_n, rep_k, configs, reference = initialize_block_scaled( - M, N, K, block_scale_type, compute_reference=True) - output = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, input_output_dtype, M, N, K, rep_m, rep_n, rep_k, configs) - torch.testing.assert_close(reference, output.to(input_output_dtype), atol=1e-3, rtol=1e-3) - print(f"✅ (pass {block_scale_type})") - - -def bench_block_scaled(K, block_scale_type="nvfp4", reps=10): - assert K % 128 == 0 - M = 8192 - N = 8192 - print(f"Problem Shape = {M}x{N}x{K}") - - a_desc, a_scale, b_desc, b_scale, rep_m, rep_n, rep_k, configs, _ = initialize_block_scaled( - M, N, K, block_scale_type, compute_reference=False) - _ = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float16, M, N, K, rep_m, rep_n, rep_k, configs) - - proton.activate(0) - for _ in range(reps): - _ = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float16, M, N, K, rep_m, rep_n, rep_k, configs) - proton.deactivate(0) - print("Done benchmarking") - - -def show_profile(profile_name): - import triton.profiler.viewer as proton_viewer - - metric_names = ["time/ms"] - metric_names = ["tflop/s"] + metric_names - file_name = f"{profile_name}.hatchet" - tree, metrics = proton_viewer.parse(metric_names, file_name) - proton_viewer.print_tree(tree, metrics) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-K", type=int, required=False, default=512) - parser.add_argument("--K_range", type=int, nargs=2) - parser.add_argument("--K_step", type=int, default=512) - parser.add_argument("--bench", action="store_true", default=True) - parser.add_argument("--format", type=str, choices=["nvfp4"], default="nvfp4") - parser.add_argument("--input_output_dtype", type=str, choices=["float32", "float16", "bfloat16"], default="bfloat16") - args = parser.parse_args() - - if not supports_block_scaling(): - print("⛔ This example requires GPU support for block scaled matmul") - else: - if args.K and args.K_range is None: - args.K_range = [args.K, args.K] - args.K_step = 1 # doesn't matter as long as it's not 0 - - torch.manual_seed(42) - dtype =torch.float32 if args.input_output_dtype == "float32" else torch.float16 if args.input_output_dtype == "float16" else torch.bfloat16 - validate_block_scaled(8192, 8192, 8192, block_scale_type=args.format, input_output_dtype=dtype) - - - if args.bench: - proton.start("block_scaled_matmul", hook="triton") - proton.deactivate(0) # Skip argument creation - for K in range(args.K_range[0], args.K_range[1] + 1, args.K_step): - bench_block_scaled(K, reps=10000, block_scale_type=args.format) - proton.finalize() - show_profile("block_scaled_matmul") \ No newline at end of file diff --git a/triton_original.py b/triton_original.py deleted file mode 100644 index b52103e6e..000000000 --- a/triton_original.py +++ /dev/null @@ -1,302 +0,0 @@ -import argparse - -import torch -import triton -import triton.language as tl -import triton.profiler as proton -from triton.tools.tensor_descriptor import TensorDescriptor -from triton.tools.mxfp import MXFP4Tensor, MXScaleTensor - - -def is_cuda(): - return triton.runtime.driver.active.get_current_target().backend == "cuda" - - -def supports_block_scaling(): - return is_cuda() and torch.cuda.get_device_capability()[0] == 10 - - -def _matmul_launch_metadata(grid, kernel, args): - ret = {} - M, N, K = args["M"], args["N"], args["K"] - kernel_name = kernel.name - if "ELEM_PER_BYTE_A" and "ELEM_PER_BYTE_B" and "VEC_SIZE" in args: - if args["ELEM_PER_BYTE_A"] == 1 and args["ELEM_PER_BYTE_B"] == 1: - kernel_name += "_mxfp8" - elif args["ELEM_PER_BYTE_A"] == 1 and args["ELEM_PER_BYTE_B"] == 2: - kernel_name += "_mixed" - elif args["ELEM_PER_BYTE_A"] == 2 and args["ELEM_PER_BYTE_B"] == 2: - if args["VEC_SIZE"] == 16: - kernel_name += "_nvfp4" - elif args["VEC_SIZE"] == 32: - kernel_name += "_mxfp4" - ret["name"] = f"{kernel_name} [M={M}, N={N}, K={K}]" - ret["flops"] = 2.0 * M * N * K - return ret - - -@triton.jit(launch_metadata=_matmul_launch_metadata) -def block_scaled_matmul_kernel( # - a_desc, # - a_scale_desc, # - b_desc, # - b_scale_desc, # - c_desc, # - M: tl.constexpr, # - N: tl.constexpr, # - K: tl.constexpr, # - output_type: tl.constexpr, # - ELEM_PER_BYTE_A: tl.constexpr, # - ELEM_PER_BYTE_B: tl.constexpr, # - VEC_SIZE: tl.constexpr, # - BLOCK_M: tl.constexpr, # - BLOCK_N: tl.constexpr, # - BLOCK_K: tl.constexpr, # - rep_m: tl.constexpr, # - rep_n: tl.constexpr, # - rep_k: tl.constexpr, # - NUM_STAGES: tl.constexpr, # -): # - if output_type == 0: - output_dtype = tl.float32 - elif output_type == 1: - output_dtype = tl.float16 - elif output_type == 2: - output_dtype = tl.float8e4nv - - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_M) - pid_m = pid % num_pid_m - pid_n = pid // num_pid_m - offs_am = pid_m * BLOCK_M - offs_bn = pid_n * BLOCK_N - offs_k_a = 0 - offs_k_b = 0 - offs_scale_m = pid_m * rep_m - offs_scale_n = pid_n * rep_n - offs_scale_k = 0 - - MIXED_PREC: tl.constexpr = ELEM_PER_BYTE_A == 1 and ELEM_PER_BYTE_B == 2 - - accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k in tl.range(0, tl.cdiv(K, BLOCK_K), num_stages=NUM_STAGES): - a = a_desc.load([offs_am, offs_k_a]) - b = b_desc.load([offs_bn, offs_k_b]) - scale_a = a_scale_desc.load([0, offs_scale_m, offs_scale_k, 0, 0]) - scale_b = b_scale_desc.load([0, offs_scale_n, offs_scale_k, 0, 0]) - - scale_a = scale_a.reshape(rep_m, rep_k, 32, 4, 4).trans(0, 3, 2, 1, 4).reshape(BLOCK_M, BLOCK_K // VEC_SIZE) - scale_b = scale_b.reshape(rep_n, rep_k, 32, 4, 4).trans(0, 3, 2, 1, 4).reshape(BLOCK_N, BLOCK_K // VEC_SIZE) - - if MIXED_PREC: - accumulator = tl.dot_scaled(a, scale_a, "e4m3", b.T, scale_b, "e2m1", accumulator) - elif ELEM_PER_BYTE_A == 2 and ELEM_PER_BYTE_B == 2: - accumulator = tl.dot_scaled(a, scale_a, "e2m1", b.T, scale_b, "e2m1", accumulator) - else: - accumulator = tl.dot_scaled(a, scale_a, "e4m3", b.T, scale_b, "e4m3", accumulator) - - offs_k_a += BLOCK_K // ELEM_PER_BYTE_A - offs_k_b += BLOCK_K // ELEM_PER_BYTE_B - offs_scale_k += rep_k - - c_desc.store([offs_am, offs_bn], accumulator.to(output_dtype)) - - -def block_scaled_matmul(a_desc, a_scale_desc, b_desc, b_scale_desc, dtype_dst, M, N, K, rep_m, rep_n, rep_k, configs): - output = torch.empty((M, N), dtype=dtype_dst, device="cuda") - if dtype_dst == torch.float32: - dtype_dst = 0 - elif dtype_dst == torch.float16: - dtype_dst = 1 - elif dtype_dst == torch.float8_e4m3fn: - dtype_dst = 2 - else: - raise ValueError(f"Unsupported dtype: {dtype_dst}") - - BLOCK_M = configs["BLOCK_SIZE_M"] - BLOCK_N = configs["BLOCK_SIZE_N"] - c_desc = TensorDescriptor.from_tensor(output, [BLOCK_M, BLOCK_N]) - - grid = (triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N), 1) - block_scaled_matmul_kernel[grid]( - a_desc, - a_scale_desc, - b_desc, - b_scale_desc, - c_desc, - M, - N, - K, - dtype_dst, - configs["ELEM_PER_BYTE_A"], - configs["ELEM_PER_BYTE_B"], - configs["VEC_SIZE"], - configs["BLOCK_SIZE_M"], - configs["BLOCK_SIZE_N"], - configs["BLOCK_SIZE_K"], - rep_m, - rep_n, - rep_k, - configs["num_stages"], - ) - return output - - -def initialize_block_scaled(M, N, K, block_scale_type="nvfp4", compute_reference=False): - BLOCK_M = 128 - BLOCK_N = 256 - BLOCK_K = 256 if "fp4" in block_scale_type else 128 - VEC_SIZE = 16 if block_scale_type == "nvfp4" else 32 - assert block_scale_type in ["nvfp4", "mxfp4", "mxfp8", "mixed"], f"Invalid block scale type: {block_scale_type}" - ELEM_PER_BYTE_A = 2 if "fp4" in block_scale_type else 1 - ELEM_PER_BYTE_B = 1 if block_scale_type == "mxfp8" else 2 - - device = "cuda" - a_ref = MXFP4Tensor(size=(M, K), device=device).random() - # Similar to Hopper's wgmma symmetric fp8 instruction, the RHS is expected - # to be in col-major layout for Blackwell's tcgen05.mma when using fp4 operands. - # To conform to the expected semantics of tl.dot_scaled, (M, K) x (K, N), - # the data is generated in col-major layout, packed along K for fp4, and then - # logically transposed. Note that if one operand is of fp8 precision, unlike Hopper, - # Blackwell supports both row-major and col-major layouts for the RHS matrix. - # For the mixed-precision case, the fp4 RHS can be either in row or col-major layout. - # But for performance reason, it is recommended to use col-major layout. If TMA is used - # for the fp4 RHS operand load in mixed-precision dot, as in this tutorial, it must be - # in col-major layout. - b_ref = MXFP4Tensor(size=(N, K), device=device).random() - if block_scale_type in ["mxfp8", "mixed"]: - a_ref = a_ref.to(torch.float32) - a = a_ref.to(torch.float8_e4m3fn) - else: - # Pack two fp4 elements per byte along K - a = a_ref.to_packed_tensor(dim=1) - - if block_scale_type == "mxfp8": - b_ref = b_ref.to(torch.float32) - b = b_ref.to(torch.float8_e4m3fn) - else: - b = b_ref.to_packed_tensor(dim=1) - - b_ref = b_ref.to(torch.float32).T - - a_desc = TensorDescriptor.from_tensor(a, [BLOCK_M, BLOCK_K // ELEM_PER_BYTE_A]) - b_desc = TensorDescriptor.from_tensor(b, [BLOCK_N, BLOCK_K // ELEM_PER_BYTE_B]) - - a_scale_shape = [M // 128, K // VEC_SIZE // 4, 32, 16] - b_scale_shape = [N // 128, K // VEC_SIZE // 4, 32, 16] - epsilon = 1e-8 - a_scale = torch.rand(a_scale_shape, device=device) + epsilon - b_scale = torch.rand(b_scale_shape, device=device) + epsilon - if block_scale_type == "nvfp4": - a_scale = a_scale.to(torch.float8_e4m3fn) - b_scale = b_scale.to(torch.float8_e4m3fn) - a_scale_ref = a_scale - b_scale_ref = b_scale - elif block_scale_type in ["mxfp4", "mxfp8", "mixed"]: - a_scale_ref = MXScaleTensor(a_scale) - b_scale_ref = MXScaleTensor(b_scale) - a_scale = a_scale_ref.data - b_scale = b_scale_ref.data - - rep_m = BLOCK_M // 128 - rep_n = BLOCK_N // 128 - rep_k = BLOCK_K // VEC_SIZE // 4 - - # Use 5D TMA descriptor [1, rep_m, rep_k, 2, 256] with uint8 elements. - # With 256 elements we better utilize the L2 and don't require the TMA - # engine to emit many small messages (16B) messages as with 32x16xu8. - a_scale_block_shape = [1, rep_m, rep_k, 2, 256] - b_scale_block_shape = [1, rep_n, rep_k, 2, 256] - a_scale = a_scale.reshape(1, a_scale_shape[0], a_scale.shape[1], 2, 256) - b_scale = b_scale.reshape(1, b_scale_shape[0], b_scale.shape[1], 2, 256) - a_scale_desc = TensorDescriptor.from_tensor(a_scale, block_shape=a_scale_block_shape) - b_scale_desc = TensorDescriptor.from_tensor(b_scale, block_shape=b_scale_block_shape) - - reference = None - if compute_reference: - a_scale_ref = a_scale_ref.to(torch.float32) - b_scale_ref = b_scale_ref.to(torch.float32) - - def unpack_scale(packed): - packed = packed.reshape(*packed.shape[:-2], 32, 4, 4) - num_chunk_m, num_chunk_k, _, _, _ = packed.shape - return packed.permute(0, 3, 2, 1, 4).reshape(num_chunk_m * 128, num_chunk_k * 4).contiguous() - - a_scale_ref = unpack_scale(a_scale_ref).repeat_interleave(VEC_SIZE, dim=1)[:M, :K] - b_scale_ref = unpack_scale(b_scale_ref).repeat_interleave(VEC_SIZE, dim=1).T.contiguous()[:K, :N] - reference = torch.matmul(a_ref.to(torch.float32) * a_scale_ref, b_ref * b_scale_ref) - - configs = { - "BLOCK_SIZE_M": BLOCK_M, - "BLOCK_SIZE_N": BLOCK_N, - "BLOCK_SIZE_K": BLOCK_K, - "num_stages": 4, - "ELEM_PER_BYTE_A": ELEM_PER_BYTE_A, - "ELEM_PER_BYTE_B": ELEM_PER_BYTE_B, - "VEC_SIZE": VEC_SIZE, - } - return a_desc, a_scale_desc, b_desc, b_scale_desc, rep_m, rep_n, rep_k, configs, reference - - -def validate_block_scaled(M, N, K, block_scale_type="nvfp4"): - a_desc, a_scale, b_desc, b_scale, rep_m, rep_n, rep_k, configs, reference = initialize_block_scaled( - M, N, K, block_scale_type, compute_reference=True) - output = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float16, M, N, K, rep_m, rep_n, rep_k, configs) - torch.testing.assert_close(reference, output.to(torch.float32), atol=1e-3, rtol=1e-3) - print(f"✅ (pass {block_scale_type})") - - -def bench_block_scaled(K, block_scale_type="nvfp4", reps=10): - assert K % 128 == 0 - M = 8192 - N = 8192 - print(f"Problem Shape = {M}x{N}x{K}") - - a_desc, a_scale, b_desc, b_scale, rep_m, rep_n, rep_k, configs, _ = initialize_block_scaled( - M, N, K, block_scale_type, compute_reference=False) - _ = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float16, M, N, K, rep_m, rep_n, rep_k, configs) - - proton.activate(0) - for _ in range(reps): - _ = block_scaled_matmul(a_desc, a_scale, b_desc, b_scale, torch.float16, M, N, K, rep_m, rep_n, rep_k, configs) - proton.deactivate(0) - print("Done benchmarking") - - -def show_profile(profile_name): - import triton.profiler.viewer as proton_viewer - - metric_names = ["time/ms"] - metric_names = ["tflop/s"] + metric_names - file_name = f"{profile_name}.hatchet" - tree, metrics = proton_viewer.parse(metric_names, file_name) - proton_viewer.print_tree(tree, metrics) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("-K", type=int, required=False, default=8192) - parser.add_argument("--K_range", type=int, nargs=2) - parser.add_argument("--K_step", type=int, default=512) - parser.add_argument("--bench", action="store_true", default=True) - parser.add_argument("--format", type=str, choices=["mxfp4", "nvfp4", "mxfp8", "mixed"], default="nvfp4") - args = parser.parse_args() - - if not supports_block_scaling(): - print("⛔ This example requires GPU support for block scaled matmul") - else: - if args.K and args.K_range is None: - args.K_range = [args.K, args.K] - args.K_step = 1 # doesn't matter as long as it's not 0 - - torch.manual_seed(42) - - validate_block_scaled(8192, 8192, 8192, block_scale_type=args.format) - - if args.bench: - proton.start("block_scaled_matmul", hook="triton") - proton.deactivate(0) # Skip argument creation - for K in range(args.K_range[0], args.K_range[1] + 1, args.K_step): - bench_block_scaled(K, reps=10000, block_scale_type=args.format) - proton.finalize() - show_profile("block_scaled_matmul") \ No newline at end of file