diff --git a/csrc/qattn/qk_int_sv_f8_cuda_sm89.cuh b/csrc/qattn/qk_int_sv_f8_cuda_sm89.cuh index f39719f..2aa334a 100644 --- a/csrc/qattn/qk_int_sv_f8_cuda_sm89.cuh +++ b/csrc/qattn/qk_int_sv_f8_cuda_sm89.cuh @@ -41,8 +41,111 @@ #define MMA_SV_N 16 #define MMA_SV_K 32 +template +__device__ __forceinline__ void temporal_pair_produce_score( + const smem_t &smem_Q, + const smem_t &smem_K, + int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], + const uint32_t Q_smem_offset_mma, const uint32_t K_smem_offset_mma, + DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], + uint32_t RS_f8[][num_tiles_k / 2][4], + const uint32_t tile, const uint32_t num_iterations, + const uint32_t Q_idx_lane_base, uint32_t &K_idx_lane_base, const uint32_t kv_len, + const float original_sm_scale, const float dequant_scale) +{ + if constexpr (num_tiles_qk_inner == 1) + { + compute_int_qk(smem_K, RS, RQ, K_smem_offset_mma); + } + else + { + uint32_t q_offset = Q_smem_offset_mma; + uint32_t k_offset = K_smem_offset_mma; + compute_int_qk(smem_Q, smem_K, RS, q_offset, k_offset); + } + + float RS_f32[num_tiles_q][num_tiles_k][8]; + const bool final_pair = tile + 2 >= num_iterations; +#pragma unroll + for (uint32_t fq = 0; fq < num_tiles_q; fq++) + { +#pragma unroll + for (uint32_t fk = 0; fk < num_tiles_k; fk++) + { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) + { + float score = __int2float_rz(RS[fq][fk][k]); + RS_f32[fq][fk][k] = final_pair ? score * dequant_scale : score; + } + } + } + + if (final_pair) + { + if constexpr (mask_mode == MaskMode::kCausal) + { + apply_causal_mask(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + if (tile + 1 == num_iterations) + { + apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + } + } + K_idx_lane_base += CTA_K; + + const float tile_sm_scale = final_pair ? original_sm_scale : original_sm_scale * dequant_scale; + if constexpr (std::is_same::value) + { + update_mdo(RS_f32, RO, m, d, tile_sm_scale); + } + else + { + update_mdo(RS_f32, RO, m, d, tile_sm_scale); + } + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) + { + accumulate_d(RS_f32, d); + } + RS_32_to_8(RS_f32, RS_f8); + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) + { + accumulate_d_f8(RS_f8, d); + } +} + +template +__device__ __forceinline__ void temporal_pair_consume_value( + const smem_t &smem_V, + uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) +{ + if constexpr (!use_inst_buffer) + { + compute_fp8_sv(smem_V, RS_f8, RO, d); + } + else if constexpr (!use_pv_fp16_accu) + { + compute_fp8_sv_inst_buf(smem_V, RS_f8, RO, d); + } + else + { + compute_fp8_sv_inst_buf_fp16_accu(smem_V, RS_f8, RO, d); + } +} + template + 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, bool temporal_pair_pipeline=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, @@ -266,6 +369,124 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr K_load_idx_lane_base += CTA_K; + if constexpr (temporal_pair_pipeline) + { + // K0 and V0 were issued by every warp above. The candidate keeps one score + // fragment per warp and temporally staggers the two warp pairs; no score, + // softmax state, or output accumulator crosses a warp boundary. + cp_async::wait_group<0>(); + __syncthreads(); + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + const bool pair_a = warp_id < (num_warps / 2); + + if (pair_a) + { + temporal_pair_produce_score( + smem_Q, smem_K, RS, RQ, Q_smem_offset_mma, K_smem_offset_mma, + RO, m, d, RS_f8, 0, num_iterations, Q_idx_lane_base, K_idx_lane_base, + kv_len, original_sm_scale, q_scale * K_scale[k_scale_idx]); + } + __syncthreads(); + + if (pair_a) + { + temporal_pair_consume_value(smem_V, RS_f8, RO, d); + } + else + { + temporal_pair_produce_score( + smem_Q, smem_K, RS, RQ, Q_smem_offset_mma, K_smem_offset_mma, + RO, m, d, RS_f8, 0, num_iterations, Q_idx_lane_base, K_idx_lane_base, + kv_len, original_sm_scale, q_scale * K_scale[k_scale_idx]); + } + __syncthreads(); + + for (uint32_t tile = 1; tile < num_iterations; tile++) + { + if (tile + 1 == num_iterations) + { + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, + K_load_idx_lane_base, kv_len); + } + else + { + load_global_to_share( + &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K); + } + cp_async::commit_group(); + cp_async::wait_group<0>(); + __syncthreads(); + + if (pair_a) + { + temporal_pair_produce_score( + smem_Q, smem_K, RS, RQ, Q_smem_offset_mma, K_smem_offset_mma, + RO, m, d, RS_f8, tile, num_iterations, Q_idx_lane_base, K_idx_lane_base, + kv_len, original_sm_scale, + q_scale * K_scale[k_scale_idx + tile * k_scale_advance_offset]); + } + else + { + temporal_pair_consume_value(smem_V, RS_f8, RO, d); + } + __syncthreads(); + + load_fp8_V_global_to_share( + &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V); + cp_async::commit_group(); + cp_async::wait_group<0>(); + __syncthreads(); + + if (pair_a) + { + temporal_pair_consume_value(smem_V, RS_f8, RO, d); + } + else + { + temporal_pair_produce_score( + smem_Q, smem_K, RS, RQ, Q_smem_offset_mma, K_smem_offset_mma, + RO, m, d, RS_f8, tile, num_iterations, Q_idx_lane_base, K_idx_lane_base, + kv_len, original_sm_scale, + q_scale * K_scale[k_scale_idx + tile * k_scale_advance_offset]); + } + __syncthreads(); + K_load_idx_lane_base += CTA_K; + } + + // Pair B deliberately carries exactly one score fragment out of the last + // phase. Drain it before the unchanged normalization and output epilogue. + if (!pair_a) + { + temporal_pair_consume_value(smem_V, RS_f8, RO, d); + } + __syncthreads(); + } + else + { #pragma unroll for (uint32_t iter = 1; iter < num_iterations - 1; iter++) { @@ -566,6 +787,7 @@ __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restr __syncthreads(); } + } // TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving. diff --git a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu index 727030f..3a72af1 100644 --- a/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu +++ b/csrc/qattn/sm89_qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf.cu @@ -152,7 +152,7 @@ torch::Tensor qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(torch::Tensor q size_t smem_max = std::max(CTA_Q * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t) + CTA_K * HEAD_DIM * sizeof(int8_t), CTA_Q * HEAD_DIM * sizeof(half)); auto kernel_func = qk_int_sv_f8_attn_kernel(QK_QUANT_GRAN), static_cast(QK_QUANT_GRAN), - float, true, DTypeOut, ComputeUnit::kCudaCore, mask_mode, RETURN_LSE, true, false, true>; + float, true, DTypeOut, ComputeUnit::kCudaCore, mask_mode, RETURN_LSE, true, false, true, true>; cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_max); @@ -184,4 +184,4 @@ torch::Tensor qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(torch::Tensor q }); return lse; -} \ No newline at end of file +}