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
17 files changed
+835 -1162
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cuda/sageattention/attn.cuh
+89 -183
View File
@@ -17,8 +17,8 @@
#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH
#define SD_CORE_CUDA_SAGEATTENTION_ATTN_CUH
#include "math.cuh"
#include "attn_utils.cuh"
#include "math.cuh"
namespace sageattention {
@@ -36,17 +36,8 @@ namespace sageattention {
#define MMA_SV_N 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,
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)
{
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>
__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.
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
// 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;
if constexpr (Q_GRAN == QuantGranularity::kPerBlock)
{
if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
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;
}
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;
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;
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);
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);
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);
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
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++)
{
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
#pragma unroll
for (uint32_t k = 0; k < 8; k++)
{
for (uint32_t k = 0; k < 8; k++) {
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
for (uint32_t k = 0; k < 4; k++)
{
for (uint32_t k = 0; k < 4; k++) {
((int32_t*)RO[fq][fv])[k] = 0;
}
}
}
}
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t k = 0; k < 2; k++)
{
for (uint32_t k = 0; k < 2; k++) {
m[fq][k] = -5000000.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 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_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
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;
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_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4;
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_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_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;
@@ -192,9 +168,9 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O);
constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O);
int8_t *Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
int8_t *K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
half *V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V) * stride_seq_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
half* V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V) * stride_seq_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V);
@@ -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
uint32_t RQ[num_tiles_q][4];
if constexpr (num_tiles_qk_inner == 1)
{
if constexpr (num_tiles_qk_inner == 1) {
#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]);
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;
#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
cp_async::wait_group<1>();
__syncthreads();
// 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>(
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>(
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];
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++)
{
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
#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]);
}
}
@@ -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
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);
}
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);
}
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
}
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);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
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>();
__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>(
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>(
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
if (num_iterations > 1)
{
if (num_iterations > 1) {
// ensure K is ready
cp_async::wait_group<1>();
__syncthreads();
// 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>(
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>(
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];
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++)
{
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
#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;
}
}
}
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_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
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);
}
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);
}
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
}
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);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
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>();
__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>(
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>(
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();
// 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>(
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>(
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];
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++)
{
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
#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;
}
}
}
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);
}
// 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);
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);
}
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);
}
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
}
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);
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
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>();
__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>(
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>(
smem_V, RS_f16, RO, d, V_smem_offset_mma);
}
__syncthreads();
}
// 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
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
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#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])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1]));
((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
if constexpr (fuse_v_mean)
{
if constexpr (fuse_v_mean) {
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
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[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16));
#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])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]);
((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
uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#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));
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
// convert RO to half
uint32_t RO_f16[4];
#pragma unroll
for (uint32_t k = 0; k < 4; k++)
{
if constexpr (std::is_same<DTypeOut, half>::value)
{
for (uint32_t k = 0; k < 4; k++) {
if constexpr (std::is_same<DTypeOut, half>::value) {
((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]);
}
}
@@ -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
((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];
}
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 + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1];
@@ -628,18 +540,15 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
__syncwarp();
// shared memory to global memory
DTypeOut *O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
DTypeOut* O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / global_to_shared_line_lanes_O, 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
for (uint32_t i = 0; i < O_smem_iters_col; i++)
{
for (uint32_t i = 0; i < O_smem_iters_col; i++) {
#pragma unroll
for (uint32_t j = 0; j < O_smem_iters_row; j++)
{
if (O_load_idx_lane_base < qo_len)
{
for (uint32_t j = 0; j < O_smem_iters_row; j++) {
if (O_load_idx_lane_base < qo_len) {
smem_O.store_128b(offset_O, O_lane_ptr);
}
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;
}
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>();
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 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]);
}
}
@@ -668,8 +575,7 @@ __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__rest
#else
RUNTIME_ASSERT("SageAttention requires SM80 or newer");
#endif
}
}
#undef PACK_SIZE_QK
#undef PACK_SIZE_V
src/core/cuda/sageattention/attn_fp8.cuh
+96 -192
View File
@@ -17,8 +17,8 @@
#ifndef SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH
#define SD_CORE_CUDA_SAGEATTENTION_ATTN_FP8_CUH
#include "math.cuh"
#include "attn_utils.cuh"
#include "math.cuh"
namespace sageattention {
@@ -36,17 +36,8 @@ namespace sageattention {
#define MMA_SV_N 16
#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,
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)
{
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>
__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.
#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 890
// 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;
if constexpr (Q_GRAN == QuantGranularity::kPerBlock)
{
if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
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;
}
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;
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;
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);
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);
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);
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
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t fv = 0; fv < num_tiles_v; fv++)
{
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
#pragma unroll
for (uint32_t k = 0; k < 8; k++)
{
for (uint32_t k = 0; k < 8; k++) {
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
for (uint32_t k = 0; k < 4; k++)
{
for (uint32_t k = 0; k < 4; k++) {
((int32_t*)RO[fq][fv])[k] = 0;
}
}
}
}
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t k = 0; k < 2; k++)
{
for (uint32_t k = 0; k < 2; k++) {
m[fq][k] = -5000000.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 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_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
// 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;
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_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4;
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_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
: 4;
// for fp16: 32
constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8;
// for fp16: 32
@@ -194,10 +170,10 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O);
constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O);
int8_t *Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
int8_t *K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
// for fp16: CTA_K / num_warps * warp_id * stride_seq_v + lane_id / global_to_shared_line_lanes_V * stride_seq_v
int8_t *V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + head_dim / num_warps * warp_id * stride_d_v + lane_id / global_to_shared_line_lanes_V * stride_d_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
int8_t* V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + head_dim / num_warps * warp_id * stride_d_v + lane_id / global_to_shared_line_lanes_V * stride_d_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V);
@@ -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
uint32_t RQ[num_tiles_q][4];
if constexpr (num_tiles_qk_inner == 1)
{
if constexpr (num_tiles_qk_inner == 1) {
#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]);
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;
#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
cp_async::wait_group<1>();
__syncthreads();
// 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>(
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>(
smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
}
float RS_f32[num_tiles_q][num_tiles_k][8];
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++)
{
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
#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]);
}
}
@@ -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;
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);
}
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);
}
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
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];
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);
}
@@ -338,18 +300,14 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
// 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>(
// 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>(
smem_V, RS_f8, RO, d);
}
else
{
if constexpr (!use_pv_fp16_accu){
} else {
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>(
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>(
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
if (num_iterations > 1)
{
if (num_iterations > 1) {
// ensure K is ready
cp_async::wait_group<1>();
__syncthreads();
// 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>(
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>(
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];
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++)
{
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
#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;
}
}
}
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_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
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);
}
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);
}
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
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];
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);
}
@@ -446,18 +391,14 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
// 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>(
// 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>(
smem_V, RS_f8, RO, d);
}
else
{
if constexpr (!use_pv_fp16_accu){
} else {
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>(
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>(
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();
// 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>(
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>(
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];
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#pragma unroll
for (uint32_t fk = 0; fk < num_tiles_k; fk++)
{
for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
#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;
}
}
}
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_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
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);
}
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);
}
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
{
if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
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];
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);
}
@@ -543,25 +472,20 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
// 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>(
// 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>(
smem_V, RS_f8, RO, d);
}
else
{
if constexpr (!use_pv_fp16_accu){
} else {
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>(
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>(
smem_V, RS_f8, RO, d);
}
}
__syncthreads();
}
// 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);
// ! 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_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
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)[1] = *((float2*)(V_scale_base_ptr + fv * 16 + 8));
#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][1] *= v_scale[1];
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_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
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)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8));
#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][1] += v_mean[1];
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
uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
#pragma unroll
for (uint32_t fq = 0; fq < num_tiles_q; fq++)
{
for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
#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));
if constexpr (std::is_same<DTypeSVAccum, float>::value)
{
if constexpr (std::is_same<DTypeSVAccum, float>::value) {
// convert RO to half
uint32_t RO_f16[4];
#pragma unroll
for (uint32_t k = 0; k < 4; k++)
{
if constexpr (std::is_same<DTypeOut, half>::value)
{
for (uint32_t k = 0; k < 4; k++) {
if constexpr (std::is_same<DTypeOut, half>::value) {
((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]);
}
}
@@ -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);
((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];
}
else if constexpr (std::is_same<DTypeSVAccum, half>::value)
{
} else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
// TODO: not implement
}
}
@@ -662,18 +571,15 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr
__syncwarp();
// shared memory to global memory
DTypeOut *O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
DTypeOut* O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / global_to_shared_line_lanes_O, 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
for (uint32_t i = 0; i < O_smem_iters_col; i++)
{
for (uint32_t i = 0; i < O_smem_iters_col; i++) {
#pragma unroll
for (uint32_t j = 0; j < O_smem_iters_row; j++)
{
if (O_load_idx_lane_base < qo_len)
{
for (uint32_t j = 0; j < O_smem_iters_row; j++) {
if (O_load_idx_lane_base < qo_len) {
smem_O.store_128b(offset_O, O_lane_ptr);
}
O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
@@ -685,23 +591,21 @@ __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;
}
if constexpr (return_lse)
{
if constexpr (return_lse) {
// ! 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>();
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 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]);
}
}
#else
RUNTIME_ASSERT("FP8 SageAttention requires SM89 or newer");
#endif
}
}
#undef PACK_SIZE_QK
#undef PACK_SIZE_V
src/core/cuda/sageattention/attn_utils.cuh
+226 -411
View File
File diff suppressed because it is too large. Load diff
src/core/cuda/sageattention/cp_async.cuh
+21 -21
View File
@@ -24,17 +24,17 @@
namespace sageattention {
namespace cp_async {
namespace cp_async {
enum class SharedMemFillMode {
enum class SharedMemFillMode {
kFillZero, // Fill zero to shared memory when predicate is false
kNoFill // Do not fill zero to shared memory when predicate is false
};
};
enum class PrefetchMode {
enum class PrefetchMode {
kNoPrefetch, // Do not fetch additional data from global memory to L2
kPrefetch // Fetch additional data from global memory to L2
};
};
#if (__CUDACC_VER_MAJOR__ >= 11)
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
@@ -42,28 +42,28 @@ enum class PrefetchMode {
#endif
#endif
/*!
/*!
* \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
* cp.async instructions to a group
*/
__device__ __forceinline__ void commit_group() {
__device__ __forceinline__ void commit_group() {
#ifdef CP_ASYNC_ENABLED
asm volatile("cp.async.commit_group;\n" ::);
#endif
}
}
/*!
/*!
* \brief Wrapper of PTX cp.async.wait_group instruction
* \tparam n Wait till most recent n groups are committed
*/
template <size_t n>
__device__ __forceinline__ void wait_group() {
template <size_t n>
__device__ __forceinline__ void wait_group() {
#ifdef CP_ASYNC_ENABLED
asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
#endif
}
}
/*!
/*!
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
* global memory to shared memory
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2
@@ -71,8 +71,8 @@ __device__ __forceinline__ void wait_group() {
* \param smem_ptr Pointer to shared memory
* \param gmem_ptr Pointer to global memory
*/
template <PrefetchMode prefetch_mode, typename T>
__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
template <PrefetchMode prefetch_mode, typename T>
__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
#ifdef CP_ASYNC_ENABLED
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
@@ -85,9 +85,9 @@ __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
#else
*((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
#endif
}
}
/*!
/*!
* \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
* global memory to shared memory with predicate.
* \tparam prefetch_mode Whether to fetch additional data from global memory to L2
@@ -98,8 +98,8 @@ __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
* \param predicate Predicate value
* \note fill zero is slower than not fill zero
*/
template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
#ifdef CP_ASYNC_ENABLED
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
@@ -139,9 +139,9 @@ __device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, b
}
}
#endif
}
}
} // namespace cp_async
} // namespace cp_async
} // namespace sageattention
src/core/cuda/sageattention/math.cuh
+45 -41
View File
@@ -26,136 +26,140 @@
namespace sageattention {
#ifndef USHORT_TYPE
#define USHORT_TYPE
typedef unsigned short ushort;
typedef unsigned short ushort;
#endif
namespace math {
namespace math {
// log2(e)
constexpr float log2e = 1.44269504088896340736f;
constexpr float log2e_recp = 1.0f / log2e;
// log2(e)
constexpr float log2e = 1.44269504088896340736f;
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
* \param x input
*/
__forceinline__ __device__ float ptx_exp2(float x) {
__forceinline__ __device__ float ptx_exp2(float x) {
float y;
asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
return y;
}
}
/*!
/*!
* \brief Wrapper of PTX lg2.approx instruction, which computes log2(x)
* \param x input
*/
__forceinline__ __device__ float ptx_log2(float x) {
__forceinline__ __device__ float ptx_log2(float x) {
float y;
asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
return y;
}
}
/*!
/*!
* \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x
* \param x input
*/
__forceinline__ __device__ half2 ptx_exp2(half2 x) {
__forceinline__ __device__ half2 ptx_exp2(half2 x) {
uint32_t y_u32;
uint32_t x_u32 = half2_as_uint32(x);
asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
return uint32_as_half2(y_u32);
}
}
/*!
/*!
* \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x
* \param x input
*/
__forceinline__ __device__ half ptx_exp2(half x) {
__forceinline__ __device__ half ptx_exp2(half x) {
ushort y_u16;
asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
return __ushort_as_half(y_u16);
}
}
/*!
/*!
* \brief Wrapper of PTX rcp.approx instruction, which computes 1/x
* \param x input
*/
__forceinline__ __device__ float ptx_rcp(float x) {
__forceinline__ __device__ float ptx_rcp(float x) {
float y;
asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
return y;
}
}
/*!
/*!
* \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle
* between threads in a warp.
* \param x The value in the source lane
* \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta]
*/
__forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) {
__forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) {
float y;
asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;"
: "=f"(y)
: "f"(x), "r"(lane_mask));
return y;
}
}
/*!
/*!
* \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly
* shuffle between threads in a warp.
* \param x The value in the source lane
* \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask]
*/
__forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) {
__forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) {
return __shfl_xor_sync(0xffffffff, x, lane_mask);
}
}
/*!
/*!
* \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x)
* \param x input
*/
__forceinline__ __device__ float rsqrt(float x) {
__forceinline__ __device__ float rsqrt(float x) {
float y;
asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
return y;
}
}
/*!
/*!
* \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x)
* \param x input
*/
__forceinline__ __device__ float tanh(float x) {
__forceinline__ __device__ float tanh(float x) {
float y;
asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x));
return y;
}
}
/*!
/*!
* \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x)
* \param x input
*/
__forceinline__ __device__ half2 tanh(half2 x) {
__forceinline__ __device__ half2 tanh(half2 x) {
uint32_t y_u32;
uint32_t x_u32 = half2_as_uint32(x);
asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
return uint32_as_half2(y_u32);
}
}
/*!
/*!
* \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x)
* \param x input
*/
__forceinline__ __device__ half tanh(half x) {
__forceinline__ __device__ half tanh(half x) {
ushort y_u16;
asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
return __ushort_as_half(y_u16);
}
}
} // namespace math
} // namespace math
} // namespace sageattention
src/core/cuda/sageattention/mma.cuh
+89 -135
View File
@@ -26,7 +26,7 @@
namespace sageattention {
namespace mma{
namespace mma {
#if (__CUDACC_VER_MAJOR__ >= 11)
#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
@@ -62,20 +62,20 @@ namespace mma{
#define RUNTIME_ASSERT(x) assert(0 && x)
#endif
enum class MMAMode {
enum class MMAMode {
kInit = 0U,
kInplaceUpdate = 1U,
};
};
/*!
/*!
* \brief Wrapper of PTX ldmatrix m8n8.x2 instruction, loads data from shared memory
* to fragment
* \tparam T data type of the fragment
* \param R pointer to the fragment
* \param smem_ptr pointer to the shared memory
*/
template <typename T>
__device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
template <typename T>
__device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
#ifdef LDMATRIX_M8N8X2_ENABLED
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
@@ -84,17 +84,17 @@ __device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory
* to fragment
* \tparam T data type of the fragment
* \param R pointer to the fragment
* \param smem_ptr pointer to the shared memory
*/
template <typename T>
__device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
template <typename T>
__device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
#ifdef LDMATRIX_M8N8X4_ENABLED
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
@@ -103,17 +103,17 @@ __device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from
* shared memory to fragment and transposes the fragment
* \tparam T data type of the fragment
* \param R pointer to the fragment
* \param smem_ptr pointer to the shared memory
*/
template <typename T>
__device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) {
template <typename T>
__device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) {
#ifdef LDMATRIX_M8N8X4_ENABLED
uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
@@ -122,9 +122,9 @@ __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr)
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
* multiplication, accumulated in f32.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -132,13 +132,11 @@ __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr)
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
// ! only support half dtype now
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%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])
: "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]));
}
else if constexpr (mma_mode == MMAMode::kInit)
{
} else if constexpr (mma_mode == MMAMode::kInit) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0, %1, %2, %3},"
@@ -164,9 +160,9 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, ui
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
* multiplication, accumulated in f32.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -174,13 +170,11 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, ui
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
// ! only support half dtype now
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%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])
: "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]));
}
else if constexpr (mma_mode == MMAMode::kInit)
{
} else if constexpr (mma_mode == MMAMode::kInit) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
"{%0, %1, %2, %3},"
@@ -224,9 +216,9 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, u
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
* multiplication, accumulated in f16.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -234,12 +226,10 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, u
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F16F16F16_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1},"
@@ -248,9 +238,7 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C,
"{%8, %9};\n"
: "=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(
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1},"
@@ -263,9 +251,9 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C,
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
* multiplication, accumulated in f16.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -273,12 +261,10 @@ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C,
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F16F16F16_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1},"
@@ -295,9 +281,7 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C
"{%8, %9};\n"
: "=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(
"mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
"{%0, %1},"
@@ -318,9 +302,9 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n8k32 instruction for row major and column major int8 matrix
* multiplication, accumulated in int32.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -328,12 +312,10 @@ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_S8S8S32_M16N8K32_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%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"(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]));
}
else if constexpr (mma_mode == MMAMode::kInit)
{
} else if constexpr (mma_mode == MMAMode::kInit) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%0, %1, %2, %3},"
@@ -359,9 +339,9 @@ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, ui
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n16k32 instruction for row major and column major int8 matrix
* multiplication, accumulated in int32.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -369,12 +349,10 @@ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, ui
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_S8S8S32_M16N8K32_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%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"(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]));
}
else if constexpr (mma_mode == MMAMode::kInit)
{
} else if constexpr (mma_mode == MMAMode::kInit) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
"{%0, %1, %2, %3},"
@@ -418,9 +394,9 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, u
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n8k32 instruction for row major and column major int4 matrix
* multiplication, accumulated in int32.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -428,12 +404,10 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, u
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_S4S4S32_M16N8K64_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
"{%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"(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]));
}
else if constexpr (mma_mode == MMAMode::kInit)
{
} else if constexpr (mma_mode == MMAMode::kInit) {
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
"{%0, %1, %2, %3},"
@@ -459,9 +431,9 @@ __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, ui
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n16k64 instruction for row major and column major int4 matrix
* multiplication, accumulated in int32.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -469,12 +441,10 @@ __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, ui
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_S4S4S32_M16N8K64_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
"{%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"(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]));
}
else if constexpr (mma_mode == MMAMode::kInit)
{
} else if constexpr (mma_mode == MMAMode::kInit) {
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
"{%0, %1, %2, %3},"
@@ -518,9 +486,9 @@ __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, u
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n8k32 instruction for row major and column major fp8 e4m3 matrix
* multiplication, accumulated in fp32.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -528,12 +496,10 @@ __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, u
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%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])
: "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]));
}
else if constexpr (mma_mode == MMAMode::kInit)
{
} else if constexpr (mma_mode == MMAMode::kInit) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0, %1, %2, %3},"
@@ -559,9 +523,9 @@ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
* multiplication, accumulated in fp16.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -569,13 +533,11 @@ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A,
uint32_t* B) {
//uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, uint32_t* B) {
// uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
#ifdef MMA_F8F8F16_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
"{%0, %1},"
@@ -593,9 +555,7 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_u
"{%8, %9};\n"
: "=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(
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
"{%0, %1},"
@@ -617,11 +577,9 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_u
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
* multiplication, accumulated in fp32.
* \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
@@ -629,12 +587,10 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_u
* \param A pointer to the fragment of matrix A
* \param B pointer to the fragment of matrix B
*/
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A,
uint32_t* B) {
template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
if constexpr (mma_mode == MMAMode::kInplaceUpdate)
{
if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%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])
: "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]));
}
else if constexpr (mma_mode == MMAMode::kInit)
{
} else if constexpr (mma_mode == MMAMode::kInit) {
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
"{%0, %1, %2, %3},"
@@ -680,12 +634,12 @@ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uin
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Use mma instructions to compute rowsum.
*/
__device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) {
__device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) {
#ifdef MMA_F16F16F32_M16N8K16_ENABLED
asm volatile(
"{\n"
@@ -701,12 +655,12 @@ __device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) {
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
/*!
/*!
* \brief Use mma instructions to compute rowsum.
*/
__device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) {
__device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) {
#ifdef MMA_F8F8F32_M16N8K16_ENABLED
asm volatile(
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
@@ -720,9 +674,9 @@ __device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) {
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
#endif
}
}
} // namespace mma
} // namespace mma
} // namespace sageattention
src/core/cuda/sageattention/numeric_conversion.cuh
+60 -67
View File
@@ -38,117 +38,110 @@ namespace sageattention {
#define RUNTIME_ASSERT(x) assert(0 && x)
#endif
__device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) {
__device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) {
uint16_t h0 = source & 0xFFFF;
uint16_t h1 = (source >> 16) & 0xFFFF;
asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0));
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
asm volatile( \
"{\n" \
".reg .b16 lo;\n" \
".reg .b16 hi;\n" \
"cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n" \
"cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n" \
"mov.b32 %0, {lo, hi};\n" \
"}" \
asm volatile(
"{\n"
".reg .b16 lo;\n"
".reg .b16 hi;\n"
"cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n"
"cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n"
"mov.b32 %0, {lo, hi};\n"
"}"
: "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#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
asm volatile( \
"{\n" \
".reg .b16 lo;\n" \
".reg .b16 hi;\n" \
"cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n" \
"cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n" \
"mov.b32 %0, {lo, hi};\n" \
"}" \
asm volatile(
"{\n"
".reg .b16 lo;\n"
".reg .b16 hi;\n"
"cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n"
"cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n"
"mov.b32 %0, {lo, hi};\n"
"}"
: "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#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
asm volatile( \
"{\n" \
".reg .b16 lo;\n" \
".reg .b16 hi;\n" \
"cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n" \
"cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n" \
"mov.b32 %0, {lo, hi};\n" \
"}" \
asm volatile(
"{\n"
".reg .b16 lo;\n"
".reg .b16 hi;\n"
"cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n"
"cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n"
"mov.b32 %0, {lo, hi};\n"
"}"
: "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#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
asm volatile( \
"{\n" \
".reg .b16 lo;\n" \
".reg .b16 hi;\n" \
"cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n" \
"cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n" \
"mov.b32 %0, {lo, hi};\n" \
"}" \
asm volatile(
"{\n"
".reg .b16 lo;\n"
".reg .b16 hi;\n"
"cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n"
"cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n"
"mov.b32 %0, {lo, hi};\n"
"}"
: "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#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
asm volatile( \
"{\n" \
".reg .b16 lo, hi;\n" \
"mov.b32 {lo, hi}, %2;\n" \
"cvt.rn.f16x2.e4m3x2 %0, lo;\n" \
"cvt.rn.f16x2.e4m3x2 %1, hi;\n" \
asm volatile(
"{\n"
".reg .b16 lo, hi;\n"
"mov.b32 {lo, hi}, %2;\n"
"cvt.rn.f16x2.e4m3x2 %0, lo;\n"
"cvt.rn.f16x2.e4m3x2 %1, hi;\n"
"}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#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
asm volatile( \
"{\n" \
".reg .b16 lo, hi;\n" \
"mov.b32 {lo, hi}, %2;\n" \
"cvt.rn.f16x2.e5m2x2 %0, lo;\n" \
"cvt.rn.f16x2.e5m2x2 %1, hi;\n" \
asm volatile(
"{\n"
".reg .b16 lo, hi;\n"
"mov.b32 {lo, hi}, %2;\n"
"cvt.rn.f16x2.e5m2x2 %0, lo;\n"
"cvt.rn.f16x2.e5m2x2 %1, hi;\n"
"}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
#else
RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
#endif
}
}
__device__ __forceinline__ int8_t float_to_int8_rn(float x)
{
__device__ __forceinline__ int8_t float_to_int8_rn(float x) {
uint32_t dst;
asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
return reinterpret_cast<const int8_t&>(dst);
}
}
} // namespace sageattention
src/core/cuda/sageattention/permuted_smem.cuh
+18 -19
View File
@@ -23,22 +23,21 @@
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include "cp_async.cuh"
#include "mma.cuh"
namespace sageattention {
enum class SwizzleMode {
enum class SwizzleMode {
k32B, // for k32B mode, a line of shared memory must have 32B (16 half value)
k64B, // for k64B mode, a line of shared memory must have 64B (32 half value)
k128B, // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B.
};
};
// Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth
using b128_t = uint4;
// Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth
using b128_t = uint4;
/*!
/*!
* \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure
* that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through
* permutation in the granularity of 128bit.
@@ -51,8 +50,8 @@ using b128_t = uint4;
* We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp.
* Under this assumption, we do not permute for k32B mode.
*/
template <SwizzleMode swizzle_mode, uint32_t stride>
struct smem_t {
template <SwizzleMode swizzle_mode, uint32_t stride>
struct smem_t {
// The base pointer.
b128_t* base;
// How many b128_t value a row contains
@@ -85,7 +84,7 @@ struct smem_t {
* \param i The row index.
* \param j The column index.
*/
static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t &i, const uint32_t &j) {
static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t& i, const uint32_t& j) {
if constexpr (swizzle_mode == SwizzleMode::k128B) {
return i * stride + (j ^ (i % 8));
} else if constexpr (swizzle_mode == SwizzleMode::k64B) {
@@ -100,7 +99,7 @@ struct smem_t {
* \param offset The current offset.
*/
template <uint32_t step_size>
static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset) {
static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset) {
if constexpr (swizzle_mode == SwizzleMode::k128B) {
static_assert(step_size % 8 == 0,
"Unsupported step size");
@@ -116,7 +115,7 @@ struct smem_t {
// ! use with care
template <uint32_t step_size>
static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t &offset, const uint32_t &step_idx) {
static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset, const uint32_t& step_idx) {
if constexpr (swizzle_mode == SwizzleMode::k128B) {
static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0,
"Unsupported step size");
@@ -141,7 +140,7 @@ struct smem_t {
}
template <uint32_t step_size>
static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t &offset) {
static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t& offset) {
if constexpr (swizzle_mode == SwizzleMode::k128B) {
static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
if constexpr (step_size == 4) {
@@ -163,39 +162,39 @@ struct smem_t {
}
}
__device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t &offset, uint32_t* R) const {
__device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t& offset, uint32_t* R) const {
b128_t* smem_ptr = base + offset;
mma::ldmatrix_m8n8x2(R, smem_ptr);
}
__device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t &offset, uint32_t* R) const {
__device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t& offset, uint32_t* R) const {
b128_t* smem_ptr = base + offset;
mma::ldmatrix_m8n8x4(R, smem_ptr);
}
__device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t &offset, uint32_t* R) const {
__device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t& offset, uint32_t* R) const {
b128_t* smem_ptr = base + offset;
mma::ldmatrix_m8n8x4_trans(R, smem_ptr);
}
template <cp_async::SharedMemFillMode fill_mode, typename T>
__device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr, bool predicate) const {
__device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr, bool predicate) const {
b128_t* smem_ptr = base + offset;
cp_async::pred_load_128b<cp_async::PrefetchMode::kPrefetch, fill_mode>(
smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
}
template <typename T>
__device__ __forceinline__ void load_128b_async(const uint32_t &offset, const T* gptr) const {
__device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr) const {
b128_t* smem_ptr = base + offset;
cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr));
}
template <typename T>
__device__ __forceinline__ void store_128b(const uint32_t &offset, T* gptr) const {
__device__ __forceinline__ void store_128b(const uint32_t& offset, T* gptr) const {
*reinterpret_cast<b128_t*>(gptr) = *(base + offset);
}
};
};
} // namespace sageattention
src/core/ggml_extend.h
+1 -1
View File
@@ -4,9 +4,9 @@
#include <cstdint>
#include <vector>
#include "core/ggml_extensions.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "core/ggml_extensions.h"
#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 = 1) const bool AFFINE = false;
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 = 1) readonly buffer M { 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[]; };
layout(binding = 0) readonly buffer X {
float data_x[];
};
layout(binding = 1) readonly buffer M {
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
layout(binding = 5) readonly buffer W { float data_w[]; };
layout(binding = 5) readonly buffer W {
float data_w[];
};
#endif
layout(push_constant) uniform parameter {
uint rows;
@@ -21,7 +33,8 @@ layout(push_constant) uniform parameter {
float eps;
uint k;
uint w_offset;
} p;
}
p;
shared float values[K];
shared vec2 sums[512];
@@ -80,7 +93,8 @@ void main() {
sums[tid] = sum;
barrier();
[[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();
}
mean = sums[0].x / p.k;
@@ -139,10 +153,7 @@ void main() {
data_s[row] = scale;
}
for (uint i = tid; i < K / 4; i += 512) {
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);
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);
}
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(push_constant) uniform parameter
{
layout(push_constant) uniform parameter {
uint n;
uint rows;
uint k;
uint has_bias;
uint row_offset;
} p;
}
p;
layout(binding = 0, std430) readonly buffer W { int data_w[]; };
layout(binding = 1, std430) readonly buffer A { int 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[]; };
layout(binding = 5, std430) writeonly buffer D { float data_d[]; };
layout(binding = 0, std430) readonly buffer W {
int data_w[];
};
layout(binding = 1, std430) readonly buffer A {
int 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[];
};
layout(binding = 5, std430) writeonly buffer D {
float data_d[];
};
shared int weights[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(push_constant) uniform parameter
{
layout(push_constant) uniform parameter {
uint n;
uint rows;
uint k;
uint has_bias;
uint row_offset;
} p;
}
p;
layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; };
layout(binding = 1, std430) readonly buffer A { 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[]; };
layout(binding = 5, std430) writeonly buffer D { float data_d[]; };
layout(binding = 0, std430) readonly buffer W {
int8_t data_w[];
};
layout(binding = 1, std430) readonly buffer A {
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[];
};
layout(binding = 5, std430) writeonly buffer D {
float data_d[];
};
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 = 4) const uint BK = 64;
layout(push_constant) uniform parameter
{
layout(push_constant) uniform parameter {
uint n;
uint rows;
uint k;
@@ -23,25 +22,45 @@ layout(push_constant) uniform parameter
uint row_offset;
float output_scale;
uint group_cols;
} p;
}
p;
layout(binding = 0, std430) readonly buffer W { int8_t data_w[]; };
layout(binding = 1, std430) readonly buffer A { 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[]; };
layout(binding = 0, std430) readonly buffer W {
int8_t data_w[];
};
layout(binding = 1, std430) readonly buffer A {
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.
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
layout(binding = 6, std430) readonly buffer UW { int8_t data_uw[]; };
layout(binding = 7, std430) readonly buffer US { float data_us[]; };
layout(binding = 6, std430) readonly buffer UW {
int8_t data_uw[];
};
layout(binding = 7, std430) readonly buffer US {
float data_us[];
};
#endif
#ifdef RESIDUAL
layout(binding = 6, std430) readonly buffer G { float data_g[]; };
layout(binding = 7, std430) readonly buffer R { float data_r[]; };
layout(binding = 6, std430) readonly buffer G {
float data_g[];
};
layout(binding = 7, std430) readonly buffer R {
float data_r[];
};
#endif
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) {
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 uint r = row_base + col;
const uint c = col_base + row;
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(push_constant) uniform parameter
{
layout(push_constant) uniform parameter {
uint k;
uint rows;
} p;
}
p;
layout(binding = 0, std430) readonly buffer A { float data_a[]; };
layout(binding = 1, std430) writeonly buffer D { uint data_d[]; };
layout(binding = 2, std430) writeonly buffer S { float data_s[]; };
layout(binding = 0, std430) readonly buffer A {
float data_a[];
};
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 maxima[64];
@@ -86,10 +92,7 @@ void quantize_row(const uint row) {
for (uint group = 0; group < groups; ++group) {
transform_group(row, group);
const uint index = tid * 4;
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);
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);
data_d[(row * p.k + group * 256) / 4 + tid] = packed;
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
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(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[]; };
layout(push_constant) uniform parameter {
uint k;
uint rows;
}
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)
layout(binding = 3, std430) readonly buffer B { float data_b[]; };
layout(binding = 3, std430) readonly buffer B {
float data_b[];
};
#endif
shared float maxima[8];
#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) {
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
float values[((K + 2047) / 2048) * 8];
[[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) {
if (gl_SubgroupID + g * 8 >= K / 256) continue;
float v[8]; transform_group(row, gl_SubgroupID + g * 8, v);
if (gl_SubgroupID + g * 8 >= K / 256)
continue;
float v[8];
transform_group(row, gl_SubgroupID + g * 8, v);
[[unroll]] for (uint j = 0; j < 8; ++j) {
values[g * 8 + j] = v[j];
local_max = max(local_max, abs(v[j]));
@@ -81,16 +98,20 @@ void main() {
}
#else
for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) {
float v[8]; transform_group(row, group, v);
[[unroll]] for (uint j = 0; j < 8; ++j) local_max = max(local_max, abs(v[j]));
float v[8];
transform_group(row, group, v);
[[unroll]] for (uint j = 0; j < 8; ++j)
local_max = max(local_max, abs(v[j]));
}
#endif
local_max = subgroupMax(local_max);
if (gl_SubgroupInvocationID == 0) maxima[gl_SubgroupID] = local_max;
if (gl_SubgroupInvocationID == 0)
maxima[gl_SubgroupID] = local_max;
barrier();
if (gl_LocalInvocationID.x == 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;
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;
#ifdef SINGLE_PASS
[[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) {
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
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) {
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
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 = 1) readonly buffer PE { float data_pe[]; };
layout(binding = 0) readonly buffer X {
float data_x[];
};
layout(binding = 1) readonly buffer PE {
float data_pe[];
};
#ifdef CAST_F16
#extension GL_EXT_shader_16bit_storage : 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
layout(binding = 2) writeonly buffer D { float data_d[]; };
layout(binding = 2) writeonly buffer D {
float data_d[];
};
#endif
layout(binding = 3) readonly buffer W { float data_w[]; };
layout(binding = 3) readonly buffer W {
float data_w[];
};
layout(push_constant) uniform parameter {
uint half_dim;
@@ -31,7 +41,8 @@ layout(push_constant) uniform parameter {
float eps;
uint w_offset;
float output_scale;
} pc;
}
pc;
float shuffle_channel(vec4 values, uint channel) {
const uint lane = channel & 31u;
src/core/vulkan/shaders/rope_apply.comp
+24 -17
View File
@@ -2,9 +2,15 @@
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 = 1) readonly buffer PE { float data_pe[]; };
layout(binding = 2) writeonly buffer D { float data_d[]; };
layout(binding = 0) readonly buffer X {
float data_x[];
};
layout(binding = 1) readonly buffer PE {
float data_pe[];
};
layout(binding = 2) writeonly buffer D {
float data_d[];
};
layout(push_constant) uniform parameter {
uint half_dim;
@@ -18,10 +24,11 @@ layout(push_constant) uniform parameter {
uint pe_offset;
uint dst_offset;
uint interleaved;
} pc;
}
pc;
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;
if (index >= pc.pairs) {
return;
}
@@ -29,30 +36,30 @@ void main() {
const uint row = index / pc.half_dim;
const uint token = row % pc.tokens;
const uint head = (row / pc.tokens) % pc.heads;
const uint batch = row / (pc.tokens*pc.heads);
const uint i0 = pc.interleaved != 0 || j >= pc.rot_half ? 2*j : j;
const uint i1 = pc.interleaved != 0 || j >= pc.rot_half ? 2*j + 1 : j + pc.rot_half;
const uint x_row = pc.x_offset + head*pc.x_stride[1] + token*pc.x_stride[2] + batch*pc.x_stride[3];
const uint out_row = pc.dst_offset + row*pc.half_dim*2;
const float a = data_x[x_row + i0*pc.x_stride[0]];
const float b = data_x[x_row + i1*pc.x_stride[0]];
const uint batch = row / (pc.tokens * pc.heads);
const uint i0 = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j : j;
const uint i1 = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j + 1 : j + pc.rot_half;
const uint x_row = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3];
const uint out_row = pc.dst_offset + row * pc.half_dim * 2;
const float a = data_x[x_row + i0 * pc.x_stride[0]];
const float b = data_x[x_row + i1 * pc.x_stride[0]];
if (j >= pc.rot_half) {
data_d[out_row + i0] = a;
data_d[out_row + i1] = b;
return;
}
const uint matrix = pc.pe_offset + j*pc.pe_stride[2] + token*pc.pe_stride[3];
const uint matrix = pc.pe_offset + j * pc.pe_stride[2] + token * pc.pe_stride[3];
const float c = data_pe[matrix];
const float s = data_pe[matrix + pc.pe_stride[1]];
// Preserve the separate multiply/add rounding of the unfused graph.
precise float y0;
precise float y1;
if (pc.interleaved != 0) {
y0 = a*c + b*data_pe[matrix + pc.pe_stride[0]];
y1 = a*s + b*data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]];
y0 = a * c + b * data_pe[matrix + pc.pe_stride[0]];
y1 = a * s + b * data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]];
} else {
y0 = a*c - b*s;
y1 = a*s + b*c;
y0 = a * c - b * s;
y1 = a * s + b * c;
}
data_d[out_row + i0] = y0;
data_d[out_row + i1] = y1;