From 40ce2d72f5c74f199033643bb67f0a7c9c280a58 Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 30 Dec 2025 04:23:10 -0600 Subject: [PATCH] [kernel] add turbodiffusion kernels (#972) --- .gitmodules | 3 + docker/Dockerfile.python3.10 | 2 +- docker/Dockerfile.python3.11 | 2 +- docker/Dockerfile.python3.12 | 2 +- docker/Dockerfile.python3.12.cuda12.9.1 | 2 +- fastvideo-kernel/CMakeLists.txt | 12 +- fastvideo-kernel/README.md | 28 +- fastvideo-kernel/build.sh | 12 +- fastvideo-kernel/csrc/common_extension.cpp | 12 + .../csrc/turbodiffusion/common/common.hpp | 85 +++ .../csrc/turbodiffusion/common/launch.hpp | 44 ++ .../csrc/turbodiffusion/common/load.hpp | 75 ++ .../csrc/turbodiffusion/common/store.hpp | 77 ++ .../csrc/turbodiffusion/gemm/gemm.cu | 77 ++ .../csrc/turbodiffusion/gemm/kernel.hpp | 522 +++++++++++++ .../csrc/turbodiffusion/gemm/launch.hpp | 64 ++ .../csrc/turbodiffusion/gemm/utils.hpp | 129 ++++ .../csrc/turbodiffusion/norm/layernorm.cu | 63 ++ .../csrc/turbodiffusion/norm/layernorm.hpp | 202 +++++ .../csrc/turbodiffusion/norm/rmsnorm.cu | 59 ++ .../csrc/turbodiffusion/norm/rmsnorm.hpp | 147 ++++ .../csrc/turbodiffusion/quant/quant.cu | 75 ++ .../csrc/turbodiffusion/quant/quant.hpp | 195 +++++ fastvideo-kernel/include/cutlass | 1 + .../python/fastvideo_kernel/__init__.py | 13 + .../fastvideo_kernel/turbodiffusion_ops.py | 716 ++++++++++++++++++ fastvideo-kernel/tests/test_turbodiffusion.py | 294 +++++++ fastvideo/tests/modal/pr_test.py | 4 + 28 files changed, 2893 insertions(+), 24 deletions(-) create mode 100644 fastvideo-kernel/csrc/turbodiffusion/common/common.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/common/launch.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/common/load.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/common/store.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu create mode 100644 fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/gemm/launch.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/gemm/utils.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu create mode 100644 fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu create mode 100644 fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp create mode 100644 fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu create mode 100644 fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp create mode 160000 fastvideo-kernel/include/cutlass create mode 100644 fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py create mode 100644 fastvideo-kernel/tests/test_turbodiffusion.py diff --git a/.gitmodules b/.gitmodules index c51738164..d0d7c9b47 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +1,6 @@ [submodule "fastvideo-kernel/include/tk"] path = fastvideo-kernel/include/tk url = https://github.com/HazyResearch/ThunderKittens.git +[submodule "fastvideo-kernel/include/cutlass"] + path = fastvideo-kernel/include/cutlass + url = https://github.com/NVIDIA/cutlass.git diff --git a/docker/Dockerfile.python3.10 b/docker/Dockerfile.python3.10 index 39e3600a6..98a1e0bd7 100644 --- a/docker/Dockerfile.python3.10 +++ b/docker/Dockerfile.python3.10 @@ -60,7 +60,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh --release + ./build.sh EXPOSE 22 diff --git a/docker/Dockerfile.python3.11 b/docker/Dockerfile.python3.11 index d8c0d2b35..fe3639a60 100644 --- a/docker/Dockerfile.python3.11 +++ b/docker/Dockerfile.python3.11 @@ -60,7 +60,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh --release + ./build.sh EXPOSE 22 diff --git a/docker/Dockerfile.python3.12 b/docker/Dockerfile.python3.12 index baa983423..64c0103c8 100644 --- a/docker/Dockerfile.python3.12 +++ b/docker/Dockerfile.python3.12 @@ -60,6 +60,6 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh --release + ./build.sh EXPOSE 22 diff --git a/docker/Dockerfile.python3.12.cuda12.9.1 b/docker/Dockerfile.python3.12.cuda12.9.1 index 764995b05..ce90131d8 100644 --- a/docker/Dockerfile.python3.12.cuda12.9.1 +++ b/docker/Dockerfile.python3.12.cuda12.9.1 @@ -60,7 +60,7 @@ RUN source $HOME/.local/bin/env && \ source /opt/venv/bin/activate && \ cd fastvideo-kernel && \ git submodule update --init --recursive && \ - ./build.sh --release + ./build.sh EXPOSE 22 diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index 356cf8e06..6c433b936 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -31,9 +31,11 @@ find_package(Torch REQUIRED) # Include directories include_directories( ${CMAKE_SOURCE_DIR}/include + ${CMAKE_SOURCE_DIR}/include/cutlass/include ${CMAKE_SOURCE_DIR}/include/tk/include ${CMAKE_SOURCE_DIR}/include/tk/prototype ${CMAKE_SOURCE_DIR}/csrc + ${CMAKE_SOURCE_DIR}/csrc/turbodiffusion ${TORCH_INCLUDE_DIRS} ) @@ -117,8 +119,14 @@ endif() if(BUILD_CXX_KERNELS) # Source files - set(EXTENSION_SOURCES csrc/common_extension.cpp) - + set(EXTENSION_SOURCES + csrc/common_extension.cpp + csrc/turbodiffusion/gemm/gemm.cu + csrc/turbodiffusion/norm/rmsnorm.cu + csrc/turbodiffusion/norm/layernorm.cu + csrc/turbodiffusion/quant/quant.cu + ) + # Conditionally add TK kernels if(ENABLE_TK_KERNELS) list(APPEND EXTENSION_SOURCES diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index 32f36a992..262788f0b 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -13,15 +13,6 @@ cd fastvideo-kernel ./build.sh ``` -### Release Build (Force Enable Kernels) -If you are building a release wheel or docker image on a machine without a GPU (e.g., CI/CD), you can force-enable the compilation of Hopper-specific ThunderKittens kernels. - -```bash -cd fastvideo-kernel -./build.sh --release -``` -*Note: The resulting wheel will contain kernels that require an H100 GPU to run, but can be built on any machine with CUDA 12.3+ toolchain.* - ### Rocm Build If you are in a rocm environment without the compilation toolchaine of CUDA. @@ -32,6 +23,10 @@ cd fastvideo-kernel ## Usage +### Sliding Tile Attention (STA) & Video Sparse Attention (VSA) + +For detailed usage, please check the [Attention Documentation](../docs/attention/index.md). + ```python from fastvideo_kernel import sliding_tile_attention, video_sparse_attn, moba_attn_varlen @@ -45,6 +40,10 @@ out = video_sparse_attn(q, k, v, block_sizes, topk=5) out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...) ``` +### TurboDiffusion Kernels + +This package also includes kernels from [TurboDiffusion](https://github.com/thu-ml/TurboDiffusion), including INT8 GEMM, Quantization, RMSNorm and LayerNorm. + ## Requirements - **Runtime**: @@ -57,3 +56,14 @@ out = moba_attn_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, ...) ## Acknowledgement This package structure and build system are based on [sgl-kernel](https://github.com/sgl-project/sglang/tree/main/sgl-kernel) from the SGLang project. + +The implementation of `turbodiffusion` kernels is adapted from [TurboDiffusion](https://github.com/thu-ml/TurboDiffusion). If you use these kernels, please cite: + +```bibtex +@article{zhang2025turbodiffusion, + title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times}, + author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun}, + journal={arXiv preprint arXiv:2512.16093}, + year={2025} +} +``` diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index 8d25cd180..9437a87d3 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -18,21 +18,15 @@ RELEASE=0 GPU_BACKEND=CUDA for arg in "$@"; do case "$arg" in - --release|-r) - RELEASE=1 - ;; --rocm) GPU_BACKEND=ROCM ;; esac done -if [ "$RELEASE" -eq 1 ]; then - # Force-enable ThunderKittens kernels and compile for Hopper. - # Intended for producing release wheels/images on machines without a GPU. - export TORCH_CUDA_ARCH_LIST="9.0a" - export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a" -fi +# Force-enable ThunderKittens kernels and compile for Hopper. +export TORCH_CUDA_ARCH_LIST="9.0a" +export CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON -DCMAKE_CUDA_ARCHITECTURES=90a" export CMAKE_ARGS="${CMAKE_ARGS:-} -DGPU_BACKEND=${GPU_BACKEND}" diff --git a/fastvideo-kernel/csrc/common_extension.cpp b/fastvideo-kernel/csrc/common_extension.cpp index 78f07007a..ad27bb26b 100644 --- a/fastvideo-kernel/csrc/common_extension.cpp +++ b/fastvideo-kernel/csrc/common_extension.cpp @@ -22,6 +22,12 @@ extern std::vector block_sparse_attention_backward( ); #endif +// TurboDiffusion kernels +void register_quant(pybind11::module_ &); +void register_rms_norm(pybind11::module_ &); +void register_layer_norm(pybind11::module_ &); +void register_gemm(pybind11::module_ &); + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "FastVideo CUDA Kernels"; @@ -33,4 +39,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("block_sparse_fwd", torch::wrap_pybind_function(block_sparse_attention_forward), "block sparse attention forward (Hopper)"); m.def("block_sparse_bwd", torch::wrap_pybind_function(block_sparse_attention_backward), "block sparse attention backward (Hopper)"); #endif + + // TurboDiffusion + register_quant(m); + register_rms_norm(m); + register_layer_norm(m); + register_gemm(m); } diff --git a/fastvideo-kernel/csrc/turbodiffusion/common/common.hpp b/fastvideo-kernel/csrc/turbodiffusion/common/common.hpp new file mode 100644 index 000000000..5793fe1c6 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/common/common.hpp @@ -0,0 +1,85 @@ +#pragma once + +#include +#include + +CUTLASS_HOST_DEVICE int64_t cdiv(int64_t const& a, int64_t const &b) { + return (a + b - 1) / b; +} + +template +CUTLASS_HOST_DEVICE T max(T a, T b) { return a > b ? a : b; } + +template +CUTLASS_HOST_DEVICE T min(T a, T b) { return a > b ? b : a; } + +#define MIN(a, b) ((a) > (b) ? (b) : (a)) +#define MAX(a, b) ((a) > (b) ? (a) : (b)) + +#define BOOL_SWITCH(COND, CONST_NAME, ...) \ +[&] { \ + if (COND) { \ + static constexpr bool CONST_NAME = true; \ + return (__VA_ARGS__)(); \ + } else { \ + static constexpr bool CONST_NAME = false; \ + return (__VA_ARGS__)(); \ + } \ +}() + +#define CUDA_CHECK(call) \ +{ \ + cudaError_t err = call; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA Error at %s:%d: %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \ + exit(err); \ + } \ +} + +#define CONFIG_SWITCH(N, ...) \ +[&] { \ + if (N <= 1024) { \ + constexpr int NUM_THR_PER_CTA = 128; \ + constexpr int MAX_HIDDEN_SIZE = 1024; \ + return (__VA_ARGS__)(); \ + } else if (N <= 2048) { \ + constexpr int NUM_THR_PER_CTA = 128; \ + constexpr int MAX_HIDDEN_SIZE = 2048; \ + return (__VA_ARGS__)(); \ + } else if (N <= 4096) { \ + constexpr int NUM_THR_PER_CTA = 128; \ + constexpr int MAX_HIDDEN_SIZE = 4096; \ + return (__VA_ARGS__)(); \ + } else if (N <= 8192) { \ + constexpr int NUM_THR_PER_CTA = 256; \ + constexpr int MAX_HIDDEN_SIZE = 8192; \ + return (__VA_ARGS__)(); \ + } else { \ + constexpr int NUM_THR_PER_CTA = 256; \ + constexpr int MAX_HIDDEN_SIZE = 16384; \ + return (__VA_ARGS__)(); \ + } \ +}() + + +template + void create_tensor( + torch::Device const &device, + std::optional &output, + std::optional &scale, + int m, int n + ) { + int num_block_m = cdiv(m, BlockSize); + int num_block_n = cdiv(n, BlockSize); + if (!output.has_value()) { + output.emplace(torch::empty( + {m, n}, + torch::TensorOptions().device(device).dtype(torch::kInt8) + )); + scale.emplace(torch::empty( + {num_block_m, num_block_n}, + torch::TensorOptions().device(device).dtype(torch::kFloat32) + )); + } + } + diff --git a/fastvideo-kernel/csrc/turbodiffusion/common/launch.hpp b/fastvideo-kernel/csrc/turbodiffusion/common/launch.hpp new file mode 100644 index 000000000..b7b2dce22 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/common/launch.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include +#include + + +template +__global__ void device_kernel( + __grid_constant__ typename Kernel::Params const params +) { + extern __shared__ char smem[]; + Kernel op; + op(params, smem); +} + +template +__global__ __launch_bounds__(Kernel::MaxThreadsPerBlock, Kernel::MinBlocksPerMultiprocessor) +void device_kernel_with_launch_bounds( + __grid_constant__ typename Kernel::Params const params +) { + extern __shared__ char smem[]; + Kernel op; + op(params, smem); +} + +template +void launch_kernel( + typename Kernel::Params const ¶ms, + dim3 grid_shape, + dim3 cta_shape, + size_t ShmSize, + cudaStream_t stream = nullptr +) { + auto func = device_kernel; + if (ShmSize >= 48 * 1024) { + CUDA_CHECK(cudaFuncSetAttribute( + func, + cudaFuncAttributeMaxDynamicSharedMemorySize, + ShmSize + )); + } + func<<>>(params); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/fastvideo-kernel/csrc/turbodiffusion/common/load.hpp b/fastvideo-kernel/csrc/turbodiffusion/common/load.hpp new file mode 100644 index 000000000..f72bcad6c --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/common/load.hpp @@ -0,0 +1,75 @@ +#pragma once + +template < + class InputDtype_, + int TileM_, + int TileN_, + int NumThrPerCta_, + bool IsEvenM, + bool IsEvenN +> +class Loader { +public: + using InputDtype = InputDtype_; + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int NumThrPerCta = NumThrPerCta_; + static constexpr int NumElementPerThread = TileM * TileN / NumThrPerCta; + static constexpr int NumThrPerRow = TileN / NumElementPerThread; + + static_assert(NumThrPerCta % TileM == 0); + static_assert(TileM * TileN % NumThrPerCta == 0); + + + CUTLASS_DEVICE void + load(void const *input_ptr, void *thr_output_reg, int64_t m, int64_t n, int blk_m, int blk_n, int tid) { + int n_alignment = (n & 31) * sizeof(InputDtype); + int thr_m_offset = tid / NumThrPerRow; + int thr_n_offset = (tid % NumThrPerRow) * NumElementPerThread; + void const *cta_input_ptr = (void*)((InputDtype*)input_ptr + blk_m * TileM * n + blk_n * TileN); + void const *thr_input_ptr = (void*)((InputDtype*)cta_input_ptr + thr_m_offset * n + thr_n_offset); + InputDtype tmp_reg[NumElementPerThread]; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + tmp_reg[i] = InputDtype(0.f); + bool pred = IsEvenM ? true : thr_m_offset + blk_m * TileM < m; + int limit = IsEvenN ? NumElementPerThread : MIN(NumElementPerThread, n - (blk_n * TileN + thr_n_offset)); + if (n_alignment % 128 == 0) + _load(thr_input_ptr, (void*)tmp_reg, limit, pred); + else if (n_alignment % 64 == 0) + _load(thr_input_ptr, (void*)tmp_reg, limit, pred); + else + _load(thr_input_ptr, (void*)tmp_reg, limit, pred); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + *((float*)thr_output_reg + i) = static_cast(reinterpret_cast(tmp_reg[i])); + } + +private: + template + CUTLASS_DEVICE void + _load(void const *thr_input_ptr, void *thr_output_reg, int limit, bool pred) { + static constexpr int NumElementPerLoad = sizeof(LoadDataType) / sizeof(InputDtype); + if (pred) { + if constexpr (IsEven) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; i += NumElementPerLoad) { + *(LoadDataType*)((InputDtype*)thr_output_reg + i) = *(LoadDataType*)((InputDtype*)thr_input_ptr + i); + } + } else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < limit; i += NumElementPerLoad) { + if (limit - i > NumElementPerLoad) + *(LoadDataType*)((InputDtype*)thr_output_reg + i) = *(LoadDataType*)((InputDtype*)thr_input_ptr + i); + else { + for (int j = 0; j < NumElementPerLoad; ++j) { + if (i + j < limit) + *((InputDtype*)thr_output_reg + i + j) = *((InputDtype*)thr_input_ptr + i + j); + } + } + } + } + } + } +}; diff --git a/fastvideo-kernel/csrc/turbodiffusion/common/store.hpp b/fastvideo-kernel/csrc/turbodiffusion/common/store.hpp new file mode 100644 index 000000000..29c3b213a --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/common/store.hpp @@ -0,0 +1,77 @@ +#pragma once +#include "common/common.hpp" + + +template < + class OutputDtype_, + int TileM_, + int TileN_, + int NumThrPerCta_, + bool IsEvenM, + bool IsEvenN, + bool Round = true, + bool SaveScale = true +> +class Saver { +public: + using OutputDtype = OutputDtype_; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int NumThrPerCta = NumThrPerCta_; + static constexpr int NumElementPerThread = TileM * TileN / NumThrPerCta; + static constexpr int NumThrPerRow = TileN / NumElementPerThread; + + static_assert(TileM * TileN % NumThrPerCta == 0); + static_assert(NumThrPerCta % TileM == 0); + + CUTLASS_DEVICE void + store(void *Optr, void *OSptr, void *reg, float scale_inv, int64_t m, int64_t n, int blk_m, int blk_n, int tid) { + int n_alignment = (n & 31) * sizeof(OutputDtype); + int thr_m_offset = tid / NumThrPerRow; + int thr_n_offset = (tid % NumThrPerRow) * NumElementPerThread; + void *cta_output_ptr = (void*)((OutputDtype*)Optr + blk_m * TileM * (Round ? cdiv(n, TileN) * TileN : n) + blk_n * TileN); + void *thr_output_ptr = (void*)((OutputDtype*)cta_output_ptr + thr_m_offset * (Round ? cdiv(n, TileN) * TileN : n) + thr_n_offset); + bool pred = IsEvenM ? true : thr_m_offset + blk_m * TileM < m; + int limit = IsEvenN ? NumElementPerThread : MIN(NumElementPerThread, n - (blk_n * TileN + thr_n_offset)); + if (n_alignment % 128 == 0) + _store(thr_output_ptr, reg, limit, pred); + else if (n_alignment % 64 == 0) + _store(thr_output_ptr, reg, limit, pred); + else + _store(thr_output_ptr, reg, limit, pred); + + if constexpr (SaveScale) { + if (tid == 0) { + *((float*)OSptr + blk_m * cdiv(n, TileN)+ blk_n) = scale_inv; + } + } + } + +private: + template + CUTLASS_DEVICE void + _store(void *thr_output_ptr, void *reg, int limit, bool pred) { + static constexpr int NumElementPerStore = sizeof(StoreDataType) / sizeof(OutputDtype); + if (pred) { + if constexpr (IsEven) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; i += NumElementPerStore) { + *(StoreDataType*)((OutputDtype*)thr_output_ptr + i) = *(StoreDataType*)((OutputDtype*)reg + i); + } + } else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < limit; i += NumElementPerStore) { + if (limit - i > NumElementPerStore) + *(StoreDataType*)((OutputDtype*)thr_output_ptr + i) = *(StoreDataType*)((OutputDtype*)reg + i); + else { + for (int j = 0; j < limit - i; ++j) { + *((OutputDtype*)thr_output_ptr + i + j) = *((OutputDtype*)reg + i + j); + } + } + } + } + } + } + +}; diff --git a/fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu b/fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu new file mode 100644 index 000000000..175ad12a6 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/gemm/gemm.cu @@ -0,0 +1,77 @@ +/* + * Copyright (c) 2025 by TurboDiffusion team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * + * Citation (please cite if you use this code): + * + * @article{zhang2025turbodiffusion, + * title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times}, + * author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun}, + * journal={arXiv preprint arXiv:2512.16093}, + * year={2025} + * } + */ + +#include +#include +#include +#include +#include +#include +#include + +#include "common/common.hpp" +#include "gemm/launch.hpp" + +void int8_gemm( + at::Tensor const& A, at::Tensor const& A_S, + at::Tensor const& B, at::Tensor const& B_S, + torch::Tensor& C +) { + + + static constexpr int swizzle_dir = 1; + static constexpr int swizzle_size_log = 5; + + int k = B.size(1); + int m = A.size(0); + int n = B.size(0); + + switch (C.scalar_type()) { + case torch::kHalf:{ + int8_gemm_ ( + (int8_t*)A.data_ptr(), A_S.data_ptr(), + (int8_t*)B.data_ptr(), B_S.data_ptr(), + (cutlass::half_t*)C.data_ptr(), + m, n, k, swizzle_dir, swizzle_size_log, at::cuda::getCurrentCUDAStream().stream() + ); + break; + } + + case torch::kBFloat16:{ + int8_gemm_ ( + (int8_t*)A.data_ptr(), A_S.data_ptr(), + (int8_t*)B.data_ptr(), B_S.data_ptr(), + (cutlass::bfloat16_t*)C.data_ptr(), + m, n, k, swizzle_dir, swizzle_size_log, at::cuda::getCurrentCUDAStream().stream() + ); + break; + } + + default: { + std::cerr << "Observing: " << C.scalar_type() << " for the output datatype which is invalid"; + throw std::runtime_error("Unsupported output data type for int8 gemm."); + } + } + +} + +void register_gemm(pybind11::module_ &m) { + m.def("gemm_cuda", &int8_gemm); +} + + + + + diff --git a/fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp b/fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp new file mode 100644 index 000000000..7a0eef8dd --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/gemm/kernel.hpp @@ -0,0 +1,522 @@ +/* + * Copyright (c) 2025 by TurboDiffusion team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * + * Citation (please cite if you use this code): + * + * @article{zhang2025turbodiffusion, + * title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times}, + * author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun}, + * journal={arXiv preprint arXiv:2512.16093}, + * year={2025} + * } + */ + +#pragma once + +#include +#include "cute/tensor.hpp" + +#include "common/common.hpp" +#include "gemm/utils.hpp" + +using namespace cute; + +template < + class OutputDtype_, + bool IsEvenM, + bool IsEvenN +> +struct GemmKernel { + using ElementA = int8_t; + using ElementB = int8_t; + using OutputDtype = OutputDtype_; + using AccumulatorDtype = int32_t; + static constexpr int BlockSize = 128; + static constexpr int TileM = 128; + static constexpr int TileN = 128; + static constexpr int TileK = 128; + static constexpr int Stage = 3; + static constexpr int EpiStage = 2; + + static_assert( + BlockSize % TileM == 0 + && BlockSize % TileN == 0 + && BlockSize % TileK == 0 + ); + + static constexpr int NumTilePerBlock = BlockSize / TileK; + + using SmemLayoutAtom = decltype( + composition( + Swizzle<3, 4, 3>{}, + make_layout( + make_shape(Int<8>{}, Int{}), + make_stride(Int{}, Int<1>{}) + ) + ) + ); + + using SmemLayoutA = decltype( + tile_to_shape( + SmemLayoutAtom{}, + make_shape(Int{}, Int{}, Int{}) + ) + ); + + using SmemLayoutB = decltype( + tile_to_shape( + SmemLayoutAtom{}, + make_shape(Int{}, Int{}, Int{}) + ) + ); + + using MmaOP = cute::SM80_16x8x32_S32S8S8S32_TN; + using TiledMma = decltype( + make_tiled_mma( + MMA_Atom>{}, + make_layout(make_shape( + _4{}, _2{}, _1{} + )), + make_tile(Int<64>{}, Int<32>{}, Int<32>{}) + ) + ); + + using G2SCopyAtomA = Copy_Atom>, ElementA>; + using G2SCopyAtomB = Copy_Atom>, ElementB>; + using G2STiledCopyA = decltype( + make_tiled_copy( + G2SCopyAtomA{}, + make_layout( + make_shape(Int<64>{}, Int<4>{}), + make_stride(Int<4>{}, Int<1>{}) + ), + make_layout(make_shape(Int<1>{}, Int<16>{})) + ) + ); + using G2STiledCopyB = decltype( + make_tiled_copy( + G2SCopyAtomB{}, + make_layout( + make_shape(Int<64>{}, Int<4>{}), + make_stride(Int<4>{}, Int<1>{}) + ), + make_layout(make_shape(Int<1>{}, Int<16>{})) + ) + ); + + using S2RCopyAtomA = Copy_Atom, ElementA>; + using S2RCopyAtomB = Copy_Atom, ElementB>; + using S2RTiledCopyA = decltype(make_tiled_copy_A(S2RCopyAtomA{}, TiledMma{})); + using S2RTiledCopyB = decltype(make_tiled_copy_B(S2RCopyAtomB{}, TiledMma{})); + + // epilogue + using SmemLayoutAtomD = decltype( + composition( + Swizzle<2, 3, 3>{}, + make_layout( + make_shape(Int<32>{}, Int<32>{}), + LayoutRight{} + ) + ) + ); + + using SmemLayoutD = decltype( + tile_to_shape( + SmemLayoutAtomD{}, + make_shape(Int<64>{}, Int<32>{}, Int{}) + ) + ); + + using R2SCopyAtomD = Copy_Atom>, OutputDtype>; + using R2STiledCopyD = decltype(make_tiled_copy_C(R2SCopyAtomD{}, TiledMma{})); + + using S2GCopyAtomD = Copy_Atom, OutputDtype>; + using S2GCopyD = decltype(make_tiled_copy( + S2GCopyAtomD{}, + make_layout(Shape<_64, _4>{}), + make_layout(Shape<_1, _8>{}) + )); + + using TileShape = decltype(make_shape(Int{}, Int{}, Int{})); + + struct SharedStorageAB: cute::aligned_struct<128> { + array_aligned, 128> smem_A; + array_aligned, 128> smem_B; + array_aligned smem_AS; + array_aligned smem_BS; + array_aligned smem_AF; + }; + + struct SharedStorageD: cute::aligned_struct<128> { + array_aligned> smem_D; + }; + + union SharedStorage { + SharedStorageAB storage_AB; + SharedStorageD storage_D; + }; + + + struct Params { + void const* Aptr; + void const* ASptr; + void const* Bptr; + void const* BSptr; + void* Dptr; + int64_t const m; + int64_t const n; + int64_t const k; + int const swizzle_dir; + int const swizzle_size; + }; + + using Arguments = Params; + + static constexpr int ThreadNum = size(TiledMma{}); + static constexpr int ShmSize = sizeof(SharedStorage); + static constexpr bool FastInt2Float = false; + + static bool can_implement(int64_t m, int64_t n, int64_t k) { + if (k % BlockSize != 0) return false; + if ((n * sizeof(OutputDtype)) % 16 != 0) + return false; + return true; + } + + static Params to_underlying_arguments(Arguments const& args) { + return args; + } + + static dim3 get_grid_size(int64_t m, int64_t n) { + return dim3(cdiv(m, TileM) * cdiv(n, TileN)); + } + + CUTLASS_HOST_DEVICE + static auto get_block_coord( + int64_t m_blocks, + int64_t n_blocks, + int const swizzle_dir, + int64_t const swizzle_size_log + ) { + int64_t blk_m; + int64_t blk_n; + + if (swizzle_dir == 1) + std::swap(m_blocks, n_blocks); + + if (swizzle_size_log == 0) { + blk_m = blockIdx.x % m_blocks; + blk_n = blockIdx.x / m_blocks; + } else { + int64_t group_size = n_blocks << swizzle_size_log; + int64_t num_groups = m_blocks >> swizzle_size_log; + int64_t group_idx = blockIdx.x / group_size; + int64_t local_idx = blockIdx.x % group_size; + if (group_idx == num_groups) { + blk_m = (num_groups << swizzle_size_log) + local_idx % (m_blocks - (num_groups << swizzle_size_log)); + blk_n = local_idx / (m_blocks - (num_groups << swizzle_size_log)); + } else { + blk_m = (local_idx & ((1LL << swizzle_size_log) - 1)) + (group_idx << swizzle_size_log); + blk_n = local_idx >> swizzle_size_log; + } + } + + if (swizzle_dir == 1) + std::swap(blk_m, blk_n); + + return make_coord(blk_m, blk_n); + } + + CUTLASS_DEVICE + void operator()( + Params const& params, char* smem_data + ) { + + SharedStorage& shared_storage = *reinterpret_cast(smem_data); + + auto t_idx = threadIdx.x; + + int64_t const m = params.m; + int64_t const n = params.n; + int64_t const k = params.k; + int const swizzle_dir = params.swizzle_dir; + int const swizzle_size = params.swizzle_size; + + Tensor A = make_tensor( + make_gmem_ptr(params.Aptr), + make_shape(m, k), + make_stride(k, _1{}) + ); + Tensor B = make_tensor( + make_gmem_ptr(params.Bptr), + make_shape(m, k), + make_stride(k, _1{}) + ); + Tensor AS = make_tensor( + make_gmem_ptr(params.ASptr), + make_shape(cdiv(m, BlockSize), cdiv(k, BlockSize)), + make_stride(cdiv(k, BlockSize), _1{}) + ); + Tensor BS = make_tensor( + make_gmem_ptr(params.BSptr), + make_shape(cdiv(n, BlockSize), cdiv(k, BlockSize)), + make_stride(cdiv(k, BlockSize), _1{}) + ); + Tensor D = make_tensor( + make_gmem_ptr(params.Dptr), + make_shape(m, n), + LayoutRight{} + ); + + auto [m_coord, n_coord] = get_block_coord( + cdiv(m, size<0>(TileShape{})), + cdiv(n, size<1>(TileShape{})), + swizzle_dir, swizzle_size + ); + + int32_t blk_m_coord = m_coord / (BlockSize / TileM); + int32_t blk_n_coord = n_coord / (BlockSize / TileN); + + // local tile + auto gA = local_tile(A, TileShape{}, make_coord(m_coord, n_coord, _), Step<_1, X, _1>{}); + auto gB = local_tile(B, TileShape{}, make_coord(m_coord, n_coord, _), Step{}); + auto gD = local_tile(D, TileShape{}, make_coord(m_coord, n_coord, _), Step<_1, _1, X>{}); + + // shared memory + Tensor sA = make_tensor( + make_smem_ptr(shared_storage.storage_AB.smem_A.data()), + SmemLayoutA{} + ); + + Tensor sB = make_tensor( + make_smem_ptr(shared_storage.storage_AB.smem_B.data()), + SmemLayoutB{} + ); + + // register + TiledMma tiled_mma; + auto thr_mma = tiled_mma.get_slice(t_idx); + auto tCrA = thr_mma.partition_fragment_A(gA(_, _, 0)); + auto tCrB = thr_mma.partition_fragment_B(gB(_, _, 0)); + auto tDrC = thr_mma.partition_fragment_C(gD); // mma accumulator + auto tDrD = make_tensor_like(tDrC); // float accumulator + + if constexpr (FastInt2Float) { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size(tDrC); ++i) + tDrC(i) = 0x4B400000; + } else { + clear(tDrC); + } + + clear(tDrD); + + + // global to shared copy + G2STiledCopyA g2s_tiled_copy_a; + auto g2s_thr_copy_a = g2s_tiled_copy_a.get_slice(t_idx); + auto tAgA = g2s_thr_copy_a.partition_S(gA); + auto tAsA = g2s_thr_copy_a.partition_D(sA); + auto cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA))); + auto tAcA = g2s_thr_copy_a.partition_S(cA); + int const m_limit = m - TileM * m_coord; + int const n_limit = n - TileN * n_coord; + + G2STiledCopyB g2s_tiled_copy_b; + auto g2s_thr_copy_b = g2s_tiled_copy_b.get_slice(t_idx); + auto tBgB = g2s_thr_copy_b.partition_S(gB); + auto tBsB = g2s_thr_copy_b.partition_D(sB); + auto cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB))); + auto tBcB = g2s_thr_copy_a.partition_S(cB); + + + // shared to register copy + S2RTiledCopyA s2r_tiled_copy_a; + auto s2r_thr_copy_a = s2r_tiled_copy_a.get_slice(t_idx); + auto tCsA = s2r_thr_copy_a.partition_S(sA); + auto tCrA_view = s2r_thr_copy_a.retile_D(tCrA); + + S2RTiledCopyB s2r_tiled_copy_b; + auto s2r_thr_copy_b = s2r_tiled_copy_b.get_slice(t_idx); + auto tCsB = s2r_thr_copy_b.partition_S(sB); + auto tCrB_view = s2r_thr_copy_b.retile_D(tCrB); + + // pipeline status + int64_t g2s_a_tile = 0; + int64_t g2s_b_tile = 0; + int g2s_a_smem = 0; + int g2s_b_smem = 0; + + int g2s_tile_in_block = 0; + int g2s_block = 0; // b block idx + + int s2r_a_smem = 0; + int s2r_b_smem = 0; + int s2r_tile_in_block = 0; + + int mma_block_a = 0; + int mma_block_b = 0; + + int ntile = k / TileK; + // load scale and fallback + // we assume all ptrs are 128bit aligned + // auto smem_fallback_A = raw_pointer_cast(make_smem_ptr(shared_storage.storage_AB.smem_AF.data())); + // auto smem_scale_A = raw_pointer_cast(make_smem_ptr(shared_storage.storage_AB.smem_AS.data())); + // auto smem_scale_B = raw_pointer_cast(make_smem_ptr(shared_storage.storage_AB.smem_BS.data())); + __syncthreads(); + + + int32_t fallbackA_load = 0; + int32_t fallbackA_mma = 0; + + // copy first Stage - 1 tile + CUTLASS_PRAGMA_UNROLL + for (int i = 0, _i = min(Stage - 1, ntile); i < _i; ++i) { + if (g2s_b_tile < ntile) { + g2s_tile_in_block = (g2s_tile_in_block + 1) % NumTilePerBlock; + copy_AB(g2s_tiled_copy_a, tAgA, tAsA, tAcA, g2s_a_tile, g2s_a_smem, m_limit); + copy_AB(g2s_tiled_copy_b, tBgB, tBsB, tBcB, g2s_b_tile, g2s_b_smem, n_limit); + ++g2s_b_tile; + ++g2s_b_smem; + ++g2s_block; + g2s_a_tile = g2s_block * NumTilePerBlock; + ++g2s_a_smem; + } + cp_async_fence(); + } + + constexpr int nk = size<2>(tCrA); + float scale_a = AS(blk_m_coord, 0); + float scale_b = BS(blk_n_coord, 0); + + CUTLASS_PRAGMA_NO_UNROLL + for (int64_t mma_b_tile = 0; mma_b_tile < ntile; ++mma_b_tile) { + s2r_tile_in_block = (s2r_tile_in_block + 1) % NumTilePerBlock; + cp_async_wait(); + __syncthreads(); + + // do mma first + CUTLASS_PRAGMA_UNROLL + for (int ik = 0; ik < nk; ++ik) { + cute::copy(s2r_tiled_copy_a, tCsA(_, _, ik, s2r_a_smem), + tCrA_view(_, _, ik)); + cute::copy(s2r_tiled_copy_b, tCsB(_, _, ik, s2r_b_smem), + tCrB_view(_, _, ik)); + cute::gemm(tiled_mma, tDrC, tCrA(_, _, ik), tCrB(_, _, ik), tDrC); + } + + // a s2r increase anyway + s2r_a_smem = (s2r_a_smem + 1) % Stage; + + // get next s2r b tile int64_t + // end of a block + + // dequant first + dequant( + tDrC.data(), tDrD.data(), scale_a * scale_b + ); + + s2r_b_smem = (s2r_b_smem + 1) % Stage; + // b advance + ++mma_block_b; + if (mma_block_b < size<1>(BS)) scale_b = BS(blk_n_coord, mma_block_b); + mma_block_a = mma_block_b; + if (mma_block_a < size<1>(AS)) scale_a = AS(blk_m_coord, mma_block_a); + + // load next stage + if (g2s_b_tile < ntile) { + g2s_tile_in_block = (g2s_tile_in_block + 1) % NumTilePerBlock; + copy_AB(g2s_tiled_copy_a, tAgA, tAsA, tAcA, g2s_a_tile, g2s_a_smem, m_limit); + copy_AB(g2s_tiled_copy_b, tBgB, tBsB, tBcB, g2s_b_tile, g2s_b_smem, n_limit); + ++g2s_b_tile; + g2s_b_smem = (g2s_b_smem + 1) % Stage; + ++g2s_block; + g2s_a_tile = g2s_block * NumTilePerBlock; + g2s_a_smem = (g2s_a_smem + 1) % Stage; + } + cp_async_fence(); + } + + + // epilogue + + Tensor sD = make_tensor( + make_smem_ptr(shared_storage.storage_D.smem_D.data()), + SmemLayoutD{} + ); + + R2STiledCopyD r2s_tiled_copy_d; + auto r2s_thr_copy_d = r2s_tiled_copy_d.get_slice(t_idx); + auto tDrD_r2s = r2s_thr_copy_d.retile_S(tDrD); + auto tDsD_r2s = r2s_thr_copy_d.partition_D(sD); + + S2GCopyD s2g_tiled_copy_d; + auto s2g_thr_copy_d = s2g_tiled_copy_d.get_slice(t_idx); + auto tDsD_s2g = s2g_thr_copy_d.partition_S(sD); + auto tDgD_s2g = s2g_thr_copy_d.partition_D(gD); + Tensor cD = make_identity_tensor(make_shape(Int{}, Int{})); + auto tDcD_s2g = s2g_thr_copy_d.partition_D(cD); + + auto tDgD_s2gx = group_modes<1, 3>(tDgD_s2g); // (CPY_, CPY_MN) + auto tDrD_r2sx = group_modes<1, 3>(tDrD_r2s); // (CPY_, CPY_MN) + auto tDcD_s2gx = group_modes<1, 3>(tDcD_s2g); + + int32_t step = size<3>(tDsD_r2s); // pipe + CUTLASS_PRAGMA_UNROLL + for (int32_t i = 0; i < size<1>(tDrD_r2sx); i += step) { + CUTLASS_PRAGMA_UNROLL + for (int32_t j = 0; j < step; ++j) { + if constexpr (std::is_same::value) { + cute::copy(r2s_tiled_copy_d, tDrD_r2sx(_, i + j), tDsD_r2s(_, 0, 0, j)); + } else { + auto t = make_tensor_like(tDrD_r2sx(_, i + j)); + cute::copy(tDrD_r2sx(_, i + j), t); + cute::copy(r2s_tiled_copy_d, t, tDsD_r2s(_, 0, 0, j)); + } + } + + __syncthreads(); + + // shm -> global + if constexpr (IsEvenM && IsEvenN) { + CUTLASS_PRAGMA_UNROLL + for (int32_t j = 0; j < step; ++j) + cute::copy(s2g_tiled_copy_d, tDsD_s2g(_, 0, 0, j), tDgD_s2gx(_, i + j)); + } else if constexpr (IsEvenN) { + CUTLASS_PRAGMA_UNROLL + for (int32_t j = 0; j < step; ++j) { + if (get<0>(tDcD_s2gx(0, i + j)) < m_limit) + cute::copy(s2g_tiled_copy_d, tDsD_s2g(_, 0, 0, j), tDgD_s2gx(_, i + j)); + } + } else if constexpr (IsEvenM) { + CUTLASS_PRAGMA_UNROLL + for (int32_t j = 0; j < step; ++j) + if (get<1>(tDcD_s2gx(size<0>(tDsD_s2g) - 1, i + j)) < n_limit) { + cute::copy(s2g_tiled_copy_d, tDsD_s2g(_, 0, 0, j), tDgD_s2gx(_, i + j)); + } else { + CUTLASS_PRAGMA_UNROLL + for (int k = 0; k < size<0>(tDsD_s2g); ++k) + if (get<1>(tDcD_s2gx(k, i + j)) < n_limit) + tDgD_s2gx(k, i + j) = tDsD_s2g(k, 0, 0, j); + } + } else { + CUTLASS_PRAGMA_UNROLL + for (int32_t j = 0; j < step; ++j) + if (get<0>(tDcD_s2gx(0, i + j)) < m_limit) { + if (get<1>(tDcD_s2gx(size<0>(tDsD_s2g) - 1, i + j)) < n_limit) { + cute::copy(s2g_tiled_copy_d, tDsD_s2g(_, 0, 0, j), tDgD_s2gx(_, i + j)); + } else { + for (int32_t k = 0; k < size<0>(tDsD_s2g); ++k) + if (get<1>(tDcD_s2gx(k, i + j)) < n_limit) + tDgD_s2gx(k, i + j) = tDsD_s2g(k, 0, 0, j); + } + } + } + __syncthreads(); + } + } + +}; + \ No newline at end of file diff --git a/fastvideo-kernel/csrc/turbodiffusion/gemm/launch.hpp b/fastvideo-kernel/csrc/turbodiffusion/gemm/launch.hpp new file mode 100644 index 000000000..d585dea25 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/gemm/launch.hpp @@ -0,0 +1,64 @@ +/* + * Copyright (c) 2025 by TurboDiffusion team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * + * Citation (please cite if you use this code): + * + * @article{zhang2025turbodiffusion, + * title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times}, + * author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun}, + * journal={arXiv preprint arXiv:2512.16093}, + * year={2025} + * } + */ + +#pragma once + +#include "common/common.hpp" +#include "common/launch.hpp" +#include "gemm/kernel.hpp" + + +template +bool int8_gemm_( + int8_t const *Aptr, float const *ASptr, + int8_t const *Bptr, float const *BSptr, + OutputDtype* Dptr, int64_t m, int64_t n, int64_t k, + int swizzle_dir = 1, int swizzle_size_log = 0, + cudaStream_t stream = nullptr +) { + BOOL_SWITCH(m % 128 == 0, IsEvenM, [&] { + BOOL_SWITCH(n % 128 == 0, IsEvenN, [&] { + using Kernel = GemmKernel; + if (!Kernel::can_implement(m, n, k)) + return false; + using Args = typename Kernel::Arguments; + Args args { + (void*)Aptr, (void*)ASptr, + (void*)Bptr, (void*)BSptr, (void*)Dptr, + m, n, k, swizzle_dir, + swizzle_size_log + }; + + auto params = Kernel::to_underlying_arguments(args); + + static constexpr size_t ShmSize = Kernel::ShmSize; + dim3 grid_shape = Kernel::get_grid_size(m, n); + dim3 block_shape = dim3(Kernel::ThreadNum); + auto func = device_kernel; + if (ShmSize >= 48 * 1024) { + cudaFuncSetAttribute( + func, + cudaFuncAttributeMaxDynamicSharedMemorySize, + ShmSize + ); + } + func<<>>( + params + ); + return true; + }); + }); + return true; +} diff --git a/fastvideo-kernel/csrc/turbodiffusion/gemm/utils.hpp b/fastvideo-kernel/csrc/turbodiffusion/gemm/utils.hpp new file mode 100644 index 000000000..b465c49d3 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/gemm/utils.hpp @@ -0,0 +1,129 @@ +/* + * Copyright (c) 2025 by TurboDiffusion team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * + * Citation (please cite if you use this code): + * + * @article{zhang2025turbodiffusion, + * title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times}, + * author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun}, + * journal={arXiv preprint arXiv:2512.16093}, + * year={2025} + * } + */ + +#pragma once + +#include +#include "cute/tensor.hpp" + +template < + bool IsEven, + class TiledCopy, + class SrcTensor, + class DstTensor, + class PrdTensor +> +CUTLASS_DEVICE void +copy_AB( + TiledCopy const& _copy, + SrcTensor const &S, + DstTensor &D, + PrdTensor const &ID, + const int64_t &i_read, + const int64_t &i_write, + const int64_t &limit +) { + using namespace cute; + if constexpr (IsEven) + cute::copy(_copy, S(_, _, _, i_read), D(_, _, _, i_write)); + else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < size<1>(ID); ++i) + if (get<0>(ID(0, i, 0)) < limit) + cute::copy(_copy, S(_, i, _, i_read), D(_, i, _, i_write)); + } +} + +template +CUTLASS_DEVICE void copy_async( + void const* gmem_src, + void* smem_dst +) { + uint32_t smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_dst));; + asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], %2;\n" + :: "r"(smem_int_ptr), + "l"(gmem_src), + "n"(N)); +} + +template +CUTLASS_DEVICE void copy_aligned(const void* src, void* dst, size_t N, int64_t thread_idx) { + static constexpr int NumElementPerLoad = sizeof(LoadType) / sizeof(T); + for (int64_t i = thread_idx * NumElementPerLoad; i < N; i += NumElementPerLoad * NumThreads) { + if (i + NumElementPerLoad <= N) { + copy_async( + (void*)((T*)src + i), + (void*)((T*)dst + i) + ); + } else { + for (int64_t j = 0; j < N - i; ++j) + copy_async( + (void*)((T*)src + i + j), + (void*)((T*)dst + i + j) + ); + } + } +} + +template +CUTLASS_DEVICE void g2s_vector_copy(const void* src, void* dst, size_t N, int64_t thread_idx) { + + uintptr_t src_addr = reinterpret_cast(src); + + if (src_addr % 16 == 0) { + copy_aligned(src, dst, N, thread_idx); + } else if (src_addr % 8 == 0) { + copy_aligned(src, dst, N, thread_idx); + } else if (src_addr % 4 == 0) { + copy_aligned(src, dst, N, thread_idx); + } else { + assert(0); + } + if constexpr (Commit) { + asm volatile("cp.async.commit_group;\n" ::); + } + if constexpr (Wait) { + asm volatile("cp.async.wait_all;\n" ::); + } +} + +template +CUTLASS_DEVICE +static void dequant( + T* mma_accum_ptr, + float* float_accum_ptr, + float scale +) { + static int const ic = 0x4B400000; + if constexpr (FastInt2Float && std::is_same_v) { + CUTLASS_PRAGMA_UNROLL + for (size_t i = 0; i < N; ++i) { + *(float_accum_ptr + i) += (__int_as_float(*(mma_accum_ptr + i)) - __int_as_float(ic)) * scale; + *(mma_accum_ptr + i) = ic; + } + } else if constexpr (std::is_same_v) { + CUTLASS_PRAGMA_UNROLL + for (size_t i = 0; i < N; ++i) { + *(float_accum_ptr + i) += __int2float_rn(*(mma_accum_ptr + i)) * scale; + *(mma_accum_ptr + i) = 0; + } + } else { + CUTLASS_PRAGMA_UNROLL + for (size_t i = 0; i < N; ++i) { + *(float_accum_ptr + i) += (*(mma_accum_ptr + i)) * scale; + *(mma_accum_ptr + i) = 0; + } + } +} \ No newline at end of file diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu new file mode 100644 index 000000000..0023774b5 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.cu @@ -0,0 +1,63 @@ +#include +#include +#include +#include +#include +#include +#include "common/common.hpp" +#include "norm/layernorm.hpp" + +auto layer_norm( + at::Tensor const Input, + float eps, + std::optional W, + std::optional const B, + std::optional Output +) { + using ElementIn = float; + using ElementOut = float; + using ElementWeight = float; + + int64_t const m = Input.size(0); + int64_t const n = Input.size(1); + torch::Device const input_device = Input.device(); + + if (!Output.has_value()) { + Output.emplace( + torch::empty( + {m, n}, + torch::TensorOptions().device(input_device).dtype(torch::kFloat32) + ) + ); + } + + + void *Iptr = Input.data_ptr(); + void *Wptr = W.has_value() ? W.value().data_ptr() : nullptr; + void *Bptr = B.has_value() ? B.value().data_ptr() : nullptr; + void *Optr = Output.value().data_ptr(); + + BOOL_SWITCH(B.has_value(), BIAS, [&]{ + BOOL_SWITCH(W.has_value(), AFFINE, [&]{ + CONFIG_SWITCH(n, [&]{ + layernorm< + ElementIn, ElementOut, ElementWeight, + AFFINE, BIAS, + MAX_HIDDEN_SIZE, NUM_THR_PER_CTA> ( + Iptr, Wptr, Bptr, + Optr, eps, m, n, + at::cuda::getCurrentCUDAStream().stream() + ); + }); + }); + }); + + + + return Output; +} + +void register_layer_norm(pybind11::module_ &m) { + m.def("layer_norm_cuda", &layer_norm); +} + diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp new file mode 100644 index 000000000..b129e0229 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/layernorm.hpp @@ -0,0 +1,202 @@ +#pragma once + +#include "common/load.hpp" +#include "common/store.hpp" +#include "common/launch.hpp" + + +template < + class InputDtype_, + class OutputDtype_, + class WeightDtype_, + bool Affine_, + bool Bias_, + int MaxHiddenSize_, + int NumThrPerCta_, + bool IsEven +> +class LayerNorm { +public: + using InputDtype = InputDtype_; + using OutputDtype = OutputDtype_; + using WeightDtype = WeightDtype_; + static constexpr int NumThrPerCta = NumThrPerCta_; + static constexpr int MaxHiddenSize = MaxHiddenSize_; + static constexpr bool Affine = Affine_; + static constexpr bool Bias = Bias_; + + static constexpr size_t ShmSize = 32; + static constexpr int NumElementPerThread = MaxHiddenSize / NumThrPerCta; + + static_assert(MaxHiddenSize % NumThrPerCta == 0); + + struct Params { + void const *Iptr; + void const *Wptr; + void const *Bptr; + void *Optr; + float eps; + int64_t m; + int64_t n; + }; + + using Arguments = Params; + + static Params to_underlying_arguments(Arguments const& args) { + return args; + } + + static dim3 get_grid_size(int64_t m, int64_t n) { + return dim3(m); + } + + static dim3 get_cta_size(int64_t m, int64_t n) { + return dim3(NumThrPerCta, 1, 1); + } + + CUTLASS_DEVICE + void operator()(Params const& params, char *shared_data) { + int const blk_m = blockIdx.x; + int const blk_n = 1; + int tidx = threadIdx.x; + float x[NumElementPerThread]; + + // load + Loader loader; + loader.load(params.Iptr, x, params.m, params.n, blk_m, 0, tidx); + + // mean reduction + float u = _reduce_sum(x, shared_data) / params.n; + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + x[i] -= u; + + __syncthreads(); + // var reduction + float v = sqrtf(_reduce_square(x, shared_data) / params.n + params.eps); + + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + x[i] /= v; + + if constexpr (Affine) { + // load weight + Loader weight_loader; + float w[NumElementPerThread]; + weight_loader.load(params.Wptr, w, 1, params.n, 0, 0, tidx); + if constexpr (Bias) { + float b[NumElementPerThread]; + weight_loader.load(params.Bptr, b, 1, params.n, 0, 0, tidx); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + x[i] = x[i] * w[i] + b[i]; + } else { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + x[i] = x[i] * w[i]; + } + } + + // save y + { + Saver saver; + if constexpr (std::is_same_v) { + saver.store(params.Optr, nullptr, x, 0, params.m, params.n, blk_m, 0, tidx); + } else { + OutputDtype tmp[NumElementPerThread]; + for (int i = 0; i < NumElementPerThread; ++i) + tmp[i] = OutputDtype(x[i]); + saver.store(params.Optr, nullptr, tmp, 0, params.m, params.n, blk_m, 0, tidx); + } + } + + } + +private: + CUTLASS_DEVICE + float _reduce_square(float *reg, char *shared_data) { + // thread + float sum_square = 0; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + sum_square += reg[i] * reg[i]; + + CUTLASS_PRAGMA_UNROLL + for (int i = 16; i >= 1; i >>= 1) { + sum_square += __shfl_down_sync(0xFFFFFFFF, sum_square, i); + } + if (threadIdx.x == 0) { + *(float*)shared_data = 0; + } + __syncthreads(); + + if (threadIdx.x % 32 == 0) { + atomicAdd((float*)shared_data, sum_square); + } + + __syncthreads(); + sum_square = *(float*)shared_data; + return sum_square; + } + + CUTLASS_DEVICE + float _reduce_sum(float *reg, char *shared_data) { + // thread + float sum = 0; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + sum += reg[i]; + + CUTLASS_PRAGMA_UNROLL + for (int i = 16; i >= 1; i >>= 1) { + sum += __shfl_down_sync(0xFFFFFFFF, sum, i); + } + if (threadIdx.x == 0) { + *(float*)shared_data = 0; + } + __syncthreads(); + + if (threadIdx.x % 32 == 0) { + atomicAdd((float*)shared_data, sum); + } + + __syncthreads(); + sum = *(float*)shared_data; + return sum; + } +}; + +template < + class InputDtype, + class OutputDtype, + class WeightDtype, + bool Affine, + bool Bias, + int MaxHiddenSize, + int NumThrPerCta +> +bool layernorm( + void const *Iptr, void const *Wptr, void const *Bptr, + void *Optr, float eps, int64_t m, int64_t n, + cudaStream_t stream = nullptr +) { + BOOL_SWITCH(n % MaxHiddenSize == 0, IsEven, [&] { + using Kernel = LayerNorm< + InputDtype, OutputDtype, WeightDtype, + Affine, Bias, + MaxHiddenSize, NumThrPerCta, + IsEven>; + using Arguments = typename Kernel::Arguments; + Arguments args = { + Iptr, Wptr, Bptr, Optr, + eps, m, n + }; + auto params = Kernel::to_underlying_arguments(args); + auto grid_shape = Kernel::get_grid_size(m, n); + auto cta_shape = Kernel::get_cta_size(m, n); + static constexpr size_t ShmSize = Kernel::ShmSize; + launch_kernel(params, grid_shape, cta_shape, ShmSize, stream); + }); + return true; +} \ No newline at end of file diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu new file mode 100644 index 000000000..dd1c4d8a8 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.cu @@ -0,0 +1,59 @@ +#include +#include +#include +#include +#include +#include +#include + +#include "common/common.hpp" +#include "norm/rmsnorm.hpp" + +auto rms_norm( + at::Tensor const& Input, + float eps, + const std::optional& Weight, + std::optional& Output +) { + + using ElementIn = float; + using ElementOut = float; + using ElementWeight = float; + + int64_t const m = Input.size(0); + int64_t const n = Input.size(1); + torch::Device const input_device = Input.device(); + + if (!Output.has_value()) { + Output.emplace( + torch::empty( + {m, n}, + torch::TensorOptions().device(input_device).dtype(torch::kFloat32) + ) + ); + } + + void *Iptr = Input.data_ptr(); + void *Wptr = Weight.has_value() ? Weight.value().data_ptr() : nullptr; + void *Optr = Output.value().data_ptr(); + + + CONFIG_SWITCH(n, [&]{ + rmsnorm< + ElementIn, ElementOut, ElementWeight, + MAX_HIDDEN_SIZE, NUM_THR_PER_CTA + > ( + Iptr, Wptr, + Optr, + eps, m, n, + at::cuda::getCurrentCUDAStream().stream() + ); +}); + + + return Output; +} + +void register_rms_norm(pybind11::module_ &m) { + m.def("rms_norm_cuda", &rms_norm); +} diff --git a/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp new file mode 100644 index 000000000..55796c92d --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/norm/rmsnorm.hpp @@ -0,0 +1,147 @@ +#pragma once + +#include "common/load.hpp" +#include "common/store.hpp" +#include "common/launch.hpp" + + +template < + class InputDtype_, + class OutputDtype_, + class WeightDtype_, + int MaxHiddenSize_, + int NumThrPerCta_, + bool IsEven +> +class RMSNorm { +public: + using InputDtype = InputDtype_; + using OutputDtype = OutputDtype_; + using WeightDtype = WeightDtype_; + static constexpr int NumThrPerCta = NumThrPerCta_; + static constexpr int MaxHiddenSize = MaxHiddenSize_; + + static constexpr size_t ShmSize = 32; + static constexpr int NumElementPerThread = MaxHiddenSize / NumThrPerCta; + + static_assert(MaxHiddenSize % NumThrPerCta == 0); + + struct Params { + void const *Iptr; + void const *Wptr; + void *Optr; + float eps; + int64_t m; + int64_t n; + }; + + using Arguments = Params; + + static Params to_underlying_arguments(Arguments const& args) { + return args; + } + + static dim3 get_grid_size(int64_t m, int64_t n) { + return dim3(m); + } + + static dim3 get_cta_size(int64_t m, int64_t n) { + return dim3(NumThrPerCta, 1, 1); + } + + CUTLASS_DEVICE + void operator()(Params const& params, char *shared_data) { + int const blk_m = blockIdx.x; + int const blk_n = 1; + int tidx = threadIdx.x; + float x[NumElementPerThread]; + + // load + Loader loader; + loader.load(params.Iptr, x, params.m, params.n, blk_m, 0, tidx); + + // rms reduction + float rms = sqrtf(_reduce_square(x, shared_data) / params.n + params.eps); + + // load weight + Loader weight_loader; + float w[NumElementPerThread]; + loader.load(params.Wptr, w, 1, params.n, 0, 0, tidx); + + // norm + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + x[i] = w[i] * x[i] / rms ; + + // save y + OutputDtype *output_reg = (OutputDtype*)x; + if constexpr (!std::is_same_v) { + output_reg = (OutputDtype*)w; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + output_reg[i] = OutputDtype(x[i]); + } + Saver saver; + saver.store(params.Optr, nullptr, output_reg, 0, params.m, params.n, blk_m, 0, tidx); + + } + +private: + CUTLASS_DEVICE + float _reduce_square(float *reg, char *shared_data) { + // thread + float sum_square = 0; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + sum_square += reg[i] * reg[i]; + + CUTLASS_PRAGMA_UNROLL + for (int i = 16; i >= 1; i >>= 1) { + sum_square += __shfl_down_sync(0xFFFFFFFF, sum_square, i); + } + if (threadIdx.x == 0) { + *(float*)shared_data = 0; + } + __syncthreads(); + + if (threadIdx.x % 32 == 0) { + atomicAdd((float*)shared_data, sum_square); + } + + __syncthreads(); + sum_square = *(float*)shared_data; + return sum_square; + } +}; + + +template < + class InputDtype, + class OutputDtype, + class WeightDtype, + int MaxHiddenSize, + int NumThrPerCta +> +bool rmsnorm( + void const *Iptr, void const *Wptr, + void *Optr, float eps, + int64_t m, int64_t n, + cudaStream_t stream = nullptr +) { + BOOL_SWITCH(n % MaxHiddenSize == 0, IsEven, [&] { + using Kernel = RMSNorm< + InputDtype, OutputDtype, WeightDtype, + MaxHiddenSize, NumThrPerCta, + IsEven>; + using Arguments = typename Kernel::Arguments; + Arguments args = { + Iptr, Wptr, Optr, eps, m, n + }; + auto params = Kernel::to_underlying_arguments(args); + auto grid_shape = Kernel::get_grid_size(m, n); + auto cta_shape = Kernel::get_cta_size(m, n); + static constexpr size_t ShmSize = Kernel::ShmSize; + launch_kernel(params, grid_shape, cta_shape, ShmSize, stream); + }); + return true; +} \ No newline at end of file diff --git a/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu new file mode 100644 index 000000000..9a56e3881 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.cu @@ -0,0 +1,75 @@ +/* + * Copyright (c) 2025 by TurboDiffusion team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * + * Citation (please cite if you use this code): + * + * @article{zhang2025turbodiffusion, + * title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times}, + * author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun}, + * journal={arXiv preprint arXiv:2512.16093}, + * year={2025} + * } + */ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "common/common.hpp" +#include "quant/quant.hpp" + +auto quant( + torch::Tensor const& Input, + std::optional& Output, + std::optional& Output_S +) { + + using ElementOut = int8_t; + static constexpr int BlockSize = 128; + static constexpr int NumThrPerCta = 256; + + int64_t m = Input.size(0); + int64_t n = Input.size(1); + torch::Device const input_device = Input.device(); + + create_tensor(input_device, Output, Output_S, m, n); + + ElementOut *Optr = (ElementOut*)Output.value().data_ptr(); + float *OSptr = Output_S.value().data_ptr(); + + switch (Input.scalar_type()) { + case torch::kHalf:{ + cutlass::half_t *Iptr = (cutlass::half_t*)Input.data_ptr(); + quantization ( + Iptr, Optr, OSptr, m, n, at::cuda::getCurrentCUDAStream().stream() + ); + break; + } + + case torch::kBFloat16:{ + cutlass::bfloat16_t *Iptr = (cutlass::bfloat16_t*)Input.data_ptr(); + quantization ( + Iptr, Optr, OSptr, m, n, at::cuda::getCurrentCUDAStream().stream() + ); + break; + } + + default: { + std::cerr << "Observing: " << Input.scalar_type() << " for the input datatype which is invalid"; + throw std::runtime_error("Unsupported input data type for quantize_to_fp4."); + } + } + + return std::make_tuple(Output, Output_S); +} + +void register_quant(pybind11::module_ &m) { + m.def("quant_cuda", &quant); +} diff --git a/fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp new file mode 100644 index 000000000..6c85b5035 --- /dev/null +++ b/fastvideo-kernel/csrc/turbodiffusion/quant/quant.hpp @@ -0,0 +1,195 @@ +/* + * Copyright (c) 2025 by TurboDiffusion team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * + * Citation (please cite if you use this code): + * + * @article{zhang2025turbodiffusion, + * title={TurboDiffusion: Accelerating Video Diffusion Models by 100-200 Times}, + * author={Zhang, Jintao and Zheng, Kaiwen and Jiang, Kai and Wang, Haoxu and Stoica, Ion and Gonzalez, Joseph E and Chen, Jianfei and Zhu, Jun}, + * journal={arXiv preprint arXiv:2512.16093}, + * year={2025} + * } + */ + +#pragma once + +#include +#include +#include "cutlass/numeric_conversion.h" + +#include "common/load.hpp" +#include "common/store.hpp" +#include "common/launch.hpp" + +template < + class InputDtype_, + int NumThrPerCta_, + bool IsEvenM, + bool IsEvenN +> +class Quantization { +public: + using InputDtype = InputDtype_; + using OutputDtype = int8_t; + using FPConverter = cutlass::NumericConverter; + + static constexpr int BlockSize = 128; + static constexpr int NumThrPerCta = NumThrPerCta_; + static constexpr int NumElementPerThread = BlockSize * BlockSize / NumThrPerCta; + static constexpr int NumThrPerRow = BlockSize / NumElementPerThread; + + static_assert(BlockSize * BlockSize % NumThrPerCta == 0); + static_assert(NumThrPerCta % BlockSize == 0); + + static constexpr size_t ShmSize = 32; + + static constexpr float int8_max = 128.f; + + struct Params { + void const *Iptr; + void *Optr; + void *OSptr; + int64_t const m; + int64_t const n; + }; + + using Arguments = Params; + + static Params to_underlying_arguments(Arguments const& args) { + return args; + } + + static dim3 get_grid_size(int64_t m, int64_t n) { + return dim3( + cdiv(n, BlockSize), + cdiv(m, BlockSize) + ); + } + + static dim3 get_cta_size(int64_t m, int64_t n) { + return dim3( + NumThrPerCta, 1, 1 + ); + } + + CUTLASS_DEVICE + void quantization( + float *float_reg, + void *Optr, void *OSptr, + int64_t const m, int64_t const n, + int blk_m, int blk_n, int tidx, + char *shared_data + ) { + + OutputDtype output_reg[NumElementPerThread]; + + + Saver saver; + + float amax = _reduce_amax(float_reg, (float*)shared_data); + + + _quantization(float_reg, output_reg, int8_max / amax); + + float scale_inv = amax / int8_max; + + saver.store(Optr, OSptr, output_reg, scale_inv, m, n, blk_m, blk_n, tidx); + + __syncthreads(); + } + + + CUTLASS_DEVICE + void operator()(Params const& params, char *shared_data) { + int blk_m = blockIdx.y; + int blk_n = blockIdx.x; + int tidx = threadIdx.x; + + float float_reg[NumElementPerThread]; + + // load float32 data + Loader loader; + loader.load(params.Iptr, float_reg, params.m, params.n, blk_m, blk_n, tidx); + quantization( + float_reg, params.Optr, params.OSptr, params.m, params.n, blk_m, blk_n, tidx, shared_data + ); + } + +private: + + CUTLASS_DEVICE float + _reduce_amax(float *reg, float *smem_ptr) { + float amax = 1e-8; + // thread reduction + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) + amax = max(amax, fabs(reg[i])); + + __syncwarp(); + + // warp reduction + CUTLASS_PRAGMA_UNROLL + for (int i = 16; i >= 1; i /= 2) { + amax = max( + __shfl_xor_sync(0xffffffff, amax, i, 32), + amax + ); + } + + // cta reduction + if (threadIdx.x == 0) { + *smem_ptr = 0; + } + __syncthreads(); + + atomicMax((uint32_t*)smem_ptr, reinterpret_cast(amax)); + + __syncthreads(); + + amax = *smem_ptr; + + return amax; + } + + CUTLASS_DEVICE void + _quantization(float *float_reg, OutputDtype *out_reg, float scale) { + FPConverter converter; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < NumElementPerThread; ++i) { + out_reg[i] = converter(float_reg[i] * scale); + } + + } +}; + +template < + class InputDtype, + int BlockSize, + int NumThrPerCta +> +bool quantization( + void const *Iptr, void *Optr, void *OSptr, + int64_t m, int64_t n, + cudaStream_t stream = nullptr +) { + BOOL_SWITCH(m % BlockSize == 0, IsEvenM, [&] { + BOOL_SWITCH(n % BlockSize == 0, IsEvenN, [&] { + using Kernel = Quantization< + InputDtype, NumThrPerCta, IsEvenM, IsEvenN>; + using Arguments = typename Kernel::Arguments; + Arguments args = { + Iptr, Optr, OSptr, + m, n + }; + auto params = Kernel::to_underlying_arguments(args); + auto grid_shape = Kernel::get_grid_size(m, n); + auto cta_shape = Kernel::get_cta_size(m, n); + static constexpr size_t ShmSize = Kernel::ShmSize; + launch_kernel(params, grid_shape, cta_shape, ShmSize, stream); + }); + }); + + return true; +} diff --git a/fastvideo-kernel/include/cutlass b/fastvideo-kernel/include/cutlass new file mode 160000 index 000000000..e67e63c33 --- /dev/null +++ b/fastvideo-kernel/include/cutlass @@ -0,0 +1 @@ +Subproject commit e67e63c331d6e4b729047c95cf6b92c8454cba89 diff --git a/fastvideo-kernel/python/fastvideo_kernel/__init__.py b/fastvideo-kernel/python/fastvideo_kernel/__init__.py index 22ede949e..560717d94 100644 --- a/fastvideo-kernel/python/fastvideo_kernel/__init__.py +++ b/fastvideo-kernel/python/fastvideo_kernel/__init__.py @@ -11,11 +11,24 @@ from fastvideo_kernel.vmoba import ( process_moba_output, ) +from fastvideo_kernel.turbodiffusion_ops import ( + Int8Linear, + FastRMSNorm, + FastLayerNorm, + int8_linear, + int8_quant, +) + __all__ = [ "sliding_tile_attention", "video_sparse_attn", "moba_attn_varlen", "process_moba_input", "process_moba_output", + "Int8Linear", + "FastRMSNorm", + "FastLayerNorm", + "int8_linear", + "int8_quant", "__version__", ] diff --git a/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py b/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py new file mode 100644 index 000000000..3d93770ad --- /dev/null +++ b/fastvideo-kernel/python/fastvideo_kernel/turbodiffusion_ops.py @@ -0,0 +1,716 @@ +from __future__ import annotations + +from typing import Optional, Tuple +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl + +# Try to load the C++ extension +try: + from fastvideo_kernel._C import fastvideo_kernel_ops + quant_cuda = getattr(fastvideo_kernel_ops, "quant_cuda", None) + gemm_cuda = getattr(fastvideo_kernel_ops, "gemm_cuda", None) + rms_norm_cuda = getattr(fastvideo_kernel_ops, "rms_norm_cuda", None) + layer_norm_cuda = getattr(fastvideo_kernel_ops, "layer_norm_cuda", None) +except ImportError: + quant_cuda = None + gemm_cuda = None + rms_norm_cuda = None + layer_norm_cuda = None + +def int8_quant(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Quantize a floating-point tensor to int8 using a custom CUDA kernel. + + Args: + x (torch.Tensor): Input tensor of type float16/bfloat16. + + Returns: + Tuple[torch.Tensor, torch.Tensor]: + - x_q: Quantized int8 tensor. + - x_scale: Per-block scale tensor used for quantization. + """ + x_q, x_scale = quant_cuda(x, None, None) + return x_q, x_scale + + +def int8_linear( + x: torch.Tensor, + w_q: torch.Tensor, + w_s: torch.Tensor, + **kwargs, +) -> torch.Tensor: + """ + Perform an int8 GEMM (matrix multiplication) using quantized weights and a + quantized version of the input. The underlying compute is performed by a + custom CUDA kernel. + + Args: + x (torch.Tensor): Input activation of shape (M, K) in float32. + w_q (torch.Tensor): Quantized int8 weight tensor of shape (N, K). + w_s (torch.Tensor): Scale tensor associated with w_q. + **kwargs: Additional options (reserved for future use). + + Returns: + torch.Tensor: Output tensor of shape (M, N) in float32. + """ + assert w_q.dtype == torch.int8, "Weight tensor must be int8." + shape = x.shape + x = x.reshape(-1, shape[-1]) + m = x.shape[0] + n = w_q.shape[0] + y = torch.zeros(m, n, dtype=x.dtype, device=x.device) + + x_q, x_s = int8_quant(x) + gemm_cuda(x_q, x_s, w_q, w_s, y) + return y.reshape(*shape[:-1], n) + +def flatten_if_batched(*tensors): + """ + Flattens all input tensors from (B, N, D_i) to (B * N, D_i) if they are batched (3D). + + Args: + *tensors: Any number of input tensors, each must have shape (B, N, D_i) or (N, D_i) + + Returns: + flat_tensors: List of flattened tensors + batched: Boolean flag indicating whether inputs were batched + batch_size: Batch size if batched, else None + """ + if not tensors: + raise ValueError("At least one tensor must be provided.") + + first = tensors[0] + assert len(first.shape) in [ + 2, + 3, + ], "Input tensors must be batched (3D) or not batched (2D)" + + if len(first.shape) == 3: # batched + batched = True + batch_size = first.shape[0] + assert all(t.shape[0] == batch_size for t in tensors), "All input tensors must have the same batch size" + assert all( + t.shape[1] == first.shape[1] for t in tensors + ), "All input tensors must have the same sequence length" + flat_tensors = [t.reshape(-1, t.shape[-1]) for t in tensors] + else: + batched = False + batch_size = None + flat_tensors = list(tensors) + + return flat_tensors, batched, batch_size + + +@triton.jit +def _rms_norm_fwd_fused( + X, + Y, + W, + Rstd, + x_stride, + y_stride, + N: tl.constexpr, # number of columns in X, + N2: tl.constexpr, + eps, # epsilon to avoid division by zero + BLOCK_M: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + pid = tl.program_id(0) + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.arange(0, N2) + mask = cols < N + + x_ptr = X + rows[:, None] * x_stride + cols[None, :] + y_ptr = Y + rows[:, None] * y_stride + cols[None, :] + + x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32) + + # Compute variance + _var = x * x + var = tl.sum(_var, axis=1) / N + rstd = 1 / tl.sqrt(var + eps) + + # Write mean / rstd + tl.store(Rstd + rows, rstd) + rstd = tl.reshape(rstd, (BLOCK_M, 1)) + + # Normalize and apply linear transformation + w = tl.load(W + cols) + x_hat = x * rstd + y = x_hat * w + + # Write output + y = y.to(Y.type.element_ty) + tl.store(y_ptr, y, mask=mask[None, :]) + + +def rmsnorm(x, w, eps): + """ + Forward pass of the RMSNorm. + + Args: + x (torch.Tensor): Input tensor, High precision. + w (torch.Tensor): RMSNorm weight tensor. + eps (float): RMSNorm epsilon value. + + Returns: + y (torch.Tensor): Output tensor, High precision. + rstd (torch.Tensor): Inverse standard deviation, needed for backward. + """ + assert x.is_contiguous(), "Input must be contiguous" + # Change batched 3D input to 2D + [x], batched, BS = flatten_if_batched(x) + + # allocate output + M, N = x.shape + y = torch.empty_like(x, dtype=x.dtype) + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + + # heuristics for number of warps + num_warps = 8 + + # Avoid illegal memory access + N2 = triton.next_power_of_2(N) + + if N <= 512: + BLOCK_M = 32 + else: + BLOCK_M = 1 + + # Call the triton kernel + _rms_norm_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( # + x, + y, + w, + rstd, # + x.stride(0), + y.stride(0), + N, + N2, + eps, + num_warps=num_warps, + BLOCK_M=BLOCK_M, + ) + + # Recover 2D to 3D + if batched: + y = y.reshape(BS, -1, y.shape[-1]) + + return y, rstd + +@triton.jit +def _layer_norm_param_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + W, # pointer to the weights + B, # pointer to the biases + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + x_stride, # how much to increase the pointer when moving by 1 row + y_stride, # how much to increase the pointer when moving by 1 row + N: tl.constexpr, # number of columns in X, + N2: tl.constexpr, # number of columns in X, + eps, # epsilon to avoid division by zero + BLOCK_M: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + pid = tl.program_id(0) + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.arange(0, N2) + mask = cols < N + + x_ptr = X + rows[:, None] * x_stride + cols[None, :] + y_ptr = Y + rows[:, None] * y_stride + cols[None, :] + + x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32) + + # Compute mean and Variance + mean = tl.sum(x, axis=1, keep_dims=True) / N + # Compute variance + _var = (x - mean) * (x - mean) + var = tl.sum(_var, axis=1, keep_dims=True) / N + rstd = 1 / tl.sqrt(var + eps) + + # Write mean / rstd + _mean = tl.reshape(mean, (BLOCK_M)) + _rstd = tl.reshape(rstd, (BLOCK_M)) + tl.store(Mean + rows, _mean) + tl.store(Rstd + rows, _rstd) + + # Normalize and apply linear transformation + x_hat = (x - mean) * rstd + + w = tl.load(W + cols) + b = tl.load(B + cols) + + x_hat = x_hat * w + b + + # Write output + x_hat = x_hat.to(Y.type.element_ty) + tl.store(y_ptr, x_hat, mask=mask[None, :]) + + +def layernorm_param(x, w, b, eps): + # Change batched 3D input to 2D + [x], batched, BS = flatten_if_batched(x) + + # allocate output + M, N = x.shape + y = torch.empty_like(x, dtype=torch.float32) + mean = torch.empty((M,), dtype=torch.float32, device=x.device) + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + # heuristics for number of warps + num_warps = 8 + + N2 = triton.next_power_of_2(N) + + if N <= 512: + BLOCK_M = 32 + else: + BLOCK_M = 1 + + # enqueue kernel + _layer_norm_param_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( # + x, + y, + w, + b, + mean, + rstd, # + x.stride(0), + y.stride(0), + N, + N2, + eps, + num_warps=num_warps, + BLOCK_M=BLOCK_M, + ) + + # Recover 2D to 3D + if batched: + y = y.reshape(BS, -1, y.shape[-1]) + + return y, mean, rstd + + +######################################################## +# Elementwise_affine=False +######################################################## + + +@triton.jit +def _layer_norm_noparam_fwd_fused( + X, # pointer to the input + Y, # pointer to the output + Mean, # pointer to the mean + Rstd, # pointer to the 1/std + x_stride, # how much to increase the pointer when moving by 1 row + y_stride, # how much to increase the pointer when moving by 1 row + N: tl.constexpr, # number of columns in X, + N2: tl.constexpr, # number of columns in X, + eps, # epsilon to avoid division by zero + BLOCK_M: tl.constexpr, +): + # Map the program id to the row of X and Y it should compute. + pid = tl.program_id(0) + rows = pid * BLOCK_M + tl.arange(0, BLOCK_M) + cols = tl.arange(0, N2) + mask = cols < N + + x_ptr = X + rows[:, None] * x_stride + cols[None, :] + y_ptr = Y + rows[:, None] * y_stride + cols[None, :] + + x = tl.load(x_ptr, mask=mask[None, :], other=0.0).to(tl.float32) + + # Compute mean and Variance + mean = tl.sum(x, axis=1, keep_dims=True) / N + # Compute variance + _var = (x - mean) * (x - mean) + var = tl.sum(_var, axis=1, keep_dims=True) / N + rstd = 1 / tl.sqrt(var + eps) + + # Write mean / rstd + _mean = tl.reshape(mean, (BLOCK_M)) + _rstd = tl.reshape(rstd, (BLOCK_M)) + tl.store(Mean + rows, _mean) + tl.store(Rstd + rows, _rstd) + + # Normalize and apply linear transformation + x_hat = (x - mean) * rstd + + # Write output + x_hat = x_hat.to(Y.type.element_ty) + tl.store(y_ptr, x_hat, mask=mask[None, :]) + + +def layernorm_noparam(x, eps): + assert x.is_contiguous(), "Input must be contiguous" + + # Change batched 3D input to 2D + [x], batched, BS = flatten_if_batched(x) + + # allocate output + M, N = x.shape + y = torch.empty_like(x, dtype=torch.float32) + mean = torch.empty((M,), dtype=torch.float32, device=x.device) + rstd = torch.empty((M,), dtype=torch.float32, device=x.device) + # heuristics for number of warps + num_warps = 8 + + N2 = triton.next_power_of_2(N) + + if N <= 512: + BLOCK_M = 32 + else: + BLOCK_M = 1 + + # enqueue kernel + _layer_norm_noparam_fwd_fused[(triton.cdiv(M, BLOCK_M),)]( # + x, + y, + mean, + rstd, # + x.stride(0), + y.stride(0), + N, + N2, + eps, + num_warps=num_warps, + BLOCK_M=BLOCK_M, + ) + + # Recover 2D to 3D + if batched: + y = y.reshape(BS, -1, y.shape[-1]) + + return y, mean, rstd + +def layernorm(x, w, b, eps, elementwise_affine=True): + if elementwise_affine: + assert w is not None and b is not None + return layernorm_param(x, w, b, eps) + else: + assert w is None and b is None + return layernorm_noparam(x, eps) + +def cdiv(a: int, b: int): + return (a + b - 1) // b + +def dequantize_weight(w_q, w_s, dtype): + # w_q: (N, K) int8 + # w_s: (NB, KB) float32 + # Block size 128 + BLOCK = 128 + N, K = w_q.shape + # Expand w_s + # Repeat interleave + w_s_exp = w_s.repeat_interleave(BLOCK, dim=0).repeat_interleave(BLOCK, dim=1) + # Crop + w_s_exp = w_s_exp[:N, :K] + return w_q.to(dtype) * w_s_exp.to(dtype) + +class Int8LinearFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w_q, w_s, bias=None): + ctx.save_for_backward(x, w_q, w_s, bias) + ctx.bias_requires_grad = bias.requires_grad if bias is not None else False + + out = int8_linear(x, w_q, w_s) + if bias is not None: + out = out + bias + return out + + @staticmethod + def backward(ctx, grad_output): + x, w_q, w_s, bias = ctx.saved_tensors + + grad_input = None + grad_weight = None + grad_scale = None + grad_bias = None + + if ctx.needs_input_grad[0]: + # grad_input = grad_output @ W + w_float = dequantize_weight(w_q, w_s, grad_output.dtype) + grad_input = torch.matmul(grad_output, w_float) + + if ctx.bias_requires_grad: + dim_to_sum = list(range(grad_output.dim() - 1)) + grad_bias = grad_output.sum(dim=dim_to_sum) + + return grad_input, grad_weight, grad_scale, grad_bias + +class RMSNormFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w, eps): + y, rstd = rmsnorm(x, w, eps) + ctx.save_for_backward(x, w, rstd) + ctx.eps = eps + return y + + @staticmethod + def backward(ctx, grad_output): + # x, w, rstd are saved + # grad_output is dL/dy + # We can implement backward using torch ops for correctness + x, w, rstd = ctx.saved_tensors + eps = ctx.eps + N = x.shape[-1] + + # dL/dy * w + dx = grad_output * w + + # Expand rstd + rstd = rstd.unsqueeze(-1) # (B, 1) or (M, 1) + if x.dim() == 3: + # Flatten if necessary or handle dimensions. + # forward flattens x to (M, N) internally but saved x might be 3D? + # rmsnorm takes x, flattens it. But does it modify x in place? No. + # But ctx.save_for_backward saves the original x (3D if input was 3D). + # rstd is (M,). + # We should flatten x and grad_output to match logic + x_flat = x.reshape(-1, N) + grad_output_flat = grad_output.reshape(-1, N) + dx = dx.reshape(-1, N) + else: + x_flat = x + grad_output_flat = grad_output + dx = dx + + # Standard RMSNorm backward + # dy = grad_output_flat + # x_hat = x * rstd + # w * dy + # c1 = mean(dy * w * x) * rstd^2 + # dx = (dy * w - c1 * x) * rstd + + # More precise: + # y = x * rstd * w + # dy = grad_output + # dw = sum(dy * x * rstd, dim=0) + + grad_w = None + if ctx.needs_input_grad[1]: # w + # grad_w = (grad_output_flat * x_flat * rstd).sum(dim=0) + # More accurate gradient for weight when considering rstd was computed from x + # Actually, for Affine part (w), it is just sum(dL/dy * x_hat). + # y = x_hat * w + # dL/dw = sum(dL/dy * x_hat) + x_hat = x_flat * rstd + grad_w = (grad_output_flat * x_hat).sum(dim=0) + + grad_input = None + if ctx.needs_input_grad[0]: # x + # x_hat = x * rstd + # dx = rstd * (w * dy - mean(w * dy * x_hat) * x_hat) + # but for RMSNorm: + # dx = rstd * (w * dy - (x * rstd^2) * mean(w * dy * x)) -> check formula + + # Using PyTorch autograd for reference logic: + # y = x / sqrt(mean(x^2) + eps) * w + # Let sigma = sqrt(...) + # dL/dx = dL/dy * w * (1/sigma) + dL/dsigma * dsigma/dx + # dsigma/dx = 1/(2*sigma) * 2x/N = x / (N * sigma) + # dL/dsigma = sum(dL/dy * w * x * (-1/sigma^2)) + # dL/dx = (dL/dy * w)/sigma - sum(dL/dy * w * x) * x / (N * sigma^3) + # = (1/sigma) * [ (dL/dy * w) - x * sum(dL/dy * w * x) / (N * sigma^2) ] + # = rstd * [ (dL/dy * w) - x * rstd^2 * mean(dL/dy * w * x) ] # Wait, sum/N is mean + + # Let's compute: + dy_w = grad_output_flat * w + # term2 = (dy_w * x_flat).sum(dim=-1, keepdim=True) / N # mean(dy*w*x) + # grad_input = rstd * (dy_w - x_flat * (rstd ** 2) * term2) + + # Re-derivation: + # x_hat = x * rstd + # y = x_hat * w + # dL/dx = dL/dy * dy/dx + # dy/dx = w * dx_hat/dx + # dx_hat/dx = rstd * (I - x * x^T * rstd^2 / N) ?? No. + # dx_hat_i / dx_j = rstd * (delta_ij - x_i * x_j * rstd^2 / N) + + # dL/dx_hat = dL/dy * w + # dL/dx = rstd * (dL/dx_hat - x * rstd^2 * mean(dL/dx_hat * x)) + + dx_hat = grad_output_flat * w + dx_hat_x_mean = (dx_hat * x_flat).mean(dim=-1, keepdim=True) + grad_input = rstd * (dx_hat - x_flat * (rstd ** 2) * dx_hat_x_mean) + + if x.dim() == 3: + grad_input = grad_input.reshape(x.shape) + + return grad_input, grad_w, None + +class LayerNormFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x, w, b, eps, elementwise_affine): + if elementwise_affine: + y, mean, rstd = layernorm_param(x, w, b, eps) + ctx.save_for_backward(x, w, b, mean, rstd) + else: + y, mean, rstd = layernorm_noparam(x, eps) + ctx.save_for_backward(x, None, None, mean, rstd) + + ctx.eps = eps + ctx.elementwise_affine = elementwise_affine + return y + + @staticmethod + def backward(ctx, grad_output): + x, w, b, mean, rstd = ctx.saved_tensors + N = x.shape[-1] + + if x.dim() == 3: + x_flat = x.reshape(-1, N) + grad_output_flat = grad_output.reshape(-1, N) + else: + x_flat = x + grad_output_flat = grad_output + + if w is None: + w_eff = 1.0 + else: + w_eff = w + + # dx calculation + # x_hat = (x - mean) * rstd + # y = x_hat * w + b + # dL/dx_hat = dL/dy * w + + dy = grad_output_flat + dx_hat = dy * w_eff + + # dL/dvar = sum(dL/dx_hat * (x-mean) * (-0.5) * (var+eps)^(-1.5)) + # = -0.5 * rstd^3 * sum(dx_hat * (x-mean)) + # dL/dmean = sum(dL/dx_hat * (-rstd)) + dL/dvar * (-2/N) * sum(x-mean) + # term sum(x-mean) is 0. So second part vanishes. + # dL/dmean = -rstd * sum(dx_hat) + + # dL/dx = dL/dx_hat * rstd + dL/dvar * 2(x-mean)/N + dL/dmean * 1/N + # = dx_hat * rstd + (-0.5 * rstd^3 * sum(dx_hat * (x-mean))) * 2(x-mean)/N + (-rstd * sum(dx_hat))/N + # = rstd * [ dx_hat - mean(dx_hat) - (x-mean)*rstd^2 * mean(dx_hat * (x-mean)) ] + + x_centered = x_flat - mean.unsqueeze(-1) + dx_hat_mean = dx_hat.mean(dim=-1, keepdim=True) + dx_hat_x_centered_mean = (dx_hat * x_centered).mean(dim=-1, keepdim=True) + + grad_input = rstd.unsqueeze(-1) * (dx_hat - dx_hat_mean - x_centered * (rstd.unsqueeze(-1)**2) * dx_hat_x_centered_mean) + + if x.dim() == 3: + grad_input = grad_input.reshape(x.shape) + + grad_w = None + grad_b = None + + if ctx.elementwise_affine: + if ctx.needs_input_grad[1]: + x_hat = x_centered * rstd.unsqueeze(-1) + grad_w = (dy * x_hat).sum(dim=0) + if ctx.needs_input_grad[2]: + grad_b = dy.sum(dim=0) + + return grad_input, grad_w, grad_b, None, None + + +class Int8Linear(nn.Module): + def __init__(self, in_features, out_features, bias=True, dtype=torch.bfloat16): + super().__init__() + self.in_features = in_features + self.out_features = out_features + + row_blocks = cdiv(out_features, b=128) + col_blocks = cdiv(in_features, b=128) + + self.register_buffer("int8_weight", torch.empty((out_features, in_features), dtype=torch.int8)) + self.register_buffer("scale", torch.empty((row_blocks, col_blocks), dtype=torch.float32)) + if bias: + self.register_buffer("bias", torch.empty(out_features, dtype=dtype)) + else: + self.bias = None + + + def forward(self, x): + return Int8LinearFunction.apply(x, self.int8_weight, self.scale, self.bias) + + @classmethod + def from_linear(cls, original_linear: nn.Linear, quantize: bool = True): + + int8_layer = cls( + original_linear.in_features, + original_linear.out_features, + bias=original_linear.bias is not None, + dtype=original_linear.weight.dtype + ) + if quantize: + w_data = original_linear.weight.data.cuda() + int8_w, scale = int8_quant(w_data) + + int8_layer.int8_weight.copy_(int8_w) + int8_layer.scale.copy_(scale) + if original_linear.bias is not None: + int8_layer.bias.data.copy_(original_linear.bias.data.cuda()) + + return int8_layer + +class FastRMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.register_buffer("weight", torch.ones(dim)) + + def forward(self, x): + return RMSNormFunction.apply(x.float(), self.weight, self.eps).to(x.dtype) + + @classmethod + def from_rmsnorm(cls, original_rmsnorm): + rmsnorm_layer = cls( + dim=original_rmsnorm.dim, + eps=original_rmsnorm.eps + ) + if original_rmsnorm.weight.device != torch.device('meta'): + rmsnorm_layer.weight.data.copy_(original_rmsnorm.weight.float().data) + return rmsnorm_layer + +class FastLayerNorm(nn.Module): + + def __init__( + self, + dim: int, + eps: float = 1e-5, + elementwise_affine: bool = False, + bias: bool = True + ) : + super().__init__() + self.dim = dim # type: ignore[arg-type] + self.eps = eps + self.elementwise_affine = elementwise_affine + if self.elementwise_affine: + self.register_buffer("weight", torch.empty(self.dim)) + if bias: + self.register_buffer("bias", torch.empty(self.dim)) + else: + self.bias = None + else: + self.register_parameter("weight", None) + self.register_parameter("bias", None) + + def forward(self, x): + return LayerNormFunction.apply(x.float(), self.weight, self.bias, self.eps, self.elementwise_affine).to(x.dtype) + + @classmethod + def from_layernorm(cls, original_layernorm): + layernorm_layer = cls( + dim=original_layernorm.normalized_shape[0], + eps=original_layernorm.eps, + elementwise_affine=False if original_layernorm.weight is None else True, + bias=original_layernorm.bias is not None + ) + if original_layernorm.weight is not None and original_layernorm.weight.device != torch.device('meta'): + layernorm_layer.weight.data.copy_(original_layernorm.weight.data) + if original_layernorm.bias is not None and original_layernorm.bias.device != torch.device('meta'): + layernorm_layer.bias.data.copy_(original_layernorm.bias.data) + return layernorm_layer diff --git a/fastvideo-kernel/tests/test_turbodiffusion.py b/fastvideo-kernel/tests/test_turbodiffusion.py new file mode 100644 index 000000000..2233097ff --- /dev/null +++ b/fastvideo-kernel/tests/test_turbodiffusion.py @@ -0,0 +1,294 @@ +import torch +import pytest +import sys +import os + +# Ensure local package is imported +# sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../python"))) + +try: + from fastvideo_kernel._C import fastvideo_kernel_ops +except ImportError: + fastvideo_kernel_ops = None + +from fastvideo_kernel import turbodiffusion_ops + +# Helper for RMS Norm reference +def rms_norm_ref(x, w, eps=1e-6): + dtype = x.dtype + x = x.float() + variance = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(variance + eps) + return (x * w.float()).to(dtype) + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +class TestTurboDiffusion: + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("shape", [(16, 128), (32, 256), (1, 1024)]) + def test_quant_correctness(self, dtype, shape): + if turbodiffusion_ops.quant_cuda is None: + pytest.skip("quant_cuda not available") + + x = torch.randn(shape, dtype=dtype, device="cuda") + x_q, x_scale = turbodiffusion_ops.int8_quant(x) + + assert x_q.dtype == torch.int8 + assert x_scale.dtype == torch.float32 + + # Simple check: dequantize and compute error + # Note: The quantization scheme details matter here (per block? per tensor?). + # Looking at quant.cu, it seems to be block-based but the output scale shape isn't immediately obvious from python signature + # without looking at C++ code deeper. + # But let's check shapes at least. + + # If we can't easily dequantize without knowing block size logic in python, + # checking that it runs and produces valid shapes is a good start. + assert x_q.shape == shape + # x_scale shape depends on block size, usually smaller than x + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_gemm_correctness(self, dtype): + if turbodiffusion_ops.gemm_cuda is None: + pytest.skip("gemm_cuda not available") + + M, N, K = 32, 64, 128 + x = torch.randn(M, K, dtype=dtype, device="cuda") + + # Create weights + # For simplicity in testing, let's create random int8 weights and scales + w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda") + + # Scale shape: The Int8Linear class uses: + # row_blocks = cdiv(out_features, b=128) + # col_blocks = cdiv(in_features, b=128) + # scale shape: (row_blocks, col_blocks) + + row_blocks = (N + 127) // 128 + col_blocks = (K + 127) // 128 + w_s = torch.randn(row_blocks, col_blocks, dtype=torch.float32, device="cuda").abs() + + # Run int8_linear + output = turbodiffusion_ops.int8_linear(x, w_q, w_s) + + assert output.shape == (M, N) + assert output.dtype == dtype + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_gemm_backward(self, dtype): + if turbodiffusion_ops.gemm_cuda is None: + pytest.skip("gemm_cuda not available") + + M, N, K = 32, 64, 128 + x = torch.randn(M, K, dtype=dtype, device="cuda", requires_grad=True) + + # Weights (frozen) + w_q = torch.randint(-127, 127, (N, K), dtype=torch.int8, device="cuda") + row_blocks = (N + 127) // 128 + col_blocks = (K + 127) // 128 + w_s = torch.randn(row_blocks, col_blocks, dtype=torch.float32, device="cuda").abs() + + bias = torch.randn(N, dtype=dtype, device="cuda", requires_grad=True) + + # Use Int8LinearFunction + output = turbodiffusion_ops.Int8LinearFunction.apply(x, w_q, w_s, bias) + loss = output.sum() + loss.backward() + + assert x.grad is not None + assert bias.grad is not None + assert x.grad.shape == (M, K) + + # Check correctness against dequantized weight + w_float = turbodiffusion_ops.dequantize_weight(w_q, w_s, dtype) + + # Reference + x_ref = x.detach().clone().requires_grad_() + bias_ref = bias.detach().clone().requires_grad_() + + # Note: Int8Linear forward uses quantized input, so output won't match exactly reference with float input. + # But backward gradients should be consistent with the logic we implemented: + # grad_input = grad_output @ w_float + + # Let's verify our backward logic matches standard matmul backward with dequantized weights + # We manually compute expected gradient given grad_output = ones + grad_output = torch.ones_like(output) + expected_x_grad = grad_output @ w_float + expected_bias_grad = grad_output.sum(0) + + torch.testing.assert_close(x.grad, expected_x_grad, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(bias.grad, expected_bias_grad, atol=1e-2, rtol=1e-2) + + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("shape", [(2, 16, 128), (4, 32, 256)]) + def test_rms_norm_triton(self, dtype, shape): + x = torch.randn(shape, dtype=dtype, device="cuda") + dim = shape[-1] + w = torch.randn(dim, dtype=dtype, device="cuda") + eps = 1e-5 + + # Triton implementation + # Note: rmsnorm returns tuple now + res = turbodiffusion_ops.rmsnorm(x, w, eps) + if isinstance(res, tuple): + out_triton = res[0] + else: + out_triton = res + + # Reference + out_ref = rms_norm_ref(x, w, eps) + + torch.testing.assert_close(out_triton, out_ref, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("shape", [(2, 16, 128)]) + def test_rms_norm_backward(self, dtype, shape): + x = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + dim = shape[-1] + w = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True) + eps = 1e-5 + + # Forward via Function + y = turbodiffusion_ops.RMSNormFunction.apply(x, w, eps) + loss = y.sum() + loss.backward() + + x_grad = x.grad + w_grad = w.grad + + # Reference + x_ref = x.detach().clone().requires_grad_() + w_ref = w.detach().clone().requires_grad_() + # Custom RMSNorm ref in pytorch + def rms_norm_ref_grad(x, w, eps): + x_float = x.float() + var = x_float.pow(2).mean(-1, keepdim=True) + rstd = torch.rsqrt(var + eps) + return (x_float * rstd * w.float()).to(x.dtype) + + y_ref = rms_norm_ref_grad(x_ref, w_ref, eps) + loss_ref = y_ref.sum() + loss_ref.backward() + + torch.testing.assert_close(x_grad, x_ref.grad, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(w_grad, w_ref.grad, atol=2e-2, rtol=2e-2) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_rms_norm_cuda(self, dtype): + if fastvideo_kernel_ops is None or not hasattr(fastvideo_kernel_ops, "rms_norm_cuda"): + pytest.skip("rms_norm_cuda not available") + + shape = (16, 128) + x = torch.randn(shape, dtype=dtype, device="cuda") + dim = shape[-1] + w = torch.randn(dim, dtype=dtype, device="cuda") + eps = 1e-5 + + # C++ implementation + # Signature: rms_norm_cuda(Input, eps, Weight, Output) -> Output + out_cuda = torch.empty_like(x) + fastvideo_kernel_ops.rms_norm_cuda(x, eps, w, out_cuda) + + # Reference + out_ref = rms_norm_ref(x, w, eps) + + torch.testing.assert_close(out_cuda, out_ref, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("shape", [(2, 16, 128), (4, 32, 256)]) + def test_layer_norm_triton(self, dtype, shape): + x = torch.randn(shape, dtype=dtype, device="cuda") + dim = shape[-1] + eps = 1e-5 + + # With affine + w = torch.randn(dim, dtype=dtype, device="cuda") + b = torch.randn(dim, dtype=dtype, device="cuda") + + # Triton implementation + # Note: layernorm returns tuple now + res = turbodiffusion_ops.layernorm(x, w, b, eps, elementwise_affine=True) + if isinstance(res, tuple): + out_triton = res[0] + else: + out_triton = res + out_triton = out_triton.to(dtype) + + # Reference + ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda() + ln.weight.data.copy_(w) + ln.bias.data.copy_(b) + out_ref = ln(x) + + torch.testing.assert_close(out_triton, out_ref, atol=1e-2, rtol=1e-2) + + # Without affine + res_no_affine = turbodiffusion_ops.layernorm(x, None, None, eps, elementwise_affine=False) + if isinstance(res_no_affine, tuple): + out_triton_no_affine = res_no_affine[0] + else: + out_triton_no_affine = res_no_affine + out_triton_no_affine = out_triton_no_affine.to(dtype) + ln_no_affine = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=False, dtype=dtype).cuda() + out_ref_no_affine = ln_no_affine(x) + + torch.testing.assert_close(out_triton_no_affine, out_ref_no_affine, atol=1e-2, rtol=1e-2) + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + @pytest.mark.parametrize("shape", [(2, 16, 128)]) + def test_layer_norm_backward(self, dtype, shape): + x = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True) + dim = shape[-1] + eps = 1e-5 + w = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True) + b = torch.randn(dim, dtype=dtype, device="cuda", requires_grad=True) + + # Triton implementation via Function + y = turbodiffusion_ops.LayerNormFunction.apply(x, w, b, eps, True) + loss = y.sum() + loss.backward() + + x_grad = x.grad + w_grad = w.grad + b_grad = b.grad + + # Reference + x_ref = x.detach().clone().requires_grad_() + ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda() + ln.weight.data.copy_(w.detach()) + ln.bias.data.copy_(b.detach()) + + y_ref = ln(x_ref) + loss_ref = y_ref.sum() + loss_ref.backward() + + torch.testing.assert_close(x_grad, x_ref.grad, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(w_grad, ln.weight.grad, atol=1e-2, rtol=1e-2) + torch.testing.assert_close(b_grad, ln.bias.grad, atol=1e-2, rtol=1e-2) + + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_layer_norm_cuda(self, dtype): + if fastvideo_kernel_ops is None or not hasattr(fastvideo_kernel_ops, "layer_norm_cuda"): + pytest.skip("layer_norm_cuda not available") + + shape = (16, 128) + x = torch.randn(shape, dtype=dtype, device="cuda") + dim = shape[-1] + eps = 1e-5 + w = torch.randn(dim, dtype=dtype, device="cuda") + b = torch.randn(dim, dtype=dtype, device="cuda") + + # C++ implementation + # Signature: layer_norm_cuda(Input, eps, W, B, Output) -> Output + out_cuda = torch.empty_like(x) + fastvideo_kernel_ops.layer_norm_cuda(x, eps, w, b, out_cuda) + + # Reference + ln = torch.nn.LayerNorm(dim, eps=eps, elementwise_affine=True, dtype=dtype).cuda() + ln.weight.data.copy_(w) + ln.bias.data.copy_(b) + out_ref = ln(x) + + torch.testing.assert_close(out_cuda, out_ref, atol=1e-2, rtol=1e-2) diff --git a/fastvideo/tests/modal/pr_test.py b/fastvideo/tests/modal/pr_test.py index c9816f33e..fa555e4b2 100644 --- a/fastvideo/tests/modal/pr_test.py +++ b/fastvideo/tests/modal/pr_test.py @@ -53,6 +53,10 @@ def run_test(pytest_command: str): git clone {git_repo} /FastVideo && cd /FastVideo && {checkout_command} && + git submodule update --init --recursive && + cd fastvideo-kernel && + ./build.sh && + cd .. && uv pip install -e .[test] && {pytest_command} """