chore: fotmat code

2026-11-12 11:56:36 +00:00 · 2026-11-01 13:40:36 +08:00
1 parent 6edeac1b20
commit 5531c8878a
@@ -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,
+    template <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_mean = false>
-        typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_mean=false>
+    __global__ void qk_int_sv_f16_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, half* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, DTypeOut* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
 __global__ void qk_int_sv_f16_attn_kernel(int8_t *__restrict__ Q, int8_t *__restrict__ K, half *__restrict__ V, DTypeOut *__restrict__ O, float *__restrict__ Lse,
                      float *__restrict__ Q_scale, float *__restrict__ K_scale, DTypeOut *__restrict__ V_mean,
                      const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups,
                      const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q,
                      const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k,
                      const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v,
                      const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o,
                      float sm_scale)
 {
        // Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes.
 #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++)
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
-    {
+                if constexpr (std::is_same<DTypeSVAccum, float>::value) {
      if constexpr (std::is_same<DTypeSVAccum, float>::value)
      {
 #pragma unroll
-        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_line_lanes_QK          = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
-  constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4;
+                                                                                                                         : 8;
        constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
                                                                                                                          : 4;
        constexpr uint32_t global_to_shared_line_lanes_V           = (V_SMEM_STRIDE == 32) ? 4 : 8;
        constexpr uint32_t global_to_shared_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* 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* 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;
+        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++)
+                    for (uint32_t k = 0; k < 4; k++) {
-        {
+                        if constexpr (std::is_same<DTypeOut, half>::value) {
          if constexpr (std::is_same<DTypeOut, half>::value)
          {
                            ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
-          }
+                        } 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++)
+            for (uint32_t j = 0; j < O_smem_iters_row; j++) {
-    {
+                if (O_load_idx_lane_base < qo_len) {
      if (O_load_idx_lane_base < qo_len)
      {
                    smem_O.store_128b(offset_O, O_lane_ptr);
                }
                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
@@ -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,
+    template <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale = false, bool fuse_v_mean = false, bool use_pv_fp16_accu = false>
-        typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale=false, bool fuse_v_mean=false, bool use_pv_fp16_accu=false>
+    __global__ void qk_int_sv_f8_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, int8_t* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, float* __restrict__ V_scale, float* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
 __global__ void qk_int_sv_f8_attn_kernel(int8_t *__restrict__ Q, int8_t *__restrict__ K, int8_t *__restrict__ V, DTypeOut *__restrict__ O, float *__restrict__ Lse,
                      float *__restrict__ Q_scale, float *__restrict__ K_scale, float *__restrict__ V_scale, float *__restrict__ V_mean,
                      const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups,
                      const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q,
                      const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k,
                      const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v,
                      const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o,
                      float sm_scale)
 {
        // Keep unsupported device passes from parsing architecture-specific intrinsics.
 #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++)
+            for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
-    {
+                if constexpr (std::is_same<DTypeSVAccum, float>::value) {
      if constexpr (std::is_same<DTypeSVAccum, float>::value)
      {
 #pragma unroll
-        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_line_lanes_QK          = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
-  constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8 : 4;
+                                                                                                                         : 8;
        constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
                                                                                                                          : 4;
        //                                                         for fp16: 32
        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* 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* 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 {
-    else
+                if constexpr (!use_pv_fp16_accu) {
    {
      if constexpr (!use_pv_fp16_accu){
                    compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
                        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 {
-    else
+                if constexpr (!use_pv_fp16_accu) {
    {
      if constexpr (!use_pv_fp16_accu){
                    compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
                        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 {
-    else
+                if constexpr (!use_pv_fp16_accu) {
    {
      if constexpr (!use_pv_fp16_accu){
                    compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
                        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++)
+                    for (uint32_t k = 0; k < 4; k++) {
-        {
+                        if constexpr (std::is_same<DTypeOut, half>::value) {
          if constexpr (std::is_same<DTypeOut, half>::value)
          {
                            ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
-          }
+                        } 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++)
+            for (uint32_t j = 0; j < O_smem_iters_row; j++) {
-    {
+                if (O_load_idx_lane_base < qo_len) {
      if (O_load_idx_lane_base < qo_len)
      {
                    smem_O.store_128b(offset_O, O_lane_ptr);
                }
                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
@@ -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>
+        template <size_t n>
-__device__ __forceinline__ void wait_group() {
+        __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>
+        template <PrefetchMode prefetch_mode, typename T>
-__device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
+        __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>
+        template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
-__device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
+        __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
@@ -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)
+        // log2(e)
-constexpr float log2e = 1.44269504088896340736f;
+        constexpr float log2e      = 1.44269504088896340736f;
-constexpr float log2e_recp = 1.0f / log2e;
+        constexpr float log2e_recp = 1.0f / log2e;
-__forceinline__ __device__ half2 uint32_as_half2(uint32_t x) { return *(half2*)&x; }
+        __forceinline__ __device__ half2 uint32_as_half2(uint32_t x) {
            return *(half2*)&x;
        }
-__forceinline__ __device__ uint32_t half2_as_uint32(half2 x) { return *(uint32_t*)&x; }
+        __forceinline__ __device__ uint32_t half2_as_uint32(half2 x) {
            return *(uint32_t*)&x;
        }
-/*!
+        /*!
         * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x
         * \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
@@ -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>
+        template <typename T>
-__device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
+        __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>
+        template <typename T>
-__device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
+        __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>
+        template <typename T>
-__device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) {
+        __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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
                                                                     uint32_t* B) {
 #ifdef MMA_F16F16F32_M16N8K16_ENABLED
            // ! 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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
                                                                     uint32_t* B) {
 #ifdef MMA_F16F16F32_M16N8K16_ENABLED
            // ! 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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
                                                                     uint32_t* B) {
 #ifdef MMA_F16F16F16_M16N8K16_ENABLED
-  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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
                                                                     uint32_t* B) {
 #ifdef MMA_F16F16F16_M16N8K16_ENABLED
-  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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
                                                                   uint32_t* B) {
 #ifdef MMA_S8S8S32_M16N8K32_ENABLED
-  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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
                                                                   uint32_t* B) {
 #ifdef MMA_S8S8S32_M16N8K32_ENABLED
-  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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
                                                                   uint32_t* B) {
 #ifdef MMA_S4S4S32_M16N8K64_ENABLED
-  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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
                                                                   uint32_t* B) {
 #ifdef MMA_S4S4S32_M16N8K64_ENABLED
-  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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
                                                                   uint32_t* B) {
 #ifdef MMA_F8F8F32_M16N8K16_ENABLED
-  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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, uint32_t* B) {
-                                                                   uint32_t* B) {
+            // uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
 //uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
 #ifdef MMA_F8F8F16_M16N8K16_ENABLED
-  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>
+        template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
-__device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A,
+        __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
                                                                   uint32_t* B) {
 #ifdef MMA_F8F8F32_M16N8K16_ENABLED
-  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
@@ -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( \
+        asm volatile(
-      "{\n" \
+            "{\n"
-      ".reg .b16 lo;\n" \
+            ".reg .b16 lo;\n"
-      ".reg .b16 hi;\n" \
+            ".reg .b16 hi;\n"
-      "cvt.rn.satfinite.e4m3x2.f32   lo, %2, %1;\n" \
+            "cvt.rn.satfinite.e4m3x2.f32   lo, %2, %1;\n"
-      "cvt.rn.satfinite.e4m3x2.f32   hi, %4, %3;\n" \
+            "cvt.rn.satfinite.e4m3x2.f32   hi, %4, %3;\n"
-      "mov.b32 %0, {lo, hi};\n" \
+            "mov.b32 %0, {lo, hi};\n"
-      "}" \
+            "}"
            : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
 #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( \
+        asm volatile(
-      "{\n" \
+            "{\n"
-      ".reg .b16 lo;\n" \
+            ".reg .b16 lo;\n"
-      ".reg .b16 hi;\n" \
+            ".reg .b16 hi;\n"
-      "cvt.rn.satfinite.e5m2x2.f32   lo, %2, %1;\n" \
+            "cvt.rn.satfinite.e5m2x2.f32   lo, %2, %1;\n"
-      "cvt.rn.satfinite.e5m2x2.f32   hi, %4, %3;\n" \
+            "cvt.rn.satfinite.e5m2x2.f32   hi, %4, %3;\n"
-      "mov.b32 %0, {lo, hi};\n" \
+            "mov.b32 %0, {lo, hi};\n"
-      "}" \
+            "}"
            : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
 #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( \
+        asm volatile(
-      "{\n" \
+            "{\n"
-      ".reg .b16 lo;\n" \
+            ".reg .b16 lo;\n"
-      ".reg .b16 hi;\n" \
+            ".reg .b16 hi;\n"
-      "cvt.rn.satfinite.e4m3x2.f16x2   lo, %1;\n" \
+            "cvt.rn.satfinite.e4m3x2.f16x2   lo, %1;\n"
-      "cvt.rn.satfinite.e4m3x2.f16x2   hi, %2;\n" \
+            "cvt.rn.satfinite.e4m3x2.f16x2   hi, %2;\n"
-      "mov.b32 %0, {lo, hi};\n" \
+            "mov.b32 %0, {lo, hi};\n"
-      "}" \
+            "}"
            : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
 #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( \
+        asm volatile(
-      "{\n" \
+            "{\n"
-      ".reg .b16 lo;\n" \
+            ".reg .b16 lo;\n"
-      ".reg .b16 hi;\n" \
+            ".reg .b16 hi;\n"
-      "cvt.rn.satfinite.e5m2x2.f16x2   lo, %1;\n" \
+            "cvt.rn.satfinite.e5m2x2.f16x2   lo, %1;\n"
-      "cvt.rn.satfinite.e5m2x2.f16x2   hi, %2;\n" \
+            "cvt.rn.satfinite.e5m2x2.f16x2   hi, %2;\n"
-      "mov.b32 %0, {lo, hi};\n" \
+            "mov.b32 %0, {lo, hi};\n"
-      "}" \
+            "}"
            : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
 #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( \
+        asm volatile(
-      "{\n" \
+            "{\n"
-      ".reg .b16 lo, hi;\n" \
+            ".reg .b16 lo, hi;\n"
-      "mov.b32 {lo, hi}, %2;\n" \
+            "mov.b32 {lo, hi}, %2;\n"
-      "cvt.rn.f16x2.e4m3x2 %0, lo;\n" \
+            "cvt.rn.f16x2.e4m3x2 %0, lo;\n"
-      "cvt.rn.f16x2.e4m3x2 %1, hi;\n" \
+            "cvt.rn.f16x2.e4m3x2 %1, hi;\n"
            "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
 #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( \
+        asm volatile(
-      "{\n" \
+            "{\n"
-      ".reg .b16 lo, hi;\n" \
+            ".reg .b16 lo, hi;\n"
-      "mov.b32 {lo, hi}, %2;\n" \
+            "mov.b32 {lo, hi}, %2;\n"
-      "cvt.rn.f16x2.e5m2x2 %0, lo;\n" \
+            "cvt.rn.f16x2.e5m2x2 %0, lo;\n"
-      "cvt.rn.f16x2.e5m2x2 %1, hi;\n" \
+            "cvt.rn.f16x2.e5m2x2 %1, hi;\n"
            "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
 #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
@@ -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
+    // Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth
-using b128_t = uint4;
+    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>
+    template <SwizzleMode swizzle_mode, uint32_t stride>
-struct smem_t {
+    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
@@ -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
@@ -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 = 0) readonly buffer X {
-layout(binding = 1) readonly buffer M { float data_m[]; };
+    float data_x[];
-layout(binding = 2) readonly buffer H { float data_h[]; };
+};
-layout(binding = 3) writeonly buffer D { uint data_d[]; };
+layout(binding = 1) readonly buffer M {
-layout(binding = 4) writeonly buffer S { float data_s[]; };
+    float data_m[];
 };
 layout(binding = 2) readonly buffer H {
    float data_h[];
 };
 layout(binding = 3) writeonly buffer D {
    uint data_d[];
 };
 layout(binding = 4) writeonly buffer S {
    float data_s[];
 };
 #ifdef RMS_AFFINE
-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)
+            data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) | (quantize(values[i * 4 + 1], inv_scale) << 8) | (quantize(values[i * 4 + 2], inv_scale) << 16) | (quantize(values[i * 4 + 3], inv_scale) << 24);
                | (quantize(values[i * 4 + 1], inv_scale) << 8)
                | (quantize(values[i * 4 + 2], inv_scale) << 16)
                | (quantize(values[i * 4 + 3], inv_scale) << 24);
        }
        barrier();
    }
@@ -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 = 0, std430) readonly buffer W {
-layout(binding = 1, std430) readonly buffer A { int data_a[]; };
+    int data_w[];
-layout(binding = 2, std430) readonly buffer WS { float data_ws[]; };
+};
-layout(binding = 3, std430) readonly buffer AS { float data_as[]; };
+layout(binding = 1, std430) readonly buffer A {
-layout(binding = 4, std430) readonly buffer B { float data_b[]; };
+    int data_a[];
-layout(binding = 5, std430) writeonly buffer D { float data_d[]; };
+};
 layout(binding = 2, std430) readonly buffer WS {
    float data_ws[];
 };
 layout(binding = 3, std430) readonly buffer AS {
    float data_as[];
 };
 layout(binding = 4, std430) readonly buffer B {
    float data_b[];
 };
 layout(binding = 5, std430) writeonly buffer D {
    float data_d[];
 };
 shared int weights[16 * 8];
 shared int activations[16 * 8];
@@ -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 = 0, std430) readonly buffer W {
-layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; };
+    int8_t data_w[];
-layout(binding = 2, std430) readonly buffer WS { float data_ws[]; };
+};
-layout(binding = 3, std430) readonly buffer AS { float data_as[]; };
+layout(binding = 1, std430) readonly buffer A {
-layout(binding = 4, std430) readonly buffer B { float data_b[]; };
+    int8_t data_a[];
-layout(binding = 5, std430) writeonly buffer D { float data_d[]; };
+};
 layout(binding = 2, std430) readonly buffer WS {
    float data_ws[];
 };
 layout(binding = 3, std430) readonly buffer AS {
    float data_as[];
 };
 layout(binding = 4, std430) readonly buffer B {
    float data_b[];
 };
 layout(binding = 5, std430) writeonly buffer D {
    float data_d[];
 };
 shared int32_t accumulators[1024];
@@ -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 = 0, std430) readonly buffer W {
-layout(binding = 1, std430) readonly buffer A { int8_t data_a[]; };
+    int8_t data_w[];
-layout(binding = 2, std430) readonly buffer WS { float data_ws[]; };
+};
-layout(binding = 3, std430) readonly buffer AS { float data_as[]; };
+layout(binding = 1, std430) readonly buffer A {
-layout(binding = 4, std430) readonly buffer B { float data_b[]; };
+    int8_t data_a[];
 };
 layout(binding = 2, std430) readonly buffer WS {
    float data_ws[];
 };
 layout(binding = 3, std430) readonly buffer AS {
    float data_as[];
 };
 layout(binding = 4, std430) readonly buffer B {
    float data_b[];
 };
 // Store FP32 bit patterns without changing the integer accumulator fragment layout.
-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 = 6, std430) readonly buffer UW {
-layout(binding = 7, std430) readonly buffer US { float data_us[]; };
+    int8_t data_uw[];
 };
 layout(binding = 7, std430) readonly buffer US {
    float data_us[];
 };
 #endif
 #ifdef RESIDUAL
-layout(binding = 6, std430) readonly buffer G { float data_g[]; };
+layout(binding = 6, std430) readonly buffer G {
-layout(binding = 7, std430) readonly buffer R { float data_r[]; };
+    float data_g[];
 };
 layout(binding = 7, std430) readonly buffer R {
    float data_r[];
 };
 #endif
-int32_t dequantize(const in uint row, const in uint col, const in int32_t acc,
+int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) {
                const in uint row_base, const in uint col_base) {
    const uint r = row_base + col;
    const uint c = col_base + row;
    if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) {
@@ -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 = 0, std430) readonly buffer A {
-layout(binding = 1, std430) writeonly buffer D { uint data_d[]; };
+    float data_a[];
-layout(binding = 2, std430) writeonly buffer S { float data_s[]; };
+};
 layout(binding = 1, std430) writeonly buffer D {
    uint data_d[];
 };
 layout(binding = 2, std430) writeonly buffer S {
    float data_s[];
 };
 shared float values[256];
 shared float 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)
+        const uint packed                           = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24);
                          | (quantize_value(values[index + 1], inv_scale) << 8)
                          | (quantize_value(values[index + 2], inv_scale) << 16)
                          | (quantize_value(values[index + 3], inv_scale) << 24);
        data_d[(row * p.k + group * 256) / 4 + tid] = packed;
        barrier();
    }
@@ -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(push_constant) uniform parameter {
-layout(binding = 0, std430) readonly buffer A { float data_a[]; };
+    uint k;
-layout(binding = 1, std430) writeonly buffer D { int8_t data_d[]; };
+    uint rows;
-layout(binding = 2, std430) writeonly buffer S { float data_s[]; };
+}
 p;
 layout(binding = 0, std430) readonly buffer A {
    float data_a[];
 };
 layout(binding = 1, std430) writeonly buffer D {
    int8_t data_d[];
 };
 layout(binding = 2, std430) writeonly buffer S {
    float data_s[];
 };
 #if defined(SWIGLU) || defined(SIGMOID_MUL)
-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;
+            if (gl_SubgroupID + g * 8 >= K / 256)
-            float v[8]; transform_group(row, gl_SubgroupID + g * 8, v);
+                continue;
            float v[8];
            transform_group(row, gl_SubgroupID + g * 8, v);
            [[unroll]] for (uint j = 0; j < 8; ++j) {
                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);
+            float v[8];
-            [[unroll]] for (uint j = 0; j < 8; ++j) local_max = max(local_max, abs(v[j]));
+            transform_group(row, group, v);
            [[unroll]] for (uint j = 0; j < 8; ++j)
                local_max = max(local_max, abs(v[j]));
        }
 #endif
        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));
            }
@@ -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 = 0) readonly buffer X {
-layout(binding = 1) readonly buffer PE { float data_pe[]; };
+    float data_x[];
 };
 layout(binding = 1) readonly buffer PE {
    float data_pe[];
 };
 #ifdef CAST_F16
 #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;
@@ -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 = 0) readonly buffer X {
-layout(binding = 1) readonly buffer PE { float data_pe[]; };
+    float data_x[];
-layout(binding = 2) writeonly buffer D { float data_d[]; };
+};
 layout(binding = 1) readonly buffer PE {
    float data_pe[];
 };
 layout(binding = 2) writeonly buffer D {
    float data_d[];
 };
 layout(push_constant) uniform parameter {
    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 batch   = row / (pc.tokens * pc.heads);
-    const uint i0    = pc.interleaved != 0 || j >= pc.rot_half ? 2*j : j;
+    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 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 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 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 a      = data_x[x_row + i0 * pc.x_stride[0]];
-    const float b = data_x[x_row + i1*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]];
+        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]];
+        y1 = a * s + b * data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]];
    } else {
-        y0 = a*c - b*s;
+        y0 = a * c - b * s;
-        y1 = a*s + b*c;
+        y1 = a * s + b * c;
    }
    data_d[out_row + i0] = y0;
    data_d[out_row + i1] = y1;