[kernel] add turbodiffusion kernels (#972)
This commit is contained in:
William Lin
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}
|
||||
}
|
||||
```
|
||||
|
||||
@@ -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}"
|
||||
|
||||
|
||||
@@ -22,6 +22,12 @@ extern std::vector<torch::Tensor> 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);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
CUTLASS_HOST_DEVICE int64_t cdiv(int64_t const& a, int64_t const &b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
template <class T>
|
||||
CUTLASS_HOST_DEVICE T max(T a, T b) { return a > b ? a : b; }
|
||||
|
||||
template <class T>
|
||||
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 <int BlockSize>
|
||||
void create_tensor(
|
||||
torch::Device const &device,
|
||||
std::optional<at::Tensor> &output,
|
||||
std::optional<at::Tensor> &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)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,44 @@
|
||||
#pragma once
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
|
||||
template <class Kernel>
|
||||
__global__ void device_kernel(
|
||||
__grid_constant__ typename Kernel::Params const params
|
||||
) {
|
||||
extern __shared__ char smem[];
|
||||
Kernel op;
|
||||
op(params, smem);
|
||||
}
|
||||
|
||||
template <class Kernel>
|
||||
__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 <class Kernel>
|
||||
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<Kernel>;
|
||||
if (ShmSize >= 48 * 1024) {
|
||||
CUDA_CHECK(cudaFuncSetAttribute(
|
||||
func,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
ShmSize
|
||||
));
|
||||
}
|
||||
func<<<grid_shape, cta_shape, ShmSize, stream>>>(params);
|
||||
CUDA_CHECK(cudaGetLastError());
|
||||
}
|
||||
@@ -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<int4, IsEvenN>(thr_input_ptr, (void*)tmp_reg, limit, pred);
|
||||
else if (n_alignment % 64 == 0)
|
||||
_load<int2, IsEvenN>(thr_input_ptr, (void*)tmp_reg, limit, pred);
|
||||
else
|
||||
_load<InputDtype, IsEvenN>(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<float>(reinterpret_cast<InputDtype const&>(tmp_reg[i]));
|
||||
}
|
||||
|
||||
private:
|
||||
template <class LoadDataType, bool IsEven>
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
@@ -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<int4, IsEvenN>(thr_output_ptr, reg, limit, pred);
|
||||
else if (n_alignment % 64 == 0)
|
||||
_store<int2, IsEvenN>(thr_output_ptr, reg, limit, pred);
|
||||
else
|
||||
_store<OutputDtype, IsEvenN>(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 <class StoreDataType, bool IsEven>
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
@@ -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 <pybind11/pybind11.h>
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/all.h>
|
||||
#include <torch/python.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cutlass/numeric_types.h>
|
||||
|
||||
#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_<cutlass::half_t> (
|
||||
(int8_t*)A.data_ptr(), A_S.data_ptr<float>(),
|
||||
(int8_t*)B.data_ptr(), B_S.data_ptr<float>(),
|
||||
(cutlass::half_t*)C.data_ptr(),
|
||||
m, n, k, swizzle_dir, swizzle_size_log, at::cuda::getCurrentCUDAStream().stream()
|
||||
);
|
||||
break;
|
||||
}
|
||||
|
||||
case torch::kBFloat16:{
|
||||
int8_gemm_<cutlass::bfloat16_t> (
|
||||
(int8_t*)A.data_ptr(), A_S.data_ptr<float>(),
|
||||
(int8_t*)B.data_ptr(), B_S.data_ptr<float>(),
|
||||
(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);
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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 <cuda.h>
|
||||
#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<TileK>{}),
|
||||
make_stride(Int<TileK>{}, Int<1>{})
|
||||
)
|
||||
)
|
||||
);
|
||||
|
||||
using SmemLayoutA = decltype(
|
||||
tile_to_shape(
|
||||
SmemLayoutAtom{},
|
||||
make_shape(Int<TileM>{}, Int<TileK>{}, Int<Stage>{})
|
||||
)
|
||||
);
|
||||
|
||||
using SmemLayoutB = decltype(
|
||||
tile_to_shape(
|
||||
SmemLayoutAtom{},
|
||||
make_shape(Int<TileN>{}, Int<TileK>{}, Int<Stage>{})
|
||||
)
|
||||
);
|
||||
|
||||
using MmaOP = cute::SM80_16x8x32_S32S8S8S32_TN;
|
||||
using TiledMma = decltype(
|
||||
make_tiled_mma(
|
||||
MMA_Atom<MMA_Traits<MmaOP>>{},
|
||||
make_layout(make_shape(
|
||||
_4{}, _2{}, _1{}
|
||||
)),
|
||||
make_tile(Int<64>{}, Int<32>{}, Int<32>{})
|
||||
)
|
||||
);
|
||||
|
||||
using G2SCopyAtomA = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>>, ElementA>;
|
||||
using G2SCopyAtomB = Copy_Atom<Copy_Traits<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>>, 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<Copy_Traits<SM75_U32x4_LDSM_N>, ElementA>;
|
||||
using S2RCopyAtomB = Copy_Atom<Copy_Traits<SM75_U32x4_LDSM_N>, 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<EpiStage>{})
|
||||
)
|
||||
);
|
||||
|
||||
using R2SCopyAtomD = Copy_Atom<UniversalCopy<std::conditional_t<sizeof(OutputDtype) == 4, int32_t, int16_t>>, OutputDtype>;
|
||||
using R2STiledCopyD = decltype(make_tiled_copy_C(R2SCopyAtomD{}, TiledMma{}));
|
||||
|
||||
using S2GCopyAtomD = Copy_Atom<UniversalCopy<uint128_t>, OutputDtype>;
|
||||
using S2GCopyD = decltype(make_tiled_copy(
|
||||
S2GCopyAtomD{},
|
||||
make_layout(Shape<_64, _4>{}),
|
||||
make_layout(Shape<_1, _8>{})
|
||||
));
|
||||
|
||||
using TileShape = decltype(make_shape(Int<TileM>{}, Int<TileN>{}, Int<TileK>{}));
|
||||
|
||||
struct SharedStorageAB: cute::aligned_struct<128> {
|
||||
array_aligned<typename TiledMma::ValTypeA, cosize_v<SmemLayoutA>, 128> smem_A;
|
||||
array_aligned<typename TiledMma::ValTypeB, cosize_v<SmemLayoutB>, 128> smem_B;
|
||||
array_aligned<float, 1> smem_AS;
|
||||
array_aligned<float, 1> smem_BS;
|
||||
array_aligned<int32_t, 1> smem_AF;
|
||||
};
|
||||
|
||||
struct SharedStorageD: cute::aligned_struct<128> {
|
||||
array_aligned<OutputDtype, cosize_v<SmemLayoutD>> 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<SharedStorage*>(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<ElementA>(params.Aptr),
|
||||
make_shape(m, k),
|
||||
make_stride(k, _1{})
|
||||
);
|
||||
Tensor B = make_tensor(
|
||||
make_gmem_ptr<ElementB>(params.Bptr),
|
||||
make_shape(m, k),
|
||||
make_stride(k, _1{})
|
||||
);
|
||||
Tensor AS = make_tensor(
|
||||
make_gmem_ptr<float>(params.ASptr),
|
||||
make_shape(cdiv(m, BlockSize), cdiv(k, BlockSize)),
|
||||
make_stride(cdiv(k, BlockSize), _1{})
|
||||
);
|
||||
Tensor BS = make_tensor(
|
||||
make_gmem_ptr<float>(params.BSptr),
|
||||
make_shape(cdiv(n, BlockSize), cdiv(k, BlockSize)),
|
||||
make_stride(cdiv(k, BlockSize), _1{})
|
||||
);
|
||||
Tensor D = make_tensor(
|
||||
make_gmem_ptr<OutputDtype>(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<X, _1, _1>{});
|
||||
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<ElementA>(shared_storage.storage_AB.smem_A.data()),
|
||||
SmemLayoutA{}
|
||||
);
|
||||
|
||||
Tensor sB = make_tensor(
|
||||
make_smem_ptr<ElementB>(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<float>(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<int32_t>(shared_storage.storage_AB.smem_AF.data()));
|
||||
// auto smem_scale_A = raw_pointer_cast(make_smem_ptr<float>(shared_storage.storage_AB.smem_AS.data()));
|
||||
// auto smem_scale_B = raw_pointer_cast(make_smem_ptr<float>(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<IsEvenM>(g2s_tiled_copy_a, tAgA, tAsA, tAcA, g2s_a_tile, g2s_a_smem, m_limit);
|
||||
copy_AB<IsEvenN>(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<Stage - 2>();
|
||||
__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<AccumulatorDtype, TileM * TileN / ThreadNum, FastInt2Float>(
|
||||
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<IsEvenM>(g2s_tiled_copy_a, tAgA, tAsA, tAcA, g2s_a_tile, g2s_a_smem, m_limit);
|
||||
copy_AB<IsEvenN>(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<OutputDtype>(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<TileM>{}, Int<TileN>{}));
|
||||
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<OutputDtype, float>::value) {
|
||||
cute::copy(r2s_tiled_copy_d, tDrD_r2sx(_, i + j), tDsD_r2s(_, 0, 0, j));
|
||||
} else {
|
||||
auto t = make_tensor_like<OutputDtype>(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();
|
||||
}
|
||||
}
|
||||
|
||||
};
|
||||
|
||||
@@ -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 <class OutputDtype>
|
||||
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<OutputDtype, IsEvenM, IsEvenN>;
|
||||
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<Kernel>;
|
||||
if (ShmSize >= 48 * 1024) {
|
||||
cudaFuncSetAttribute(
|
||||
func,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize,
|
||||
ShmSize
|
||||
);
|
||||
}
|
||||
func<<<grid_shape, block_shape, ShmSize, stream>>>(
|
||||
params
|
||||
);
|
||||
return true;
|
||||
});
|
||||
});
|
||||
return true;
|
||||
}
|
||||
@@ -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 <cuda.h>
|
||||
#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 <int N>
|
||||
CUTLASS_DEVICE void copy_async(
|
||||
void const* gmem_src,
|
||||
void* smem_dst
|
||||
) {
|
||||
uint32_t smem_int_ptr = static_cast<uint32_t>(__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<class LoadType, class T, int NumThreads>
|
||||
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<sizeof(LoadType)>(
|
||||
(void*)((T*)src + i),
|
||||
(void*)((T*)dst + i)
|
||||
);
|
||||
} else {
|
||||
for (int64_t j = 0; j < N - i; ++j)
|
||||
copy_async<sizeof(T)>(
|
||||
(void*)((T*)src + i + j),
|
||||
(void*)((T*)dst + i + j)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<class T, int NumThreads, bool Wait = true, bool Commit = true>
|
||||
CUTLASS_DEVICE void g2s_vector_copy(const void* src, void* dst, size_t N, int64_t thread_idx) {
|
||||
|
||||
uintptr_t src_addr = reinterpret_cast<uintptr_t>(src);
|
||||
|
||||
if (src_addr % 16 == 0) {
|
||||
copy_aligned<int4, T, NumThreads>(src, dst, N, thread_idx);
|
||||
} else if (src_addr % 8 == 0) {
|
||||
copy_aligned<int2, T, NumThreads>(src, dst, N, thread_idx);
|
||||
} else if (src_addr % 4 == 0) {
|
||||
copy_aligned<int, T, NumThreads>(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<class T, int N, bool FastInt2Float>
|
||||
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<T, int32_t>) {
|
||||
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<T, int32_t>) {
|
||||
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;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,63 @@
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/all.h>
|
||||
#include <torch/python.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include "common/common.hpp"
|
||||
#include "norm/layernorm.hpp"
|
||||
|
||||
auto layer_norm(
|
||||
at::Tensor const Input,
|
||||
float eps,
|
||||
std::optional<at::Tensor const> W,
|
||||
std::optional<at::Tensor const> const B,
|
||||
std::optional<at::Tensor> 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);
|
||||
}
|
||||
|
||||
@@ -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<InputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> 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<WeightDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> 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<OutputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven, false, false> saver;
|
||||
if constexpr (std::is_same_v<OutputDtype, float>) {
|
||||
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<Kernel>(params, grid_shape, cta_shape, ShmSize, stream);
|
||||
});
|
||||
return true;
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/all.h>
|
||||
#include <torch/python.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
|
||||
#include "common/common.hpp"
|
||||
#include "norm/rmsnorm.hpp"
|
||||
|
||||
auto rms_norm(
|
||||
at::Tensor const& Input,
|
||||
float eps,
|
||||
const std::optional<at::Tensor>& Weight,
|
||||
std::optional<at::Tensor>& 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);
|
||||
}
|
||||
@@ -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<InputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> 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<WeightDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven> 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<OutputDtype, float>) {
|
||||
output_reg = (OutputDtype*)w;
|
||||
CUTLASS_PRAGMA_UNROLL
|
||||
for (int i = 0; i < NumElementPerThread; ++i)
|
||||
output_reg[i] = OutputDtype(x[i]);
|
||||
}
|
||||
Saver<OutputDtype, 1, MaxHiddenSize, NumThrPerCta, true, IsEven, false, false> 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<Kernel>(params, grid_shape, cta_shape, ShmSize, stream);
|
||||
});
|
||||
return true;
|
||||
}
|
||||
@@ -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 <pybind11/pybind11.h>
|
||||
#include <torch/extension.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/all.h>
|
||||
#include <torch/python.h>
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cutlass/numeric_types.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
|
||||
#include "common/common.hpp"
|
||||
#include "quant/quant.hpp"
|
||||
|
||||
auto quant(
|
||||
torch::Tensor const& Input,
|
||||
std::optional<torch::Tensor>& Output,
|
||||
std::optional<torch::Tensor>& 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<BlockSize>(input_device, Output, Output_S, m, n);
|
||||
|
||||
ElementOut *Optr = (ElementOut*)Output.value().data_ptr();
|
||||
float *OSptr = Output_S.value().data_ptr<float>();
|
||||
|
||||
switch (Input.scalar_type()) {
|
||||
case torch::kHalf:{
|
||||
cutlass::half_t *Iptr = (cutlass::half_t*)Input.data_ptr();
|
||||
quantization<cutlass::half_t, BlockSize, NumThrPerCta> (
|
||||
Iptr, Optr, OSptr, m, n, at::cuda::getCurrentCUDAStream().stream()
|
||||
);
|
||||
break;
|
||||
}
|
||||
|
||||
case torch::kBFloat16:{
|
||||
cutlass::bfloat16_t *Iptr = (cutlass::bfloat16_t*)Input.data_ptr();
|
||||
quantization<cutlass::bfloat16_t, BlockSize, NumThrPerCta> (
|
||||
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);
|
||||
}
|
||||
@@ -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 <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#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<int8_t, float, cutlass::FloatRoundStyle::round_to_nearest>;
|
||||
|
||||
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<OutputDtype, BlockSize, BlockSize, NumThrPerCta, IsEvenM, IsEvenN> 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<InputDtype, BlockSize, BlockSize, NumThrPerCta, IsEvenM, IsEvenN> 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<const uint32_t&>(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<Kernel>(params, grid_shape, cta_shape, ShmSize, stream);
|
||||
});
|
||||
});
|
||||
|
||||
return true;
|
||||
}
|
||||
Submodule
+1
Submodule fastvideo-kernel/include/cutlass added at e67e63c331
@@ -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__",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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}
|
||||
"""
|
||||
|
||||
Reference in New Issue
Block a user