indigo/LTX-2
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Automated PR - 2026-07-07

Browse Source
This commit is contained in:
github-actions[bot]
2026-07-07 16:57:50 +00:00
parent 780984275f
commit 63fd9a4f86
157 changed files with 15976 additions and 5043 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ltx-kernels/csrc/ops/fp6_bitpack.cpp
+73
View File
@@ -0,0 +1,73 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <torch/python.h>
#include <vector>
void fp6_pack_cuda(
at::Tensor& x,
at::Tensor& out,
cudaStream_t stream
);
void fp6_unpack_cuda(
at::Tensor& x,
at::Tensor& out,
cudaStream_t stream
);
at::Tensor fp6_pack(at::Tensor &x) {
// TORCH_CHECK(x.dtype() == torch::kUInt8, "Input tensor must be uint8");
TORCH_CHECK(x.is_cuda(), "Input tensor must be on CUDA");
TORCH_CHECK(x.is_contiguous(), "Input tensor must be contiguous");
TORCH_CHECK(x.dim() == 2, "Input tensor must be 2D [m, n]");
int64_t m = x.size(0);
int64_t n = x.size(1);
TORCH_CHECK(n % 8 == 0, "n must be divisible by 8, got ", n);
// Output shape: [m, n*3/4] since 4 elements of 8-bit = 32 bits, 4 elements of 6-bit = 24 bits = 3 bytes
int64_t n_packed = n * 3 / 4;
auto options = torch::TensorOptions()
.dtype(torch::kUInt8)
.device(x.device());
at::Tensor out = torch::empty({m, n_packed}, options);
at::cuda::CUDAGuard device_guard{x.get_device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
fp6_pack_cuda(x, out, stream);
return out;
}
at::Tensor fp6_unpack(at::Tensor &x, int64_t original_n) {
TORCH_CHECK(x.dtype() == torch::kUInt8, "Input tensor must be uint8");
TORCH_CHECK(x.is_cuda(), "Input tensor must be on CUDA");
TORCH_CHECK(x.is_contiguous(), "Input tensor must be contiguous");
TORCH_CHECK(x.dim() == 2, "Input tensor must be 2D [m, n_packed]");
TORCH_CHECK(original_n % 8 == 0, "original_n must be divisible by 8, got ", original_n);
int64_t m = x.size(0);
int64_t n_packed = x.size(1);
TORCH_CHECK(n_packed == original_n * 3 / 4,
"Packed size mismatch: expected ", original_n * 3 / 4, " got ", n_packed);
auto options = torch::TensorOptions()
.dtype(torch::kUInt8)
.device(x.device());
at::Tensor out = torch::empty({m, original_n}, options);
at::cuda::CUDAGuard device_guard{x.get_device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
fp6_unpack_cuda(x, out, stream);
return out;
}
packages/ltx-kernels/csrc/ops/fp6_pack.cu
+211
View File
@@ -0,0 +1,211 @@
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cuda.h>
#include <ATen/ATen.h>
#include <torch/types.h>
// Device function to pack 8-bit to 6-bit
// 8-bit layout: s e_1 e_2 e_3 m_1 m_2 m_3 m_4 (bits 7-0)
// 6-bit layout: s e_3 m_1 m_2 m_3 m_4 (bits 5-0)
// Drop e_1 (bit 6) and e_2 (bit 5)
__device__ __forceinline__ uint8_t pack_8bit_to_6bit(uint8_t input) {
// Extract the sign bit (bit 7)
uint8_t sign = (input >> 7) & 0x1;
// Extract e_3 (bit 4)
uint8_t e_3 = (input >> 4) & 0x1;
// Extract mantissa bits (bits 3-0)
uint8_t mantissa = input & 0x0F;
// Pack into 6-bit format: s e_3 m_1 m_2 m_3 m_4
uint8_t result = (sign << 5) | (e_3 << 4) | mantissa;
return result & 0x3F; // Mask to 6 bits
}
// Device function to pack 4 x 6-bit values into 3 bytes
__device__ __forceinline__ void pack_4x6bit_to_3bytes(const uint8_t* input_6bit, uint8_t* output_3bytes) {
uint8_t v0 = input_6bit[0] & 0x3F;
uint8_t v1 = input_6bit[1] & 0x3F;
uint8_t v2 = input_6bit[2] & 0x3F;
uint8_t v3 = input_6bit[3] & 0x3F;
// Pack: [v0: 6 bits][v1: 6 bits][v2: 6 bits][v3: 6 bits] = 24 bits = 3 bytes
output_3bytes[0] = (v0 << 2) | (v1 >> 4);
output_3bytes[1] = (v1 << 4) | (v2 >> 2);
output_3bytes[2] = (v2 << 6) | v3;
}
// CUDA kernel for packing 2D tensor
// Input: [m, n] uint8 tensor
// Output: [m, n*3/4] uint8 tensor
__global__ void fp6_pack_kernel(
const uint8_t* __restrict__ input,
uint8_t* __restrict__ output,
int m,
int n,
int n_packed
) {
// Each thread processes one row and 4 elements at a time
int row = blockIdx.x;
int col_group = blockIdx.y * blockDim.x + threadIdx.x;
if (row >= m) return;
// Calculate input and output positions
int input_col = col_group * 4;
if (input_col >= n) return;
int output_col = col_group * 3;
const uint8_t* input_row = input + row * n;
uint8_t* output_row = output + row * n_packed;
uint8_t temp_6bit[4];
// Pack 4 elements
#pragma unroll
for (int i = 0; i < 4; i++) {
if (input_col + i < n) {
temp_6bit[i] = pack_8bit_to_6bit(input_row[input_col + i]);
} else {
temp_6bit[i] = 0;
}
}
// Write 3 bytes to output
uint8_t temp_3bytes[3];
pack_4x6bit_to_3bytes(temp_6bit, temp_3bytes);
if (output_col < n_packed) output_row[output_col] = temp_3bytes[0];
if (output_col + 1 < n_packed) output_row[output_col + 1] = temp_3bytes[1];
if (output_col + 2 < n_packed) output_row[output_col + 2] = temp_3bytes[2];
}
// Device function to unpack 6-bit to 8-bit
__device__ __forceinline__ uint8_t unpack_6bit_to_8bit(uint8_t input) {
input = input & 0x3F; // Ensure only 6 bits
uint8_t sign = (input >> 5) & 0x1;
uint8_t e_3 = (input >> 4) & 0x1;
uint8_t mantissa = input & 0x0F;
// Reconstruct 8-bit with e_1 and e_2 set to 0
uint8_t result = (sign << 7) | (e_3 << 4) | mantissa;
return result;
}
// Device function to unpack 3 bytes into 4 x 6-bit values
__device__ __forceinline__ void unpack_3bytes_to_4x6bit(const uint8_t* input_3bytes, uint8_t* output_6bit) {
output_6bit[0] = (input_3bytes[0] >> 2) & 0x3F;
output_6bit[1] = ((input_3bytes[0] << 4) | (input_3bytes[1] >> 4)) & 0x3F;
output_6bit[2] = ((input_3bytes[1] << 2) | (input_3bytes[2] >> 6)) & 0x3F;
output_6bit[3] = input_3bytes[2] & 0x3F;
}
// CUDA kernel for unpacking 2D tensor
// Input: [m, n_packed] uint8 tensor
// Output: [m, n] uint8 tensor
__global__ void fp6_unpack_kernel(
const uint8_t* __restrict__ input,
uint8_t* __restrict__ output,
int m,
int n_packed,
int n
) {
// Each thread processes one row and 4 elements at a time
int row = blockIdx.x;
int col_group = blockIdx.y * blockDim.x + threadIdx.x;
if (row >= m) return;
// Calculate input and output positions
int input_col = col_group * 3;
if (input_col >= n_packed) return;
int output_col = col_group * 4;
const uint8_t* input_row = input + row * n_packed;
uint8_t* output_row = output + row * n;
// Read 3 bytes
uint8_t temp_3bytes[3];
temp_3bytes[0] = (input_col < n_packed) ? input_row[input_col] : 0;
temp_3bytes[1] = (input_col + 1 < n_packed) ? input_row[input_col + 1] : 0;
temp_3bytes[2] = (input_col + 2 < n_packed) ? input_row[input_col + 2] : 0;
// Unpack to 4 x 6-bit values
uint8_t temp_6bit[4];
unpack_3bytes_to_4x6bit(temp_3bytes, temp_6bit);
// Convert to 8-bit and write
#pragma unroll
for (int i = 0; i < 4; i++) {
if (output_col + i < n) {
output_row[output_col + i] = unpack_6bit_to_8bit(temp_6bit[i]);
}
}
}
// Host function to launch pack kernel
void fp6_pack_cuda(
at::Tensor& x,
at::Tensor& out,
cudaStream_t stream
) {
int m = x.size(0);
int n = x.size(1);
int n_packed = out.size(1);
const uint8_t* input_ptr = (uint8_t*)x.data_ptr();
uint8_t* output_ptr = (uint8_t*)out.data_ptr();
// Each thread handles 4 input elements -> 3 output bytes
int num_groups = (n + 3) / 4;
int threads = 256;
dim3 blocks(m, (num_groups + threads - 1) / threads);
fp6_pack_kernel<<<blocks, threads, 0, stream>>>(
input_ptr,
output_ptr,
m,
n,
n_packed
);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
// Host function to launch unpack kernel
void fp6_unpack_cuda(
at::Tensor& x,
at::Tensor& out,
cudaStream_t stream
) {
int m = x.size(0);
int n_packed = x.size(1);
int n = out.size(1);
const uint8_t* input_ptr = (uint8_t*)x.data_ptr();
uint8_t* output_ptr = (uint8_t*)out.data_ptr();
// Each thread handles 3 input bytes -> 4 output elements
int num_groups = (n + 3) / 4;
int threads = 256;
dim3 blocks(m, (num_groups + threads - 1) / threads);
fp6_unpack_kernel<<<blocks, threads, 0, stream>>>(
input_ptr,
output_ptr,
m,
n_packed,
n
);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
packages/ltx-kernels/csrc/ops/include/fast_hadamard_transform.h
+163
View File
@@ -0,0 +1,163 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#pragma once
////////////////////////////////////////////////////////////////////////////////////////////////////
struct HadamardParamsBase {
using index_t = int64_t;
int batch, dim, log_N;
index_t x_batch_stride;
index_t out_batch_stride;
float scale;
// Common data pointers.
void *__restrict__ x_ptr;
void *__restrict__ out_ptr;
};
struct UnifiedHadamardParamsBase{
using index_t = int64_t;
int batch, dim, log_N;
int batch_fma_change;
index_t x_batch_stride;
index_t out_batch_stride;
index_t fma_batch_stride;
index_t cos_freq_batch_stride;
index_t sin_freq_batch_stride;
float scale;
// Common data pointers.
void *__restrict__ x_ptr;
void *__restrict__ out_ptr;
void *__restrict__ out_scales_ptr;
void *__restrict__ y_scale_ptr;
void *__restrict__ z_shift_ptr;
void *__restrict__ weights_ptr;
void *__restrict__ cos_freq_ptr;
void *__restrict__ sin_freq_ptr;
};
struct DequantHadamardParamsBase {
using index_t = int64_t;
int batch, dim, log_N;
index_t x_batch_stride;
index_t out_batch_stride;
float scale;
// Common data pointers.
void *__restrict__ x_ptr;
void *__restrict__ scales_ptr;
void *__restrict__ out_ptr;
};
struct QuantHadamardParamsBase {
using index_t = int64_t;
int batch, dim, log_N;
index_t x_batch_stride;
index_t out_batch_stride;
float scale;
// Common data pointers.
void *__restrict__ x_ptr;
void *__restrict__ out_ptr;
void *__restrict__ out_scales_ptr;
};
struct NormFMAHadamardParamsBase {
using index_t = int64_t;
int batch, dim, log_N;
int seqlen;
index_t x_batch_stride;
index_t out_batch_stride;
index_t fma_batch_stride;
float scale;
// Common data pointers.
void *__restrict__ x_ptr;
void *__restrict__ out_ptr;
void *__restrict__ y_scale_ptr;
void *__restrict__ z_shift_ptr;
void *__restrict__ weights_ptr;
};
struct NormRopeHadamardParamsBase {
using index_t = int64_t;
int batch, dim, log_N;
index_t x_batch_stride;
index_t out_batch_stride;
index_t cos_freq_batch_stride;
index_t sin_freq_batch_stride;
float scale;
// Common data pointers.
void *__restrict__ x_ptr;
void *__restrict__ out_ptr;
void *__restrict__ cos_freq_ptr;
void *__restrict__ sin_freq_ptr;
void *__restrict__ weights_ptr;
};
struct NormHadamardParamsBase {
using index_t = int64_t;
int batch, dim, log_N;
index_t x_batch_stride;
index_t out_batch_stride;
float scale;
// Common data pointers.
void *__restrict__ x_ptr;
void *__restrict__ out_ptr;
void *__restrict__ weights_ptr;
};
struct RopeHadamardParamsBase {
using index_t = int64_t;
int batch, dim, log_N;
index_t x_batch_stride;
index_t out_batch_stride;
index_t cos_freq_batch_stride;
index_t sin_freq_batch_stride;
float scale;
// Common data pointers.
void *__restrict__ x_ptr;
void *__restrict__ out_ptr;
void *__restrict__ cos_freq_ptr;
void *__restrict__ sin_freq_ptr;
};
packages/ltx-kernels/csrc/ops/include/fast_hadamard_transform_common.h
+319
View File
@@ -0,0 +1,319 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
#pragma once
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#define FULL_MASK 0xffffffff
////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename TYPE> struct QuantMax {};
template<> struct QuantMax<int8_t> { static constexpr float value = 127.0; };
template<> struct QuantMax<at::Float8_e4m3fn> { static constexpr float value = 256.0; };
struct uint8 {
uint4 u;
uint4 v;
};
template<int BYTES> struct BytesToType {};
template<>
struct BytesToType<32> {
using Type = uint8;
static_assert(sizeof(Type) == 32);
};
template<> struct BytesToType<16> {
using Type = uint4;
static_assert(sizeof(Type) == 16);
};
template<> struct BytesToType<8> {
using Type = uint64_t;
static_assert(sizeof(Type) == 8);
};
template<> struct BytesToType<4> {
using Type = uint32_t;
static_assert(sizeof(Type) == 4);
};
template<> struct BytesToType<2> {
using Type = uint16_t;
static_assert(sizeof(Type) == 2);
};
template<> struct BytesToType<1> {
using Type = uint8_t;
static_assert(sizeof(Type) == 1);
};
////////////////////////////////////////////////////////////////////////////////////////////////////
template<typename T>
struct SumOp {
__device__ inline T operator()(T const & x, T const & y) { return x + y; }
};
template<typename T>
struct MaxOp {
__device__ inline T operator()(T const & x, T const & y) { return max(x, y); }
};
template <>
struct MaxOp<float> {
// This is slightly faster
__device__ inline float operator()(float const &x, float const &y) { return max(x, y); }
};
template<int THREADS>
struct Allreduce {
static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
template<typename T, typename Operator>
static __device__ inline T run(T x, Operator &op) {
constexpr int OFFSET = THREADS / 2;
x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
return Allreduce<OFFSET>::run(x, op);
}
};
template<>
struct Allreduce<2> {
template<typename T, typename Operator>
static __device__ inline T run(T x, Operator &op) {
x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
return x;
}
};
////////////////////////////////////////////////////////////////////////////////////////////////////
// https://stackoverflow.com/questions/35311711/whats-the-right-way-to-compute-integral-base-2-logarithms-at-compile-time
constexpr int cilog2(int val) { return val > 0 ? 1 + cilog2(val >> 1) : -1; }
////////////////////////////////////////////////////////////////////////////////////////////////////
template<int kLogN, int kNChunks>
__device__ __forceinline__ void hadamard_mult_thread(float x[kNChunks][1 << kLogN]) {
constexpr int N = 1 << kLogN;
#pragma unroll
for (int i = 0; i < kLogN; ++i) {
const int stride = 1 << i;
#pragma unroll
for (int j = 0; j < N / 2; ++j) {
const int lo = j & (stride - 1);
const int idx = (j - lo) * 2 + lo;
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
const float a = x[c][idx];
const float b = x[c][idx + stride];
x[c][idx] = a + b;
x[c][idx + stride] = a - b;
}
}
}
}
template<int kLogWarpSize, int kStepStart, int kNChunks, int kNItems>
__device__ __forceinline__ void hadamard_mult_warp(float x[kNChunks][kNItems]) {
constexpr int N = 1 << kLogWarpSize;
int lane_id = threadIdx.x % N;
#pragma unroll
for (int step = kStepStart; step < kLogWarpSize; ++step) {
const int lane_mask = 1 << step;
const float sign = (lane_id & lane_mask) ? -1.f : 1.f;
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
#pragma unroll
for (int i = 0; i < kNItems; ++i) {
float x_val_other = __shfl_xor_sync(FULL_MASK, x[c][i], lane_mask);
x[c][i] = sign * x[c][i] + x_val_other;
}
}
}
}
////////////////////////////////////////////////////////////////////////////////////////////////////
template <int kNChunks, int kNElts, typename input_t>
inline __device__ void load_input(input_t *x, float x_vals[kNChunks][kNElts], int dim) {
using vec_t = typename BytesToType<sizeof(input_t) * kNElts>::Type;
input_t x_vals_load[kNChunks][kNElts] = {0};
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
if ((c * blockDim.x + threadIdx.x) * kNElts < dim) {
reinterpret_cast<vec_t*>(x_vals_load)[c] = reinterpret_cast<const vec_t*>(x)[c * blockDim.x + threadIdx.x];
}
}
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
#pragma unroll
for (int i = 0; i < kNElts; ++i) { x_vals[c][i] = float(x_vals_load[c][i]); }
}
}
template <int kNChunks, int kNElts, typename output_t, bool do_round>
inline __device__ void store_output(output_t *out, float out_vals[kNChunks][kNElts], int dim, float scale=1.f) {
using vec_t = typename BytesToType<sizeof(output_t) * kNElts>::Type;
output_t out_vals_store[kNChunks][kNElts];
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
#pragma unroll
for (int i = 0; i < kNElts; ++i) {
if constexpr (do_round){
out_vals_store[c][i] = round(out_vals[c][i] * scale);
} else {
out_vals_store[c][i] = out_vals[c][i] * scale;
}
}
}
#pragma unroll
for (int c = 0; c < kNChunks; ++c) {
if ((c * blockDim.x + threadIdx.x) * kNElts < dim) {
reinterpret_cast<vec_t*>(out)[c * blockDim.x + threadIdx.x] = reinterpret_cast<const vec_t*>(out_vals_store)[c];
}
}
}
////////////////////////////////////////////////////////////////////////////////////////////////////
// Pre=true means the exchange before the hadamard_mult_warp, Pre=false means after.
template <int kNChunks, int kChunksPerExchange, int kNElts, int kWarpSize, int kNWarps, bool Pre, typename vec_t>
inline __device__ void exchange_smem_pre(float x_vals[kNChunks][kNElts], vec_t *smem) {
constexpr int kNThreads = kWarpSize * kNWarps;
constexpr int kNExchangePerVec = kNElts / (sizeof(vec_t) / sizeof(float));
const int warp_id = threadIdx.x / kWarpSize;
const int lane_id = threadIdx.x % kWarpSize;
const int row_t = threadIdx.x % kNWarps;
const int col_t = threadIdx.x / kNWarps;
// We use the XOR swizzle trick (new_col = col ^ row) to avoid / reduce smem bank conflicts.
#pragma unroll
for (int c0 = 0; c0 < kNChunks / kChunksPerExchange; ++c0) {
__syncthreads();
#pragma unroll
for (int c1 = 0; c1 < kChunksPerExchange; ++c1) {
#pragma unroll
for (int r = 0; r < kNExchangePerVec; ++r) {
smem[(c1 * kNExchangePerVec + r) * kNThreads + (Pre ? warp_id * kWarpSize + lane_id ^ warp_id : row_t * kWarpSize + col_t ^ row_t)] = reinterpret_cast<vec_t*>(x_vals[c0 * kChunksPerExchange + c1])[r];
}
}
__syncthreads();
#pragma unroll
for (int c1 = 0; c1 < kChunksPerExchange; ++c1) {
#pragma unroll
for (int r = 0; r < kNExchangePerVec; ++r) {
reinterpret_cast<vec_t*>(x_vals[c0 * kChunksPerExchange + c1])[r] = smem[(c1 * kNExchangePerVec + r) * kNThreads + (Pre ? row_t * kWarpSize + col_t ^ row_t : warp_id * kWarpSize + lane_id ^ warp_id)];
}
}
}
}
inline __device__ float gelu_approximate(float x){
constexpr float sqrthalfpi2 = 0.7978845608028653558798921198687637369517172623298693153318516593f;
constexpr float factor = 0.044715f;
return 0.5f*x*(1.0f + tanhf(sqrthalfpi2*(x + factor*x*x*x)));
}
template <int kNChunks, int kNElts>
inline __device__ void fused_gelu(float x_vals[kNChunks][kNElts]){
#pragma unroll
for (size_t c = 0; c < kNChunks; c++)
{
#pragma unroll
for (size_t i = 0; i < kNElts; i++)
{
x_vals[c][i] = gelu_approximate(x_vals[c][i]);
}
}
}
template <int kNChunks, int kNElts, int kNWarps, bool norm_affine>
inline __device__ void fused_rms_norm(float x_vals[kNChunks][kNElts], float weights_vals[kNChunks][kNElts], float* smem_sum, float dim){
float thread_squared_sum = 0.0f;
const int warp_id = threadIdx.x / 32;
#pragma unroll
for (size_t c = 0; c < kNChunks; c++)
{
#pragma unroll
for (size_t i = 0; i < kNElts; i++)
{
thread_squared_sum += x_vals[c][i] * x_vals[c][i];
}
}
SumOp<float> sum_op;
float warp_sum = Allreduce<32>::run(thread_squared_sum, sum_op);
if(threadIdx.x % 32 == 0){
smem_sum[warp_id] = warp_sum;
}
__syncthreads();
float norm = 0.0f;
#pragma unroll
for (size_t i = 0; i < kNWarps; i++)
{
norm += smem_sum[i];
}
norm *= 1.0f/dim;
norm = rsqrtf(norm + 0.0000001f);
#pragma unroll
for (size_t c = 0; c < kNChunks; c++)
{
#pragma unroll
for (size_t i = 0; i < kNElts; i++)
{
if constexpr (norm_affine){
x_vals[c][i] *= (norm * weights_vals[c][i]);
} else {
x_vals[c][i] *= norm;
}
}
}
}
template <int kNChunks, int kNElts>
inline __device__ void fused_rope(float x_vals[kNChunks][kNElts], float sin_freqs_vals[kNChunks][kNElts], float cos_freqs_vals[kNChunks][kNElts]){
#pragma unroll
for (size_t c = 0; c < kNChunks; c++)
{
#pragma unroll
for (size_t i = 0; i < kNElts; i+=2)
{
float x_1 = x_vals[c][i];
float x_2 = x_vals[c][i+1];
x_vals[c][i] = -x_2*sin_freqs_vals[c][i] + x_1*cos_freqs_vals[c][i];
x_vals[c][i+1] = x_1*sin_freqs_vals[c][i+1] + x_2*cos_freqs_vals[c][i+1];
}
}
}
template <int kNChunks, int kNElts, bool add_one_scale>
inline __device__ void fused_multiply_add(float x_vals[kNChunks][kNElts], float y_scale_vals[kNChunks][kNElts], float z_shift_vals[kNChunks][kNElts]) {
#pragma unroll
for (size_t c = 0; c < kNChunks; c++)
{
#pragma unroll
for (size_t i = 0; i < kNElts; i++)
{
if constexpr (add_one_scale){
x_vals[c][i] = x_vals[c][i] * (1.0f + y_scale_vals[c][i]) + z_shift_vals[c][i];
} else {
x_vals[c][i] = x_vals[c][i] * y_scale_vals[c][i] + z_shift_vals[c][i];
}
}
}
}
packages/ltx-kernels/csrc/ops/include/fast_hadamard_transform_special.h
+138
View File
@@ -0,0 +1,138 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// This file is auto-generated. See "code_gen.py"
#pragma once
__device__ __forceinline__ void hadamard_mult_thread_12(float x[12]) {
float out[12];
out[0] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] + x[6] + x[7] + x[8] + x[9] + x[10] + x[11];
out[1] = - x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] - x[11];
out[2] = + x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11];
out[3] = + x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11];
out[4] = + x[0] + x[1] + x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11];
out[5] = + x[0] - x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] + x[9] - x[10] + x[11];
out[6] = + x[0] + x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11];
out[7] = + x[0] - x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11];
out[8] = + x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] + x[8] - x[9] + x[10] + x[11];
out[9] = + x[0] - x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11];
out[10] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11];
out[11] = + x[0] - x[1] + x[2] - x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11];
#pragma unroll
for (int i = 0; i < 12; i++) { x[i] = out[i]; }
}
__device__ __forceinline__ void hadamard_mult_thread_20(float x[20]) {
float out[20];
out[0] = + x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19];
out[1] = - x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] + x[19];
out[2] = - x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19];
out[3] = - x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] + x[13] + x[14] + x[15] + x[16] - x[17] + x[18] - x[19];
out[4] = - x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] + x[19];
out[5] = - x[0] + x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19];
out[6] = + x[0] - x[1] + x[2] + x[3] + x[4] - x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - x[14] + x[15] + x[16] + x[17] - x[18] - x[19];
out[7] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] + x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] + x[16] + x[17] + x[18] - x[19];
out[8] = + x[0] + x[1] + x[2] - x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] + x[17] + x[18] + x[19];
out[9] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] - x[17] + x[18] + x[19];
out[10] = - x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] - x[15] + x[16] + x[17] + x[18] + x[19];
out[11] = - x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19];
out[12] = + x[0] - x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19];
out[13] = + x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] - x[18] + x[19];
out[14] = - x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] + x[15] + x[16] + x[17] + x[18] - x[19];
out[15] = - x[0] + x[1] - x[2] - x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] - x[17] - x[18] - x[19];
out[16] = + x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] - x[18] - x[19];
out[17] = - x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] - x[19];
out[18] = - x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19];
out[19] = + x[0] - x[1] - x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19];
#pragma unroll
for (int i = 0; i < 20; i++) { x[i] = out[i]; }
}
__device__ __forceinline__ void hadamard_mult_thread_28(float x[28]) {
float out[28];
out[0] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] - x[27];
out[1] = - x[0] + x[1] - x[2] - x[3] - x[4] - x[5] - x[6] + x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] - x[27];
out[2] = - x[0] - x[1] + x[2] - x[3] - x[4] - x[5] - x[6] - x[7] + x[8] + x[9] + x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] + x[27];
out[3] = - x[0] - x[1] - x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] + x[9] + x[10] + x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] + x[27];
out[4] = - x[0] - x[1] - x[2] - x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] + x[10] + x[11] + x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] - x[27];
out[5] = - x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] - x[7] - x[8] - x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] - x[27];
out[6] = - x[0] - x[1] - x[2] - x[3] - x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] - x[16] - x[17] + x[18] - x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] + x[27];
out[7] = - x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] + x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] - x[24] - x[25] + x[26] - x[27];
out[8] = - x[0] - x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] + x[8] - x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] + x[27];
out[9] = + x[0] - x[1] - x[2] - x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] - x[27];
out[10] = + x[0] + x[1] - x[2] - x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] - x[27];
out[11] = + x[0] + x[1] + x[2] - x[3] - x[4] - x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] + x[27];
out[12] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] - x[22] - x[23] + x[24] - x[25] + x[26] - x[27];
out[13] = - x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] - x[26] + x[27];
out[14] = - x[0] + x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27];
out[15] = + x[0] - x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] - x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] + x[26] + x[27];
out[16] = - x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] - x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] - x[23] - x[24] + x[25] + x[26] + x[27];
out[17] = + x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] + x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] - x[24] - x[25] + x[26] + x[27];
out[18] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] + x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] - x[25] - x[26] + x[27];
out[19] = - x[0] + x[1] + x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] - x[27];
out[20] = + x[0] - x[1] + x[2] + x[3] - x[4] + x[5] - x[6] - x[7] - x[8] + x[9] + x[10] - x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] - x[27];
out[21] = - x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] - x[26] - x[27];
out[22] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] - x[27];
out[23] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] - x[26] - x[27];
out[24] = - x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] + x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] + x[24] - x[25] - x[26] - x[27];
out[25] = - x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] + x[10] - x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] - x[27];
out[26] = + x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] - x[10] + x[11] - x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] + x[26] - x[27];
out[27] = + x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] - x[11] + x[12] - x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] + x[27];
#pragma unroll
for (int i = 0; i < 28; i++) { x[i] = out[i]; }
}
__device__ __forceinline__ void hadamard_mult_thread_40(float x[40]) {
float out[40];
out[0] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] + x[20] - x[21] - x[22] - x[23] - x[24] - x[25] - x[26] - x[27] - x[28] - x[29] - x[30] - x[31] - x[32] - x[33] - x[34] - x[35] - x[36] - x[37] - x[38] - x[39];
out[1] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] + x[20] + x[21] - x[22] + x[23] + x[24] - x[25] - x[26] - x[27] - x[28] + x[29] - x[30] + x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] - x[38] + x[39];
out[2] = + x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] + x[20] + x[21] + x[22] - x[23] + x[24] + x[25] - x[26] - x[27] - x[28] - x[29] + x[30] - x[31] + x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] - x[39];
out[3] = + x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] + x[18] - x[19] + x[20] - x[21] + x[22] + x[23] - x[24] + x[25] + x[26] - x[27] - x[28] - x[29] - x[30] + x[31] - x[32] + x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39];
out[4] = + x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] + x[16] + x[17] + x[18] + x[19] + x[20] - x[21] - x[22] + x[23] + x[24] - x[25] + x[26] + x[27] - x[28] - x[29] - x[30] - x[31] + x[32] - x[33] + x[34] - x[35] + x[36] + x[37] + x[38] + x[39];
out[5] = + x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19] + x[20] + x[21] - x[22] - x[23] + x[24] + x[25] - x[26] + x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] - x[34] + x[35] - x[36] + x[37] + x[38] + x[39];
out[6] = + x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] + x[19] + x[20] + x[21] + x[22] - x[23] - x[24] + x[25] + x[26] - x[27] + x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] - x[35] + x[36] - x[37] + x[38] + x[39];
out[7] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] + x[20] + x[21] + x[22] + x[23] - x[24] - x[25] + x[26] + x[27] - x[28] + x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] - x[36] + x[37] - x[38] + x[39];
out[8] = + x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] + x[20] + x[21] + x[22] + x[23] + x[24] - x[25] - x[26] + x[27] + x[28] - x[29] + x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] - x[37] + x[38] - x[39];
out[9] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] + x[19] + x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] - x[27] + x[28] + x[29] - x[30] + x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] - x[38] + x[39];
out[10] = + x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] + x[20] + x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27] - x[28] + x[29] + x[30] - x[31] + x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] - x[39];
out[11] = + x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] + x[20] - x[21] + x[22] - x[23] + x[24] + x[25] + x[26] + x[27] - x[28] - x[29] + x[30] + x[31] - x[32] + x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39];
out[12] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] + x[20] + x[21] - x[22] + x[23] - x[24] + x[25] + x[26] + x[27] + x[28] - x[29] - x[30] + x[31] + x[32] - x[33] + x[34] + x[35] - x[36] - x[37] - x[38] - x[39];
out[13] = + x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] - x[17] - x[18] - x[19] + x[20] - x[21] + x[22] - x[23] + x[24] - x[25] + x[26] + x[27] + x[28] + x[29] - x[30] - x[31] + x[32] + x[33] - x[34] + x[35] + x[36] - x[37] - x[38] - x[39];
out[14] = + x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] + x[20] - x[21] - x[22] + x[23] - x[24] + x[25] - x[26] + x[27] + x[28] + x[29] + x[30] - x[31] - x[32] + x[33] + x[34] - x[35] + x[36] + x[37] - x[38] - x[39];
out[15] = + x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] + x[20] - x[21] - x[22] - x[23] + x[24] - x[25] + x[26] - x[27] + x[28] + x[29] + x[30] + x[31] - x[32] - x[33] + x[34] + x[35] - x[36] + x[37] + x[38] - x[39];
out[16] = + x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] + x[20] - x[21] - x[22] - x[23] - x[24] + x[25] - x[26] + x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] - x[34] + x[35] + x[36] - x[37] + x[38] + x[39];
out[17] = + x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] + x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] - x[27] + x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] - x[35] + x[36] + x[37] - x[38] + x[39];
out[18] = + x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] + x[20] + x[21] + x[22] - x[23] - x[24] - x[25] - x[26] + x[27] - x[28] + x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] - x[36] + x[37] + x[38] - x[39];
out[19] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] + x[20] - x[21] + x[22] + x[23] - x[24] - x[25] - x[26] - x[27] + x[28] - x[29] + x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] - x[37] + x[38] + x[39];
out[20] = + x[0] - x[1] - x[2] - x[3] - x[4] - x[5] - x[6] - x[7] - x[8] - x[9] - x[10] - x[11] - x[12] - x[13] - x[14] - x[15] - x[16] - x[17] - x[18] - x[19] - x[20] + x[21] + x[22] + x[23] + x[24] + x[25] + x[26] + x[27] + x[28] + x[29] + x[30] + x[31] + x[32] + x[33] + x[34] + x[35] + x[36] + x[37] + x[38] + x[39];
out[21] = + x[0] + x[1] - x[2] + x[3] + x[4] - x[5] - x[6] - x[7] - x[8] + x[9] - x[10] + x[11] - x[12] + x[13] + x[14] + x[15] + x[16] - x[17] - x[18] + x[19] - x[20] - x[21] + x[22] - x[23] - x[24] + x[25] + x[26] + x[27] + x[28] - x[29] + x[30] - x[31] + x[32] - x[33] - x[34] - x[35] - x[36] + x[37] + x[38] - x[39];
out[22] = + x[0] + x[1] + x[2] - x[3] + x[4] + x[5] - x[6] - x[7] - x[8] - x[9] + x[10] - x[11] + x[12] - x[13] + x[14] + x[15] + x[16] + x[17] - x[18] - x[19] - x[20] - x[21] - x[22] + x[23] - x[24] - x[25] + x[26] + x[27] + x[28] + x[29] - x[30] + x[31] - x[32] + x[33] - x[34] - x[35] - x[36] - x[37] + x[38] + x[39];
out[23] = + x[0] - x[1] + x[2] + x[3] - x[4] + x[5] + x[6] - x[7] - x[8] - x[9] - x[10] + x[11] - x[12] + x[13] - x[14] + x[15] + x[16] + x[17] + x[18] - x[19] - x[20] + x[21] - x[22] - x[23] + x[24] - x[25] - x[26] + x[27] + x[28] + x[29] + x[30] - x[31] + x[32] - x[33] + x[34] - x[35] - x[36] - x[37] - x[38] + x[39];
out[24] = + x[0] - x[1] - x[2] + x[3] + x[4] - x[5] + x[6] + x[7] - x[8] - x[9] - x[10] - x[11] + x[12] - x[13] + x[14] - x[15] + x[16] + x[17] + x[18] + x[19] - x[20] + x[21] + x[22] - x[23] - x[24] + x[25] - x[26] - x[27] + x[28] + x[29] + x[30] + x[31] - x[32] + x[33] - x[34] + x[35] - x[36] - x[37] - x[38] - x[39];
out[25] = + x[0] + x[1] - x[2] - x[3] + x[4] + x[5] - x[6] + x[7] + x[8] - x[9] - x[10] - x[11] - x[12] + x[13] - x[14] + x[15] - x[16] + x[17] + x[18] + x[19] - x[20] - x[21] + x[22] + x[23] - x[24] - x[25] + x[26] - x[27] - x[28] + x[29] + x[30] + x[31] + x[32] - x[33] + x[34] - x[35] + x[36] - x[37] - x[38] - x[39];
out[26] = + x[0] + x[1] + x[2] - x[3] - x[4] + x[5] + x[6] - x[7] + x[8] + x[9] - x[10] - x[11] - x[12] - x[13] + x[14] - x[15] + x[16] - x[17] + x[18] + x[19] - x[20] - x[21] - x[22] + x[23] + x[24] - x[25] - x[26] + x[27] - x[28] - x[29] + x[30] + x[31] + x[32] + x[33] - x[34] + x[35] - x[36] + x[37] - x[38] - x[39];
out[27] = + x[0] + x[1] + x[2] + x[3] - x[4] - x[5] + x[6] + x[7] - x[8] + x[9] + x[10] - x[11] - x[12] - x[13] - x[14] + x[15] - x[16] + x[17] - x[18] + x[19] - x[20] - x[21] - x[22] - x[23] + x[24] + x[25] - x[26] - x[27] + x[28] - x[29] - x[30] + x[31] + x[32] + x[33] + x[34] - x[35] + x[36] - x[37] + x[38] - x[39];
out[28] = + x[0] + x[1] + x[2] + x[3] + x[4] - x[5] - x[6] + x[7] + x[8] - x[9] + x[10] + x[11] - x[12] - x[13] - x[14] - x[15] + x[16] - x[17] + x[18] - x[19] - x[20] - x[21] - x[22] - x[23] - x[24] + x[25] + x[26] - x[27] - x[28] + x[29] - x[30] - x[31] + x[32] + x[33] + x[34] + x[35] - x[36] + x[37] - x[38] + x[39];
out[29] = + x[0] - x[1] + x[2] + x[3] + x[4] + x[5] - x[6] - x[7] + x[8] + x[9] - x[10] + x[11] + x[12] - x[13] - x[14] - x[15] - x[16] + x[17] - x[18] + x[19] - x[20] + x[21] - x[22] - x[23] - x[24] - x[25] + x[26] + x[27] - x[28] - x[29] + x[30] - x[31] - x[32] + x[33] + x[34] + x[35] + x[36] - x[37] + x[38] - x[39];
out[30] = + x[0] + x[1] - x[2] + x[3] + x[4] + x[5] + x[6] - x[7] - x[8] + x[9] + x[10] - x[11] + x[12] + x[13] - x[14] - x[15] - x[16] - x[17] + x[18] - x[19] - x[20] - x[21] + x[22] - x[23] - x[24] - x[25] - x[26] + x[27] + x[28] - x[29] - x[30] + x[31] - x[32] - x[33] + x[34] + x[35] + x[36] + x[37] - x[38] + x[39];
out[31] = + x[0] - x[1] + x[2] - x[3] + x[4] + x[5] + x[6] + x[7] - x[8] - x[9] + x[10] + x[11] - x[12] + x[13] + x[14] - x[15] - x[16] - x[17] - x[18] + x[19] - x[20] + x[21] - x[22] + x[23] - x[24] - x[25] - x[26] - x[27] + x[28] + x[29] - x[30] - x[31] + x[32] - x[33] - x[34] + x[35] + x[36] + x[37] + x[38] - x[39];
out[32] = + x[0] + x[1] - x[2] + x[3] - x[4] + x[5] + x[6] + x[7] + x[8] - x[9] - x[10] + x[11] + x[12] - x[13] + x[14] + x[15] - x[16] - x[17] - x[18] - x[19] - x[20] - x[21] + x[22] - x[23] + x[24] - x[25] - x[26] - x[27] - x[28] + x[29] + x[30] - x[31] - x[32] + x[33] - x[34] - x[35] + x[36] + x[37] + x[38] + x[39];
out[33] = + x[0] - x[1] + x[2] - x[3] + x[4] - x[5] + x[6] + x[7] + x[8] + x[9] - x[10] - x[11] + x[12] + x[13] - x[14] + x[15] + x[16] - x[17] - x[18] - x[19] - x[20] + x[21] - x[22] + x[23] - x[24] + x[25] - x[26] - x[27] - x[28] - x[29] + x[30] + x[31] - x[32] - x[33] + x[34] - x[35] - x[36] + x[37] + x[38] + x[39];
out[34] = + x[0] - x[1] - x[2] + x[3] - x[4] + x[5] - x[6] + x[7] + x[8] + x[9] + x[10] - x[11] - x[12] + x[13] + x[14] - x[15] + x[16] + x[17] - x[18] - x[19] - x[20] + x[21] + x[22] - x[23] + x[24] - x[25] + x[26] - x[27] - x[28] - x[29] - x[30] + x[31] + x[32] - x[33] - x[34] + x[35] - x[36] - x[37] + x[38] + x[39];
out[35] = + x[0] - x[1] - x[2] - x[3] + x[4] - x[5] + x[6] - x[7] + x[8] + x[9] + x[10] + x[11] - x[12] - x[13] + x[14] + x[15] - x[16] + x[17] + x[18] - x[19] - x[20] + x[21] + x[22] + x[23] - x[24] + x[25] - x[26] + x[27] - x[28] - x[29] - x[30] - x[31] + x[32] + x[33] - x[34] - x[35] + x[36] - x[37] - x[38] + x[39];
out[36] = + x[0] - x[1] - x[2] - x[3] - x[4] + x[5] - x[6] + x[7] - x[8] + x[9] + x[10] + x[11] + x[12] - x[13] - x[14] + x[15] + x[16] - x[17] + x[18] + x[19] - x[20] + x[21] + x[22] + x[23] + x[24] - x[25] + x[26] - x[27] + x[28] - x[29] - x[30] - x[31] - x[32] + x[33] + x[34] - x[35] - x[36] + x[37] - x[38] - x[39];
out[37] = + x[0] + x[1] - x[2] - x[3] - x[4] - x[5] + x[6] - x[7] + x[8] - x[9] + x[10] + x[11] + x[12] + x[13] - x[14] - x[15] + x[16] + x[17] - x[18] + x[19] - x[20] - x[21] + x[22] + x[23] + x[24] + x[25] - x[26] + x[27] - x[28] + x[29] - x[30] - x[31] - x[32] - x[33] + x[34] + x[35] - x[36] - x[37] + x[38] - x[39];
out[38] = + x[0] + x[1] + x[2] - x[3] - x[4] - x[5] - x[6] + x[7] - x[8] + x[9] - x[10] + x[11] + x[12] + x[13] + x[14] - x[15] - x[16] + x[17] + x[18] - x[19] - x[20] - x[21] - x[22] + x[23] + x[24] + x[25] + x[26] - x[27] + x[28] - x[29] + x[30] - x[31] - x[32] - x[33] - x[34] + x[35] + x[36] - x[37] - x[38] + x[39];
out[39] = + x[0] - x[1] + x[2] + x[3] - x[4] - x[5] - x[6] - x[7] + x[8] - x[9] + x[10] - x[11] + x[12] + x[13] + x[14] + x[15] - x[16] - x[17] + x[18] + x[19] - x[20] + x[21] - x[22] - x[23] + x[24] + x[25] + x[26] + x[27] - x[28] + x[29] - x[30] + x[31] - x[32] - x[33] - x[34] - x[35] + x[36] + x[37] - x[38] - x[39];
#pragma unroll
for (int i = 0; i < 40; i++) { x[i] = out[i]; }
}
packages/ltx-kernels/csrc/ops/include/static_switch.h
+23
View File
@@ -0,0 +1,23 @@
// Inspired by https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
#pragma once
/// @param COND - a boolean expression to switch by
/// @param CONST_NAME - a name given for the constexpr bool variable.
/// @param ... - code to execute for true and false
///
/// Usage:
/// ```
/// BOOL_SWITCH(flag, BoolConst, [&] {
/// some_function<BoolConst>(...);
/// });
/// ```
#define BOOL_SWITCH(COND, CONST_NAME, ...) \
if (COND) { \
constexpr static bool CONST_NAME = true; \
__VA_ARGS__ \
} else { \
constexpr static bool CONST_NAME = false; \
__VA_ARGS__ \
} \
packages/ltx-kernels/csrc/ops/ops_api.cpp
+31
View File
@@ -0,0 +1,31 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <torch/python.h>
at::Tensor rms_norm_rope(at::Tensor &x, c10::optional<at::Tensor>& weights_, at::Tensor &cos_freqs, at::Tensor &sin_freqs, bool out_16bit);
at::Tensor fp6_pack(at::Tensor &x);
at::Tensor fp6_unpack(at::Tensor &x, int64_t original_n);
at::Tensor rms_norm_split_rope(
at::Tensor &x,
at::Tensor &sin_freqs,
at::Tensor &cos_freqs,
at::Tensor &weights,
bool out_fp8
);
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("rms_norm_rope", &rms_norm_rope,
"fused norm + rope + cvt");
m.def("fp6_pack", &fp6_pack,
"Pack 8-bit to 6-bit by dropping e_1 and e_2 bits");
m.def("fp6_unpack", &fp6_unpack,
"Unpack 6-bit to 8-bit (with e_1 and e_2 set to 0)");
m.def("rms_norm_split_rope", &rms_norm_split_rope,
"RMS norm + split RoPE with optional FP8 output");
}
packages/ltx-kernels/csrc/ops/rms_norm_rope.cpp
+114
View File
@@ -0,0 +1,114 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// Host entry point for the fused RMS-norm + RoPE kernel.
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <vector>
#include "fast_hadamard_transform.h"
#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
template<typename input_t, typename output_t, bool norm_affine>
void rms_norm_rope_cuda(NormRopeHadamardParamsBase &params, cudaStream_t stream);
void set_norm_rope_hadamard_params(NormRopeHadamardParamsBase &params,
// sizes
const size_t batch,
const size_t dim,
const size_t multiple,
// device pointers
const at::Tensor x,
const at::Tensor cos_freqs,
const at::Tensor sin_freqs,
const at::Tensor weights,
const at::Tensor out,
bool norm_affine,
float scale
) {
// Reset the parameters
memset(&params, 0, sizeof(params));
params.batch = batch;
params.dim = dim;
params.log_N = int(ceil(std::log2(dim / multiple)));
// Set the pointers and strides.
params.x_ptr = x.data_ptr();
params.out_ptr = out.data_ptr();
params.cos_freq_ptr = cos_freqs.data_ptr();
params.sin_freq_ptr = sin_freqs.data_ptr();
if (norm_affine){
params.weights_ptr = weights.data_ptr();
} else {
params.weights_ptr = nullptr;
}
// All stride are in elements, not bytes.
params.x_batch_stride = x.stride(0);
params.out_batch_stride = out.stride(0);
params.cos_freq_batch_stride = cos_freqs.stride(0);
params.sin_freq_batch_stride = sin_freqs.stride(0);
params.scale = scale;
}
at::Tensor rms_norm_rope(at::Tensor &x, c10::optional<at::Tensor>& weights_, at::Tensor &cos_freqs, at::Tensor &sin_freqs, bool out_16bit) {
auto input_type = x.scalar_type();
float scale = 1.0f; // :D
TORCH_CHECK(input_type == at::ScalarType::BFloat16);
TORCH_CHECK(x.is_cuda());
const auto shapes_og = x.sizes();
const int dim_og = x.size(-1);
x = x.reshape({-1, dim_og});
if (x.stride(-1) != 1) { x = x.contiguous(); }
const auto sizes = x.sizes();
const int batch_size = sizes[0];
cos_freqs = cos_freqs.reshape({-1, dim_og});
sin_freqs = sin_freqs.reshape({-1, dim_og});
at::Tensor weights;
bool norm_affine = false;
if(weights_.has_value()){
weights = weights_.value();
norm_affine = true;
}
CHECK_SHAPE(x, batch_size, dim_og);
TORCH_CHECK(x.stride(1) == 1);
if (dim_og % 8 != 0) {
x = torch::nn::functional::pad(x, torch::nn::functional::PadFuncOptions({0, 8 - dim_og % 8}));
}
const int dim = x.size(1);
at::Tensor out;
if (out_16bit){
out = torch::empty(x.sizes(), x.options().dtype(torch::kBFloat16));
} else {
out = torch::empty(x.sizes(), x.options().dtype(torch::kFloat8_e4m3fn));
}
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
NormRopeHadamardParamsBase params;
set_norm_rope_hadamard_params(params, batch_size, dim, 1, x, cos_freqs, sin_freqs, weights, out, norm_affine, scale);
TORCH_CHECK(dim % 8 == 0, "fast_hadamard_transform only supports hidden dimension divisible by 8 for now");
TORCH_CHECK(dim <= 32768, "fast_hadamard_transform only supports hidden dimension at most 32768 for now");
if (norm_affine){
if (out_16bit){
rms_norm_rope_cuda<at::BFloat16, at::BFloat16, true>(params, stream);
} else {
rms_norm_rope_cuda<at::BFloat16, at::Float8_e4m3fn, true>(params, stream);
}
} else {
if (out_16bit){
rms_norm_rope_cuda<at::BFloat16, at::BFloat16, false>(params, stream);
} else {
rms_norm_rope_cuda<at::BFloat16, at::Float8_e4m3fn, false>(params, stream);
}
}
return out.reshape(shapes_og);
}
packages/ltx-kernels/csrc/ops/rms_norm_rope_cuda.cu
+228
View File
@@ -0,0 +1,228 @@
/******************************************************************************
* Copyright (c) 2023, Tri Dao.
******************************************************************************/
// #pragma once
#include <c10/util/BFloat16.h>
#include <c10/util/Half.h>
#include <c10/util/Float8_e4m3fn.h>
#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
#include "fast_hadamard_transform.h"
#include "fast_hadamard_transform_common.h"
#include "fast_hadamard_transform_special.h"
#include "static_switch.h"
template<int kNThreads_, int kLogN_, typename input_t_, typename output_t_, bool norm_affine_>
struct norm_rope_kernel_traits {
using input_t = input_t_;
using output_t = output_t_;
static constexpr int kNThreads = kNThreads_;
static constexpr int kLogN = kLogN_;
static constexpr int N = 1 << kLogN;
static constexpr int kNBytes = sizeof(input_t);
static constexpr int OutkNBytes = sizeof(output_t);
static constexpr bool norm_affine = norm_affine_;
static_assert(kNBytes == 1 || kNBytes == 2 || kNBytes == 4);
static constexpr int kNElts = kNBytes == 4 ? 4 : kNBytes == 2 ? 8 : 8;
// It's possible that we need to do 2 rounds of exchange if input_t is 16 bits
// (since then we'd have 8 values of float, and each round we can exchange 4 floats).
static constexpr int kNExchangePerVec = sizeof(float) / sizeof(input_t);
using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
using vec_t_out = typename BytesToType<OutkNBytes * kNElts>::Type;
static constexpr int kNChunks = N / (kNElts * kNThreads);
// We don't want to use more than 32 KB of shared memory.
static constexpr int kSmemExchangeSize = std::min(N * 4, 32 * 1024);
static constexpr int kNExchangeRounds = N * 4 / kSmemExchangeSize;
static_assert(kNExchangeRounds * kSmemExchangeSize == N * 4);
static constexpr int kSmemSize = kSmemExchangeSize;
};
template<typename Ktraits>
__global__ __launch_bounds__(Ktraits::kNThreads)
void norm_rope_cvt_kernel(NormRopeHadamardParamsBase params) {
constexpr int kNThreads = Ktraits::kNThreads;
constexpr int kNElts = Ktraits::kNElts;
constexpr int kNExchangePerVec = Ktraits::kNExchangePerVec;
constexpr int kNExchangeRounds = Ktraits::kNExchangeRounds;
constexpr int kNChunks = Ktraits::kNChunks;
constexpr bool norm_affine = Ktraits::norm_affine;
using input_t = typename Ktraits::input_t;
using output_t = typename Ktraits::output_t;
using vec_t = typename Ktraits::vec_t;
using out_vec_t = typename Ktraits::vec_t_out;
using weights_t = typename Ktraits::input_t;
using freqs_t = typename Ktraits::input_t;
constexpr int kLogNElts = cilog2(Ktraits::kNElts);
static_assert(1 << kLogNElts == kNElts, "kNElts must be a power of 2");
constexpr int kWarpSize = std::min(kNThreads, 32);
constexpr int kLogWarpSize = cilog2(kWarpSize);
static_assert(1 << kLogWarpSize == kWarpSize, "Warp size must be a power of 2");
constexpr int kNWarps = kNThreads / kWarpSize;
constexpr int kLogNWarps = cilog2(kNWarps);
static_assert(1 << kLogNWarps == kNWarps, "kNWarps must be a power of 2");
constexpr int kLoadsPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNThreads);
static_assert(kLoadsPerExchange * sizeof(vec_t) * kNThreads == Ktraits::kSmemExchangeSize, "kSmemExchangeSize should be a power of 2");
static_assert(kNExchangeRounds * kLoadsPerExchange * sizeof(vec_t) == kNChunks * kNElts * sizeof(float));
constexpr int kChunksPerExchange = Ktraits::kSmemExchangeSize / (sizeof(vec_t) * kNExchangePerVec * kNThreads);
static_assert(kChunksPerExchange * sizeof(vec_t) * kNExchangePerVec * kNThreads == Ktraits::kSmemExchangeSize);
constexpr int kNExchanges = kNChunks / kChunksPerExchange;
static_assert(kNExchanges * kChunksPerExchange == kNChunks);
// Shared memory.
extern __shared__ char smem_[];
vec_t *smem_exchange = reinterpret_cast<vec_t *>(smem_);
const int batch_id = blockIdx.x;
const int warp_id = threadIdx.x / 32;
input_t *x = reinterpret_cast<input_t *>(params.x_ptr) + batch_id * params.x_batch_stride;
output_t *out = reinterpret_cast<output_t *>(params.out_ptr) + batch_id * params.out_batch_stride;
weights_t *weights = norm_affine ? reinterpret_cast<weights_t*>(params.weights_ptr) : nullptr;
float x_vals[kNChunks][kNElts];
float weights_vals[kNChunks][kNElts];
load_input<kNChunks, kNElts, input_t>(x, x_vals, params.dim);
//RMS Norm START
float thread_squared_sum = 0.0f;
#pragma unroll
for (size_t c = 0; c < kNChunks; c++)
{
#pragma unroll
for (size_t i = 0; i < kNElts; i++)
{
thread_squared_sum += x_vals[c][i] * x_vals[c][i];
}
}
SumOp<float> sum_op;
float warp_sum = Allreduce<32>::run(thread_squared_sum, sum_op);
float *smem_sum = reinterpret_cast<float*>(smem_);
if(threadIdx.x % 32 == 0){
smem_sum[warp_id] = warp_sum;
}
__syncthreads();
float norm = 0.0f;
#pragma unroll
for (size_t i = 0; i < kNWarps; i++)
{
norm += smem_sum[i];
}
norm *= 1.0f/params.dim;
norm = rsqrtf(norm);
if constexpr (norm_affine){
load_input<kNChunks, kNElts, weights_t>(weights, weights_vals, params.dim);
}
#pragma unroll
for (size_t c = 0; c < kNChunks; c++)
{
#pragma unroll
for (size_t i = 0; i < kNElts; i++)
{
if constexpr (norm_affine){
x_vals[c][i] *= (norm * weights_vals[c][i]);
} else {
x_vals[c][i] *= norm;
}
}
}
//RMS NORM END
//ROPE START
float sin_freqs_vals[kNChunks][kNElts];
float cos_freqs_vals[kNChunks][kNElts];
freqs_t *cos_freqs = reinterpret_cast<freqs_t*>(params.cos_freq_ptr) + batch_id * params.cos_freq_batch_stride;
freqs_t *sin_freqs = reinterpret_cast<freqs_t*>(params.sin_freq_ptr) + batch_id * params.sin_freq_batch_stride;
load_input<kNChunks, kNElts, freqs_t>(cos_freqs, cos_freqs_vals, params.dim);
load_input<kNChunks, kNElts, freqs_t>(sin_freqs, sin_freqs_vals, params.dim);
#pragma unroll
for (size_t c = 0; c < kNChunks; c++)
{
#pragma unroll
for (size_t i = 0; i < kNElts; i+=2)
{
float x_1 = x_vals[c][i];
float x_2 = x_vals[c][i+1];
x_vals[c][i] = -x_2*sin_freqs_vals[c][i] + x_1*cos_freqs_vals[c][i];
x_vals[c][i+1] = x_1*sin_freqs_vals[c][i+1] + x_2*cos_freqs_vals[c][i+1];
}
}
//ROPE END
store_output<kNChunks, kNElts, output_t, false>(out, x_vals, params.dim, params.scale);
}
template<int kNThreads, int kLogN, typename input_t, typename output_t, bool norm_affine>
void norm_rope_cvt_launch(NormRopeHadamardParamsBase &params, cudaStream_t stream) {
using Ktraits = norm_rope_kernel_traits<kNThreads, kLogN, input_t, output_t, norm_affine>;
constexpr int kSmemSize = Ktraits::kSmemSize;
dim3 grid(params.batch);
auto kernel = &norm_rope_cvt_kernel<Ktraits>;
if (kSmemSize >= 48 * 1024) {
C10_CUDA_CHECK(cudaFuncSetAttribute(
kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
}
kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
template<typename input_t, typename output_t, bool norm_affine>
void rms_norm_rope_cuda(NormRopeHadamardParamsBase &params, cudaStream_t stream) {
if (params.log_N == 3) {
norm_rope_cvt_launch<1, 3, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 4) {
norm_rope_cvt_launch<2, 4, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 5) {
norm_rope_cvt_launch<4, 5, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 6) {
norm_rope_cvt_launch<8, 6, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 7) {
norm_rope_cvt_launch<16, 7, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 8) {
norm_rope_cvt_launch<32, 8, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 9) {
norm_rope_cvt_launch<32, 9, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 10) {
norm_rope_cvt_launch<128, 10, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 11) {
norm_rope_cvt_launch<256, 11, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 12) {
norm_rope_cvt_launch<256, 12, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 13) {
norm_rope_cvt_launch<256, 13, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 14) {
norm_rope_cvt_launch<256, 14, input_t, output_t, norm_affine>(params, stream);
} else if (params.log_N == 15) {
norm_rope_cvt_launch<256, 15, input_t, output_t, norm_affine>(params, stream);
}
}
template void rms_norm_rope_cuda<at::BFloat16, at::Float8_e4m3fn, false>(NormRopeHadamardParamsBase &params, cudaStream_t stream);
template void rms_norm_rope_cuda<at::BFloat16, at::Float8_e4m3fn, true>(NormRopeHadamardParamsBase &params, cudaStream_t stream);
template void rms_norm_rope_cuda<at::BFloat16, at::BFloat16, false>(NormRopeHadamardParamsBase &params, cudaStream_t stream);
template void rms_norm_rope_cuda<at::BFloat16, at::BFloat16, true>(NormRopeHadamardParamsBase &params, cudaStream_t stream);
// template void fast_hadamard_transform_cuda<at::BFloat16, at::BFloat16>(HadamardParamsBase &params, cudaStream_t stream);
// template void fast_hadamard_transform_cuda<at::Float8_e4m3fn, at::BFloat16>(HadamardParamsBase &params, cudaStream_t stream);
// template void fast_hadamard_transform_cuda<at::Float8_e4m3fn, at::Float8_e4m3fn>(HadamardParamsBase &params, cudaStream_t stream);
packages/ltx-kernels/csrc/ops/rms_norm_split_rope.cpp
+111
View File
@@ -0,0 +1,111 @@
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <torch/python.h>
#include <vector>
// Forward declaration of CUDA kernel template
template<typename out_t>
void rms_norm_split_rope_cuda(
void* x,
void* sin_freqs,
void* cos_freqs,
void* weights,
int b,
int s,
int n,
int h,
long cos_sb, long cos_sn, long cos_ss,
long sin_sb, long sin_sn, long sin_ss,
void* out,
cudaStream_t stream
);
at::Tensor rms_norm_split_rope(
at::Tensor &x,
at::Tensor &sin_freqs,
at::Tensor &cos_freqs,
at::Tensor &weights,
bool out_fp8
) {
TORCH_CHECK(x.scalar_type() == at::ScalarType::BFloat16, "Input must be BFloat16");
TORCH_CHECK(sin_freqs.scalar_type() == at::ScalarType::BFloat16, "sin_freqs must be BFloat16");
TORCH_CHECK(cos_freqs.scalar_type() == at::ScalarType::BFloat16, "cos_freqs must be BFloat16");
TORCH_CHECK(x.is_cuda(), "Input must be on CUDA");
TORCH_CHECK(sin_freqs.is_cuda(), "sin_freqs must be on CUDA");
TORCH_CHECK(cos_freqs.is_cuda(), "cos_freqs must be on CUDA");
// Get dimensions
// x: [b, s, h]
// cos, sin: [b, n, s, d] where n*d = h/2
int b = x.size(0);
int s = x.size(1);
int h = x.size(2);
TORCH_CHECK(cos_freqs.dim() == 4, "cos_freqs must be 4D");
TORCH_CHECK(sin_freqs.dim() == 4, "sin_freqs must be 4D");
int n = cos_freqs.size(1);
int d = h / n;
// Require a contiguous innermost (d/2) dim for the vectorized int4 freq load,
// but keep the outer (b, n, s) strides: apply_split_rotary_emb hands us a
// swapaxes view (logical [b, n, s, d/2], physical [b, s, n, d/2]) whose inner
// stride is already 1, so this never copies it. The strides are forwarded to
// the kernel so the read is correct regardless of the physical layout.
if (x.stride(-1) != 1) { x = x.contiguous(); }
if (cos_freqs.stride(-1) != 1) { cos_freqs = cos_freqs.contiguous(); }
if (sin_freqs.stride(-1) != 1) { sin_freqs = sin_freqs.contiguous(); }
long cos_sb = cos_freqs.stride(0), cos_sn = cos_freqs.stride(1), cos_ss = cos_freqs.stride(2);
long sin_sb = sin_freqs.stride(0), sin_sn = sin_freqs.stride(1), sin_ss = sin_freqs.stride(2);
// Create output tensor
at::Tensor out;
if (out_fp8) {
out = torch::empty(x.sizes(), x.options().dtype(torch::kFloat8_e4m3fn));
} else {
out = torch::empty(x.sizes(), x.options().dtype(torch::kBFloat16));
}
// Setup CUDA
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
auto stream = at::cuda::getCurrentCUDAStream().stream();
// Launch kernel
if (out_fp8) {
rms_norm_split_rope_cuda<at::Float8_e4m3fn>(
x.data_ptr(),
sin_freqs.data_ptr(),
cos_freqs.data_ptr(),
weights.data_ptr(), // weights (optional, not used yet)
b,
s,
n,
h,
cos_sb, cos_sn, cos_ss,
sin_sb, sin_sn, sin_ss,
(void*)out.data_ptr(),
stream
);
} else {
rms_norm_split_rope_cuda<at::BFloat16>(
x.data_ptr(),
sin_freqs.data_ptr(),
cos_freqs.data_ptr(),
weights.data_ptr(), // weights (optional, not used yet)
b,
s,
n,
h,
cos_sb, cos_sn, cos_ss,
sin_sb, sin_sn, sin_ss,
(void*)out.data_ptr(),
stream
);
}
return out;
}
packages/ltx-kernels/csrc/ops/rms_norm_split_rope_cuda.cu
+202
View File
@@ -0,0 +1,202 @@
#include <c10/util/BFloat16.h>
#include <c10/util/Float8_e4m3fn.h>
#include <c10/cuda/CUDAException.h>
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp8.h>
// CUDA kernel template for RMS norm + split RoPE
// out_t can be at::Float8_e4m3fn or at::BFloat16
using bf16 = __nv_bfloat16;
using fp8 = __nv_fp8_e4m3;
__device__ __forceinline__ void _load_x(bf16* x, float x_vals[8], int h){
bf16 x_tmp[8];
*reinterpret_cast<int4*>(x_tmp) = reinterpret_cast<int4*>(x + blockIdx.x * h)[threadIdx.x];
#pragma unroll
for(int i = 0; i < 8; i++){
x_vals[i] = float(x_tmp[i]);
}
}
// Load 8 freq values for this thread from a table laid out logically as
// [b, n, s, d/2] (what apply_split_rotary_emb produces -- a swapaxes view whose
// physical layout is [b, s, n, d/2]). The strides (sb, sn, ss; inner d/2 stride
// is 1) are forwarded from the host so the read is correct for both that
// non-contiguous view and a genuinely contiguous [b, n, s, d/2] tensor. Both
// head-halves map to the same freq element (mirrors the eager cos.unsqueeze(-2)).
__device__ __forceinline__ void _load_freqs(
const bf16* freqs, float x_vals[8], int s, int d, long sb, long sn, long ss
){
bf16 x_tmp[8];
int threads_per_head = d / 8;
int head_idx = threadIdx.x / threads_per_head;
int lane = threadIdx.x % (threads_per_head / 2);
int b_idx = blockIdx.x / s;
int t_idx = blockIdx.x % s;
long off = b_idx * sb + head_idx * sn + t_idx * ss + (long)lane * 8;
*reinterpret_cast<int4*>(x_tmp) = *reinterpret_cast<const int4*>(freqs + off);
#pragma unroll
for(int i = 0; i < 8; i++){
x_vals[i] = float(x_tmp[i]);
}
}
template<typename out_t>
__global__ void _rms_norm_split_rope_kernel(bf16* x, bf16* sin_freqs, bf16* cos_freqs, void* out, bf16* weights, int b, int s, int n, int h,
long cos_sb, long cos_sn, long cos_ss, long sin_sb, long sin_sn, long sin_ss){
int token_idx = blockIdx.x;
int tid = threadIdx.x;
int lane_id = tid % 32;
// freqs have shape [b, s, h/2]
// each thread block calculate one row
// there are h/8 threads in thread block, each thread processes 8 values
// num_of_rows = b * s
// freqs have h/2 dim
// gridDim is (num_of_rows, 1, 1)
// calculate rms norm x_normed = x/x_norm * weights. x_norm is calculated across row, it means thread block wide sum reduction
extern __shared__ float smem[];
// Step 1: Load input values (8 per thread)
float x_vals[8];
_load_x(x, x_vals, h);
float sum_sq = 0.0f;
#pragma unroll
for(int i = 0; i < 8; i++){
sum_sq += x_vals[i] * x_vals[i];
}
// Warp-level reduction
#pragma unroll
for(int offset = 16; offset > 0; offset >>= 1){
sum_sq += __shfl_xor_sync(0xffffffff, sum_sq, offset);
}
if(tid % 32 == 0){
smem[tid / 32] = sum_sq;
}
__syncthreads();
// Final reduction across warps
if(tid == 0){
float total_sum = 0.0f;
int num_warps = blockDim.x / 32;
for(int i = 0; i < num_warps; i++){
total_sum += smem[i];
}
// RMS: sqrt(mean(x^2))
float rms = rsqrtf(total_sum / h + 1e-6f); // Add epsilon for numerical stability
smem[0] = rms;
}
__syncthreads();
float inv_rms = smem[0];
// Step 3: Apply RMS normalization (and weights if provided)
#pragma unroll
for(int i = 0; i < 8; i++){
x_vals[i] *= inv_rms;
// TODO: Apply weights if provided
if(weights != nullptr) x_vals[i] *= float(weights[tid * 8 + i]);
}
// Step 4: Calculate dimensions for split RoPE
// Conceptually: [b, s, h] -> [b, s, n, 2*d] -> [b, s, n, 2, d]
// where h = n * 2 * d
int d = h / n;
float x_other_vals[8];
int threads_per_head = d / 8;
int head_idx = tid / threads_per_head;
int idx_in_head = tid % threads_per_head;
bool is_first_half = idx_in_head < (threads_per_head / 2);
// LT-PATCH: full-warp mask. The original (1u << threads_per_head) - 1 only marks
// the first head's lanes active, so lanes belonging to heads beyond the first are
// not in the mask -> __shfl_xor_sync result is undefined and can corrupt RoPE. The
// XOR pattern keeps data within each power-of-two head group, so a full-warp mask
// is correct for every lane.
const unsigned mask = 0xffffffffu;
const int laneMask = threads_per_head / 2; // 4, 8, or 16
#pragma unroll
for (int i = 0; i < 8; i++) {
x_other_vals[i] = __shfl_xor_sync(mask, x_vals[i], laneMask);
}
float cos_vals[8], sin_vals[8];
_load_freqs(cos_freqs, cos_vals, s, d, cos_sb, cos_sn, cos_ss);
_load_freqs(sin_freqs, sin_vals, s, d, sin_sb, sin_sn, sin_ss);
#pragma unroll
for(int i = 0; i < 8; i++){
x_vals[i] = cos_vals[i]*x_vals[i];
}
float sign = is_first_half ? -1.0f : 1.0f;
for(int i = 0; i < 8; i++){
x_vals[i] += sign*sin_vals[i]*x_other_vals[i];
}
// Step 6: Convert and store output
if constexpr (std::is_same_v<out_t, at::Float8_e4m3fn>){
fp8 out_tmp[8];
#pragma unroll
for(int i = 0; i < 8; i++){
out_tmp[i] = fp8(x_vals[i]);
}
*reinterpret_cast<int64_t*>((fp8*)out + token_idx * h + tid * 8) = *reinterpret_cast<int64_t*>(out_tmp);
} else {
bf16 out_tmp[8];
#pragma unroll
for(int i = 0; i < 8; i++){
out_tmp[i] = __float2bfloat16(x_vals[i]);
}
*reinterpret_cast<int4*>((bf16*)out + token_idx * h + tid * 8) = *reinterpret_cast<int4*>(out_tmp);
}
}
template<typename out_t>
void rms_norm_split_rope_cuda(
void* x, // Input: [b, s, h]
void* sin_freqs, // Sin frequencies: [b, n, s, d]
void* cos_freqs, // Cos frequencies: [b, n, s, d]
void* weights,
int b, // Batch size
int s, // Sequence length
int n, // Number of heads (32)
int h, // Hidden dimension (2048, 4096, or 8192)
long cos_sb, long cos_sn, long cos_ss, // cos_freqs strides (b, n, s)
long sin_sb, long sin_sn, long sin_ss, // sin_freqs strides (b, n, s)
void* out, // Output: [b, s, h]
cudaStream_t stream
) {
int num_tokens = b * s;
int num_threads = h / 8; // Each thread processes 8 elements
int smem_size = (num_threads / 32 + 1) * sizeof(float); // Shared memory for reductions
dim3 grid(num_tokens);
dim3 block(num_threads);
_rms_norm_split_rope_kernel<out_t><<<grid, block, smem_size, stream>>>(
reinterpret_cast<bf16*>(x),
reinterpret_cast<bf16*>(sin_freqs),
reinterpret_cast<bf16*>(cos_freqs),
out,
reinterpret_cast<bf16*>(weights),
b, s, n, h,
cos_sb, cos_sn, cos_ss,
sin_sb, sin_sn, sin_ss
);
C10_CUDA_KERNEL_LAUNCH_CHECK();
}
// Explicit template instantiations
template void rms_norm_split_rope_cuda<at::BFloat16>(
void*, void*, void*, void*, int, int, int, int, long, long, long, long, long, long, void*, cudaStream_t
);
template void rms_norm_split_rope_cuda<at::Float8_e4m3fn>(
void*, void*, void*, void*, int, int, int, int, long, long, long, long, long, long, void*, cudaStream_t
);