DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-12 11:56:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: fotmat code

Browse Source
This commit is contained in:
leejet committed 2026-11-01 13:40:36 +08:00
1 parent 6edeac1b20
commit 5531c8878a
16 files changed
+591 -918
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cuda/sageattention/attn.cuh
+82 -176
View File
@@ -17,8 +17,8 @@
#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH #ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH
#define SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH #define SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH
#include "math.cuh"
#include "attn_utils.cuh" #include "attn_utils.cuh"
#include "math.cuh"
namespace sageattention { namespace sageattention {
@@ -36,17 +36,8 @@ namespace sageattention {
#define MMA_SV_N 16 #define MMA_SV_N 16
#define MMA_SV_K 16 #define MMA_SV_K 16
template<uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, template <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_mean = false>
typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_mean=false> __global__ void qk_int_sv_f16_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, half* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, DTypeOut* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
__global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__restrict__ K, half *__restrict__ V, DTypeOut *__restrict__ O, float *__restrict__ Lse,
float *__restrict__ Q_scale, float *__restrict__ K_scale, DTypeOut *__restrict__ V_mean,
const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups,
const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q,
const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k,
const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v,
const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o,
float sm_scale)
{
// Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes. // Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes.
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800 #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
// compile time check // compile time check
@@ -96,71 +87,53 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
uint32_t q_scale_idx, k_scale_idx; uint32_t q_scale_idx, k_scale_idx;
if constexpr (Q_GRAN == QuantGranularity::kPerBlock) if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
{
const uint32_t num_block_q = gridDim.x; const uint32_t num_block_q = gridDim.x;
q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx; q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx;
} } else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) {
else if constexpr (Q_GRAN == QuantGranularity::kPerWarp)
{
const uint32_t num_warp_block_q = gridDim.x * num_warps_q; const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>(); q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>();
} } else if constexpr (Q_GRAN == QuantGranularity::kPerThread) {
else if constexpr (Q_GRAN == QuantGranularity::kPerThread)
{
const uint32_t num_warp_block_q = gridDim.x * num_warps_q; const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4; q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4;
} }
if constexpr (K_GRAN == QuantGranularity::kPerBlock) if constexpr (K_GRAN == QuantGranularity::kPerBlock) {
{
const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K); const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K);
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k; k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k;
} } else if constexpr (K_GRAN == QuantGranularity::kPerWarp) {
else if constexpr (K_GRAN == QuantGranularity::kPerWarp)
{
const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>(); k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>();
} } else if constexpr (K_GRAN == QuantGranularity::kPerThread) {
else if constexpr (K_GRAN == QuantGranularity::kPerThread)
{
const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4; k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4;
} }
constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K) : (CTA_K / WARP_K) * 4; constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K)
: (CTA_K / WARP_K) * 4;
// initialize o, m, d // initialize o, m, d
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
RO[fq][fv][k] = 0.0f; RO[fq][fv][k] = 0.0f;
} }
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 4; k++) for (uint32_t k = 0; k < 4; k++) {
{
((int32_t*)RO[fq][fv])[k] = 0; ((int32_t*)RO[fq][fv])[k] = 0;
} }
} }
} }
} }
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 2; k++) for (uint32_t k = 0; k < 2; k++) {
{
m[fq][k] = -5000000.0f; m[fq][k] = -5000000.0f;
d[fq][k] = 1.0f; d[fq][k] = 1.0f;
} }
@@ -169,7 +142,8 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
constexpr uint32_t K_smem_idx_offset = CTA_Q; constexpr uint32_t K_smem_idx_offset = CTA_Q;
constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K; constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K;
constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B
: SwizzleMode::k128B;
smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem); smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE); smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
@@ -177,8 +151,10 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem); smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem);
constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4 : 8; constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4; : 8;
constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
: 4;
constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 32) ? 4 : 8; constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 32) ? 4 : 8;
constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 32) ? 8 : 4; constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 32) ? 8 : 4;
constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8; constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8;
@@ -227,11 +203,9 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
// for num_tiles_qk_inner = 1, we load all Qs in register // for num_tiles_qk_inner = 1, we load all Qs in register
uint32_t RQ[num_tiles_q][4]; uint32_t RQ[num_tiles_q][4];
if constexpr (num_tiles_qk_inner == 1) if constexpr (num_tiles_qk_inner == 1) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]); smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]);
Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma); Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma);
} }
@@ -258,20 +232,16 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
V_load_idx_lane_base += CTA_K; V_load_idx_lane_base += CTA_K;
#pragma unroll #pragma unroll
for (uint32_t iter = 1; iter < num_iterations - 1; iter++) for (uint32_t iter = 1; iter < num_iterations - 1; iter++) {
{
// ensure K is ready // ensure K is ready
cp_async::wait_group<1>(); cp_async::wait_group<1>();
__syncthreads(); __syncthreads();
// compute QK^T // compute QK^T
if constexpr (num_tiles_qk_inner == 1) if constexpr (num_tiles_qk_inner == 1) {
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_K, RS, RQ, K_smem_offset_mma); smem_K, RS, RQ, K_smem_offset_mma);
} } else {
else
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
} }
@@ -279,14 +249,11 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
float RS_f32[num_tiles_q][num_tiles_k][8]; float RS_f32[num_tiles_q][num_tiles_k][8];
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]); RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]);
} }
} }
@@ -295,25 +262,20 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
// do not apply causal mask and out of bound mask for these iterations // do not apply causal mask and out of bound mask for these iterations
K_idx_lane_base += CTA_K; K_idx_lane_base += CTA_K;
if constexpr (std::is_same<DTypeSVAccum, float>::value) if constexpr (std::is_same<DTypeSVAccum, float>::value) {
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, sm_scale);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, sm_scale);
} }
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
} }
uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16); RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
} }
@@ -331,13 +293,10 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
cp_async::wait_group<1>(); cp_async::wait_group<1>();
__syncthreads(); __syncthreads();
if constexpr (!use_inst_buffer) if constexpr (!use_inst_buffer) {
{
compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
smem_V, RS_f16, RO, d, V_smem_offset_mma); smem_V, RS_f16, RO, d, V_smem_offset_mma);
} } else {
else
{
compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
smem_V, RS_f16, RO, d, V_smem_offset_mma); smem_V, RS_f16, RO, d, V_smem_offset_mma);
} }
@@ -352,20 +311,16 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
} }
// second last iter, apply causal mask // second last iter, apply causal mask
if (num_iterations > 1) if (num_iterations > 1) {
{
// ensure K is ready // ensure K is ready
cp_async::wait_group<1>(); cp_async::wait_group<1>();
__syncthreads(); __syncthreads();
// compute QK^T // compute QK^T
if constexpr (num_tiles_qk_inner == 1) if constexpr (num_tiles_qk_inner == 1) {
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_K, RS, RQ, K_smem_offset_mma); smem_K, RS, RQ, K_smem_offset_mma);
} } else {
else
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
} }
@@ -373,45 +328,36 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
float RS_f32[num_tiles_q][num_tiles_k][8]; float RS_f32[num_tiles_q][num_tiles_k][8];
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale; RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
} }
} }
} }
if constexpr (mask_mode == MaskMode::kCausal) if constexpr (mask_mode == MaskMode::kCausal) {
{
apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32); apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
} }
// apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len); // apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
K_idx_lane_base += CTA_K; K_idx_lane_base += CTA_K;
if constexpr (std::is_same<DTypeSVAccum, float>::value) if constexpr (std::is_same<DTypeSVAccum, float>::value) {
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
} }
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
} }
uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16); RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
} }
@@ -429,13 +375,10 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
cp_async::wait_group<1>(); cp_async::wait_group<1>();
__syncthreads(); __syncthreads();
if constexpr (!use_inst_buffer) if constexpr (!use_inst_buffer) {
{
compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
smem_V, RS_f16, RO, d, V_smem_offset_mma); smem_V, RS_f16, RO, d, V_smem_offset_mma);
} } else {
else
{
compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
smem_V, RS_f16, RO, d, V_smem_offset_mma); smem_V, RS_f16, RO, d, V_smem_offset_mma);
} }
@@ -456,13 +399,10 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
__syncthreads(); __syncthreads();
// compute QK^T // compute QK^T
if constexpr (num_tiles_qk_inner == 1) if constexpr (num_tiles_qk_inner == 1) {
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_K, RS, RQ, K_smem_offset_mma); smem_K, RS, RQ, K_smem_offset_mma);
} } else {
else
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
} }
@@ -470,46 +410,37 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
float RS_f32[num_tiles_q][num_tiles_k][8]; float RS_f32[num_tiles_q][num_tiles_k][8];
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale; RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
} }
} }
} }
if constexpr (mask_mode == MaskMode::kCausal) if constexpr (mask_mode == MaskMode::kCausal) {
{
apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32); apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
} }
// check out of bound in the last iter // check out of bound in the last iter
apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len); apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
K_idx_lane_base += CTA_K; K_idx_lane_base += CTA_K;
if constexpr (std::is_same<DTypeSVAccum, float>::value) if constexpr (std::is_same<DTypeSVAccum, float>::value) {
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
} }
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
} }
uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16); RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
} }
@@ -517,19 +448,15 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
cp_async::wait_group<0>(); cp_async::wait_group<0>();
__syncthreads(); __syncthreads();
if constexpr (!use_inst_buffer) if constexpr (!use_inst_buffer) {
{
compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
smem_V, RS_f16, RO, d, V_smem_offset_mma); smem_V, RS_f16, RO, d, V_smem_offset_mma);
} } else {
else
{
compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
smem_V, RS_f16, RO, d, V_smem_offset_mma); smem_V, RS_f16, RO, d, V_smem_offset_mma);
} }
__syncthreads(); __syncthreads();
} }
// TODO: thread block sync mdo state for num_warps_k > 0 // TODO: thread block sync mdo state for num_warps_k > 0
@@ -541,14 +468,11 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
// { // {
// convert half to bfloat16 // convert half to bfloat16
if constexpr (std::is_same<DTypeSVAccum, half>::value && std::is_same<DTypeOut, nv_bfloat16>::value) if constexpr (std::is_same<DTypeSVAccum, half>::value && std::is_same<DTypeOut, nv_bfloat16>::value) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
((nv_bfloat162*)RO[fq][fv])[0] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[0])); ((nv_bfloat162*)RO[fq][fv])[0] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[0]));
((nv_bfloat162*)RO[fq][fv])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1])); ((nv_bfloat162*)RO[fq][fv])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1]));
((nv_bfloat162*)RO[fq][fv])[2] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[2])); ((nv_bfloat162*)RO[fq][fv])[2] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[2]));
@@ -558,18 +482,15 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
} }
// add v_mean // add v_mean
if constexpr (fuse_v_mean) if constexpr (fuse_v_mean) {
{
DTypeOut2 v_mean[2]; DTypeOut2 v_mean[2];
DTypeOut* V_mean_lane_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + lane_id % 4 * 2; DTypeOut* V_mean_lane_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + lane_id % 4 * 2;
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
v_mean[0] = *((DTypeOut2*)(V_mean_lane_ptr + fv * 16)); v_mean[0] = *((DTypeOut2*)(V_mean_lane_ptr + fv * 16));
v_mean[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16)); v_mean[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16));
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
((DTypeOut2*)RO[fq][fv])[0] = __hadd2(((DTypeOut2*)RO[fq][fv])[0], v_mean[0]); ((DTypeOut2*)RO[fq][fv])[0] = __hadd2(((DTypeOut2*)RO[fq][fv])[0], v_mean[0]);
((DTypeOut2*)RO[fq][fv])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]); ((DTypeOut2*)RO[fq][fv])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]);
((DTypeOut2*)RO[fq][fv])[2] = __hadd2(((DTypeOut2*)RO[fq][fv])[2], v_mean[1]); ((DTypeOut2*)RO[fq][fv])[2] = __hadd2(((DTypeOut2*)RO[fq][fv])[2], v_mean[1]);
@@ -581,26 +502,19 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
// save the result to shared memory // save the result to shared memory
uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4; uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O)); uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O));
if constexpr (std::is_same<DTypeSVAccum, float>::value) if constexpr (std::is_same<DTypeSVAccum, float>::value) {
{
// convert RO to half // convert RO to half
uint32_t RO_f16[4]; uint32_t RO_f16[4];
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 4; k++) for (uint32_t k = 0; k < 4; k++) {
{ if constexpr (std::is_same<DTypeOut, half>::value) {
if constexpr (std::is_same<DTypeOut, half>::value)
{
((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
} } else if constexpr (std::is_same<DTypeOut, nv_bfloat16>::value) {
else if constexpr (std::is_same<DTypeOut, nv_bfloat16>::value)
{
((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]); ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]);
} }
} }
@@ -611,9 +525,7 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
// ! permuted, make sure you know what you are doing // ! permuted, make sure you know what you are doing
((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = RO_f16[2]; ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = RO_f16[2];
((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3]; ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3];
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[0]; ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[0];
((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1]; ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1];
@@ -633,13 +545,10 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O; uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O;
#pragma unroll #pragma unroll
for (uint32_t i = 0; i < O_smem_iters_col; i++) for (uint32_t i = 0; i < O_smem_iters_col; i++) {
{
#pragma unroll #pragma unroll
for (uint32_t j = 0; j < O_smem_iters_row; j++) for (uint32_t j = 0; j < O_smem_iters_row; j++) {
{ if (O_load_idx_lane_base < qo_len) {
if (O_load_idx_lane_base < qo_len)
{
smem_O.store_128b(offset_O, O_lane_ptr); smem_O.store_128b(offset_O, O_lane_ptr);
} }
O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O); O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
@@ -651,15 +560,13 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O; O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O;
} }
if constexpr (return_lse) if constexpr (return_lse) {
{
uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>(); uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>();
float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx; float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx;
uint32_t fq = (lane_id % 4) / 2; uint32_t fq = (lane_id % 4) / 2;
uint32_t k = (lane_id % 4) % 2; uint32_t k = (lane_id % 4) % 2;
if (lse_idx < qo_len && (lane_id % 4) < 2 * num_tiles_q) if (lse_idx < qo_len && (lane_id % 4) < 2 * num_tiles_q) {
{
lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]); lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]);
} }
} }
@@ -670,7 +577,6 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
#endif #endif
} }
#undef PACK_SIZE_QK #undef PACK_SIZE_QK
#undef PACK_SIZE_V #undef PACK_SIZE_V
#undef PACK_SIZE_O #undef PACK_SIZE_O
src/core/cuda/sageattention/attn_fp8.cuh
+85 -181
View File
@@ -17,8 +17,8 @@
#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH #ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH
#define SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH #define SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH
#include "math.cuh"
#include "attn_utils.cuh" #include "attn_utils.cuh"
#include "math.cuh"
namespace sageattention { namespace sageattention {
@@ -36,17 +36,8 @@ namespace sageattention {
#define MMA_SV_N 16 #define MMA_SV_N 16
#define MMA_SV_K 32 #define MMA_SV_K 32
template<uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, template <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale = false, bool fuse_v_mean = false, bool use_pv_fp16_accu = false>
typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale=false, bool fuse_v_mean=false, bool use_pv_fp16_accu=false> __global__ void qk_int_sv_f8_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, int8_t* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, float* __restrict__ V_scale, float* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
__global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restrict__ K, int8_t *__restrict__ V, DTypeOut *__restrict__ O, float *__restrict__ Lse,
float *__restrict__ Q_scale, float *__restrict__ K_scale, float *__restrict__ V_scale, float *__restrict__ V_mean,
const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups,
const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q,
const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k,
const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v,
const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o,
float sm_scale)
{
// Keep unsupported device passes from parsing architecture-specific intrinsics. // Keep unsupported device passes from parsing architecture-specific intrinsics.
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 890 #if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 890
// compile time check // compile time check
@@ -94,71 +85,53 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
uint32_t q_scale_idx, k_scale_idx; uint32_t q_scale_idx, k_scale_idx;
if constexpr (Q_GRAN == QuantGranularity::kPerBlock) if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
{
const uint32_t num_block_q = gridDim.x; const uint32_t num_block_q = gridDim.x;
q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx; q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx;
} } else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) {
else if constexpr (Q_GRAN == QuantGranularity::kPerWarp)
{
const uint32_t num_warp_block_q = gridDim.x * num_warps_q; const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>(); q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>();
} } else if constexpr (Q_GRAN == QuantGranularity::kPerThread) {
else if constexpr (Q_GRAN == QuantGranularity::kPerThread)
{
const uint32_t num_warp_block_q = gridDim.x * num_warps_q; const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4; q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4;
} }
if constexpr (K_GRAN == QuantGranularity::kPerBlock) if constexpr (K_GRAN == QuantGranularity::kPerBlock) {
{
const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K); const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K);
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k; k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k;
} } else if constexpr (K_GRAN == QuantGranularity::kPerWarp) {
else if constexpr (K_GRAN == QuantGranularity::kPerWarp)
{
const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>(); k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>();
} } else if constexpr (K_GRAN == QuantGranularity::kPerThread) {
else if constexpr (K_GRAN == QuantGranularity::kPerThread)
{
const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K); const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4; k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4;
} }
constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K) : (CTA_K / WARP_K) * 4; constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K)
: (CTA_K / WARP_K) * 4;
// initialize o, m, d // initialize o, m, d
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
RO[fq][fv][k] = 0.0f; RO[fq][fv][k] = 0.0f;
} }
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 4; k++) for (uint32_t k = 0; k < 4; k++) {
{
((int32_t*)RO[fq][fv])[k] = 0; ((int32_t*)RO[fq][fv])[k] = 0;
} }
} }
} }
} }
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 2; k++) for (uint32_t k = 0; k < 2; k++) {
{
m[fq][k] = -5000000.0f; m[fq][k] = -5000000.0f;
d[fq][k] = 1.0f; d[fq][k] = 1.0f;
} }
@@ -167,7 +140,8 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
constexpr uint32_t K_smem_idx_offset = CTA_Q; constexpr uint32_t K_smem_idx_offset = CTA_Q;
constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K; constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K;
constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B; constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B
: SwizzleMode::k128B;
smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem); smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE); smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
// for fp16: 32 // for fp16: 32
@@ -176,8 +150,10 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B; constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem); smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem);
constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4 : 8; constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4; : 8;
constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
: 4;
// for fp16: 32 // for fp16: 32
constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8; constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8;
// for fp16: 32 // for fp16: 32
@@ -231,11 +207,9 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
// for num_tiles_qk_inner = 1, we load all Qs in register // for num_tiles_qk_inner = 1, we load all Qs in register
uint32_t RQ[num_tiles_q][4]; uint32_t RQ[num_tiles_q][4];
if constexpr (num_tiles_qk_inner == 1) if constexpr (num_tiles_qk_inner == 1) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]); smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]);
Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma); Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma);
} }
@@ -264,34 +238,27 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
K_load_idx_lane_base += CTA_K; K_load_idx_lane_base += CTA_K;
#pragma unroll #pragma unroll
for (uint32_t iter = 1; iter < num_iterations - 1; iter++) for (uint32_t iter = 1; iter < num_iterations - 1; iter++) {
{
// ensure K is ready // ensure K is ready
cp_async::wait_group<1>(); cp_async::wait_group<1>();
__syncthreads(); __syncthreads();
// compute QK^T // compute QK^T
if constexpr (num_tiles_qk_inner == 1) if constexpr (num_tiles_qk_inner == 1) {
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_K, RS, RQ, K_smem_offset_mma); smem_K, RS, RQ, K_smem_offset_mma);
} } else {
else
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
} }
float RS_f32[num_tiles_q][num_tiles_k][8]; float RS_f32[num_tiles_q][num_tiles_k][8];
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]); RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]);
} }
} }
@@ -299,25 +266,20 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
K_idx_lane_base += CTA_K; K_idx_lane_base += CTA_K;
if constexpr (std::is_same<DTypeSVAccum, float>::value) if constexpr (std::is_same<DTypeSVAccum, float>::value) {
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, sm_scale);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, sm_scale);
} }
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
} }
uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8); RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
{
accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d); accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
} }
@@ -338,18 +300,14 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
// for fp16: // for fp16:
// compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
// smem_V, RS_f16, RO, d, V_smem_offset_mma); // smem_V, RS_f16, RO, d, V_smem_offset_mma);
if constexpr (!use_inst_buffer) if constexpr (!use_inst_buffer) {
{
compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} } else {
else
{
if constexpr (!use_pv_fp16_accu) { if constexpr (!use_pv_fp16_accu) {
compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} } else {
else{
compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} }
@@ -366,20 +324,16 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
} }
// second last iter, apply causal mask // second last iter, apply causal mask
if (num_iterations > 1) if (num_iterations > 1) {
{
// ensure K is ready // ensure K is ready
cp_async::wait_group<1>(); cp_async::wait_group<1>();
__syncthreads(); __syncthreads();
// compute QK^T // compute QK^T
if constexpr (num_tiles_qk_inner == 1) if constexpr (num_tiles_qk_inner == 1) {
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_K, RS, RQ, K_smem_offset_mma); smem_K, RS, RQ, K_smem_offset_mma);
} } else {
else
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
} }
@@ -387,45 +341,36 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
float RS_f32[num_tiles_q][num_tiles_k][8]; float RS_f32[num_tiles_q][num_tiles_k][8];
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale; RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
} }
} }
} }
if constexpr (mask_mode == MaskMode::kCausal) if constexpr (mask_mode == MaskMode::kCausal) {
{
apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32); apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
} }
// apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len); // apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
K_idx_lane_base += CTA_K; K_idx_lane_base += CTA_K;
if constexpr (std::is_same<DTypeSVAccum, float>::value) if constexpr (std::is_same<DTypeSVAccum, float>::value) {
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
} }
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
} }
uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8); RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
{
accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d); accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
} }
@@ -446,18 +391,14 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
// for fp16: // for fp16:
// compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
// smem_V, RS_f16, RO, d, V_smem_offset_mma); // smem_V, RS_f16, RO, d, V_smem_offset_mma);
if constexpr (!use_inst_buffer) if constexpr (!use_inst_buffer) {
{
compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} } else {
else
{
if constexpr (!use_pv_fp16_accu) { if constexpr (!use_pv_fp16_accu) {
compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} } else {
else{
compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} }
@@ -480,13 +421,10 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
__syncthreads(); __syncthreads();
// compute QK^T // compute QK^T
if constexpr (num_tiles_qk_inner == 1) if constexpr (num_tiles_qk_inner == 1) {
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_K, RS, RQ, K_smem_offset_mma); smem_K, RS, RQ, K_smem_offset_mma);
} } else {
else
{
compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>( compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma); smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
} }
@@ -494,45 +432,36 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
float RS_f32[num_tiles_q][num_tiles_k][8]; float RS_f32[num_tiles_q][num_tiles_k][8];
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale; RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
} }
} }
} }
if constexpr (mask_mode == MaskMode::kCausal) if constexpr (mask_mode == MaskMode::kCausal) {
{
apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32); apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
} }
apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len); apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
K_idx_lane_base += CTA_K; K_idx_lane_base += CTA_K;
if constexpr (std::is_same<DTypeSVAccum, float>::value) if constexpr (std::is_same<DTypeSVAccum, float>::value) {
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale); update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
} }
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
{
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d); accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
} }
uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8); RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
{
accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d); accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
} }
@@ -543,25 +472,20 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
// for fp16: // for fp16:
// compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>( // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
// smem_V, RS_f16, RO, d, V_smem_offset_mma); // smem_V, RS_f16, RO, d, V_smem_offset_mma);
if constexpr (!use_inst_buffer) if constexpr (!use_inst_buffer) {
{
compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} } else {
else
{
if constexpr (!use_pv_fp16_accu) { if constexpr (!use_pv_fp16_accu) {
compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} } else {
else{
compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>( compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
smem_V, RS_f8, RO, d); smem_V, RS_f8, RO, d);
} }
} }
__syncthreads(); __syncthreads();
} }
// TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving. // TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving.
@@ -569,18 +493,15 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
normalize_d<num_tiles_q, num_tiles_v, ComputeUnit::kCudaCore>(RO, m, d); normalize_d<num_tiles_q, num_tiles_v, ComputeUnit::kCudaCore>(RO, m, d);
// ! here we just implement the case for fp32 acumulation // ! here we just implement the case for fp32 acumulation
if constexpr (fuse_v_scale) if constexpr (fuse_v_scale) {
{
float v_scale[4]; float v_scale[4];
float* V_scale_base_ptr = V_scale + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2; float* V_scale_base_ptr = V_scale + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2;
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
((float2*)v_scale)[0] = *((float2*)(V_scale_base_ptr + fv * 16)); ((float2*)v_scale)[0] = *((float2*)(V_scale_base_ptr + fv * 16));
((float2*)v_scale)[1] = *((float2*)(V_scale_base_ptr + fv * 16 + 8)); ((float2*)v_scale)[1] = *((float2*)(V_scale_base_ptr + fv * 16 + 8));
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
RO[fq][fv][0] *= v_scale[0]; RO[fq][fv][0] *= v_scale[0];
RO[fq][fv][1] *= v_scale[1]; RO[fq][fv][1] *= v_scale[1];
RO[fq][fv][2] *= v_scale[0]; RO[fq][fv][2] *= v_scale[0];
@@ -593,18 +514,15 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
} }
} }
if constexpr (fuse_v_mean) if constexpr (fuse_v_mean) {
{
float v_mean[4]; float v_mean[4];
float* V_mean_base_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2; float* V_mean_base_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2;
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
((float2*)v_mean)[0] = *((float2*)(V_mean_base_ptr + fv * 16)); ((float2*)v_mean)[0] = *((float2*)(V_mean_base_ptr + fv * 16));
((float2*)v_mean)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8)); ((float2*)v_mean)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8));
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
RO[fq][fv][0] += v_mean[0]; RO[fq][fv][0] += v_mean[0];
RO[fq][fv][1] += v_mean[1]; RO[fq][fv][1] += v_mean[1];
RO[fq][fv][2] += v_mean[0]; RO[fq][fv][2] += v_mean[0];
@@ -620,26 +538,19 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
// save the result to shared memory // save the result to shared memory
uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4; uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O)); uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O));
if constexpr (std::is_same<DTypeSVAccum, float>::value) if constexpr (std::is_same<DTypeSVAccum, float>::value) {
{
// convert RO to half // convert RO to half
uint32_t RO_f16[4]; uint32_t RO_f16[4];
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 4; k++) for (uint32_t k = 0; k < 4; k++) {
{ if constexpr (std::is_same<DTypeOut, half>::value) {
if constexpr (std::is_same<DTypeOut, half>::value)
{
((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
} } else {
else
{
((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]); ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]);
} }
} }
@@ -650,9 +561,7 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O) + 1); offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O) + 1);
((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[2]; ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[2];
((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3]; ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3];
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
// TODO: not implement // TODO: not implement
} }
} }
@@ -667,13 +576,10 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O; uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O;
#pragma unroll #pragma unroll
for (uint32_t i = 0; i < O_smem_iters_col; i++) for (uint32_t i = 0; i < O_smem_iters_col; i++) {
{
#pragma unroll #pragma unroll
for (uint32_t j = 0; j < O_smem_iters_row; j++) for (uint32_t j = 0; j < O_smem_iters_row; j++) {
{ if (O_load_idx_lane_base < qo_len) {
if (O_load_idx_lane_base < qo_len)
{
smem_O.store_128b(offset_O, O_lane_ptr); smem_O.store_128b(offset_O, O_lane_ptr);
} }
O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O); O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
@@ -685,16 +591,14 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O; O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O;
} }
if constexpr (return_lse) if constexpr (return_lse) {
{
// ! this only works for num_tiles_q = 2 // ! this only works for num_tiles_q = 2
uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>(); uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>();
float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx; float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx;
uint32_t fq = (lane_id % 4) / 2; uint32_t fq = (lane_id % 4) / 2;
uint32_t k = (lane_id % 4) % 2; uint32_t k = (lane_id % 4) % 2;
if (lse_idx < qo_len) if (lse_idx < qo_len) {
{
lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]); lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]);
} }
} }
src/core/cuda/sageattention/attn_utils.cuh
+156 -341
View File
@@ -22,8 +22,8 @@
#include "cp_async.cuh" #include "cp_async.cuh"
#include "math.cuh" #include "math.cuh"
#include "mma.cuh" #include "mma.cuh"
#include "permuted_smem.cuh"
#include "numeric_conversion.cuh" #include "numeric_conversion.cuh"
#include "permuted_smem.cuh"
namespace sageattention { namespace sageattention {
@@ -60,45 +60,35 @@ enum class ComputeUnit {
kCudaCore, kCudaCore,
}; };
__device__ __forceinline__ uint32_t get_warp_id() __device__ __forceinline__ uint32_t get_warp_id() {
{
return threadIdx.y; return threadIdx.y;
} }
__device__ __forceinline__ uint32_t get_lane_id() __device__ __forceinline__ uint32_t get_lane_id() {
{
return threadIdx.x; return threadIdx.x;
} }
template <uint32_t num_warps_q, uint32_t num_warps_k> template <uint32_t num_warps_q, uint32_t num_warps_k>
__device__ __forceinline__ uint32_t get_warp_idx_q() __device__ __forceinline__ uint32_t get_warp_idx_q() {
{
return get_warp_id() / num_warps_k; return get_warp_id() / num_warps_k;
} }
template <uint32_t num_warps_q, uint32_t num_warps_k> template <uint32_t num_warps_q, uint32_t num_warps_k>
__device__ __forceinline__ uint32_t get_warp_idx_k() __device__ __forceinline__ uint32_t get_warp_idx_k() {
{
return get_warp_id() % num_warps_k; return get_warp_id() % num_warps_k;
} }
template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T> __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem) {
__device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &smem_offset,
const uint32_t &gmem_stride,
const smem_t<swizzle_mode, stride> &smem)
{
static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value); static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16; constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16;
#pragma unroll #pragma unroll
for (uint32_t i = 0; i < smem_iters_col; i++) for (uint32_t i = 0; i < smem_iters_col; i++) {
{
#pragma unroll #pragma unroll
for (uint32_t j = 0; j < smem_iters_row; j++) for (uint32_t j = 0; j < smem_iters_row; j++) {
{
smem.load_128b_async(smem_offset, *lane_ptr); smem.load_128b_async(smem_offset, *lane_ptr);
*lane_ptr += (global_to_shared_line_lanes * pack_size); *lane_ptr += (global_to_shared_line_lanes * pack_size);
smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset); smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
@@ -112,23 +102,17 @@ __device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &sme
} }
// with predicate // with predicate
template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T> __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem, uint32_t base_idx, uint32_t max_len) {
__device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &smem_offset,
const uint32_t &gmem_stride,
const smem_t<swizzle_mode, stride> &smem, uint32_t base_idx, uint32_t max_len)
{
static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value); static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16; constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16;
#pragma unroll #pragma unroll
for (uint32_t i = 0; i < smem_iters_col; i++) for (uint32_t i = 0; i < smem_iters_col; i++) {
{
#pragma unroll #pragma unroll
for (uint32_t j = 0; j < smem_iters_row; j++) for (uint32_t j = 0; j < smem_iters_row; j++) {
{
// Masked probabilities still propagate NaNs from stale V lanes through 0 * NaN. // Masked probabilities still propagate NaNs from stale V lanes through 0 * NaN.
smem.load_128b_async<std::is_same<T, half>::value ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len); smem.load_128b_async<std::is_same<T, half>::value ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len);
*lane_ptr += (global_to_shared_line_lanes * pack_size); *lane_ptr += (global_to_shared_line_lanes * pack_size);
@@ -143,21 +127,15 @@ __device__ __forceinline__ void load_global_to_share(T **lane_ptr, uint32_t &sme
*lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
} }
template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA>
uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA> __device__ __forceinline__ void load_fp8_V_global_to_share(int8_t** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem) {
__device__ __forceinline__ void load_fp8_V_global_to_share(int8_t **lane_ptr, uint32_t &smem_offset,
const uint32_t &gmem_stride,
const smem_t<swizzle_mode, stride> &smem)
{
static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
constexpr uint32_t pack_size_fp8 = 16; constexpr uint32_t pack_size_fp8 = 16;
#pragma unroll #pragma unroll
for (uint32_t i = 0; i < smem_iters_col; i++) for (uint32_t i = 0; i < smem_iters_col; i++) {
{
#pragma unroll #pragma unroll
for (uint32_t j = 0; j < smem_iters_row; j++) for (uint32_t j = 0; j < smem_iters_row; j++) {
{
smem.load_128b_async(smem_offset, *lane_ptr); smem.load_128b_async(smem_offset, *lane_ptr);
*lane_ptr += (global_to_shared_line_lanes * pack_size_fp8); *lane_ptr += (global_to_shared_line_lanes * pack_size_fp8);
smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset); smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
@@ -172,11 +150,8 @@ __device__ __forceinline__ void load_fp8_V_global_to_share(int8_t **lane_ptr, ui
*lane_ptr -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride; *lane_ptr -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
} }
template <uint32_t num_warps_q, uint32_t num_warps_k, template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride>& smem_Q, const smem_t<swizzle_mode, stride>& smem_K, int32_t RS[][num_tiles_k][8], uint32_t& offset_Q, uint32_t& offset_K) {
SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
__device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride> &smem_Q, const smem_t<swizzle_mode, stride> &smem_K, int32_t RS[][num_tiles_k][8], uint32_t &offset_Q, uint32_t &offset_K)
{
static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4); static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4);
uint32_t RQ[num_tiles_q][4]; uint32_t RQ[num_tiles_q][4];
@@ -184,12 +159,10 @@ __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride
// the first iteration, mma mode is kInit // the first iteration, mma mode is kInit
#pragma unroll #pragma unroll
for (uint32_t iter = 0; iter < 1; iter++) for (uint32_t iter = 0; iter < 1; iter++) {
{
// load RQ // load RQ
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]); smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]);
offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q); offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q);
} }
@@ -197,22 +170,17 @@ __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride
offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter); offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter);
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
// load RK // load RK
smem_K.ldmatrix_m8n8x4(offset_K, RK); smem_K.ldmatrix_m8n8x4(offset_K, RK);
offset_K = smem_K.advance_offset_by_row<16>(offset_K); offset_K = smem_K.advance_offset_by_row<16>(offset_K);
// mma // mma
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (DTypeQK == DataType::kInt8) {
if constexpr (DTypeQK == DataType::kInt8)
{
mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK); mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
} } else if constexpr (DTypeQK == DataType::kInt4) {
else if constexpr (DTypeQK == DataType::kInt4)
{
mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK); mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
} }
} }
@@ -222,34 +190,27 @@ __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride
// following iteration, mma mode is kInplace // following iteration, mma mode is kInplace
#pragma unroll #pragma unroll
for (uint32_t iter = 1; iter < num_tiles_qk_inner; iter++) for (uint32_t iter = 1; iter < num_tiles_qk_inner; iter++) {
{
// load RQ // load RQ
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]); smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]);
offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q); offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q);
} }
offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter); offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter);
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
// load RK // load RK
smem_K.ldmatrix_m8n8x4(offset_K, RK); smem_K.ldmatrix_m8n8x4(offset_K, RK);
offset_K = smem_K.advance_offset_by_row<16>(offset_K); offset_K = smem_K.advance_offset_by_row<16>(offset_K);
// mma // mma
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (DTypeQK == DataType::kInt8) {
if constexpr (DTypeQK == DataType::kInt8)
{
mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK); mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
} } else if constexpr (DTypeQK == DataType::kInt4) {
else if constexpr (DTypeQK == DataType::kInt4)
{
mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK); mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
} }
} }
@@ -262,11 +223,8 @@ __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride
} }
// for case when num_tiles_qk_inner = 1 // for case when num_tiles_qk_inner = 1
template <uint32_t num_warps_q, uint32_t num_warps_k, template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride>& smem_K, int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], uint32_t offset_K) {
SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
__device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride> &smem_K, int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], uint32_t offset_K)
{
static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4); static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4);
static_assert(num_tiles_qk_inner == 1); static_assert(num_tiles_qk_inner == 1);
@@ -274,22 +232,17 @@ __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride
// mma mode is kInit // mma mode is kInit
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
// load RK // load RK
smem_K.ldmatrix_m8n8x4(offset_K, RK); smem_K.ldmatrix_m8n8x4(offset_K, RK);
offset_K = smem_K.advance_offset_by_row<16>(offset_K); offset_K = smem_K.advance_offset_by_row<16>(offset_K);
// mma // mma
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (DTypeQK == DataType::kInt8) {
if constexpr (DTypeQK == DataType::kInt8)
{
mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK); mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
} } else if constexpr (DTypeQK == DataType::kInt4) {
else if constexpr (DTypeQK == DataType::kInt4)
{
mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK); mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
} }
} }
@@ -297,27 +250,20 @@ __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride
} }
template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum> template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
__device__ __forceinline__ void apply_causal_mask(const uint32_t &Q_idx_lane_base, const uint32_t &K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8]) __device__ __forceinline__ void apply_causal_mask(const uint32_t& Q_idx_lane_base, const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8]) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
const uint32_t q_idx = Q_idx_lane_base + fq * 16 + 8 * ((k % 4) / 2); const uint32_t q_idx = Q_idx_lane_base + fq * 16 + 8 * ((k % 4) / 2);
const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2; const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2;
const bool out_of_boundary = (kv_idx > q_idx); const bool out_of_boundary = (kv_idx > q_idx);
if constexpr (std::is_same<DTypeQKAccum, float>::value) if constexpr (std::is_same<DTypeQKAccum, float>::value) {
{
RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
} } else if constexpr (std::is_same<DTypeQKAccum, half>::value) {
else if constexpr (std::is_same<DTypeQKAccum, half>::value)
{
RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
} }
} }
@@ -326,26 +272,19 @@ __device__ __forceinline__ void apply_causal_mask(const uint32_t &Q_idx_lane_bas
} }
template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum> template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
__device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t &K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8], const uint32_t &kv_len) __device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8], const uint32_t& kv_len) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{
const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2; const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2;
const bool out_of_boundary = (kv_idx >= kv_len); const bool out_of_boundary = (kv_idx >= kv_len);
if constexpr (std::is_same<DTypeQKAccum, float>::value) if constexpr (std::is_same<DTypeQKAccum, float>::value) {
{
RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
} } else if constexpr (std::is_same<DTypeQKAccum, half>::value) {
else if constexpr (std::is_same<DTypeQKAccum, half>::value)
{
RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
} }
} }
@@ -355,39 +294,29 @@ __device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t &K_idx_la
// for DTypeQKAccum float // for DTypeQKAccum float
template <uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, bool use_half_o_scale, bool exp_offset, bool fuse_scale = false, typename DTypeSVAccum> template <uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, bool use_half_o_scale, bool exp_offset, bool fuse_scale = false, typename DTypeSVAccum>
__device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], const float &sm_scale) __device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], const float& sm_scale) {
{
static_assert(std::is_same<DTypeSVAccum, half>::value || (!use_half_o_scale)); static_assert(std::is_same<DTypeSVAccum, half>::value || (!use_half_o_scale));
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 2; k++) for (uint32_t k = 0; k < 2; k++) {
{
// assign the smallest value possible // assign the smallest value possible
float m_prev = m[fq][k]; float m_prev = m[fq][k];
float m_temp = -5000000.0f; float m_temp = -5000000.0f;
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
float m_local = max(max(RS[fq][fk][k * 2 + 0], RS[fq][fk][k * 2 + 1]), float m_local = max(max(RS[fq][fk][k * 2 + 0], RS[fq][fk][k * 2 + 1]),
max(RS[fq][fk][k * 2 + 4], RS[fq][fk][k * 2 + 5])); max(RS[fq][fk][k * 2 + 4], RS[fq][fk][k * 2 + 5]));
m_temp = max(m_temp, m_local); m_temp = max(m_temp, m_local);
} }
if constexpr (!fuse_scale) if constexpr (!fuse_scale) {
{ if constexpr (exp_offset) {
if constexpr (exp_offset)
{
m_temp = fmaf(m_temp, sm_scale, -S_FP8_OFFSET); m_temp = fmaf(m_temp, sm_scale, -S_FP8_OFFSET);
} } else {
else
{
m_temp *= sm_scale; m_temp *= sm_scale;
} }
} } else if constexpr (exp_offset) {
else if constexpr (exp_offset)
{
m_temp += (-S_FP8_OFFSET); m_temp += (-S_FP8_OFFSET);
} }
@@ -403,31 +332,23 @@ __device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAc
d[fq][k] *= o_scale; d[fq][k] *= o_scale;
half2 o_scale2; half2 o_scale2;
if constexpr (use_half_o_scale) if constexpr (use_half_o_scale) {
{
o_scale2 = __floats2half2_rn(o_scale, o_scale); o_scale2 = __floats2half2_rn(o_scale, o_scale);
} }
// update RO // update RO
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
RO[fq][fv][k * 2 + 0] *= o_scale; RO[fq][fv][k * 2 + 0] *= o_scale;
RO[fq][fv][k * 2 + 1] *= o_scale; RO[fq][fv][k * 2 + 1] *= o_scale;
RO[fq][fv][k * 2 + 4] *= o_scale; RO[fq][fv][k * 2 + 4] *= o_scale;
RO[fq][fv][k * 2 + 5] *= o_scale; RO[fq][fv][k * 2 + 5] *= o_scale;
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value) if constexpr (use_half_o_scale) {
{
if constexpr (use_half_o_scale)
{
((half2*)RO[fq][fv])[k] = __hmul2(((half2*)RO[fq][fv])[k], o_scale2); ((half2*)RO[fq][fv])[k] = __hmul2(((half2*)RO[fq][fv])[k], o_scale2);
((half2*)RO[fq][fv])[k + 2] = __hmul2(((half2*)RO[fq][fv])[k + 2], o_scale2); ((half2*)RO[fq][fv])[k + 2] = __hmul2(((half2*)RO[fq][fv])[k + 2], o_scale2);
} } else {
else
{
RO[fq][fv][k * 2 + 0] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 0]) * o_scale); RO[fq][fv][k * 2 + 0] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 0]) * o_scale);
RO[fq][fv][k * 2 + 1] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 1]) * o_scale); RO[fq][fv][k * 2 + 1] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 1]) * o_scale);
RO[fq][fv][k * 2 + 4] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 4]) * o_scale); RO[fq][fv][k * 2 + 4] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 4]) * o_scale);
@@ -439,17 +360,13 @@ __device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAc
// raise RS to exponent // raise RS to exponent
float negative_m = -m[fq][k]; float negative_m = -m[fq][k];
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{ if constexpr (fuse_scale) {
if constexpr (fuse_scale)
{
RS[fq][fk][k * 2 + 0] = math::ptx_exp2(RS[fq][fk][k * 2 + 0] + negative_m); RS[fq][fk][k * 2 + 0] = math::ptx_exp2(RS[fq][fk][k * 2 + 0] + negative_m);
RS[fq][fk][k * 2 + 1] = math::ptx_exp2(RS[fq][fk][k * 2 + 1] + negative_m); RS[fq][fk][k * 2 + 1] = math::ptx_exp2(RS[fq][fk][k * 2 + 1] + negative_m);
RS[fq][fk][k * 2 + 4] = math::ptx_exp2(RS[fq][fk][k * 2 + 4] + negative_m); RS[fq][fk][k * 2 + 4] = math::ptx_exp2(RS[fq][fk][k * 2 + 4] + negative_m);
RS[fq][fk][k * 2 + 5] = math::ptx_exp2(RS[fq][fk][k * 2 + 5] + negative_m); RS[fq][fk][k * 2 + 5] = math::ptx_exp2(RS[fq][fk][k * 2 + 5] + negative_m);
} } else {
else
{
RS[fq][fk][k * 2 + 0] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 0], sm_scale, negative_m)); RS[fq][fk][k * 2 + 0] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 0], sm_scale, negative_m));
RS[fq][fk][k * 2 + 1] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 1], sm_scale, negative_m)); RS[fq][fk][k * 2 + 1] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 1], sm_scale, negative_m));
RS[fq][fk][k * 2 + 4] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 4], sm_scale, negative_m)); RS[fq][fk][k * 2 + 4] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 4], sm_scale, negative_m));
@@ -461,15 +378,12 @@ __device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAc
} }
template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename T> template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename T>
__device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_16[][num_tiles_k][4]) __device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_16[][num_tiles_k][4]) {
{
static_assert(sizeof(T) == 4); static_assert(sizeof(T) == 4);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
((half2*)RS_16[fq][fk])[0] = __float22half2_rn(((float2*)RS[fq][fk])[0]); ((half2*)RS_16[fq][fk])[0] = __float22half2_rn(((float2*)RS[fq][fk])[0]);
((half2*)RS_16[fq][fk])[1] = __float22half2_rn(((float2*)RS[fq][fk])[1]); ((half2*)RS_16[fq][fk])[1] = __float22half2_rn(((float2*)RS[fq][fk])[1]);
((half2*)RS_16[fq][fk])[2] = __float22half2_rn(((float2*)RS[fq][fk])[2]); ((half2*)RS_16[fq][fk])[2] = __float22half2_rn(((float2*)RS[fq][fk])[2]);
@@ -479,14 +393,11 @@ __device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_
} }
template <uint32_t num_tiles_q, uint32_t num_tiles_k> template <uint32_t num_tiles_q, uint32_t num_tiles_k>
__device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t RS_8[][num_tiles_k / 2][4]) __device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t RS_8[][num_tiles_k / 2][4]) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
{
floatx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 4); floatx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 4);
floatx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 2, RS[fq][fk * 2 + 0] + 6); floatx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 2, RS[fq][fk * 2 + 0] + 6);
floatx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 4); floatx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 4);
@@ -496,14 +407,11 @@ __device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t
} }
template <uint32_t num_tiles_q, uint32_t num_tiles_k> template <uint32_t num_tiles_q, uint32_t num_tiles_k>
__device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32_t RS_8[][num_tiles_k / 2][4]) __device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32_t RS_8[][num_tiles_k / 2][4]) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
{
halfx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2); halfx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2);
halfx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3); halfx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3);
halfx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2); halfx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2);
@@ -513,14 +421,11 @@ __device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32
} }
template <uint32_t num_tiles_q, uint32_t num_tiles_k> template <uint32_t num_tiles_q, uint32_t num_tiles_k>
__device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4], uint32_t RS[][num_tiles_k][4]) __device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4], uint32_t RS[][num_tiles_k][4]) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
{
e4m3x4_to_halfx4(RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2, RS_8[fq][fk]); e4m3x4_to_halfx4(RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2, RS_8[fq][fk]);
e4m3x4_to_halfx4(RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3, RS_8[fq][fk] + 1); e4m3x4_to_halfx4(RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3, RS_8[fq][fk] + 1);
e4m3x4_to_halfx4(RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2, RS_8[fq][fk] + 2); e4m3x4_to_halfx4(RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2, RS_8[fq][fk] + 2);
@@ -530,26 +435,20 @@ __device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4],
} }
template <uint32_t num_tiles_q, uint32_t num_tiles_k, ComputeUnit compute_unit = ComputeUnit::kTensorCore, typename T> template <uint32_t num_tiles_q, uint32_t num_tiles_k, ComputeUnit compute_unit = ComputeUnit::kTensorCore, typename T>
__device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit == ComputeUnit::kTensorCore)? 4 : 8], float d[][2]) __device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit == ComputeUnit::kTensorCore) ? 4 : 8], float d[][2]) {
{
// for compute unit cuda core, RS is float // for compute unit cuda core, RS is float
// for compute unit tensor core, RS is packed half // for compute unit tensor core, RS is packed half
static_assert((std::is_same<T, float>::value && compute_unit == ComputeUnit::kCudaCore) || static_assert((std::is_same<T, float>::value && compute_unit == ComputeUnit::kCudaCore) ||
(std::is_same<T, uint32_t>::value && compute_unit == ComputeUnit::kTensorCore)); (std::is_same<T, uint32_t>::value && compute_unit == ComputeUnit::kTensorCore));
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{ if constexpr (compute_unit == ComputeUnit::kTensorCore) {
if constexpr (compute_unit == ComputeUnit::kTensorCore)
{
// full accumulate with tensor core // full accumulate with tensor core
mma::rowsum_f16f16f32(d[fq], (uint32_t*)(RS[fq][fk])); mma::rowsum_f16f16f32(d[fq], (uint32_t*)(RS[fq][fk]));
} } else if constexpr (compute_unit == ComputeUnit::kCudaCore) {
else if constexpr (compute_unit == ComputeUnit::kCudaCore)
{
// partial accumulate with cuda core // partial accumulate with cuda core
d[fq][0] += RS[fq][fk][0] + RS[fq][fk][1] + RS[fq][fk][4] + RS[fq][fk][5]; d[fq][0] += RS[fq][fk][0] + RS[fq][fk][1] + RS[fq][fk][4] + RS[fq][fk][5];
d[fq][1] += RS[fq][fk][2] + RS[fq][fk][3] + RS[fq][fk][6] + RS[fq][fk][7]; d[fq][1] += RS[fq][fk][2] + RS[fq][fk][3] + RS[fq][fk][6] + RS[fq][fk][7];
@@ -559,45 +458,33 @@ __device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit =
} }
template <uint32_t num_tiles_q, uint32_t num_tiles_k> template <uint32_t num_tiles_q, uint32_t num_tiles_k>
__device__ __forceinline__ void accumulate_d_f8(uint32_t RS[][num_tiles_k / 2][4], float d[][2]) __device__ __forceinline__ void accumulate_d_f8(uint32_t RS[][num_tiles_k / 2][4], float d[][2]) {
{
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
{
mma::rowsum_f8f8f32(d[fq], RS[fq][fk]); mma::rowsum_f8f8f32(d[fq], RS[fq][fk]);
} }
} }
} }
template <uint32_t num_warps_q, uint32_t num_warps_k, template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, __device__ __forceinline__ void compute_fp16_sv(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f16[][num_tiles_k][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
__device__ __forceinline__ void compute_fp16_sv(const smem_t<swizzle_mode, stride> &smem_V, uint32_t RS_f16[][num_tiles_k][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2])
{
uint32_t smem_V_row_base = get_warp_idx_k<num_warps_q, num_warps_k>() * (num_tiles_k * 16) + get_lane_id() % 16; uint32_t smem_V_row_base = get_warp_idx_k<num_warps_q, num_warps_k>() * (num_tiles_k * 16) + get_lane_id() % 16;
uint32_t smem_V_col_base = get_lane_id() / 16; uint32_t smem_V_col_base = get_lane_id() / 16;
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fk * 16, smem_V_col_base + fv * 2); uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fk * 16, smem_V_col_base + fv * 2);
smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV); mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV); mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV);
} }
} }
@@ -605,32 +492,23 @@ __device__ __forceinline__ void compute_fp16_sv(const smem_t<swizzle_mode, strid
} }
} }
template <uint32_t num_warps_q, uint32_t num_warps_k, template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width = 4, typename T, typename DTypeSVAccum>
uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, __device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t<swizzle_mode, stride>& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) {
SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width=4, typename T, typename DTypeSVAccum>
__device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t<swizzle_mode, stride> &smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t &offset_V)
{
static_assert(sizeof(T) == 4); static_assert(sizeof(T) == 4);
// ! be sure you know what you are doing // ! be sure you know what you are doing
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++) for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
} }
} }
@@ -644,11 +522,8 @@ __device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t<swizzle_mo
offset_V -= (16 * num_tiles_k * stride); offset_V -= (16 * num_tiles_k * stride);
} }
template <uint32_t num_warps_q, uint32_t num_warps_k, template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width = 4, typename T, typename DTypeSVAccum>
uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<swizzle_mode, stride>& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) {
SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width=4, typename T, typename DTypeSVAccum>
__device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<swizzle_mode, stride> &smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t &offset_V)
{
static_assert(sizeof(T) == 4); static_assert(sizeof(T) == 4);
static_assert(std::is_same<DTypeSVAccum, float>::value); static_assert(std::is_same<DTypeSVAccum, float>::value);
@@ -656,17 +531,14 @@ __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<s
// ! be sure you know what you are doing // ! be sure you know what you are doing
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < 1; fk++) for (uint32_t fk = 0; fk < 1; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
{ {
mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInit>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInit>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
} }
@@ -678,17 +550,14 @@ __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<s
} }
#pragma unroll #pragma unroll
for (uint32_t fk = 1; fk < num_tiles_k; fk++) for (uint32_t fk = 1; fk < num_tiles_k; fk++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
smem_V.ldmatrix_m8n8x4_trans(offset_V, RV); smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
{ {
mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInplaceUpdate>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInplaceUpdate>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
} }
@@ -701,11 +570,9 @@ __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<s
// accumulate into RO // accumulate into RO
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
RO[fq][fv][0] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].x); RO[fq][fv][0] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].x);
RO[fq][fv][1] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].y); RO[fq][fv][1] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].y);
RO[fq][fv][2] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].x); RO[fq][fv][2] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].x);
@@ -723,19 +590,16 @@ __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<s
template <uint32_t num_tiles_q, uint32_t num_tiles_v, template <uint32_t num_tiles_q, uint32_t num_tiles_v,
ComputeUnit compute_unit = ComputeUnit::kTensorCore, // compute unit for accumulate_d ComputeUnit compute_unit = ComputeUnit::kTensorCore, // compute unit for accumulate_d
typename DTypeQKAccum, typename DTypeSVAccum> typename DTypeQKAccum,
__device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], DTypeQKAccum m[][2], float d[][2]) typename DTypeSVAccum>
{ __device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], DTypeQKAccum m[][2], float d[][2]) {
if constexpr (compute_unit == ComputeUnit::kCudaCore) if constexpr (compute_unit == ComputeUnit::kCudaCore) {
{
// accumulate_d performs partial accumulation with cuda core // accumulate_d performs partial accumulation with cuda core
// aggregate d // aggregate d
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 2; k++) for (uint32_t k = 0; k < 2; k++) {
{
d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x1); // sum 0 and 1, 2 and 3 d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x1); // sum 0 and 1, 2 and 3
d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x2); // sum 0 and 2, 1 and 3 d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x2); // sum 0 and 2, 1 and 3
} }
@@ -745,31 +609,23 @@ __device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], D
// divide O by d // divide O by d
float d_rcp[num_tiles_q][2]; float d_rcp[num_tiles_q][2];
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 2; k++) for (uint32_t k = 0; k < 2; k++) {
{
// TODO: check m to prevent nan // TODO: check m to prevent nan
d_rcp[fq][k] = math::ptx_rcp(d[fq][k]); d_rcp[fq][k] = math::ptx_rcp(d[fq][k]);
} }
} }
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
#pragma unroll #pragma unroll
for (uint32_t k = 0; k < 8; k++) for (uint32_t k = 0; k < 8; k++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2]; RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2];
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]); RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]);
} }
} }
@@ -777,34 +633,25 @@ __device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], D
} }
} }
template <uint32_t num_warps_q, uint32_t num_warps_k, template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, __device__ __forceinline__ void compute_fp8_sv(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
__device__ __forceinline__ void compute_fp8_sv(const smem_t<swizzle_mode, stride> &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2])
{
uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
// uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
{
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
smem_V.ldmatrix_m8n8x4(offset_V, RV); smem_V.ldmatrix_m8n8x4(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV); mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
// ! Not Implemented // ! Not Implemented
} }
} }
@@ -813,11 +660,8 @@ __device__ __forceinline__ void compute_fp8_sv(const smem_t<swizzle_mode, stride
} }
} }
template <uint32_t num_warps_q, uint32_t num_warps_k, template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
__device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mode, stride> &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2])
{
uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
// uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
@@ -825,25 +669,19 @@ __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mod
float RO_inst_buf[num_tiles_q][num_tiles_v][8]; float RO_inst_buf[num_tiles_q][num_tiles_v][8];
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < 1; fk++) for (uint32_t fk = 0; fk < 1; fk++) {
{
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
smem_V.ldmatrix_m8n8x4(offset_V, RV); smem_V.ldmatrix_m8n8x4(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
// ! Not Implemented // ! Not Implemented
} }
} }
@@ -852,25 +690,19 @@ __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mod
} }
#pragma unroll #pragma unroll
for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) {
{
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
smem_V.ldmatrix_m8n8x4(offset_V, RV); smem_V.ldmatrix_m8n8x4(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
// ! Not Implemented // ! Not Implemented
} }
} }
@@ -879,11 +711,9 @@ __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mod
} }
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
RO[fq][fv][0] += RO_inst_buf[fq][fv][0]; RO[fq][fv][0] += RO_inst_buf[fq][fv][0];
RO[fq][fv][1] += RO_inst_buf[fq][fv][1]; RO[fq][fv][1] += RO_inst_buf[fq][fv][1];
RO[fq][fv][2] += RO_inst_buf[fq][fv][2]; RO[fq][fv][2] += RO_inst_buf[fq][fv][2];
@@ -896,11 +726,8 @@ __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mod
} }
} }
template <uint32_t num_warps_q, uint32_t num_warps_k, template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, __device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
__device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t<swizzle_mode, stride> &smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2])
{
uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8; uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
// uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2; // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
uint32_t smem_V_col_base = (get_lane_id() / 8) % 2; uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
@@ -908,26 +735,20 @@ __device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t<s
uint32_t RO_int32[num_tiles_q][num_tiles_v][4]; uint32_t RO_int32[num_tiles_q][num_tiles_v][4];
#pragma unroll #pragma unroll
for (uint32_t fk = 0; fk < 1; fk++) for (uint32_t fk = 0; fk < 1; fk++) {
{
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
smem_V.ldmatrix_m8n8x4(offset_V, RV); smem_V.ldmatrix_m8n8x4(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
// mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); // mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInit>(RO_int32[fq][fv], RS_f8[fq][fk], RV); mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInit>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
// ! Not Implemented // ! Not Implemented
} }
} }
@@ -936,26 +757,20 @@ __device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t<s
} }
#pragma unroll #pragma unroll
for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) {
{
uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2); uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
#pragma unroll #pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++) for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
{
// load RV // load RV
uint32_t RV[4]; uint32_t RV[4];
// uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2); // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
smem_V.ldmatrix_m8n8x4(offset_V, RV); smem_V.ldmatrix_m8n8x4(offset_V, RV);
#pragma unroll #pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++) for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
{ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
// mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); // mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInplaceUpdate>(RO_int32[fq][fv], RS_f8[fq][fk], RV); mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInplaceUpdate>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
} } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
// ! Not Implemented // ! Not Implemented
} }
} }
src/core/cuda/sageattention/math.cuh
+6 -2
View File
@@ -35,9 +35,13 @@ namespace math {
constexpr float log2e = 1.44269504088896340736f; constexpr float log2e = 1.44269504088896340736f;
constexpr float log2e_recp = 1.0f / log2e; constexpr float log2e_recp = 1.0f / log2e;
__forceinline__ __device__ half2 uint32_as_half2(uint32_t x) { return *(half2*)&x; } __forceinline__ __device__ half2 uint32_as_half2(uint32_t x) {
return *(half2*)&x;
}
__forceinline__ __device__ uint32_t half2_as_uint32(half2 x) { return *(uint32_t*)&x; } __forceinline__ __device__ uint32_t half2_as_uint32(half2 x) {
return *(uint32_t*)&x;
}
/*! /*!
* \brief Wrapper of PTX ex2.approx instruction, which computes 2^x * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x
src/core/cuda/sageattention/mma.cuh
+33 -79
View File
@@ -133,12 +133,10 @@ __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr)
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_F16F16F32_M16N8K16_ENABLED #ifdef MMA_F16F16F32_M16N8K16_ENABLED
// ! only support half dtype now // ! only support half dtype now
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -148,9 +146,7 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, ui
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
"f"(C[2]), "f"(C[3])); "f"(C[2]), "f"(C[3]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -175,12 +171,10 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, ui
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_F16F16F32_M16N8K16_ENABLED #ifdef MMA_F16F16F32_M16N8K16_ENABLED
// ! only support half dtype now // ! only support half dtype now
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -199,9 +193,7 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, u
: "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
"f"(C[6]), "f"(C[7])); "f"(C[6]), "f"(C[7]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -235,11 +227,9 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, u
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_F16F16F16_M16N8K16_ENABLED #ifdef MMA_F16F16F16_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1}," "{%0, %1},"
@@ -248,9 +238,7 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C,
"{%8, %9};\n" "{%8, %9};\n"
: "=r"(C[0]), "=r"(C[1]) : "=r"(C[0]), "=r"(C[1])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1])); : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1}," "{%0, %1},"
@@ -274,11 +262,9 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C,
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_F16F16F16_M16N8K16_ENABLED #ifdef MMA_F16F16F16_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1}," "{%0, %1},"
@@ -295,9 +281,7 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C
"{%8, %9};\n" "{%8, %9};\n"
: "=r"(C[2]), "=r"(C[3]) : "=r"(C[2]), "=r"(C[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3])); : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 " "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1}," "{%0, %1},"
@@ -329,11 +313,9 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_S8S8S32_M16N8K32_ENABLED #ifdef MMA_S8S8S32_M16N8K32_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -343,9 +325,7 @@ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, ui
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
"r"(C[2]), "r"(C[3])); "r"(C[2]), "r"(C[3]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -370,11 +350,9 @@ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, ui
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_S8S8S32_M16N8K32_ENABLED #ifdef MMA_S8S8S32_M16N8K32_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -393,9 +371,7 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, u
: "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]), : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
"r"(C[6]), "r"(C[7])); "r"(C[6]), "r"(C[7]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 " "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -429,11 +405,9 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, u
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_S4S4S32_M16N8K64_ENABLED #ifdef MMA_S4S4S32_M16N8K64_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -443,9 +417,7 @@ __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, ui
: "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3]) : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]), : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
"r"(C[2]), "r"(C[3])); "r"(C[2]), "r"(C[3]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -470,11 +442,9 @@ __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, ui
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_S4S4S32_M16N8K64_ENABLED #ifdef MMA_S4S4S32_M16N8K64_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -493,9 +463,7 @@ __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, u
: "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7]) : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]), : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
"r"(C[6]), "r"(C[7])); "r"(C[6]), "r"(C[7]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 " "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -529,11 +497,9 @@ __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, u
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_F8F8F32_M16N8K16_ENABLED #ifdef MMA_F8F8F32_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -543,9 +509,7 @@ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint
: "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3]) : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]), : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
"f"(C[2]), "f"(C[3])); "f"(C[2]), "f"(C[3]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -570,12 +534,10 @@ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, uint32_t* B) {
uint32_t* B) {
// uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C); // uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
#ifdef MMA_F8F8F16_M16N8K16_ENABLED #ifdef MMA_F8F8F16_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
"{%0, %1}," "{%0, %1},"
@@ -593,9 +555,7 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_u
"{%8, %9};\n" "{%8, %9};\n"
: "=r"(C_uint32[2]), "=r"(C_uint32[3]) : "=r"(C_uint32[2]), "=r"(C_uint32[3])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3])); : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 " "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
"{%0, %1}," "{%0, %1},"
@@ -619,8 +579,6 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_u
#endif #endif
} }
/*! /*!
* \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
* multiplication, accumulated in fp32. * multiplication, accumulated in fp32.
@@ -630,11 +588,9 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_u
* \param B pointer to the fragment of matrix B * \param B pointer to the fragment of matrix B
*/ */
template <MMAMode mma_mode = MMAMode::kInplaceUpdate> template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
uint32_t* B) {
#ifdef MMA_F8F8F32_M16N8K16_ENABLED #ifdef MMA_F8F8F32_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate) if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
@@ -654,9 +610,7 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uin
: "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7]) : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
: "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]), : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
"f"(C[6]), "f"(C[7])); "f"(C[6]), "f"(C[7]));
} } else if constexpr (mma_mode == MMAMode::kInit) {
else if constexpr (mma_mode == MMAMode::kInit)
{
asm volatile( asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 " "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0, %1, %2, %3}," "{%0, %1, %2, %3},"
src/core/cuda/sageattention/numeric_conversion.cuh
+51 -58
View File
@@ -45,106 +45,99 @@ __device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, u
asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1)); asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1));
} }
__device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t *dest, float *source0, float *source1) __device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t* dest, float* source0, float* source1) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo;\n" \ ".reg .b16 lo;\n"
".reg .b16 hi;\n" \ ".reg .b16 hi;\n"
"cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" \ "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n"
"cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" \ "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n"
"mov.b32 %0, {lo, hi};\n" \ "mov.b32 %0, {lo, hi};\n"
"}" \ "}"
: "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1])); : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t *dest, float *source0, float *source1) __device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t* dest, float* source0, float* source1) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo;\n" \ ".reg .b16 lo;\n"
".reg .b16 hi;\n" \ ".reg .b16 hi;\n"
"cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n" \ "cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n"
"cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n" \ "cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n"
"mov.b32 %0, {lo, hi};\n" \ "mov.b32 %0, {lo, hi};\n"
"}" \ "}"
: "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1])); : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t *dest, uint32_t *source0, uint32_t *source1) __device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo;\n" \ ".reg .b16 lo;\n"
".reg .b16 hi;\n" \ ".reg .b16 hi;\n"
"cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n" \ "cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n"
"cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n" \ "cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n"
"mov.b32 %0, {lo, hi};\n" \ "mov.b32 %0, {lo, hi};\n"
"}" \ "}"
: "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t *dest, uint32_t *source0, uint32_t *source1) __device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo;\n" \ ".reg .b16 lo;\n"
".reg .b16 hi;\n" \ ".reg .b16 hi;\n"
"cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n" \ "cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n"
"cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n" \ "cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n"
"mov.b32 %0, {lo, hi};\n" \ "mov.b32 %0, {lo, hi};\n"
"}" \ "}"
: "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0])); : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source) __device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo, hi;\n" \ ".reg .b16 lo, hi;\n"
"mov.b32 {lo, hi}, %2;\n" \ "mov.b32 {lo, hi}, %2;\n"
"cvt.rn.f16x2.e4m3x2 %0, lo;\n" \ "cvt.rn.f16x2.e4m3x2 %0, lo;\n"
"cvt.rn.f16x2.e4m3x2 %1, hi;\n" \ "cvt.rn.f16x2.e4m3x2 %1, hi;\n"
"}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t *dest0, uint32_t *dest1, uint32_t *source) __device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
{
#ifdef FP8_CAST_ENABLED #ifdef FP8_CAST_ENABLED
asm volatile( \ asm volatile(
"{\n" \ "{\n"
".reg .b16 lo, hi;\n" \ ".reg .b16 lo, hi;\n"
"mov.b32 {lo, hi}, %2;\n" \ "mov.b32 {lo, hi}, %2;\n"
"cvt.rn.f16x2.e5m2x2 %0, lo;\n" \ "cvt.rn.f16x2.e5m2x2 %0, lo;\n"
"cvt.rn.f16x2.e5m2x2 %1, hi;\n" \ "cvt.rn.f16x2.e5m2x2 %1, hi;\n"
"}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0])); "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
#else #else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction"); RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif #endif
} }
__device__ __forceinline__ int8_t float_to_int8_rn(float x) __device__ __forceinline__ int8_t float_to_int8_rn(float x) {
{
uint32_t dst; uint32_t dst;
asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x)); asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
return reinterpret_cast<const int8_t&>(dst); return reinterpret_cast<const int8_t&>(dst);
src/core/cuda/sageattention/permuted_smem.cuh
-1
View File
@@ -23,7 +23,6 @@
#include <cuda_fp16.h> #include <cuda_fp16.h>
#include <cuda_runtime.h> #include <cuda_runtime.h>
#include "cp_async.cuh" #include "cp_async.cuh"
#include "mma.cuh" #include "mma.cuh"
src/core/ggml_extend.h
+1 -1
View File
@@ -4,9 +4,9 @@
#include <cstdint> #include <cstdint>
#include <vector> #include <vector>
#include "core/ggml_extensions.h"
#include "ggml-backend.h" #include "ggml-backend.h"
#include "ggml.h" #include "ggml.h"
#include "core/ggml_extensions.h"
#define EPS 1e-05f #define EPS 1e-05f
src/core/vulkan/shaders/modulate_convrot.comp
+23 -12
View File
@@ -5,13 +5,25 @@
layout(constant_id = 0) const uint K = 4096; layout(constant_id = 0) const uint K = 4096;
layout(constant_id = 1) const bool AFFINE = false; layout(constant_id = 1) const bool AFFINE = false;
layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer X { float data_x[]; }; layout(binding = 0) readonly buffer X {
layout(binding = 1) readonly buffer M { float data_m[]; }; float data_x[];
layout(binding = 2) readonly buffer H { float data_h[]; }; };
layout(binding = 3) writeonly buffer D { uint data_d[]; }; layout(binding = 1) readonly buffer M {
layout(binding = 4) writeonly buffer S { float data_s[]; }; float data_m[];
};
layout(binding = 2) readonly buffer H {
float data_h[];
};
layout(binding = 3) writeonly buffer D {
uint data_d[];
};
layout(binding = 4) writeonly buffer S {
float data_s[];
};
#ifdef RMS_AFFINE #ifdef RMS_AFFINE
layout(binding = 5) readonly buffer W { float data_w[]; }; layout(binding = 5) readonly buffer W {
float data_w[];
};
#endif #endif
layout(push_constant) uniform parameter { layout(push_constant) uniform parameter {
uint rows; uint rows;
@@ -21,7 +33,8 @@ layout(push_constant) uniform parameter {
float eps; float eps;
uint k; uint k;
uint w_offset; uint w_offset;
} p; }
p;
shared float values[K]; shared float values[K];
shared vec2 sums[512]; shared vec2 sums[512];
@@ -80,7 +93,8 @@ void main() {
sums[tid] = sum; sums[tid] = sum;
barrier(); barrier();
[[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) { [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) {
if (tid < stride) sums[tid] += sums[tid + stride]; if (tid < stride)
sums[tid] += sums[tid + stride];
barrier(); barrier();
} }
mean = sums[0].x / p.k; mean = sums[0].x / p.k;
@@ -139,10 +153,7 @@ void main() {
data_s[row] = scale; data_s[row] = scale;
} }
for (uint i = tid; i < K / 4; i += 512) { for (uint i = tid; i < K / 4; i += 512) {
data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) | (quantize(values[i * 4 + 1], inv_scale) << 8) | (quantize(values[i * 4 + 2], inv_scale) << 16) | (quantize(values[i * 4 + 3], inv_scale) << 24);
| (quantize(values[i * 4 + 1], inv_scale) << 8)
| (quantize(values[i * 4 + 2], inv_scale) << 16)
| (quantize(values[i * 4 + 3], inv_scale) << 24);
} }
barrier(); barrier();
} }
src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp
+21 -9
View File
@@ -5,21 +5,33 @@
layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in; layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in;
layout(push_constant) uniform parameter layout(push_constant) uniform parameter {
{
uint n; uint n;
uint rows; uint rows;
uint k; uint k;
uint has_bias; uint has_bias;
uint row_offset; uint row_offset;
} p; }
p;
layout(binding = 0, std430) readonly buffer W { int data_w[]; }; layout(binding = 0, std430) readonly buffer W {
layout(binding = 1, std430) readonly buffer A { int data_a[]; }; int data_w[];
layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; };
layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; layout(binding = 1, std430) readonly buffer A {
layout(binding = 4, std430) readonly buffer B { float data_b[]; }; int data_a[];
layout(binding = 5, std430) writeonly buffer D { float data_d[]; }; };
layout(binding = 2, std430) readonly buffer WS {
float data_ws[];
};
layout(binding = 3, std430) readonly buffer AS {
float data_as[];
};
layout(binding = 4, std430) readonly buffer B {
float data_b[];
};
layout(binding = 5, std430) writeonly buffer D {
float data_d[];
};
shared int weights[16 * 8]; shared int weights[16 * 8];
shared int activations[16 * 8]; shared int activations[16 * 8];
src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp
+21 -9
View File
@@ -17,21 +17,33 @@ const uint TILE_N = 4;
layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
layout(push_constant) uniform parameter layout(push_constant) uniform parameter {
{
uint n; uint n;
uint rows; uint rows;
uint k; uint k;
uint has_bias; uint has_bias;
uint row_offset; uint row_offset;
} p; }
p;
layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; }; layout(binding = 0, std430) readonly buffer W {
layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; }; int8_t data_w[];
layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; };
layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; layout(binding = 1, std430) readonly buffer A {
layout(binding = 4, std430) readonly buffer B { float data_b[]; }; int8_t data_a[];
layout(binding = 5, std430) writeonly buffer D { float data_d[]; }; };
layout(binding = 2, std430) readonly buffer WS {
float data_ws[];
};
layout(binding = 3, std430) readonly buffer AS {
float data_as[];
};
layout(binding = 4, std430) readonly buffer B {
float data_b[];
};
layout(binding = 5, std430) writeonly buffer D {
float data_d[];
};
shared int32_t accumulators[1024]; shared int32_t accumulators[1024];
src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp
+34 -15
View File
@@ -14,8 +14,7 @@ layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;
layout(constant_id = 3) const uint BM = 128; layout(constant_id = 3) const uint BM = 128;
layout(constant_id = 4) const uint BK = 64; layout(constant_id = 4) const uint BK = 64;
layout(push_constant) uniform parameter layout(push_constant) uniform parameter {
{
uint n; uint n;
uint rows; uint rows;
uint k; uint k;
@@ -23,25 +22,45 @@ layout(push_constant) uniform parameter
uint row_offset; uint row_offset;
float output_scale; float output_scale;
uint group_cols; uint group_cols;
} p; }
p;
layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; }; layout(binding = 0, std430) readonly buffer W {
layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; }; int8_t data_w[];
layout(binding = 2, std430) readonly buffer WS { float data_ws[]; }; };
layout(binding = 3, std430) readonly buffer AS { float data_as[]; }; layout(binding = 1, std430) readonly buffer A {
layout(binding = 4, std430) readonly buffer B { float data_b[]; }; int8_t data_a[];
};
layout(binding = 2, std430) readonly buffer WS {
float data_ws[];
};
layout(binding = 3, std430) readonly buffer AS {
float data_as[];
};
layout(binding = 4, std430) readonly buffer B {
float data_b[];
};
// Store FP32 bit patterns without changing the integer accumulator fragment layout. // Store FP32 bit patterns without changing the integer accumulator fragment layout.
layout(binding = 5, std430) writeonly buffer D { int32_t data_d[]; }; layout(binding = 5, std430) writeonly buffer D {
int32_t data_d[];
};
#ifdef SPLIT_SWIGLU #ifdef SPLIT_SWIGLU
layout(binding = 6, std430) readonly buffer UW { int8_t data_uw[]; }; layout(binding = 6, std430) readonly buffer UW {
layout(binding = 7, std430) readonly buffer US { float data_us[]; }; int8_t data_uw[];
};
layout(binding = 7, std430) readonly buffer US {
float data_us[];
};
#endif #endif
#ifdef RESIDUAL #ifdef RESIDUAL
layout(binding = 6, std430) readonly buffer G { float data_g[]; }; layout(binding = 6, std430) readonly buffer G {
layout(binding = 7, std430) readonly buffer R { float data_r[]; }; float data_g[];
};
layout(binding = 7, std430) readonly buffer R {
float data_r[];
};
#endif #endif
int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) {
const in uint row_base, const in uint col_base) {
const uint r = row_base + col; const uint r = row_base + col;
const uint c = col_base + row; const uint c = col_base + row;
if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) { if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) {
src/core/vulkan/shaders/quantize_i8_convrot.comp
+13 -10
View File
@@ -4,15 +4,21 @@
layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in;
layout(push_constant) uniform parameter layout(push_constant) uniform parameter {
{
uint k; uint k;
uint rows; uint rows;
} p; }
p;
layout(binding = 0, std430) readonly buffer A { float data_a[]; }; layout(binding = 0, std430) readonly buffer A {
layout(binding = 1, std430) writeonly buffer D { uint data_d[]; }; float data_a[];
layout(binding = 2, std430) writeonly buffer S { float data_s[]; }; };
layout(binding = 1, std430) writeonly buffer D {
uint data_d[];
};
layout(binding = 2, std430) writeonly buffer S {
float data_s[];
};
shared float values[256]; shared float values[256];
shared float maxima[64]; shared float maxima[64];
@@ -86,10 +92,7 @@ void quantize_row(const uint row) {
for (uint group = 0; group < groups; ++group) { for (uint group = 0; group < groups; ++group) {
transform_group(row, group); transform_group(row, group);
const uint index = tid * 4; const uint index = tid * 4;
const uint packed = quantize_value(values[index], inv_scale) const uint packed = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24);
| (quantize_value(values[index + 1], inv_scale) << 8)
| (quantize_value(values[index + 2], inv_scale) << 16)
| (quantize_value(values[index + 3], inv_scale) << 24);
data_d[(row * p.k + group * 256) / 4 + tid] = packed; data_d[(row * p.k + group * 256) / 4 + tid] = packed;
barrier(); barrier();
} }
src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp
+37 -14
View File
@@ -8,12 +8,24 @@
#extension GL_EXT_control_flow_attributes : require #extension GL_EXT_control_flow_attributes : require
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
layout(push_constant) uniform parameter { uint k; uint rows; } p; layout(push_constant) uniform parameter {
layout(binding = 0, std430) readonly buffer A { float data_a[]; }; uint k;
layout(binding = 1, std430) writeonly buffer D { int8_t data_d[]; }; uint rows;
layout(binding = 2, std430) writeonly buffer S { float data_s[]; }; }
p;
layout(binding = 0, std430) readonly buffer A {
float data_a[];
};
layout(binding = 1, std430) writeonly buffer D {
int8_t data_d[];
};
layout(binding = 2, std430) writeonly buffer S {
float data_s[];
};
#if defined(SWIGLU) || defined(SIGMOID_MUL) #if defined(SWIGLU) || defined(SIGMOID_MUL)
layout(binding = 3, std430) readonly buffer B { float data_b[]; }; layout(binding = 3, std430) readonly buffer B {
float data_b[];
};
#endif #endif
shared float maxima[8]; shared float maxima[8];
#ifdef SINGLE_PASS #ifdef SINGLE_PASS
@@ -62,7 +74,10 @@ void transform_group(uint row, uint group, out float v[8]) {
} }
[[unroll]] for (uint j = 0; j < 2; ++j) { [[unroll]] for (uint j = 0; j < 2; ++j) {
vec4 r = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]); vec4 r = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]);
v[j] = r.x; v[j + 2] = r.y; v[j + 4] = r.z; v[j + 6] = r.w; v[j] = r.x;
v[j + 2] = r.y;
v[j + 4] = r.z;
v[j + 6] = r.w;
} }
} }
@@ -72,8 +87,10 @@ void main() {
#ifdef SINGLE_PASS #ifdef SINGLE_PASS
float values[((K + 2047) / 2048) * 8]; float values[((K + 2047) / 2048) * 8];
[[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) { [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
if (gl_SubgroupID + g * 8 >= K / 256) continue; if (gl_SubgroupID + g * 8 >= K / 256)
float v[8]; transform_group(row, gl_SubgroupID + g * 8, v); continue;
float v[8];
transform_group(row, gl_SubgroupID + g * 8, v);
[[unroll]] for (uint j = 0; j < 8; ++j) { [[unroll]] for (uint j = 0; j < 8; ++j) {
values[g * 8 + j] = v[j]; values[g * 8 + j] = v[j];
local_max = max(local_max, abs(v[j])); local_max = max(local_max, abs(v[j]));
@@ -81,16 +98,20 @@ void main() {
} }
#else #else
for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) { for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
float v[8]; transform_group(row, group, v); float v[8];
[[unroll]] for (uint j = 0; j < 8; ++j) local_max = max(local_max, abs(v[j])); transform_group(row, group, v);
[[unroll]] for (uint j = 0; j < 8; ++j)
local_max = max(local_max, abs(v[j]));
} }
#endif #endif
local_max = subgroupMax(local_max); local_max = subgroupMax(local_max);
if (gl_SubgroupInvocationID == 0) maxima[gl_SubgroupID] = local_max; if (gl_SubgroupInvocationID == 0)
maxima[gl_SubgroupID] = local_max;
barrier(); barrier();
if (gl_LocalInvocationID.x == 0) { if (gl_LocalInvocationID.x == 0) {
float row_max = 0.0; float row_max = 0.0;
for (uint j = 0; j < 8; ++j) row_max = max(row_max, maxima[j]); for (uint j = 0; j < 8; ++j)
row_max = max(row_max, maxima[j]);
maxima[0] = row_max; maxima[0] = row_max;
data_s[row] = row_max / 127.0; data_s[row] = row_max / 127.0;
} }
@@ -99,14 +120,16 @@ void main() {
const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale; const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale;
#ifdef SINGLE_PASS #ifdef SINGLE_PASS
[[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) { [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
if (gl_SubgroupID + g * 8 >= K / 256) continue; if (gl_SubgroupID + g * 8 >= K / 256)
continue;
[[unroll]] for (uint j = 0; j < 8; ++j) { [[unroll]] for (uint j = 0; j < 8; ++j) {
data_d[row * p.k + (gl_SubgroupID + g * 8) * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(values[g * 8 + j] * inv_scale)), -127, 127)); data_d[row * p.k + (gl_SubgroupID + g * 8) * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(values[g * 8 + j] * inv_scale)), -127, 127));
} }
} }
#else #else
for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) { for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
float v[8]; transform_group(row, group, v); float v[8];
transform_group(row, group, v);
[[unroll]] for (uint j = 0; j < 8; ++j) { [[unroll]] for (uint j = 0; j < 8; ++j) {
data_d[row * p.k + group * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(v[j] * inv_scale)), -127, 127)); data_d[row * p.k + group * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(v[j] * inv_scale)), -127, 127));
} }
src/core/vulkan/shaders/rms_rope_apply.comp
+17 -6
View File
@@ -5,16 +5,26 @@
#extension GL_EXT_control_flow_attributes : require #extension GL_EXT_control_flow_attributes : require
layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer X { float data_x[]; }; layout(binding = 0) readonly buffer X {
layout(binding = 1) readonly buffer PE { float data_pe[]; }; float data_x[];
};
layout(binding = 1) readonly buffer PE {
float data_pe[];
};
#ifdef CAST_F16 #ifdef CAST_F16
#extension GL_EXT_shader_16bit_storage : require #extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require #extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
layout(binding = 2) writeonly buffer D { float16_t data_d[]; }; layout(binding = 2) writeonly buffer D {
float16_t data_d[];
};
#else #else
layout(binding = 2) writeonly buffer D { float data_d[]; }; layout(binding = 2) writeonly buffer D {
float data_d[];
};
#endif #endif
layout(binding = 3) readonly buffer W { float data_w[]; }; layout(binding = 3) readonly buffer W {
float data_w[];
};
layout(push_constant) uniform parameter { layout(push_constant) uniform parameter {
uint half_dim; uint half_dim;
@@ -31,7 +41,8 @@ layout(push_constant) uniform parameter {
float eps; float eps;
uint w_offset; uint w_offset;
float output_scale; float output_scale;
} pc; }
pc;
float shuffle_channel(vec4 values, uint channel) { float shuffle_channel(vec4 values, uint channel) {
const uint lane = channel & 31u; const uint lane = channel & 31u;
src/core/vulkan/shaders/rope_apply.comp
+11 -4
View File
@@ -2,9 +2,15 @@
layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in;
layout(binding = 0) readonly buffer X { float data_x[]; }; layout(binding = 0) readonly buffer X {
layout(binding = 1) readonly buffer PE { float data_pe[]; }; float data_x[];
layout(binding = 2) writeonly buffer D { float data_d[]; }; };
layout(binding = 1) readonly buffer PE {
float data_pe[];
};
layout(binding = 2) writeonly buffer D {
float data_d[];
};
layout(push_constant) uniform parameter { layout(push_constant) uniform parameter {
uint half_dim; uint half_dim;
@@ -18,7 +24,8 @@ layout(push_constant) uniform parameter {
uint pe_offset; uint pe_offset;
uint dst_offset; uint dst_offset;
uint interleaved; uint interleaved;
} pc; }
pc;
void main() { void main() {
const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x * gl_WorkGroupSize.x * gl_GlobalInvocationID.y; const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x * gl_WorkGroupSize.x * gl_GlobalInvocationID.y;