daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
h3-blackwell-runtime/research/sage2_temporal_pair/patches/0001-sm89-temporal-two-pair-qk-pv.patch
Daniel Maddern ea14b98544 Preserve Vortex performance research history
2026-08-25 20:30:22 +07:00

270 lines
12 KiB
Diff
Raw Blame History

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<uint32_t CTA_K, uint32_t num_warps_q, uint32_t num_warps_k,
+ uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, uint32_t num_tiles_v,
+ SwizzleMode swizzle_mode_QK, uint32_t qk_stride, DataType DTypeQK,
+ typename DTypeSVAccum, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode>
+__device__ __forceinline__ void temporal_pair_produce_score(
+ const smem_t<swizzle_mode_QK, qk_stride> &smem_Q,
+ const smem_t<swizzle_mode_QK, qk_stride> &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner,
+ swizzle_mode_QK, qk_stride, DTypeQK>(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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner,
+ swizzle_mode_QK, qk_stride, DTypeQK>(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<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+ }
+ if (tile + 1 == num_iterations)
+ {
+ 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;
+
+ const float tile_sm_scale = final_pair ? original_sm_scale : original_sm_scale * dequant_scale;
+ 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, tile_sm_scale);
+ }
+ else
+ {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, tile_sm_scale);
+ }
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore)
+ {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+ }
+ RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore)
+ {
+ accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+ }
+}
+
+template<uint32_t num_warps_q, uint32_t num_warps_k,
+ uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v,
+ SwizzleMode swizzle_mode_V, uint32_t v_stride,
+ typename DTypeSVAccum, bool use_inst_buffer, bool use_pv_fp16_accu>
+__device__ __forceinline__ void temporal_pair_consume_value(
+ const smem_t<swizzle_mode_V, v_stride> &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v,
+ swizzle_mode_V, v_stride>(smem_V, RS_f8, RO, d);
+ }
+ 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_stride>(smem_V, RS_f8, RO, d);
+ }
+ 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_stride>(smem_V, RS_f8, RO, d);
+ }
+}
+
template<uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN,
- typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale=false, bool fuse_v_mean=false, bool use_pv_fp16_accu=false>
+ typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale=false, bool fuse_v_mean=false, bool use_pv_fp16_accu=false, 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<CTA_K, num_warps_q, num_warps_k, num_tiles_q, num_tiles_k,
+ num_tiles_qk_inner, num_tiles_v, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK,
+ DTypeQK, DTypeSVAccum, DenominatorAccumUnit, mask_mode>(
+ 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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k,
+ num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V,
+ DTypeSVAccum, use_inst_buffer, use_pv_fp16_accu>(smem_V, RS_f8, RO, d);
+ }
+ else
+ {
+ temporal_pair_produce_score<CTA_K, num_warps_q, num_warps_k, num_tiles_q, num_tiles_k,
+ num_tiles_qk_inner, num_tiles_v, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK,
+ DTypeQK, DTypeSVAccum, DenominatorAccumUnit, mask_mode>(
+ 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<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK,
+ QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK,
+ QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+ &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<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK,
+ QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK,
+ QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+ &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<CTA_K, num_warps_q, num_warps_k, num_tiles_q, num_tiles_k,
+ num_tiles_qk_inner, num_tiles_v, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK,
+ DTypeQK, DTypeSVAccum, DenominatorAccumUnit, mask_mode>(
+ 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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k,
+ num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V,
+ DTypeSVAccum, use_inst_buffer, use_pv_fp16_accu>(smem_V, RS_f8, RO, d);
+ }
+ __syncthreads();
+
+ load_fp8_V_global_to_share<global_to_shared_line_lanes_V,
+ global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col,
+ swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+ &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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k,
+ num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V,
+ DTypeSVAccum, use_inst_buffer, use_pv_fp16_accu>(smem_V, RS_f8, RO, d);
+ }
+ else
+ {
+ temporal_pair_produce_score<CTA_K, num_warps_q, num_warps_k, num_tiles_q, num_tiles_k,
+ num_tiles_qk_inner, num_tiles_v, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK,
+ DTypeQK, DTypeSVAccum, DenominatorAccumUnit, mask_mode>(
+ 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<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k,
+ num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V,
+ DTypeSVAccum, use_inst_buffer, use_pv_fp16_accu>(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<CTA_Q, CTA_K, WARP_Q, WARP_K, HEAD_DIM, DataType::kInt8, static_cast<QuantGranularity>(QK_QUANT_GRAN), static_cast<QuantGranularity>(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
+}
Reference in a new issue View git blame Copy permalink