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/blockwise/kernels/geforce/gemm.cu
+287
View File
@@ -0,0 +1,287 @@
#include "cutlass/cutlass.h"
#include "cutlass/layout/layout.h"
#include <cute/tensor.hpp>
#include <c10/cuda/CUDAException.h>
#include <torch/extension.h>
#include <torch/python.h>
#include <cuda_runtime.h>
#include <iostream>
#include "kernel_traits.cuh"
#include "static_switch.h"
namespace sm89{
using namespace cute;
__device__ static void copy_1d(float* gmem_src, float* smem_dst)
{
uint32_t smem_int_ptr = cast_smem_ptr_to_uint((void*)smem_dst);
asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], %2;\n"
:: "r"(smem_int_ptr),
"l"(gmem_src),
"n"(sizeof(float)));
}
template <typename KernelTraits=gemm_traits<128, 256, 2, 4096, 2, 4, true, half_t, bfloat16_t>>
__global__ void gemm_fp8_kernel(float_e4m3_t* Aptr, float* sfa, float_e4m3_t* Bptr, float* sfb, float* bias_ptr, void* out, int M, int N, int K, int TMA_ALIGNED_M){
using output_t = typename KernelTraits::out_t;
using SmemLayoutA = typename KernelTraits::SmemLayoutA;
using SmemLayoutB = typename KernelTraits::SmemLayoutB;
using SmemLayoutC = typename KernelTraits::SmemLayoutC;
constexpr int BM = KernelTraits::BM;
constexpr int BN = KernelTraits::BN;
constexpr int BK = KernelTraits::BK;
constexpr int Ksfa = KernelTraits::KSF;
constexpr bool has_bias = KernelTraits::HasBias;
extern __shared__ float smem_[];
float *bias_shm = smem_;
float *sfa_shm = reinterpret_cast<float*>(bias_shm + cosize(typename KernelTraits::SmemLayoutBias{}));
output_t* C_shm = reinterpret_cast<output_t*>(sfa_shm + cosize(typename KernelTraits::SmemLayoutSFA{}));
float_e4m3_t* A_shm = reinterpret_cast<float_e4m3_t*>(sfa_shm + cosize(typename KernelTraits::SmemLayoutSFA{}));
float_e4m3_t* B_shm = reinterpret_cast<float_e4m3_t*>(A_shm + cosize(SmemLayoutA{}));
int idx = threadIdx.x;
int ix = blockIdx.x;
int iy = blockIdx.y;
// sfa += BM * iy;
sfb += KernelTraits::NUM_SFB_PER_STEP * ix * Ksfa;
output_t* Cptr = reinterpret_cast<output_t*>(out);
Tensor A = make_tensor(make_gmem_ptr(Aptr), make_shape(M, K), make_stride(K, Int<1>{}));
Tensor B = make_tensor(make_gmem_ptr(Bptr), make_shape(N, K), make_stride(K, Int<1>{}));
Tensor D = make_tensor(make_gmem_ptr(Cptr), make_shape(M, N), make_stride(N, Int<1>{}));
Tensor SFA = make_tensor(make_gmem_ptr(sfa), make_shape(M, Ksfa), make_stride(Int<1>{}, TMA_ALIGNED_M));
Tensor gA = local_tile(A, make_tile(Int<BM>{}, Int<BK>{}), make_coord(iy, _));
Tensor gB = local_tile(B, make_tile(Int<BN>{}, Int<BK>{}), make_coord(ix, _));
Tensor gD = local_tile(D, make_tile(Int<BM>{}, Int<BN>{}), make_coord(iy, ix));
Tensor gSFA = local_tile(SFA, make_tile(Int<BM>{}, Int<1>{}), make_coord(iy, _));
auto sBias = make_tensor(make_smem_ptr(bias_shm), typename KernelTraits::SmemLayoutBias{});
if constexpr (has_bias){
Tensor Bias = make_tensor(make_gmem_ptr(bias_ptr), make_shape(_1{}, N), make_stride(N, Int<1>{}));
Tensor gBias = local_tile(Bias, make_tile(Int<1>{}, Int<BN>{}), make_coord(_, ix));
typename KernelTraits::G2SBiasCopy g2s_bias_copy;
auto g2s_bias_thr_copy = g2s_bias_copy.get_slice(idx);
auto tCBiasgBias = g2s_bias_thr_copy.partition_S(gBias);
auto tCBiassBias = g2s_bias_thr_copy.partition_D(sBias);
if(idx < BN){
copy_1d((float*)&gBias(0) + idx, (float*)&sBias(0) + idx);
}
}
auto sSFA = make_tensor(make_smem_ptr(sfa_shm), typename KernelTraits::SmemLayoutSFA{});
auto sA = make_tensor(make_smem_ptr(A_shm), SmemLayoutA{});
auto sB = make_tensor(make_smem_ptr(B_shm), SmemLayoutB{});
typename KernelTraits::MMATile tiled_mma;
auto thr_mma = tiled_mma.get_slice(threadIdx.x);
auto tCrA = thr_mma.partition_fragment_A(gA(_, _, 0));
auto tCrB = thr_mma.partition_fragment_B(gB(_, _, 0));
auto tCrD = thr_mma.partition_fragment_C(gD);
clear(tCrD);
auto tCrD_fp32 = make_tensor_like<float>(tCrD);
clear(tCrD_fp32);
typename KernelTraits::G2STiledCopy g2s_tiled_copy;
auto g2s_thr_copy = g2s_tiled_copy.get_slice(idx);
auto tAgA_copy = g2s_thr_copy.partition_S(gA);
auto tAsA_copy = g2s_thr_copy.partition_D(sA);
auto tBgB_copy = g2s_thr_copy.partition_S(gB);
auto tBsB_copy = g2s_thr_copy.partition_D(sB);
auto s2r_tiled_copy_a = make_tiled_copy_A(typename KernelTraits::S2RCopyAtomA{}, tiled_mma);
auto s2r_thr_copy_a = s2r_tiled_copy_a.get_slice(idx);
auto tAsA = s2r_thr_copy_a.partition_S(sA);
auto tCrA_view = s2r_thr_copy_a.retile_D(tCrA);
auto s2r_tiled_copy_b = make_tiled_copy_B(typename KernelTraits::S2RCopyAtomB{}, tiled_mma);
auto s2r_thr_copy_b = s2r_tiled_copy_b.get_slice(idx);
auto tBsB = s2r_thr_copy_b.partition_S(sB);
auto tCrB_view = s2r_thr_copy_b.retile_D(tCrB);
auto cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA)));
auto tAcA = g2s_thr_copy.partition_S(cA);
int residual = M - iy*BM;
int itile_to_read = 0;
int ismem_read = 0;
int ismem_write = 0;
int ismem_read_sfa = 0;
constexpr int kStages = KernelTraits::KStages;
#pragma unroll
for(int istage=0; istage<kStages - 1; ++istage){
for (size_t m = 0; m < size<1>(tAsA_copy); m++)
{
for (size_t k = 0; k < size<2>(tAsA_copy); k++)
{
if(get<0>(tAcA(0, m, k)) < residual){
cute::copy(g2s_tiled_copy, tAgA_copy(_, m, k, istage), tAsA_copy(_, m, k, istage));
}
}
}
if(idx < KernelTraits::THREADS_SFA_COPY && (BM * iy + idx * KernelTraits::SFA_ELEMS_PER_COPY < M)) {
copy_1d((float*)&gSFA(0, 0, istage) + idx*KernelTraits::SFA_ELEMS_PER_COPY, (float*)&sSFA(0, istage) + idx*KernelTraits::SFA_ELEMS_PER_COPY);
}
cute::copy(g2s_tiled_copy, tBgB_copy(_, _, _, istage), tBsB_copy(_, _, _, istage));
cp_async_fence();
++itile_to_read;
++ismem_write;
}
cp_async_wait<kStages - 2>();
__syncthreads();
cute::copy(s2r_tiled_copy_a, tAsA(_, _, 0, ismem_read), tCrA_view(_, _, 0));
cute::copy(s2r_tiled_copy_b, tBsB(_, _, 0, ismem_read), tCrB_view(_, _, 0));
static constexpr int nk = size<2>(tCrA);
auto sfa_tv = typename KernelTraits::SFAThreadLayout{};
static constexpr int NTILES = KernelTraits::NTiles;
#pragma unroll
for(int itile = 0; itile < NTILES; itile++){
clear(tCrD);
#pragma unroll
for(int ik = 0; ik < nk; ik++){
int ik_next = (ik + 1) % nk;
if(ik == nk - 1) {
cp_async_wait<kStages - 2>();
__syncthreads();
ismem_read = (ismem_read + 1) % kStages;
}
cute::copy(s2r_tiled_copy_a, tAsA(_, _, ik_next, ismem_read), tCrA_view(_, _, ik_next));
cute::copy(s2r_tiled_copy_b, tBsB(_, _, ik_next, ismem_read), tCrB_view(_, _, ik_next));
if(ik == 0){
if(itile_to_read < NTILES){
for (size_t m = 0; m < size<1>(tAsA_copy); m++)
{
for (size_t k = 0; k < size<2>(tAsA_copy); k++)
{
if(get<0>(tAcA(0, m, k)) < residual){
cute::copy(g2s_tiled_copy, tAgA_copy(_, m, k, itile_to_read), tAsA_copy(_, m, k, ismem_write));
}
}
}
cute::copy(g2s_tiled_copy, tBgB_copy(_, _, _, itile_to_read), tBsB_copy(_, _, _, ismem_write));
if(idx < KernelTraits::THREADS_SFA_COPY && (BM * iy + idx * KernelTraits::SFA_ELEMS_PER_COPY < M)) {
copy_1d((float*)&gSFA(0, 0, itile_to_read) + idx * KernelTraits::SFA_ELEMS_PER_COPY, (float*)&sSFA(0, ismem_write) + idx*KernelTraits::SFA_ELEMS_PER_COPY);
}
++itile_to_read;
ismem_write = (ismem_write + 1) % kStages;
}
cp_async_fence();
}
cute::gemm(tiled_mma, tCrD, tCrA(_, _, ik), tCrB(_, _, ik), tCrD);
}
int sf_ind = itile / KernelTraits::TILES_PER_BLOCK;
float sfb_val = sfb[sf_ind];
#pragma unroll
for(int i = 0; i < size<1>(tCrD); i++){ // (MMA, MMA_M, MMA_N) = (4, 4, 4)
float sfa_val_1 = sSFA(sfa_tv(idx) + i * KernelTraits::MMA_WARP_M, ismem_read_sfa);
float sfa_val_2 = sSFA(sfa_tv(idx) + 8 + i * KernelTraits::MMA_WARP_M, ismem_read_sfa);
#pragma unroll
for(int j = 0; j < size<2>(tCrD); j++){
tCrD_fp32(0, i, j) += sfa_val_1 * sfb_val * float(tCrD(0, i, j));
tCrD_fp32(1, i, j) += sfa_val_1 * sfb_val * float(tCrD(1, i, j));
tCrD_fp32(2, i, j) += sfa_val_2 * sfb_val * float(tCrD(2, i, j));
tCrD_fp32(3, i, j) += sfa_val_2 * sfb_val * float(tCrD(3, i, j));
}
}
ismem_read_sfa = (ismem_read_sfa + 1) % kStages;
__syncthreads();
}
auto tCrBias = make_tensor<float>(Layout<Shape<_2, Int<size<2>(tCrD_fp32)>>>{});
auto bias_threads = typename KernelTraits::BiasThreadLayout{};
if constexpr (has_bias){
#pragma unroll
for(int i = 0; i<size<2>(tCrD_fp32); i++){
tCrBias(0, i) = sBias(bias_threads(idx) + i * KernelTraits::MMA_WARP_N);
tCrBias(1, i) = sBias(1 + bias_threads(idx) + i * KernelTraits::MMA_WARP_N);
}
#pragma unroll
for(int i = 0; i<size<1>(tCrD_fp32); i++){
#pragma unroll
for (int j = 0; j < size<2>(tCrD_fp32) ; j++)
{
tCrD_fp32(0, i, j) += tCrBias(0, j);
tCrD_fp32(1, i, j) += tCrBias(1, j);
tCrD_fp32(2, i, j) += tCrBias(0, j);
tCrD_fp32(3, i, j) += tCrBias(1, j);
}
}
}
auto sC = make_tensor(make_smem_ptr(C_shm), SmemLayoutC{});
auto r2s_tiled_copy_c = make_tiled_copy_C(typename KernelTraits::R2SCopyAtomC{}, tiled_mma);
auto r2s_thr_copy_c = r2s_tiled_copy_c.get_slice(idx);
auto tCrC_r2s = r2s_thr_copy_c.retile_S(tCrD_fp32);
auto tCsC_r2s = r2s_thr_copy_c.partition_D(sC);
typename KernelTraits::S2GCopyC s2g_tiled_copy_c;
auto s2g_thr_copy_c = s2g_tiled_copy_c.get_thread_slice(idx);
auto tCsC_s2g = s2g_thr_copy_c.partition_S(sC);
auto tCgC_s2g = s2g_thr_copy_c.partition_D(gD);
int pipe = size<2>(tCsC_r2s);
auto cC = make_identity_tensor(make_shape(size<0>(gD), size<1>(gD)));
auto tCcC = s2g_thr_copy_c.partition_D(cC);
for(int i = 0; i< size<1>(tCrC_r2s); i++){
for(int j = 0; j < size<2>(tCrC_r2s); j+=pipe){
for(int step = 0; step < pipe; ++step){
auto fragment = make_tensor_like<output_t>(tCrC_r2s(_, i, j + step));
cute::copy(tCrC_r2s(_, i, j + step), fragment);
cute::copy(r2s_tiled_copy_c, fragment, tCsC_r2s(_, 0, step));
}
__syncthreads();
if (get<0>(tCcC(0, i, j / pipe)) < residual){
cute::copy(s2g_tiled_copy_c, tCsC_s2g(_, 0, 0), tCgC_s2g(_, i, j / pipe));
}
__syncthreads();
}
}
}
template <bool has_bias, typename accum_type>
void fp8_kernel_launch(void* Aptr, void* sfa, void* Bptr, void* sfb, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream) {
int TMA_ALIGNED_M = ((M + sizeof(float) - 1) / sizeof(float)) * sizeof(float); // SIZEOF(float) = 4
BLOCK_K_SWITCH(K_, M_SWITCH(
using KernelTraits = gemm_traits<BM, BN, 3, K_, WARP_ROW, WARP_COL, has_bias, accum_type, bfloat16_t>;
auto kernel = &gemm_fp8_kernel<KernelTraits>;
int BX = (N + KernelTraits::BN - 1) / KernelTraits::BN;
int BY = (M + KernelTraits::BM - 1) / KernelTraits::BM;
dim3 block(KernelTraits::NUM_THREADS);
dim3 gridDim(BX, BY);
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, KernelTraits::SmemSize);
kernel<<<gridDim, KernelTraits::NUM_THREADS, KernelTraits::SmemSize, stream>>>((float_e4m3_t*)Aptr, (float*)sfa, (float_e4m3_t*)Bptr, (float*)sfb, (float*)bias_ptr, out, M, N, K, TMA_ALIGNED_M);
C10_CUDA_KERNEL_LAUNCH_CHECK();))
}
template<bool use_fast_accum>
void fp8_bias_gemm_cuda(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream){
using accum_type = std::conditional_t<use_fast_accum, half_t, float>;
fp8_kernel_launch<true, accum_type>(Aptr, SFA, Bptr, SFB, bias_ptr, out, M, N, K, stream);
}
// template<bool use_fast_accum>
// void fp8_gemm_cuda(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream){
// using accum_type = std::conditional_t<use_fast_accum, half_t, float>;
// BLOCK_K_SWITCH(num_acc_upcast_steps, fp8_kernel_launch<false, num_acc_upcast_steps, accum_type>(Aptr, SFA, Bptr, SFB, nullptr, out, M, N, K, stream);)
// }
// template void fp8_gemm_cuda<true>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream);
// template void fp8_gemm_cuda<false>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream);
template void fp8_bias_gemm_cuda<true>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream);
template void fp8_bias_gemm_cuda<false>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream);
}; // namespace sm89
packages/ltx-kernels/csrc/blockwise/kernels/geforce/kernel_traits.cuh
+121
View File
@@ -0,0 +1,121 @@
#pragma once
#include <cute/tensor.hpp>
#include <cutlass/cutlass.h>
#include <cutlass/layout/layout.h>
#include <cutlass/numeric_types.h>
#include "mma_sm89_fp16.hpp"
#include "mma_traits_sm89_fp16.hpp"
using namespace cute;
template<int BYTES> struct BytesToType {};
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<int BM_, int BN_, int KStages_, int K_, int WARP_ROW_=2, int WARP_COL_=2, bool HasBias_=false, typename accum_t_=cutlass::half_t, typename out_t_=cutlass::bfloat16_t>
struct gemm_traits {
static constexpr int BLOCK_SIZE = 128;
static constexpr int K = K_;
static constexpr int BM = BM_;
static constexpr int BN = BN_;
static constexpr int BK = 128;
static constexpr int TILES_PER_BLOCK = BLOCK_SIZE / BK;
static constexpr int NUM_SFB_PER_STEP = BN / BLOCK_SIZE;
static constexpr int NTiles = K / BK;
static constexpr int KSF = K / BLOCK_SIZE;
static constexpr int KStages = KStages_;
static constexpr int WARP_ROW = WARP_ROW_;
static constexpr int WARP_COL = WARP_COL_;
static constexpr int NUM_WARPS = WARP_ROW * WARP_COL;
static constexpr int NUM_THREADS = NUM_WARPS * 32;
static constexpr int MMA_WARP_M = WARP_ROW * 16;
static constexpr int MMA_WARP_N = WARP_COL * 8;
static constexpr int MMA_WARP_K = 32;
using accum_t = accum_t_;
using out_t = out_t_;
using SwizzleLayoutO = std::conditional_t<
std::is_same_v<out_t_, cutlass::bfloat16_t>,
Swizzle<3, 3, 3>,
Swizzle<2, 4, 3>
>;
using SwizzleLayoutAB = Swizzle<2, 4, 3>;
using MMA_Atom_SM89 = std::conditional_t<
std::is_same_v<accum_t, cutlass::half_t>,
MMA_Atom<SM89_16x8x32_F16E4M3E4M3F16_TN>,
MMA_Atom<SM89_16x8x32_F32E4M3E4M3F32_TN>
>;
static constexpr int INPUT_ELEMS_PER_COPY = sizeof(uint128_t) / sizeof(float_e4m3_t);
static constexpr int OUTPUT_ELEMS_PER_COPY = sizeof(uint128_t) / sizeof(out_t_);
static constexpr int THREADS_PER_ROW = BK / INPUT_ELEMS_PER_COPY;
using GMEMLayout = Layout< Shape <Int<NUM_THREADS / THREADS_PER_ROW>, Int<THREADS_PER_ROW>>, Stride<Int<THREADS_PER_ROW>, _1>>;
using G2SCopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, float_e4m3_t>;
using G2STiledCopy = decltype(
make_tiled_copy(
G2SCopyAtom{},
GMEMLayout{},
Layout<Shape<_1, Int<INPUT_ELEMS_PER_COPY>>>{}
)
);
using S2RCopyAtomA = Copy_Atom<SM75_U32x4_LDSM_N, float_e4m3_t>;
using S2RCopyAtomB = Copy_Atom<SM75_U32x2_LDSM_N, float_e4m3_t>;
using SmemLayoutAtom = decltype(composition(
Swizzle<2, 4, 3>{},
make_layout(make_shape(Int<8>{}, Int<BK>{}),
make_stride(Int<BK>{}, Int<1>{}))));
using SmemLayoutA = decltype(
tile_to_shape(SmemLayoutAtom{}, make_shape(Int<BM>{}, Int<BK>{}, Int<KStages>{}))
);
using SmemLayoutB = decltype(
tile_to_shape(SmemLayoutAtom{}, make_shape(Int<BN>{}, Int<BK>{}, Int<KStages>{}))
);
using MMATile = decltype(
make_tiled_mma(
MMA_Atom_SM89{},
Layout<Shape<Int<WARP_ROW>, Int<WARP_COL>, _1>>{},
Tile<Int<MMA_WARP_M>, Int<MMA_WARP_N>, Int<MMA_WARP_K>>{}
)
);
static constexpr int ELEMS_PER_TILE = MMA_WARP_M * MMA_WARP_N;
static constexpr int NUM_ELEMS_PER_WRITE = NUM_THREADS * sizeof(cute::uint128_t) / sizeof(out_t_);
static constexpr int OUT_PIPE = NUM_ELEMS_PER_WRITE / ELEMS_PER_TILE;
// using SmemLayoutC = Layout<Shape<Int<BM>, Int<BN>>, Stride<Int<BN>, Int<1>>>;
using SmemLayoutC = decltype(
make_layout(
make_shape(Int<MMA_WARP_M>{}, Int<MMA_WARP_N*OUT_PIPE>{}),
make_stride(Int<MMA_WARP_N*OUT_PIPE>{}, Int<1>{})
)
);
static constexpr int THREADS_PER_ROW_WRITE = MMA_WARP_N * OUT_PIPE / OUTPUT_ELEMS_PER_COPY;
using R2SCopyAtomC = Copy_Atom<UniversalCopy<typename BytesToType<2*sizeof(out_t)>::Type>, out_t>;
using S2GCopyAtomC = Copy_Atom<UniversalCopy<cute::uint128_t>, out_t>;
using S2GCopyC = decltype(make_tiled_copy(S2GCopyAtomC{},
make_layout(make_shape(Int<NUM_THREADS / THREADS_PER_ROW_WRITE>{}, Int<THREADS_PER_ROW_WRITE>{}),
make_stride(Int<THREADS_PER_ROW_WRITE>{}, Int<1>{})),
make_layout(make_shape(Int<1>{}, Int<OUTPUT_ELEMS_PER_COPY>{}))));
using G2SBiasCopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<float>, float>;
using G2SBiasCopy = decltype(make_tiled_copy(G2SBiasCopyAtom{}, make_layout(
make_shape(Int<1>{},Int<BN>{}), make_stride(Int<BN>{}, Int<1>{})),
make_layout(make_shape(Int<1>{},Int<1>{}), make_stride(Int<1>{}, Int<1>{}))));
using sfa_copy_vtype = float;
static constexpr int SFA_ELEMS_PER_COPY = sizeof(sfa_copy_vtype)/sizeof(float);
static constexpr int THREADS_SFA_COPY = BM * sizeof(float) / sizeof(sfa_copy_vtype);
// using G2SSFACopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, float>;
static constexpr bool HasBias = HasBias_;
using SmemLayoutBias = Layout<Shape<Int<1>, Int<BN>>, Stride<Int<BN>, Int<1>>>;
using SmemLayoutSFA = Layout<Shape<Int<BM>, Int<KStages>>, Stride<Int<1>, Int<BM>>>;
using BiasThreadLayout = Layout<Shape<Shape<_4, _8>, Shape<Int<WARP_ROW>, Int<WARP_COL>>>, Stride<Stride<_2, _0>, Stride<_0, _8>>>;
using SFAThreadLayout = Layout<Shape<Shape<_4, _8>, Shape<Int<WARP_ROW>, Int<WARP_COL>>>, Stride<Stride<_0, _1>, Stride<_16, _0>>>;
static constexpr int SmemSize = cute::max(cute::cosize(SmemLayoutA{})+cute::cosize(SmemLayoutB{}), cute::cosize(SmemLayoutC{})*sizeof(out_t)) + cute::cosize(SmemLayoutBias{}) * sizeof(float) + cute::cosize(SmemLayoutSFA{})*sizeof(float);
};
packages/ltx-kernels/csrc/blockwise/kernels/geforce/mma_sm89_fp16.hpp
+84
View File
@@ -0,0 +1,84 @@
#pragma once
#include <cute/config.hpp>
#include <cute/arch/mma.hpp>
////////////////////////////////////////////////////////////////////////////////
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 4)
# define CUTE_ARCH_MMA_F32_SM89_SUPPORTED
#endif
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8)
# define CUTE_ARCH_MMA_F16_SM89_SUPPORTED
#endif
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
# if defined(CUTE_ARCH_MMA_F32_SM89_SUPPORTED)
# define CUTE_ARCH_MMA_F32_SM89_ENABLED
# endif
# if defined(CUTE_ARCH_MMA_F16_SM89_SUPPORTED)
# define CUTE_ARCH_MMA_F16_SM89_ENABLED
# endif
#endif
namespace cute {
struct SM89_16x8x32_F32E4M3E4M3F32_TN
{
using DRegisters = float[4];
using ARegisters = uint32_t[4];
using BRegisters = uint32_t[2];
using CRegisters = float[4];
CUTE_HOST_DEVICE static void
fma(float & d0, float & d1, float & d2, float & d3,
uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3,
uint32_t const& b0, uint32_t const& b1,
float const& c0, float const& c1, float const& c2, float const& c3)
{
#if defined(CUTE_ARCH_MMA_F32_SM89_ENABLED)
asm(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
: "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3)
:
"r"(a0), "r"(a1), "r"(a2), "r"(a3),
"r"(b0), "r"(b1),
"f"(c0), "f"(c1), "f"(c2), "f"(c3)
);
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM89_16x8x32_F32E4M3E4M3F32_TN without CUTE_ARCH_MMA_F32_SM89_ENABLED");
#endif
}
};
// MMA 16x8x32 TN
struct SM89_16x8x32_F16E4M3E4M3F16_TN
{
using DRegisters = uint32_t[2];
using ARegisters = uint32_t[4];
using BRegisters = uint32_t[2];
using CRegisters = uint32_t[2];
CUTE_HOST_DEVICE static void
fma(uint32_t & d0, uint32_t & d1,
uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3,
uint32_t const& b0, uint32_t const& b1,
uint32_t const& c0, uint32_t const& c1)
{
#if defined(CUTE_ARCH_MMA_F16_SM89_ENABLED)
asm(
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
"{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
: "=r"(d0), "=r"(d1)
:
"r"(a0), "r"(a1), "r"(a2), "r"(a3),
"r"(b0), "r"(b1),
"r"(c0), "r"(c1)
);
#else
CUTE_INVALID_CONTROL_PATH("Attempting to use SM89_16x8x32_F32E4M3E4M3F32_TN without CUTE_ARCH_MMA_F16_SM89_ENABLED");
#endif
}
};
}
packages/ltx-kernels/csrc/blockwise/kernels/geforce/mma_traits_sm89_fp16.hpp
+50
View File
@@ -0,0 +1,50 @@
#pragma once
#include <cute/atom/mma_traits.hpp>
#include <cute/layout.hpp>
#include <cute/numeric/numeric_types.hpp>
#include "mma_sm89_fp16.hpp"
namespace cute
{
namespace {
// (T32,V4) -> (M16,N8)
using SM80_16x8_Row = Layout<Shape <Shape < _4,_8>,Shape < _2,_2>>,
Stride<Stride<_32,_1>,Stride<_16,_8>>>;
}
template <>
struct MMA_Traits<SM89_16x8x32_F32E4M3E4M3F32_TN> {
using ValTypeD = float;
using ValTypeA = float_e4m3_t;
using ValTypeB = float_e4m3_t;
using ValTypeC = float;
using Shape_MNK = Shape<_16,_8,_32>;
using ThrID = Layout<_32>;
using ALayout = Layout<Shape <Shape < _4,_8>,Shape < _4,_2, _2>>,
Stride<Stride<_64,_1>,Stride<_16,_8,_256>>>;
using BLayout = Layout<Shape <Shape < _4,_8>,Shape <_4, _2>>,
Stride<Stride<_32,_1>,Stride<_8,_128>>>;
using CLayout = SM80_16x8_Row;
};
template <>
struct MMA_Traits<SM89_16x8x32_F16E4M3E4M3F16_TN> {
using ValTypeD = half_t;
using ValTypeA = float_e4m3_t;
using ValTypeB = float_e4m3_t;
using ValTypeC = half_t;
using Shape_MNK = Shape<_16,_8,_32>;
using ThrID = Layout<_32>;
using ALayout = Layout<Shape <Shape < _4,_8>,Shape < _4,_2, _2>>,
Stride<Stride<_64,_1>,Stride<_16,_8,_256>>>;
using BLayout = Layout<Shape <Shape < _4,_8>,Shape <_4, _2>>,
Stride<Stride<_32,_1>,Stride<_8,_128>>>;
using CLayout = SM80_16x8_Row;
};
}
packages/ltx-kernels/csrc/blockwise/kernels/geforce/static_switch.h
+35
View File
@@ -0,0 +1,35 @@
#pragma once
#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__ \
}
//K/128
#define BLOCK_K_SWITCH(COSNT_NAME, ...) \
if (K == 2048) { \
constexpr static int COSNT_NAME = 2048; \
__VA_ARGS__ \
} \
else if (K == 4096) { \
constexpr static int COSNT_NAME = 4096; \
__VA_ARGS__ \
} else if (K == 8192) { \
constexpr static int COSNT_NAME = 8192; \
__VA_ARGS__ \
} else if (K == 16384) { \
constexpr static int COSNT_NAME = 16384; \
__VA_ARGS__ \
} else { \
TORCH_CHECK(false, "Unsupported K value: ", K); \
}
#define M_SWITCH(...) \
constexpr static int BM = 64; \
constexpr static int BN = 128; \
constexpr static int WARP_ROW = 2; \
constexpr static int WARP_COL = 4; \
__VA_ARGS__